[clang] [flang] [llvm] [OpenMP][LLVM] Update alloca IP after `PrivCB` in `OMPIRBUIlder` (PR #93920)

[Kareem Ergawy via llvm-commits]

https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/93920

>From 926cf8d19c625880c303aff0527e2e6e8a1629bd Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Thu, 30 May 2024 23:16:39 -0500
Subject: [PATCH 1/3] [OpenMP][][LLVM] Update alloca IP after `PrivCB` in
 `OMPIRBUIlder`

Fixes a crash uncovered by [pr77666.f90](https://github.com/llvm/llvm-test-suite/blob/main/Fortran/gfortran/regression/gomp/pr77666.f90) in the test suite.

In particular, whenever `PrivCB` (the callback responsible for
generating privatizaiton logic for an OMP variable) generates a
multi-block privatization region, the insertion point diverges: the BB
component of the IP can become a different BB from the parent block of
the instruction iterator component of the IP. This PR updates the IP to
make sure that the BB is the parent block of the instruction iterator.
---
 ...rivatization-lower-allocatable-to-llvm.f90 | 23 +++++++++++++++++++
 llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp     |  3 +++
 2 files changed, 26 insertions(+)
 create mode 100644 flang/test/Lower/OpenMP/delayed-privatization-lower-allocatable-to-llvm.f90

diff --git a/flang/test/Lower/OpenMP/delayed-privatization-lower-allocatable-to-llvm.f90 b/flang/test/Lower/OpenMP/delayed-privatization-lower-allocatable-to-llvm.f90
new file mode 100644
index 0000000000000..ac9a6d8746cf2
--- /dev/null
+++ b/flang/test/Lower/OpenMP/delayed-privatization-lower-allocatable-to-llvm.f90
@@ -0,0 +1,23 @@
+! Tests the OMPIRBuilder can handle multiple privatization regions that contain
+! multiple BBs (for example, for allocatables).
+
+! RUN: %flang -S -emit-llvm -fopenmp -mmlir --openmp-enable-delayed-privatization \
+! RUN:   -o - %s 2>&1 | FileCheck %s
+
+subroutine foo(x)
+  integer, allocatable :: x, y
+!$omp parallel private(x, y)
+  x = y
+!$omp end parallel
+end
+
+! CHECK-LABEL: define void @foo_
+! CHECK:         ret void
+! CHECK-NEXT:  }
+
+! CHECK-LABEL: define internal void @foo_..omp_par
+! CHECK-DAG:     call ptr @malloc
+! CHECK-DAG:     call ptr @malloc
+! CHECK-DAG:     call void @free
+! CHECK-DAG:     call void @free
+! CHECK:       }
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index cb4de9c8876dc..eab41eb8a35b2 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -1583,6 +1583,9 @@ IRBuilder<>::InsertPoint OpenMPIRBuilder::createParallel(
     } else {
       Builder.restoreIP(
           PrivCB(InnerAllocaIP, Builder.saveIP(), V, *Inner, ReplacementValue));
+      InnerAllocaIP = {InnerAllocaIP.getPoint()->getParent(),
+                       InnerAllocaIP.getPoint()};
+
       assert(ReplacementValue &&
              "Expected copy/create callback to set replacement value!");
       if (ReplacementValue == &V)

>From 659eec2a468902cf1654394f3eccdab16e92a027 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Mon, 3 Jun 2024 09:05:35 -0500
Subject: [PATCH 2/3] update ip

---
 llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index eab41eb8a35b2..2c4b45255d059 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -1583,8 +1583,9 @@ IRBuilder<>::InsertPoint OpenMPIRBuilder::createParallel(
     } else {
       Builder.restoreIP(
           PrivCB(InnerAllocaIP, Builder.saveIP(), V, *Inner, ReplacementValue));
-      InnerAllocaIP = {InnerAllocaIP.getPoint()->getParent(),
-                       InnerAllocaIP.getPoint()};
+      InnerAllocaIP = {
+          InnerAllocaIP.getBlock(),
+          InnerAllocaIP.getBlock()->getTerminator()->getIterator()};
 
       assert(ReplacementValue &&
              "Expected copy/create callback to set replacement value!");

>From d5403868ef43729aceaf763d8fa7e8e784938948 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Mon, 3 Jun 2024 22:13:05 -0500
Subject: [PATCH 3/3] fix clang expectations

---
 clang/test/OpenMP/parallel_codegen.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/clang/test/OpenMP/parallel_codegen.cpp b/clang/test/OpenMP/parallel_codegen.cpp
index d545b4a9d9fa8..9082f1c3232af 100644
--- a/clang/test/OpenMP/parallel_codegen.cpp
+++ b/clang/test/OpenMP/parallel_codegen.cpp
@@ -822,8 +822,8 @@ int main (int argc, char **argv) {
 // CHECK3-NEXT:    [[TMP1:%.*]] = load i32, ptr [[TID_ADDR]], align 4
 // CHECK3-NEXT:    store i32 [[TMP1]], ptr [[TID_ADDR_LOCAL]], align 4
 // CHECK3-NEXT:    [[TID:%.*]] = load i32, ptr [[TID_ADDR_LOCAL]], align 4
-// CHECK3-NEXT:    [[TMP2:%.*]] = load i64, ptr [[LOADGEP__RELOADED]], align 8
 // CHECK3-NEXT:    [[VAR:%.*]] = alloca ptr, align 8
+// CHECK3-NEXT:    [[TMP2:%.*]] = load i64, ptr [[LOADGEP__RELOADED]], align 8
 // CHECK3-NEXT:    br label [[OMP_PAR_REGION:%.*]]
 // CHECK3:       omp.par.region:
 // CHECK3-NEXT:    [[TMP3:%.*]] = load ptr, ptr [[LOADGEP_ARGC_ADDR]], align 8
@@ -966,8 +966,8 @@ int main (int argc, char **argv) {
 // CHECK4-NEXT:    [[TMP1:%.*]] = load i32, ptr [[TID_ADDR]], align 4
 // CHECK4-NEXT:    store i32 [[TMP1]], ptr [[TID_ADDR_LOCAL]], align 4
 // CHECK4-NEXT:    [[TID:%.*]] = load i32, ptr [[TID_ADDR_LOCAL]], align 4
-// CHECK4-NEXT:    [[TMP2:%.*]] = load i64, ptr [[LOADGEP__RELOADED]], align 8
 // CHECK4-NEXT:    [[VAR:%.*]] = alloca ptr, align 8
+// CHECK4-NEXT:    [[TMP2:%.*]] = load i64, ptr [[LOADGEP__RELOADED]], align 8
 // CHECK4-NEXT:    br label [[OMP_PAR_REGION:%.*]]
 // CHECK4:       omp.par.region:
 // CHECK4-NEXT:    [[TMP3:%.*]] = load ptr, ptr [[LOADGEP_ARGC_ADDR]], align 8, !dbg [[DBG58:![0-9]+]]