[mlir][sparse] add test for DimOp folding

[aartbik]

Summary

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.

[bixia]

mlir/test/Dialect/SparseTensor/conversion.mlir

66

can we make the irrelevant dims unknown, such as ?x20x?xf64?

[aartbik]

mlir/test/Dialect/SparseTensor/conversion.mlir

66

I could, but why do you think that is better? If you want, I can add another test with that.

[bixia]

mlir/test/Dialect/SparseTensor/conversion.mlir

66

10x20x30 is a static shape but ?x20x? is a dynamic shape.

[aartbik]

rebase with main

mlir/test/Dialect/SparseTensor/conversion.mlir

-// 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"]
 }>
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"]
 }>

 }
 
 // CHECK-LABEL: func @sparse_dim3d(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: %[[C:.*]] = constant 2 : 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>
 //   CHECK-DAG: %[[V:.*]] = constant dense<128> : tensor<1xi64>
 //   CHECK-DAG: %[[W:.*]] = constant dense<0> : tensor<1xi64>
 //   CHECK-DAG: %[[X:.*]] = tensor.cast %[[U]] : tensor<1xi8> to tensor<?xi8>
 //   CHECK-DAG: %[[Y:.*]] = tensor.cast %[[V]] : tensor<1xi64> to tensor<?xi64>
 //   CHECK-DAG: %[[Z:.*]] = tensor.cast %[[W]] : tensor<1xi64> to tensor<?xi64>