[Mlir-commits] [mlir] [mlir][Bufferization] Do not have read semantics for destination of `tensor.parallel_insert_slice`. (PR #134169)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Apr 2 15:53:17 PDT 2025
https://github.com/MaheshRavishankar updated https://github.com/llvm/llvm-project/pull/134169
>From e167d941df444ffb4f1759564669be7c2be3826f Mon Sep 17 00:00:00 2001
From: MaheshRavishankar <mahesh.ravishankar at gmail.com>
Date: Tue, 1 Apr 2025 20:15:53 -0700
Subject: [PATCH 1/3] [mlir][Bufferization] Do not have read semantics for
destination of `tensor.parallel_insert_slice`.
Signed-off-by: MaheshRavishankar <mahesh.ravishankar at gmail.com>
---
.../Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
index 4ac6eca586961..31014172a9555 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -930,8 +930,7 @@ struct ParallelInsertSliceOpInterface
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
- return insertSliceOpRequiresRead(cast<tensor::ParallelInsertSliceOp>(op),
- opOperand);
+ return opOperand == cast<ParallelInsertSliceOp>(op).getSourceMutable();
}
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
>From 2cd1361c18f025cbe952d7a6d849f940bf391ebb Mon Sep 17 00:00:00 2001
From: MaheshRavishankar <mahesh.ravishankar at gmail.com>
Date: Wed, 2 Apr 2025 14:44:20 -0700
Subject: [PATCH 2/3] Make shared outs of scf.forall have read semantics.
Signed-off-by: MaheshRavishankar <mahesh.ravishankar at gmail.com>
---
.../BufferizableOpInterfaceImpl.cpp | 28 ++++---------------
1 file changed, 5 insertions(+), 23 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
index f48d2a2df9c3c..cf62ee8bc45b5 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -1186,18 +1186,6 @@ struct YieldOpInterface
}
};
-/// Return `true` if the given loop may have 0 iterations.
-bool mayHaveZeroIterations(scf::ForallOp forallOp) {
- for (auto [lb, ub] : llvm::zip(forallOp.getMixedLowerBound(),
- forallOp.getMixedUpperBound())) {
- std::optional<int64_t> lbConst = getConstantIntValue(lb);
- std::optional<int64_t> ubConst = getConstantIntValue(ub);
- if (!lbConst.has_value() || !ubConst.has_value() || *lbConst >= *ubConst)
- return true;
- }
- return false;
-}
-
/// Bufferization of ForallOp. This also bufferizes the terminator of the
/// region. There are op interfaces for the terminators (InParallelOp
/// and ParallelInsertSliceOp), but these are only used during analysis. Not
@@ -1207,17 +1195,11 @@ struct ForallOpInterface
ForallOp> {
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
- auto forallOp = cast<ForallOp>(op);
-
- // If the loop has zero iterations, the results of the op are their
- // corresponding shared_outs, meaning that the shared_outs bufferize to a
- // read.
- if (mayHaveZeroIterations(forallOp))
- return true;
-
- // scf::ForallOp alone doesn't bufferize to a memory read, one of the
- // uses of its matching bbArg may.
- return state.isValueRead(forallOp.getTiedBlockArgument(&opOperand));
+ // All tensor operands to `scf.forall` are `shared_outs` and all
+ // shared outs are assumed to be read by the loop. This does not
+ // account for the case where the entire value is over-written,
+ // but being conservative here.
+ return true;
}
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
>From c61a1c23ceecedb8b93e3a7e65ab3efa85e74d72 Mon Sep 17 00:00:00 2001
From: MaheshRavishankar <mahesh.ravishankar at gmail.com>
Date: Wed, 2 Apr 2025 15:44:41 -0700
Subject: [PATCH 3/3] Add test.
Signed-off-by: MaheshRavishankar <mahesh.ravishankar at gmail.com>
---
mlir/test/Dialect/SCF/one-shot-bufferize.mlir | 35 +++++++++++++++++++
1 file changed, 35 insertions(+)
diff --git a/mlir/test/Dialect/SCF/one-shot-bufferize.mlir b/mlir/test/Dialect/SCF/one-shot-bufferize.mlir
index bb9f7dfdba83f..a1067ec3ba05f 100644
--- a/mlir/test/Dialect/SCF/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/SCF/one-shot-bufferize.mlir
@@ -946,3 +946,38 @@ func.func @index_switch(%pred: index, %b: tensor<5xf32>, %c: tensor<5xf32>) -> t
// CHECK: return %[[r]]
return %0 : tensor<5xf32>
}
+
+// -----
+
+// See Issue https://github.com/llvm/llvm-project/issues/133964 . Checks that
+// tensor.parallel_insert_slice dest operand does not have read semantics.
+func.func @check_scfforall_inplace_bufferizer(%arg0 : tensor<?x?xf32>,
+ %arg1 : tensor<?x?xf32>,
+ %arg2 : tensor<?xf32> {bufferization.writable = true}) -> tensor<?xf32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %d0 = tensor.dim %arg2, %c0 : tensor<?xf32>
+ %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+ %0 = scf.forall (%arg3) in (%c1) shared_outs(%arg4 = %arg2) -> (tensor<?xf32>) {
+ %1 = tensor.extract_slice %arg0[0, 0][%d0, %d1][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+ %2 = tensor.extract_slice %arg1[0, 0][%d0, %d1][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
+ %3 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0, d1)>,
+ affine_map<(d0, d1) -> (d0)>],
+ iterator_types = ["parallel", "reduction"]}
+ ins(%1, %2 : tensor<?x?xf32>, tensor<?x?xf32>)
+ outs(%arg4 : tensor<?xf32>) {
+ ^bb0(%b0 : f32, %b1: f32, %b2 : f32):
+ %4 = arith.mulf %b0, %b1 : f32
+ %5 = arith.addf %4, %b2 : f32
+ linalg.yield %5 : f32
+ } -> tensor<?xf32>
+ scf.forall.in_parallel {
+ tensor.parallel_insert_slice %3 into %arg4[0] [%d0] [1] : tensor<?xf32> into tensor<?xf32>
+ }
+ }
+ return %0 : tensor<?xf32>
+}
+// CHECK-LABEL: func @check_scfforall_inplace_bufferizer
+// CHECK-NOT: memref.alloc
More information about the Mlir-commits
mailing list