[openmp][mlir] Lower parallel if to new fork_call_if function.

[DavidTruby]

Summary

This patch adds a new runtime function fork_call_if and uses that
to lower parallel if statements when going through OpenMPIRBuilder.

This fixes an issue where the OpenMPIRBuilder passes all arguments to
fork_call as a struct but this struct is not filled corretly in the
non-if branch by handling the fork inside the runtime.

[DavidTruby]

This is intended to fix https://github.com/llvm/llvm-project/issues/58233 .
Does this look like the right sort of approach to push the fork handling into the runtime?

[jdoerfert]

Nice. Thanks for fixing this the right way.

[kiranchandramohan]

Looks fine to me.

There are clang-format issues.

changed files:
    llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
    openmp/runtime/src/kmp.h

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

967

Does this need an update since there is no ThenBB?

1030

Nit: This comment needs an update.

mlir/test/Target/LLVMIR/openmp-llvm.mlir

197

Was there any particular reason to remove this test? Can you update this test instead?

openmp/runtime/src/kmp_csupport.cpp

344

Would we ever want a variable length parameter option like the kmpc_fork_call?

347

Nit: Use braces here to match the else.

351

Would this launch a separate thread?

[DavidTruby]

openmp/runtime/src/kmp_csupport.cpp

344

As long as this is used only by the OpenMPIRBuilder this parameter will always be a struct. Possibly there would be some use to changing it to a VA_ARGS function like the others for consistency, but I tend to prefer to avoid them whenever possible :)

351

No, this will just call the outlined function directly from the main thread (or I guess from whichever thread called __kmpc_fork_call_if)

[DavidTruby]

mlir/test/Target/LLVMIR/openmp-llvm.mlir

197

This test was testing whether the tid.addr and zero.addr values were calculated in the correct basic block when the blocks were nested. Since these are now calculated inside __kmpc_fork_call_if when needed this test doesn't actually test anything any more.

[DavidTruby]

Fix failing unit test

[kiranchandramohan]

LG. Please check whether a unit test is required for kmpc_fork_call_if before submitting.

[DavidTruby]

Make compatible with C and add unit test.

Adding the unit test also exposed an additional bug where no arguments are
passed if the region doesn't capture any variables. Instead we create a void
pointer to pass to fork_call_if in that case, and then check if that argument
exists before passing it on.

[DavidTruby]

Fix failing MLIR lit tests

[DavidTruby]

@jdoerfert I've had to change this quite a bit to cover some corner cases that came up, are you still happy with it?

[mstorsjo]

openmp/runtime/src/kmp.h

3904

This new function needs to be added to the dllexports file to make it available and usable on Windows.

[DavidTruby]

openmp/runtime/src/kmp.h

3904

Ok, I can add that; do you have any clue how the numbers for functions are chosen in that file? I don't really understand how dll exporting works on Windows

[mstorsjo]

openmp/runtime/src/kmp.h

3904

In general cases, it would be enough to just add the name (or mark the symbol with a dllexport attribute in the source code), but here in OpenMP, the exports also are given ordinal numbers (which need to be unique). I think the logic so far is to just pick the next free number.

D131830 is adding 292 and 293 (if it ends up reapplied in its current form), so I guess we could go with 294 to avoid conflicts - unless anyone of the actual OpenMP maintainers have a better suggestion?

[mstorsjo]

openmp/runtime/src/kmp.h

3904

CC @tianshilei1992 who has sorted out such issues before too.

I'd like to have this issue solved soonish in one way or another - this has had my nightly builds broken for many days already.

[mstorsjo]

openmp/runtime/src/kmp.h

3904

What do the maintainers, who I hope are reading this, prefer here - that I push a patch that just adds this function to dllexports with a free ordinal number picked for the function - or revert the patch that was applied last week so that we get it fixed with proper process?

[kiranchandramohan]

openmp/runtime/src/kmp.h

3904

@DavidTruby was/is not working today. If you can add a fix that would be great.

[mstorsjo]

openmp/runtime/src/kmp.h

3904

I pushed a fix for this in 15151315f76b0840423bd3bd62c5f9fc647d31c6 now.

