[OpenMP] Further adjustments of nvptx runtime functions

[Hahnfeld]

Summary

Pass in default value of 1, similar to previous commit rL318836.

[ABataev]

LG

lib/CodeGen/CGOpenMPRuntimeNVPTX.cpp

enum OpenMPRTLFunctionNVPTX {
   /// \brief Call to void __kmpc_kernel_deinit(int16_t IsOMPRuntimeInitialized);
   OMPRTL_NVPTX__kmpc_kernel_deinit,
   /// \brief Call to void __kmpc_spmd_kernel_init(kmp_int32 thread_limit,
   /// int16_t RequiresOMPRuntime, int16_t RequiresDataSharing);
   OMPRTL_NVPTX__kmpc_spmd_kernel_init,
   /// \brief Call to void __kmpc_spmd_kernel_deinit();
   OMPRTL_NVPTX__kmpc_spmd_kernel_deinit,
   /// \brief Call to void __kmpc_kernel_prepare_parallel(void
-  /// *outlined_function, void ***args, kmp_int32 nArgs);
+  /// *outlined_function, void ***args, kmp_int32 nArgs, int16_t
+  /// IsOMPRuntimeInitialized);
   OMPRTL_NVPTX__kmpc_kernel_prepare_parallel,
   /// \brief Call to bool __kmpc_kernel_parallel(void **outlined_function, void
-  /// ***args);
+  /// ***args, int16_t IsOMPRuntimeInitialized);
   OMPRTL_NVPTX__kmpc_kernel_parallel,
   /// \brief Call to void __kmpc_kernel_end_parallel();
   OMPRTL_NVPTX__kmpc_kernel_end_parallel,
   /// Call to void __kmpc_serialized_parallel(ident_t *loc, kmp_int32
   /// global_tid);
   OMPRTL_NVPTX__kmpc_serialized_parallel,
   /// Call to void __kmpc_end_serialized_parallel(ident_t *loc, kmp_int32
   /// global_tid);
void CGOpenMPRuntimeNVPTX::emitWorkerLoop(CodeGenFunction &CGF,
   Address ExecStatus =
       CGF.CreateDefaultAlignTempAlloca(CGF.Int8Ty, /*Name=*/"exec_status");
   CGF.InitTempAlloca(ExecStatus, Bld.getInt8(/*C=*/0));
   CGF.InitTempAlloca(WorkFn, llvm::Constant::getNullValue(CGF.Int8PtrTy));
 
   // Set up shared arguments
   Address SharedArgs =
       CGF.CreateDefaultAlignTempAlloca(CGF.Int8PtrPtrTy, "shared_args");
-  llvm::Value *Args[] = {WorkFn.getPointer(), SharedArgs.getPointer()};
+  // TODO: Optimize runtime initialization and pass in correct value.
+  llvm::Value *Args[] = {WorkFn.getPointer(), SharedArgs.getPointer(),
+                         /*RequiresOMPRuntime=*/Bld.getInt16(1)};
   llvm::Value *Ret = CGF.EmitRuntimeCall(
       createNVPTXRuntimeFunction(OMPRTL_NVPTX__kmpc_kernel_parallel), Args);
   Bld.CreateStore(Bld.CreateZExt(Ret, CGF.Int8Ty), ExecStatus);
 
   // On termination condition (workid == 0), exit loop.
   llvm::Value *ShouldTerminate =
       Bld.CreateIsNull(Bld.CreateLoad(WorkFn), "should_terminate");
   Bld.CreateCondBr(ShouldTerminate, ExitBB, SelectWorkersBB);
case OMPRTL_NVPTX__kmpc_spmd_kernel_deinit: {
     // Build void __kmpc_spmd_kernel_deinit();
     llvm::FunctionType *FnTy =
         llvm::FunctionType::get(CGM.VoidTy, llvm::None, /*isVarArg*/ false);
     RTLFn = CGM.CreateRuntimeFunction(FnTy, "__kmpc_spmd_kernel_deinit");
     break;
   }
   case OMPRTL_NVPTX__kmpc_kernel_prepare_parallel: {
     /// Build void __kmpc_kernel_prepare_parallel(
-    /// void *outlined_function, void ***args, kmp_int32 nArgs);
+    /// void *outlined_function, void ***args, kmp_int32 nArgs, int16_t
+    /// IsOMPRuntimeInitialized);
     llvm::Type *TypeParams[] = {CGM.Int8PtrTy,
-        CGM.Int8PtrPtrTy->getPointerTo(0), CGM.Int32Ty};
+                                CGM.Int8PtrPtrTy->getPointerTo(0), CGM.Int32Ty,
+                                CGM.Int16Ty};
     llvm::FunctionType *FnTy =
         llvm::FunctionType::get(CGM.VoidTy, TypeParams, /*isVarArg*/ false);
     RTLFn = CGM.CreateRuntimeFunction(FnTy, "__kmpc_kernel_prepare_parallel");
     break;
   }
   case OMPRTL_NVPTX__kmpc_kernel_parallel: {
-    /// Build bool __kmpc_kernel_parallel(void **outlined_function, void ***args);
+    /// Build bool __kmpc_kernel_parallel(void **outlined_function, void
+    /// ***args, int16_t IsOMPRuntimeInitialized);
     llvm::Type *TypeParams[] = {CGM.Int8PtrPtrTy,
-        CGM.Int8PtrPtrTy->getPointerTo(0)};
+                                CGM.Int8PtrPtrTy->getPointerTo(0), CGM.Int16Ty};
     llvm::Type *RetTy = CGM.getTypes().ConvertType(CGM.getContext().BoolTy);
     llvm::FunctionType *FnTy =
         llvm::FunctionType::get(RetTy, TypeParams, /*isVarArg*/ false);
     RTLFn = CGM.CreateRuntimeFunction(FnTy, "__kmpc_kernel_parallel");
     break;
   }
   case OMPRTL_NVPTX__kmpc_kernel_end_parallel: {
     /// Build void __kmpc_kernel_end_parallel();
auto &&L0ParallelGen = [this, WFn, &CapturedVars](CodeGenFunction &CGF,
     llvm::Value *ID = Bld.CreateBitOrPointerCast(WFn, CGM.Int8PtrTy);
 
     if (!CapturedVars.empty()) {
       // Prepare for parallel region. Indicate the outlined function.
       Address SharedArgs =
           CGF.CreateDefaultAlignTempAlloca(CGF.VoidPtrPtrTy,
               "shared_args");
       llvm::Value *SharedArgsPtr = SharedArgs.getPointer();
+      // TODO: Optimize runtime initialization and pass in correct value.
       llvm::Value *Args[] = {ID, SharedArgsPtr,
-                             Bld.getInt32(CapturedVars.size())};
+                             Bld.getInt32(CapturedVars.size()),
+                             /*RequiresOMPRuntime=*/Bld.getInt16(1)};
 
       CGF.EmitRuntimeCall(
           createNVPTXRuntimeFunction(OMPRTL_NVPTX__kmpc_kernel_prepare_parallel),
           Args);
 
       unsigned Idx = 0;
       ASTContext &Ctx = CGF.getContext();
       for (llvm::Value *V : CapturedVars) {
         Address Dst = Bld.CreateConstInBoundsGEP(
             CGF.EmitLoadOfPointer(SharedArgs,
             Ctx.getPointerType(
                 Ctx.getPointerType(Ctx.VoidPtrTy)).castAs<PointerType>()),
             Idx, CGF.getPointerSize());
         llvm::Value *PtrV = Bld.CreateBitCast(V, CGF.VoidPtrTy);
         CGF.EmitStoreOfScalar(PtrV, Dst, /*Volatile=*/false,
             Ctx.getPointerType(Ctx.VoidPtrTy));
         Idx++;
       }
     } else {
-      llvm::Value *Args[] = {ID,
-          llvm::ConstantPointerNull::get(CGF.VoidPtrPtrTy->getPointerTo(0)),
-          /*nArgs=*/Bld.getInt32(0)};
+      // TODO: Optimize runtime initialization and pass in correct value.
+      llvm::Value *Args[] = {
+          ID, llvm::ConstantPointerNull::get(CGF.VoidPtrPtrTy->getPointerTo(0)),
+          /*nArgs=*/Bld.getInt32(0), /*RequiresOMPRuntime=*/Bld.getInt16(1)};
       CGF.EmitRuntimeCall(
           createNVPTXRuntimeFunction(OMPRTL_NVPTX__kmpc_kernel_prepare_parallel),
           Args);
     }
 
     // Activate workers. This barrier is used by the master to signal
     // work for the workers.
     syncCTAThreads(CGF);

test/OpenMP/nvptx_data_sharing.cpp

#pragma omp target
     }
   }
 }
 
 /// ========= In the worker function ========= ///
 
 // CK1: define internal void @__omp_offloading_{{.*}}test_ds{{.*}}worker(){{.*}}{
 // CK1: [[SHAREDARGS:%.+]] = alloca i8**
