[Mlir-commits] [flang] [llvm] [mlir] Fix for changed code at the end of AllocaIP. (PR #92430)

[Mats Petersson]

https://github.com/Leporacanthicus created https://github.com/llvm/llvm-project/pull/92430

Some of the OpenMP code can change the instruction pointed at by the insertion point. This leads to an assert in the compiler about BB->getParent() and IP->getParent() not matching or something like that.

The fix is to rebuild the insertionpoint from the block, rather than use builder.restoreIP.

Also, move some of the alloca generation, rather than skipping back and forth between insert points (and ensure all the allocas are done before their users are created).

A simple test, mainly to ensure the minimal reproducer doesn't fail to compile in the future is also added.

>From 807c33831fe340893a1baa9d5ae79d2c33cbba9e Mon Sep 17 00:00:00 2001
From: Mats Petersson <mats.petersson at arm.com>
Date: Fri, 19 Apr 2024 18:00:58 +0100
Subject: [PATCH] Fix for changed code at the end of AllocaIP.

Some of the OpenMP code can change the instruction pointed at by
the insertion point. This leads to an assert in the compiler about
BB->getParent() and IP->getParent() not matching or something like
that.

The fix is to rebuild the insertionpoint from the block, rather
than use builder.restoreIP.

Also, move some of the alloca generation, rather than skipping back
and forth between insert points (and ensure all the allocas are done
before their users are created).

A simple test, mainly to ensure the minimal reproducer doesn't fail
to compile in the future is also added.
---
 .../OpenMP/parallel-reduction-allocate.f90    | 23 +++++++++++++++++++
 llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp     | 11 ++++++---
 .../OpenMP/OpenMPToLLVMIRTranslation.cpp      | 20 ++++++++++------
 3 files changed, 44 insertions(+), 10 deletions(-)
 create mode 100644 flang/test/Lower/OpenMP/parallel-reduction-allocate.f90

diff --git a/flang/test/Lower/OpenMP/parallel-reduction-allocate.f90 b/flang/test/Lower/OpenMP/parallel-reduction-allocate.f90
new file mode 100644
index 0000000000000..fddce25ae22cc
--- /dev/null
+++ b/flang/test/Lower/OpenMP/parallel-reduction-allocate.f90
@@ -0,0 +1,23 @@
+!! The main point of this test is to check that the code compiles at all, so the
+!! checking is not very detailed. Not hitting an assert, crashing or otherwise failing
+!! to compile is the key point. Also, emitting llvm is required for this to happen.
+! RUN: %flang_fc1 -emit-llvm -fopenmp -o - %s 2>&1 | FileCheck %s
+subroutine proc
+  implicit none
+  real(8),allocatable :: F(:)
+  real(8),allocatable :: A(:)
+
+!$omp parallel private(A) reduction(+:F)
+  allocate(A(10))
+!$omp end parallel
+end subroutine proc
+
+!CHECK-LABEL: define void @proc_()
+!CHECK: call void
+!CHECK-SAME: @__kmpc_fork_call(ptr {{.*}}, i32 1, ptr @[[OMP_PAR:.*]], {{.*}})
+
+!CHECK: define internal void @[[OMP_PAR]]
+!CHECK: omp.par.region8:
+!CHECK-NEXT: call ptr @malloc
+!CHECK-SAME: i64 10
+
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 9f0c07ef0da91..c0f4b5b066097 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -1391,7 +1391,8 @@ IRBuilder<>::InsertPoint OpenMPIRBuilder::createParallel(
 
   // Change the location to the outer alloca insertion point to create and
   // initialize the allocas we pass into the parallel region.
-  Builder.restoreIP(OuterAllocaIP);
+  InsertPointTy NewOuter(OuterAllocaBlock, OuterAllocaBlock->begin());
+  Builder.restoreIP(NewOuter);
   AllocaInst *TIDAddrAlloca = Builder.CreateAlloca(Int32, nullptr, "tid.addr");
   AllocaInst *ZeroAddrAlloca =
       Builder.CreateAlloca(Int32, nullptr, "zero.addr");
@@ -2135,7 +2136,8 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createReductions(
   // values.
   unsigned NumReductions = ReductionInfos.size();
   Type *RedArrayTy = ArrayType::get(Builder.getPtrTy(), NumReductions);
-  Builder.restoreIP(AllocaIP);
+  Builder.SetInsertPoint(AllocaIP.getBlock()->getTerminator());
+  //  Builder.restoreIP(AllocaIP);
   Value *RedArray = Builder.CreateAlloca(RedArrayTy, nullptr, "red.array");
 
   Builder.SetInsertPoint(InsertBlock, InsertBlock->end());
@@ -2536,7 +2538,10 @@ OpenMPIRBuilder::applyStaticWorkshareLoop(DebugLoc DL, CanonicalLoopInfo *CLI,
       getOrCreateRuntimeFunction(M, omp::OMPRTL___kmpc_for_static_fini);
 
   // Allocate space for computed loop bounds as expected by the "init" function.
-  Builder.restoreIP(AllocaIP);
+
+  //  Builder.restoreIP(AllocaIP);
+  Builder.SetInsertPoint(AllocaIP.getBlock()->getTerminator());
+
   Type *I32Type = Type::getInt32Ty(M.getContext());
   Value *PLastIter = Builder.CreateAlloca(I32Type, nullptr, "p.lastiter");
   Value *PLowerBound = Builder.CreateAlloca(IVTy, nullptr, "p.lowerbound");
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index bfd7d65912bdb..b5f46b641f596 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -1194,6 +1194,17 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
     MutableArrayRef<BlockArgument> reductionArgs =
         opInst.getRegion().getArguments().take_back(
             opInst.getNumReductionVars());
+
+    SmallVector<llvm::Value *> byRefVars;
+    if (isByRef) {
+      for (unsigned i = 0; i < opInst.getNumReductionVars(); ++i) {
+        // Allocate reduction variable (which is a pointer to the real reduciton
+        // variable allocated in the inlined region)
+        byRefVars.push_back(builder.CreateAlloca(
+            moduleTranslation.convertType(reductionDecls[i].getType())));
+      }
+    }
+
     for (unsigned i = 0; i < opInst.getNumReductionVars(); ++i) {
       SmallVector<llvm::Value *> phis;
 
@@ -1206,18 +1217,13 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
       assert(phis.size() == 1 &&
              "expected one value to be yielded from the "
              "reduction neutral element declaration region");
-      builder.restoreIP(allocaIP);
 
       if (isByRef) {
-        // Allocate reduction variable (which is a pointer to the real reduciton
-        // variable allocated in the inlined region)
-        llvm::Value *var = builder.CreateAlloca(
-            moduleTranslation.convertType(reductionDecls[i].getType()));
         // Store the result of the inlined region to the allocated reduction var
         // ptr
-        builder.CreateStore(phis[0], var);
+        builder.CreateStore(phis[0], byRefVars[i]);
 
-        privateReductionVariables.push_back(var);
+        privateReductionVariables.push_back(byRefVars[i]);
         moduleTranslation.mapValue(reductionArgs[i], phis[0]);
         reductionVariableMap.try_emplace(opInst.getReductionVars()[i], phis[0]);
       } else {