[mlir][sparse] Replace the folding of nop convert with a codegen rule.

[bixia]

Summary

This is to allow the use of a nop convert to express that the sparse tensor
allocated through bufferization::AllocTensorOp will be expanded to sparse
tensor storage by sparse tensor codegen.

[bixia]

Rebase. Fix format.

[aartbik]

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

724

we now have a difference between codegen ("folded") and conversion (not "folded").
You probably want to add it to conversion too as new case (and test), so we are consistent again

mlir/test/Dialect/SparseTensor/fold.mlir

4–5

please remove this test completely from this file
(fold.mlir is to test folding, and we no longer do that)

[bixia]

Remove trivial convert from fold.mlir.

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

724

Actually, the conversion pass already folds this, see code

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td

string description = [{
 
     // The following conversion is not allowed (since it would require a
     // runtime assertion that the source's dimension size is actually 100).
     %4 = sparse_tensor.convert %d : tensor<?xf64> to tensor<100xf64, #SV>
     ```
 
   }];
   let assemblyFormat = "$source attr-dict `:` type($source) `to` type($dest)";
-  let hasFolder = 1;
   let hasVerifier = 1;
 }
 
 def SparseTensor_ToPointersOp : SparseTensor_Op<"pointers", [Pure]>,
     Arguments<(ins AnySparseTensor:$tensor, IndexAttr:$dimension)>,
     Results<(outs AnyStridedMemRefOfRank<1>:$result)> {
   let summary = "Extracts pointers array at given dimension from a tensor";
   let description = [{

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

if (auto tp2 = getDest().getType().dyn_cast<RankedTensorType>()) {
         if (shape1[d] != shape2[d] && shape2[d] != ShapedType::kDynamicSize)
           return emitError("unexpected conversion mismatch in dimension ") << d;
       return success();
     }
   }
   return emitError("unexpected type in convert");
 }
 
-OpFoldResult ConvertOp::fold(ArrayRef<Attribute> operands) {
-  if (getType() == getSource().getType())
-    return getSource();
-  return {};
-}
-
 LogicalResult ToPointersOp::verify() {
   auto e = getSparseTensorEncoding(getTensor().getType());
   if (failed(isInBounds(getDimension().getZExtValue(), getTensor())))
     return emitError("requested pointers dimension out of bounds");
   if (failed(isMatchingWidth(getResult(), e.getPointerBitWidth())))
     return emitError("unexpected type for pointers");
   return success();
 }

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

public:
   // Callback for SparseGetterOpConverter.
   static unsigned getIndexForOp(UnrealizedConversionCastOp tuple,
                                 ToValuesOp /*op*/) {
     // The last field holds the value buffer.
     return tuple.getInputs().size() - 1;
   }
 };
 
+/// Sparse codegen rule for the convert operator.
+class SparseConvertConverter : public OpConversionPattern<ConvertOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    if (op.getType() != op.getSource().getType()) {
+      // This should be handled by rewriting before codegen.
+      return failure();
+    }
+    rewriter.replaceOp(op, adaptor.getSource());
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Sparse tensor type conversion into an actual buffer.
 //===----------------------------------------------------------------------===//
 
 mlir::SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
   addConversion([](Type type) { return type; });
   addConversion(convertSparseTensorType);
 }
 
 //===----------------------------------------------------------------------===//
 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
                                                RewritePatternSet &patterns) {
   patterns.add<SparseReturnConverter, SparseCallConverter, SparseDimOpConverter,
                SparseCastConverter, SparseTensorAllocConverter,
                SparseTensorDeallocConverter, SparseTensorLoadConverter,
                SparseExpandConverter, SparseCompressConverter,
                SparseInsertConverter, SparseToPointersConverter,
-               SparseToIndicesConverter, SparseToValuesConverter>(
-      typeConverter, patterns.getContext());
+               SparseToIndicesConverter, SparseToValuesConverter,
+               SparseConvertConverter>(typeConverter, patterns.getContext());
 }

mlir/test/Dialect/SparseTensor/codegen.mlir

 //       CHECK: %[[S2:.*]] = arith.index_cast %[[T3]] : index to i32
 //       CHECK: %[[T4:.*]] = sparse_tensor.push_back %[[A1]], %[[T0]], %[[S2]]
 //       CHECK: return %[[A0]], %[[A1]], %[[T4]], %[[T1]], %[[T2]] : memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf64>
 func.func @sparse_insert_typed(%arg0: tensor<128xf64, #SparseVector>, %arg1: index, %arg2: f64) -> tensor<128xf64, #SparseVector> {
   %0 = sparse_tensor.insert %arg2 into %arg0[%arg1] : tensor<128xf64, #SparseVector>
   %1 = sparse_tensor.load %0 hasInserts : tensor<128xf64, #SparseVector>
   return %1 : tensor<128xf64, #SparseVector>
 }
+
+// CHECK-LABEL:   func.func @sparse_nop_convert(
+//  CHECK-SAME:   %[[A0:.*]]: memref<1xindex>,
+//  CHECK-SAME:   %[[A1:.*]]: memref<3xindex>,
+//  CHECK-SAME:   %[[A2:.*]]: memref<?xi32>,
+//  CHECK-SAME:   %[[A3:.*]]: memref<?xi64>,
+//  CHECK-SAME:   %[[A4:.*]]: memref<?xf32>)
+//       CHECK:   return %[[A0]], %[[A1]], %[[A2]], %[[A3]], %[[A4]] : memref<1xindex>, memref<3xindex>, memref<?xi32>, memref<?xi64>, memref<?xf32>
+func.func @sparse_nop_convert(%arg0: tensor<?xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> {
+  %0 = sparse_tensor.convert %arg0 : tensor<?xf32, #SparseVector> to tensor<?xf32, #SparseVector>
+  return %0 : tensor<?xf32, #SparseVector>
+}

mlir/test/Dialect/SparseTensor/fold.mlir

 // RUN: mlir-opt %s  --canonicalize --cse | FileCheck %s
 
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
-// CHECK-LABEL: func @sparse_nop_convert(
-//  CHECK-SAME: %[[A:.*]]: tensor<64xf32, #sparse_tensor.encoding<{{{.*}}}>>)
-//   CHECK-NOT: sparse_tensor.convert
-//       CHECK: return %[[A]] : tensor<64xf32, #sparse_tensor.encoding<{{{.*}}}>>
-func.func @sparse_nop_convert(%arg0: tensor<64xf32, #SparseVector>) -> tensor<64xf32, #SparseVector> {
-  %0 = sparse_tensor.convert %arg0 : tensor<64xf32, #SparseVector> to tensor<64xf32, #SparseVector>
-  return %0 : tensor<64xf32, #SparseVector>
-}
-
 // CHECK-LABEL: func @sparse_dce_convert(
 //  CHECK-SAME: %[[A:.*]]: tensor<64xf32>)
 //   CHECK-NOT: sparse_tensor.convert
 //       CHECK: return
 func.func @sparse_dce_convert(%arg0: tensor<64xf32>) {
   %0 = sparse_tensor.convert %arg0 : tensor<64xf32> to tensor<64xf32, #SparseVector>
   return
 }