-// CK1: call i1 @__kmpc_kernel_parallel(i8** %work_fn, i8*** [[SHAREDARGS]])
+// CK1: call i1 @__kmpc_kernel_parallel(i8** %work_fn, i8*** [[SHAREDARGS]], i16 1)
 // CK1: [[SHARGSTMP:%.+]] = load i8**, i8*** [[SHAREDARGS]]
 // CK1: call void @__omp_outlined___wrapper{{.*}}({{.*}}, i8** [[SHARGSTMP]])
 
 /// ========= In the kernel function ========= ///
 
 // CK1: {{.*}}define void @__omp_offloading{{.*}}test_ds{{.*}}()
 // CK1: [[SHAREDARGS1:%.+]] = alloca i8**
-// CK1: call void @__kmpc_kernel_prepare_parallel({{.*}}, i8*** [[SHAREDARGS1]], i32 1)
+// CK1: call void @__kmpc_kernel_prepare_parallel({{.*}}, i8*** [[SHAREDARGS1]], i32 1, i16 1)
 // CK1: [[SHARGSTMP1:%.+]] = load i8**, i8*** [[SHAREDARGS1]]
 // CK1: [[SHARGSTMP2:%.+]] = getelementptr inbounds i8*, i8** [[SHARGSTMP1]]
 // CK1: [[SHAREDVAR:%.+]] = bitcast i32* {{.*}} to i8*
 // CK1: store i8* [[SHAREDVAR]], i8** [[SHARGSTMP2]]
 
 /// ========= In the data sharing wrapper function ========= ///
 
 // CK1: {{.*}}define internal void @__omp_outlined___wrapper({{.*}}i8**){{.*}}{

