[Mlir-commits] [mlir] [mlir][bufferization] Skip parallel analysis without parallel regions (PR #227202)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Ziliang Zhang (ziliangzl)

<details>
<summary>Changes</summary>

One-Shot Bufferize currently runs parallel-region conflict analysis even when the analyzed IR has no parallel regions. This can repeatedly walk long reverse use-def chains without being able to find a parallel-region conflict.

Record whether the analyzed IR contains any parallel region while performing the existing initialization walk, and skip parallel analysis when it does not. When parallel analysis is needed, also move the write-side parallel-region lookup out of the definition loop.

A synthetic no-parallel-region stress test with 16 long tensor.cast chains reduced median analysis time by 60.0% to 64.5% across chain lengths from 64 to 256. At the largest size, runtime decreased from 11.62 seconds to 4.13 seconds. This deliberately amplifies reverse use-def traversal and is not intended to predict typical workload speedups.

Testing:
- ninja -C build -j10 check-mlir
- clang-format --dry-run --Werror on the modified files


---
Full diff: https://github.com/llvm/llvm-project/pull/227202.diff


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h (+6) 
- (modified) mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp (+18-3) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
index af5d2bd9a2db44..26d7d64e852c2c 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
@@ -91,6 +91,9 @@ class OneShotAnalysisState : public AnalysisState {
   /// the options when set; otherwise computed by walking the IR.
   bool mayHaveUnstructuredControlFlow() const { return mayHaveUnstructuredCF; }
 
+  /// True if any allowed op in the analyzed IR has a parallel region.
+  bool hasParallelRegion() const { return containsParallelRegion; }
+
   /// Analyze the given op and its nested ops.
   LogicalResult analyzeOp(Operation *op, const DominanceInfo &domInfo);
 
@@ -267,6 +270,9 @@ class OneShotAnalysisState : public AnalysisState {
   /// True if any region has more than one block.
   bool mayHaveUnstructuredCF = false;
 
+  /// True if any allowed op has a parallel region.
+  bool containsParallelRegion = false;
+
   /// Cached CFG reachability. Defined out-of-line to keep BitVector out of
   /// this header.
   class CFGReachabilityCache;
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
index 1649a84c41674a..d1de35d41c9e5f 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
@@ -153,6 +153,19 @@ OneShotAnalysisState::OneShotAnalysisState(
   op->walk([&](BufferizableOpInterface bufferizableOp) {
     if (!options.isOpAllowed(bufferizableOp))
       return WalkResult::skip();
+
+    // Reuse this existing walk to determine whether parallel-region analysis
+    // can ever find a conflict. Avoid querying the interface once a parallel
+    // region was found.
+    if (options.checkParallelRegions && !containsParallelRegion) {
+      for (Region &region : bufferizableOp->getRegions()) {
+        if (bufferizableOp.isParallelRegion(region.getRegionNumber())) {
+          containsParallelRegion = true;
+          break;
+        }
+      }
+    }
+
     for (OpOperand &opOperand : bufferizableOp->getOpOperands())
       if (isa<TensorLikeType>(opOperand.get().getType()))
         if (bufferizableOp.mustBufferizeInPlace(opOperand, *this))
@@ -841,7 +854,8 @@ hasReadAfterWriteInterference(const DenseSet<OpOperand *> &usesRead,
   // Before going through the main RaW analysis, find cases where a buffer must
   // be privatized due to parallelism. If the result of a write is never read,
   // privatization is not necessary (and large parts of the IR are likely dead).
-  if (options.checkParallelRegions && !usesRead.empty()) {
+  if (options.checkParallelRegions && state.hasParallelRegion() &&
+      !usesRead.empty()) {
     for (OpOperand *uConflictingWrite : usesWrite) {
       // Find the allocation point or last write (definition) of the buffer.
       // Note: In contrast to `findDefinitions`, this also returns results of
@@ -858,10 +872,11 @@ hasReadAfterWriteInterference(const DenseSet<OpOperand *> &usesRead,
 
       // The writing op must bufferize out-of-place if the definition is in a
       // different parallel region than this write.
+      Region *writeParallelRegion = getParallelRegion(
+          uConflictingWrite->getOwner()->getParentRegion(), options);
       for (Value def : definitionsOrLeaves) {
         if (getParallelRegion(def.getParentRegion(), options) !=
-            getParallelRegion(uConflictingWrite->getOwner()->getParentRegion(),
-                              options)) {
+            writeParallelRegion) {
           LDBG() << "\n- bufferizes out-of-place due to parallel region:\n"
                  << "  unConflictingWrite = operand "
                  << uConflictingWrite->getOperandNumber() << " of "

``````````

</details>


https://github.com/llvm/llvm-project/pull/227202