[Mlir-commits] [mlir] 14aba20 - [mlir][sparse] Improve the rewriting for dense-to-sparse conversion.
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jan 11 08:27:10 PST 2023
Author: bixia1
Date: 2023-01-11T08:27:04-08:00
New Revision: 14aba2084d977c0dbfeb289633ca86a6df636da3
URL: https://github.com/llvm/llvm-project/commit/14aba2084d977c0dbfeb289633ca86a6df636da3
DIFF: https://github.com/llvm/llvm-project/commit/14aba2084d977c0dbfeb289633ca86a6df636da3.diff
LOG: [mlir][sparse] Improve the rewriting for dense-to-sparse conversion.
Reviewed By: Peiming
Differential Revision: https://reviews.llvm.org/D141335
Added:
Modified:
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
mlir/test/Dialect/SparseTensor/convert_dense2sparse.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
index 324c5b32f505..5b8ebdb08ace 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
@@ -695,22 +695,36 @@ struct ConvertRewriter : public OpRewritePattern<ConvertOp> {
}
}
- RankedTensorType cooTp = getUnorderedCOOFromType(dstTp);
- auto cooBuffer =
- rewriter.create<AllocTensorOp>(loc, cooTp, dynSizes).getResult();
+ SparseTensorEncodingAttr encDst = getSparseTensorEncoding(dstTp);
+ // We don't need a temporary COO tensor if the destination has an identity
+ // ordering. Otherwise, we use the destination ordering for the temporary
+ // COO tensor.
+ // TODO: enhance foreachOp to take ordering to remove the need of a
+ // temporary COO tensor here.
+ RankedTensorType bufferTp = encDst.hasIdDimOrdering()
+ ? dstTp
+ : getUnorderedCOOFromTypeWithOrdering(
+ dstTp, encDst.getDimOrdering());
+ auto buffer =
+ rewriter.create<AllocTensorOp>(loc, bufferTp, dynSizes).getResult();
auto foreachOp = rewriter.create<ForeachOp>(
- loc, src, cooBuffer,
+ loc, src, buffer,
[&](OpBuilder &builder, Location loc, ValueRange indices, Value v,
ValueRange reduc) {
Value input = reduc.front();
+ uint64_t rank = dstTp.getRank();
+ SmallVector<Value> indicesArray(rank, Value());
+ for (uint64_t i = 0; i < rank; i++)
+ indicesArray[toStoredDim(encDst, i)] = indices[i];
if (fromSparseConst) {
- input = builder.create<InsertOp>(loc, v, input, indices);
+ input = builder.create<InsertOp>(loc, v, input, indicesArray);
} else {
Value cond = genIsNonzero(builder, loc, v);
auto ifOp = builder.create<scf::IfOp>(
loc, TypeRange(input.getType()), cond, /*else*/ true);
builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
- Value insert = builder.create<InsertOp>(loc, v, input, indices);
+ Value insert =
+ builder.create<InsertOp>(loc, v, input, indicesArray);
builder.create<scf::YieldOp>(loc, insert);
builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
builder.create<scf::YieldOp>(loc, input);
@@ -721,8 +735,12 @@ struct ConvertRewriter : public OpRewritePattern<ConvertOp> {
});
rewriter.setInsertionPointAfter(op);
src = rewriter.create<LoadOp>(loc, foreachOp.getResult(0), true);
- rewriter.replaceOpWithNewOp<ConvertOp>(op, dstTp, src);
- rewriter.create<DeallocTensorOp>(loc, src);
+ if (bufferTp != dstTp) {
+ rewriter.replaceOpWithNewOp<ConvertOp>(op, dstTp, src);
+ rewriter.create<DeallocTensorOp>(loc, src);
+ } else {
+ rewriter.replaceOp(op, src);
+ }
return success();
}
diff --git a/mlir/test/Dialect/SparseTensor/convert_dense2sparse.mlir b/mlir/test/Dialect/SparseTensor/convert_dense2sparse.mlir
index 4eb16a6ba3f8..6bdd5cd34694 100644
--- a/mlir/test/Dialect/SparseTensor/convert_dense2sparse.mlir
+++ b/mlir/test/Dialect/SparseTensor/convert_dense2sparse.mlir
@@ -10,6 +10,11 @@
dimLevelType = ["dense", "compressed"]
}>
+#CSC = #sparse_tensor.encoding<{
+ dimLevelType = [ "dense", "compressed" ],
+ dimOrdering = affine_map<(i, j) -> (j, i)>
+}>
+
#SparseTensor = #sparse_tensor.encoding<{
dimLevelType = ["dense", "compressed", "compressed"],
dimOrdering = affine_map<(i,j,k) -> (k,i,j)>
@@ -121,21 +126,8 @@ func.func @sparse_convert_complex(%arg0: tensor<100xcomplex<f64>>) -> tensor<100
// CHECK-RWT: }
// CHECK-RWT: sparse_tensor.yield %[[IFR]]
// CHECK-RWT: }
-// CHECK-RWT: %[[COO:.*]] = sparse_tensor.load %[[T2]] hasInserts
-// CHECK-RWT: %[[NNZ:.*]] = sparse_tensor.number_of_entries %[[COO]]
-// CHECK-RWT: %[[V:.*]] = sparse_tensor.values %[[COO]]
-// CHECK-RWT: %[[I:.*]] = sparse_tensor.indices_buffer %[[COO]]
-// CHECK-RWT: sparse_tensor.sort_coo %[[NNZ]], %[[I]] jointly %[[V]] {nx = 2 : index, ny = 0 : index}
-// CHECK-RWT: %[[T3:.*]] = bufferization.alloc_tensor()
-// CHECK-RWT: %[[T4:.*]] = sparse_tensor.foreach in %[[COO]] init(%[[T3]])
-// CHECK-RWT: ^bb0(%[[L1I0:.*]]: index, %[[L1I1:.*]]: index, %[[L1V:.*]]: f64, %[[L1T:.*]]: tensor
-// CHECK-RWT: %[[L1T2:.*]] = sparse_tensor.insert %[[L1V]] into %[[L1T]]{{\[}}%[[L1I0]], %[[L1I1]]]
-// CHECK-RWT: sparse_tensor.yield %[[L1T2]]
-// CHECK-RWT: }
-// CHECK-RWT: %[[T5:.*]] = sparse_tensor.load %[[T4]] hasInserts
-// CHECK-RWT: %[[T6:.*]] = sparse_tensor.convert %[[T5]]
-// CHECK-RWT: bufferization.dealloc_tensor %[[COO]]
-// CHECK-RWT: return %[[T6]]
+// CHECK-RWT: %[[R:.*]] = sparse_tensor.load %[[T2]] hasInserts
+// CHECK-RWT: return %[[R]]
// CHECK-RWT: }
func.func @sparse_convert_2d(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #CSR> {
%0 = sparse_tensor.convert %arg0 : tensor<2x4xf64> to tensor<2x4xf64, #CSR>
@@ -180,6 +172,25 @@ func.func @sparse_convert_2d(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #CSR> {
// CHECK-RWT: %[[L0T2:.*]] = sparse_tensor.insert %[[L0V]] into %[[L0T]]{{\[}}%[[L0I0]], %[[L0I1]]]
// CHECK-RWT: sparse_tensor.yield %[[L0T2]]
// CHECK-RWT: }
+// CHECK-RWT: %[[R:.*]] = sparse_tensor.load %[[T1]] hasInserts
+// CHECK-RWT: return %[[R]]
+// CHECK-RWT: }
+func.func @sparse_constant() -> tensor<8x7xf32, #CSR>{
+ // Initialize a tensor.
+ %0 = arith.constant sparse<[[0, 0], [1, 6]], [1.0, 5.0]> : tensor<8x7xf32>
+ // Convert the tensor to a sparse tensor.
+ %1 = sparse_tensor.convert %0 : tensor<8x7xf32> to tensor<8x7xf32, #CSR>
+ return %1 : tensor<8x7xf32, #CSR>
+}
+
+// CHECK-RWT-LABEL: func.func @sparse_constant_csc() -> tensor<8x7xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)> }>> {
+// CHECK-RWT: %[[F0:.*]] = arith.constant sparse<{{\[\[}}0, 0], [1, 6]], [1.000000e+00, 5.000000e+00]> : tensor<8x7xf32>
+// CHECK-RWT: %[[T0:.*]] = bufferization.alloc_tensor()
+// CHECK-RWT: %[[T1:.*]] = sparse_tensor.foreach in %[[F0]] init(%[[T0]])
+// CHECK-RWT: ^bb0(%[[L0I0:.*]]: index, %[[L0I1:.*]]: index, %[[L0V:.*]]: f32, %[[L0T:.*]]: tensor
+// CHECK-RWT: %[[L0T2:.*]] = sparse_tensor.insert %[[L0V]] into %[[L0T]]{{\[}}%[[L0I1]], %[[L0I0]]]
+// CHECK-RWT: sparse_tensor.yield %[[L0T2]]
+// CHECK-RWT: }
// CHECK-RWT: %[[COO:.*]] = sparse_tensor.load %[[T1]] hasInserts
// CHECK-RWT: %[[NNZ:.*]] = sparse_tensor.number_of_entries %[[COO]]
// CHECK-RWT: %[[V:.*]] = sparse_tensor.values %[[COO]]
@@ -188,7 +199,7 @@ func.func @sparse_convert_2d(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #CSR> {
// CHECK-RWT: %[[T3:.*]] = bufferization.alloc_tensor()
// CHECK-RWT: %[[T4:.*]] = sparse_tensor.foreach in %[[COO]] init(%[[T3]])
// CHECK-RWT: ^bb0(%[[L1I0:.*]]: index, %[[L1I1:.*]]: index, %[[L1V:.*]]: f32, %[[L1T:.*]]: tensor
-// CHECK-RWT: %[[L1T2:.*]] = sparse_tensor.insert %[[L1V]] into %[[L1T]]{{\[}}%[[L1I0]], %[[L1I1]]]
+// CHECK-RWT: %[[L1T2:.*]] = sparse_tensor.insert %[[L1V]] into %[[L1T]]{{\[}}%[[L1I1]], %[[L1I0]]]
// CHECK-RWT: sparse_tensor.yield %[[L1T2]]
// CHECK-RWT: }
// CHECK-RWT: %[[T5:.*]] = sparse_tensor.load %[[T4]] hasInserts
@@ -196,12 +207,12 @@ func.func @sparse_convert_2d(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #CSR> {
// CHECK-RWT: bufferization.dealloc_tensor %[[COO]]
// CHECK-RWT: return %[[T6]]
// CHECK-RWT: }
-func.func @sparse_constant() -> tensor<8x7xf32, #CSR>{
+func.func @sparse_constant_csc() -> tensor<8x7xf32, #CSC>{
// Initialize a tensor.
%0 = arith.constant sparse<[[0, 0], [1, 6]], [1.0, 5.0]> : tensor<8x7xf32>
// Convert the tensor to a sparse tensor.
- %1 = sparse_tensor.convert %0 : tensor<8x7xf32> to tensor<8x7xf32, #CSR>
- return %1 : tensor<8x7xf32, #CSR>
+ %1 = sparse_tensor.convert %0 : tensor<8x7xf32> to tensor<8x7xf32, #CSC>
+ return %1 : tensor<8x7xf32, #CSC>
}
// CHECK-LABEL: func @sparse_convert_3d(
More information about the Mlir-commits
mailing list