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