[OMPIRBuilder] Fix shared clause for task construct

[psoni2628]

Summary

This patch fixes the shared clause for the task construct with multiple shared variables. The shareds field in the kmp_task_t is not an inline array in the struct, rather it is a pointer to an array. With an inline array, the pointer dereference to the outlined function body of the task would segmentation fault when accessed by the runtime.

[kiranchandramohan]

Thanks @psoni2628 for the patch. This is great.
Could you upload the patch with full context for ease of review?
Are the Clang OpenMP test failures related?

[psoni2628]

In D158462#4609718, @kiranchandramohan wrote:

Thanks @psoni2628 for the patch. This is great.
Could you upload the patch with full context for ease of review?
Are the Clang OpenMP test failures related?

Yes. So for context, the following Fortran test case was failing previously, but now it runs.

subroutine foo()
  implicit none
  integer::x,y
  x=0
   !$omp task shared(x,y)
   x=2
   y=3

   !$omp end task
  !$omp taskwait        
  print *, x, y
end subroutine foo

program p
implicit none

!$omp parallel
!$omp single
call foo()
!$omp end single
!$omp end parallel
end program p

As for the clang failures, it looks like they are related. I think I just need to update the CHECK lines, but I'm still figuring that out.

[kiranchandramohan]

In D158462#4610054, @psoni2628 wrote:

In D158462#4609718, @kiranchandramohan wrote:

Thanks @psoni2628 for the patch. This is great.
Could you upload the patch with full context for ease of review?
Are the Clang OpenMP test failures related?

Yes. So for context, the following Fortran test case was failing previously, but now it runs.

subroutine foo()
  implicit none
  integer::x,y
  x=0
   !$omp task shared(x,y)
   x=2
   y=3

   !$omp end task
  !$omp taskwait        
  print *, x, y
end subroutine foo

program p
implicit none

!$omp parallel
!$omp single
call foo()
!$omp end single
!$omp end parallel
end program p

As for the clang failures, it looks like they are related. I think I just need to update the CHECK lines, but I'm still figuring that out.

@psoni2628 By context, i meant showing the full diff for the files where changes are made.

To get a full diff, use one of the following commands (or just use Arcanist to upload your patch):

git show HEAD -U999999 > mypatch.patch
git diff -U999999 @{u} > mypatch.patch
git diff HEAD~1 -U999999 > mypatch.patch

[psoni2628]

• Upload full diff

[psoni2628]

In D158462#4610054, @psoni2628 wrote:

In D158462#4609718, @kiranchandramohan wrote:

Thanks @psoni2628 for the patch. This is great.
Could you upload the patch with full context for ease of review?
Are the Clang OpenMP test failures related?

Yes. So for context, the following Fortran test case was failing previously, but now it runs.

subroutine foo()
  implicit none
  integer::x,y
  x=0
   !$omp task shared(x,y)
   x=2
   y=3

   !$omp end task
  !$omp taskwait        
  print *, x, y
end subroutine foo

program p
implicit none

!$omp parallel
!$omp single
call foo()
!$omp end single
!$omp end parallel
end program p

As for the clang failures, it looks like they are related. I think I just need to update the CHECK lines, but I'm still figuring that out.

In D158462#4610054, @psoni2628 wrote:

In D158462#4609718, @kiranchandramohan wrote:

Thanks @psoni2628 for the patch. This is great.
Could you upload the patch with full context for ease of review?
Are the Clang OpenMP test failures related?

Yes. So for context, the following Fortran test case was failing previously, but now it runs.

subroutine foo()
  implicit none
  integer::x,y
  x=0
   !$omp task shared(x,y)
   x=2
   y=3

   !$omp end task
  !$omp taskwait        
  print *, x, y
end subroutine foo

program p
implicit none

!$omp parallel
!$omp single
call foo()
!$omp end single
!$omp end parallel
end program p

As for the clang failures, it looks like they are related. I think I just need to update the CHECK lines, but I'm still figuring that out.

In D158462#4610148, @kiranchandramohan wrote:

In D158462#4610054, @psoni2628 wrote:

In D158462#4609718, @kiranchandramohan wrote:

Thanks @psoni2628 for the patch. This is great.
Could you upload the patch with full context for ease of review?
Are the Clang OpenMP test failures related?

Yes. So for context, the following Fortran test case was failing previously, but now it runs.

subroutine foo()
  implicit none
  integer::x,y
  x=0
   !$omp task shared(x,y)
   x=2
   y=3

   !$omp end task
  !$omp taskwait        
  print *, x, y
end subroutine foo

program p
implicit none

!$omp parallel
!$omp single
call foo()
!$omp end single
!$omp end parallel
end program p

As for the clang failures, it looks like they are related. I think I just need to update the CHECK lines, but I'm still figuring that out.

@psoni2628 By context, i meant showing the full diff for the files where changes are made.

To get a full diff, use one of the following commands (or just use Arcanist to upload your patch):

git show HEAD -U999999 > mypatch.patch
git diff -U999999 @{u} > mypatch.patch
git diff HEAD~1 -U999999 > mypatch.patch

Done. I forgot to do that, sorry.

[psoni2628]

In D158462#4610054, @psoni2628 wrote:

In D158462#4609718, @kiranchandramohan wrote:

Thanks @psoni2628 for the patch. This is great.
Could you upload the patch with full context for ease of review?
Are the Clang OpenMP test failures related?

Yes. So for context, the following Fortran test case was failing previously, but now it runs.

subroutine foo()
  implicit none
  integer::x,y
  x=0
   !$omp task shared(x,y)
   x=2
   y=3

   !$omp end task
  !$omp taskwait        
  print *, x, y
