Make it work for both MLIR and Flang tests

Author: Leporacanthicus

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -2157,10 +2157,9 @@ OpenMPIRBuilder::createReductions(const LocationDescription &Loc,
   unsigned NumReductions = ReductionInfos.size();
   Type *RedArrayTy = ArrayType::get(Builder.getPtrTy(), NumReductions);
   Builder.SetInsertPoint(AllocaIP.getBlock()->getTerminator());
-  //  Builder.restoreIP(AllocaIP);
   Value *RedArray = Builder.CreateAlloca(RedArrayTy, nullptr, "red.array");
 
-  Builder.SetInsertPoint(InsertBlock, InsertBlock->end());
+  Builder.SetInsertPoint(InsertBlock, InsertBlock->begin());
 
   for (auto En : enumerate(ReductionInfos)) {
     unsigned Index = En.index();
@@ -2558,8 +2557,6 @@ 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.SetInsertPoint(AllocaIP.getBlock()->getTerminator());
 
   Type *I32Type = Type::getInt32Ty(M.getContext());

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -1155,6 +1155,7 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
         opInst.getRegion().getArguments().take_back(
             opInst.getNumReductionVars());
 
+    llvm::BasicBlock *initBlock = nullptr;
     SmallVector<llvm::Value *> byRefVars;
     if (isByRef) {
       for (unsigned i = 0; i < opInst.getNumReductionVars(); ++i) {
@@ -1163,6 +1164,9 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
         byRefVars.push_back(builder.CreateAlloca(
             moduleTranslation.convertType(reductionDecls[i].getType())));
       }
+
+      initBlock = splitBB(builder, true, "omp.reduction.init");
+      allocaIP = InsertPointTy(allocaIP.getBlock(), allocaIP.getBlock()->end());
     }
 
     for (unsigned i = 0; i < opInst.getNumReductionVars(); ++i) {
@@ -1177,7 +1181,10 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
       assert(phis.size() == 1 &&
              "expected one value to be yielded from the "
              "reduction neutral element declaration region");
-
+      if (initBlock)
+        builder.SetInsertPoint(initBlock->getTerminator());
+      else
+        builder.restoreIP(allocaIP);
       if (isByRef) {
         // Store the result of the inlined region to the allocated reduction var
         // ptr

mlir/test/Target/LLVMIR/openmp-parallel-reduction-cleanup.mlir

@@ -55,11 +55,11 @@
 
 // Private reduction variable and its initialization.
 // CHECK: %tid.addr.local = alloca i32
-// CHECK: %[[MALLOC_I:.+]] = call ptr @malloc(i64 4)
 // CHECK: %[[PRIV_PTR_I:.+]] = alloca ptr
+// CHECK: %[[PRIV_PTR_J:.+]] = alloca ptr
+// CHECK: %[[MALLOC_I:.+]] = call ptr @malloc(i64 4)
 // CHECK: store ptr %[[MALLOC_I]], ptr %[[PRIV_PTR_I]]
 // CHECK: %[[MALLOC_J:.+]] = call ptr @malloc(i64 4)
-// CHECK: %[[PRIV_PTR_J:.+]] = alloca ptr
 // CHECK: store ptr %[[MALLOC_J]], ptr %[[PRIV_PTR_J]]
 
 // Call to the reduction function.

mlir/test/Target/LLVMIR/openmp-reduction-init-arg.mlir

@@ -59,13 +59,15 @@ module {
 // CHECK:         %[[VAL_17:.*]] = load i32, ptr %[[VAL_18:.*]], align 4
 // CHECK:         store i32 %[[VAL_17]], ptr %[[VAL_16]], align 4
 // CHECK:         %[[VAL_19:.*]] = load i32, ptr %[[VAL_16]], align 4
-// CHECK:         %[[VAL_20:.*]] = load { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, ptr %[[VAL_13]], align 8
 // CHECK:         %[[VAL_21:.*]] = alloca ptr, align 8
+// CHECK:         %[[VAL_23:.*]] = alloca ptr, align 8
+// CHECK:         %[[VAL_20:.*]] = load { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, ptr %[[VAL_13]], align 8
+// CHECK:         %[[VAL_24:.*]] = alloca [2 x ptr], align 8
+// CHECK:         br label %[[INIT_LABEL:.*]]
+// CHECK: [[INIT_LABEL]]:
 // CHECK:         store ptr %[[VAL_13]], ptr %[[VAL_21]], align 8
 // CHECK:         %[[VAL_22:.*]] = load { ptr, i64, i32, i8, i8, i8, i8, [1 x [3 x i64]] }, ptr %[[VAL_15]], align 8
-// CHECK:         %[[VAL_23:.*]] = alloca ptr, align 8
 // CHECK:         store ptr %[[VAL_15]], ptr %[[VAL_23]], align 8
-// CHECK:         %[[VAL_24:.*]] = alloca [2 x ptr], align 8
 // CHECK:         br label %[[VAL_25:.*]]
 // CHECK:       omp.par.region:                                   ; preds = %[[VAL_26:.*]]
 // CHECK:         br label %[[VAL_27:.*]]