[mlir][sparse] add sparse tensor type conversion operation

[aartbik]

Summary

Introduces a conversion from one (sparse) tensor type to another
(sparse) tensor type. See the operation doc for details. Actual
codegen for all cases is still TBD.

[ThomasRaoux]

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

59

It can also convert from dense to sparse and sparse to dense right? The summary should be updated too.

[aartbik]

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

59

The explanation below talks about sparse vs non-sparse, but I tried to make it a bit more clear in description

[aartbik]

updated doc

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]>,
+    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.
+
+     When converting between two different sparse tensor types, only explicitly
+     stored values are moved from one underlying sparse storage format to
+     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.
+
+    Examples:
+
+    ```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)";
+}
+
 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
      given dimension for the given sparse tensor. This is similar to the
      `memref.buffer_cast` operation in the sense that it provides a bridge

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

 static LogicalResult isMatchingWidth(Value result, unsigned width) {
   Type etp = result.getType().cast<MemRefType>().getElementType();
   if ((width == 0 && etp.isIndex()) || (width > 0 && etp.isInteger(width)))
     return success();
   return failure();
 }
 
 static LogicalResult verify(NewOp op) {
-  if (!getSparseTensorEncoding(op.getResult().getType()))
+  if (!getSparseTensorEncoding(op.result().getType()))
     return op.emitError("expected a sparse tensor result");
   return success();
 }
 
+static LogicalResult verify(ConvertOp op) {
+  if (auto tp1 = op.source().getType().dyn_cast<RankedTensorType>()) {
+    if (auto tp2 = op.dest().getType().dyn_cast<RankedTensorType>()) {
+      assert(tp1.getRank() == tp2.getRank());
+      auto shape1 = tp1.getShape();
+      auto shape2 = tp2.getShape();
+      for (unsigned d = 0, rank = tp1.getRank(); d < rank; d++)
+        if (shape1[d] != shape2[d])
+          return op.emitError()
+                 << "unexpected conversion mismatch in dimension " << d;
+      return success();
+    }
+  }
+  return op.emitError("unexpected type in convert");
+}
+
 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/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 namespace {
 
-/// Internal encoding of primary storage. Keep this enum consistent
-/// with the equivalent enum in the sparse runtime support library.
-enum PrimaryTypeEnum : uint64_t {
-  kF64 = 1,
-  kF32 = 2,
-  kI64 = 3,
-  kI32 = 4,
-  kI16 = 5,
-  kI8 = 6
-};
+/// Returns internal type encoding for primary storage. Keep these
+/// values consistent with the sparse runtime support library.
+static unsigned getPrimaryTypeEncoding(Type tp) {
+  if (tp.isF64())
+    return 1;
+  if (tp.isF32())
+    return 2;
+  if (tp.isInteger(64))
+    return 3;
+  if (tp.isInteger(32))
+    return 4;
+  if (tp.isInteger(16))
+    return 5;
+  if (tp.isInteger(8))
+    return 6;
+  return 0;
+}
 
 /// Returns internal type encoding for overhead storage. Keep these
 /// values consistent with the sparse runtime support library.
 static unsigned getOverheadTypeEncoding(unsigned width) {
   switch (width) {
   default:
     return 1;
   case 32:
matchAndRewrite(NewOp op, ArrayRef<Value> operands,
     } else {
       for (unsigned i = 0; i < sz; i++)
         perm[i] = APInt(64, i);
     }
     params.push_back(getTensor(rewriter, 64, loc, perm));
     // Secondary and primary types encoding.
     unsigned secPtr = getOverheadTypeEncoding(enc.getPointerBitWidth());
     unsigned secInd = getOverheadTypeEncoding(enc.getIndexBitWidth());
-    unsigned primary;
-    if (eltType.isF64())
-      primary = kF64;
-    else if (eltType.isF32())
-      primary = kF32;
-    else if (eltType.isInteger(64))
-      primary = kI64;
-    else if (eltType.isInteger(32))
-      primary = kI32;
-    else if (eltType.isInteger(16))
-      primary = kI16;
-    else if (eltType.isInteger(8))
-      primary = kI8;
-    else
+    unsigned primary = getPrimaryTypeEncoding(eltType);
+    if (!primary)
       return failure();
     params.push_back(
         rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(secPtr)));
     params.push_back(
         rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(secInd)));
     params.push_back(
         rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(primary)));
     // Generate the call to create new tensor.
     Type ptrType = LLVM::LLVMPointerType::get(IntegerType::get(context, 8));
     StringRef name = "newSparseTensor";
     rewriter.replaceOpWithNewOp<CallOp>(
         op, ptrType, getFunc(op, name, ptrType, params), params);
     return success();
   }
 };
 
+/// Sparse conversion rule for the convert operator.
+class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(ConvertOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    // TODO: implement conversions lowering
+    return failure();
+  }
+};
+
 /// Sparse conversion rule for pointer accesses.
 class SparseTensorToPointersConverter
     : public OpConversionPattern<ToPointersOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToPointersOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {

 
 } // namespace
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorConversionPatterns(TypeConverter &typeConverter,
                                                   RewritePatternSet &patterns) {
   patterns.add<SparseReturnConverter, SparseTensorToDimSizeConverter,
-               SparseTensorNewConverter, SparseTensorToPointersConverter,
-               SparseTensorToIndicesConverter, SparseTensorToValuesConverter,
-               SparseTensorToTensorConverter>(typeConverter,
-                                              patterns.getContext());
+               SparseTensorNewConverter, SparseTensorConvertConverter,
+               SparseTensorToPointersConverter, SparseTensorToIndicesConverter,
+               SparseTensorToValuesConverter, SparseTensorToTensorConverter>(
+      typeConverter, patterns.getContext());
 }

mlir/test/Dialect/SparseTensor/invalid.mlir

 
 // -----
 
 func @sparse_to_unannotated_tensor(%arg0: memref<?xf64>) -> tensor<16x32xf64> {
   // expected-error@+1 {{expected a sparse tensor as result}}
   %0 = sparse_tensor.tensor %arg0 : memref<?xf64> to tensor<16x32xf64>
   return %0 : tensor<16x32xf64>
 }
+
+// -----
+
+func @sparse_convert_unranked(%arg0: tensor<*xf32>) -> tensor<10xf32> {
+  // expected-error@+1 {{unexpected type in convert}}
+  %0 = sparse_tensor.convert %arg0 : tensor<*xf32> to tensor<10xf32>
+  return %0 : tensor<10xf32>
+}
+
+// -----
+
+#CSR = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"]}>
+
+func @sparse_convert_mismatch(%arg0: tensor<10x10xf32>) -> tensor<10x?xf32, #CSR> {
+  // expected-error@+1 {{unexpected conversion mismatch in dimension 1}}
+  %0 = sparse_tensor.convert %arg0 : tensor<10x10xf32> to tensor<10x?xf32, #CSR>
+  return %0 : tensor<10x?xf32, #CSR>
+}

mlir/test/Dialect/SparseTensor/roundtrip.mlir

func @sparse_new(%arg0: !llvm.ptr<i8>) -> tensor<128xf64, #SparseVector> {
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<128xf64, #SparseVector>
   return %0 : tensor<128xf64, #SparseVector>
 }
 
 // -----
 
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
+// CHECK-LABEL: func @sparse_convert_1d_to_sparse(
+// CHECK-SAME: %[[A:.*]]: tensor<64xf32>)
+//       CHECK: %[[T:.*]] = sparse_tensor.convert %[[A]] : tensor<64xf32> to tensor<64xf32, #{{.*}}>
+//       CHECK: return %[[T]] : tensor<64xf32, #{{.*}}>
+func @sparse_convert_1d_to_sparse(%arg0: tensor<64xf32>) -> tensor<64xf32, #SparseVector> {
+  %0 = sparse_tensor.convert %arg0 : tensor<64xf32> to tensor<64xf32, #SparseVector>
+  return %0 : tensor<64xf32, #SparseVector>
+}
+
+// -----
+
+#SparseTensor = #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>
+
+// CHECK-LABEL: func @sparse_convert_3d_from_sparse(
+// CHECK-SAME: %[[A:.*]]: tensor<8x8x8xf64, #{{.*}}>)
+//       CHECK: %[[T:.*]] = sparse_tensor.convert %[[A]] : tensor<8x8x8xf64, #{{.*}}> to tensor<8x8x8xf64>
+//       CHECK: return %[[T]] : tensor<8x8x8xf64>
+func @sparse_convert_3d_from_sparse(%arg0: tensor<8x8x8xf64, #SparseTensor>) -> tensor<8x8x8xf64> {
+  %0 = sparse_tensor.convert %arg0 : tensor<8x8x8xf64, #SparseTensor> to tensor<8x8x8xf64>
+  return %0 : tensor<8x8x8xf64>
+}
+
+// -----
+
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+
 // CHECK-LABEL: func @sparse_pointers(
 //  CHECK-SAME: %[[A:.*]]: tensor<128xf64, #{{.*}}>)
 //       CHECK: %[[C:.*]] = constant 0 : index
 //       CHECK: %[[T:.*]] = sparse_tensor.pointers %[[A]], %[[C]] : tensor<128xf64, #{{.*}}> to memref<?xindex>
 //       CHECK: return %[[T]] : memref<?xindex>
 func @sparse_pointers(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xindex> {
   %c = constant 0 : index
   %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector> to memref<?xindex>