[llvm-branch-commits] [flang] [Flang][OpenMP] Prevent allocas from being inserted into loop wrappers (PR #97563)

[via llvm-branch-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-flang-fir-hlfir

Author: Sergio Afonso (skatrak)

<details>
<summary>Changes</summary>

This patch updates the `FirOpBuilder::getAllocaBlock()` method to avoid returning blocks which are part of a region owned by a loop wrapper operation.

This avoids introducing `fir.alloca` operations inside of loop wrappers, which would cause verifier errors due to strict restrictions on their allowed contents. These allocations will be created in the entry block of the first non-loop wrapper parent region.

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


1 Files Affected:

- (modified) flang/lib/Optimizer/Builder/FIRBuilder.cpp (+23) 


``````````diff
diff --git a/flang/lib/Optimizer/Builder/FIRBuilder.cpp b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
index 2ea302d188018..fbe915b9fd9e3 100644
--- a/flang/lib/Optimizer/Builder/FIRBuilder.cpp
+++ b/flang/lib/Optimizer/Builder/FIRBuilder.cpp
@@ -259,6 +259,29 @@ mlir::Block *fir::FirOpBuilder::getAllocaBlock() {
     return recipeIface.getAllocaBlock(getRegion());
   }
 
+  // All allocations associated with an OpenMP loop wrapper must happen outside
+  // of all wrappers.
+  mlir::Operation *currentOp = getRegion().getParentOp();
+  auto wrapperIface =
+      llvm::isa<mlir::omp::LoopNestOp>(currentOp)
+          ? llvm::cast<mlir::omp::LoopWrapperInterface>(
+                currentOp->getParentOp())
+          : llvm::dyn_cast<mlir::omp::LoopWrapperInterface>(currentOp);
+  if (wrapperIface) {
+    // Cannot use LoopWrapperInterface methods here because the whole nest may
+    // not have been created at this point. Manually traverse parents instead.
+    mlir::omp::LoopWrapperInterface lastWrapperOp = wrapperIface;
+    while (true) {
+      if (auto nextWrapper =
+              llvm::dyn_cast_if_present<mlir::omp::LoopWrapperInterface>(
+                  lastWrapperOp->getParentOp()))
+        lastWrapperOp = nextWrapper;
+      else
+        break;
+    }
+    return &lastWrapperOp->getParentRegion()->front();
+  }
+
   return getEntryBlock();
 }
 

``````````

</details>


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