[kiranchandramohan]

openmp/runtime/src/kmp.h

3904

Thanks @mstorsjo.

[tianshilei1992]

openmp/runtime/src/kmp_csupport.cpp

356

This doesn't work on ARM64. The last argument is not pushed to stack but it is accessed in the callee.

    0x10021a84c <+188>: 0x910053e0   add    x0, sp, #0x14
    0x10021a850 <+192>: 0x910043e1   add    x1, sp, #0x10
    0x10021a854 <+196>: 0xd63f0100   blr    x8
->  0x10021a858 <+200>: 0x14000006   b      0x10021a870               ; <+224> at kmp_csupport.cpp:360:36

x2 is supposed to contain the last argument, but for some reason it is not pushed. I'm looking at this right now. This one breaks tests on Apple Silicone.

[tianshilei1992]

openmp/runtime/src/kmp_csupport.cpp

356

Sorry, not "pushed to stack". It should be store to x2, based on ARM64 calling convention.

llvm/include/llvm/Frontend/OpenMP/OMPKinds.def

 
 __OMP_RTL(__kmpc_barrier, false, Void, IdentPtr, Int32)
 __OMP_RTL(__kmpc_cancel, false, Int32, IdentPtr, Int32, Int32)
 __OMP_RTL(__kmpc_cancel_barrier, false, Int32, IdentPtr, Int32)
 __OMP_RTL(__kmpc_flush, false, Void, IdentPtr)
 __OMP_RTL(__kmpc_global_thread_num, false, Int32, IdentPtr)
 __OMP_RTL(__kmpc_get_hardware_thread_id_in_block, false, Int32, )
 __OMP_RTL(__kmpc_fork_call, true, Void, IdentPtr, Int32, ParallelTaskPtr)
+__OMP_RTL(__kmpc_fork_call_if, true, Void, IdentPtr, Int32, ParallelTaskPtr,
+          Int1)
 __OMP_RTL(__kmpc_omp_taskwait, false, Int32, IdentPtr, Int32)
 __OMP_RTL(__kmpc_omp_taskyield, false, Int32, IdentPtr, Int32, /* Int */ Int32)
 __OMP_RTL(__kmpc_push_num_threads, false, Void, IdentPtr, Int32,
           /* Int */ Int32)
 __OMP_RTL(__kmpc_push_proc_bind, false, Void, IdentPtr, Int32, /* Int */ Int32)
 __OMP_RTL(__kmpc_omp_reg_task_with_affinity, false, Int32, IdentPtr, Int32,
           /* kmp_task_t */ VoidPtr, Int32,
           /* kmp_task_affinity_info_t */ VoidPtr)

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

