[OpenMPOpt][SplitMemTransfer] Refactored API declarations

[hamax97]

Summary

Refactored __tgt_target_data_begin_mapper_<issue|wait> to receive the handle as an input/output argument as discussed in previous meeting. This given the compiler warning of returning the handle as copy.

[jdoerfert]

LGTM

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

__OMP_RTL(__tgt_target_teams_mapper, false, Int32, Int64, VoidPtr, Int32,
           VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr, Int32, Int32)
 __OMP_RTL(__tgt_target_teams_nowait_mapper, false, Int32, Int64, VoidPtr, Int32,
           VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr, Int32, Int32)
 __OMP_RTL(__tgt_register_requires, false, Void, Int64)
 __OMP_RTL(__tgt_target_data_begin_mapper, false, Void, Int64, Int32, VoidPtrPtr,
           VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr)
 __OMP_RTL(__tgt_target_data_begin_nowait_mapper, false, Void, Int64, Int32,
           VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr)
-__OMP_RTL(__tgt_target_data_begin_mapper_issue, false, AsyncInfo, Int64, Int32,
-          VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr)
-__OMP_RTL(__tgt_target_data_begin_mapper_wait, false, Void, Int64, AsyncInfo)
+__OMP_RTL(__tgt_target_data_begin_mapper_issue, false, Void, Int64, Int32,
+          VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr, AsyncInfoPtr)
+__OMP_RTL(__tgt_target_data_begin_mapper_wait, false, Void, Int64, AsyncInfoPtr)
 __OMP_RTL(__tgt_target_data_end_mapper, false, Void, Int64, Int32, VoidPtrPtr,
           VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr)
 __OMP_RTL(__tgt_target_data_end_nowait_mapper, false, Void, Int64, Int32,
           VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr)
 __OMP_RTL(__tgt_target_data_update_mapper, false, Void, Int64, Int32,
           VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr)
 __OMP_RTL(__tgt_target_data_update_nowait_mapper, false, Void, Int64, Int32,
           VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, VoidPtrPtr)

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