end subroutine foo

program p
implicit none

!$omp parallel
!$omp single
call foo()
!$omp end single
!$omp end parallel
end program p

As for the clang failures, it looks like they are related. I think I just need to update the CHECK lines, but I'm still figuring that out.

The clang test failures are being caused by my addition of kmp_task_t to OMPKinds.def. It is conflicting with Clang's definition of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp, so the struct gets renamed to kmp_task_t.0 and kmp_task_t.1. I don't think it is that simple to update the usage of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp. Should I rename my addition of kmp_task_t in OMPKinds.def to kmp_task, or should I just fix the tests to allow kmp_task_t.0 and kmp_task_t.1?

[kiranchandramohan]

In D158462#4610220, @psoni2628 wrote:

In D158462#4610054, @psoni2628 wrote:

In D158462#4609718, @kiranchandramohan wrote:

Thanks @psoni2628 for the patch. This is great.
Could you upload the patch with full context for ease of review?
Are the Clang OpenMP test failures related?

Yes. So for context, the following Fortran test case was failing previously, but now it runs.

subroutine foo()
  implicit none
  integer::x,y
  x=0
   !$omp task shared(x,y)
   x=2
   y=3

   !$omp end task
  !$omp taskwait        
  print *, x, y
end subroutine foo

program p
implicit none

!$omp parallel
!$omp single
call foo()
!$omp end single
!$omp end parallel
end program p

As for the clang failures, it looks like they are related. I think I just need to update the CHECK lines, but I'm still figuring that out.

The clang test failures are being caused by my addition of kmp_task_t to OMPKinds.def. It is conflicting with Clang's definition of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp, so the struct gets renamed to kmp_task_t.0 and kmp_task_t.1. I don't think it is that simple to update the usage of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp. Should I rename my addition of kmp_task_t in OMPKinds.def to kmp_task, or should I just fix the tests to allow kmp_task_t.0 and kmp_task_t.1?

ping @jdoerfert

[jdoerfert]

In D158462#4610257, @kiranchandramohan wrote:

In D158462#4610220, @psoni2628 wrote:

As for the clang failures, it looks like they are related. I think I just need to update the CHECK lines, but I'm still figuring that out.

The clang test failures are being caused by my addition of kmp_task_t to OMPKinds.def. It is conflicting with Clang's definition of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp, so the struct gets renamed to kmp_task_t.0 and kmp_task_t.1. I don't think it is that simple to update the usage of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp. Should I rename my addition of kmp_task_t in OMPKinds.def to kmp_task, or should I just fix the tests to allow kmp_task_t.0 and kmp_task_t.1?

Can we make use of the new kmp_task_t in Clang? That would be the correct way to resolve this.

[psoni2628]

In D158462#4610370, @jdoerfert wrote:

In D158462#4610257, @kiranchandramohan wrote:

In D158462#4610220, @psoni2628 wrote:

As for the clang failures, it looks like they are related. I think I just need to update the CHECK lines, but I'm still figuring that out.

The clang test failures are being caused by my addition of kmp_task_t to OMPKinds.def. It is conflicting with Clang's definition of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp, so the struct gets renamed to kmp_task_t.0 and kmp_task_t.1. I don't think it is that simple to update the usage of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp. Should I rename my addition of kmp_task_t in OMPKinds.def to kmp_task, or should I just fix the tests to allow kmp_task_t.0 and kmp_task_t.1?

Can we make use of the new kmp_task_t in Clang? That would be the correct way to resolve this.

I am not sure how to go about doing that. I think the kmp_task_t struct is actually being created in llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp:initializeTypes from the macro defined in OMPKinds.def. Do you know how I can refer to that struct in Clang?

[psoni2628]

In D158462#4610644, @psoni2628 wrote:

In D158462#4610370, @jdoerfert wrote:

In D158462#4610257, @kiranchandramohan wrote:

In D158462#4610220, @psoni2628 wrote:

As for the clang failures, it looks like they are related. I think I just need to update the CHECK lines, but I'm still figuring that out.

The clang test failures are being caused by my addition of kmp_task_t to OMPKinds.def. It is conflicting with Clang's definition of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp, so the struct gets renamed to kmp_task_t.0 and kmp_task_t.1. I don't think it is that simple to update the usage of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp. Should I rename my addition of kmp_task_t in OMPKinds.def to kmp_task, or should I just fix the tests to allow kmp_task_t.0 and kmp_task_t.1?

Can we make use of the new kmp_task_t in Clang? That would be the correct way to resolve this.

I am not sure how to go about doing that. I think the kmp_task_t struct is actually being created in llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp:initializeTypes from the macro defined in OMPKinds.def. Do you know how I can refer to that struct in Clang?

I looked at this more closely, and I don't think it's possible to refer to kmp_task_t in Clang. The only way I can think of doing this is doing similarly to D123460, which is basically building the structure in OMPIRBuilder, and calling it in Clang. This is not a straightforward change to make though, because Clang builds several RecordDecls on top of kmp_task_t such as the taskloop_task and task_with_privates. We would also have to build those structs in OMPIRBuilder, but that expands the scope of this patch quite significantly. Is this what you meant @jdoerfert?

[kiranchandramohan]

In D158462#4614095, @psoni2628 wrote:

In D158462#4610644, @psoni2628 wrote:

In D158462#4610370, @jdoerfert wrote:

In D158462#4610257, @kiranchandramohan wrote:

In D158462#4610220, @psoni2628 wrote:

As for the clang failures, it looks like they are related. I think I just need to update the CHECK lines, but I'm still figuring that out.

