[mlir][GPU] Extend gpu.barrier with scope and named-barrier support - #195692
Conversation
|
@llvm/pr-subscribers-backend-amdgpu @llvm/pr-subscribers-mlir-gpu Author: Krzysztof Drewniak (krzysz00) ChangesThis commit adds two features to gpu.barrier that are supported on targets like recent AMDGPU chips and SPIR-V. The first of these is named barriers, which allow creating a barrier object that is initialized with the number of subgroups that must arrive at it before those subgroups are released. These are represented in MLIR with a new The other change is adding a "scope" enum and using it to specify the execution scope of barriers. This allows for rerpresenting cluster- and subgroup-wide barriers (the latter exists on AMDGPU and Nvidia, and while I suspect Nvidia has cluster-scope barriers, I didn't go looking) and allows us to fully lower to SPIR-V's OpControlBarrier. While these are two different features, I figured I'd land them in one PR so the full API for Patch is 44.64 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/195692.diff 19 Files Affected:
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..68f06b625170e 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUBase.td
@@ -115,6 +115,29 @@ def GPU_AddressSpaceAttr :
def GPU_AddressSpaceAttrArray : TypedArrayAttrBase<GPU_AddressSpaceAttr, "GPU Address Space array">;
+def GPU_ScopeThread : I32EnumCase<"Thread", 0, "thread">;
+def GPU_ScopeSubgroup : I32EnumCase<"Subgroup", 1, "subgroup">;
+def GPU_ScopeWorkgroup : I32EnumCase<"Workgroup", 2, "workgroup">;
+def GPU_ScopeCluster : I32EnumCase<"Cluster", 3, "cluster">;
+def GPU_ScopeDevice : I32EnumCase<"Device", 4, "device">;
+def GPU_ScopeCrossDevice : I32EnumCase<"CrossDevice", 5, "cross_device">;
+
+def GPU_ScopeEnum : I32Enum<
+ "Scope", "barrier execution scope", [
+ GPU_ScopeThread,
+ GPU_ScopeSubgroup,
+ GPU_ScopeWorkgroup,
+ GPU_ScopeCluster,
+ GPU_ScopeDevice,
+ GPU_ScopeCrossDevice
+ ]> {
+ let cppNamespace = "::mlir::gpu";
+}
+
+def GPU_ScopeAttr : EnumAttr<GPU_Dialect, GPU_ScopeEnum, "scope"> {
+ let assemblyFormat = "`<` $value `>`";
+}
+
def GPU_Dimension : GPU_I32Enum<"Dimension",
"a dimension, either 'x', 'y', or 'z'",
[
@@ -169,6 +192,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<NamedBarrierType, Type, TypeStorage> {
+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..a9df774ba76c7 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<GPU_AddressSpaceAttrArray> :$address_spaces)> {
- let summary = "Synchronizes all work items of a workgroup.";
+ Arguments<(ins
+ OptionalAttr<GPU_AddressSpaceAttrArray>:$address_spaces,
+ Optional<GPU_NamedBarrier>:$named_barrier,
+ DefaultValuedAttr<GPU_ScopeAttr,
+ "::mlir::gpu::Scope::Workgroup">:$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 <subgroup>
+ // Synchronize across the entire device.
+ gpu.barrier scope <device>
+ ```
+
+ 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<workgroup>]
// No memory accesses required to be visible.
gpu.barrier memfence []
@@ -1461,17 +1487,56 @@ 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<workgroup>]
+ // Subgroup barrier with a global fence.
+ gpu.barrier memfence [#gpu.address_space<global>] scope <subgroup>
+ ```
+
+ 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 must execute the same dynamic instance of `gpu.barrier` before
+ executing any other instance of it.
+ }];
+ 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<DefaultResource>]>]> {
+ 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/GPUToNVVM/GPUToNVVM.td b/mlir/lib/Conversion/GPUToNVVM/GPUToNVVM.td
index 0fcda38631a9b..b2288938fec22 100644
--- a/mlir/lib/Conversion/GPUToNVVM/GPUToNVVM.td
+++ b/mlir/lib/Conversion/GPUToNVVM/GPUToNVVM.td
@@ -17,6 +17,4 @@ include "mlir/IR/PatternBase.td"
include "mlir/Dialect/GPU/IR/GPUOps.td"
include "mlir/Dialect/LLVMIR/NVVMOps.td"
-def : Pat<(GPU_BarrierOp : $op $memory_fence), (NVVM_Barrier0Op)>;
-
#endif // MLIR_CONVERSION_GPUTONVVM_TD
diff --git a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
index 44a4cafcc224b..02b71a638c108 100644
--- a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
+++ b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
@@ -370,6 +370,37 @@ struct AssertOpToAssertfailLowering
/// Import the GPU Ops to NVVM Patterns.
#include "GPUToNVVM.cpp.inc"
+struct GPUBarrierOpNVVMLowering final
+ : public ConvertOpToLLVMPattern<gpu::BarrierOp> {
+ using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
+
+ LogicalResult
+ matchAndRewrite(gpu::BarrierOp op, gpu::BarrierOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ if (op.getNamedBarrier())
+ return rewriter.notifyMatchFailure(
+ op, "named barriers are not supported on NVVM");
+
+ gpu::Scope scope = op.getScope();
+ switch (scope) {
+ case gpu::Scope::Workgroup:
+ rewriter.replaceOpWithNewOp<NVVM::Barrier0Op>(op);
+ return success();
+ case gpu::Scope::Subgroup: {
+ // Emit __syncwarp(0xFFFFFFFF) for full-warp sync.
+ Value mask = LLVM::ConstantOp::create(
+ rewriter, op.getLoc(), rewriter.getI32Type(),
+ rewriter.getI32IntegerAttr(0xFFFFFFFF));
+ rewriter.replaceOpWithNewOp<NVVM::SyncWarpOp>(op, mask);
+ return success();
+ }
+ default:
+ return rewriter.notifyMatchFailure(
+ op, "unsupported scope for NVVM barrier lowering");
+ }
+ }
+};
+
/// A pass that replaces all occurrences of GPU device operations with their
/// corresponding NVVM equivalent.
///
@@ -506,8 +537,8 @@ void mlir::populateGpuToNVVMConversionPatterns(
// TODO: Pass benefit to generated patterns.
populateWithGenerated(patterns);
- patterns.add<GPUPrintfOpToVPrintfLowering, AssertOpToAssertfailLowering>(
- converter, benefit);
+ patterns.add<GPUBarrierOpNVVMLowering, GPUPrintfOpToVPrintfLowering,
+ AssertOpToAssertfailLowering>(converter, benefit);
patterns.add<
gpu::index_lowering::OpLowering<gpu::ThreadIdOp, NVVM::ThreadIdXOp,
NVVM::ThreadIdYOp, NVVM::ThreadIdZOp>>(
diff --git a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
index 58ab1c799b574..39810e498f6a7 100644
--- a/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
+++ b/mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
@@ -36,6 +36,7 @@
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Matchers.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -517,6 +518,51 @@ struct GPUShuffleOpLowering : public ConvertOpToLLVMPattern<gpu::ShuffleOp> {
}
};
+/// 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<ArrayAttr> addrSpaces,
+ ConversionPatternRewriter &rewriter, Location loc,
+ StringRef scope, bool before) {
+ bool fenceGlobal = false;
+ bool fenceLDS = false;
+
+ if (addrSpaces) {
+ for (auto spaceAttr : addrSpaces->getAsRange<gpu::AddressSpaceAttr>()) {
+ 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<LLVM::MMRATagAttr>("amdgpu-synchronize-as", "local");
+ else if (fenceGlobal && !fenceLDS)
+ mmra = rewriter.getAttr<LLVM::MMRATagAttr>("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);
+}
+
struct GPUBarrierOpLowering final : ConvertOpToLLVMPattern<gpu::BarrierOp> {
GPUBarrierOpLowering(const LLVMTypeConverter &converter,
amdgpu::Chipset chipset)
@@ -528,71 +574,147 @@ struct GPUBarrierOpLowering final : ConvertOpToLLVMPattern<gpu::BarrierOp> {
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<ArrayAttr> addrSpacesToFence = op.getAddressSpaces();
-
- if (addrSpacesToFence) {
- for (auto spaceAttr :
- addrSpacesToFence->getAsRange<gpu::AddressSpaceAttr>()) {
- 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::Scope scope = op.getScope();
+
+ // Subgroup (wave) scope.
+ if (scope == gpu::Scope::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<LLVM::MMRATagAttr>("amdgpu-synchronize-as", "local");
- } else if (fenceGlobal && !fenceLDS) {
- mmra = rewriter.getAttr<LLVM::MMRATagAttr>("amdgpu-synchronize-as",
- "global");
+ // Device and CrossDevice scopes are not directly representable.
+ if (scope == gpu::Scope::Device || scope == gpu::Scope::CrossDevice)
+ return op.emitOpError("device/cross_device scope barriers are not "
+ "supported on AMDGPU");
+
+ // Cluster scope: gfx1250+ only, signal/wait with constant -3.
+ if (scope == gpu::Scope::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, -3);
+ ROCDL::BarrierWaitOp::create(rewriter, loc,
+ static_cast<int16_t>(-3));
+ 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::Scope::Workgroup);
+
+ // 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<int16_t>(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::BarrierWaitOp::create(rewriter, loc, static_cast<int16_t>(-1));
}
+ emitFences(op.getAddressSpaces(), rewriter, loc, "workgroup",
+ /*before=*/false);
+ rewriter.eraseOp(op);
+ return success();
+ }
+};
+
+struct GPUInitializeNamedBarrierOpLowering final
+ : ConvertOpToLLVMPattern<gpu::InitializeNamedBarrierOp> {
+ GPUInitializeNamedBarrierOpLowering(const LLVMTypeConverter &converter,
+ amdgpu::Chipset chipset)
+ : ConvertOpToLLVMPattern<gpu::InitializeNamedBarrierOp>(converter),
+ chipset(chipset) {}
+
+ 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();
- if (emitFences) {
- auto acqFence = LLVM::FenceOp::create(
- rewriter, loc, LLVM::AtomicOrdering::acquire, scope);
- if (mmra)
- acqFence->setDiscardableAttr(LLVM::LLVMDialect::getMmraAttrName(),
- mmra);
+ // The count must be a constant for rocdl.s.barrier.init.
+ IntegerAttr countAttr;
+ if (!matchPattern(op.getMemberCount(), m_Const...
[truncated]
|
|
( @grypp You're here for the Nvidia changes, and I hope I grabbed the right person for those ) |
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
| OpBuilder<(ins "Value":$memrefToFence)>]; | ||
| } | ||
|
|
||
| def GPU_InitializeNamedBarrierOp |
There was a problem hiding this comment.
PTX has had named barriers via bar.sync a, b for a long time, and the nvvm.barrier op already exposes them as barrierId + numberOfThreads. Has lowering gpu.barrier named(%nb) to that been considered?
What would gpu.initialize_named_barrier lower to on NVVM, given there's no init op in PTX?
There was a problem hiding this comment.
Does y'all's LLVM backend not have a thing where it'll allocate the IDs for you?
There was a problem hiding this comment.
If not, we'd probably want to lower to {i32, i32} where the first i32is the value of some constant global symbol we conjure up and the second is the member count * the size of of the subgroup.
There was a problem hiding this comment.
Nope, named barriers are pre-allocated in hardware, so we have a single instruction for that.
There was a problem hiding this comment.
The problem with the current design is that it requires use-def analysis to recover the barrier id and thread count. This isn’t always possible: a named barrier type can be produced in one function and passed to another, at which point the def-chain is broken and the information is lost.
There was a problem hiding this comment.
Once could have id and participants on the IR, like
gpu.named_barrier.init -> !gpu.named_barrier<id = 3, count = 256>
...
func.func @consumer(%nb : !gpu.named_barrier<id = 3, count = 256>)
{
// lowers to nvvm.barrier %c3, %c256
gpu.barrier %nb
}
This way you lose the ability to express a runtime-computed participant count, because types can only carry attributes (constants), not SSA values. But this info is typically known at codegen time.
There was a problem hiding this comment.
Is my implementation proof enough that the current design works?
This commit adds two features to gpu.barrier that are supported on targets like recent AMDGPU chips and SPIR-V. The first of these is named barriers, which allow creating a barrier object that is initialized with the number of subgroups that must arrive at it before those subgroups are released. These are represented in MLIR with a new `!gpu.named_barrier` type and created by `gpu.initialized_named_barrier` operation. These named barriers then become arguments to `gpu.barrier`. The other change is adding a "scope" enum and using it to specify the execution scope of barriers. This allows for rerpresenting cluster- and subgroup-wide barriers (the latter exists on AMDGPU and Nvidia, and while I suspect Nvidia has cluster-scope barriers, I didn't go looking) and allows us to fully lower to SPIR-V's OpControlBarrier. While there are two different features, I figured I'd land them in one PR so the full API for `gpu.barrier` only changes once. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
|
||
| def GPU_ScopeThread : I32EnumCase<"Thread", 0, "thread">; | ||
| def GPU_ScopeSubgroup : I32EnumCase<"Subgroup", 1, "subgroup">; | ||
| def GPU_ScopeWorkgroup : I32EnumCase<"Workgroup", 2, "workgroup">; |
There was a problem hiding this comment.
We could rename this ThreadBlock to be consistent with the remaining of GPU dialect
There was a problem hiding this comment.
Except that we've got some workgroup running around - in the address spaces, for one thing
There was a problem hiding this comment.
But I don't have a strong opinion here (my weak opinion is we should rename blocks to workgroups but I don't want to break the world about it)
| OpBuilder<(ins "Value":$memrefToFence)>]; | ||
| } | ||
|
|
||
| def GPU_InitializeNamedBarrierOp |
There was a problem hiding this comment.
Nope, named barriers are pre-allocated in hardware, so we have a single instruction for that.
| OpBuilder<(ins "Value":$memrefToFence)>]; | ||
| } | ||
|
|
||
| def GPU_InitializeNamedBarrierOp |
There was a problem hiding this comment.
The problem with the current design is that it requires use-def analysis to recover the barrier id and thread count. This isn’t always possible: a named barrier type can be produced in one function and passed to another, at which point the def-chain is broken and the information is lost.
| OpBuilder<(ins "Value":$memrefToFence)>]; | ||
| } | ||
|
|
||
| def GPU_InitializeNamedBarrierOp |
There was a problem hiding this comment.
Once could have id and participants on the IR, like
gpu.named_barrier.init -> !gpu.named_barrier<id = 3, count = 256>
...
func.func @consumer(%nb : !gpu.named_barrier<id = 3, count = 256>)
{
// lowers to nvvm.barrier %c3, %c256
gpu.barrier %nb
}
This way you lose the ability to express a runtime-computed participant count, because types can only carry attributes (constants), not SSA values. But this info is typically known at codegen time.
| class NamedBarrierType | ||
| : public Type::TypeBase<NamedBarrierType, Type, TypeStorage> { | ||
| public: | ||
| using Base::Base; | ||
|
|
||
| static constexpr StringLiteral name = "gpu.named_barrier"; | ||
| }; | ||
|
|
There was a problem hiding this comment.
Why do you define it in C++ and not td?
There was a problem hiding this comment.
That's how the dialect does things?
| def GPU_NamedBarrier : DialectType< | ||
| GPU_Dialect, | ||
| CPred<"::llvm::isa<::mlir::gpu::NamedBarrierType>($_self)">, | ||
| "named barrier type">, | ||
| BuildableType<"mlir::gpu::NamedBarrierType::get($_builder.getContext())">; |
There was a problem hiding this comment.
Using dialect type here is weird, either define the type in ODS or use simpler TypeConstraint.
| that have been associated with a named barrier handle. Named barriers | ||
| require workgroup scope. |
There was a problem hiding this comment.
In view of https://developer.nvidia.com/blog/cooperative-groups/ the requirement of workgroup seems arbitrary.
There was a problem hiding this comment.
If we've got named barriers at cluster scope, I'm down to withdraw that restriction.
There was a problem hiding this comment.
What I meant is, NVIDIA appears to support thread barriers that are smaller than a warp/wave. @grypp , thoughts?
There was a problem hiding this comment.
I think nvvm.barrier requires the member count to be a multiple of the subgorup size? Unless that's been cleaned up since the documentation I looked at?
There was a problem hiding this comment.
Yes, @fabianmcg is right. At this point, you could synchronize almost anything—sub-warp, cluster, sub-thread block, etc.—using mbarrier or some strange atomic tricks.
I think we should focus on supporting named barriers, which are a single instruction (bar.sync), and error out for everything else for the time being
| 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<workgroup>] | ||
| // Subgroup barrier with a global fence. | ||
| gpu.barrier memfence [#gpu.address_space<global>] scope <subgroup> | ||
| ``` |
There was a problem hiding this comment.
This looks too generic, can we scope it a bit? eg. named barriers cannot be used in conjunction with scopes. We can relax this requirement when it proves necessary.
There was a problem hiding this comment.
I'd prefer to keep the options, and would like to to note that on AMD, named barriers can darn well be used with memfence lists (the only current catch is that they need to be workgroup-scoped)
There was a problem hiding this comment.
I was objecting named barriers and scope. Not the memfence + (named or scope).
There was a problem hiding this comment.
Named barriers + scope is going to be cluster-scoped named barriers on gfx[eventually] (I forget if it's a gfx1250 or not)
81dada2 to
bf035d9
Compare
|
@grypp The current design doesn't require use-def analysis to get at the barrier ID - see the changes I pushed. We can allocate barrier IDs at lowering time, and then |
|
And half the point of this setup and of having an "initialize" setup is to, following the AMDGCN backend and SPIR-V, not require users to count |
Ok nice. Can we also support attributes on this type so we don’t need to pass SSA values when calling a function? |
|
Hm, yeah, not a bad idea - though I'm going to push back on On the other hand, Nvidia are the only people that are making the user count before hitting LLVM, so I'm tempted to go with the general abstraction. Is that something you'd be willing to take as a followup? |
fabianmcg
left a comment
fabianmcg
left a comment
There was a problem hiding this comment.
LGTM! Please wait a bit for any of the people that previously left comments.
…lvm#195692) This commit adds two features to gpu.barrier that are supported on targets like recent AMDGPU chips, Nvidia's hardware, and SPIR-V. The first of these is named barriers, which allow creating a barrier object that is initialized with the number of subgroups that must arrive at it before those subgroups are released. These are represented in MLIR with a new `!gpu.named_barrier` type and created by `gpu.initialized_named_barrier` operation. These named barriers then become arguments to `gpu.barrier`. The other change is adding a "scope" enum and using it to specify the execution scope of barriers. This allows for rerpresenting cluster- and subgroup-wide barriers (the latter exists on AMDGPU and Nvidia, and while I suspect Nvidia has cluster-scope barriers, I didn't go looking) and allows us to fully lower to SPIR-V's OpControlBarrier. While these are two different features, I figured I'd land them in one PR so the full API for `gpu.barrier` only changes once. --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
5 participants
This commit adds two features to gpu.barrier that are supported on targets like recent AMDGPU chips and SPIR-V.
The first of these is named barriers, which allow creating a barrier object that is initialized with the number of subgroups that must arrive at it before those subgroups are released. These are represented in MLIR with a new
!gpu.named_barriertype and created bygpu.initialized_named_barrieroperation. These named barriers then become arguments togpu.barrier.The other change is adding a "scope" enum and using it to specify the execution scope of barriers. This allows for rerpresenting cluster- and subgroup-wide barriers (the latter exists on AMDGPU and Nvidia, and while I suspect Nvidia has cluster-scope barriers, I didn't go looking) and allows us to fully lower to SPIR-V's OpControlBarrier.
While these are two different features, I figured I'd land them in one PR so the full API for
gpu.barrieronly changes once.