From: bixia1 Date: Tue, 10 Jan 2023 17:49:24 +0000 (-0800) Subject: [mlir][sparse] Improve the rewriting for dense-to-sparse conversion. X-Git-Tag: upstream/17.0.6~21362 X-Git-Url: http://review.tizen.org/git/?a=commitdiff_plain;h=14aba2084d977c0dbfeb289633ca86a6df636da3;p=platform%2Fupstream%2Fllvm.git [mlir][sparse] Improve the rewriting for dense-to-sparse conversion. Reviewed By: Peiming Differential Revision: https://reviews.llvm.org/D141335 --- diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp index 324c5b3..5b8ebdb 100644 --- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp +++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp @@ -695,22 +695,36 @@ private: } } - RankedTensorType cooTp = getUnorderedCOOFromType(dstTp); - auto cooBuffer = - rewriter.create(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(loc, bufferTp, dynSizes).getResult(); auto foreachOp = rewriter.create( - 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 indicesArray(rank, Value()); + for (uint64_t i = 0; i < rank; i++) + indicesArray[toStoredDim(encDst, i)] = indices[i]; if (fromSparseConst) { - input = builder.create(loc, v, input, indices); + input = builder.create(loc, v, input, indicesArray); } else { Value cond = genIsNonzero(builder, loc, v); auto ifOp = builder.create( loc, TypeRange(input.getType()), cond, /*else*/ true); builder.setInsertionPointToStart(&ifOp.getThenRegion().front()); - Value insert = builder.create(loc, v, input, indices); + Value insert = + builder.create(loc, v, input, indicesArray); builder.create(loc, insert); builder.setInsertionPointToStart(&ifOp.getElseRegion().front()); builder.create(loc, input); @@ -721,8 +735,12 @@ private: }); rewriter.setInsertionPointAfter(op); src = rewriter.create(loc, foreachOp.getResult(0), true); - rewriter.replaceOpWithNewOp(op, dstTp, src); - rewriter.create(loc, src); + if (bufferTp != dstTp) { + rewriter.replaceOpWithNewOp(op, dstTp, src); + rewriter.create(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 4eb16a6..6bdd5cd 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>) -> 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(