[Mlir-commits] [mlir] [MLIR][NVVM] Update nvvm.barrier.arrive Op (PR #202608)
Srinivasa Ravi
llvmlistbot at llvm.org
Tue Jun 9 21:59:34 PDT 2026
https://github.com/Wolfram70 updated https://github.com/llvm/llvm-project/pull/202608
>From 151c376ecd01f9be336e38592c9a795be75d3462 Mon Sep 17 00:00:00 2001
From: Srinivasa Ravi <srinivasar at nvidia.com>
Date: Tue, 9 Jun 2026 13:15:14 +0000
Subject: [PATCH 1/3] [MLIR][NVVM] Update nvvm.barrier.arrive Op
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
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 18 +++++++-----------
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 15 +++++++++++++++
mlir/test/Target/LLVMIR/nvvm/barrier.mlir | 18 ++++++++++++++++++
3 files changed, 40 insertions(+), 11 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index dabb519aa1395..7f9e9d7fe1f55 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -1215,9 +1215,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
@@ -1226,19 +1227,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 addfd63946758..3c5692befd819 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3467,6 +3467,21 @@ 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::Intrinsic::ID id =
+ thisOp.getAligned()
+ ? llvm::Intrinsic::nvvm_barrier_cta_arrive_aligned_count
+ : llvm::Intrinsic::nvvm_barrier_cta_arrive_count;
+ llvm::SmallVector<llvm::Value *> args = {
+ barrierId, mt.lookupValue(thisOp.getNumberOfThreads())};
+ return {id, std::move(args)};
+}
+
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/Target/LLVMIR/nvvm/barrier.mlir b/mlir/test/Target/LLVMIR/nvvm/barrier.mlir
index a6a40fb4fa430..f74071e9c075e 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
+}
>From 06795933852aa28678c11837a92c78bac47f3056 Mon Sep 17 00:00:00 2001
From: Srinivasa Ravi <srinivasar at nvidia.com>
Date: Tue, 9 Jun 2026 14:36:53 +0000
Subject: [PATCH 2/3] remove old test
---
mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir | 13 -------------
1 file changed, 13 deletions(-)
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(
>From 587090ac056abbd746440e97a3d75ef1cbe57174 Mon Sep 17 00:00:00 2001
From: Srinivasa Ravi <srinivasar at nvidia.com>
Date: Tue, 9 Jun 2026 14:39:26 +0000
Subject: [PATCH 3/3] address comments
---
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 5 ++---
1 file changed, 2 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 3c5692befd819..eb2a076fb559d 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3473,13 +3473,12 @@ mlir::NVVM::IDArgPair NVVM::BarrierArriveOp::getIntrinsicIDAndArgs(
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;
- llvm::SmallVector<llvm::Value *> args = {
- barrierId, mt.lookupValue(thisOp.getNumberOfThreads())};
- return {id, std::move(args)};
+ return {id, {barrierId, numThreads}};
}
mlir::NVVM::IDArgPair NVVM::BarrierReductionOp::getIntrinsicIDAndArgs(
More information about the Mlir-commits
mailing list