[llvm-branch-commits] [mlir] [mlir][OpenMP] Pack task private variables into a heap-allocated context struct (PR #125307)

[Tom Eccles via llvm-branch-commits]

================
@@ -1796,36 +1918,110 @@ convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder,
   // Allocate and initialize private variables
   // TODO: package private variables up in a structure
   builder.SetInsertPoint(initBlock->getTerminator());
-  for (auto [privDecl, mlirPrivVar, blockArg] :
-       llvm::zip_equal(privateDecls, mlirPrivateVars, privateBlockArgs)) {
-    llvm::Type *llvmAllocType =
-        moduleTranslation.convertType(privDecl.getType());
 
-    // Allocations:
-    builder.SetInsertPoint(allocaIP.getBlock()->getTerminator());
-    llvm::Value *llvmPrivateVar = builder.CreateAlloca(
-        llvmAllocType, /*ArraySize=*/nullptr, "omp.private.alloc");
+  // Create task variable structure
+  llvm::SmallVector<llvm::Value *> privateVarAllocations;
+  taskStructMgr.generateTaskContextStruct();
+  // GEPs so that we can initialize the variables. Don't use these GEPs inside
+  // of the body otherwise it will be the GEP not the struct which is fowarded
+  // to the outlined function. GEPs forwarded in this way are passed in a
+  // stack-allocated (by OpenMPIRBuilder) structure which is not safe for tasks
+  // which may not be executed until after the current stack frame goes out of
+  // scope.
+  taskStructMgr.createGEPsToPrivateVars(privateVarAllocations);
+
+  for (auto [privDecl, mlirPrivVar, blockArg, llvmPrivateVarAlloc] :
+       llvm::zip_equal(privateDecls, mlirPrivateVars, privateBlockArgs,
+                       privateVarAllocations)) {
+    if (!llvmPrivateVarAlloc)
+      // to be handled inside the task
+      continue;
 
-    // builder.SetInsertPoint(initBlock->getTerminator());
-    auto err =
+    llvm::Expected<llvm::Value *> privateVarOrErr =
         initPrivateVar(builder, moduleTranslation, privDecl, mlirPrivVar,
-                       blockArg, llvmPrivateVar, llvmPrivateVars, initBlock);
-    if (err)
+                       blockArg, llvmPrivateVarAlloc, initBlock);
+    if (auto err = privateVarOrErr.takeError())
       return handleError(std::move(err), *taskOp.getOperation());
+
+    llvm::IRBuilderBase::InsertPointGuard guard(builder);
+    builder.SetInsertPoint(builder.GetInsertBlock()->getTerminator());
----------------
tblah wrote:

Yes. I create all of the block structure before doing anything else so that it is in a known state.

The builder's insertion point is set here https://github.com/llvm/llvm-project/blob/d6d50cd96025d1aa74cbdb7f5da604cf104d9afc/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp#L1920

Then `initPrivateVar` might create one or more blocks for the init region but these will always be terminated.

This line is needed exactly because the block is terminated. If new init blocks were created in `initPrivateVar`, the builder will currently have an insertion point pointing right after the terminator. We want to add any new instructions for the hack for character boxes before the terminator, which is what this line does (I don't like the naming of `SetInsertPoint` in this case: it should be `SetInsertPointBefore`).

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