[Mlir-commits] [mlir] [mlir] [bufferization] Default implementation of BufferizableOpInterface::isParallelRegion() based on new trait OpTrait::HasParallelRegion (PR #91184)

[Rafael Ubal]

https://github.com/rafaelubalmw updated https://github.com/llvm/llvm-project/pull/91184

>From efe98cf906dc59548e6b321805a9eb4e639d50f2 Mon Sep 17 00:00:00 2001
From: Rafael Ubal Tena <rubal at mathworks.com>
Date: Mon, 6 May 2024 06:24:54 -0400
Subject: [PATCH] Changed default implementation of
 BufferizableOpInterface::isParallelRegion()

---
 .../Dialect/Bufferization/IR/BufferizableOpInterface.h   | 6 ++++++
 .../Dialect/Bufferization/IR/BufferizableOpInterface.td  | 9 +++++++--
 .../Dialect/Bufferization/IR/BufferizableOpInterface.cpp | 7 +++++++
 3 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index 2d8add82383be..8dbbe1141aea1 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -701,6 +701,12 @@ bool defaultResultBufferizesToMemoryWrite(OpResult opResult,
 bool defaultIsRepetitiveRegion(BufferizableOpInterface bufferizableOp,
                                unsigned index);
 
+/// This is the default implementation of
+/// BufferizableOpInterface::isParallelRegion. Should not be called from other
+/// places.
+bool defaultIsParallelRegion(BufferizableOpInterface bufferizableOp,
+                             unsigned index);
+
 /// This is the default implementation of getAliasingOpOperands in case the
 /// defining op does not implement the BufferizableOpInterface.
 AliasingOpOperandList unknownGetAliasingOpOperands(Value value);
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
index 007c05adc30b5..3ab16c70b14b6 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
@@ -565,14 +565,19 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
           The RaW conflict detection of One-Shot Analysis is more strict inside
           parallel regions: Buffer may have to be privatized.
 
-          By default, regions are assumed to be sequential.
+          By default, an op region is considered parallel if the containing op
+          has trait `HasParallelRegion`. While this default implementation is
+          generally sufficient, a specific op may relax this condition by
+          marking a region as non-parallel when it is detected to execute
+          exactly once, and in spite of its parallel semantics.
         }],
         /*retType=*/"bool",
         /*methodName=*/"isParallelRegion",
         /*args=*/(ins "unsigned":$index),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
-          return false;
+          return ::mlir::bufferization::detail::defaultIsParallelRegion(
+              ::llvm::cast<BufferizableOpInterface>($_op.getOperation()), index);
         }]
       >,
       InterfaceMethod<
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index d51d63f243ea0..cb36ad8f0f964 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -18,6 +18,7 @@
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Debug.h"
 
@@ -956,6 +957,12 @@ bool bufferization::detail::defaultIsRepetitiveRegion(
   return regionInterface.isRepetitiveRegion(index);
 }
 
+bool bufferization::detail::defaultIsParallelRegion(
+    BufferizableOpInterface bufferizableOp, unsigned index) {
+  assert(index < bufferizableOp->getNumRegions() && "invalid region index");
+  return bufferizableOp->hasTrait<OpTrait::HasParallelRegion>();
+}
+
 AliasingOpOperandList
 bufferization::detail::unknownGetAliasingOpOperands(Value value) {
   // TODO: Take into account successor blocks.