[mlir][sparse] Fix the codegen for the convert operator to handle hidden nop convert.

[bixia]

Partially restore the folder to fold dense-to-dense convert.

[bixia]

Fix comment.

[bixia]

Fix comment.

[bixia]

Add description about trivial convert op in the td file.

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

string description = [{
     the other. When converting from an unannotated dense tensor type to a
     sparse tensor type, an explicit test for nonzero values is used. When
     converting to an unannotated dense tensor type, implicit zeroes in the
     sparse storage format are made explicit. Note that the conversions can have
     non-trivial costs associated with them, since they may involve elaborate
     data structure transformations. Also, conversions from sparse tensor types
     into dense tensor types may be infeasible in terms of storage requirements.
 
+    Trivial dense-to-dense convert will be removed by canonicalization while
+    trivial sparse-to-sparse convert will be removed by the sparse codegen. This
+    is because we use trivial sparse-to-sparse convert to tell bufferization
+    that the sparse codegen will expand the tensor buffer into sparse tensor
+    storage.
+
     Examples:
 
     ```mlir
     %0 = sparse_tensor.convert %a : tensor<32x32xf32> to tensor<32x32xf32, #CSR>
     %1 = sparse_tensor.convert %a : tensor<32x32xf32> to tensor<?x?xf32, #CSR>
     %2 = sparse_tensor.convert %b : tensor<8x8xi32, #CSC> to tensor<8x8xi32, #CSR>
     %3 = sparse_tensor.convert %c : tensor<4x8xf64, #CSR> to tensor<4x?xf64, #CSC>
 
     // 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) {
+  Type dstType = getType();
+  // Fold trivial dense-to-dense convert and leave trivial sparse-to-sparse
+  // convert for codegen to remove. This is because we use trivial
+  // sparse-to-sparse convert to tell bufferization that the sparse codegen
+  // will expand the tensor buffer into sparse tensor storage.
+  if (!getSparseTensorEncoding(dstType) && dstType == 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

 
 /// 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()) {
+    SparseTensorEncodingAttr encDst = getSparseTensorEncoding(op.getType());
+    SparseTensorEncodingAttr encSrc =
+        getSparseTensorEncoding(op.getSource().getType());
+    if (encDst != encSrc) {
       // This should be handled by rewriting before codegen.
       return failure();
     }
     rewriter.replaceOp(op, adaptor.getSource());
     return success();
   }
 };
 

mlir/test/Dialect/SparseTensor/codegen.mlir

 
 // 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>
+func.func @sparse_nop_convert(%arg0: tensor<32xf32, #SparseVector>) -> tensor<?xf32, #SparseVector> {
+  %0 = sparse_tensor.convert %arg0 : tensor<32xf32, #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_dense2dense_convert(
+//  CHECK-SAME: %[[A:.*]]: tensor<64xf32>)
+//   CHECK-NOT: sparse_tensor.convert
+//       CHECK: return %[[A]] : tensor<64xf32>
+func.func @sparse_nop_dense2dense_convert(%arg0: tensor<64xf32>) -> tensor<64xf32> {
+  %0 = sparse_tensor.convert %arg0 : tensor<64xf32> to tensor<64xf32>
+  return %0 : tensor<64xf32>
+}
+
 // 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
 }