[Mlir-commits] [llvm] [mlir] [OMPIRBuilder] - Fix emitTargetTaskProxyFunc to not generate empty functions (PR #126958)

[Kareem Ergawy]

================
@@ -7229,11 +7233,28 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::emitTargetTask(
       Builder, AllocaIP, ToBeDeleted, TargetTaskAllocaIP, "global.tid", false));
 
   Builder.restoreIP(TargetTaskBodyIP);
-
   if (Error Err = TaskBodyCB(DeviceID, RTLoc, TargetTaskAllocaIP))
     return Err;
 
-  OI.ExitBB = Builder.saveIP().getBlock();
+  // The outliner (CodeExtractor) extract a sequence or vector of blocks that
+  // it is given. These blocks are enumerated by
+  // OpenMPIRBuilder::OutlineInfo::collectBlocks which expects the OI.ExitBlock
+  // to be outside the region. In other words, OI.ExitBlock is expected to be
+  // the start of the region after the outlining. We used to set OI.ExitBlock
+  // to the InsertBlock after TaskBodyCB is done. This is fine in most cases
+  // except when the task body is a single basic block. In that case,
+  // OI.ExitBlock is set to the single task body block and will get left out of
+  // the outlining process. So, simply create a new empty block to which we
+  // uncoditionally branch from where TaskBodyCB left off
+  BasicBlock *TargetTaskContBlock =
+      BasicBlock::Create(Builder.getContext(), "target.task.cont");
+
+  auto *CurFn = Builder.GetInsertBlock()->getParent();
+  emitBranch(TargetTaskContBlock);
+  emitBlock(TargetTaskContBlock, CurFn, /*IsFinished=*/true);
+
+  OI.ExitBB = TargetTaskContBlock;
----------------
ergawy wrote:

I think we can use `splitBB` here.

```suggestion
  OI.ExitBB = splitBB(Builder, /*CreateBranch=*/true, "target.task.cont");
  emitBlock(OI.ExitBB, OI.ExitBB->getParent(), /*IsFinished=*/true);
```

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