[llvm] [OpenMPIRBuilder] Added `createTeams` (PR #65767)

[Johannes Doerfert via llvm-commits]

================
@@ -6106,6 +6108,148 @@ void OpenMPIRBuilder::loadOffloadInfoMetadata(Module &M) {
   }
 }
 
+OpenMPIRBuilder::InsertPointTy
+OpenMPIRBuilder::createTeams(const LocationDescription &Loc,
+                             BodyGenCallbackTy BodyGenCB) {
+  if (!updateToLocation(Loc)) {
+    return Loc.IP;
+  }
+
+  uint32_t SrcLocStrSize;
+  Constant *SrcLocStr = getOrCreateSrcLocStr(Loc, SrcLocStrSize);
+  Value *Ident = getOrCreateIdent(SrcLocStr, SrcLocStrSize);
+
+  // Splitting a basic block expects a terminator. Hence, creating an
+  // unreachable instruction, which will be deleted later.
+  UnreachableInst *UI = Builder.CreateUnreachable();
+  BasicBlock *CurrentBasicBlock = Builder.GetInsertBlock();
+
+  // The current basic block is split into four basic blocks. After outlining,
+  // they will be mapped as follows:
+  // ```
+  // def current_fn() {
+  //   current_basic_block:
+  //     br label %teams.exit
+  //   teams.exit:
+  //     ; instructions after task
+  // }
+  // def outlined_fn() {
+  //   teams.alloca:
+  //     br label %teams.body
+  //   teams.body:
+  //     ; instructions within teams body
+  // }
+  // ```
+  BasicBlock *AllocaBB = CurrentBasicBlock->splitBasicBlock(UI, "teams.alloca");
+  BasicBlock *BodyBB = AllocaBB->splitBasicBlock(UI, "teams.body");
+  BasicBlock *ExitBB = BodyBB->splitBasicBlock(UI, "teams.exit");
+
+  UI->eraseFromParent();
+
+  // Generate the body of teams.
+  InsertPointTy CodeGenIP(BodyBB, BodyBB->begin());
+  InsertPointTy AllocaIP(AllocaBB, AllocaBB->begin());
+  BodyGenCB(AllocaIP, CodeGenIP);
+
+  OutlineInfo OI;
+  OI.EntryBB = AllocaBB;
+  OI.ExitBB = ExitBB;
+  OI.PostOutlineCB = [this, Ident](Function &OutlinedFn) {
+    // The input IR here looks like the following-
+    // ```
+    // func @current_fn() {
+    //   outlined_fn(%args)
+    // }
+    // func @outlined_fn(%args) {
+    //   ; teams body
+    // }
+    // ```
+    //
+    // This is changed to the following-
+    //
+    // ```
+    // func @current_fn() {
+    //   runtime_call(..., wrapper_fn, ...)
+    // }
+    // func @wrapper_fn(..., %args) {
+    //   ; teams body
+    // }
+    // ```
+
+    // The outlined function has different inputs than what is expected from it.
+    // So, a wrapper function with expected signature is created and the
+    // required arguments are passed to the outlined function. The stale call
+    // instruction in current function will be replaced with a new call
+    // instruction for runtime call with the wrapper function. The outlined
+    // function is then inlined in the wrapper function and the call from the
+    // current function is removed.
+
+    assert(OutlinedFn.getNumUses() == 1 &&
+           "there must be a single user for the outlined function");
+    CallInst *StaleCI = cast<CallInst>(OutlinedFn.user_back());
+    assert(StaleCI && "Error while outlining - no CallInst user found for the "
+                      "outlined function.");
+    OutlinedFn.addFnAttr(Attribute::AttrKind::AlwaysInline);
+
+    // Create the wrapper function.
+    Builder.SetInsertPoint(StaleCI);
+    SmallVector<Type *> WrapperArgTys{Builder.getPtrTy(), Builder.getPtrTy()};
+    for (auto &Arg : OutlinedFn.args()) {
+      WrapperArgTys.push_back(Arg.getType());
+    }
+    FunctionCallee WrapperFuncVal = M.getOrInsertFunction(
+        "outlined_omp_teams",
+        FunctionType::get(Builder.getVoidTy(), WrapperArgTys, false));
----------------
jdoerfert wrote:

Where is this function generated? I assume it's internal, if so, we cannot use the name.

https://github.com/llvm/llvm-project/pull/65767