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

[Ziliang Zhang]

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

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


>From f74c2cdc4dd53df0089ec097bcde6ea0dbdb0dd3 Mon Sep 17 00:00:00 2001
From: Ziliang Zhang <zzl.coding at gmail.com>
Date: Mon, 28 Sep 2026 17:33:16 +0800
Subject: [PATCH] [mlir][bufferization] Skip parallel analysis without parallel
 regions

Skip parallel privatization analysis when the analyzed IR has no parallel
regions. When needed, hoist the write-side parallel-region lookup out of
the definition loop.
---
 .../Transforms/OneShotAnalysis.h              |  6 ++++++
 .../Transforms/OneShotAnalysis.cpp            | 21 ++++++++++++++++---
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h
index af5d2bd9a2db4..26d7d64e852c2 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 1649a84c41674..d1de35d41c9e5 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 "