commit b01c71bbd4f7f7565313c38ac7a3b7377eada798
author Uday Bondhugula <uday@polymagelabs.com> Sat Mar 08 20:23:56 2025 +0530
committer GitHub <noreply@github.com> Sat Mar 08 20:23:56 2025 +0530
tree 813e1af453a1f5412eec66f4f1e0993882b37d64
parent a779af3f882e1500ee6b9cb973e60c10ea0819b4

[MLIR][Affine] Fix crash in loop unswitching/hoistAffineIfOp (#130401)

Fix obvious crash as a result of missing affine.parallel handling. Also,
fix bug exposed in a helper method used by hoistAffineIfOp.

Fixes: https://github.com/llvm/llvm-project/issues/62323

mlir/lib/Dialect/Affine/Utils/Utils.cpp

diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 009b1b2..2723cff 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -263,16 +263,17 @@
 static Operation *getOutermostInvariantForOp(AffineIfOp ifOp) {
   // Walk up the parents past all for op that this conditional is invariant on.
   auto ifOperands = ifOp.getOperands();
-  auto *res = ifOp.getOperation();
-  while (!isa<func::FuncOp>(res->getParentOp())) {
+  Operation *res = ifOp;
+  while (!res->getParentOp()->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
     auto *parentOp = res->getParentOp();
     if (auto forOp = dyn_cast<AffineForOp>(parentOp)) {
       if (llvm::is_contained(ifOperands, forOp.getInductionVar()))
         break;
     } else if (auto parallelOp = dyn_cast<AffineParallelOp>(parentOp)) {
-      for (auto iv : parallelOp.getIVs())
-        if (llvm::is_contained(ifOperands, iv))
-          break;
+      if (llvm::any_of(parallelOp.getIVs(), [&](Value iv) {
+            return llvm::is_contained(ifOperands, iv);
+          }))
+        break;
     } else if (!isa<AffineIfOp>(parentOp)) {
       // Won't walk up past anything other than affine.for/if ops.
       break;
@@ -438,11 +439,10 @@
   if (folded)
     *folded = false;
 
-  // The folding above should have ensured this, but the affine.if's
-  // canonicalization is missing composition of affine.applys into it.
+  // The folding above should have ensured this.
   assert(llvm::all_of(ifOp.getOperands(),
                       [](Value v) {
-                        return isTopLevelValue(v) || isAffineForInductionVar(v);
+                        return isTopLevelValue(v) || isAffineInductionVar(v);
                       }) &&
          "operands not composed");
 

mlir/test/Dialect/Affine/loop-unswitch.mlir

diff --git a/mlir/test/Dialect/Affine/loop-unswitch.mlir b/mlir/test/Dialect/Affine/loop-unswitch.mlir
index 5a58941..c94f551 100644
--- a/mlir/test/Dialect/Affine/loop-unswitch.mlir
+++ b/mlir/test/Dialect/Affine/loop-unswitch.mlir
@@ -254,3 +254,31 @@
 // CHECK-NEXT: return
 
 func.func private @external()
+
+// Check to ensure affine.parallel ops are handled as well.
+
+#set = affine_set<(d0) : (-d0 + 3 >= 0)>
+// CHECK-LABEL: affine_parallel
+func.func @affine_parallel(%arg0: memref<35xf32>) -> memref<35xf32> {
+  %0 = llvm.mlir.constant(1.000000e+00 : f32) : f32
+  %alloc = memref.alloc() {alignment = 64 : i64} : memref<35xf32>
+  // CHECK: affine.parallel
+  affine.parallel (%arg1) = (0) to (35) step (32) {
+    // This can't be hoisted further.
+    // CHECK-NEXT: affine.if
+    affine.if #set(%arg1) {
+      affine.parallel (%arg2) = (%arg1) to (%arg1 + 32) {
+        %1 = affine.load %arg0[%arg2] : memref<35xf32>
+        %2 = llvm.fdiv %0, %1 : f32
+        affine.store %2, %alloc[%arg2] : memref<35xf32>
+      }
+    } else {
+      affine.parallel (%arg2) = (%arg1) to (min(%arg1 + 32, 35)) {
+        %1 = affine.load %arg0[%arg2] : memref<35xf32>
+        %2 = llvm.fdiv %0, %1 : f32
+        affine.store %2, %alloc[%arg2] : memref<35xf32>
+      }
+    }
+  }
+  return %alloc : memref<35xf32>
+}