[Mlir-commits] [mlir] Avoid buffer hoisting from parallel loops (PR #90735)
Rafael Ubal
llvmlistbot at llvm.org
Fri May 3 04:47:11 PDT 2024
https://github.com/rafaelubalmw updated https://github.com/llvm/llvm-project/pull/90735
>From 66fb97782a1bce476b4cd9e8474b01fec96f5cdf Mon Sep 17 00:00:00 2001
From: Rafael Ubal Tena <rubal at mathworks.com>
Date: Wed, 1 May 2024 10:34:32 -0400
Subject: [PATCH 1/2] Avoiding buffer hoisting from parallel loops
---
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 6 ++--
mlir/include/mlir/IR/OpBase.td | 2 ++
mlir/include/mlir/IR/OpDefinition.h | 12 +++++++
.../Transforms/BufferOptimizations.cpp | 17 +++++++---
.../Transforms/buffer-loop-hoisting.mlir | 32 +++++++++++++++++++
5 files changed, 62 insertions(+), 7 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index b3d085bfff1af9..41a0b67c42a0bf 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -307,7 +307,8 @@ def ForallOp : SCF_Op<"forall", [
RecursiveMemoryEffects,
SingleBlockImplicitTerminator<"scf::InParallelOp">,
DeclareOpInterfaceMethods<RegionBranchOpInterface>,
- DestinationStyleOpInterface
+ DestinationStyleOpInterface,
+ HasParallelRegion
]> {
let summary = "evaluate a block multiple times in parallel";
let description = [{
@@ -764,7 +765,8 @@ def ParallelOp : SCF_Op<"parallel",
"getSingleLowerBound", "getSingleUpperBound", "getSingleStep"]>,
RecursiveMemoryEffects,
DeclareOpInterfaceMethods<RegionBranchOpInterface>,
- SingleBlockImplicitTerminator<"scf::ReduceOp">]> {
+ SingleBlockImplicitTerminator<"scf::ReduceOp">,
+ HasParallelRegion]> {
let summary = "parallel for operation";
let description = [{
The "scf.parallel" operation represents a loop nest taking 4 groups of SSA
diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 7866ac24c1ccba..6cd0d38a906a45 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -79,6 +79,8 @@ def Idempotent : NativeOpTrait<"IsIdempotent">;
def Involution : NativeOpTrait<"IsInvolution">;
// Op behaves like a constant.
def ConstantLike : NativeOpTrait<"ConstantLike">;
+// Op contains a region with parallel execution semantics
+def HasParallelRegion : NativeOpTrait<"HasParallelRegion">;
// Op is isolated from above.
def IsolatedFromAbove : NativeOpTrait<"IsIsolatedFromAbove">;
// Op results are float or vectors/tensors thereof.
diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index 59f094d6690991..38b286c01b57c8 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -1242,6 +1242,18 @@ class ConstantLike : public TraitBase<ConcreteType, ConstantLike> {
}
};
+// A trait indicating that the single region contained in the operation has
+// parallel execution semantics. This may have implications in a certain pass.
+// For example, buffer hoisting is illegal in parallel loops, and local buffers
+// may be accessed by parallel threads simultaneously.
+template <typename ConcreteType>
+class HasParallelRegion : public TraitBase<ConcreteType, HasParallelRegion> {
+public:
+ static LogicalResult verifyTrait(Operation *op) {
+ return impl::verifyOneRegion(op);
+ }
+};
+
/// This class provides the API for ops that are known to be isolated from
/// above.
template <typename ConcreteType>
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 9dc2f262a51161..d7056f35cbc8d5 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -59,6 +59,12 @@ static bool isLoop(Operation *op) {
return regionInterface.hasLoop();
}
+/// Return whether the given operation is a loop with sequential execution
+/// semantics.
+static bool isSequentialLoop(Operation *op) {
+ return !op->hasTrait<OpTrait::HasParallelRegion>() && isLoop(op);
+}
+
/// Returns true if the given operation implements the AllocationOpInterface
/// and it supports the dominate block hoisting.
static bool allowAllocDominateBlockHoisting(Operation *op) {
@@ -338,12 +344,13 @@ struct BufferAllocationLoopHoistingState : BufferAllocationHoistingStateBase {
return dependencyBlock ? dependencyBlock : nullptr;
}
- /// Returns true if the given operation represents a loop and one of the
- /// aliases caused the `aliasDominatorBlock` to be "above" the block of the
- /// given loop operation. If this is the case, it indicates that the
- /// allocation is passed via a back edge.
+ /// Returns true if the given operation represents a loop with sequential
+ /// execution semantics and one of the aliases caused the
+ /// `aliasDominatorBlock` to be "above" the block of the given loop operation.
+ /// If this is the case, it indicates that the allocation is passed via a back
+ /// edge.
bool isLegalPlacement(Operation *op) {
- return isLoop(op) &&
+ return isSequentialLoop(op) &&
!dominators->dominates(aliasDominatorBlock, op->getBlock());
}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
index 5f8bc18b64da1e..ee7571b2c84669 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-hoisting.mlir
@@ -461,6 +461,38 @@ func.func @partial_hoist_multiple_loop_dependency(
// -----
+// CHECK-LABEL: func @no_hoist_parallel
+func.func @no_hoist_parallel(
+ %lb: index,
+ %ub: index,
+ %step: index) {
+ scf.parallel (%i) = (%lb) to (%ub) step (%step) {
+ %0 = memref.alloc() : memref<2xf32>
+ scf.reduce
+ }
+ return
+}
+
+// CHECK: memref.alloc
+// CHECK-NEXT: scf.reduce
+
+// -----
+
+func.func @no_hoist_forall(
+ %lb: index,
+ %ub: index,
+ %step: index) {
+ scf.forall (%i) = (%lb) to (%ub) step (%step) {
+ %1 = memref.alloc() : memref<2xf32>
+ }
+ return
+}
+
+// CHECK: scf.forall
+// CHECK-NEXT: memref.alloc
+
+// -----
+
// Test with allocas to ensure that op is also considered.
// CHECK-LABEL: func @hoist_alloca
>From 531782cd361f1975f4dd210e18bd6c8a9af83854 Mon Sep 17 00:00:00 2001
From: Rafael Ubal Tena <rubal at mathworks.com>
Date: Fri, 3 May 2024 07:46:59 -0400
Subject: [PATCH 2/2] Moved definition of trait 'HasParallelRegion' to
'Interfaces/LoopLikeInterface.h|td'
---
mlir/include/mlir/IR/OpBase.td | 2 --
mlir/include/mlir/IR/OpDefinition.h | 12 ----------
.../mlir/Interfaces/LoopLikeInterface.h | 23 +++++++++++++++++++
.../mlir/Interfaces/LoopLikeInterface.td | 11 +++++++++
4 files changed, 34 insertions(+), 14 deletions(-)
diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 6cd0d38a906a45..7866ac24c1ccba 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -79,8 +79,6 @@ def Idempotent : NativeOpTrait<"IsIdempotent">;
def Involution : NativeOpTrait<"IsInvolution">;
// Op behaves like a constant.
def ConstantLike : NativeOpTrait<"ConstantLike">;
-// Op contains a region with parallel execution semantics
-def HasParallelRegion : NativeOpTrait<"HasParallelRegion">;
// Op is isolated from above.
def IsolatedFromAbove : NativeOpTrait<"IsIsolatedFromAbove">;
// Op results are float or vectors/tensors thereof.
diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index 38b286c01b57c8..59f094d6690991 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -1242,18 +1242,6 @@ class ConstantLike : public TraitBase<ConcreteType, ConstantLike> {
}
};
-// A trait indicating that the single region contained in the operation has
-// parallel execution semantics. This may have implications in a certain pass.
-// For example, buffer hoisting is illegal in parallel loops, and local buffers
-// may be accessed by parallel threads simultaneously.
-template <typename ConcreteType>
-class HasParallelRegion : public TraitBase<ConcreteType, HasParallelRegion> {
-public:
- static LogicalResult verifyTrait(Operation *op) {
- return impl::verifyOneRegion(op);
- }
-};
-
/// This class provides the API for ops that are known to be isolated from
/// above.
template <typename ConcreteType>
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.h b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
index 7c7d378d0590ab..42609e824c86ac 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.h
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
@@ -29,8 +29,31 @@ namespace detail {
/// Verify invariants of the LoopLikeOpInterface.
LogicalResult verifyLoopLikeOpInterface(Operation *op);
} // namespace detail
+
+//===----------------------------------------------------------------------===//
+// Traits
+//===----------------------------------------------------------------------===//
+
+namespace OpTrait {
+// A trait indicating that the single region contained in the operation has
+// parallel execution semantics. This may have implications in a certain pass.
+// For example, buffer hoisting is illegal in parallel loops, and local buffers
+// may be accessed by parallel threads simultaneously.
+template <typename ConcreteType>
+class HasParallelRegion : public TraitBase<ConcreteType, HasParallelRegion> {
+public:
+ static LogicalResult verifyTrait(Operation *op) {
+ return impl::verifyOneRegion(op);
+ }
+};
+
+} // namespace OpTrait
} // namespace mlir
+//===----------------------------------------------------------------------===//
+// Interfaces
+//===----------------------------------------------------------------------===//
+
/// Include the generated interface declarations.
#include "mlir/Interfaces/LoopLikeInterface.h.inc"
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index e2ac85a3f7725d..f0dc6e60eba584 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -15,6 +15,10 @@
include "mlir/IR/OpBase.td"
+//===----------------------------------------------------------------------===//
+// Interfaces
+//===----------------------------------------------------------------------===//
+
def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
let description = [{
Contains helper functions to query properties and perform transformations
@@ -371,4 +375,11 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
}];
}
+//===----------------------------------------------------------------------===//
+// Traits
+//===----------------------------------------------------------------------===//
+
+// Op contains a region with parallel execution semantics
+def HasParallelRegion : NativeOpTrait<"HasParallelRegion">;
+
#endif // MLIR_INTERFACES_LOOPLIKEINTERFACE
More information about the Mlir-commits
mailing list