diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h index 1965f7b983805..0d33c9baaa903 100644 --- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h +++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h @@ -3272,6 +3272,13 @@ class OpenMPIRBuilder { /// \param HintInst Hint Instruction for hint clause associated with critical /// /// \returns The insertion position *after* the critical. + /// Emit a critical region once per thread of the block on a target device, so + /// only one lane of a wavefront is inside it at a time. + InsertPointOrErrorTy emitDeviceSerializedCritical(Instruction *EntryCall, + Instruction *ExitCall, + BodyGenCallbackTy BodyGenCB, + FinalizeCallbackTy FiniCB); + LLVM_ABI InsertPointOrErrorTy createCritical(const LocationDescription &Loc, BodyGenCallbackTy BodyGenCB, FinalizeCallbackTy FiniCB, diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp index 5a363d0ac3dbd..b46298300f129 100644 --- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp +++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp @@ -8050,10 +8050,97 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createCritical( getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_end_critical); Instruction *ExitCall = createRuntimeFunctionCall(ExitRTLFn, Args); + // On a GPU the runtime lock is acquired by a single lane of a wavefront, so + // every other lane would enter the region unsynchronized. Give each thread of + // the block its own turn, as clang does in + // CGOpenMPRuntimeGPU::emitCriticalRegion, and let the lock serialize the + // wavefronts against each other. + if (Config.IsGPU.value_or(false)) + return emitDeviceSerializedCritical(EntryCall, ExitCall, BodyGenCB, FiniCB); + return EmitOMPInlinedRegion(OMPD, EntryCall, ExitCall, BodyGenCB, FiniCB, /*Conditional*/ false, /*hasFinalize*/ true); } +/// Emit the critical region once per thread of the block, so only one lane of a +/// wavefront is inside it at a time: +/// +/// mask = __kmpc_warp_active_thread_mask(); +/// for (i = 0; i < __kmpc_get_hardware_num_threads_in_block(); ++i) { +/// if (__kmpc_get_hardware_thread_id_in_block() == i) +/// +/// __kmpc_syncwarp(mask); +/// } +OpenMPIRBuilder::InsertPointOrErrorTy +OpenMPIRBuilder::emitDeviceSerializedCritical(Instruction *EntryCall, + Instruction *ExitCall, + BodyGenCallbackTy BodyGenCB, + FinalizeCallbackTy FiniCB) { + Function *CurFn = Builder.GetInsertBlock()->getParent(); + LLVMContext &Ctx = CurFn->getContext(); + Type *I32 = Type::getInt32Ty(Ctx); + + // The mask must describe the lanes that are active here; an all-ones mask + // makes the reconvergence wait on lanes that are not present. + Value *Mask = createRuntimeFunctionCall( + getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_warp_active_thread_mask), {}); + Value *TId = createRuntimeFunctionCall( + getOrCreateRuntimeFunctionPtr( + OMPRTL___kmpc_get_hardware_thread_id_in_block), + {}); + Value *NumThreads = createRuntimeFunctionCall( + getOrCreateRuntimeFunctionPtr( + OMPRTL___kmpc_get_hardware_num_threads_in_block), + {}); + + // The insertion block has no terminator yet, so use the OMPIRBuilder + // splitter rather than BasicBlock::splitBasicBlock. + BasicBlock *ExitBB = + splitBB(Builder, /*CreateBranch=*/false, "omp.critical.serial.exit"); + BasicBlock *EntryBB = Builder.GetInsertBlock(); + BasicBlock *HeaderBB = + BasicBlock::Create(Ctx, "omp.critical.serial.header", CurFn, ExitBB); + BasicBlock *TurnBB = + BasicBlock::Create(Ctx, "omp.critical.serial.turn", CurFn, ExitBB); + BasicBlock *RegionBB = + BasicBlock::Create(Ctx, "omp.critical.serial.region", CurFn, ExitBB); + BasicBlock *SyncBB = + BasicBlock::Create(Ctx, "omp.critical.serial.sync", CurFn, ExitBB); + + Builder.SetInsertPoint(EntryBB); + Builder.CreateBr(HeaderBB); + + Builder.SetInsertPoint(HeaderBB); + PHINode *Turn = Builder.CreatePHI(I32, 2, "omp.critical.turn"); + Turn->addIncoming(ConstantInt::get(I32, 0), EntryBB); + Builder.CreateCondBr(Builder.CreateICmpSLT(Turn, NumThreads), TurnBB, ExitBB); + + Builder.SetInsertPoint(TurnBB); + Builder.CreateCondBr(Builder.CreateICmpEQ(TId, Turn), RegionBB, SyncBB); + + Builder.SetInsertPoint(RegionBB); + Builder.CreateBr(SyncBB); + Builder.SetInsertPoint(RegionBB->getTerminator()); + // EmitOMPInlinedRegion only relocates the exit call; the entry call stays + // where it was built, which is before the loop. + EntryCall->moveBefore(*RegionBB, RegionBB->getTerminator()->getIterator()); + InsertPointOrErrorTy AfterIP = EmitOMPInlinedRegion( + Directive::OMPD_critical, EntryCall, ExitCall, BodyGenCB, FiniCB, + /*Conditional*/ false, /*hasFinalize*/ true); + if (!AfterIP) + return AfterIP.takeError(); + + Builder.SetInsertPoint(SyncBB); + createRuntimeFunctionCall( + getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_syncwarp), {Mask}); + Value *Next = Builder.CreateAdd(Turn, ConstantInt::get(I32, 1)); + Turn->addIncoming(Next, SyncBB); + Builder.CreateBr(HeaderBB); + + Builder.SetInsertPoint(ExitBB, ExitBB->getFirstInsertionPt()); + return Builder.saveIP(); +} + OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createOrderedDepend(const LocationDescription &Loc, InsertPointTy AllocaIP, unsigned NumLoops, diff --git a/mlir/test/Target/LLVMIR/omptarget-critical-device.mlir b/mlir/test/Target/LLVMIR/omptarget-critical-device.mlir new file mode 100644 index 0000000000000..c3466fd4d7c72 --- /dev/null +++ b/mlir/test/Target/LLVMIR/omptarget-critical-device.mlir @@ -0,0 +1,46 @@ +// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s + +// On a GPU the runtime lock is taken by one lane of a wavefront, so the region +// must additionally be given to one thread of the block at a time. + +module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true, omp.is_target_device = true} { + llvm.func @critical_device(%x : !llvm.ptr, %xval : i32) attributes {omp.declare_target = #omp.declaretarget} { + omp.critical { + llvm.store %xval, %x : i32, !llvm.ptr + omp.terminator + } + llvm.return + } +} + +// CHECK-LABEL: define hidden void @critical_device( +// CHECK: %[[MASK:.*]] = call i64 @__kmpc_warp_active_thread_mask() +// CHECK: %[[TID:.*]] = call i32 @__kmpc_get_hardware_thread_id_in_block() +// CHECK: %[[NTHREADS:.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block() +// CHECK: br label %omp.critical.serial.header + +// CHECK: omp.critical.serial.header: +// CHECK: %[[TURN:.*]] = phi i32 [ 0, %{{.*}} ], [ %[[NEXT:.*]], %omp.critical.serial.sync ] +// CHECK: %[[GO:.*]] = icmp slt i32 %[[TURN]], %[[NTHREADS]] +// CHECK: br i1 %[[GO]], label %omp.critical.serial.turn, label %omp.critical.serial.exit + +// CHECK: omp.critical.serial.turn: +// CHECK: %[[MINE:.*]] = icmp eq i32 %[[TID]], %[[TURN]] +// CHECK: br i1 %[[MINE]], label %omp.critical.serial.region, label %omp.critical.serial.sync + +// The lock must be acquired and released inside the region a single thread +// enters, not around the loop. +// CHECK: omp.critical.serial.region: +// CHECK: call void @__kmpc_critical( +// CHECK: omp.critical.region: +// CHECK: store i32 %{{.*}}, ptr %{{.*}} +// CHECK: omp_region.finalize: +// CHECK: call void @__kmpc_end_critical( +// CHECK: br label %omp.critical.serial.sync + +// CHECK: omp.critical.serial.sync: +// CHECK: call void @__kmpc_syncwarp(i64 %[[MASK]]) +// CHECK: %[[NEXT]] = add i32 %[[TURN]], 1 +// CHECK: br label %omp.critical.serial.header + +// CHECK: omp.critical.serial.exit: