[Mlir-commits] [mlir] c1991da - [MLIR][NVVM] Update nvvm.barrier.arrive Op (#202608)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jun 11 20:56:55 PDT 2026
Author: Srinivasa Ravi
Date: 2026-06-12T09:26:50+05:30
New Revision: c1991da69a5f6bd7117224c68f2f66546264a743
URL: https://github.com/llvm/llvm-project/commit/c1991da69a5f6bd7117224c68f2f66546264a743
DIFF: https://github.com/llvm/llvm-project/commit/c1991da69a5f6bd7117224c68f2f66546264a743.diff
LOG: [MLIR][NVVM] Update nvvm.barrier.arrive Op (#202608)
This change updates the `nvvm.barrier.arrive` Op to lower using
intrinsics instead of inline PTX. It also adds a new `aligned` attribute
to the Op to lower to both aligned and unaligned forms.
PTX Spec Reference:
https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-bar
Added:
Modified:
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
mlir/test/Target/LLVMIR/nvvm/barrier.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 8cc83216ca723..16994d487802f 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -1217,9 +1217,10 @@ def NVVM_BarrierReductionOp :
}];
}
-def NVVM_BarrierArriveOp : NVVM_PTXBuilder_Op<"barrier.arrive">
+def NVVM_BarrierArriveOp : NVVM_VoidIntrinsicOp<"barrier.arrive">
{
- let arguments = (ins Optional<I32>:$barrierId, I32:$numberOfThreads);
+ let arguments = (ins Optional<I32>:$barrierId, I32:$numberOfThreads,
+ DefaultValuedAttr<BoolAttr, "true">:$aligned);
let description = [{
Thread that executes this op announces their arrival at the barrier with
@@ -1228,19 +1229,14 @@ def NVVM_BarrierArriveOp : NVVM_PTXBuilder_Op<"barrier.arrive">
The default barrier id is 0 that is similar to `nvvm.barrier` Op. When
`barrierId` is not present, the default barrier id is used.
+ The `aligned` attribute, which defaults to `true`, generates the aligned
+ form of the barrier (all threads in the CTA execute the same barrier
+ instruction). When set to `false`, the unaligned form is generated.
+
[For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-bar)
}];
let assemblyFormat = "(`id` `=` $barrierId^)? `number_of_threads` `=` $numberOfThreads attr-dict";
-
- let extraClassDefinition = [{
- std::string $cppClass::getPtx() {
- std::string ptx = "bar.arrive ";
- if (getBarrierId()) { ptx += "%0, %1;"; }
- else { ptx += "0, %0;"; }
- return ptx;
- }
- }];
}
def NVVM_ClusterArriveOp : NVVM_Op<"cluster.arrive"> {
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 7e9d550cec12e..13de1e7f2b2ad 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3468,6 +3468,20 @@ mlir::NVVM::IDArgPair NVVM::BarrierOp::getIntrinsicIDAndArgs(
return {id, std::move(args)};
}
+mlir::NVVM::IDArgPair NVVM::BarrierArriveOp::getIntrinsicIDAndArgs(
+ Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
+ auto thisOp = cast<NVVM::BarrierArriveOp>(op);
+ llvm::Value *barrierId = thisOp.getBarrierId()
+ ? mt.lookupValue(thisOp.getBarrierId())
+ : builder.getInt32(0);
+ llvm::Value *numThreads = mt.lookupValue(thisOp.getNumberOfThreads());
+ llvm::Intrinsic::ID id =
+ thisOp.getAligned()
+ ? llvm::Intrinsic::nvvm_barrier_cta_arrive_aligned_count
+ : llvm::Intrinsic::nvvm_barrier_cta_arrive_count;
+ return {id, {barrierId, numThreads}};
+}
+
mlir::NVVM::IDArgPair NVVM::BarrierReductionOp::getIntrinsicIDAndArgs(
Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
auto thisOp = cast<NVVM::BarrierReductionOp>(op);
diff --git a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
index 5a381ce1e679e..b012e695072a8 100644
--- a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
+++ b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
@@ -582,19 +582,6 @@ func.func @cp_async_bulk_wait_group() {
func.return
}
-// -----
-
-// CHECK-LABEL: @llvm_nvvm_barrier_arrive
-// CHECK-SAME: (%[[barId:.*]]: i32, %[[numberOfThreads:.*]]: i32)
-llvm.func @llvm_nvvm_barrier_arrive(%barID : i32, %numberOfThreads : i32) {
- // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "bar.arrive 0, $0;", "r" %[[numberOfThreads]] : (i32) -> ()
- nvvm.barrier.arrive number_of_threads = %numberOfThreads
- // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "bar.arrive $0, $1;", "r,r" %[[barId]], %[[numberOfThreads]] : (i32, i32) -> ()
- nvvm.barrier.arrive id = %barID number_of_threads = %numberOfThreads
- llvm.return
-}
-
-
// -----
llvm.func @init_mbarrier(
diff --git a/mlir/test/Target/LLVMIR/nvvm/barrier.mlir b/mlir/test/Target/LLVMIR/nvvm/barrier.mlir
index 6e44bdcc135d9..994daa24048e6 100644
--- a/mlir/test/Target/LLVMIR/nvvm/barrier.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/barrier.mlir
@@ -57,3 +57,21 @@ llvm.func @llvm_nvvm_barrier(%barID : i32, %numberOfThreads : i32, %redOperand :
llvm.return
}
+
+// LLVM-LABEL: @llvm_nvvm_barrier_arrive(
+// LLVM-SAME: i32 %[[barId:.*]], i32 %[[numThreads:.*]])
+llvm.func @llvm_nvvm_barrier_arrive(%barID : i32, %numberOfThreads : i32) {
+ // LLVM: call void @llvm.nvvm.barrier.cta.arrive.aligned.count(i32 0, i32 %[[numThreads]])
+ // CHECK: nvvm.barrier.arrive number_of_threads = %{{.*}}
+ nvvm.barrier.arrive number_of_threads = %numberOfThreads
+ // LLVM: call void @llvm.nvvm.barrier.cta.arrive.aligned.count(i32 %[[barId]], i32 %[[numThreads]])
+ // CHECK: nvvm.barrier.arrive id = %{{.*}} number_of_threads = %{{.*}}
+ nvvm.barrier.arrive id = %barID number_of_threads = %numberOfThreads
+ // LLVM: call void @llvm.nvvm.barrier.cta.arrive.count(i32 0, i32 %[[numThreads]])
+ // CHECK: nvvm.barrier.arrive number_of_threads = %{{.*}} {aligned = false}
+ nvvm.barrier.arrive number_of_threads = %numberOfThreads {aligned = false}
+ // LLVM: call void @llvm.nvvm.barrier.cta.arrive.count(i32 %[[barId]], i32 %[[numThreads]])
+ // CHECK: nvvm.barrier.arrive id = %{{.*}} number_of_threads = %{{.*}} {aligned = false}
+ nvvm.barrier.arrive id = %barID number_of_threads = %numberOfThreads {aligned = false}
+ llvm.return
+}
More information about the Mlir-commits
mailing list