[OpenMP][Clang] Pass global thread ID to outlined function

[gtbercea]

Summary

The data sharing wrapper function needs to pass a valid global thread ID to the parallel outlined function when the parallel is combined with a for directive.

[ABataev]

Tests?

[gtbercea]

After some internal discussion with @ABataev he is going to replace the manual computation of the thread ID with a call to the runtime in a new patch.

lib/CodeGen/CGOpenMPRuntimeNVPTX.cpp

public:
   /// Checks if the escaped local variable is actually a parameter passed by
   /// value.
   const llvm::SmallPtrSetImpl<const Decl *> &getEscapedParameters() const {
     return EscapedParameters;
   }
 };
 } // anonymous namespace
 
+///
+/// NVPTX API calls.
+///
+
 /// Get the GPU warp size.
 static llvm::Value *getNVPTXWarpSize(CodeGenFunction &CGF) {
   return CGF.EmitRuntimeCall(
       llvm::Intrinsic::getDeclaration(
           &CGF.CGM.getModule(), llvm::Intrinsic::nvvm_read_ptx_sreg_warpsize),
       "nvptx_warp_size");
 }
 
+/// Get the id of the current block on the GPU.
+static llvm::Value *getNVPTXBlockID(CodeGenFunction &CGF) {
+  return CGF.EmitRuntimeCall(
+      llvm::Intrinsic::getDeclaration(
+          &CGF.CGM.getModule(), llvm::Intrinsic::nvvm_read_ptx_sreg_ctaid_x),
+      "nvptx_block_id");
+}
+
 /// Get the id of the current thread on the GPU.
 static llvm::Value *getNVPTXThreadID(CodeGenFunction &CGF) {
   return CGF.EmitRuntimeCall(
       llvm::Intrinsic::getDeclaration(
           &CGF.CGM.getModule(), llvm::Intrinsic::nvvm_read_ptx_sreg_tid_x),
       "nvptx_tid");
 }
 
static llvm::Value *getMasterThreadID(CodeGenFunction &CGF) {
 
   // We assume that the warp size is a power of 2.
   llvm::Value *Mask = Bld.CreateSub(getNVPTXWarpSize(CGF), Bld.getInt32(1));
 
   return Bld.CreateAnd(Bld.CreateSub(NumThreads, Bld.getInt32(1)),
                        Bld.CreateNot(Mask), "master_tid");
 }
 
+/// Get number of OMP workers for parallel region after subtracting
+/// the master warp.
+static llvm::Value *getNumWorkers(CodeGenFunction &CGF) {
+  CGBuilderTy &Bld = CGF.Builder;
+  return Bld.CreateNUWSub(getNVPTXNumThreads(CGF), Bld.getInt32(32),
+                          "num_workers");
+}
+
+/// Get thread id in team.
+/// FIXME: Remove the expensive remainder operation.
+static llvm::Value *getTeamThreadId(CodeGenFunction &CGF) {
+  CGBuilderTy &Bld = CGF.Builder;
+  // N % M = N & (M-1) it M is a power of 2. The master Id is expected to be a
+  // power fo two in all cases.
+  auto *Mask = Bld.CreateNUWSub(getMasterThreadID(CGF), Bld.getInt32(1));
+  return Bld.CreateAnd(getNVPTXThreadID(CGF), Mask, "team_tid");
+}
+
+/// Get global thread id.
+static llvm::Value *getGlobalThreadId(CodeGenFunction &CGF) {
+  assert(CGF.CurFn && "No function in current CodeGenFunction.");
+  CGBuilderTy &Bld = CGF.Builder;
+  return Bld.CreateAdd(Bld.CreateMul(getNVPTXBlockID(CGF), getNumWorkers(CGF)),
+                       getTeamThreadId(CGF), "global_tid");
+}
+
 CGOpenMPRuntimeNVPTX::WorkerFunctionState::WorkerFunctionState(
     CodeGenModule &CGM, SourceLocation Loc)
     : WorkerFn(nullptr), CGFI(nullptr), Loc(Loc) {
   createWorkerFunction(CGM);
 }
 
 void CGOpenMPRuntimeNVPTX::WorkerFunctionState::createWorkerFunction(
     CodeGenModule &CGM) {
CGF.StartFunction(GlobalDecl(), Ctx.VoidTy, Fn, CGFI, WrapperArgs,
                     D.getLocStart(), D.getLocStart());
 
   const auto *RD = CS.getCapturedRecordDecl();
   auto CurField = RD->field_begin();
 
   // Get the array of arguments.
   SmallVector<llvm::Value *, 8> Args;
 
-  // TODO: suppport SIMD and pass actual values
-  Args.emplace_back(
-      llvm::ConstantPointerNull::get(CGM.Int32Ty->getPointerTo()));
+  // First argument is the global thread ID.
+  Address GlobalThreadIDAddr =
+      CGF.CreateDefaultAlignTempAlloca(CGF.Int32Ty, "global_tid");
+  CGF.EmitStoreOfScalar(getGlobalThreadId(CGF), GlobalThreadIDAddr,
+                        /*Volatile=*/false,
+                        Ctx.getPointerType(Ctx.VoidPtrTy));
+  Args.emplace_back(GlobalThreadIDAddr.getPointer());
+
+  // TODO: suppport SIMD and pass actual value
   Args.emplace_back(
       llvm::ConstantPointerNull::get(CGM.Int32Ty->getPointerTo()));
 
   CGBuilderTy &Bld = CGF.Builder;
   auto CI = CS.capture_begin();
 
   // Use global memory for data sharing.
   // Handle passing of global args to workers.