The clang test failures are being caused by my addition of kmp_task_t to OMPKinds.def. It is conflicting with Clang's definition of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp, so the struct gets renamed to kmp_task_t.0 and kmp_task_t.1. I don't think it is that simple to update the usage of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp. Should I rename my addition of kmp_task_t in OMPKinds.def to kmp_task, or should I just fix the tests to allow kmp_task_t.0 and kmp_task_t.1?

Can we make use of the new kmp_task_t in Clang? That would be the correct way to resolve this.

I am not sure how to go about doing that. I think the kmp_task_t struct is actually being created in llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp:initializeTypes from the macro defined in OMPKinds.def. Do you know how I can refer to that struct in Clang?

I looked at this more closely, and I don't think it's possible to refer to kmp_task_t in Clang. The only way I can think of doing this is doing similarly to D123460, which is basically building the structure in OMPIRBuilder, and calling it in Clang. This is not a straightforward change to make though, because Clang builds several RecordDecls on top of kmp_task_t such as the taskloop_task and task_with_privates. We would also have to build those structs in OMPIRBuilder, but that expands the scope of this patch quite significantly. Is this what you meant @jdoerfert?

I think we should proceed with either of the following two options:
-> Reuse Clang's kmp_task_t in the OpenMPIRBuilder and use only the relevant fields that we support now. I am assuming even if we have to move the definition of the Structure to OpenMPIRBuilder, there will not be much change in clang since the definition of the Structure remains the same. It is just a bit of additional logical to correctly dereference the portion of the structure in the OPenMPIRBuilder.
-> Define a separate kmp_task_ompbuilder_t and use that in the OpenMPIRBuilder.

[psoni2628]

In D158462#4631899, @kiranchandramohan wrote:

In D158462#4614095, @psoni2628 wrote:

In D158462#4610644, @psoni2628 wrote:

In D158462#4610370, @jdoerfert wrote:

In D158462#4610257, @kiranchandramohan wrote:

In D158462#4610220, @psoni2628 wrote:

As for the clang failures, it looks like they are related. I think I just need to update the CHECK lines, but I'm still figuring that out.

The clang test failures are being caused by my addition of kmp_task_t to OMPKinds.def. It is conflicting with Clang's definition of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp, so the struct gets renamed to kmp_task_t.0 and kmp_task_t.1. I don't think it is that simple to update the usage of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp. Should I rename my addition of kmp_task_t in OMPKinds.def to kmp_task, or should I just fix the tests to allow kmp_task_t.0 and kmp_task_t.1?

Can we make use of the new kmp_task_t in Clang? That would be the correct way to resolve this.

I am not sure how to go about doing that. I think the kmp_task_t struct is actually being created in llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp:initializeTypes from the macro defined in OMPKinds.def. Do you know how I can refer to that struct in Clang?

I looked at this more closely, and I don't think it's possible to refer to kmp_task_t in Clang. The only way I can think of doing this is doing similarly to D123460, which is basically building the structure in OMPIRBuilder, and calling it in Clang. This is not a straightforward change to make though, because Clang builds several RecordDecls on top of kmp_task_t such as the taskloop_task and task_with_privates. We would also have to build those structs in OMPIRBuilder, but that expands the scope of this patch quite significantly. Is this what you meant @jdoerfert?

I think we should proceed with either of the following two options:
-> Reuse Clang's kmp_task_t in the OpenMPIRBuilder and use only the relevant fields that we support now. I am assuming even if we have to move the definition of the Structure to OpenMPIRBuilder, there will not be much change in clang since the definition of the Structure remains the same. It is just a bit of additional logical to correctly dereference the portion of the structure in the OPenMPIRBuilder.
-> Define a separate kmp_task_ompbuilder_t and use that in the OpenMPIRBuilder.

I don't think option 1 is possible. The struct gets created in https://github.com/llvm/llvm-project/blob/c987f9d7fdc7b22c9bf68d7b3f0df10b68c679be/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp#L5680. IMO, the only way this makes sense is to create a function in OMPIRBuilder to populate the kmp_task_t, and then call that in Clang (similar to D123460). That is straightforward, but the issue is that we currently only populate the shareds field of kmp_task_t in OMPIRBuilder. We need to populate the other fields as well, otherwise we will lose the functionality in Clang. Once we have the full support for all the features of kmp_task_t in OMPIRBuilder, then we can use it. Until that happens we can just use a different struct name (i.e. the second option).

[psoni2628]

• Change struct name to kmp_task_ompbuilder_t from kmp_task_t

[jdoerfert]

In D158462#4614095, @psoni2628 wrote:

In D158462#4610644, @psoni2628 wrote:

In D158462#4610370, @jdoerfert wrote:

In D158462#4610257, @kiranchandramohan wrote:

In D158462#4610220, @psoni2628 wrote:

As for the clang failures, it looks like they are related. I think I just need to update the CHECK lines, but I'm still figuring that out.

The clang test failures are being caused by my addition of kmp_task_t to OMPKinds.def. It is conflicting with Clang's definition of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp, so the struct gets renamed to kmp_task_t.0 and kmp_task_t.1. I don't think it is that simple to update the usage of kmp_task_t in clang/lib/CodeGen/CGOpenMPRuntime.cpp. Should I rename my addition of kmp_task_t in OMPKinds.def to kmp_task, or should I just fix the tests to allow kmp_task_t.0 and kmp_task_t.1?

Can we make use of the new kmp_task_t in Clang? That would be the correct way to resolve this.

