[OPENMP][NVPTX]Fix critical region codegen.

[ABataev]

Summary

Previously critical regions were emitted with the barrier making it a
worksharing construct though it is not. Also, it leads to incorrect
behavior in Cuda9+. Patch fixes this problem.

[jdoerfert]

I guess IR test should be affected already and it would be good to have the run time test that breaks with barriers.

[ABataev]

Fix the test

[ABataev]

In D66673#1643544, @jdoerfert wrote:

I guess IR test should be affected already and it would be good to have the run time test that breaks with barriers.

Runtime test is libomptarget/deviceRTLs/nvptx/test/parallel/spmd_parallel_regions.cpp

[jdoerfert]

LGTM

lib/CodeGen/CGOpenMPRuntimeNVPTX.cpp

enum OpenMPRTLFunctionNVPTX {
   /// Call to void __kmpc_restore_team_static_memory(int16_t
   /// isSPMDExecutionMode, int16_t is_shared);
   OMPRTL_NVPTX__kmpc_restore_team_static_memory,
   /// Call to void __kmpc_barrier(ident_t *loc, kmp_int32 global_tid);
   OMPRTL__kmpc_barrier,
   /// Call to void __kmpc_barrier_simple_spmd(ident_t *loc, kmp_int32
   /// global_tid);
   OMPRTL__kmpc_barrier_simple_spmd,
+  /// Call to int32_t __kmpc_warp_active_thread_mask(void);
+  OMPRTL_NVPTX__kmpc_warp_active_thread_mask,
+  /// Call to void __kmpc_syncwarp(int32_t Mask);
+  OMPRTL_NVPTX__kmpc_syncwarp,
 };
 
 /// Pre(post)-action for different OpenMP constructs specialized for NVPTX.
 class NVPTXActionTy final : public PrePostActionTy {
   llvm::FunctionCallee EnterCallee = nullptr;
   ArrayRef<llvm::Value *> EnterArgs;
   llvm::FunctionCallee ExitCallee = nullptr;
   ArrayRef<llvm::Value *> ExitArgs;
case OMPRTL__kmpc_barrier_simple_spmd: {
     auto *FnTy =
         llvm::FunctionType::get(CGM.VoidTy, TypeParams, /*isVarArg*/ false);
     RTLFn =
         CGM.CreateRuntimeFunction(FnTy, /*Name*/ "__kmpc_barrier_simple_spmd");
     cast<llvm::Function>(RTLFn.getCallee())
         ->addFnAttr(llvm::Attribute::Convergent);
     break;
   }
+  case OMPRTL_NVPTX__kmpc_warp_active_thread_mask: {
+    // Build int32_t __kmpc_warp_active_thread_mask(void);
+    auto *FnTy =
+        llvm::FunctionType::get(CGM.Int32Ty, llvm::None, /*isVarArg=*/false);
+    RTLFn = CGM.CreateRuntimeFunction(FnTy, "__kmpc_warp_active_thread_mask");
+    break;
+  }
+  case OMPRTL_NVPTX__kmpc_syncwarp: {
+    // Build void __kmpc_syncwarp(kmp_int32 Mask);
+    auto *FnTy =
+        llvm::FunctionType::get(CGM.VoidTy, CGM.Int32Ty, /*isVarArg=*/false);
+    RTLFn = CGM.CreateRuntimeFunction(FnTy, "__kmpc_syncwarp");
+    break;
+  }
   }
   return RTLFn;
 }
 
 void CGOpenMPRuntimeNVPTX::createOffloadEntry(llvm::Constant *ID,
                                               llvm::Constant *Addr,
                                               uint64_t Size, int32_t,
                                               llvm::GlobalValue::LinkageTypes) {
void CGOpenMPRuntimeNVPTX::emitCriticalRegion(
     const RegionCodeGenTy &CriticalOpGen, SourceLocation Loc,
     const Expr *Hint) {
   llvm::BasicBlock *LoopBB = CGF.createBasicBlock("omp.critical.loop");
   llvm::BasicBlock *TestBB = CGF.createBasicBlock("omp.critical.test");
   llvm::BasicBlock *SyncBB = CGF.createBasicBlock("omp.critical.sync");
   llvm::BasicBlock *BodyBB = CGF.createBasicBlock("omp.critical.body");
   llvm::BasicBlock *ExitBB = CGF.createBasicBlock("omp.critical.exit");
 
+  // Get the mask of active threads in the warp.
+  llvm::Value *Mask = CGF.EmitRuntimeCall(
+      createNVPTXRuntimeFunction(OMPRTL_NVPTX__kmpc_warp_active_thread_mask));
   // Fetch team-local id of the thread.
   llvm::Value *ThreadID = getNVPTXThreadID(CGF);
 
   // Get the width of the team.
   llvm::Value *TeamWidth = getNVPTXNumThreads(CGF);
 
   // Initialize the counter variable for the loop.
   QualType Int32Ty =
void CGOpenMPRuntimeNVPTX::emitCriticalRegion(
   CGOpenMPRuntime::emitCriticalRegion(CGF, CriticalName, CriticalOpGen, Loc,
                                       Hint);
 
   // After the body surrounded by the critical region, the single executing
   // thread will jump to the synchronisation point.
   // Block waits for all threads in current team to finish then increments the
   // counter variable and returns to the loop.
   CGF.EmitBlock(SyncBB);
-  emitBarrierCall(CGF, Loc, OMPD_unknown, /*EmitChecks=*/false,
-                  /*ForceSimpleCall=*/true);
+  // Reconverge active threads in the warp.
+  (void)CGF.EmitRuntimeCall(
+      createNVPTXRuntimeFunction(OMPRTL_NVPTX__kmpc_syncwarp), Mask);
 
   llvm::Value *IncCounterVal =
       CGF.Builder.CreateNSWAdd(CounterVal, CGF.Builder.getInt32(1));
   CGF.EmitStoreOfScalar(IncCounterVal, CounterLVal);
   CGF.EmitBranch(LoopBB);
 
   // Block that is reached when  all threads in the team complete the region.
   CGF.EmitBlock(ExitBB, /*IsFinished=*/true);

test/OpenMP/nvptx_parallel_codegen.cpp

 // CHECK-64: [[A:%.+]] = load i32, i32* [[CONV]],
 // CHECK: [[GLOBAL_A_ADDR:%.+]] = getelementptr inbounds %struct._globalized_locals_ty, %struct._globalized_locals_ty* [[BC]], i{{[0-9]+}} 0, i{{[0-9]+}} 0
 // CHECK: store i32 [[A]], i32* [[GLOBAL_A_ADDR]],
 // CHECK: [[IS_SHARED:%.+]] = load i16, i16* [[KERNEL_SHARED]],
 // CHECK: call void @__kmpc_restore_team_static_memory(i16 0, i16 [[IS_SHARED]])
 
 // CHECK-LABEL: define internal void @{{.+}}(i32* noalias %{{.+}}, i32* noalias %{{.+}}, i32* dereferenceable{{.*}})
 // CHECK:  [[CC:%.+]] = alloca i32,
+// CHECK:  [[MASK:%.+]] = call i32 @__kmpc_warp_active_thread_mask()
 // CHECK:  [[TID:%.+]] = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
 // CHECK:  [[NUM_THREADS:%.+]] = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x()
 // CHECK:  store i32 0, i32* [[CC]],
 // CHECK:  br label
 
 // CHECK:  [[CC_VAL:%.+]] = load i32, i32* [[CC]],
 // CHECK:  [[RES:%.+]] = icmp slt i32 [[CC_VAL]], [[NUM_THREADS]]
 // CHECK:  br i1 [[RES]], label
 
 // CHECK:  [[CC_VAL:%.+]] = load i32, i32* [[CC]],
 // CHECK:  [[RES:%.+]] = icmp eq i32 [[TID]], [[CC_VAL]]
 // CHECK:  br i1 [[RES]], label
 
 // CHECK:  call void @__kmpc_critical(
 // CHECK:  load i32, i32*
 // CHECK:  add nsw i32
 // CHECK:  store i32
 // CHECK:  call void @__kmpc_end_critical(
 
-// CHECK:  call void @__kmpc_barrier(%struct.ident_t* @{{.+}}, i32 %{{.+}})
+// CHECK:  call void @__kmpc_syncwarp(i32 [[MASK]])
 // CHECK:  [[NEW_CC_VAL:%.+]] = add nsw i32 [[CC_VAL]], 1
 // CHECK:  store i32 [[NEW_CC_VAL]], i32* [[CC]],
 // CHECK:  br label
 
 // CHECK: attributes #[[#BARRIER_ATTRS]] = {{.*}} convergent {{.*}}
 
 #endif