IRBuilder<>::InsertPoint OpenMPIRBuilder::createParallel(
   SmallVector<Instruction *, 4> ToBeDeleted;
 
   // Change the location to the outer alloca insertion point to create and
   // initialize the allocas we pass into the parallel region.
   Builder.restoreIP(OuterAllocaIP);
   AllocaInst *TIDAddr = Builder.CreateAlloca(Int32, nullptr, "tid.addr");
   AllocaInst *ZeroAddr = Builder.CreateAlloca(Int32, nullptr, "zero.addr");
 
-  // If there is an if condition we actually use the TIDAddr and ZeroAddr in the
-  // program, otherwise we only need them for modeling purposes to get the
-  // associated arguments in the outlined function. In the former case,
-  // initialize the allocas properly, in the latter case, delete them later.
-  if (IfCondition) {
-    Builder.CreateStore(Constant::getNullValue(Int32), TIDAddr);
-    Builder.CreateStore(Constant::getNullValue(Int32), ZeroAddr);
-  } else {
-    ToBeDeleted.push_back(TIDAddr);
-    ToBeDeleted.push_back(ZeroAddr);
-  }
+  // We only need TIDAddr and ZeroAddr for modeling purposes to get the
+  // associated arguments in the outlined function, so we delete them later.
 
   // Create an artificial insertion point that will also ensure the blocks we
   // are about to split are not degenerated.
   auto *UI = new UnreachableInst(Builder.getContext(), InsertBB);
 
-  Instruction *ThenTI = UI, *ElseTI = nullptr;
-  if (IfCondition)
-    SplitBlockAndInsertIfThenElse(IfCondition, UI, &ThenTI, &ElseTI);
-
-  BasicBlock *ThenBB = ThenTI->getParent();
-  BasicBlock *PRegEntryBB = ThenBB->splitBasicBlock(ThenTI, "omp.par.entry");
+  BasicBlock *BB = UI->getParent();
+  BasicBlock *PRegEntryBB = BB->splitBasicBlock(UI, "omp.par.entry");
   BasicBlock *PRegBodyBB =
-      PRegEntryBB->splitBasicBlock(ThenTI, "omp.par.region");
+      PRegEntryBB->splitBasicBlock(UI, "omp.par.region");
   BasicBlock *PRegPreFiniBB =
-      PRegBodyBB->splitBasicBlock(ThenTI, "omp.par.pre_finalize");
+      PRegBodyBB->splitBasicBlock(UI, "omp.par.pre_finalize");
   BasicBlock *PRegExitBB =
-      PRegPreFiniBB->splitBasicBlock(ThenTI, "omp.par.exit");
+      PRegPreFiniBB->splitBasicBlock(UI, "omp.par.exit");
 
   auto FiniCBWrapper = [&](InsertPointTy IP) {
     // Hide "open-ended" blocks from the given FiniCB by setting the right jump
     // target to the region exit block.
     if (IP.getBlock()->end() == IP.getPoint()) {
       IRBuilder<>::InsertPointGuard IPG(Builder);
       Builder.restoreIP(IP);
       Instruction *I = Builder.CreateBr(PRegExitBB);
IRBuilder<>::InsertPoint OpenMPIRBuilder::createParallel(
   Instruction *PrivTID = Builder.CreateLoad(Int32, PrivTIDAddr, "tid");
 
   // Add some fake uses for OpenMP provided arguments.
   ToBeDeleted.push_back(Builder.CreateLoad(Int32, TIDAddr, "tid.addr.use"));
   Instruction *ZeroAddrUse =
       Builder.CreateLoad(Int32, ZeroAddr, "zero.addr.use");
   ToBeDeleted.push_back(ZeroAddrUse);
 
   // ThenBB
   //   |
   //   V
   // PRegionEntryBB         <- Privatization allocas are placed here.
   //   |
   //   V
   // PRegionBodyBB          <- BodeGen is invoked here.
   //   |
   //   V
   // PRegPreFiniBB          <- The block we will start finalization from.
   //   |
   //   V
   // PRegionExitBB          <- A common exit to simplify block collection.
   //
 
   LLVM_DEBUG(dbgs() << "Before body codegen: " << *OuterFn << "\n");
 
   // Let the caller create the body.
   assert(BodyGenCB && "Expected body generation callback!");
   InsertPointTy CodeGenIP(PRegBodyBB, PRegBodyBB->begin());
   BodyGenCB(InnerAllocaIP, CodeGenIP);
 
   LLVM_DEBUG(dbgs() << "After  body codegen: " << *OuterFn << "\n");
+  FunctionCallee RTLFn;
+  if (IfCondition)
+    RTLFn = getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_fork_call_if);
+  else
+    RTLFn = getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_fork_call);
 
-  FunctionCallee RTLFn = getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_fork_call);
   if (auto *F = dyn_cast<llvm::Function>(RTLFn.getCallee())) {
     if (!F->hasMetadata(llvm::LLVMContext::MD_callback)) {
       llvm::LLVMContext &Ctx = F->getContext();
       MDBuilder MDB(Ctx);
       // Annotate the callback behavior of the __kmpc_fork_call:
       //  - The callback callee is argument number 2 (microtask).
       //  - The first two arguments of the callback callee are unknown (-1).
       //  - All variadic arguments to the __kmpc_fork_call are passed to the
assert(OutlinedFn.arg_size() >= 2 &&
            "Expected at least tid and bounded tid as arguments");
     unsigned NumCapturedVars =
         OutlinedFn.arg_size() - /* tid & bounded tid */ 2;
 
     CallInst *CI = cast<CallInst>(OutlinedFn.user_back());
     CI->getParent()->setName("omp_parallel");
     Builder.SetInsertPoint(CI);
 
     // Build call __kmpc_fork_call(Ident, n, microtask, var1, .., varn);
     Value *ForkCallArgs[] = {
         Ident, Builder.getInt32(NumCapturedVars),
         Builder.CreateBitCast(&OutlinedFn, ParallelTaskPtr)};
 
     SmallVector<Value *, 16> RealArgs;
     RealArgs.append(std::begin(ForkCallArgs), std::end(ForkCallArgs));
+    if (IfCondition)
+      RealArgs.push_back(IfCondition);
     RealArgs.append(CI->arg_begin() + /* tid & bound tid */ 2, CI->arg_end());
 
     Builder.CreateCall(RTLFn, RealArgs);
 
     LLVM_DEBUG(dbgs() << "With fork_call placed: "
                       << *Builder.GetInsertBlock()->getParent() << "\n");
 
     InsertPointTy ExitIP(PRegExitBB, PRegExitBB->end());
 
     // Initialize the local TID stack location with the argument value.
     Builder.SetInsertPoint(PrivTID);
     Function::arg_iterator OutlinedAI = OutlinedFn.arg_begin();
     Builder.CreateStore(Builder.CreateLoad(Int32, OutlinedAI), PrivTIDAddr);
 
-    // If no "if" clause was present we do not need the call created during
-    // outlining, otherwise we reuse it in the serialized parallel region.
-    if (!ElseTI) {
-      CI->eraseFromParent();
-    } else {
-
-      // If an "if" clause was present we are now generating the serialized
-      // version into the "else" branch.
-      Builder.SetInsertPoint(ElseTI);
-
-      // Build calls __kmpc_serialized_parallel(&Ident, GTid);
-      Value *SerializedParallelCallArgs[] = {Ident, ThreadID};
-      Builder.CreateCall(
-          getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_serialized_parallel),
-          SerializedParallelCallArgs);
-
-      // OutlinedFn(&GTid, &zero, CapturedStruct);
-      CI->removeFromParent();
-      Builder.Insert(CI);
-
-      // __kmpc_end_serialized_parallel(&Ident, GTid);
-      Value *EndArgs[] = {Ident, ThreadID};
-      Builder.CreateCall(
-          getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_end_serialized_parallel),
-          EndArgs);
-
-      LLVM_DEBUG(dbgs() << "With serialized parallel region: "
-                        << *Builder.GetInsertBlock()->getParent() << "\n");
-    }
 
     for (Instruction *I : ToBeDeleted)
       I->eraseFromParent();
   };
 
   // Adjust the finalization stack, verify the adjustment, and call the
   // finalize function a last time to finalize values between the pre-fini
   // block and the exit block if we left the parallel "the normal way".

