[Mlir-commits] [mlir] [MLIR][NVVM] Update nvvm.barrier.arrive Op (PR #202608)

[Srinivasa Ravi]

https://github.com/Wolfram70 created https://github.com/llvm/llvm-project/pull/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

>From 5a61b345177db2cfecececa96fcd411af3313314 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] [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 9cbf76b9210be..79cc9c05e8ee5 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -1210,9 +1210,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 
@@ -1221,19 +1222,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 2d929f740f137..db59616b7db8b 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3443,6 +3443,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 ecfee955904bd..331b1d9d6f154 100644
--- a/mlir/test/Target/LLVMIR/nvvm/barrier.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/barrier.mlir
@@ -35,3 +35,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
+}