Skip to content

[mlir][sparse] Migrate to sparse_tensor.print #83946

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Mar 7, 2024
Merged

[mlir][sparse] Migrate to sparse_tensor.print #83946

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
9a9c8d8
migrate to sparse_tensor.print
yinying-lisa-li Mar 5, 2024
d3e9136
add more tests
yinying-lisa-li Mar 6, 2024
6730ae5
Merge branch 'main' into migrate
yinying-lisa-li Mar 6, 2024
b72ab32
migrate to sparse_tensor.print
yinying-lisa-li Mar 5, 2024
43da99a
add more tests
yinying-lisa-li Mar 6, 2024
bd6ab89
Merge branch 'main' into migrate
yinying-lisa-li Mar 6, 2024
489fba1
Merge branch 'migrate' of github.com:yinying-lisa-li/llvm-project int…
yinying-lisa-li Mar 6, 2024
48a85a9
address comments
yinying-lisa-li Mar 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
// DEFINE: %{run_opts} = -e entry -entry-point-result=void
// DEFINE: %{run_opts} = -e main -entry-point-result=void
// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
//
Expand Down Expand Up @@ -88,7 +88,7 @@ module {
//
// Main driver that reads matrix from file and calls the sparse kernel.
//
func.func @entry() {
func.func @main() {
%i0 = arith.constant 0 : i32
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
Expand Down
4 changes: 2 additions & 2 deletions mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
// DEFINE: %{run_opts} = -e entry -entry-point-result=void
// DEFINE: %{run_opts} = -e main -entry-point-result=void
// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
//
Expand Down Expand Up @@ -84,7 +84,7 @@ module {
//
// Main driver that reads matrix from file and calls the sparse kernel.
//
func.func @entry() {
func.func @main() {
%f0 = arith.constant 0.0 : f64
%cst0 = arith.constant 0 : index
%cst1 = arith.constant 1 : index
Expand Down
19 changes: 13 additions & 6 deletions mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_mult_elt.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
// DEFINE: %{run_opts} = -e entry -entry-point-result=void
// DEFINE: %{run_opts} = -e main -entry-point-result=void
// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
//
Expand Down Expand Up @@ -60,7 +60,7 @@ module {
}

// Driver method to call and verify kernel.
func.func @entry() {
func.func @main() {
%c0 = arith.constant 0 : index
%f0 = arith.constant 0.0 : f32

Expand All @@ -84,11 +84,18 @@ module {
//
// Verify results. Only two entries stored in result!
//
// CHECK: ( 14, 20, 0, 0 )
// CHECK: ---- Sparse Tensor ----
// CHECK-NEXT: nse = 2
// CHECK-NEXT: dim = ( 32, 16 )
// CHECK-NEXT: lvl = ( 32, 16 )
// CHECK-NEXT: pos[0] : ( 0, 2
// CHECK-NEXT: crd[0] : ( 2, 31
// CHECK-NEXT: pos[1] : ( 0, 1, 2
// CHECK-NEXT: crd[1] : ( 2, 0
// CHECK-NEXT: values : ( 14, 20
// CHECK-NEXT: ----
//
%val = sparse_tensor.values %0 : tensor<32x16xf32, #DCSR> to memref<?xf32>
%vv = vector.transfer_read %val[%c0], %f0: memref<?xf32>, vector<4xf32>
vector.print %vv : vector<4xf32>
sparse_tensor.print %0 : tensor<32x16xf32, #DCSR>

// Release the resources.
bufferization.dealloc_tensor %sta : tensor<32x16xf32, #DCSR>
Expand Down
29 changes: 13 additions & 16 deletions mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
// DEFINE: %{run_opts} = -e entry -entry-point-result=void
// DEFINE: %{run_opts} = -e main -entry-point-result=void
// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
//
Expand Down Expand Up @@ -70,10 +70,7 @@ module {
}

// Driver method to call and verify tensor kernel.
func.func @entry() {
%c0 = arith.constant 0 : index
%i0 = arith.constant 0 : i32

func.func @main() {
// Setup very sparse 3-d tensors.
%t1 = arith.constant sparse<
[ [1,1,3], [2,0,0], [2,2,1], [2,2,2], [2,2,3] ], [ 1, 2, 3, 4, 5 ]
Expand All @@ -94,23 +91,23 @@ module {
//
// Verify results. Only two entries stored in result. Correct structure.
//
// CHECK: ( 7, 69, 0, 0 )
// CHECK-NEXT: ( ( 0, 0, 0 ), ( 0, 7, 0 ), ( 0, 0, 69 ) )
// CHECK: ---- Sparse Tensor ----
// CHECK-NEXT: nse = 2
// CHECK-NEXT: dim = ( 3, 3 )
// CHECK-NEXT: lvl = ( 3, 3 )
// CHECK-NEXT: pos[0] : ( 0, 2
// CHECK-NEXT: crd[0] : ( 1, 2
// CHECK-NEXT: pos[1] : ( 0, 1, 2
// CHECK-NEXT: crd[1] : ( 1, 2
// CHECK-NEXT: values : ( 7, 69
// CHECK-NEXT: ----
//
%val = sparse_tensor.values %0
: tensor<?x?xi32, #SparseMatrix> to memref<?xi32>
%vv = vector.transfer_read %val[%c0], %i0: memref<?xi32>, vector<4xi32>
vector.print %vv : vector<4xi32>
%dm = sparse_tensor.convert %0
: tensor<?x?xi32, #SparseMatrix> to tensor<?x?xi32>
%vm = vector.transfer_read %dm[%c0, %c0], %i0: tensor<?x?xi32>, vector<3x3xi32>
vector.print %vm : vector<3x3xi32>
sparse_tensor.print %0 : tensor<?x?xi32, #SparseMatrix>

// Release the resources.
bufferization.dealloc_tensor %st1 : tensor<?x?x?xi32, #SparseTensor>
bufferization.dealloc_tensor %st2 : tensor<?x?x?xi32, #SparseTensor>
bufferization.dealloc_tensor %0 : tensor<?x?xi32, #SparseMatrix>
bufferization.dealloc_tensor %dm : tensor<?x?xi32>

return
}
Expand Down
19 changes: 13 additions & 6 deletions mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
// DEFINE: %{run_opts} = -e entry -entry-point-result=void
// DEFINE: %{run_opts} = -e main -entry-point-result=void
// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
//
Expand Down Expand Up @@ -70,7 +70,7 @@ module {
//
// Main driver that reads matrix from file and calls the sparse kernel.
//
func.func @entry() {
func.func @main() {
%d0 = arith.constant 0.0 : f64
%c0 = arith.constant 0 : index

Expand All @@ -83,11 +83,18 @@ module {

// Print the result for verification.
//
// CHECK: ( 1, 1.96, 4, 6.25, 9, 16.81, 16, 27.04, 25 )
// CHECK: ---- Sparse Tensor ----
// CHECK-NEXT: nse = 9
// CHECK-NEXT: dim = ( 5, 5 )
// CHECK-NEXT: lvl = ( 5, 5 )
// CHECK-NEXT: pos[0] : ( 0, 5
// CHECK-NEXT: crd[0] : ( 0, 1, 2, 3, 4
// CHECK-NEXT: pos[1] : ( 0, 2, 4, 5, 7, 9
// CHECK-NEXT: crd[1] : ( 0, 3, 1, 4, 2, 0, 3, 1, 4
// CHECK-NEXT: values : ( 1, 1.96, 4, 6.25, 9, 16.81, 16, 27.04, 25
// CHECK-NEXT: ----
//
%m = sparse_tensor.values %0 : tensor<?x?xf64, #DCSR> to memref<?xf64>
%v = vector.transfer_read %m[%c0], %d0: memref<?xf64>, vector<9xf64>
vector.print %v : vector<9xf64>
sparse_tensor.print %0 : tensor<?x?xf64, #DCSR>

// Release the resources.
bufferization.dealloc_tensor %x : tensor<?x?xf64, #DCSR>
Expand Down
4 changes: 2 additions & 2 deletions mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
// DEFINE: %{run_opts} = -e entry -entry-point-result=void
// DEFINE: %{run_opts} = -e main -entry-point-result=void
// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
//
Expand Down Expand Up @@ -54,7 +54,7 @@ module {
//
// Main driver.
//
func.func @entry() {
func.func @main() {
%c0 = arith.constant 0 : index
%f0 = arith.constant 0.0 : f64
%i0 = arith.constant 0 : i32
Expand Down
53 changes: 33 additions & 20 deletions mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pack_d.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
// DEFINE: %{run_opts} = -e entry -entry-point-result=void
// DEFINE: %{run_opts} = -e main -entry-point-result=void
// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
//
Expand Down Expand Up @@ -48,7 +48,7 @@ module {
//
// Main driver.
//
func.func @entry() {
func.func @main() {
%c0 = arith.constant 0 : index
%f0 = arith.constant 0.0 : f32

Expand Down Expand Up @@ -107,26 +107,39 @@ module {
//
// Verify.
//
// CHECK: ( ( ( 1, 2 ), ( 3, 4 ), ( 0, 0 ) ), ( ( 0, 0 ), ( 0, 0 ), ( 0, 0 ) ), ( ( 0, 0 ), ( 5, 0 ), ( 6, 7 ) ), ( ( 0, 0 ), ( 8, 0 ), ( 0, 0 ) ) )
// CHECK: ( ( ( 1, 2 ), ( 0, 3 ), ( 4, 0 ) ), ( ( 5, 6 ), ( 0, 0 ), ( 0, 7 ) ), ( ( 8, 9 ), ( 10, 11 ), ( 12, 13 ) ), ( ( 14, 0 ), ( 0, 15 ), ( 0, 16 ) ) )
// CHECK: ( ( ( 1, 2 ), ( 0, 3 ), ( 4, 0 ) ), ( ( 5, 6 ), ( 0, 0 ), ( 0, 7 ) ), ( ( 8, 9 ), ( 10, 11 ), ( 12, 13 ) ), ( ( 14, 0 ), ( 0, 15 ), ( 0, 16 ) ) )
// CHECK: ---- Sparse Tensor ----
// CHECK-NEXT: nse = 8
// CHECK-NEXT: dim = ( 4, 3, 2 )
// CHECK-NEXT: lvl = ( 4, 3, 2 )
// CHECK-NEXT: pos[0] : ( 0, 3
// CHECK-NEXT: crd[0] : ( 0, 2, 3
// CHECK-NEXT: pos[1] : ( 0, 2, 4, 5
// CHECK-NEXT: crd[1] : ( 0, 1, 1, 2, 1
// CHECK-NEXT: pos[2] : ( 0, 2, 4, 5, 7, 8
// CHECK-NEXT: crd[2] : ( 0, 1, 0, 1, 0, 0, 1, 0
// CHECK-NEXT: values : ( 1, 2, 3, 4, 5, 6, 7, 8
// CHECK-NEXT: ----
// CHECK: ---- Sparse Tensor ----
// CHECK-NEXT: nse = 16
// CHECK-NEXT: dim = ( 4, 3, 2 )
// CHECK-NEXT: lvl = ( 4, 3, 2 )
// CHECK-NEXT: pos[2] : ( 0, 2, 3, 4, 6, 6, 7, 9, 11, 13, 14, 15, 16
// CHECK-NEXT: crd[2] : ( 0, 1, 1, 0, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1
// CHECK-NEXT: values : ( 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
// CHECK-NEXT: ----
// CHECK: ---- Sparse Tensor ----
// CHECK-NEXT: nse = 22
// CHECK-NEXT: dim = ( 4, 3, 2 )
// CHECK-NEXT: lvl = ( 4, 3, 2 )
// CHECK-NEXT: pos[1] : ( 0, 3, 5, 8, 11
// CHECK-NEXT: crd[1] : ( 0, 1, 2, 0, 2, 0, 1, 2, 0, 1, 2
// CHECK-NEXT: values : ( 1, 2, 0, 3, 4, 0, 5, 6, 0, 7, 8, 9, 10, 11, 12, 13, 14, 0, 0, 15, 0, 16
// CHECK-NEXT: ----
//
sparse_tensor.print %s0 : tensor<4x3x2xf32, #CCC>
sparse_tensor.print %s1 : tensor<4x3x2xf32, #BatchedCSR>
sparse_tensor.print %s2 : tensor<4x3x2xf32, #CSRDense>

%d0 = sparse_tensor.convert %s0 : tensor<4x3x2xf32, #CCC> to tensor<4x3x2xf32>
%v0 = vector.transfer_read %d0[%c0, %c0, %c0], %f0 : tensor<4x3x2xf32>, vector<4x3x2xf32>
vector.print %v0 : vector<4x3x2xf32>

%d1 = sparse_tensor.convert %s1 : tensor<4x3x2xf32, #BatchedCSR> to tensor<4x3x2xf32>
%v1 = vector.transfer_read %d1[%c0, %c0, %c0], %f0 : tensor<4x3x2xf32>, vector<4x3x2xf32>
vector.print %v1 : vector<4x3x2xf32>

%d2 = sparse_tensor.convert %s2 : tensor<4x3x2xf32, #CSRDense> to tensor<4x3x2xf32>
%v2 = vector.transfer_read %d1[%c0, %c0, %c0], %f0 : tensor<4x3x2xf32>, vector<4x3x2xf32>
vector.print %v2 : vector<4x3x2xf32>

bufferization.dealloc_tensor %d0 : tensor<4x3x2xf32>
bufferization.dealloc_tensor %d1 : tensor<4x3x2xf32>
bufferization.dealloc_tensor %d2 : tensor<4x3x2xf32>
// FIXME: doing this explicitly crashes runtime
// bufferization.dealloc_tensor %s0 : tensor<4x3x2xf32, #CCC>
// bufferization.dealloc_tensor %s1 : tensor<4x3x2xf32, #BatchedCSR>
Expand Down
27 changes: 18 additions & 9 deletions mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_pooling_nhwc.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
// DEFINE: %{run_opts} = -e entry -entry-point-result=void
// DEFINE: %{run_opts} = -e main -entry-point-result=void
// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
//
Expand Down Expand Up @@ -47,7 +47,7 @@ func.func @pooling_nhwc_sum(%input: tensor<1x4x4x1xf32>, %filter: tensor<2x2xf32
}


func.func @entry() {
func.func @main() {
%c0 = arith.constant 0 : index
%zero = arith.constant 0.00000e+00 : f32

Expand Down Expand Up @@ -76,17 +76,26 @@ func.func @entry() {
//
// Sparse pooling should have the same output.
//

// CHECK-NEXT: ( ( ( ( 6 ), ( 6 ), ( 6 ) ), ( ( 6 ), ( 6 ), ( 6 ) ), ( ( 6 ), ( 6 ), ( 6 ) ) ) )
%s1 = sparse_tensor.convert %CCCC_ret : tensor<1x3x3x1xf32, #CCCC> to tensor<1x3x3x1xf32>
%v1 = vector.transfer_read %s1[%c0, %c0, %c0, %c0], %zero
: tensor<1x3x3x1xf32>, vector<1x3x3x1xf32>
vector.print %v1 : vector<1x3x3x1xf32>
// CHECK: ---- Sparse Tensor ----
// CHECK-NEXT: nse = 9
// CHECK-NEXT: dim = ( 1, 3, 3, 1 )
// CHECK-NEXT: lvl = ( 1, 3, 3, 1 )
// CHECK-NEXT: pos[0] : ( 0, 1
// CHECK-NEXT: crd[0] : ( 0
// CHECK-NEXT: pos[1] : ( 0, 3
// CHECK-NEXT: crd[1] : ( 0, 1, 2
// CHECK-NEXT: pos[2] : ( 0, 3, 6, 9
// CHECK-NEXT: crd[2] : ( 0, 1, 2, 0, 1, 2, 0, 1, 2
// CHECK-NEXT: pos[3] : ( 0, 1, 2, 3, 4, 5, 6, 7, 8, 9
// CHECK-NEXT: crd[3] : ( 0, 0, 0, 0, 0, 0, 0, 0, 0
// CHECK-NEXT: values : ( 6, 6, 6, 6, 6, 6, 6, 6, 6
// CHECK-NEXT: ----
//
sparse_tensor.print %CCCC_ret : tensor<1x3x3x1xf32, #CCCC>

// Releases resources.
bufferization.dealloc_tensor %in_CCCC : tensor<1x4x4x1xf32, #CCCC>
bufferization.dealloc_tensor %CCCC_ret : tensor<1x3x3x1xf32, #CCCC>
bufferization.dealloc_tensor %dense_ret : tensor<1x3x3x1xf32>
bufferization.dealloc_tensor %s1 : tensor<1x3x3x1xf32>
return
}
4 changes: 2 additions & 2 deletions mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
// DEFINE: %{run_opts} = -e entry -entry-point-result=void
// DEFINE: %{run_opts} = -e main -entry-point-result=void
// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
//
Expand Down Expand Up @@ -49,7 +49,7 @@ module {
return %0: tensor<5x6xi32>
}

func.func @entry() {
func.func @main() {
%c0 = arith.constant 0 : index
%i0 = arith.constant 0 : i32

Expand Down
Loading