mlir/test/Target/LLVMIR/openmp-llvm.mlir

// CHECK: define internal void @[[OMP_OUTLINED_FN_NUM_THREADS_3_2]]
   // CHECK: call void @__kmpc_barrier
 
 // CHECK: define internal void @[[OMP_OUTLINED_FN_NUM_THREADS_3_1]]
   // CHECK: call void @__kmpc_barrier
 
 // CHECK: define void @test_omp_parallel_if_1(i32 %[[IF_VAR_1:.*]])
 llvm.func @test_omp_parallel_if_1(%arg0: i32) -> () {
 
-// Check that the allocas are emitted by the OpenMPIRBuilder at the top of the
-// function, before the condition. Allocas are only emitted by the builder when
-// the `if` clause is present. We match specific SSA value names since LLVM
-// actually produces those names.
-// CHECK: %tid.addr{{.*}} = alloca i32
-// CHECK: %zero.addr{{.*}} = alloca i32
-
-// CHECK: %[[IF_COND_VAR_1:.*]] = icmp slt i32 %[[IF_VAR_1]], 0
   %0 = llvm.mlir.constant(0 : index) : i32
   %1 = llvm.icmp "slt" %arg0, %0 : i32
+// CHECK: %[[IF_COND_VAR_1:.*]] = icmp slt i32 %[[IF_VAR_1]], 0
+
 
 // CHECK: %[[GTN_IF_1:.*]] = call i32 @__kmpc_global_thread_num(ptr @[[SI_VAR_IF_1:.*]])
-// CHECK: br i1 %[[IF_COND_VAR_1]], label %[[IF_COND_TRUE_BLOCK_1:.*]], label %[[IF_COND_FALSE_BLOCK_1:.*]]
-// CHECK: [[IF_COND_TRUE_BLOCK_1]]:
 // CHECK: br label %[[OUTLINED_CALL_IF_BLOCK_1:.*]]
 // CHECK: [[OUTLINED_CALL_IF_BLOCK_1]]:
-// CHECK: call void {{.*}} @__kmpc_fork_call(ptr @[[SI_VAR_IF_1]], {{.*}} @[[OMP_OUTLINED_FN_IF_1:.*]])
+// CHECK: call void {{.*}} @__kmpc_fork_call_if(ptr @[[SI_VAR_IF_1]], {{.*}} @[[OMP_OUTLINED_FN_IF_1:.*]], i1 %[[IF_COND_VAR_1]])
 // CHECK: br label %[[OUTLINED_EXIT_IF_1:.*]]
 // CHECK: [[OUTLINED_EXIT_IF_1]]:
-// CHECK: br label %[[OUTLINED_EXIT_IF_2:.*]]
-// CHECK: [[OUTLINED_EXIT_IF_2]]:
 // CHECK: br label %[[RETURN_BLOCK_IF_1:.*]]
-// CHECK: [[IF_COND_FALSE_BLOCK_1]]:
-// CHECK: call void @__kmpc_serialized_parallel(ptr @[[SI_VAR_IF_1]], i32 %[[GTN_IF_1]])
-// CHECK: call void @[[OMP_OUTLINED_FN_IF_1]]
-// CHECK: call void @__kmpc_end_serialized_parallel(ptr @[[SI_VAR_IF_1]], i32 %[[GTN_IF_1]])
-// CHECK: br label %[[RETURN_BLOCK_IF_1]]
   omp.parallel if(%1 : i1) {
     omp.barrier
     omp.terminator
   }
 
 // CHECK: [[RETURN_BLOCK_IF_1]]:
 // CHECK: ret void
   llvm.return
 }
 
 // CHECK: define internal void @[[OMP_OUTLINED_FN_IF_1]]
   // CHECK: call void @__kmpc_barrier
 
 // -----
 
