[Mlir-commits] [mlir] [MLIR][Affine] Fix crash in hoistAffineIfOp (PR #130401)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Uday Bondhugula (bondhugula)

<details>
<summary>Changes</summary>

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


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


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Utils/Utils.cpp (+7-6) 
- (modified) mlir/test/Dialect/Affine/loop-unswitch.mlir (+28) 


``````````diff
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 7ef016f88be37..dbad46821d039 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -263,16 +263,17 @@ static void promoteIfBlock(AffineIfOp ifOp, bool elseBlock) {
 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;
@@ -442,7 +443,7 @@ LogicalResult mlir::affine::hoistAffineIfOp(AffineIfOp ifOp, bool *folded) {
   // canonicalization is missing composition of affine.applys into it.
   assert(llvm::all_of(ifOp.getOperands(),
                       [](Value v) {
-                        return isTopLevelValue(v) || isAffineForInductionVar(v);
+                        return isTopLevelValue(v) || isAffineInductionVar(v);
                       }) &&
          "operands not composed");
 
diff --git a/mlir/test/Dialect/Affine/loop-unswitch.mlir b/mlir/test/Dialect/Affine/loop-unswitch.mlir
index 5a58941937bf5..ea94d0b72d823 100644
--- a/mlir/test/Dialect/Affine/loop-unswitch.mlir
+++ b/mlir/test/Dialect/Affine/loop-unswitch.mlir
@@ -254,3 +254,31 @@ func.func @multiple_if(%N : index) {
 // CHECK-NEXT: return
 
 func.func private @external()
+
+// Check if affine.parallel ops don't lead to any breakage.
+
+#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>
+}

``````````

</details>


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