Instruction *canBeMovedDownwards(CallInst &RuntimeCall) {
 
     // Return end of BasicBlock.
     return RuntimeCall.getParent()->getTerminator();
   }
 
   /// Splits \p RuntimeCall into its "issue" and "wait" counterparts.
   bool splitTargetDataBeginRTC(CallInst &RuntimeCall,
                                Instruction &WaitMovementPoint) {
+    // Create stack allocated handle (__tgt_async_info) at the beginning of the
+    // function. Used for storing information of the async transfer, allowing to
+    // wait on it later.
     auto &IRBuilder = OMPInfoCache.OMPBuilder;
+    auto *F = RuntimeCall.getCaller();
+    Instruction *FirstInst = &(F->getEntryBlock().front());
+    AllocaInst *Handle = new AllocaInst(
+        IRBuilder.AsyncInfo, F->getAddressSpace(), "handle", FirstInst);
+
     // Add "issue" runtime call declaration:
     // declare %struct.tgt_async_info @__tgt_target_data_begin_issue(i64, i32,
     //   i8**, i8**, i64*, i64*)
     FunctionCallee IssueDecl = IRBuilder.getOrCreateRuntimeFunction(
         M, OMPRTL___tgt_target_data_begin_mapper_issue);
 
     // Change RuntimeCall call site for its asynchronous version.
     SmallVector<Value *, 8> Args;
     for (auto &Arg : RuntimeCall.args())
       Args.push_back(Arg.get());
+    Args.push_back(Handle);
 
     CallInst *IssueCallsite =
-        CallInst::Create(IssueDecl, Args, "handle", &RuntimeCall);
+        CallInst::Create(IssueDecl, Args, /*NameStr=*/"", &RuntimeCall);
     RuntimeCall.eraseFromParent();
 
     // Add "wait" runtime call declaration:
     // declare void @__tgt_target_data_begin_wait(i64, %struct.__tgt_async_info)
     FunctionCallee WaitDecl = IRBuilder.getOrCreateRuntimeFunction(
         M, OMPRTL___tgt_target_data_begin_mapper_wait);
 
     // Add call site to WaitDecl.
+    const unsigned DeviceIDArgNum = 0;
     Value *WaitParams[2] = {
-        IssueCallsite->getArgOperand(0), // device_id.
-        IssueCallsite // returned handle.
+        IssueCallsite->getArgOperand(DeviceIDArgNum), // device_id.
+        Handle                                        // handle to wait on.
     };
     CallInst::Create(WaitDecl, WaitParams, /*NameStr=*/"", &WaitMovementPoint);
 
     return true;
   }
 
   static Value *combinedIdentStruct(Value *CurrentIdent, Value *NextIdent,
                                     bool GlobalOnly, bool &SingleChoice) {

llvm/test/Transforms/OpenMP/hide_mem_transfer_latency.ll

 ;  #pragma omp target teams
 ;  for (int i = 0; i < 1000; ++i) {
 ;    a *= i*i / 2;
 ;  }
 ;
 ;  return random + a;
 ;}
 define dso_local double @heavyComputation1() {
-; CHECK-LABEL: define {{[^@]+}}@heavyComputation1()
+; CHECK-LABEL: define {{[^@]+}}@heavyComputation1() {
 ; CHECK-NEXT:  entry:
+
+; CHECK-NEXT:    %handle = alloca %struct.__tgt_async_info, align 8
+
 ; CHECK-NEXT:    %a = alloca double, align 8
 ; CHECK-NEXT:    %.offload_baseptrs = alloca [1 x i8*], align 8
 ; CHECK-NEXT:    %.offload_ptrs = alloca [1 x i8*], align 8
 ; CHECK-NEXT:    %.offload_baseptrs4 = alloca [1 x i8*], align 8
 ; CHECK-NEXT:    %.offload_ptrs5 = alloca [1 x i8*], align 8
 ; CHECK-NEXT:    %0 = bitcast double* %a to i8*
 ; CHECK-NEXT:    %call = tail call i32 (...) @rand()
 ; CHECK-NEXT:    %rem = srem i32 %call, 777
 ; CHECK-NEXT:    %conv = sitofp i32 %rem to double
 ; CHECK-NEXT:    store double %conv, double* %a, align 8
 ; CHECK-NEXT:    %call1 = tail call i32 (...) @rand()
 ; CHECK-NEXT:    %1 = getelementptr inbounds [1 x i8*], [1 x i8*]* %.offload_baseptrs, i64 0, i64 0
 ; CHECK-NEXT:    %2 = bitcast [1 x i8*]* %.offload_baseptrs to double**
 ; CHECK-NEXT:    store double* %a, double** %2, align 8
 ; CHECK-NEXT:    %3 = getelementptr inbounds [1 x i8*], [1 x i8*]* %.offload_ptrs, i64 0, i64 0
 ; CHECK-NEXT:    %4 = bitcast [1 x i8*]* %.offload_ptrs to double**
 ; CHECK-NEXT:    store double* %a, double** %4, align 8
 
-; CHECK-NEXT:    %handle = call %struct.__tgt_async_info @__tgt_target_data_begin_mapper_issue(i64 -1, i32 1, i8** %1, i8** %3, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @.offload_sizes.1, i64 0, i64 0), i64* getelementptr inbounds ([1 x i64], [1 x i64]* @.offload_maptypes, i64 0, i64 0), i8** null)
+; CHECK-NEXT:    call void @__tgt_target_data_begin_mapper_issue(i64 -1, i32 1, i8** %1, i8** %3, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @.offload_sizes.1, i64 0, i64 0), i64* getelementptr inbounds ([1 x i64], [1 x i64]* @.offload_maptypes, i64 0, i64 0), i8** null, %struct.__tgt_async_info* %handle)
 
 ; CHECK-NEXT:    %5 = bitcast double* %a to i64*
 
