[Mlir-commits] [llvm] [mlir] [OMPIRBuilder] - Make offloading input data persist for deferred target tasks (PR #133499)

[Sergio Afonso]

================
@@ -7378,13 +7530,43 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::emitTargetTask(
 
     TaskData = Builder.CreateCall(TaskAllocFn, TaskAllocArgs);
 
+    Align Alignment = TaskData->getPointerAlignment(M.getDataLayout());
     if (HasShareds) {
-      Value *Shareds = StaleCI->getArgOperand(1);
-      Align Alignment = TaskData->getPointerAlignment(M.getDataLayout());
-      Value *TaskShareds = Builder.CreateLoad(VoidPtr, TaskData);
+      Value *Shareds = StaleCI->getArgOperand(SharedArgOperandNo);
+      Value *TaskT = Builder.CreateStructGEP(TaskWithPrivatesTy, TaskData, 0);
+      Value *TaskSharedsPtr = TaskT;
+      if (TaskWithPrivatesTy != Task) {
+        TaskSharedsPtr = Builder.CreateStructGEP(Task, TaskT, 0);
+      }
+      Value *TaskShareds = Builder.CreateLoad(VoidPtr, TaskSharedsPtr);
+
       Builder.CreateMemCpy(TaskShareds, Alignment, Shareds, Alignment,
                            SharedsSize);
     }
+    if (OffloadingArraysToPrivatize.size()) {
+      Value *Privates =
+          Builder.CreateStructGEP(TaskWithPrivatesTy, TaskData, 1);
+      for (unsigned int i = 0; i < OffloadingArraysToPrivatize.size(); ++i) {
+        Value *PtrToPrivatize = OffloadingArraysToPrivatize[i];
+        Type *ArrayType = nullptr;
+        if (auto *GEP = dyn_cast<GetElementPtrInst>(PtrToPrivatize))
+          ArrayType = GEP->getSourceElementType();
+        else if (auto *Alloca = dyn_cast<AllocaInst>(PtrToPrivatize))
+          ArrayType = Alloca->getAllocatedType();
+        else
+          llvm_unreachable("Unhandled Instruction type");
+        assert(ArrayType && "ArrayType cannot be nullptr");
----------------
skatrak wrote:

Nit: This is duplicated from `createTaskWithPrivatesTy` and I believe it would be a good idea to refactor and reuse. That should make things much easier whenever adding support for the missing instruction types.

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