-// CHECK-LABEL: @test_nested_alloca_ip
-llvm.func @test_nested_alloca_ip(%arg0: i32) -> () {
-
-  // Check that the allocas are emitted by the OpenMPIRBuilder at the top of
-  // the function, before the condition. Allocas are only emitted by the
-  // builder when the `if` clause is present. We match specific SSA value names
-  // since LLVM actually produces those names and ensure they come before the
-  // "icmp" that is the first operation we emit.
-  // CHECK: %tid.addr{{.*}} = alloca i32
-  // CHECK: %zero.addr{{.*}} = alloca i32
-  // CHECK: icmp slt i32 %{{.*}}, 0
-  %0 = llvm.mlir.constant(0 : index) : i32
-  %1 = llvm.icmp "slt" %arg0, %0 : i32
-
-  omp.parallel if(%1 : i1) {
-    // The "parallel" operation will be outlined, check the the function is
-    // produced. Inside that function, further allocas should be placed before
-    // another "icmp".
-    // CHECK: define
-    // CHECK: %tid.addr{{.*}} = alloca i32
-    // CHECK: %zero.addr{{.*}} = alloca i32
-    // CHECK: icmp slt i32 %{{.*}}, 1
-    %2 = llvm.mlir.constant(1 : index) : i32
-    %3 = llvm.icmp "slt" %arg0, %2 : i32
-
-    omp.parallel if(%3 : i1) {
-      // One more nesting level.
-      // CHECK: define
-      // CHECK: %tid.addr{{.*}} = alloca i32
-      // CHECK: %zero.addr{{.*}} = alloca i32
-      // CHECK: icmp slt i32 %{{.*}}, 2
-
-      %4 = llvm.mlir.constant(2 : index) : i32
-      %5 = llvm.icmp "slt" %arg0, %4 : i32
-
-      omp.parallel if(%5 : i1) {
-        omp.barrier
-        omp.terminator
-      }
-
-      omp.barrier
-      omp.terminator
-    }
-    omp.barrier
-    omp.terminator
-  }
-
-  llvm.return
-}
-
-// -----
-
 // CHECK-LABEL: define void @test_omp_parallel_3()
 llvm.func @test_omp_parallel_3() -> () {
   // CHECK: [[OMP_THREAD_3_1:%.*]] = call i32 @__kmpc_global_thread_num(ptr @{{[0-9]+}})
   // CHECK: call void @__kmpc_push_proc_bind(ptr @{{[0-9]+}}, i32 [[OMP_THREAD_3_1]], i32 2)
   // CHECK: call void{{.*}}@__kmpc_fork_call{{.*}}@[[OMP_OUTLINED_FN_3_1:.*]])
   omp.parallel proc_bind(master) {
     omp.barrier
     omp.terminator

openmp/runtime/src/kmp.h

 KMP_EXPORT kmp_int32 __kmpc_global_thread_num(ident_t *);
 KMP_EXPORT kmp_int32 __kmpc_global_num_threads(ident_t *);
 KMP_EXPORT kmp_int32 __kmpc_bound_thread_num(ident_t *);
 KMP_EXPORT kmp_int32 __kmpc_bound_num_threads(ident_t *);
 
 KMP_EXPORT kmp_int32 __kmpc_ok_to_fork(ident_t *);
 KMP_EXPORT void __kmpc_fork_call(ident_t *, kmp_int32 nargs,
                                  kmpc_micro microtask, ...);
+KMP_EXPORT void __kmpc_fork_call_if(ident_t *loc, kmp_int32 nargs,
+                                    kmpc_micro microtask, bool cond, void *args);
 
 KMP_EXPORT void __kmpc_serialized_parallel(ident_t *, kmp_int32 global_tid);
 KMP_EXPORT void __kmpc_end_serialized_parallel(ident_t *, kmp_int32 global_tid);
 
 KMP_EXPORT void __kmpc_flush(ident_t *);
 KMP_EXPORT void __kmpc_barrier(ident_t *, kmp_int32 global_tid);
 KMP_EXPORT kmp_int32 __kmpc_master(ident_t *, kmp_int32 global_tid);
 KMP_EXPORT void __kmpc_end_master(ident_t *, kmp_int32 global_tid);

openmp/runtime/src/kmp_csupport.cpp

#if KMP_STATS_ENABLED
   } else {
     KMP_POP_PARTITIONED_TIMER();
   }
 #endif // KMP_STATS_ENABLED
 }
 
 /*!
 @ingroup PARALLEL
-@param loc source location information
+@param loc  source location information
+@param microtask  pointer to callback routine consisting of outlined parallel
+construct
+@param cond  condition for running in parallel
+@param args  struct of pointers to shared variables that aren't global
+
+Perform a fork only if the condition is true.
+*/
+void __kmpc_fork_call_if(ident_t *loc, kmp_int32 argc, kmpc_micro microtask,
+                         bool cond, void *args) {
+  int gtid = __kmp_entry_gtid();
+  int zero = 0;
+  if (cond)
+    __kmpc_fork_call(loc, argc, microtask, args);
+  else {
+    __kmpc_serialized_parallel(loc, gtid);
+    microtask(&gtid, &zero, args);
+    __kmpc_end_serialized_parallel(loc, gtid);
+  }
+}
+
+/*!
+@ingroup PARALLEL
+@param loc source location information
 @param global_tid global thread number
 @param num_teams number of teams requested for the teams construct
 @param num_threads number of threads per team requested for the teams construct
 
 Set the number of teams to be used by the teams construct.
 This call is only required if the teams construct has a `num_teams` clause
 or a `thread_limit` clause (or both).
 */