diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp @@ -1058,9 +1058,11 @@ SparseTensorEncodingAttr encDst = getSparseTensorEncoding(op.getType()); SparseTensorEncodingAttr encSrc = getSparseTensorEncoding(op.getSource().getType()); - // Different encoding (except for different bitwidth) should be handled by - // rewriting. - if (encDst.withoutBitWidths() != encSrc.withoutBitWidths()) { + // Different encoding (except for different bitwidth or input is a slice) + // should be handled by rewriting. + assert(!encDst.isSlice() && "Can not convert to a sparse tensor slices."); + if (encDst.withoutBitWidths() != encSrc.withoutBitWidths() || + encSrc.isSlice()) { return failure(); } diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp @@ -618,7 +618,7 @@ 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. diff --git a/mlir/test/Dialect/SparseTensor/convert_sparse2sparse.mlir b/mlir/test/Dialect/SparseTensor/convert_sparse2sparse.mlir --- a/mlir/test/Dialect/SparseTensor/convert_sparse2sparse.mlir +++ b/mlir/test/Dialect/SparseTensor/convert_sparse2sparse.mlir @@ -25,6 +25,10 @@ dimLevelType = ["compressed"] }> +#SortedCOO2D = #sparse_tensor.encoding<{ + dimLevelType = [ "compressed-nu", "singleton" ], +}> + #SortedCOO3D = #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton-nu", "singleton" ] @@ -35,6 +39,11 @@ 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) -> !llvm.ptr // CHECK: return %[[A]] : !llvm.ptr @@ -185,3 +194,20 @@ %0 = sparse_tensor.convert %arg0 : tensor to tensor return %0 : tensor } + +// 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> +}