[mlir][sparse] add folding to sparse_tensor.convert

[aartbik]

Summary

folds conversion between identical types (with tests)

[aartbik]

rebase with diff in main

[ThomasRaoux]

mlir/test/Dialect/SparseTensor/conversion.mlir

115–121

this looks like the same test as in fold.mlir? Is that on purpose?

[aartbik]

mlir/test/Dialect/SparseTensor/conversion.mlir

115–121

Almost, but it has also converted the type into the opaque pointer! Just to make sure we don't fold "too late".

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

string description = [{
 
     ```mlir
     sparse_tensor.new %source : !Source to tensor<1024x1024xf64, #CSR>
     ```
   }];
   let assemblyFormat = "$source attr-dict `:` type($source) `to` type($result)";
 }
 
-def SparseTensor_ConvertOp : SparseTensor_Op<"convert", [SameOperandsAndResultType]>,
+def SparseTensor_ConvertOp : SparseTensor_Op<"convert",
+  [NoSideEffect, SameOperandsAndResultType]>,
     Arguments<(ins AnyTensor:$source)>,
     Results<(outs AnyTensor:$dest)> {
   string summary = "Converts between different tensor types";
   string description = [{
     Converts one sparse or dense tensor type to another tensor type. The rank
     and dimensions of the source and destination types must match exactly,
     only the sparse encoding of these types may be different. The name `convert`
     was preferred over `cast`, since the operation may incur a non-trivial cost.
string description = [{
     ```mlir
     %0 = sparse_tensor.convert %1 : tensor<32x32xf32> to tensor<32x32xf32, #CSR>
 
     %2 = sparse_tensor.convert %3 : tensor<8x8xi32, #CSC> to tensor<8x8xi32, #CSR>
     ```
 
   }];
   let assemblyFormat = "$source attr-dict `:` type($source) `to` type($dest)";
+  let hasFolder = 1;
 }
 
 def SparseTensor_ToPointersOp : SparseTensor_Op<"pointers", [NoSideEffect]>,
     Arguments<(ins AnyTensor:$tensor, Index:$dim)>,
     Results<(outs AnyStridedMemRefOfRank<1>:$result)> {
   let summary = "Extract pointers array at given dimension from a tensor";
   let description = [{
     Returns the pointers array of the sparse storage scheme at the

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

if (auto tp2 = op.dest().getType().dyn_cast<RankedTensorType>()) {
                  << "unexpected conversion mismatch in dimension " << d;
       }
       return success();
     }
   }
   return op.emitError("unexpected type in convert");
 }
 
+OpFoldResult ConvertOp::fold(ArrayRef<Attribute> operands) {
+  if (getType() == source().getType())
+    return source();
+  return {};
+}
+
 static LogicalResult verify(ToPointersOp op) {
   if (auto e = getSparseTensorEncoding(op.tensor().getType())) {
     if (failed(isInBounds(op.dim(), op.tensor())))
       return op.emitError("requested pointers dimension out of bounds");
     if (failed(isMatchingWidth(op.result(), e.getPointerBitWidth())))
       return op.emitError("unexpected type for pointers");
     return success();
   }

mlir/test/Dialect/SparseTensor/conversion.mlir

 //   CHECK-DAG: %[[Z:.*]] = tensor.cast %[[W]] : tensor<3xi64> to tensor<?xi64>
 //       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
 func @sparse_new3d(%arg0: !llvm.ptr<i8>) -> tensor<?x?x?xf32, #SparseTensor> {
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?x?xf32, #SparseTensor>
   return %0 : tensor<?x?x?xf32, #SparseTensor>
 }
 
+// CHECK-LABEL: func @sparse_nop_convert(
+//  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8>
+//       CHECK: return %[[A]] : !llvm.ptr<i8>
+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_convert_1d(
 //  CHECK-SAME: %[[A:.*]]: tensor<?xi32>) -> !llvm.ptr<i8>
 //   CHECK-DAG: %[[C0:.*]] = constant 0 : index
 //   CHECK-DAG: %[[C1:.*]] = constant 1 : index
 //   CHECK-DAG: %[[D0:.*]] = constant dense<0> : tensor<1xi64>
 //   CHECK-DAG: %[[D1:.*]] = constant dense<1> : tensor<1xi8>
 //   CHECK-DAG: %[[X:.*]] = tensor.cast %[[D1]] : tensor<1xi8> to tensor<?xi8>
 //   CHECK-DAG: %[[Y:.*]] = tensor.cast %[[D0]] : tensor<1xi64> to tensor<?xi64>

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