diff options
Diffstat (limited to 'mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reshape.mlir')
| -rwxr-xr-x[-rw-r--r--] | mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reshape.mlir | 30 |
1 files changed, 20 insertions, 10 deletions
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reshape.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reshape.mlir index eed880866884..57d2a931fad5 100644..100755 --- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reshape.mlir +++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reshape.mlir @@ -32,11 +32,10 @@ module { return %0 : tensor<3x4xf64, #SparseMatrix> } -// TODO: make this work -// func.func @expand_sparse2sparse(%arg0: tensor<12xf64, #SparseVector>) -> tensor<3x4xf64, #SparseMatrix> { -// %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64, #SparseVector> into tensor<3x4xf64, #SparseMatrix> -// return %0 : tensor<3x4xf64, #SparseMatrix> -// } + func.func @expand_sparse2sparse(%arg0: tensor<12xf64, #SparseVector>) -> tensor<3x4xf64, #SparseMatrix> { + %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64, #SparseVector> into tensor<3x4xf64, #SparseMatrix> + return %0 : tensor<3x4xf64, #SparseMatrix> + } func.func @collapse_dense(%arg0: tensor<3x4xf64>) -> tensor<12xf64> { %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64> into tensor<12xf64> @@ -53,11 +52,10 @@ module { return %0 : tensor<12xf64, #SparseVector> } -// TODO: make this work -// func.func @collapse_sparse2sparse(%arg0: tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64, #SparseVector> { -// %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64, #SparseMatrix> into tensor<12xf64, #SparseVector> -// return %0 : tensor<12xf64, #SparseVector> -// } + func.func @collapse_sparse2sparse(%arg0: tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64, #SparseVector> { + %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64, #SparseMatrix> into tensor<12xf64, #SparseVector> + return %0 : tensor<12xf64, #SparseVector> + } // @@ -81,10 +79,12 @@ module { %expand0 = call @expand_dense(%v) : (tensor<12xf64>) -> tensor<3x4xf64> %expand1 = call @expand_from_sparse(%sv) : (tensor<12xf64, #SparseVector>) -> tensor<3x4xf64> %expand2 = call @expand_to_sparse(%v) : (tensor<12xf64>) -> tensor<3x4xf64, #SparseMatrix> + %expand3 = call @expand_sparse2sparse(%sv) : (tensor<12xf64, #SparseVector>) -> tensor<3x4xf64, #SparseMatrix> %collapse0 = call @collapse_dense(%m) : (tensor<3x4xf64>) -> tensor<12xf64> %collapse1 = call @collapse_from_sparse(%sm) : (tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64> %collapse2 = call @collapse_to_sparse(%m) : (tensor<3x4xf64>) -> tensor<12xf64, #SparseVector> + %collapse3 = call @collapse_sparse2sparse(%sm) : (tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64, #SparseVector> // // Verify result. @@ -92,9 +92,11 @@ module { // CHECK: ( ( 1, 2, 3, 4 ), ( 5, 6, 7, 8 ), ( 9, 10, 11, 12 ) ) // CHECK-NEXT: ( ( 1, 2, 3, 4 ), ( 5, 6, 7, 8 ), ( 9, 10, 11, 12 ) ) // CHECK-NEXT: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, -1, -1, -1, -1 ) + // CHECK-NEXT: ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, -1, -1, -1, -1 ) // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4 ) // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4 ) // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4, -1, -1, -1, -1 ) + // CHECK-NEXT: ( 1.1, 1.2, 1.3, 1.4, 2.1, 2.2, 2.3, 2.4, 3.1, 3.2, 3.3, 3.4, -1, -1, -1, -1 ) // %m0 = vector.transfer_read %expand0[%c0, %c0], %df: tensor<3x4xf64>, vector<3x4xf64> vector.print %m0 : vector<3x4xf64> @@ -103,6 +105,9 @@ module { %a2 = sparse_tensor.values %expand2 : tensor<3x4xf64, #SparseMatrix> to memref<?xf64> %m2 = vector.transfer_read %a2[%c0], %df: memref<?xf64>, vector<16xf64> vector.print %m2 : vector<16xf64> + %a3 = sparse_tensor.values %expand3 : tensor<3x4xf64, #SparseMatrix> to memref<?xf64> + %m3 = vector.transfer_read %a3[%c0], %df: memref<?xf64>, vector<16xf64> + vector.print %m3 : vector<16xf64> %v0 = vector.transfer_read %collapse0[%c0], %df: tensor<12xf64>, vector<12xf64> vector.print %v0 : vector<12xf64> @@ -111,12 +116,17 @@ module { %b2 = sparse_tensor.values %collapse2 : tensor<12xf64, #SparseVector> to memref<?xf64> %v2 = vector.transfer_read %b2[%c0], %df: memref<?xf64>, vector<16xf64> vector.print %v2 : vector<16xf64> + %b3 = sparse_tensor.values %collapse3 : tensor<12xf64, #SparseVector> to memref<?xf64> + %v3 = vector.transfer_read %b3[%c0], %df: memref<?xf64>, vector<16xf64> + vector.print %v3 : vector<16xf64> // Release sparse resources. sparse_tensor.release %sv : tensor<12xf64, #SparseVector> sparse_tensor.release %sm : tensor<3x4xf64, #SparseMatrix> sparse_tensor.release %expand2 : tensor<3x4xf64, #SparseMatrix> + sparse_tensor.release %expand3 : tensor<3x4xf64, #SparseMatrix> sparse_tensor.release %collapse2 : tensor<12xf64, #SparseVector> + sparse_tensor.release %collapse3 : tensor<12xf64, #SparseVector> // Release dense resources. %meme1 = bufferization.to_memref %expand1 : memref<3x4xf64> |