I am not sure how to go about doing that. I think the kmp_task_t struct is actually being created in llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp:initializeTypes from the macro defined in OMPKinds.def. Do you know how I can refer to that struct in Clang?

I looked at this more closely, and I don't think it's possible to refer to kmp_task_t in Clang. The only way I can think of doing this is doing similarly to D123460, which is basically building the structure in OMPIRBuilder, and calling it in Clang. This is not a straightforward change to make though, because Clang builds several RecordDecls on top of kmp_task_t such as the taskloop_task and task_with_privates. We would also have to build those structs in OMPIRBuilder, but that expands the scope of this patch quite significantly. Is this what you meant @jdoerfert?

Sorry for the delay. Yes, we should create all the structs in the OMPIRBuilder (or similar) and stop doing records in clang.
That said, I do not block this, but having two places creating "the same struct" is not a good idea and should be avoided as soon as possible.

[kiranchandramohan]

LG.

@shraiysh You might not remember much of this. But would you have any comments here?

llvm/include/llvm/Frontend/OpenMP/OMPKinds.def

 __OMP_STRUCT_TYPE(Ident, ident_t, false, Int32, Int32, Int32, Int32, Int8Ptr)
 __OMP_STRUCT_TYPE(OffloadEntry, __tgt_offload_entry, false, Int8Ptr, Int8Ptr, SizeTy,
                   Int32, Int32)
 __OMP_STRUCT_TYPE(KernelArgs, __tgt_kernel_arguments, false, Int32, Int32, VoidPtrPtr,
       VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr, VoidPtrPtr,
       Int64, Int64, Int32Arr3Ty, Int32Arr3Ty, Int32)
 __OMP_STRUCT_TYPE(AsyncInfo, __tgt_async_info, false, Int8Ptr)
 __OMP_STRUCT_TYPE(DependInfo, kmp_dep_info, false, SizeTy, SizeTy, Int8)
+__OMP_STRUCT_TYPE(Task, kmp_task_ompbuilder_t, false, VoidPtr, VoidPtr, Int32, VoidPtr, VoidPtr)
 __OMP_STRUCT_TYPE(ConfigurationEnvironment, ConfigurationEnvironmentTy, false,
                   Int8, Int8, Int8)
 __OMP_STRUCT_TYPE(DynamicEnvironment, DynamicEnvironmentTy, false, Int16)
 __OMP_STRUCT_TYPE(KernelEnvironment, KernelEnvironmentTy, false,
                   ConfigurationEnvironment, IdentPtr, DynamicEnvironmentPtr)
 
 #undef __OMP_STRUCT_TYPE
 #undef OMP_STRUCT_TYPE

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

OI.PostOutlineCB = [this, Ident, Tied, Final, IfCondition,
     // ```
 
     // The stale call instruction will be replaced with a new call instruction
     // for runtime call with a wrapper function.
     assert(OutlinedFn.getNumUses() == 1 &&
            "there must be a single user for the outlined function");
     CallInst *StaleCI = cast<CallInst>(OutlinedFn.user_back());
 
-    // HasTaskData is true if any variables are captured in the outlined region,
+    // HasShareds is true if any variables are captured in the outlined region,
     // false otherwise.
-    bool HasTaskData = StaleCI->arg_size() > 0;
+    bool HasShareds = StaleCI->arg_size() > 0;
     Builder.SetInsertPoint(StaleCI);
 
     // Gather the arguments for emitting the runtime call for
     // @__kmpc_omp_task_alloc
     Function *TaskAllocFn =
         getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_omp_task_alloc);
 
     // Arguments - `loc_ref` (Ident) and `gtid` (ThreadID)
if (Final) {
       Value *FinalFlag =
           Builder.CreateSelect(Final, Builder.getInt32(2), Builder.getInt32(0));
       Flags = Builder.CreateOr(FinalFlag, Flags);
     }
 
     // Argument - `sizeof_kmp_task_t` (TaskSize)
     // Tasksize refers to the size in bytes of kmp_task_t data structure
     // including private vars accessed in task.