-; CHECK-NEXT:    call void @__tgt_target_data_begin_mapper_wait(i64 -1, %struct.__tgt_async_info %handle)
+; CHECK-NEXT:    call void @__tgt_target_data_begin_mapper_wait(i64 -1, %struct.__tgt_async_info* %handle)
 
 ; CHECK-NEXT:    %6 = load i64, i64* %5, align 8
 ; CHECK-NEXT:    %7 = getelementptr inbounds [1 x i8*], [1 x i8*]* %.offload_baseptrs4, i64 0, i64 0
 ; CHECK-NEXT:    %8 = bitcast [1 x i8*]* %.offload_baseptrs4 to i64*
 ; CHECK-NEXT:    store i64 %6, i64* %8, align 8
 ; CHECK-NEXT:    %9 = getelementptr inbounds [1 x i8*], [1 x i8*]* %.offload_ptrs5, i64 0, i64 0
 ; CHECK-NEXT:    %10 = bitcast [1 x i8*]* %.offload_ptrs5 to i64*
 ; CHECK-NEXT:    store i64 %6, i64* %10, align 8

 ;  #pragma omp target teams
 ;  for (int i = 0; i < size; ++i) {
 ;    a[i] = ++a[i] * 3.141624;
 ;  }
 ;
 ;  return random;
 ;}
 define dso_local i32 @heavyComputation2(double* %a, i32 %size) {
-; CHECK-LABEL: define {{[^@]+}}@heavyComputation2(double* %a, i32 %size)
+; CHECK-LABEL: define {{[^@]+}}@heavyComputation2(double* %a, i32 %size) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    %size.addr = alloca i32, align 4
 ; CHECK-NEXT:    %.offload_baseptrs = alloca [2 x i8*], align 8
 ; CHECK-NEXT:    %.offload_ptrs = alloca [2 x i8*], align 8
 ; CHECK-NEXT:    %.offload_sizes = alloca [2 x i64], align 8
 ; CHECK-NEXT:    %.offload_baseptrs2 = alloca [2 x i8*], align 8
 ; CHECK-NEXT:    %.offload_ptrs3 = alloca [2 x i8*], align 8
 ; CHECK-NEXT:    store i32 %size, i32* %size.addr, align 4

 ;  #pragma omp target teams
 ;  for (int i = 0; i < size; ++i) {
 ;    a[i] = ++a[i] * 3.141624;
 ;  }
 ;
 ;  return random;
 ;}
 define dso_local i32 @heavyComputation3(double* noalias %a, i32 %size) {
-; CHECK-LABEL: define {{[^@]+}}@heavyComputation3(double* noalias %a, i32 %size)
+; CHECK-LABEL: define {{[^@]+}}@heavyComputation3(double* noalias %a, i32 %size) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    %size.addr = alloca i32, align 4
 ; CHECK-NEXT:    %.offload_baseptrs = alloca [2 x i8*], align 8
 ; CHECK-NEXT:    %.offload_ptrs = alloca [2 x i8*], align 8
 ; CHECK-NEXT:    %.offload_sizes = alloca [2 x i64], align 8
 ; CHECK-NEXT:    %.offload_baseptrs2 = alloca [2 x i8*], align 8
 ; CHECK-NEXT:    %.offload_ptrs3 = alloca [2 x i8*], align 8
 ; CHECK-NEXT:    store i32 %size, i32* %size.addr, align 4

 ;  args[0] = &a;
 ;  __tgt_target_data_begin(..., args, ...)
 ;
 ;  // Random computation.
 ;  random %= size;
 ;  return random;
 ;}
 define dso_local i32 @dataTransferOnly1(double* noalias %a, i32 %size) {
-; CHECK-LABEL: define {{[^@]+}}@dataTransferOnly1(double* noalias %a, i32 %size)
+; CHECK-LABEL: define {{[^@]+}}@dataTransferOnly1(double* noalias %a, i32 %size) {
 ; CHECK-NEXT:  entry:
+
+; CHECK-NEXT:    %handle = alloca %struct.__tgt_async_info, align 8
+
 ; CHECK-NEXT:    %.offload_baseptrs = alloca [1 x i8*], align 8
 ; CHECK-NEXT:    %.offload_ptrs = alloca [1 x i8*], align 8
 ; CHECK-NEXT:    %.offload_sizes = alloca [1 x i64], align 8
 ; CHECK-NEXT:    %call = tail call i32 (...) @rand()
 ; CHECK-NEXT:    %conv = zext i32 %size to i64
 ; CHECK-NEXT:    %0 = shl nuw nsw i64 %conv, 3
 ; CHECK-NEXT:    %1 = getelementptr inbounds [1 x i8*], [1 x i8*]* %.offload_baseptrs, i64 0, i64 0
 ; CHECK-NEXT:    %2 = bitcast [1 x i8*]* %.offload_baseptrs to double**
 ; CHECK-NEXT:    store double* %a, double** %2, align 8
 ; CHECK-NEXT:    %3 = getelementptr inbounds [1 x i8*], [1 x i8*]* %.offload_ptrs, i64 0, i64 0
 ; CHECK-NEXT:    %4 = bitcast [1 x i8*]* %.offload_ptrs to double**
 ; CHECK-NEXT:    store double* %a, double** %4, align 8
 ; CHECK-NEXT:    %5 = getelementptr inbounds [1 x i64], [1 x i64]* %.offload_sizes, i64 0, i64 0
 ; CHECK-NEXT:    store i64 %0, i64* %5, align 8
 
-; CHECK-NEXT:    %handle = call %struct.__tgt_async_info @__tgt_target_data_begin_mapper_issue(i64 -1, i32 1, i8** %1, i8** %3, i64* %5, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @.offload_maptypes.5, i64 0, i64 0), i8** null)
+; CHECK-NEXT:    call void @__tgt_target_data_begin_mapper_issue(i64 -1, i32 1, i8** %1, i8** %3, i64* %5, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @.offload_maptypes.5, i64 0, i64 0), i8** null, %struct.__tgt_async_info* %handle)
 
 ; CHECK-NEXT:    %rem = urem i32 %call, %size
 