test/OpenMP/nvptx_target_teams_codegen.cpp

int bar(int n){
   // CHECK-DAG: [[OMP_EXEC_STATUS:%.+]] = alloca i8,
   // CHECK-DAG: [[OMP_WORK_FN:%.+]] = alloca i8*,
   // CHECK: store i8* null, i8** [[OMP_WORK_FN]],
   // CHECK: store i8 0, i8* [[OMP_EXEC_STATUS]],
   // CHECK: br label {{%?}}[[AWAIT_WORK:.+]]
   //
   // CHECK: [[AWAIT_WORK]]
   // CHECK: call void @llvm.nvvm.barrier0()
-  // CHECK: [[KPR:%.+]] = call i1 @__kmpc_kernel_parallel(i8** [[OMP_WORK_FN]], i8*** %shared_args)
+  // CHECK: [[KPR:%.+]] = call i1 @__kmpc_kernel_parallel(i8** [[OMP_WORK_FN]], i8*** %shared_args, i16 1)
   // CHECK: [[KPRB:%.+]] = zext i1 [[KPR]] to i8
   // store i8 [[KPRB]], i8* [[OMP_EXEC_STATUS]], align 1
   // CHECK: [[WORK:%.+]] = load i8*, i8** [[OMP_WORK_FN]],
   // CHECK: [[SHOULD_EXIT:%.+]] = icmp eq i8* [[WORK]], null
   // CHECK: br i1 [[SHOULD_EXIT]], label {{%?}}[[EXIT:.+]], label {{%?}}[[SEL_WORKERS:.+]]
   //
   // CHECK: [[SEL_WORKERS]]
   // CHECK: [[ST:%.+]] = load i8, i8* [[OMP_EXEC_STATUS]]
int bar(int n){
   // CHECK-DAG: [[OMP_EXEC_STATUS:%.+]] = alloca i8,
   // CHECK-DAG: [[OMP_WORK_FN:%.+]] = alloca i8*,
   // CHECK: store i8* null, i8** [[OMP_WORK_FN]],
   // CHECK: store i8 0, i8* [[OMP_EXEC_STATUS]],
   // CHECK: br label {{%?}}[[AWAIT_WORK:.+]]
   //
   // CHECK: [[AWAIT_WORK]]
   // CHECK: call void @llvm.nvvm.barrier0()
-  // CHECK: [[KPR:%.+]] = call i1 @__kmpc_kernel_parallel(i8** [[OMP_WORK_FN]], i8*** %shared_args)
+  // CHECK: [[KPR:%.+]] = call i1 @__kmpc_kernel_parallel(i8** [[OMP_WORK_FN]], i8*** %shared_args, i16 1)
   // CHECK: [[KPRB:%.+]] = zext i1 [[KPR]] to i8
   // store i8 [[KPRB]], i8* [[OMP_EXEC_STATUS]], align 1
   // CHECK: [[WORK:%.+]] = load i8*, i8** [[OMP_WORK_FN]],
   // CHECK: [[SHOULD_EXIT:%.+]] = icmp eq i8* [[WORK]], null
   // CHECK: br i1 [[SHOULD_EXIT]], label {{%?}}[[EXIT:.+]], label {{%?}}[[SEL_WORKERS:.+]]
   //
   // CHECK: [[SEL_WORKERS]]
   // CHECK: [[ST:%.+]] = load i8, i8* [[OMP_EXEC_STATUS]]