[mlir][GPU] Extend gpu.barrier with scope and named-barrier support

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

mlir/include/mlir/Dialect/GPU/IR/GPUBase.td

@@ -115,6 +115,24 @@ def GPU_AddressSpaceAttr :
 
 def GPU_AddressSpaceAttrArray : TypedArrayAttrBase<GPU_AddressSpaceAttr, "GPU Address Space array">;
 
+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<GPU_Dialect, GPU_BarrierScopeEnum, "barrier_scope"> {
+  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.

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 {

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_BarrierScopeAttr,
+        "::mlir::gpu::BarrierScope::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,35 +1451,94 @@ 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 within a cluster.
+    gpu.barrier scope <cluster>
+    ```
+
+    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 []
     // All memory accesses required to be visible.
     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 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<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> {

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(

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<StringAttr>
+createNVVMNamedBarrierIdGlobal(gpu::InitializeNamedBarrierOp op,
+                               ConversionPatternRewriter &rewriter) {
+  auto funcOp = op->getParentOfType<FunctionOpInterface>();
+  if (!funcOp) {
+    op.emitOpError("must be inside a function-like op");
+    return failure();
+  }
+  Operation *symbolTableOp = funcOp->getParentWithTrait<OpTrait::SymbolTable>();
+  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<LLVM::GlobalOp>())
+    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<gpu::BarrierOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct GPUBarrierOpToNVVMLowering final
+    : public ConvertOpToLLVMPattern<gpu::BarrierOp> {
+  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<NVVM::BarrierOp>(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<NVVM::SyncWarpOp>(op, mask);
+      return success();
+    }
+    default:
+      return rewriter.notifyMatchFailure(
+          op, "unsupported scope for NVVM barrier lowering");
+    }
+  }
+};
+
+struct GPUInitializeNamedBarrierOpToNVVMLowering final
+    : public ConvertOpToLLVMPattern<gpu::InitializeNamedBarrierOp> {
+  using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
 
-  LogicalResult matchAndRewrite(gpu::BarrierOp op,
-                                PatternRewriter &rewriter) const override {
-    rewriter.replaceOpWithNewOp<NVVM::BarrierOp>(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<StringAttr> 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<GPUBarrierToNVVMLowering>(patterns.getContext(), benefit);
-  patterns.add<GPUPrintfOpToVPrintfLowering, AssertOpToAssertfailLowering>(
+  patterns.add<GPUBarrierOpToNVVMLowering,
+               GPUInitializeNamedBarrierOpToNVVMLowering,
+               GPUPrintfOpToVPrintfLowering, AssertOpToAssertfailLowering>(
       converter, benefit);
   patterns.add<
       gpu::index_lowering::OpLowering<gpu::ThreadIdOp, NVVM::ThreadIdXOp,