[mlir][sparse] relax constraints on tensor.cast with pre-rewriting

[wrengr]

This CL looks incomplete, since you're not actually using isSameTypesWithoutEncoding anywhere...

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

105

Should be just "Type" not "Types"

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

105

hrmm, it looks like that's the name used by lib/Dialect/Tensor/IR/TensorOps.cpp... So I guess consistency is probably better than grammatical correctness...

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

105

Actually, now that I've seen that other version of this function, I think it would be better to have this CL change the lib/Dialect/Tensor/IR/TensorOps.cpp version from being static to being exported via include/Dialect/Tensor/IR/TensorOps.h. And then you should use that public version for whatever you needed here

[aartbik]

add missing file

[aartbik]

addressed comments

[wrengr]

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

113 ↗

(On Diff #518538)

Any particular reason why you moved the definition from where the old one was? Just asking for the sake of minimizing the diff

113 ↗

(On Diff #518538)

(i.e., the line numbers; not the making it public)

[wrengr]

mlir/include/mlir/Dialect/Tensor/IR/Tensor.h

145 ↗

(On Diff #518538)

I would use "when" here rather than "but"

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

378 ↗

(On Diff #518538)

Is there any way to have this case emit a warning message to the user? (or "note" or whatever severity-level seems most appropriate)

[aartbik]

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

378 ↗

(On Diff #518538)

We actually wanted to add this to verify() proper (and fully reject such cases). But there was also some push back with adding sparse specific checks in tensor dialect. This walks a bit of a middle ground until we come up with a full design that everybody likes followed by full audit of "dense" code.

[aartbik]

addressed comment

[aartbik]

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

113 ↗

(On Diff #518538)

I am usually all +1 on minimizing diffs. But here the original static method was in a section marked " Reassociative reshape ops" so I decided to move it up (where the other public methods are defined, and in same order as header file).

[dcaballe]

The non-sparse part looks good to me! I'll leave the final approval to @wrengr. Thanks!

[wrengr]

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

378 ↗

(On Diff #518538)

I'm fine with leaving this rewriting case here, I was just wondering whether MLIR had some way to emit warnings despite also returning success. That is, since we consider this rewrite to be dubious, it'd be nice to give users some sort of notification whenever it fires. (perhaps only doing so whenever they set the debugging verbosity high enough)

[aartbik]

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

378 ↗

(On Diff #518538)

For sure, I don't believe we have any good precedent for that at the moment. If we get this, we may add it here (or, hopefully one day, just reject the op in dense verify() case ;-)

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

 inline MemRefType getMemRefType(T t) {
   return t.getType().template cast<MemRefType>();
 }
 
 /// Convenience method to get a sparse encoding attribute from a type.
 /// Returns null-attribute for any type without an encoding.
 SparseTensorEncodingAttr getSparseTensorEncoding(Type type);
 
+/// Tests if types are the same, but ignoring encoding on ranked tensors.
+bool isSameTypesWithoutEncoding(Type tp1, Type tp2);
+
 /// Returns true iff the given sparse tensor encoding attribute has a trailing
 /// COO region starting at the given level.
 bool isCOOType(SparseTensorEncodingAttr enc, Level startLvl, bool isUnique);
 
 /// Returns true iff the given type is a COO type where the last level
 /// is unique.
 bool isUniqueCOOType(Type tp);
 

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

 mlir::sparse_tensor::getSparseTensorEncoding(Type type) {
   if (auto ttp = type.dyn_cast<RankedTensorType>())
     return ttp.getEncoding().dyn_cast_or_null<SparseTensorEncodingAttr>();
   if (auto mdtp = type.dyn_cast<StorageSpecifierType>())
     return mdtp.getEncoding();
   return nullptr;
 }
 
-/// Returns true iff the given sparse tensor encoding attribute has a trailing
-/// COO region starting at the given level.
+bool mlir::sparse_tensor::isSameTypesWithoutEncoding(Type tp1, Type tp2) {
+  if (auto rtp1 = tp1.dyn_cast<RankedTensorType>()) {
+    if (auto rtp2 = tp2.dyn_cast<RankedTensorType>())
+      return rtp1.getShape() == rtp2.getShape() &&
+             rtp1.getElementType() == rtp2.getElementType();
+    return false;
+  }
+  return tp1 == tp2; // default test
+}
+
 bool mlir::sparse_tensor::isCOOType(SparseTensorEncodingAttr enc,
                                     Level startLvl, bool isUnique) {
   if (!enc ||
       !(enc.isCompressedLvl(startLvl) || enc.isCompressedWithHiLvl(startLvl)))
     return false;
   const Level lvlRank = enc.getLvlRank();
   for (Level l = startLvl + 1; l < lvlRank; ++l)
     if (!enc.isSingletonLvl(l))

mlir/test/Dialect/SparseTensor/post_rewriting.mlir

 // RUN: mlir-opt %s -post-sparsification-rewrite | FileCheck %s
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"]
 }>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = ["compressed", "compressed"]
 }>
 
 // CHECK-LABEL: func.func @expand_dense(
 // CHECK-SAME:    %[[A:.*]]: tensor<12xf64>) -> tensor<3x4xf64> {
 // CHECK:         %[[E:.*]] = tensor.expand_shape %[[A]] {{.*}} : tensor<12xf64> into tensor<3x4xf64>
 // CHECK:         return %[[E]] : tensor<3x4xf64>
 // CHECK:       }
 func.func @expand_dense(%arg0: tensor<12xf64>) -> tensor<3x4xf64> {
   %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64> into tensor<3x4xf64>
   return %0 : tensor<3x4xf64>
 }
 
 // CHECK-LABEL: func.func @expand_from_sparse(
 // CHECK-SAME:    %[[A:.*]]: tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<3x4xf64> {
 // CHECK:         %[[C:.*]] = sparse_tensor.convert %[[A]] : tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>> to tensor<12xf64>
 // CHECK:         %[[E:.*]] = tensor.expand_shape %[[C]] {{.*}} : tensor<12xf64> into tensor<3x4xf64>
 // CHECK:         return %[[E]] : tensor<3x4xf64>
 // CHECK:       }
 func.func @expand_from_sparse(%arg0: tensor<12xf64, #SparseVector>) -> tensor<3x4xf64> {
   %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64, #SparseVector> into tensor<3x4xf64>
   return %0 : tensor<3x4xf64>
 }
 
 // CHECK-LABEL: func.func @expand_to_sparse(
 // CHECK-SAME:    %[[A:.*]]: tensor<12xf64>) -> tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>> {
 // CHECK:         %[[E:.*]] = tensor.expand_shape %[[A]] {{.*}} : tensor<12xf64> into tensor<3x4xf64>
 // CHECK:         %[[C:.*]] = sparse_tensor.convert %[[E]] : tensor<3x4xf64> to tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>>
 // CHECK:         return %[[C]] : tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>>
 // CHECK:       }
 func.func @expand_to_sparse(%arg0: tensor<12xf64>) -> tensor<3x4xf64, #SparseMatrix> {
   %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<12xf64> into tensor<3x4xf64, #SparseMatrix>
   return %0 : tensor<3x4xf64, #SparseMatrix>
 }
 
 //
 // Not rewritten, needs conversion.
 //
 // CHECK-LABEL:   func.func @expand_sparse2sparse(
 // CHECK-SAME:    %[[A:.*]]: tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>> {
 // CHECK:         %[[E:.*]] = tensor.expand_shape %[[A]] {{.*}} : tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>> into tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>>
 // CHECK:         return %[[E]] : tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>>
 // CHECK:       }
 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>
 }
 
 // CHECK-LABEL: func.func @collapse_dense(
 // CHECK-SAME:    %[[A:.*]]: tensor<3x4xf64>) -> tensor<12xf64> {
 // CHECK:         %[[R:.*]] = tensor.collapse_shape %[[A]] {{.*}} : tensor<3x4xf64> into tensor<12xf64>
 // CHECK:         return %[[R]] : tensor<12xf64>
 // CHECK:       }
 func.func @collapse_dense(%arg0: tensor<3x4xf64>) -> tensor<12xf64> {
   %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64> into tensor<12xf64>
   return %0 : tensor<12xf64>
 }
 
 // CHECK-LABEL: func.func @collapse_from_sparse(
 // CHECK-SAME:    %[[A:.*]]: tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<12xf64> {
 // CHECK:         %[[C:.*]] = sparse_tensor.convert %[[A]] : tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>> to tensor<3x4xf64>
 // CHECK:         %[[R:.*]] = tensor.collapse_shape %[[C]] {{.*}} : tensor<3x4xf64> into tensor<12xf64>
 // CHECK:         return %[[R]] : tensor<12xf64>
 // CHECK:       }
 func.func @collapse_from_sparse(%arg0: tensor<3x4xf64, #SparseMatrix>) -> tensor<12xf64> {
   %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64, #SparseMatrix> into tensor<12xf64>
   return %0 : tensor<12xf64>
 }
 
 // CHECK-LABEL: func.func @collapse_to_sparse(
 // CHECK-SAME:    %[[A:.*]]: tensor<3x4xf64>) -> tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>> {
 // CHECK:         %[[R:.*]] = tensor.collapse_shape %[[A]] {{.*}} : tensor<3x4xf64> into tensor<12xf64>
 // CHECK:         %[[C:.*]] = sparse_tensor.convert %[[R]] : tensor<12xf64> to tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>>
 // CHECK:         return %[[C]] : tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>>
 // CHECK:       }
 func.func @collapse_to_sparse(%arg0: tensor<3x4xf64>) -> tensor<12xf64, #SparseVector> {
   %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<3x4xf64> into tensor<12xf64, #SparseVector>
   return %0 : tensor<12xf64, #SparseVector>
 }
 
 //
 // Not rewritten, needs conversion.
 //
 // CHECK-LABEL:   func.func @collapse_sparse2sparse(
 // CHECK-SAME:    %[[A:.*]]: tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>> {
 // CHECK:         %[[C:.*]] = tensor.collapse_shape %[[A]] {{.*}} : tensor<3x4xf64, #sparse_tensor.encoding<{{{.*}}}>> into tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>>
 // CHECK:         return %[[C]] : tensor<12xf64, #sparse_tensor.encoding<{{{.*}}}>>
 // CHECK:       }
 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>
 }

mlir/test/Dialect/SparseTensor/pre_rewriting.mlir

+// RUN: mlir-opt %s -pre-sparsification-rewrite | FileCheck %s
+
+#SparseVector = #sparse_tensor.encoding<{
+  dimLevelType = ["compressed"]
+}>
+
+#SortedCOO = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton" ]
+}>
+
+#Slice = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton" ],
+  slice = [ (?, 1, 1), (?, 3, 1) ]
+}>
+
+// CHECK-LABEL: func @sparse_nop_cast(
+//  CHECK-SAME: %[[A:.*]]: tensor<?xf32, #sparse_tensor.encoding<{{{.*}}}>>)
+//       CHECK: return %[[A]] : tensor<?xf32, #sparse_tensor.encoding<{{{.*}}}>>
+func.func @sparse_nop_cast(%a : tensor<?xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> {
+  %0 = tensor.cast %a : tensor<?xf32, #SparseVector> to tensor<?xf32, #SparseVector>
+  %1 = tensor.cast %0 : tensor<?xf32, #SparseVector> to tensor<?xf32, #SparseVector>
+  %2 = tensor.cast %1 : tensor<?xf32, #SparseVector> to tensor<?xf32, #SparseVector>
+  return %2 : tensor<?xf32, #SparseVector>
+}
+
+// CHECK-LABEL: func @sparse_repair_cast(
+//  CHECK-SAME: %[[A:.*]]: tensor<?xf32>)
+//       CHECK: %[[C:.*]] = sparse_tensor.convert %[[A]] : tensor<?xf32> to tensor<?xf32, #sparse_tensor.encoding<{{{.*}}}>
+//       CHECK: return %[[C]] : tensor<?xf32, #sparse_tensor.encoding<{{{.*}}}>>
+func.func @sparse_repair_cast(%a : tensor<?xf32>) -> tensor<?xf32, #SparseVector> {
+  %0 = tensor.cast %a : tensor<?xf32> to tensor<?xf32, #SparseVector>
+  return %0 : tensor<?xf32, #SparseVector>
+}
+
+// CHECK-LABEL: func @sparse_fuse_slice(
+//  CHECK-SAME: %[[A:.*]]: tensor<2x3xi64, #sparse_tensor.encoding<{{{.*}}}>>)
+//       CHECK: %[[E:.*]] = tensor.extract_slice %[[A]][1, 0] [1, 3] [1, 1] : tensor<2x3xi64, #sparse_tensor.encoding<{{{.*}}}>> to tensor<1x3xi64, #sparse_tensor.encoding<{{{.*}}}>>
+//       CHECK: %[[C:.*]] = sparse_tensor.convert %[[E]] : tensor<1x3xi64, #sparse_tensor.encoding<{{{.*}}}>> to tensor<1x3xi64, #sparse_tensor.encoding<{{{.*}}}>>
+//       CHECK: return %[[C]] : tensor<1x3xi64, #sparse_tensor.encoding<{{{.*}}}>>
+func.func @sparse_fuse_slice(%a : tensor<2x3xi64, #SortedCOO>) -> tensor<1x3xi64, #SortedCOO> {
+  %extracted_slice = tensor.extract_slice %a[1, 0] [1, 3] [1, 1] : tensor<2x3xi64, #SortedCOO> to tensor<1x3xi64>
+  %cast = tensor.cast %extracted_slice : tensor<1x3xi64> to tensor<1x3xi64, #Slice>
+  %0 = sparse_tensor.convert %cast : tensor<1x3xi64, #Slice> to tensor<1x3xi64, #SortedCOO>
+  return %0 : tensor<1x3xi64, #SortedCOO>
+}