[Mlir-commits] [mlir] [mlir] Fix crash in ForwardDominanceIterator when encountering graph regions (PR #185043)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

ForwardDominanceIterator<NoGraphRegions=true> was asserting when it encountered a region without SSA dominance (a "graph region"), such as scf.forall.in_parallel's body. This crash was triggered by -test-ir-visitors when walking functions that contain graph-region ops.

Change the behavior of ForwardDominanceIterator<true> and ReverseDominanceIterator<true> to silently skip graph regions instead of asserting, and update the documentation accordingly. This matches the intended semantics of the NoGraphRegions flag: the traversal simply does not enumerate blocks/ops inside such regions.

Fixes #<!-- -->116370

Assisted-by: Claude Code

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


2 Files Affected:

- (modified) mlir/include/mlir/IR/Iterators.h (+16-14) 
- (modified) mlir/test/IR/visitors.mlir (+22) 


``````````diff
diff --git a/mlir/include/mlir/IR/Iterators.h b/mlir/include/mlir/IR/Iterators.h
index 2c6137c72cf5d..b44ba7ddfdd87 100644
--- a/mlir/include/mlir/IR/Iterators.h
+++ b/mlir/include/mlir/IR/Iterators.h
@@ -40,10 +40,10 @@ struct ReverseIterator {
 /// enumerated according to their successor relationship. Unreachable blocks are
 /// not enumerated. Blocks may not be erased during the traversal.
 ///
-/// Note: If `NoGraphRegions` is set to "true", this iterator asserts that each
-/// visited region has SSA dominance. In either case, the ops in such regions
-/// are visited in forward order, but for regions without SSA dominance this
-/// does not guarantee that defining ops are visited before their users.
+/// Note: If `NoGraphRegions` is set to "true", graph regions (regions without
+/// SSA dominance) are skipped during the traversal. If set to "false" (the
+/// default), graph regions are visited but without dominance guarantees (i.e.,
+/// defining ops are not guaranteed to be visited before their users).
 template <bool NoGraphRegions = false>
 struct ForwardDominanceIterator {
   static Block &makeIterable(Block &range) {
@@ -51,14 +51,15 @@ struct ForwardDominanceIterator {
   }
 
   static auto makeIterable(Region &region) {
-    if (NoGraphRegions) {
-      // Only regions with SSA dominance are allowed.
-      assert(mayHaveSSADominance(region) && "graph regions are not allowed");
+    Block *null = nullptr;
+    if (NoGraphRegions && !mayHaveSSADominance(region)) {
+      // Skip graph regions.
+      return llvm::make_pointee_range(
+          llvm::make_range(llvm::df_end(null), llvm::df_end(null)));
     }
 
     // Create DFS iterator. Blocks are enumerated according to their successor
     // relationship.
-    Block *null = nullptr;
     auto it = region.empty()
                   ? llvm::make_range(llvm::df_end(null), llvm::df_end(null))
                   : llvm::depth_first(&region.front());
@@ -79,8 +80,8 @@ struct ForwardDominanceIterator {
 /// Cycles in the block graph are broken in an unspecified way. Unreachable
 /// blocks are not enumerated. Blocks may not be erased during the traversal.
 ///
-/// Note: If `NoGraphRegions` is set to "true", this iterator asserts that each
-/// visited region has SSA dominance.
+/// Note: If `NoGraphRegions` is set to "true", graph regions (regions without
+/// SSA dominance) are skipped during the traversal.
 template <bool NoGraphRegions = false>
 struct ReverseDominanceIterator {
   // llvm::reverse uses RangeT::rbegin and RangeT::rend.
@@ -93,14 +94,15 @@ struct ReverseDominanceIterator {
   }
 
   static auto makeIterable(Region &region) {
-    if (NoGraphRegions) {
-      // Only regions with SSA dominance are allowed.
-      assert(mayHaveSSADominance(region) && "graph regions are not allowed");
+    Block *null = nullptr;
+    if (NoGraphRegions && !mayHaveSSADominance(region)) {
+      // Skip graph regions.
+      return llvm::make_pointee_range(
+          llvm::make_range(llvm::po_end(null), llvm::po_end(null)));
     }
 
     // Create post-order iterator. Blocks are enumerated according to their
     // successor relationship.
-    Block *null = nullptr;
     auto it = region.empty()
                   ? llvm::make_range(llvm::po_end(null), llvm::po_end(null))
                   : llvm::post_order(&region.front());
diff --git a/mlir/test/IR/visitors.mlir b/mlir/test/IR/visitors.mlir
index 24929f9fb6cf1..278d9fa07694f 100644
--- a/mlir/test/IR/visitors.mlir
+++ b/mlir/test/IR/visitors.mlir
@@ -422,3 +422,25 @@ func.func @test_no_skip_block_erasure_block_args(%arg0: i32, %arg1: i32) -> i32
 module {}
 // CHECK-LABEL: Block post-order erasures (no skip)
 // CHECK-NEXT:  Erasing block ^bb0 from region 0 from operation 'builtin.module'
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/116370:
+// ForwardDominanceIterator<NoGraphRegions=true> should skip graph regions
+// (such as scf.forall.in_parallel's body) instead of asserting.
+// CHECK-LABEL: Op forward dominance post-order visits
+// CHECK: Visiting op 'scf.forall'
+// CHECK-NOT: Visiting op 'tensor.parallel_insert_slice'
+// CHECK: Op reverse dominance post-order visits
+func.func @graph_region_skip(%fill: tensor<2xf32>, %output: tensor<2xf32>) {
+  %c0 = arith.constant 0.0 : f32
+  %0 = linalg.fill ins(%c0 : f32) outs(%fill : tensor<2xf32>) -> tensor<2xf32>
+  %1 = scf.forall (%i) in (2) shared_outs(%arg1 = %output) -> (tensor<2xf32>) {
+    %2 = tensor.extract_slice %0[%i][1][1] : tensor<2xf32> to tensor<1xf32>
+    %3 = tensor.extract_slice %arg1[%i][1][1] : tensor<2xf32> to tensor<1xf32>
+    scf.forall.in_parallel {
+      tensor.parallel_insert_slice %3 into %arg1[%i][1][1] : tensor<1xf32> into tensor<2xf32>
+    }
+  }
+  return
+}

``````````

</details>


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