[Mlir-commits] [mlir] 758ccf8 - [mlir][sparse] add test for DimOp folding

[Aart Bik]

Author: Aart Bik
Date: 2021-08-20T11:24:09-07:00
New Revision: 758ccf8506c387f03a2c42d303de1343f2c4022b

URL: https://github.com/llvm/llvm-project/commit/758ccf8506c387f03a2c42d303de1343f2c4022b
DIFF: https://github.com/llvm/llvm-project/commit/758ccf8506c387f03a2c42d303de1343f2c4022b.diff

LOG: [mlir][sparse] add test for DimOp folding

Folding in the MLIR uses the order of the type directly
but folding in the underlying implementation must take
the dim ordering into account. These tests clarify that
behavior and verify it is done right.

Reviewed By: bixia

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

Added: 
    

Modified: 
    mlir/test/Dialect/SparseTensor/conversion.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/test/Dialect/SparseTensor/conversion.mlir b/mlir/test/Dialect/SparseTensor/conversion.mlir
index 2b479b53eb1e..5a2e3b135672 100644
--- a/mlir/test/Dialect/SparseTensor/conversion.mlir
+++ b/mlir/test/Dialect/SparseTensor/conversion.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s --sparse-tensor-conversion --canonicalize | FileCheck %s
+// RUN: mlir-opt %s --sparse-tensor-conversion --canonicalize --cse | FileCheck %s
 
 #DenseVector = #sparse_tensor.encoding<{
   dimLevelType = ["dense"]
@@ -46,12 +46,27 @@ func @sparse_dim1d(%arg0: tensor<?xf64, #SparseVector>) -> index {
 //       CHECK: %[[D:.*]] = call @sparseDimSize(%[[A]], %[[C]])
 //       CHECK: return %[[D]] : index
 func @sparse_dim3d(%arg0: tensor<?x?x?xf64, #SparseTensor>) -> index {
-  // Needs permuting 1 into 2.
+  // Querying for dimension 1 in the tensor type needs to be
+  // permuted into querying for dimension 2 in the stored sparse
+  // tensor scheme, since the latter honors the dimOrdering.
   %c = constant 1 : index
   %0 = tensor.dim %arg0, %c : tensor<?x?x?xf64, #SparseTensor>
   return %0 : index
 }
 
+// CHECK-LABEL: func @sparse_dim3d_const(
+//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
+//       CHECK: %[[C:.*]] = constant 20 : index
+//       CHECK: return %[[C]] : index
+func @sparse_dim3d_const(%arg0: tensor<10x20x30xf64, #SparseTensor>) -> index {
+  // Querying for dimension 1 in the tensor type can be directly
+  // folded into the right value (even though it corresponds
+  // to dimension 2 in the stored sparse tensor scheme).
+  %c = constant 1 : index
+  %0 = tensor.dim %arg0, %c : tensor<10x20x30xf64, #SparseTensor>
+  return %0 : index
+}
+
 // CHECK-LABEL: func @sparse_new1d(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8>
 //   CHECK-DAG: %[[U:.*]] = constant dense<1> : tensor<1xi8>