diff --git a/mlir/include/mlir/Conversion/GPUToSPIRV/GPUToSPIRV.h b/mlir/include/mlir/Conversion/GPUToSPIRV/GPUToSPIRV.h index 825142d120587..75d3050bc95d9 100644 --- a/mlir/include/mlir/Conversion/GPUToSPIRV/GPUToSPIRV.h +++ b/mlir/include/mlir/Conversion/GPUToSPIRV/GPUToSPIRV.h @@ -35,6 +35,10 @@ void populateGpuWMMAToSPIRVCoopMatrixKHRConversionPatterns( /// conversion to the type converter. void populateMMAToSPIRVCoopMatrixTypeConversion( SPIRVTypeConverter &typeConverter); + +/// Adds `gpu::NamedBarrierType` to `spirv::NamedBarrierType` conversion. +void populateGPUNamedBarrierToSPIRVTypeConversion( + SPIRVTypeConverter &typeConverter); } // namespace mlir #endif // MLIR_CONVERSION_GPUTOSPIRV_GPUTOSPIRV_H diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td b/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td index 49ecc7f6c9b95..d57312849313a 100644 --- a/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td +++ b/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td @@ -115,6 +115,24 @@ def GPU_AddressSpaceAttr : def GPU_AddressSpaceAttrArray : TypedArrayAttrBase; +def GPU_BarrierScopeSubgroup : I32EnumCase<"Subgroup", 1, "subgroup">; +def GPU_BarrierScopeWorkgroup : I32EnumCase<"Workgroup", 2, "workgroup">; +def GPU_BarrierScopeCluster : I32EnumCase<"Cluster", 3, "cluster">; + +def GPU_BarrierScopeEnum : I32Enum< + "BarrierScope", "barrier execution scope", [ + GPU_BarrierScopeSubgroup, + GPU_BarrierScopeWorkgroup, + GPU_BarrierScopeCluster + ]> { + let cppNamespace = "::mlir::gpu"; +} + +def GPU_BarrierScopeAttr : + EnumAttr { + let assemblyFormat = "`<` $value `>`"; +} + def GPU_Dimension : GPU_I32Enum<"Dimension", "a dimension, either 'x', 'y', or 'z'", [ @@ -169,6 +187,11 @@ def GPU_SparseDnTensorHandle : GPU_SparseHandle<"SparseDnTensorHandleType", "den def GPU_SparseSpGEMMOpHandle : GPU_SparseHandle<"SparseSpGEMMOpHandleType", "SpGEMM operation">; def GPU_SparseSpMatHandle : GPU_SparseHandle<"SparseSpMatHandleType", "sparse matrix">; +def GPU_NamedBarrier : DialectType< + GPU_Dialect, + CPred<"::llvm::isa<::mlir::gpu::NamedBarrierType>($_self)">, + "named barrier type">, + BuildableType<"mlir::gpu::NamedBarrierType::get($_builder.getContext())">; //===----------------------------------------------------------------------===// // GPU Interfaces. diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h index 6e9cf709c7585..0e135781d9d92 100644 --- a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h +++ b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h @@ -51,6 +51,14 @@ class AsyncTokenType static constexpr StringLiteral name = "gpu.async_token"; }; +class NamedBarrierType + : public Type::TypeBase { +public: + using Base::Base; + + static constexpr StringLiteral name = "gpu.named_barrier"; +}; + /// MMAMatrixType storage and uniquing. Array is uniqued based on its shape /// and type. struct MMAMatrixStorageType : public TypeStorage { diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td index a5525580fb320..a9032820f5646 100644 --- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td +++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td @@ -1427,13 +1427,21 @@ def GPU_RotateOp : GPU_Op< } def GPU_BarrierOp : GPU_Op<"barrier">, - Arguments<(ins OptionalAttr :$address_spaces)> { - let summary = "Synchronizes all work items of a workgroup."; + Arguments<(ins + OptionalAttr:$address_spaces, + Optional:$named_barrier, + DefaultValuedAttr:$scope + )> { + let summary = "Synchronizes work items within an execution scope."; let description = [{ - The `barrier` op synchronizes all work items of a workgroup. It is used - to coordinate communication between the work items of the workgroup. + The `barrier` op synchronizes work items within the specified execution + scope. By default, the scope is `workgroup`, synchronizing all work items + in a workgroup. ```mlir + // Synchronize all work items in the workgroup, making all prior + // memory accesses visible. gpu.barrier ``` @@ -1443,17 +1451,35 @@ def GPU_BarrierOp : GPU_Op<"barrier">, accessing the same memory can be avoided by synchronizing work items in-between these accesses. - If the `memfence` attribute is specified, the set of memory accesses that must - by completed after the barrier resolves is limited to only those accesses that - read from or write to the specified address spaces (though accesses to other - address spaces may be completed as well, especially if a particular combination - of address spaces is not supported on a given backend). In particular, - specifying `memfence []` creates a barrier that is not required to affect - the visibility of any memory operations and is purely used for synchronizing - work items. + The `scope` attribute controls the execution scope of the barrier: ```mlir - // Only workgroup address spaces accesses required to be visible. + // Synchronize within a subgroup (warp/wavefront). + gpu.barrier scope + // Synchronize within a cluster. + gpu.barrier scope + ``` + + A `named` barrier allows synchronizing a specific subset of subgroups + that have been associated with a named barrier handle. Named barriers + require workgroup scope. + + ```mlir + // Initialize a named barrier for 4 participating members. + %nb = gpu.initialize_named_barrier %c4 : i32 -> !gpu.named_barrier + // Wait on the named barrier. + gpu.barrier named(%nb : !gpu.named_barrier) + ``` + + If the `memfence` attribute is specified, the set of memory accesses that + must be completed after the barrier resolves is limited to only those + accesses that read from or write to the specified address spaces. In + particular, specifying `memfence []` creates a barrier that is not required + to affect the visibility of any memory operations and is purely used for + synchronizing work items. + + ```mlir + // Only workgroup address space accesses required to be visible. gpu.barrier memfence [#gpu.address_space] // No memory accesses required to be visible. gpu.barrier memfence [] @@ -1461,17 +1487,58 @@ def GPU_BarrierOp : GPU_Op<"barrier">, gpu.barrier ``` - Either none or all work items of a workgroup need to execute this op - in convergence. + The three clauses can be combined in any order, but not all combinations may + be supported on a given target: + + ```mlir + // Named barrier with a workgroup-only memory fence. + gpu.barrier named(%nb : !gpu.named_barrier) memfence [#gpu.address_space] + // Subgroup barrier with a global fence. + gpu.barrier memfence [#gpu.address_space] scope + ``` + + Once one thread of execution in a given scope (say, thread in a workgroup) + has executed a particular dynamic instance of `gpu.barrier`, all other threads + in that scope are required to execute the same dynamic instance of `gpu.barrier` + before any thread executes any other instance of it. That is, you cannot, for + example, have the two subgroups of a workgroup arrive at `gpu.barrier` ops in + different branches of an if statement and have this work. + }]; + let assemblyFormat = [{ + oilist( + `named` `(` $named_barrier `:` type($named_barrier) `)` + | `memfence` $address_spaces + | `scope` $scope + ) attr-dict }]; - let assemblyFormat = "(`memfence` $address_spaces^)? attr-dict"; let hasCanonicalizer = 1; + let hasVerifier = 1; let builders = [OpBuilder<( ins CArg<"std::optional<::mlir::gpu::AddressSpace>", "std::nullopt">:$addressSpace)>, OpBuilder<(ins "Value":$memrefToFence)>]; } +def GPU_InitializeNamedBarrierOp + : GPU_Op<"initialize_named_barrier", + [MemoryEffects<[MemAlloc]>]> { + let summary = "Initialize a named barrier with a member count."; + let description = [{ + Initializes a named barrier object with the given number of participating + members (subgroups) and returns a handle to it. All members that will + synchronize on this barrier must be accounted for in the count. + + ```mlir + %nb = gpu.initialize_named_barrier %num_members : i32 -> !gpu.named_barrier + ``` + }]; + let arguments = (ins I32:$member_count); + let results = (outs GPU_NamedBarrier:$result); + let assemblyFormat = [{ + $member_count attr-dict `:` type($member_count) `->` type($result) + }]; +} + def GPU_GPUModuleOp : GPU_Op<"module", [ IsolatedFromAbove, DataLayoutOpInterface, HasDefaultDLTIDataLayout, NoRegionArguments, SymbolTable, Symbol] # GraphRegionNoTerminator.traits> { diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp index 8b9b3e194adf9..adc4f0422a2cb 100644 --- a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp +++ b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp @@ -4407,6 +4407,10 @@ void mlir::amdgpu::populateCommonGPUTypeAndAttributeConversions( } llvm_unreachable("unknown address space enum value"); }); + typeConverter.addConversion([](gpu::NamedBarrierType type) { + return LLVM::LLVMPointerType::get( + type.getContext(), ROCDL::ROCDLDialect::kSharedMemoryAddressSpace); + }); } void mlir::populateAMDGPUTypeAndAttributeConversions( diff --git a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp index 7a10f7f79d596..dab68fd734236 100644 --- a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp +++ b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp @@ -30,6 +30,8 @@ #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h" #include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h" +#include "mlir/IR/SymbolTable.h" +#include "mlir/Interfaces/FunctionInterfaces.h" #include "mlir/Transforms/DialectConversion.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" @@ -94,6 +96,49 @@ convertToNVVMReductionKind(gpu::AllReduceOperation mode) { return std::nullopt; } +static constexpr llvm::StringLiteral kNVVMNamedBarrierIdPrefix = + "__named_barrier_id"; +static constexpr int32_t kNVVMFirstNamedBarrierId = 1; +static constexpr int32_t kNVVMLastNamedBarrierId = 15; +static constexpr int32_t kNVVMWarpSize = 32; + +static FailureOr +createNVVMNamedBarrierIdGlobal(gpu::InitializeNamedBarrierOp op, + ConversionPatternRewriter &rewriter) { + auto funcOp = op->getParentOfType(); + if (!funcOp) { + op.emitOpError("must be inside a function-like op"); + return failure(); + } + Operation *symbolTableOp = funcOp->getParentWithTrait(); + if (!symbolTableOp) { + op.emitOpError( + "enclosing function-like op must have a symbol-table parent"); + return failure(); + } + + int32_t numNamedBarriers = 0; + for (auto globalOp : + symbolTableOp->getRegion(0).front().getOps()) + if (globalOp.getSymName().starts_with(kNVVMNamedBarrierIdPrefix)) + ++numNamedBarriers; + + int32_t barrierId = kNVVMFirstNamedBarrierId + numNamedBarriers; + if (barrierId > kNVVMLastNamedBarrierId) { + op.emitOpError("NVVM supports at most 15 named barriers per CTA"); + return failure(); + } + + OpBuilder detachedBuilder(rewriter.getContext()); + Type i32 = rewriter.getI32Type(); + auto globalOp = LLVM::GlobalOp::create( + detachedBuilder, op.getLoc(), i32, /*isConstant=*/true, + LLVM::Linkage::Internal, kNVVMNamedBarrierIdPrefix, + rewriter.getI32IntegerAttr(barrierId), /*alignment=*/0, + /*addrSpace=*/0); + return SymbolTable(symbolTableOp).insert(globalOp); +} + /// This pass lowers gpu.subgroup_reduce op into to the nvvm.redux op. The op /// must be run by the entire subgroup, otherwise it is undefined behaviour. struct GPUSubgroupReduceOpLowering @@ -367,13 +412,85 @@ struct AssertOpToAssertfailLowering } }; -/// Lowering of gpu.barrier to nvvm.barrier (defaults to barrier id 0). -struct GPUBarrierToNVVMLowering : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; +struct GPUBarrierOpToNVVMLowering final + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; + + LogicalResult + matchAndRewrite(gpu::BarrierOp op, gpu::BarrierOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + if (Value namedBarrier = adaptor.getNamedBarrier()) { + Location loc = op.getLoc(); + Value barrierId = + LLVM::ExtractValueOp::create(rewriter, loc, namedBarrier, 0); + Value numberOfThreads = + LLVM::ExtractValueOp::create(rewriter, loc, namedBarrier, 1); + NVVM::BarrierOp::create(rewriter, loc, barrierId, numberOfThreads, + NVVM::BarrierReductionAttr{}, Value{}); + rewriter.eraseOp(op); + return success(); + } + + gpu::BarrierScope scope = op.getScope(); + switch (scope) { + case gpu::BarrierScope::Workgroup: + rewriter.replaceOpWithNewOp(op); + return success(); + case gpu::BarrierScope::Subgroup: { + // Emit __syncwarp(0xFFFFFFFF) for full-warp sync. + Value mask = + LLVM::ConstantOp::create(rewriter, op.getLoc(), rewriter.getI32Type(), + rewriter.getI32IntegerAttr(0xFFFFFFFF)); + rewriter.replaceOpWithNewOp(op, mask); + return success(); + } + default: + return rewriter.notifyMatchFailure( + op, "unsupported scope for NVVM barrier lowering"); + } + } +}; + +struct GPUInitializeNamedBarrierOpToNVVMLowering final + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; - LogicalResult matchAndRewrite(gpu::BarrierOp op, - PatternRewriter &rewriter) const override { - rewriter.replaceOpWithNewOp(op); + LogicalResult + matchAndRewrite(gpu::InitializeNamedBarrierOp op, + gpu::InitializeNamedBarrierOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Location loc = op.getLoc(); + MLIRContext *ctx = rewriter.getContext(); + Type i32 = rewriter.getI32Type(); + Type namedBarrierType = + getTypeConverter()->convertType(op.getResult().getType()); + if (!namedBarrierType) + return rewriter.notifyMatchFailure(op, "failed to convert result type"); + + FailureOr maybeGlobalName = + createNVVMNamedBarrierIdGlobal(op, rewriter); + if (failed(maybeGlobalName)) + return failure(); + + auto addressOf = LLVM::AddressOfOp::create( + rewriter, loc, LLVM::LLVMPointerType::get(ctx), *maybeGlobalName); + Value barrierId = + LLVM::LoadOp::create(rewriter, loc, i32, addressOf.getResult()); + + Value warpSize = LLVM::ConstantOp::create( + rewriter, loc, i32, rewriter.getI32IntegerAttr(kNVVMWarpSize)); + Value numberOfThreads = + LLVM::MulOp::create(rewriter, loc, adaptor.getMemberCount(), warpSize); + + Value namedBarrier = + LLVM::PoisonOp::create(rewriter, loc, namedBarrierType); + DenseI64ArrayAttr barrierIdPos = rewriter.getDenseI64ArrayAttr({0}); + DenseI64ArrayAttr numberOfThreadsPos = rewriter.getDenseI64ArrayAttr({1}); + namedBarrier = LLVM::InsertValueOp::create(rewriter, loc, namedBarrier, + barrierId, barrierIdPos); + namedBarrier = LLVM::InsertValueOp::create( + rewriter, loc, namedBarrier, numberOfThreads, numberOfThreadsPos); + rewriter.replaceOp(op, namedBarrier); return success(); } }; @@ -493,6 +610,11 @@ void mlir::configureGpuToNVVMConversionLegality(ConversionTarget &target) { void mlir::configureGpuToNVVMTypeConverter(LLVMTypeConverter &converter) { nvgpu::populateCommonGPUTypeAndAttributeConversions(converter); + converter.addConversion([&](gpu::NamedBarrierType type) -> Type { + Type i32 = IntegerType::get(type.getContext(), 32); + return LLVM::LLVMStructType::getLiteral(type.getContext(), {i32, i32}); + }); + // Lowering for MMAMatrixType. converter.addConversion([&](gpu::MMAMatrixType type) -> Type { return convertMMAToLLVMType(type); @@ -511,8 +633,9 @@ void mlir::populateGpuToNVVMConversionPatterns( using gpu::index_lowering::IndexKind; using gpu::index_lowering::IntrType; - patterns.add(patterns.getContext(), benefit); - patterns.add( + patterns.add( converter, benefit); patterns.add< gpu::index_lowering::OpLowering { } }; +/// Emit an LLVM fence with MMRA metadata based on the given address spaces. +/// If `addrSpaces` is nullopt, all memory is fenced (global + LDS). +static void emitFences(std::optional addrSpaces, + ConversionPatternRewriter &rewriter, Location loc, + StringRef scope, bool before) { + bool fenceGlobal = false; + bool fenceLDS = false; + + if (addrSpaces) { + for (auto spaceAttr : addrSpaces->getAsRange()) { + switch (spaceAttr.getValue()) { + case gpu::AddressSpace::Global: + fenceGlobal = true; + break; + case gpu::AddressSpace::Workgroup: + fenceLDS = true; + break; + case gpu::AddressSpace::Private: + case gpu::AddressSpace::Constant: + break; + } + } + } else { + fenceGlobal = true; + fenceLDS = true; + } + + if (!fenceGlobal && !fenceLDS) + return; + + Attribute mmra; + if (fenceLDS && !fenceGlobal) + mmra = + rewriter.getAttr("amdgpu-synchronize-as", "local"); + else if (fenceGlobal && !fenceLDS) + mmra = + rewriter.getAttr("amdgpu-synchronize-as", "global"); + + auto ordering = + before ? LLVM::AtomicOrdering::release : LLVM::AtomicOrdering::acquire; + auto fence = LLVM::FenceOp::create(rewriter, loc, ordering, scope); + if (mmra) + fence->setDiscardableAttr(LLVM::LLVMDialect::getMmraAttrName(), mmra); +} + +static constexpr int32_t kWholeClusterBarrierId = -3; +static constexpr int32_t kWholeWorkgroupBarrierId = -1; struct GPUBarrierOpLowering final : ConvertOpToLLVMPattern { GPUBarrierOpLowering(const LLVMTypeConverter &converter, amdgpu::Chipset chipset) @@ -528,71 +576,142 @@ struct GPUBarrierOpLowering final : ConvertOpToLLVMPattern { matchAndRewrite(gpu::BarrierOp op, gpu::BarrierOp::Adaptor adaptor, ConversionPatternRewriter &rewriter) const override { Location loc = op.getLoc(); - - // Analyze the address_spaces attribute to determine fence behavior. - bool fenceGlobal = false; - bool fenceLDS = false; - std::optional addrSpacesToFence = op.getAddressSpaces(); - - if (addrSpacesToFence) { - for (auto spaceAttr : - addrSpacesToFence->getAsRange()) { - switch (spaceAttr.getValue()) { - case gpu::AddressSpace::Global: - fenceGlobal = true; - break; - case gpu::AddressSpace::Workgroup: - fenceLDS = true; - break; - case gpu::AddressSpace::Private: - case gpu::AddressSpace::Constant: - // Private is thread-local, constant is read-only; no fencing needed. - break; - } - } - } else { - // Default semantics match __syncthreads() and fence both global and LDS. - fenceGlobal = true; - fenceLDS = true; + gpu::BarrierScope scope = op.getScope(); + + // Subgroup (wave) scope. + if (scope == gpu::BarrierScope::Subgroup) { + emitFences(op.getAddressSpaces(), rewriter, loc, "wavefront", + /*before=*/true); + ROCDL::WaveBarrierOp::create(rewriter, loc); + emitFences(op.getAddressSpaces(), rewriter, loc, "wavefront", + /*before=*/false); + rewriter.eraseOp(op); + return success(); } - Attribute mmra; - if (fenceLDS && !fenceGlobal) { - mmra = - rewriter.getAttr("amdgpu-synchronize-as", "local"); - } else if (fenceGlobal && !fenceLDS) { - mmra = rewriter.getAttr("amdgpu-synchronize-as", - "global"); + // Cluster scope: gfx1250+ only, signal/wait with constant -3. + if (scope == gpu::BarrierScope::Cluster) { + if (chipset < amdgpu::Chipset(12, 5, 0)) + return op.emitOpError("cluster scope barriers require gfx1250+"); + emitFences(op.getAddressSpaces(), rewriter, loc, "cluster", + /*before=*/true); + ROCDL::BarrierSignalOp::create(rewriter, loc, kWholeClusterBarrierId); + ROCDL::BarrierWaitOp::create( + rewriter, loc, static_cast(kWholeClusterBarrierId)); + emitFences(op.getAddressSpaces(), rewriter, loc, "cluster", + /*before=*/false); + rewriter.eraseOp(op); + return success(); } - constexpr llvm::StringLiteral scope = "workgroup"; - - bool emitFences = fenceGlobal || fenceLDS; - // Emit release fence if needed. - if (emitFences) { - auto relFence = LLVM::FenceOp::create( - rewriter, loc, LLVM::AtomicOrdering::release, scope); - if (mmra) - relFence->setDiscardableAttr(LLVM::LLVMDialect::getMmraAttrName(), - mmra); + // Workgroup scope (default). + assert(scope == gpu::BarrierScope::Workgroup && "unsupported scope"); + + // Named barrier path. + if (Value namedBarrier = adaptor.getNamedBarrier()) { + if (chipset.majorVersion < 12) + return op.emitOpError("named barriers require gfx12+"); + + emitFences(op.getAddressSpaces(), rewriter, loc, "workgroup", + /*before=*/true); + // A wave must join the named barrier before it may signal it. + ROCDL::BarrierJoinOp::create(rewriter, loc, namedBarrier); + // Signal with memberCnt=0 retains the count from s.barrier.init. + ROCDL::BarrierSignalVarOp::create(rewriter, loc, namedBarrier, + /*memberCnt=*/0); + // id=1 selects the named-barrier wait class; the actual barrier waited + // on is the last one this wave joined. + ROCDL::BarrierWaitOp::create(rewriter, loc, static_cast(1)); + emitFences(op.getAddressSpaces(), rewriter, loc, "workgroup", + /*before=*/false); + rewriter.eraseOp(op); + return success(); } + // Regular workgroup barrier. + emitFences(op.getAddressSpaces(), rewriter, loc, "workgroup", + /*before=*/true); if (chipset.majorVersion < 12) { ROCDL::SBarrierOp::create(rewriter, loc); } else { - ROCDL::BarrierSignalOp::create(rewriter, loc, -1); - ROCDL::BarrierWaitOp::create(rewriter, loc, -1); + ROCDL::BarrierSignalOp::create(rewriter, loc, kWholeWorkgroupBarrierId); + ROCDL::BarrierWaitOp::create( + rewriter, loc, static_cast(kWholeWorkgroupBarrierId)); } + emitFences(op.getAddressSpaces(), rewriter, loc, "workgroup", + /*before=*/false); + rewriter.eraseOp(op); + return success(); + } +}; + +struct GPUInitializeNamedBarrierOpLowering final + : ConvertOpToLLVMPattern { + GPUInitializeNamedBarrierOpLowering(const LLVMTypeConverter &converter, + amdgpu::Chipset chipset) + : ConvertOpToLLVMPattern(converter), + chipset(chipset) {} - if (emitFences) { - auto acqFence = LLVM::FenceOp::create( - rewriter, loc, LLVM::AtomicOrdering::acquire, scope); - if (mmra) - acqFence->setDiscardableAttr(LLVM::LLVMDialect::getMmraAttrName(), - mmra); + amdgpu::Chipset chipset; + + LogicalResult + matchAndRewrite(gpu::InitializeNamedBarrierOp op, + gpu::InitializeNamedBarrierOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + if (chipset.majorVersion < 12) + return op.emitOpError("named barriers require gfx12+"); + + Location loc = op.getLoc(); + + // The count must be a constant for rocdl.s.barrier.init. + IntegerAttr countAttr; + if (!matchPattern(op.getMemberCount(), m_Constant(&countAttr))) + return op.emitOpError( + "named barrier member count must be a constant for ROCDL lowering"); + int32_t count = countAttr.getInt(); + + // Place the global in the symbol-table scope enclosing the function-like + // op that contains this barrier (typically a module). + auto funcOp = op->getParentOfType(); + if (!funcOp) + return op.emitOpError("must be inside a function-like op"); + Operation *symbolTableOp = + funcOp->getParentWithTrait(); + if (!symbolTableOp) + return op.emitOpError( + "enclosing function-like op must have a symbol-table parent"); + + auto targetTy = LLVM::LLVMTargetExtType::get( + rewriter.getContext(), "amdgcn.named.barrier", {}, {0}); + auto ptrTy = LLVM::LLVMPointerType::get(rewriter.getContext(), 3); + + // Build the global detached so SymbolTable::insert can both place it and + // rename it as needed without creating a transient name conflict in IR. + OpBuilder detachedBuilder(rewriter.getContext()); + auto globalOp = LLVM::GlobalOp::create( + detachedBuilder, loc, targetTy, /*isConstant=*/false, + LLVM::Linkage::Internal, "__named_barrier", /*value=*/Attribute(), + /*alignment=*/0, /*addrSpace=*/3); + // Initialize with poison. + { + Region ®ion = globalOp.getInitializerRegion(); + Block *block = detachedBuilder.createBlock(®ion); + detachedBuilder.setInsertionPointToStart(block); + auto poison = LLVM::PoisonOp::create(detachedBuilder, loc, targetTy); + LLVM::ReturnOp::create(detachedBuilder, loc, poison); } + // SymbolTable::insert places the op in the symbol-table body and renames + // the symbol to avoid collisions with any existing entries. + StringAttr globalName = SymbolTable(symbolTableOp).insert(globalOp); - rewriter.eraseOp(op); + // Get address of the global. + rewriter.setInsertionPoint(op); + auto addrOf = LLVM::AddressOfOp::create(rewriter, loc, ptrTy, globalName); + + // Initialize the barrier. + ROCDL::BarrierInitOp::create(rewriter, loc, addrOf, count); + + rewriter.replaceOp(op, addrOf.getResult()); return success(); } }; @@ -789,7 +908,8 @@ void mlir::populateGpuToROCDLConversionPatterns( patterns.add(converter); patterns.add(converter, chipset); + GPUBarrierOpLowering, GPUInitializeNamedBarrierOpLowering>( + converter, chipset); populateMathToROCDLConversionPatterns(converter, patterns, chipset); } diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp index 0a8d80e232456..87a032bfd766b 100644 --- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp +++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp @@ -100,7 +100,8 @@ class GPUReturnOpConversion final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override; }; -/// Pattern to convert a gpu.barrier op into a spirv.ControlBarrier op. +/// Pattern to convert a gpu.barrier op into a spirv.ControlBarrier or +/// spirv.MemoryNamedBarrier op. class GPUBarrierConversion final : public OpConversionPattern { public: using Base::Base; @@ -110,6 +111,18 @@ class GPUBarrierConversion final : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override; }; +/// Pattern to convert a gpu.initialize_named_barrier into +/// spirv.NamedBarrierInitialize. +class GPUInitializeNamedBarrierConversion final + : public OpConversionPattern { +public: + using Base::Base; + + LogicalResult + matchAndRewrite(gpu::InitializeNamedBarrierOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + /// Pattern to convert a gpu.shuffle op into a spirv.GroupNonUniformShuffle op. class GPUShuffleConversion final : public OpConversionPattern { public: @@ -437,18 +450,58 @@ LogicalResult GPUReturnOpConversion::matchAndRewrite( // Barrier. //===----------------------------------------------------------------------===// +/// Map gpu::BarrierScope to spirv::Scope. +static FailureOr +mapGPUBarrierScopeToSPIRV(gpu::BarrierScope gpuScope) { + switch (gpuScope) { + case gpu::BarrierScope::Subgroup: + return spirv::Scope::Subgroup; + case gpu::BarrierScope::Workgroup: + return spirv::Scope::Workgroup; + case gpu::BarrierScope::Cluster: + return failure(); + } + return failure(); +} + LogicalResult GPUBarrierConversion::matchAndRewrite( gpu::BarrierOp barrierOp, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const { MLIRContext *context = getContext(); - // Both execution and memory scope should be workgroup. - auto scope = spirv::ScopeAttr::get(context, spirv::Scope::Workgroup); + + // Map GPU scope to SPIR-V scope. + auto spirvScope = mapGPUBarrierScopeToSPIRV(barrierOp.getScope()); + if (failed(spirvScope)) + return rewriter.notifyMatchFailure( + barrierOp, "cluster scope is not supported in SPIR-V"); + + auto scopeAttr = spirv::ScopeAttr::get(context, *spirvScope); + auto memoryScopeAttr = + spirv::ScopeAttr::get(context, spirv::Scope::Workgroup); + // Require acquire and release memory semantics for workgroup memory. auto memorySemantics = spirv::MemorySemanticsAttr::get( context, spirv::MemorySemantics::WorkgroupMemory | spirv::MemorySemantics::AcquireRelease); - rewriter.replaceOpWithNewOp(barrierOp, scope, scope, - memorySemantics); + + if (adaptor.getNamedBarrier()) { + spirv::MemoryNamedBarrierOp::create(rewriter, barrierOp.getLoc(), + adaptor.getNamedBarrier(), + memoryScopeAttr, memorySemantics); + rewriter.eraseOp(barrierOp); + } else { + rewriter.replaceOpWithNewOp( + barrierOp, scopeAttr, memoryScopeAttr, memorySemantics); + } + return success(); +} + +LogicalResult GPUInitializeNamedBarrierConversion::matchAndRewrite( + gpu::InitializeNamedBarrierOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto nbType = spirv::NamedBarrierType::get(getContext()); + rewriter.replaceOpWithNewOp( + op, nbType, adaptor.getMemberCount()); return success(); } @@ -921,9 +974,10 @@ LogicalResult GPUPrintfConversion::matchAndRewrite( void mlir::populateGPUToSPIRVPatterns(const SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns) { patterns.add< - GPUBarrierConversion, GPUBallotConversion, GPUFuncOpConversion, - GPUModuleConversion, GPUReturnOpConversion, GPUShuffleConversion, - GPURotateConversion, GPUSubgroupBroadcastConversion, + GPUBarrierConversion, GPUInitializeNamedBarrierConversion, + GPUBallotConversion, GPUFuncOpConversion, GPUModuleConversion, + GPUReturnOpConversion, GPUShuffleConversion, GPURotateConversion, + GPUSubgroupBroadcastConversion, LaunchConfigConversion, LaunchConfigConversion, LaunchConfigConversion, @@ -943,3 +997,10 @@ void mlir::populateGPUToSPIRVPatterns(const SPIRVTypeConverter &typeConverter, GPUSubgroupReduceConversion, GPUPrintfConversion>(typeConverter, patterns.getContext()); } + +void mlir::populateGPUNamedBarrierToSPIRVTypeConversion( + SPIRVTypeConverter &typeConverter) { + typeConverter.addConversion([](gpu::NamedBarrierType type) { + return spirv::NamedBarrierType::get(type.getContext()); + }); +} diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp index 5eab05742d401..1b49e9d6305f1 100644 --- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp +++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRVPass.cpp @@ -144,6 +144,7 @@ void GPUToSPIRVPass::runOnOperation() { options.use64bitIndex = this->use64bitIndex; SPIRVTypeConverter typeConverter(targetAttr, options); populateMMAToSPIRVCoopMatrixTypeConversion(typeConverter); + populateGPUNamedBarrierToSPIRVTypeConversion(typeConverter); RewritePatternSet patterns(context); populateGPUToSPIRVPatterns(typeConverter, patterns); diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp index f776129c77405..5d75a6265ea85 100644 --- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp +++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp @@ -285,6 +285,7 @@ struct GPUInlinerInterface : public DialectInlinerInterface { void GPUDialect::initialize() { addTypes(); addTypes(); + addTypes(); addTypes(); addTypes(); addTypes(); @@ -365,6 +366,9 @@ Type GPUDialect::parseType(DialectAsmParser &parser) const { shape, elementType, operand); } + if (keyword == "named_barrier") + return NamedBarrierType::get(context); + if (keyword == getSparseHandleKeyword(SparseHandleKind::DnTensor)) return SparseDnTensorHandleType::get(context); if (keyword == getSparseHandleKeyword(SparseHandleKind::SpMat)) @@ -380,6 +384,7 @@ Type GPUDialect::parseType(DialectAsmParser &parser) const { void GPUDialect::printType(Type type, DialectAsmPrinter &os) const { TypeSwitch(type) .Case([&](Type) { os << "async.token"; }) + .Case([&](Type) { os << "named_barrier"; }) .Case([&](Type) { os << getSparseHandleKeyword(SparseHandleKind::DnTensor); }) @@ -1516,6 +1521,15 @@ LogicalResult RotateOp::verify() { // BarrierOp //===----------------------------------------------------------------------===// +LogicalResult BarrierOp::verify() { + BarrierScope scope = getScope(); + + if (getNamedBarrier() && scope != BarrierScope::Workgroup) + return emitOpError("named barriers require workgroup scope"); + + return success(); +} + /// Remove gpu.barrier after gpu.barrier, the threads are already synchronized! static LogicalResult eraseRedundantGpuBarrierOps(BarrierOp op, PatternRewriter &rewriter) { @@ -1523,6 +1537,14 @@ static LogicalResult eraseRedundantGpuBarrierOps(BarrierOp op, if (!nextOp) return failure(); + // Cannot merge barriers of different scopes. + if (op.getScope() != nextOp.getScope()) + return failure(); + + // Cannot merge named barriers unless both refer to the same handle. + if (op.getNamedBarrier() != nextOp.getNamedBarrier()) + return failure(); + std::optional thisMemfence = op.getAddressSpaces(); std::optional nextMemfence = nextOp.getAddressSpaces(); @@ -1562,7 +1584,9 @@ void BarrierOp::build(mlir::OpBuilder &odsBuilder, if (addressSpace) addressSpacesAttr = odsBuilder.getArrayAttr( AddressSpaceAttr::get(odsBuilder.getContext(), addressSpace.value())); - build(odsBuilder, odsState, addressSpacesAttr); + build( + odsBuilder, odsState, addressSpacesAttr, /*named_barrier=*/Value{}, + BarrierScopeAttr::get(odsBuilder.getContext(), BarrierScope::Workgroup)); } /// Builds a barrier that causes memory operations affecting `memrefToFence` to diff --git a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp index bbc82e59c1428..9591b76a5330c 100644 --- a/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp +++ b/mlir/lib/Dialect/GPU/TransformOps/GPUTransformOps.cpp @@ -13,7 +13,6 @@ #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h" #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h" #include "mlir/Conversion/LLVMCommon/TypeConverter.h" -#include "mlir/Conversion/NVGPUToNVVM/NVGPUToNVVM.h" #include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h" #include "mlir/Dialect/AMDGPU/Utils/Chipset.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -63,11 +62,7 @@ using namespace mlir::transform::gpu; void transform::ApplyGPUToNVVMConversionPatternsOp::populatePatterns( TypeConverter &typeConverter, RewritePatternSet &patterns) { auto &llvmTypeConverter = static_cast(typeConverter); - nvgpu::populateCommonGPUTypeAndAttributeConversions(llvmTypeConverter); - // Used in GPUToNVVM/WmmaOpsToNvvm.cpp so attaching here for now. - // TODO: We should have a single to_nvvm_type_converter. - llvmTypeConverter.addConversion( - [&](MMAMatrixType type) -> Type { return convertMMAToLLVMType(type); }); + configureGpuToNVVMTypeConverter(llvmTypeConverter); // Set higher benefit, so patterns will run before generic LLVM lowering. populateGpuToNVVMConversionPatterns(llvmTypeConverter, patterns, getBenefit()); diff --git a/mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp b/mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp index 821311a1df5e6..21162546fe1b2 100644 --- a/mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp +++ b/mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp @@ -670,6 +670,17 @@ class BarrierElimination final : public OpRewritePattern { LDBG() << "checking the necessity of: " << barrier << " " << barrier.getLoc(); + // Named barriers have precise arrival-count semantics; never eliminate. + if (barrier.getNamedBarrier()) { + LDBG() << "barrier is a named barrier, retain it"; + return failure(); + } + + if (barrier.getScope() != gpu::BarrierScope::Workgroup) { + LDBG() << "barrier has non-workgroup scope, retain it"; + return failure(); + } + std::optional fencedMemSpaces = barrier.getAddressSpaces(); if (fencedMemSpaces && fencedMemSpaces->empty()) { LDBG() diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-invalid-named-barrier.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-invalid-named-barrier.mlir new file mode 100644 index 0000000000000..faadafa6c4912 --- /dev/null +++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm-invalid-named-barrier.mlir @@ -0,0 +1,26 @@ +// RUN: mlir-opt %s -convert-gpu-to-nvvm -split-input-file -verify-diagnostics + +gpu.module @test_module { + func.func @too_many_named_barriers() { + %c1 = arith.constant 1 : i32 + %nb0 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + %nb1 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + %nb2 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + %nb3 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + %nb4 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + %nb5 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + %nb6 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + %nb7 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + %nb8 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + %nb9 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + %nb10 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + %nb11 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + %nb12 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + %nb13 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + %nb14 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + // expected-error@+2 {{NVVM supports at most 15 named barriers per CTA}} + // expected-error@+1 {{failed to legalize operation 'gpu.initialize_named_barrier'}} + %nb15 = gpu.initialize_named_barrier %c1 : i32 -> !gpu.named_barrier + func.return + } +} diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir index d5aad8321cb9f..b96069ac41a44 100644 --- a/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir +++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-nvvm.mlir @@ -189,6 +189,58 @@ gpu.module @test_module_5 { gpu.barrier func.return } + + // CHECK-LABEL: func @gpu_sync_subgroup() + func.func @gpu_sync_subgroup() { + // CHECK: %[[WARP_MASK:.*]] = llvm.mlir.constant(-1 : i32) : i32 + // CHECK: nvvm.bar.warp.sync %[[WARP_MASK]] + gpu.barrier scope + func.return + } + + // CHECK-LABEL: func @gpu_named_barriers + // CHECK-SAME: (%[[MEMBER_COUNT:.*]]: i32) + func.func @gpu_named_barriers(%member_count : i32) { + // CHECK: %[[ID0_ADDR:.*]] = llvm.mlir.addressof @[[$NB0:__named_barrier_id[_0-9]*]] : !llvm.ptr + // CHECK: %[[ID0:.*]] = llvm.load %[[ID0_ADDR]] : !llvm.ptr -> i32 + // CHECK: %[[WARP_SIZE0:.*]] = llvm.mlir.constant(32 : i32) : i32 + // CHECK: %[[THREADS0:.*]] = llvm.mul %[[MEMBER_COUNT]], %[[WARP_SIZE0]] : i32 + // CHECK: %[[DESC0:.*]] = llvm.mlir.poison : !llvm.struct<(i32, i32)> + // CHECK: %[[DESC1:.*]] = llvm.insertvalue %[[ID0]], %[[DESC0]][0] : !llvm.struct<(i32, i32)> + // CHECK: %[[DESC2:.*]] = llvm.insertvalue %[[THREADS0]], %[[DESC1]][1] : !llvm.struct<(i32, i32)> + %nb0 = gpu.initialize_named_barrier %member_count : i32 -> !gpu.named_barrier + %c2 = arith.constant 2 : i32 + // CHECK: %[[ID1_ADDR:.*]] = llvm.mlir.addressof @[[$NB1:__named_barrier_id[_0-9]*]] : !llvm.ptr + // CHECK: %[[ID1:.*]] = llvm.load %[[ID1_ADDR]] : !llvm.ptr -> i32 + // CHECK: %[[WARP_SIZE1:.*]] = llvm.mlir.constant(32 : i32) : i32 + // CHECK: %[[THREADS1:.*]] = llvm.mul %{{.*}}, %[[WARP_SIZE1]] : i32 + // CHECK: %[[DESC3:.*]] = llvm.mlir.poison : !llvm.struct<(i32, i32)> + // CHECK: %[[DESC4:.*]] = llvm.insertvalue %[[ID1]], %[[DESC3]][0] : !llvm.struct<(i32, i32)> + // CHECK: %[[DESC5:.*]] = llvm.insertvalue %[[THREADS1]], %[[DESC4]][1] : !llvm.struct<(i32, i32)> + %nb1 = gpu.initialize_named_barrier %c2 : i32 -> !gpu.named_barrier + // CHECK: %[[BARRIER_ID0:.*]] = llvm.extractvalue %[[DESC2]][0] : !llvm.struct<(i32, i32)> + // CHECK: %[[BARRIER_THREADS0:.*]] = llvm.extractvalue %[[DESC2]][1] : !llvm.struct<(i32, i32)> + // CHECK: nvvm.barrier id = %[[BARRIER_ID0]] number_of_threads = %[[BARRIER_THREADS0]] + gpu.barrier named(%nb0 : !gpu.named_barrier) + // CHECK: %[[BARRIER_ID1:.*]] = llvm.extractvalue %[[DESC5]][0] : !llvm.struct<(i32, i32)> + // CHECK: %[[BARRIER_THREADS1:.*]] = llvm.extractvalue %[[DESC5]][1] : !llvm.struct<(i32, i32)> + // CHECK: nvvm.barrier id = %[[BARRIER_ID1]] number_of_threads = %[[BARRIER_THREADS1]] + gpu.barrier named(%nb1 : !gpu.named_barrier) + func.return + } + + // CHECK-LABEL: func @gpu_named_barrier_arg + // CHECK-SAME: (%[[NB:.*]]: !llvm.struct<(i32, i32)>) + func.func @gpu_named_barrier_arg(%nb : !gpu.named_barrier) { + // CHECK: %[[BARRIER_ID:.*]] = llvm.extractvalue %[[NB]][0] : !llvm.struct<(i32, i32)> + // CHECK: %[[BARRIER_THREADS:.*]] = llvm.extractvalue %[[NB]][1] : !llvm.struct<(i32, i32)> + // CHECK: nvvm.barrier id = %[[BARRIER_ID]] number_of_threads = %[[BARRIER_THREADS]] + gpu.barrier named(%nb : !gpu.named_barrier) + func.return + } + + // CHECK: llvm.mlir.global internal constant @[[$NB0]](1 : i32) {addr_space = 0 : i32} : i32 + // CHECK: llvm.mlir.global internal constant @[[$NB1]](2 : i32) {addr_space = 0 : i32} : i32 } @@ -1214,4 +1266,3 @@ module attributes {gpu.container_module} { } } } - diff --git a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-barrier.mlir b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-barrier.mlir index 11a874cfb89e7..618d1889b8478 100644 --- a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-barrier.mlir +++ b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-barrier.mlir @@ -77,4 +77,22 @@ func.func @barrier_constant_only() { gpu.barrier memfence [#gpu.address_space] func.return } + +// CHECK-LABEL: func @barrier_subgroup_scope +func.func @barrier_subgroup_scope() { + // CHECK-NEXT: llvm.fence syncscope("wavefront") release + // CHECK-NEXT: rocdl.wave.barrier + // CHECK-NEXT: llvm.fence syncscope("wavefront") acquire + gpu.barrier scope + func.return +} + +// CHECK-LABEL: func @barrier_subgroup_scope_no_fence +func.func @barrier_subgroup_scope_no_fence() { + // CHECK-NEXT: rocdl.wave.barrier + // CHECK-NOT: llvm.fence + gpu.barrier scope memfence [] + func.return +} + } diff --git a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-barriers-gfx12.mlir b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-barriers-gfx12.mlir new file mode 100644 index 0000000000000..c6a9574ca43c1 --- /dev/null +++ b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-barriers-gfx12.mlir @@ -0,0 +1,54 @@ +// RUN: mlir-opt %s -convert-gpu-to-rocdl='chipset=gfx1250' --mlir-print-local-scope | FileCheck %s + +gpu.module @test_module { + +// CHECK-LABEL: func @named_barrier +func.func @named_barrier() { + %member_count = arith.constant 4 : i32 + // CHECK: %[[ADDR:.*]] = llvm.mlir.addressof @[[NB:__named_barrier[_0-9]*]] : !llvm.ptr<3> + // CHECK: rocdl.s.barrier.init %[[ADDR]] member_cnt = 4 + %nb = gpu.initialize_named_barrier %member_count : i32 -> !gpu.named_barrier + // CHECK: llvm.fence syncscope("workgroup") release + // CHECK: rocdl.s.barrier.join %[[ADDR]] + // CHECK: rocdl.s.barrier.signal.var %[[ADDR]] member_cnt = 0 + // CHECK: rocdl.s.barrier.wait id = 1 + // CHECK: llvm.fence syncscope("workgroup") acquire + gpu.barrier named(%nb : !gpu.named_barrier) + func.return +} + +// CHECK-LABEL: func @two_named_barriers +func.func @two_named_barriers() { + %c4 = arith.constant 4 : i32 + %c8 = arith.constant 8 : i32 + // CHECK: %[[ADDR0:.*]] = llvm.mlir.addressof @[[NB0:__named_barrier[_0-9]*]] : !llvm.ptr<3> + // CHECK: rocdl.s.barrier.init %[[ADDR0]] member_cnt = 4 + %nb0 = gpu.initialize_named_barrier %c4 : i32 -> !gpu.named_barrier + // CHECK: %[[ADDR1:.*]] = llvm.mlir.addressof @[[NB1:__named_barrier[_0-9]*]] : !llvm.ptr<3> + // CHECK: rocdl.s.barrier.init %[[ADDR1]] member_cnt = 8 + %nb1 = gpu.initialize_named_barrier %c8 : i32 -> !gpu.named_barrier + // CHECK: rocdl.s.barrier.join %[[ADDR0]] + // CHECK: rocdl.s.barrier.signal.var %[[ADDR0]] member_cnt = 0 + // CHECK: rocdl.s.barrier.wait id = 1 + gpu.barrier named(%nb0 : !gpu.named_barrier) + // CHECK: rocdl.s.barrier.join %[[ADDR1]] + // CHECK: rocdl.s.barrier.signal.var %[[ADDR1]] member_cnt = 0 + // CHECK: rocdl.s.barrier.wait id = 1 + gpu.barrier named(%nb1 : !gpu.named_barrier) + func.return +} + +// CHECK-LABEL: func @cluster_scope +func.func @cluster_scope() { + // CHECK: llvm.fence syncscope("cluster") release + // CHECK-NEXT: rocdl.s.barrier.signal id = -3 + // CHECK-NEXT: rocdl.s.barrier.wait id = -3 + // CHECK-NEXT: llvm.fence syncscope("cluster") acquire + gpu.barrier scope + func.return +} + +// One LDS global per gpu.initialize_named_barrier. +// CHECK-COUNT-3: llvm.mlir.global internal @__named_barrier{{[_0-9]*}}() {addr_space = 3 : i32} : !llvm.target<"amdgcn.named.barrier", 0> + +} diff --git a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-invalid-named-barrier.mlir b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-invalid-named-barrier.mlir new file mode 100644 index 0000000000000..3f39f4abcf396 --- /dev/null +++ b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-invalid-named-barrier.mlir @@ -0,0 +1,10 @@ +// RUN: mlir-opt %s -convert-gpu-to-rocdl='chipset=gfx1100' -split-input-file -verify-diagnostics + +gpu.module @test_module { + func.func @initialize_named_barrier_pre_gfx12(%count : i32) { + // expected-error@+2 {{named barriers require gfx12+}} + // expected-error@+1 {{failed to legalize}} + %nb = gpu.initialize_named_barrier %count : i32 -> !gpu.named_barrier + func.return + } +} diff --git a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-named-barrier-non-const.mlir b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-named-barrier-non-const.mlir new file mode 100644 index 0000000000000..c9ce2794f1422 --- /dev/null +++ b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl-named-barrier-non-const.mlir @@ -0,0 +1,10 @@ +// RUN: mlir-opt %s -convert-gpu-to-rocdl='chipset=gfx1250' -split-input-file -verify-diagnostics + +gpu.module @test_module { + func.func @non_constant_member_count(%count : i32) { + // expected-error@+2 {{named barrier member count must be a constant for ROCDL lowering}} + // expected-error@+1 {{failed to legalize}} + %nb = gpu.initialize_named_barrier %count : i32 -> !gpu.named_barrier + func.return + } +} diff --git a/mlir/test/Conversion/GPUToSPIRV/gpu-to-spirv.mlir b/mlir/test/Conversion/GPUToSPIRV/gpu-to-spirv.mlir index 7bf6f8419be0d..5bcd180a37c16 100644 --- a/mlir/test/Conversion/GPUToSPIRV/gpu-to-spirv.mlir +++ b/mlir/test/Conversion/GPUToSPIRV/gpu-to-spirv.mlir @@ -128,3 +128,41 @@ module attributes {gpu.container_module} { return } } + +// ----- + +module attributes { + gpu.container_module, + spirv.target_env = #spirv.target_env< + #spirv.vce, #spirv.resource_limits<>> +} { + gpu.module @kernels { + // CHECK-LABEL: spirv.func @barrier_subgroup_scope + gpu.func @barrier_subgroup_scope() kernel + attributes {spirv.entry_point_abi = #spirv.entry_point_abi} { + // CHECK: spirv.ControlBarrier , , + gpu.barrier scope + gpu.return + } + } +} + +// ----- + +module attributes { + gpu.container_module, + spirv.target_env = #spirv.target_env< + #spirv.vce, #spirv.resource_limits<>> +} { + gpu.module @kernels { + // CHECK-LABEL: spirv.func @named_barrier + gpu.func @named_barrier(%member_count : i32) kernel + attributes {spirv.entry_point_abi = #spirv.entry_point_abi} { + // CHECK: %[[NB:.*]] = spirv.NamedBarrierInitialize %[[MEMBER_COUNT:.*]] : i32 -> !spirv.named_barrier + %nb = gpu.initialize_named_barrier %member_count : i32 -> !gpu.named_barrier + // CHECK: spirv.MemoryNamedBarrier %[[NB]], , : !spirv.named_barrier + gpu.barrier named(%nb : !gpu.named_barrier) + gpu.return + } + } +} diff --git a/mlir/test/Dialect/GPU/barrier-elimination.mlir b/mlir/test/Dialect/GPU/barrier-elimination.mlir index b9ceb8a8d424b..dd134c150d746 100644 --- a/mlir/test/Dialect/GPU/barrier-elimination.mlir +++ b/mlir/test/Dialect/GPU/barrier-elimination.mlir @@ -343,6 +343,14 @@ attributes {__parallel_region_boundary_for_test} { return } +// CHECK-LABEL: @non_workgroup_barrier_retained +func.func @non_workgroup_barrier_retained() +attributes {__parallel_region_boundary_for_test} { + // CHECK: gpu.barrier scope + gpu.barrier scope + return +} + // CHECK-LABEL: @empty_barrrier_retained func.func @empty_barrrier_retained() attributes {__parallel_region_boundary_for_test} { diff --git a/mlir/test/Dialect/GPU/canonicalize.mlir b/mlir/test/Dialect/GPU/canonicalize.mlir index 1283c1465ca47..7627af11c636c 100644 --- a/mlir/test/Dialect/GPU/canonicalize.mlir +++ b/mlir/test/Dialect/GPU/canonicalize.mlir @@ -67,6 +67,51 @@ func.func @erase_barriers_empty_memfence() { return } +// CHECK-LABEL: func @erase_barriers_same_scope +// CHECK-NEXT: gpu.barrier scope +// CHECK-NEXT: return +func.func @erase_barriers_same_scope() { + gpu.barrier scope + gpu.barrier scope + return +} + +// CHECK-LABEL: func @no_fold_different_scope +// CHECK-NEXT: gpu.barrier scope +// CHECK-NEXT: gpu.barrier scope +// CHECK-NEXT: return +func.func @no_fold_different_scope() { + gpu.barrier scope + gpu.barrier scope + return +} + +// CHECK-LABEL: func @no_fold_different_named_barriers +// CHECK-NEXT: gpu.initialize_named_barrier +// CHECK-NEXT: gpu.initialize_named_barrier +// CHECK-NEXT: gpu.barrier named +// CHECK-NEXT: gpu.barrier named +// CHECK-NEXT: return +func.func @no_fold_different_named_barriers(%member_count : i32) { + %nb1 = gpu.initialize_named_barrier %member_count : i32 -> !gpu.named_barrier + %nb2 = gpu.initialize_named_barrier %member_count : i32 -> !gpu.named_barrier + gpu.barrier named(%nb1 : !gpu.named_barrier) + gpu.barrier named(%nb2 : !gpu.named_barrier) + return +} + +// CHECK-LABEL: func @fold_same_named_barrier +// CHECK-NEXT: gpu.initialize_named_barrier +// CHECK-NEXT: gpu.barrier named +// CHECK-NOT: gpu.barrier +// CHECK-NEXT: return +func.func @fold_same_named_barrier(%member_count : i32) { + %nb = gpu.initialize_named_barrier %member_count : i32 -> !gpu.named_barrier + gpu.barrier named(%nb : !gpu.named_barrier) + gpu.barrier named(%nb : !gpu.named_barrier) + return +} + // ----- // Replace uses of gpu.wait op with its async dependency. diff --git a/mlir/test/Dialect/GPU/named-barrier.mlir b/mlir/test/Dialect/GPU/named-barrier.mlir new file mode 100644 index 0000000000000..6aef0648d16f7 --- /dev/null +++ b/mlir/test/Dialect/GPU/named-barrier.mlir @@ -0,0 +1,69 @@ +// RUN: mlir-opt %s --split-input-file -verify-diagnostics | FileCheck %s + +// CHECK-LABEL: func @barrier_default +// CHECK-NEXT: gpu.barrier{{$}} +// CHECK-NEXT: return +func.func @barrier_default() { + gpu.barrier + return +} + +// ----- + +// CHECK-LABEL: func @barrier_with_scope +// CHECK-NEXT: gpu.barrier scope +// CHECK-NEXT: return +func.func @barrier_with_scope() { + gpu.barrier scope + return +} + +// ----- + +// CHECK-LABEL: func @barrier_workgroup_scope_not_printed +// CHECK-NEXT: gpu.barrier{{$}} +// CHECK-NEXT: return +func.func @barrier_workgroup_scope_not_printed() { + gpu.barrier scope + return +} + +// ----- + +// CHECK-LABEL: func @barrier_memfence_and_scope +// CHECK-NEXT: gpu.barrier memfence [#gpu.address_space] scope +// CHECK-NEXT: return +func.func @barrier_memfence_and_scope() { + gpu.barrier memfence [#gpu.address_space] scope + return +} + +// ----- + +// CHECK-LABEL: func @initialize_named_barrier +// CHECK: %[[NB:.*]] = gpu.initialize_named_barrier %[[MEMBER_COUNT:.*]] : i32 -> !gpu.named_barrier +// CHECK: gpu.barrier named(%[[NB]] : !gpu.named_barrier) +func.func @initialize_named_barrier(%member_count : i32) { + %nb = gpu.initialize_named_barrier %member_count : i32 -> !gpu.named_barrier + gpu.barrier named(%nb : !gpu.named_barrier) + return +} + +// ----- + +// CHECK-LABEL: func @named_barrier_with_memfence +// CHECK: gpu.barrier named(%{{.*}} : !gpu.named_barrier) memfence [#gpu.address_space] +func.func @named_barrier_with_memfence(%member_count : i32) { + %nb = gpu.initialize_named_barrier %member_count : i32 -> !gpu.named_barrier + gpu.barrier named(%nb : !gpu.named_barrier) memfence [#gpu.address_space] + return +} + +// ----- + +func.func @named_barrier_non_workgroup_scope(%member_count : i32) { + %nb = gpu.initialize_named_barrier %member_count : i32 -> !gpu.named_barrier + // expected-error @+1 {{named barriers require workgroup scope}} + gpu.barrier named(%nb : !gpu.named_barrier) scope + return +}