[flang] [llvm] [mlir] [Flang][OpenMP][Taskloop] Translation support for taskloop construct (PR #166903)
Kaviya Rajendiran via llvm-commits
llvm-commits at lists.llvm.org
Tue Nov 18 07:37:26 PST 2025
================
@@ -2419,6 +2466,207 @@ convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder,
return success();
}
+// Converts an OpenMP taskloop construct into LLVM IR using OpenMPIRBuilder.
+static LogicalResult
+convertOmpTaskloopOp(Operation &opInst, llvm::IRBuilderBase &builder,
+ LLVM::ModuleTranslation &moduleTranslation) {
+ using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
+ auto taskloopOp = cast<omp::TaskloopOp>(opInst);
+ if (failed(checkImplementationStatus(opInst)))
+ return failure();
+
+ // It stores the pointer of allocated firstprivate copies,
+ // which can be used later for freeing the allocated space.
+ SmallVector<llvm::Value *> llvmFirstPrivateVars;
+ PrivateVarsInfo privateVarsInfo(taskloopOp);
+ TaskContextStructManager taskStructMgr{builder, moduleTranslation,
+ privateVarsInfo.privatizers};
+
+ llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
+ findAllocaInsertPoint(builder, moduleTranslation);
+
+ assert(builder.GetInsertPoint() == builder.GetInsertBlock()->end());
+ llvm::BasicBlock *taskloopStartBlock = llvm::BasicBlock::Create(
+ builder.getContext(), "omp.taskloop.start",
+ /*Parent=*/builder.GetInsertBlock()->getParent());
+ llvm::Instruction *branchToTaskloopStartBlock =
+ builder.CreateBr(taskloopStartBlock);
+ builder.SetInsertPoint(branchToTaskloopStartBlock);
+
+ llvm::BasicBlock *copyBlock =
+ splitBB(builder, /*CreateBranch=*/true, "omp.private.copy");
+ llvm::BasicBlock *initBlock =
+ splitBB(builder, /*CreateBranch=*/true, "omp.private.init");
+
+ LLVM::ModuleTranslation::SaveStack<OpenMPAllocaStackFrame> frame(
+ moduleTranslation, allocaIP);
+
+ // Allocate and initialize private variables
+ builder.SetInsertPoint(initBlock->getTerminator());
+
+ taskStructMgr.generateTaskContextStruct();
+ taskStructMgr.createGEPsToPrivateVars();
+
+ llvmFirstPrivateVars.resize(privateVarsInfo.blockArgs.size());
+ int index = 0;
+
+ for (auto [privDecl, mlirPrivVar, blockArg, llvmPrivateVarAlloc] :
+ llvm::zip_equal(privateVarsInfo.privatizers, privateVarsInfo.mlirVars,
+ privateVarsInfo.blockArgs,
+ taskStructMgr.getLLVMPrivateVarGEPs())) {
+ // To be handled inside the taskloop.
+ if (!privDecl.readsFromMold())
+ continue;
+ assert(llvmPrivateVarAlloc &&
+ "reads from mold so shouldn't have been skipped");
+
+ llvm::Expected<llvm::Value *> privateVarOrErr =
+ initPrivateVar(builder, moduleTranslation, privDecl, mlirPrivVar,
+ blockArg, llvmPrivateVarAlloc, initBlock);
+ if (!privateVarOrErr)
+ return handleError(privateVarOrErr, *taskloopOp.getOperation());
+
+ llvmFirstPrivateVars[index++] = privateVarOrErr.get();
+
+ llvm::IRBuilderBase::InsertPointGuard guard(builder);
+ builder.SetInsertPoint(builder.GetInsertBlock()->getTerminator());
+
+ if ((privateVarOrErr.get() != llvmPrivateVarAlloc) &&
+ !mlir::isa<LLVM::LLVMPointerType>(blockArg.getType())) {
+ builder.CreateStore(privateVarOrErr.get(), llvmPrivateVarAlloc);
+ // Load it so we have the value pointed to by the GEP
+ llvmPrivateVarAlloc = builder.CreateLoad(privateVarOrErr.get()->getType(),
+ llvmPrivateVarAlloc);
+ }
+ assert(llvmPrivateVarAlloc->getType() ==
+ moduleTranslation.convertType(blockArg.getType()));
+ }
+
+ // firstprivate copy region
+ setInsertPointForPossiblyEmptyBlock(builder, copyBlock);
+ if (failed(copyFirstPrivateVars(
+ taskloopOp, builder, moduleTranslation, privateVarsInfo.mlirVars,
+ taskStructMgr.getLLVMPrivateVarGEPs(), privateVarsInfo.privatizers,
+ taskloopOp.getPrivateNeedsBarrier())))
+ return llvm::failure();
+
+ // Set up inserttion point for call to createTaskloop()
+ builder.SetInsertPoint(taskloopStartBlock);
+
+ auto bodyCB = [&](InsertPointTy allocaIP,
+ InsertPointTy codegenIP) -> llvm::Error {
+ // Save the alloca insertion point on ModuleTranslation stack for use in
+ // nested regions.
+ LLVM::ModuleTranslation::SaveStack<OpenMPAllocaStackFrame> frame(
+ moduleTranslation, allocaIP);
+
+ // translate the body of the taskloop:
+ builder.restoreIP(codegenIP);
+
+ llvm::BasicBlock *privInitBlock = nullptr;
+ privateVarsInfo.llvmVars.resize(privateVarsInfo.blockArgs.size());
+ for (auto [i, zip] : llvm::enumerate(llvm::zip_equal(
+ privateVarsInfo.blockArgs, privateVarsInfo.privatizers,
+ privateVarsInfo.mlirVars))) {
+ auto [blockArg, privDecl, mlirPrivVar] = zip;
+ // This is handled before the task executes
+ if (privDecl.readsFromMold())
+ continue;
+
+ llvm::IRBuilderBase::InsertPointGuard guard(builder);
+ llvm::Type *llvmAllocType =
+ moduleTranslation.convertType(privDecl.getType());
+ builder.SetInsertPoint(allocaIP.getBlock()->getTerminator());
+ llvm::Value *llvmPrivateVar = builder.CreateAlloca(
+ llvmAllocType, /*ArraySize=*/nullptr, "omp.private.alloc");
+
+ llvm::Expected<llvm::Value *> privateVarOrError =
+ initPrivateVar(builder, moduleTranslation, privDecl, mlirPrivVar,
+ blockArg, llvmPrivateVar, privInitBlock);
+ if (!privateVarOrError)
+ return privateVarOrError.takeError();
+ moduleTranslation.mapValue(blockArg, privateVarOrError.get());
+ privateVarsInfo.llvmVars[i] = privateVarOrError.get();
+ // Add private var to llvmFirstPrivateVars
+ llvmFirstPrivateVars[index++] = privateVarOrError.get();
+ }
+
+ taskStructMgr.createGEPsToPrivateVars();
+ for (auto [i, llvmPrivVar] :
+ llvm::enumerate(taskStructMgr.getLLVMPrivateVarGEPs())) {
+ if (!llvmPrivVar) {
+ assert(privateVarsInfo.llvmVars[i] &&
+ "This is added in the loop above");
+ continue;
+ }
+ privateVarsInfo.llvmVars[i] = llvmPrivVar;
+ }
+
+ // Find and map the addresses of each variable within the taskloop context
+ // structure
+ for (auto [blockArg, llvmPrivateVar, privateDecl] :
+ llvm::zip_equal(privateVarsInfo.blockArgs, privateVarsInfo.llvmVars,
+ privateVarsInfo.privatizers)) {
+ // This was handled above.
+ if (!privateDecl.readsFromMold())
+ continue;
+ // Fix broken pass-by-value case for Fortran character boxes
+ if (!mlir::isa<LLVM::LLVMPointerType>(blockArg.getType())) {
+ llvmPrivateVar = builder.CreateLoad(
+ moduleTranslation.convertType(blockArg.getType()), llvmPrivateVar);
+ }
+ assert(llvmPrivateVar->getType() ==
+ moduleTranslation.convertType(blockArg.getType()));
+ moduleTranslation.mapValue(blockArg, llvmPrivateVar);
+ }
+
+ auto continuationBlockOrError =
+ convertOmpOpRegions(taskloopOp.getRegion(), "omp.taskloop.region",
+ builder, moduleTranslation);
+ ;
+ if (failed(handleError(continuationBlockOrError, opInst)))
+ return llvm::make_error<PreviouslyReportedError>();
+
+ builder.SetInsertPoint(continuationBlockOrError.get()->getTerminator());
+
+ // dummy check to ensure that the task context structure is accessed inside
+ // the outlined fn.
+ llvm::Value *cond = taskStructMgr.isAllocated();
----------------
kaviya2510 wrote:
Thankyou @tblah for taking the time to go through all my questions and respond so thoroughly.
During discussion with Tom on slack, I had proposed an another approach where we make use of the StructArg to store the loop bounds.
Here is the overview of new design:
1. When we generate a task, the outlined function must see its inputs from StructArg. In the body callback (before outlining), we must generate two GEPs:
1. A GEP, into the StructArg itself to record field (lb,ub,step) that the outlined function will later need.
```llvmir
%structArg = alloca { i64,i64,i64,ptr }, align 8
%task_lb = getelementptr { ptr, ptr, ptr, ptr }, ptr %1, i32 0, i32 0
%task_ub = getelementptr { ptr, ptr, ptr, ptr }, ptr %1, i32 0, i32 1
%task_st = getelementptr { ptr, ptr, ptr, ptr }, ptr %1, i32 0, i32 2
```
2. A GEP into the structArg so you can store the task-context pointer into the struct.
```
%gep_omp.task.context_ptr = getelementptr { ptr,ptr,ptr,ptr }, ptr %1, i32 0, i32 3
%loadgep_omp.task.context_ptr = load ptr, ptr %gep_omp.task.context_ptr, align 8, !align !5
```
3.`__kmpc_omp_task_alloc` returns a pointer to a newly allocated `kmp_task_t` object.
4. Copy the StructArg into the allocated task of type `kmp_task_t`
5. During the call to `@__kmpc_taskloop(..)`, the runtime access the `kmp_task_t` object and update the runtime-computed lb,ub and step value.
6. Later in the outlined function, we can extract the runtime bounds by GEP into the StructArg and task_context struct and adjust the loop-nest values.
With this approach, we can pack the loop bounds with StructArg and there is no need of `%struct.kmp_task_info = type { ptr, ptr, i32, ptr, ptr, i64, i64, i64 }`.
> llvm::Value *cond = taskStructMgr.isAllocated();
The dummy check is no longer required under the new approach. Because the loop bounds are packed into StructArg, their usage is guaranteed within the outlined function, which ensures that the second argument can be passed correctly.
I am not sure that I can spend enough time to implement this new design, as it will require a significant effort. If either @tblah or @Stylie777 is available, please feel free to take it up. If you both have other priority work, I will discuss with @kiranktp and decide how to proceed.
Thankyou!
https://github.com/llvm/llvm-project/pull/166903
More information about the llvm-commits
mailing list