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

[Sergio Afonso via llvm-branch-commits]

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

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.

>From 5cce160a7ee0b673e6e108502b9be8771b730e3b Mon Sep 17 00:00:00 2001
From: Sergio Afonso <safonsof at amd.com>
Date: Wed, 3 Jul 2024 11:46:00 +0100
Subject: [PATCH] [Flang][OpenMP] Prevent allocas from being inserted into loop
 wrappers

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.
---
 flang/lib/Optimizer/Builder/FIRBuilder.cpp | 23 ++++++++++++++++++++++
 1 file changed, 23 insertions(+)

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();
 }