-; CHECK-NEXT:    call void @__tgt_target_data_begin_mapper_wait(i64 -1, %struct.__tgt_async_info %handle)
+; CHECK-NEXT:    call void @__tgt_target_data_begin_mapper_wait(i64 -1, %struct.__tgt_async_info* %handle)
 
 ; CHECK-NEXT:    call void @__tgt_target_data_end_mapper(i64 -1, i32 1, i8** nonnull %1, i8** nonnull %3, i64* nonnull %5, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @.offload_maptypes.5, i64 0, i64 0), i8** null)
 ; CHECK-NEXT:    ret i32 %rem
 ;
 entry:
   %.offload_baseptrs = alloca [1 x i8*], align 8
   %.offload_ptrs = alloca [1 x i8*], align 8
   %.offload_sizes = alloca [1 x i64], align 8

 }
 
 declare void @__tgt_target_data_begin_mapper(i64, i32, i8**, i8**, i64*, i64*, i8**)
 declare i32 @__tgt_target_teams_mapper(i64, i8*, i32, i8**, i8**, i64*, i64*, i8**, i32, i32)
 declare void @__tgt_target_data_end_mapper(i64, i32, i8**, i8**, i64*, i64*, i8**)
 
 declare dso_local i32 @rand(...)
 
-; CHECK: declare %struct.__tgt_async_info @__tgt_target_data_begin_mapper_issue(i64, i32, i8**, i8**, i64*, i64*, i8**)
-; CHECK: declare void @__tgt_target_data_begin_mapper_wait(i64, %struct.__tgt_async_info)
+; CHECK: declare void @__tgt_target_data_begin_mapper_issue(i64, i32, i8**, i8**, i64*, i64*, i8**, %struct.__tgt_async_info*)
+; CHECK: declare void @__tgt_target_data_begin_mapper_wait(i64, %struct.__tgt_async_info*)