[mlir][sparse] convert a sparse tensor slice to sparse tensor correctly.

[Peiming]

minor fixes.

[aartbik]

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

1064

space after closing )

1065

add comment on slice properties (of source and dest)

1066–1068

Cannot

[wrengr]

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

1064

You should leave the space there

1066–1068

Does ConvertOp::verify actually guarantee this condition? If not, then this should be a failure or emitError rather than an assert

[Peiming]

address comments.

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

LogicalResult UnpackOp::verify() {
   return verifyPackUnPack(*this, false, srcTp, valuesTp, coordinatesTp);
 }
 
 LogicalResult ConvertOp::verify() {
   if (auto tp1 = getSource().getType().dyn_cast<RankedTensorType>()) {
     if (auto tp2 = getDest().getType().dyn_cast<RankedTensorType>()) {
       if (tp1.getRank() != tp2.getRank())
         return emitError("unexpected conversion mismatch in rank");
+      auto dstEnc =
+          tp2.getEncoding().dyn_cast_or_null<SparseTensorEncodingAttr>();
+      if (dstEnc && dstEnc.isSlice())
+        return emitError("cannot convert to a sparse tensor slice");
+
       auto shape1 = tp1.getShape();
       auto shape2 = tp2.getShape();
       // Accept size matches between the source and the destination type
       // (e.g. 10 vs. 10, 10 vs. ?, or ? vs. ?), but reject direct mismatches or
       // matches that would need a runtime assert (e.g. 10 vs. 20 or ? vs. 10).
       for (Dimension d = 0, dimRank = tp1.getRank(); d < dimRank; d++)
         if (shape1[d] != shape2[d] && shape2[d] != ShapedType::kDynamic)
           return emitError("unexpected conversion mismatch in dimension ") << d;

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

 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     SparseTensorEncodingAttr encDst = getSparseTensorEncoding(op.getType());
     SparseTensorEncodingAttr encSrc =
         getSparseTensorEncoding(op.getSource().getType());
+    // The output tensor can not be a slice and those cases should have been
+    // rejected by ConvertOp::verify() already.
+    assert(!encDst.isSlice() && "Cannot convert to a sparse tensor slices.");
     // Different encoding (except for different bitwidth) should be handled by
     // rewriting.
-    if (encDst.withoutBitWidths() != encSrc.withoutBitWidths()) {
+    // We need further rewrites if the input tensor is a slice too.
+    if (encDst.withoutBitWidths() != encSrc.withoutBitWidths() ||
+        encSrc.isSlice()) {
       return failure();
     }
 
     Type retElemTp = op.getResult().getType().getElementType();
     Type srcElemTp = op.getSource().getType().getElementType();
     // Fold the trivial cases.
     if (retElemTp == srcElemTp && encDst == encSrc) {
       rewriter.replaceOp(op, adaptor.getSource());

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

 
 /// Sparse rewriting rule for the convert operator.
 struct ConvertRewriter : public OpRewritePattern<ConvertOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(ConvertOp op,
                                 PatternRewriter &rewriter) const override {
     auto encDst = getSparseTensorEncoding(op.getType());
     auto encSrc = getSparseTensorEncoding(op.getSource().getType());
-    if (encDst && encSrc &&
+    if (encDst && encSrc && !encSrc.isSlice() &&
         encSrc.withoutBitWidths() == encDst.withoutBitWidths()) {
       // Trivial tensor conversion and simple element type conversion is handled
       // in codegen.
       return failure();
     }
     // TODO: Add a cast before generating InsertOp.
     assert(op.getSource().getType().getElementType() ==
            op.getDest().getType().getElementType());

mlir/test/Dialect/SparseTensor/convert_sparse2sparse.mlir

#SparseVector32 = #sparse_tensor.encoding<{
   posWidth = 32,
   crdWidth = 32
 }>
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"]
 }>
 
+#SortedCOO2D = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton" ],
+}>
+
 #SortedCOO3D = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed-nu", "singleton-nu", "singleton" ]
 
 }>
 
 #TsssPermuted = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed", "compressed" ],
   dimOrdering = affine_map<(i,j,k) -> (k,i,j)>
 }>
 
+#COOSlice = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed-nu", "singleton" ],
+  slice = [ (2, 2, 1), (12, 13, 1) ]
+}>
+
 // CHECK-LABEL: func @sparse_nop_convert(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8>
 //       CHECK: return %[[A]] : !llvm.ptr<i8>
 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-RWT: bufferization.dealloc_tensor %[[VAL_16]]
 //       CHECK-RWT: %[[VAL_28:.*]] = sparse_tensor.load %[[VAL_29:.*]] hasInserts
 //       CHECK-RWT: %[[VAL_30:.*]] = sparse_tensor.convert %[[VAL_28]]
 //       CHECK-RWT: return %[[VAL_30]]
 func.func @sparse_convert_permuted(%arg0: tensor<?x?x?xf32, #SortedCOO3D>) -> tensor<?x?x?xf32, #TsssPermuted> {
   %0 = sparse_tensor.convert %arg0 : tensor<?x?x?xf32, #SortedCOO3D> to tensor<?x?x?xf32, #TsssPermuted>
   return %0 : tensor<?x?x?xf32, #TsssPermuted>
 }
+
+// CHECK-RWT-LABEL: func.func @sparse_convert_slice(
+//  CHECK-RWT-SAME: %[[VAL_0:.*]]: tensor<2x13xi32, #{{.*}}>) -> tensor<2x13xi32, #{{.*}}> {
+//       CHECK-RWT: %[[VAL_1:.*]] = sparse_tensor.number_of_entries %[[VAL_0]] : tensor<2x13xi32, #{{.*}}>
+//       CHECK-RWT: %[[VAL_2:.*]] = bufferization.alloc_tensor() size_hint=%[[VAL_1]] : tensor<2x13xi32, #{{.*}}>
+//       CHECK-RWT: %[[VAL_3:.*]] = sparse_tensor.foreach in %[[VAL_0]] init(%[[VAL_2]]) : tensor<2x13xi32, #{{.*}}>, tensor<2x13xi32, #{{.*}}> -> tensor<2x13xi32, #{{.*}}> do {
+//       CHECK-RWT: ^bb0(%[[VAL_4:.*]]: index, %[[VAL_5:.*]]: index, %[[VAL_6:.*]]: i32, %[[VAL_7:.*]]: tensor<2x13xi32, #{{.*}}>):
+//       CHECK-RWT:   %[[VAL_8:.*]] = sparse_tensor.insert %[[VAL_6]] into %[[VAL_7]]{{\[}}%[[VAL_4]], %[[VAL_5]]] : tensor<2x13xi32, #{{.*}}>
+//       CHECK-RWT:   sparse_tensor.yield %[[VAL_8]] : tensor<2x13xi32, #{{.*}}>
+//       CHECK-RWT: }
+//       CHECK-RWT: %[[VAL_9:.*]] = sparse_tensor.load %[[VAL_10:.*]] hasInserts : tensor<2x13xi32, #{{.*}}>
+//       CHECK-RWT: %[[VAL_11:.*]] = sparse_tensor.convert %[[VAL_9]] : tensor<2x13xi32, #{{.*}}> to tensor<2x13xi32, #{{.*}}>
+//       CHECK-RWT: return %[[VAL_11]] : tensor<2x13xi32, #{{.*}}>
+func.func @sparse_convert_slice(%arg0: tensor<2x13xi32, #COOSlice>) -> (tensor<2x13xi32, #SortedCOO2D>)  {
+  %0 = sparse_tensor.convert %arg0 : tensor<2x13xi32, #COOSlice> to tensor<2x13xi32, #SortedCOO2D>
+  return %0 : tensor<2x13xi32, #SortedCOO2D>
+}

mlir/test/Dialect/SparseTensor/invalid.mlir

 func.func @invalid_out_dense(%arg0: tensor<10xf64>, %arg1: !llvm.ptr<i8>) {
   // expected-error@+1 {{'sparse_tensor.out' op operand #0 must be sparse tensor of any type values, but got 'tensor<10xf64>'}}
   sparse_tensor.out %arg0, %arg1 : tensor<10xf64>, !llvm.ptr<i8>
   return
 }
 
 // -----
 
+#CSR = #sparse_tensor.encoding<{
+  dimLevelType = ["dense", "compressed"],
+  slice = [ (1, 4, 1), (1, 4, 2) ]
+}>
+
+func.func @sparse_convert_to_slice(%arg0: tensor<10x?xf32>) -> tensor<10x10xf32, #CSR> {
+  // expected-error@+1 {{cannot convert to a sparse tensor slice}}
+  %0 = sparse_tensor.convert %arg0 : tensor<10x?xf32> to tensor<10x10xf32, #CSR>
+  return %0 : tensor<10x10xf32, #CSR>
+}
+
+// -----
+
 func.func @invalid_binary_num_args_mismatch_overlap(%arg0: f64, %arg1: f64) -> f64 {
   // expected-error@+1 {{overlap region must have exactly 2 arguments}}
   %r = sparse_tensor.binary %arg0, %arg1 : f64, f64 to f64
     overlap={
       ^bb0(%x: f64):
         sparse_tensor.yield %x : f64
     }
     left={}