-    Value *TaskSize = Builder.getInt64(0);
-    if (HasTaskData) {
+    // TODO: add kmp_task_t_with_privates (privates)
+    Value *TaskSize = Builder.getInt64(
+        divideCeil(M.getDataLayout().getTypeSizeInBits(Task), 8));
+
+    // Argument - `sizeof_shareds` (SharedsSize)
+    // SharedsSize refers to the shareds array size in the kmp_task_t data
+    // structure.
+    Value *SharedsSize = Builder.getInt64(0);
+    if (HasShareds) {
       AllocaInst *ArgStructAlloca =
           dyn_cast<AllocaInst>(StaleCI->getArgOperand(0));
       assert(ArgStructAlloca &&
              "Unable to find the alloca instruction corresponding to arguments "
              "for extracted function");
       StructType *ArgStructType =
           dyn_cast<StructType>(ArgStructAlloca->getAllocatedType());
       assert(ArgStructType && "Unable to find struct type corresponding to "
                               "arguments for extracted function");
-      TaskSize =
+      SharedsSize =
           Builder.getInt64(M.getDataLayout().getTypeStoreSize(ArgStructType));
     }
 
-    // TODO: Argument - sizeof_shareds
-
     // Argument - task_entry (the wrapper function)
-    // If the outlined function has some captured variables (i.e. HasTaskData is
+    // If the outlined function has some captured variables (i.e. HasShareds is
     // true), then the wrapper function will have an additional argument (the
     // struct containing captured variables). Otherwise, no such argument will
     // be present.
     SmallVector<Type *> WrapperArgTys{Builder.getInt32Ty()};
-    if (HasTaskData)
+    if (HasShareds)
       WrapperArgTys.push_back(OutlinedFn.getArg(0)->getType());
     FunctionCallee WrapperFuncVal = M.getOrInsertFunction(
         (Twine(OutlinedFn.getName()) + ".wrapper").str(),
         FunctionType::get(Builder.getInt32Ty(), WrapperArgTys, false));
     Function *WrapperFunc = dyn_cast<Function>(WrapperFuncVal.getCallee());
 
     // Emit the @__kmpc_omp_task_alloc runtime call
     // The runtime call returns a pointer to an area where the task captured
-    // variables must be copied before the task is run (NewTaskData)
-    CallInst *NewTaskData = Builder.CreateCall(
-        TaskAllocFn,
-        {/*loc_ref=*/Ident, /*gtid=*/ThreadID, /*flags=*/Flags,
-         /*sizeof_task=*/TaskSize, /*sizeof_shared=*/Builder.getInt64(0),
+    // variables must be copied before the task is run (TaskData)
+    CallInst *TaskData = Builder.CreateCall(
+        TaskAllocFn, {/*loc_ref=*/Ident, /*gtid=*/ThreadID, /*flags=*/Flags,
+                      /*sizeof_task=*/TaskSize, /*sizeof_shared=*/SharedsSize,
                       /*task_func=*/WrapperFunc});
 
     // Copy the arguments for outlined function
-    if (HasTaskData) {
-      Value *TaskData = StaleCI->getArgOperand(0);
+    if (HasShareds) {
+      Value *Shareds = StaleCI->getArgOperand(0);
       Align Alignment = TaskData->getPointerAlignment(M.getDataLayout());
-      Builder.CreateMemCpy(NewTaskData, Alignment, TaskData, Alignment,
-                           TaskSize);
+      Value *TaskShareds = Builder.CreateLoad(VoidPtr, TaskData);
+      Builder.CreateMemCpy(TaskShareds, Alignment, Shareds, Alignment,
+                           SharedsSize);
     }
 
     Value *DepArrayPtr = nullptr;
     if (Dependencies.size()) {
       InsertPointTy OldIP = Builder.saveIP();
       Builder.SetInsertPoint(
           &OldIP.getBlock()->getParent()->getEntryBlock().back());
 
if (IfCondition) {
       Builder.SetInsertPoint(IfTerminator);
       SplitBlockAndInsertIfThenElse(IfCondition, IfTerminator, &ThenTI,
                                     &ElseTI);
       Builder.SetInsertPoint(ElseTI);
       Function *TaskBeginFn =
           getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_omp_task_begin_if0);
       Function *TaskCompleteFn =
           getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_omp_task_complete_if0);
-      Builder.CreateCall(TaskBeginFn, {Ident, ThreadID, NewTaskData});
-      if (HasTaskData)
-        Builder.CreateCall(WrapperFunc, {ThreadID, NewTaskData});
+      Builder.CreateCall(TaskBeginFn, {Ident, ThreadID, TaskData});
+      if (HasShareds)
+        Builder.CreateCall(WrapperFunc, {ThreadID, TaskData});
       else
         Builder.CreateCall(WrapperFunc, {ThreadID});
-      Builder.CreateCall(TaskCompleteFn, {Ident, ThreadID, NewTaskData});
+      Builder.CreateCall(TaskCompleteFn, {Ident, ThreadID, TaskData});
       Builder.SetInsertPoint(ThenTI);
     }
 
     if (Dependencies.size()) {
       Function *TaskFn =
           getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_omp_task_with_deps);
       Builder.CreateCall(
           TaskFn,
-          {Ident, ThreadID, NewTaskData, Builder.getInt32(Dependencies.size()),
+          {Ident, ThreadID, TaskData, Builder.getInt32(Dependencies.size()),
            DepArrayPtr, ConstantInt::get(Builder.getInt32Ty(), 0),
            ConstantPointerNull::get(Type::getInt8PtrTy(M.getContext()))});
 
     } else {
       // Emit the @__kmpc_omp_task runtime call to spawn the task
       Function *TaskFn = getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_omp_task);
-      Builder.CreateCall(TaskFn, {Ident, ThreadID, NewTaskData});
+      Builder.CreateCall(TaskFn, {Ident, ThreadID, TaskData});
     }
 
     StaleCI->eraseFromParent();
 
     // Emit the body for wrapper function
     BasicBlock *WrapperEntryBB =
         BasicBlock::Create(M.getContext(), "", WrapperFunc);
     Builder.SetInsertPoint(WrapperEntryBB);
-    if (HasTaskData)
-      Builder.CreateCall(&OutlinedFn, {WrapperFunc->getArg(1)});
-    else
+    if (HasShareds) {
+      llvm::Value *Shareds =
+          Builder.CreateLoad(VoidPtr, WrapperFunc->getArg(1));
+      Builder.CreateCall(&OutlinedFn, {Shareds});
+    } else {
       Builder.CreateCall(&OutlinedFn);
+    }
     Builder.CreateRet(Builder.getInt32(0));
   };
 
   addOutlineInfo(std::move(OI));
 
   InsertPointTy TaskAllocaIP =
       InsertPointTy(TaskAllocaBB, TaskAllocaBB->begin());
   InsertPointTy TaskBodyIP = InsertPointTy(TaskBodyBB, TaskBodyBB->begin());

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

TEST_F(OpenMPIRBuilderTest, CreateTask) {
   ConstantInt *Flags = dyn_cast<ConstantInt>(TaskAllocCall->getArgOperand(2));
   ASSERT_NE(Flags, nullptr);
   EXPECT_EQ(Flags->getSExtValue(), 1);
 
   // Verify the data size
   ConstantInt *DataSize =
       dyn_cast<ConstantInt>(TaskAllocCall->getArgOperand(3));
   ASSERT_NE(DataSize, nullptr);
-  EXPECT_EQ(DataSize->getSExtValue(), 24); // 64-bit pointer + 128-bit integer
+  EXPECT_EQ(DataSize->getSExtValue(), 40);
 
-  // TODO: Verify size of shared clause variables
+  ConstantInt *SharedsSize =
+      dyn_cast<ConstantInt>(TaskAllocCall->getOperand(4));
+  EXPECT_EQ(SharedsSize->getSExtValue(),
+            24); // 64-bit pointer + 128-bit integer
 
   // Verify Wrapper function
   Function *WrapperFunc =
       dyn_cast<Function>(TaskAllocCall->getArgOperand(5)->stripPointerCasts());
   ASSERT_NE(WrapperFunc, nullptr);
+
+  LoadInst *SharedsLoad = dyn_cast<LoadInst>(WrapperFunc->begin()->begin());
+  ASSERT_NE(SharedsLoad, nullptr);
+  EXPECT_EQ(SharedsLoad->getPointerOperand(), WrapperFunc->getArg(1));
+
   EXPECT_FALSE(WrapperFunc->isDeclaration());
-  CallInst *OutlinedFnCall = dyn_cast<CallInst>(WrapperFunc->begin()->begin());
+  CallInst *OutlinedFnCall =
+      dyn_cast<CallInst>(++WrapperFunc->begin()->begin());
   ASSERT_NE(OutlinedFnCall, nullptr);
   EXPECT_EQ(WrapperFunc->getArg(0)->getType(), Builder.getInt32Ty());
-  EXPECT_EQ(OutlinedFnCall->getArgOperand(0), WrapperFunc->getArg(1));
+  EXPECT_EQ(OutlinedFnCall->getArgOperand(0),
+            WrapperFunc->getArg(1)->uses().begin()->getUser());
 
   // Verify the presence of `trunc` and `icmp` instructions in Outlined function
   Function *OutlinedFn = OutlinedFnCall->getCalledFunction();
   ASSERT_NE(OutlinedFn, nullptr);
   EXPECT_TRUE(any_of(instructions(OutlinedFn),
                      [](Instruction &inst) { return isa<TruncInst>(&inst); }));
   EXPECT_TRUE(any_of(instructions(OutlinedFn),
                      [](Instruction &inst) { return isa<ICmpInst>(&inst); }));

mlir/test/Target/LLVMIR/openmp-llvm.mlir

 
 // -----
 
 // CHECK-LABEL: define void @omp_task
 // CHECK-SAME: (i32 %[[x:.+]], i32 %[[y:.+]], ptr %[[zaddr:.+]])
 llvm.func @omp_task(%x: i32, %y: i32, %zaddr: !llvm.ptr<i32>) {
   // CHECK: %[[omp_global_thread_num:.+]] = call i32 @__kmpc_global_thread_num({{.+}})
   // CHECK: %[[task_data:.+]] = call ptr @__kmpc_omp_task_alloc
-  // CHECK-SAME: (ptr @{{.+}}, i32 %[[omp_global_thread_num]], i32 1, i64 0,
+  // CHECK-SAME: (ptr @{{.+}}, i32 %[[omp_global_thread_num]], i32 1, i64 40,
   // CHECK-SAME:  i64 0, ptr @[[wrapper_fn:.+]])
   // CHECK: call i32 @__kmpc_omp_task(ptr @{{.+}}, i32 %[[omp_global_thread_num]], ptr %[[task_data]])
   omp.task {
     %n = llvm.mlir.constant(1 : i64) : i64
     %valaddr = llvm.alloca %n x i32 : (i64) -> !llvm.ptr<i32>
     %val = llvm.load %valaddr : !llvm.ptr<i32>
     %double = llvm.add %val, %val : i32
     llvm.store %double, %valaddr : !llvm.ptr<i32>

 // CHECK:  store i64 %[[dep_arr_addr_0_val_int]], ptr %[[dep_arr_addr_0_val]], align 4
 // CHECK:  %[[dep_arr_addr_0_size:.+]] = getelementptr inbounds %struct.kmp_dep_info, ptr %[[dep_arr_addr_0]], i32 0, i32 1
 // CHECK:  store i64 8, ptr %[[dep_arr_addr_0_size]], align 4
 // CHECK:  %[[dep_arr_addr_0_kind:.+]] = getelementptr inbounds %struct.kmp_dep_info, ptr %[[dep_arr_addr_0]], i32 0, i32 2
 // CHECK: store i8 1, ptr %[[dep_arr_addr_0_kind]], align 1
 llvm.func @omp_task_with_deps(%zaddr: !llvm.ptr<i32>) {
   // CHECK: %[[omp_global_thread_num:.+]] = call i32 @__kmpc_global_thread_num({{.+}})
   // CHECK: %[[task_data:.+]] = call ptr @__kmpc_omp_task_alloc
-  // CHECK-SAME: (ptr @{{.+}}, i32 %[[omp_global_thread_num]], i32 1, i64 0,
+  // CHECK-SAME: (ptr @{{.+}}, i32 %[[omp_global_thread_num]], i32 1, i64 40,
   // CHECK-SAME:  i64 0, ptr @[[wrapper_fn:.+]])
   // CHECK: call i32 @__kmpc_omp_task_with_deps(ptr @{{.+}}, i32 %[[omp_global_thread_num]], ptr %[[task_data]], {{.*}})
   omp.task depend(taskdependin -> %zaddr : !llvm.ptr<i32>) {
     %n = llvm.mlir.constant(1 : i64) : i64
     %valaddr = llvm.alloca %n x i32 : (i64) -> !llvm.ptr<i32>
     %val = llvm.load %valaddr : !llvm.ptr<i32>
     %double = llvm.add %val, %val : i32
     llvm.store %double, %valaddr : !llvm.ptr<i32>

 module attributes {llvm.target_triple = "x86_64-unknown-linux-gnu"} {
   llvm.func @omp_task(%x: i32, %y: i32, %zaddr: !llvm.ptr<i32>) {
     // CHECK: %[[diff:.+]] = sub i32 %[[x]], %[[y]]
     %diff = llvm.sub %x, %y : i32
     // CHECK: store i32 %[[diff]], ptr %2
     llvm.store %diff, %zaddr : !llvm.ptr<i32>
     // CHECK: %[[omp_global_thread_num:.+]] = call i32 @__kmpc_global_thread_num({{.+}})
     // CHECK: %[[task_data:.+]] = call ptr @__kmpc_omp_task_alloc
-    // CHECK-SAME: (ptr @{{.+}}, i32 %[[omp_global_thread_num]], i32 1, i64 16, i64 0,
+    // CHECK-SAME: (ptr @{{.+}}, i32 %[[omp_global_thread_num]], i32 1, i64 40, i64 16,
     // CHECK-SAME: ptr @[[wrapper_fn:.+]])
-    // CHECK: call void @llvm.memcpy.p0.p0.i64(ptr {{.+}} %[[task_data]], ptr {{.+}}, i64 16, i1 false)
+    // CHECK: %[[shareds:.+]] = load ptr, ptr %[[task_data]]
+    // CHECK: call void @llvm.memcpy.p0.p0.i64(ptr {{.+}} %[[shareds]], ptr {{.+}}, i64 16, i1 false)
     // CHECK: call i32 @__kmpc_omp_task(ptr @{{.+}}, i32 %[[omp_global_thread_num]], ptr %[[task_data]])
     omp.task {
       %z = llvm.add %x, %y : i32
       llvm.store %z, %zaddr : !llvm.ptr<i32>
       omp.terminator
     }
     // CHECK: %[[prod:.+]] = mul i32 %[[x]], %[[y]]
     %b = llvm.mul %x, %y : i32

 // CHECK:   %[[sum:.+]] = add i32 %{{.+}}, %{{.+}}
 // CHECK:   store i32 %[[sum]], ptr %{{.+}}
 // CHECK:   br label %[[exit_stub:[^, ]+]]
 // CHECK: [[exit_stub]]:
 // CHECK:   ret void
 
 
 // CHECK: define i32 @[[wrapper_fn]](i32 %{{.+}}, ptr %[[task_data:.+]]) {
-// CHECK:   call void @[[outlined_fn]](ptr %[[task_data]])
+// CHECK:  %[[shareds:.+]] = load ptr, ptr %1, align 8
+// CHECK:   call void @[[outlined_fn]](ptr %[[shareds]])
 // CHECK:   ret i32 0
 // CHECK: }
 
 // -----
 
 llvm.func @par_task_(%arg0: !llvm.ptr<i32> {fir.bindc_name = "a"}) {
   %0 = llvm.mlir.constant(1 : i32) : i32
   omp.task   {

 // CHECK:       [[omp_taskgroup_region1:.+]]:
 // CHECK:         call void @foo()
 // CHECK:         br label %[[omp_region_cont:[^,]+]]
 // CHECK:       [[omp_taskgroup_region]]:
 // CHECK:         %{{.+}} = alloca i8, align 1
 // CHECK:         br label %[[codeRepl:[^,]+]]
 // CHECK:       [[codeRepl]]:
 // CHECK:         %[[omp_global_thread_num_t1:.+]] = call i32 @__kmpc_global_thread_num(ptr @{{.+}})
-// CHECK:         %[[t1_alloc:.+]] = call ptr @__kmpc_omp_task_alloc(ptr @{{.+}}, i32 %[[omp_global_thread_num_t1]], i32 1, i64 0, i64 0, ptr @omp_taskgroup_task..omp_par.wrapper)
+// CHECK:         %[[t1_alloc:.+]] = call ptr @__kmpc_omp_task_alloc(ptr @{{.+}}, i32 %[[omp_global_thread_num_t1]], i32 1, i64 40, i64 0, ptr @omp_taskgroup_task..omp_par.wrapper)
 // CHECK:         %{{.+}} = call i32 @__kmpc_omp_task(ptr @{{.+}}, i32 %[[omp_global_thread_num_t1]], ptr %[[t1_alloc]])
 // CHECK:         br label %[[task_exit:[^,]+]]
 // CHECK:       [[task_exit]]:
 // CHECK:         br label %[[codeRepl9:[^,]+]]
 // CHECK:       [[codeRepl9]]:
 // CHECK:         %[[gep1:.+]] = getelementptr { i32, i32, ptr }, ptr %[[structArg]], i32 0, i32 0
 // CHECK:         store i32 %[[x]], ptr %[[gep1]], align 4
 // CHECK:         %[[gep2:.+]] = getelementptr { i32, i32, ptr }, ptr %[[structArg]], i32 0, i32 1
 // CHECK:         store i32 %[[y]], ptr %[[gep2]], align 4
 // CHECK:         %[[gep3:.+]] = getelementptr { i32, i32, ptr }, ptr %[[structArg]], i32 0, i32 2
 // CHECK:         store ptr %[[zaddr]], ptr %[[gep3]], align 8
 // CHECK:         %[[omp_global_thread_num_t2:.+]] = call i32 @__kmpc_global_thread_num(ptr @{{.+}})
-// CHECK:         %[[t2_alloc:.+]] = call ptr @__kmpc_omp_task_alloc(ptr @{{.+}}, i32 %[[omp_global_thread_num_t2]], i32 1, i64 16, i64 0, ptr @omp_taskgroup_task..omp_par.1.wrapper)
-// CHECK:         call void @llvm.memcpy.p0.p0.i64(ptr align 8 %[[t2_alloc]], ptr align 8 %[[structArg]], i64 16, i1 false)
+// CHECK:         %[[t2_alloc:.+]] = call ptr @__kmpc_omp_task_alloc(ptr @{{.+}}, i32 %[[omp_global_thread_num_t2]], i32 1, i64 40, i64 16, ptr @omp_taskgroup_task..omp_par.1.wrapper)
+// CHECK:         %[[shareds:.+]] = load ptr, ptr %[[t2_alloc]]
+// CHECK:         call void @llvm.memcpy.p0.p0.i64(ptr align 1 %[[shareds]], ptr align 1 %[[structArg]], i64 16, i1 false)
 // CHECK:         %{{.+}} = call i32 @__kmpc_omp_task(ptr @{{.+}}, i32 %[[omp_global_thread_num_t2]], ptr %[[t2_alloc]])
 // CHECK:         br label %[[task_exit3:[^,]+]]
 // CHECK:       [[task_exit3]]:
 // CHECK:         br label %[[omp_taskgroup_region1]]
 // CHECK:       [[omp_region_cont]]:
 // CHECK:         br label %[[taskgroup_exit:[^,]+]]
 // CHECK:       [[taskgroup_exit]]:
 // CHECK:         call void @__kmpc_end_taskgroup(ptr @{{.+}}, i32 %[[omp_global_thread_num]])

 // CHECK:         call void @foo_before()
 // CHECK:         br label %[[entry:[^,]+]]
 // CHECK:       [[entry]]:
 // CHECK:         br label %[[codeRepl:[^,]+]]
 // CHECK:       [[codeRepl]]:                                         ; preds = %entry
 // CHECK:         %[[omp_global_thread_num:.+]] = call i32 @__kmpc_global_thread_num(ptr @{{.+}})
 // CHECK:         %[[final_flag:.+]] = select i1 %[[boolexpr]], i32 2, i32 0
 // CHECK:         %[[task_flags:.+]] = or i32 %[[final_flag]], 1
-// CHECK:         %[[task_data:.+]] = call ptr @__kmpc_omp_task_alloc(ptr @{{.+}}, i32 %[[omp_global_thread_num]], i32 %[[task_flags]], i64 0, i64 0, ptr @omp_task_final..omp_par.wrapper)
+// CHECK:         %[[task_data:.+]] = call ptr @__kmpc_omp_task_alloc(ptr @{{.+}}, i32 %[[omp_global_thread_num]], i32 %[[task_flags]], i64 40, i64 0, ptr @omp_task_final..omp_par.wrapper)
 // CHECK:         %{{.+}} = call i32 @__kmpc_omp_task(ptr @{{.+}}, i32 %[[omp_global_thread_num]], ptr %[[task_data]])
 // CHECK:         br label %[[task_exit:[^,]+]]
 // CHECK:       [[task_exit]]:
 // CHECK:         call void @foo_after()
 // CHECK:         ret void
 
 // -----
 

 // CHECK-LABEL: define void @omp_task_if(
 // CHECK-SAME:    i1 %[[boolexpr:.+]]) {
 // CHECK:         call void @foo_before()
 // CHECK:         br label %[[entry:[^,]+]]
 // CHECK:       [[entry]]:
 // CHECK:         br label %[[codeRepl:[^,]+]]
 // CHECK:       [[codeRepl]]:
 // CHECK:         %[[omp_global_thread_num:.+]] = call i32 @__kmpc_global_thread_num(ptr @{{.+}})
-// CHECK:         %[[task_data:.+]] = call ptr @__kmpc_omp_task_alloc(ptr @{{.+}}, i32 %[[omp_global_thread_num]], i32 1, i64 0, i64 0, ptr @omp_task_if..omp_par.wrapper)
+// CHECK:         %[[task_data:.+]] = call ptr @__kmpc_omp_task_alloc(ptr @{{.+}}, i32 %[[omp_global_thread_num]], i32 1, i64 40, i64 0, ptr @omp_task_if..omp_par.wrapper)
 // CHECK:         br i1 %[[boolexpr]], label %[[true_label:[^,]+]], label %[[false_label:[^,]+]]
 // CHECK:       [[true_label]]:
 // CHECK:         %{{.+}} = call i32 @__kmpc_omp_task(ptr @{{.+}}, i32 %[[omp_global_thread_num]], ptr %[[task_data]])
 // CHECK:         br label %[[if_else_exit:[^,]+]]
 // CHECK:       [[false_label:[^,]+]]:                                                ; preds = %codeRepl
 // CHECK:         call void @__kmpc_omp_task_begin_if0(ptr @{{.+}}, i32 %[[omp_global_thread_num]], ptr %[[task_data]])
 // CHECK:         %{{.+}} = call i32 @omp_task_if..omp_par.wrapper(i32 %[[omp_global_thread_num]])
 // CHECK:         call void @__kmpc_omp_task_complete_if0(ptr @{{.+}}, i32 %[[omp_global_thread_num]], ptr %[[task_data]])