[Mlir-commits] [mlir] 6df483c - [mlir][sparse] add a check test for foreach operation on constant sparse tensor

[Peiming Liu]

Author: Peiming Liu
Date: 2023-03-09T21:25:37Z
New Revision: 6df483c9a03b40a3e29bff7f3b0c0859c45ae882

URL: https://github.com/llvm/llvm-project/commit/6df483c9a03b40a3e29bff7f3b0c0859c45ae882
DIFF: https://github.com/llvm/llvm-project/commit/6df483c9a03b40a3e29bff7f3b0c0859c45ae882.diff

LOG: [mlir][sparse] add a check test for foreach operation on constant sparse tensor

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D145728

Added: 
    mlir/test/Dialect/SparseTensor/sparse_foreach.mlir

Modified: 
    

Removed: 
    


################################################################################
diff  --git a/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir b/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir
new file mode 100644
index 000000000000..ad203abba088
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir
@@ -0,0 +1,29 @@
+// RUN: mlir-opt %s --post-sparsification-rewrite="enable-runtime-library=false enable-foreach=true" | FileCheck %s
+
+// CHECK-LABEL: func.func @sparse_foreach_constant
+// CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG:   %[[V1:.*]] = arith.constant 5.000000e+00 : f32
+// CHECK-DAG:   %[[C2:.*]] = arith.constant 2 : index
+// CHECK-DAG:   %[[V3:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK-DAG:   %[[V4:.*]] = arith.constant 6.000000e+00 : f32
+//               (1, 1) -> (2, 1) -> (2, 2)
+// CHECK-NEXT:  "test.use"(%[[C1]], %[[C1]], %[[V1]])
+// CHECK-NEXT:  "test.use"(%[[C2]], %[[C1]], %[[V3]])
+// CHECK-NEXT:  "test.use"(%[[C1]], %[[C2]], %[[V4]])
+//               (1, 1) -> (1, 2) -> (2, 1)
+// CHECK-NEXT:  "test.use"(%[[C1]], %[[C1]], %[[V1]])
+// CHECK-NEXT:  "test.use"(%[[C1]], %[[C2]], %[[V4]])
+// CHECK-NEXT:  "test.use"(%[[C2]], %[[C1]], %[[V3]])
+func.func @sparse_foreach_constant() -> () {
+  %cst = arith.constant sparse<[[2, 1], [1, 1], [1, 2]], [1.0, 5.0, 6.0]> : tensor<8x7xf32>
+  // Make use the sparse constant are properly sorted based on the requested order.
+  sparse_tensor.foreach in %cst { order = affine_map<(d0, d1) -> (d1, d0)> } : tensor<8x7xf32> do {
+  ^bb0(%arg0: index, %arg1: index, %arg2: f32):
+    "test.use" (%arg0, %arg1, %arg2): (index,index,f32)->()
+  }
+  sparse_tensor.foreach in %cst : tensor<8x7xf32> do {
+  ^bb0(%arg0: index, %arg1: index, %arg2: f32):
+    "test.use" (%arg0, %arg1, %arg2): (index,index,f32)->()
+  }
+  return
+}