[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 ↗

(On Diff #518065)

Should be just "Type" not "Types"

[wrengr]

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

105 ↗

(On Diff #518065)

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 ↗

(On Diff #518065)

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

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

113

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

[wrengr]

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

145

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

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

378

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

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

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

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

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/Tensor/IR/Tensor.h

 FailureOr<Value> getOrCreateDestination(OpBuilder &b, Location loc,
                                         OpResult opResult);
 
 /// This is a helper function for DestinationStyleOpInterface. Get or create
 /// destinations for every tensor OpResult of the given op.
 LogicalResult getOrCreateDestinations(OpBuilder &b, Location loc, Operation *op,
                                       SmallVector<Value> &result);
 
+/// Tests if types are the same when ignoring encoding on ranked tensors.
+bool isSameTypeWithoutEncoding(Type tp1, Type tp2);
+
 /// Function to control the folding of constant and extract slice.
 using ControlConstantExtractSliceFusionFn = std::function<bool(ExtractSliceOp)>;
 
 /// Patterns to fold the extract slice op with its constant operand.
 void populateFoldConstantExtractSlicePatterns(
     RewritePatternSet &patterns,
     const ControlConstantExtractSliceFusionFn &controlFn =
         [](ExtractSliceOp op) {

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::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/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

 
 private:
   // Helper to add argument and record the mapping.
   static void addArg(IRMapping &mapper, Block *b, BlockArgument a) {
     mapper.map(a, b->addArgument(a.getType(), a.getLoc()));
   }
 };
 
+// Fuse a tensor cast into producing operation. Note that a tensor.cast
+// should really not be used to convert between sparse encodings. Since
+// the pattern currently appears as a result of some prior rewriting
+// we make an attempt to repair very obvious cases.
+// TODO: audit the pure tensor dialect rewriting rules
+struct FuseTensorCast : public OpRewritePattern<tensor::CastOp> {
+public:
+  using OpRewritePattern<tensor::CastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::CastOp op,
+                                PatternRewriter &rewriter) const override {
+    Type srcType = op.getSource().getType();
+    Type dstType = op.getDest().getType();
+    // A nop cast simply folds away.
+    if (srcType == dstType) {
+      rewriter.replaceOp(op, op->getResults());
+      return success();
+    }
+    // See if a sparsity changing cast can be fused into producer.
+    if (tensor::isSameTypeWithoutEncoding(srcType, dstType)) {
+      if (Operation *def = op.getSource().getDefiningOp()) {
+        if (def->hasOneUse() && isa<tensor::ExtractSliceOp>(def)) {
+          def->getResult(0).setType(op->getResultTypes()[0]);
+          rewriter.replaceOp(op, def->getResult(0));
+          return success();
+        }
+      }
+    }
+    // Repair tensor casts with at least one sparse operand into the
+    // the properly supported sparse_tensor.convert.
+    if (getSparseTensorEncoding(srcType) || getSparseTensorEncoding(dstType)) {
+      rewriter.replaceOpWithNewOp<ConvertOp>(op, dstType, op.getSource());
+      return success();
+    }
+    // Fail otherwise.
+    return failure();
+  }
+};
+
 /// Sparse rewriting rule for sparse-to-sparse reshape operator.
 template <typename ReshapeOp>
 struct Sparse2SparseReshapeRewriter : public OpRewritePattern<ReshapeOp> {
 public:
   using OpRewritePattern<ReshapeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ReshapeOp op,
                                 PatternRewriter &rewriter) const override {

 
 } // namespace
 
 //===---------------------------------------------------------------------===//
 // Methods that add patterns described in this file to a pattern list.
 //===---------------------------------------------------------------------===//
 
 void mlir::populatePreSparsificationRewriting(RewritePatternSet &patterns) {
-  patterns.add<FoldInvariantYield, FuseSparseMultiplyOverAdd>(
+  patterns.add<FoldInvariantYield, FuseSparseMultiplyOverAdd, FuseTensorCast>(
       patterns.getContext());
 }
 
 void mlir::populatePostSparsificationRewriting(RewritePatternSet &patterns,
                                                bool enableRT,
                                                bool enableForeach,
                                                bool enableConvert) {
   patterns.add<ReshapeRewriter<tensor::ExpandShapeOp>,

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

if (opResult.getType().isa<TensorType>()) {
       if (failed(destination))
         return failure();
       result.push_back(*destination);
     }
   }
   return success();
 }
 
+bool tensor::isSameTypeWithoutEncoding(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 implementation
+}
+
 /// Compute the dropped dimensions of a rank-reducing tensor.extract_slice op or
 /// rank-extending tensor.insert_slice op.
 static llvm::SmallBitVector getDroppedDims(ArrayRef<int64_t> reducedShape,
                                            ArrayRef<OpFoldResult> mixedSizes) {
   llvm::SmallBitVector droppedDims(mixedSizes.size());
   int64_t shapePos = 0;
 
   for (const auto &size : enumerate(mixedSizes)) {
auto resultType = inferCollapsedType(
       src.getType().cast<RankedTensorType>(),
       getSymbolLessAffineMaps(
           convertReassociationIndicesToExprs(b.getContext(), reassociation)));
   build(b, result, resultType, src, attrs);
   result.addAttribute(getReassociationAttrStrName(),
                       getReassociationIndicesAttribute(b, reassociation));
 }
 
-// Checks if types are the same, but ignoring encoding on ranked tensors.
-static bool 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;
-  }
-  // Default implementation.
-  return tp1 == tp2;
-}
-
 template <typename TensorReshapeOp, bool isExpansion = std::is_same<
                                         TensorReshapeOp, ExpandShapeOp>::value>
 static LogicalResult verifyTensorReshapeOp(TensorReshapeOp op,
                                            RankedTensorType expandedType,
                                            RankedTensorType collapsedType) {
   if (failed(
           verifyReshapeLikeTypes(op, expandedType, collapsedType, isExpansion)))
     return failure();
 
   auto maps = op.getReassociationMaps();
   RankedTensorType expectedType =
       CollapseShapeOp::inferCollapsedType(expandedType, maps);
-  if (!isSameTypesWithoutEncoding(collapsedType, expectedType))
+  if (!isSameTypeWithoutEncoding(collapsedType, expectedType))
     return op.emitOpError("expected collapsed type to be ")
            << expectedType << ", but got " << collapsedType;
   return success();
 }
 
 LogicalResult ExpandShapeOp::verify() {
   auto srcType = getSrcType();
   auto resultType = getResultType();

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>
+}