[Mlir-commits] [mlir] ebeb871 - [MLIR][XeGPU] Restrict XeGPUPropagateLayout analysis scope (#221793)

[llvmlistbot at llvm.org]

Author: Tuomas Kärnä
Date: 2026-09-09T14:08:23+03:00
New Revision: ebeb871f9a5fd31f088b7fb952c99cc4da17eeb1

URL: https://github.com/llvm/llvm-project/commit/ebeb871f9a5fd31f088b7fb952c99cc4da17eeb1
DIFF: https://github.com/llvm/llvm-project/commit/ebeb871f9a5fd31f088b7fb952c99cc4da17eeb1.diff

LOG: [MLIR][XeGPU] Restrict XeGPUPropagateLayout analysis scope (#221793)

This PR prevents concurrent access to traversal and mutation of the
shared IR by restricting the walk scope to determine program order
within each gpu.module.

Previously, although the pass is scoped to a single gpu.module, this
walk climbs all the way up to the enclosing top-level builtin.module and
traverses the entire module — including sibling gpu.modules. When there
is multiple gpu.modules, it causes concurrent traversal and mutation of
the shared IR and eventually leads to a segfault.

Assisted-by: Claude

Added: 
    

Modified: 
    mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 2903e19335a17..ca1c6501a4856 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -207,12 +207,17 @@ class LayoutInfoPropagation
   xegpu::LayoutKind layoutKind;
   unsigned indexBitWidth;
 
+  // The op this analysis runs on; program order is numbered within this scope
+  // only (not the enclosing module, which may be mutated concurrently by the
+  // parallel pass manager running this pass on sibling gpu.modules).
+  Operation *scopeRoot = nullptr;
+
   // Program-order index of every op, built lazily on first use via a pre-order
-  // walk of the top-level module/function (matching printed-IR order). Used to
-  // tell which consumer of a value is nearer to its producer.
+  // walk of `scopeRoot` (matching printed-IR order). Used to tell which
+  // consumer of a value is nearer to its producer.
   DenseMap<Operation *, int64_t> programOrder;
   // Returns the program-order index of `op`, populating `programOrder` from
-  // `op`'s top-level ancestor on first call.
+  // `scopeRoot` on first call.
   int64_t getProgramOrder(Operation *op);
 
   int64_t currentProgramOrder = std::numeric_limits<int64_t>::max();
@@ -323,9 +328,11 @@ class LayoutInfoPropagation
 
   LayoutInfoPropagation(DataFlowSolver &solver,
                         SymbolTableCollection &symbolTable,
-                        xegpu::LayoutKind layoutKind, unsigned indexBitWidth)
+                        xegpu::LayoutKind layoutKind, unsigned indexBitWidth,
+                        Operation *scopeRoot)
       : SparseBackwardDataFlowAnalysis(solver, symbolTable),
-        layoutKind(layoutKind), indexBitWidth(indexBitWidth) {}
+        layoutKind(layoutKind), indexBitWidth(indexBitWidth),
+        scopeRoot(scopeRoot) {}
   using SparseBackwardDataFlowAnalysis::SparseBackwardDataFlowAnalysis;
 
   LogicalResult
@@ -355,15 +362,14 @@ int64_t LayoutInfoPropagation::getProgramOrder(Operation *op) {
   auto it = programOrder.find(op);
   if (it != programOrder.end())
     return it->second;
-  // First time we see this op's tree: number every op under its top-level
-  // ancestor in pre-order (i.e. printed-IR order). Nested ops (e.g. inside an
-  // scf.for body) get an index between their parent and the parent's next
-  // sibling, so a use inside a loop is "nearer" than a use after it.
-  Operation *root = op;
-  while (root->getParentOp())
-    root = root->getParentOp();
+  // First time we number the tree: number every op under the analysis scope in
+  // pre-order (i.e. printed-IR order). Nested ops (e.g. inside an scf.for body)
+  // get an index between their parent and the parent's next sibling, so a use
+  // inside a loop is "nearer" than a use after it. Numbering is confined to
+  // `scopeRoot` (the op this pass runs on) rather than the enclosing module,
+  // which may be mutated concurrently by the parallel pass manager.
   int64_t counter = 0;
-  root->walk<WalkOrder::PreOrder>(
+  scopeRoot->walk<WalkOrder::PreOrder>(
       [&](Operation *o) { programOrder[o] = counter++; });
   return programOrder.lookup(op);
 }
@@ -1454,7 +1460,7 @@ class RunLayoutInfoPropagation {
     SymbolTableCollection symbolTable;
     loadBaselineAnalyses(solver);
     analysis = solver.load<LayoutInfoPropagation>(symbolTable, layoutKind,
-                                                  indexBitWidth);
+                                                  indexBitWidth, op);
     (void)solver.initializeAndRun(op);
   }