[Mlir-commits] [mlir] [MLIR][NVVM] Fix the lowering of legacy mbar.arrive (PR #172476)

[Durgadoss R]

https://github.com/durga4github created https://github.com/llvm/llvm-project/pull/172476

We have the most basic mbarrier.arrive supported on sm_80.
It supports: Space=cta, scope=cta, No relaxed, No explicit count.

This patch updates the lowering to the legacy intrinsic when all
these conditions are met, addressing a lowering issue for sm_80.

>From 24a1340523a5fba1374ac5b9c23ef134b33e3a9a Mon Sep 17 00:00:00 2001
From: Durgadoss R <durgadossr at nvidia.com>
Date: Tue, 16 Dec 2025 18:12:40 +0530
Subject: [PATCH] [MLIR][NVVM] Fix the lowering of legacy mbar.arrive

We have the most basic mbarrier.arrive supported on sm_80.
It supports: Space=cta, scope=cta, No relaxed, No explicit count.

This patch updates the lowering to the legacy intrinsic
when all these conditions are met, addressing a lowering
issue for sm_80.

Signed-off-by: Durgadoss R <durgadossr at nvidia.com>
---
 mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp    | 10 ++++++++--
 mlir/test/Target/LLVMIR/nvvm/mbar_arrive.mlir |  4 ++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index a3ff904146b92..55771555171c7 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3282,13 +3282,19 @@ mlir::NVVM::IDArgPair MBarrierArriveDropOp::getIntrinsicIDAndArgs(
   if (needCast)
     mbar = castPtrToAddrSpace(builder, mbar, NVVMMemorySpace::Shared);
 
+  // We have the most basic mbarrier.arrive supported on sm_80.
+  // It supports: Space=cta, scope=cta, No relaxed, No explicit count.
+  // So, only for this combination use the legacy intrinsic.
+  bool hasCount = static_cast<bool>(thisOp.getCount());
+  if (!hasCount &&
+      (id == llvm::Intrinsic::nvvm_mbarrier_arrive_scope_cta_space_cta))
+    return {llvm::Intrinsic::nvvm_mbarrier_arrive_shared, {mbar}};
+
   // When count is not explicitly specified, the default is 1.
   llvm::LLVMContext &ctx = mt.getLLVMContext();
-  bool hasCount = static_cast<bool>(thisOp.getCount());
   llvm::Value *count =
       hasCount ? mt.lookupValue(thisOp.getCount())
                : llvm::ConstantInt::get(llvm::Type::getInt32Ty(ctx), 1);
-
   return {id, {mbar, count}};
 }
 
diff --git a/mlir/test/Target/LLVMIR/nvvm/mbar_arrive.mlir b/mlir/test/Target/LLVMIR/nvvm/mbar_arrive.mlir
index 3a4e06636e96f..f406922ea3873 100644
--- a/mlir/test/Target/LLVMIR/nvvm/mbar_arrive.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/mbar_arrive.mlir
@@ -3,7 +3,7 @@
 llvm.func @mbarrier_arrive_generic(%barrier: !llvm.ptr, %count : i32) {
   // CHECK-LABEL: define void @mbarrier_arrive_generic(ptr %0, i32 %1) {
   // CHECK-NEXT: %3 = addrspacecast ptr %0 to ptr addrspace(3)
-  // CHECK-NEXT: %4 = call i64 @llvm.nvvm.mbarrier.arrive.scope.cta.space.cta(ptr addrspace(3) %3, i32 1)
+  // CHECK-NEXT: %4 = call i64 @llvm.nvvm.mbarrier.arrive.shared(ptr addrspace(3) %3)
   // CHECK-NEXT: %5 = addrspacecast ptr %0 to ptr addrspace(3)
   // CHECK-NEXT: %6 = call i64 @llvm.nvvm.mbarrier.arrive.scope.cta.space.cta(ptr addrspace(3) %5, i32 %1)
   // CHECK-NEXT: %7 = addrspacecast ptr %0 to ptr addrspace(3)
@@ -34,7 +34,7 @@ llvm.func @mbarrier_arrive_generic(%barrier: !llvm.ptr, %count : i32) {
 
 llvm.func @mbarrier_arrive_shared(%barrier: !llvm.ptr<3>, %count : i32) {
   // CHECK-LABEL: define void @mbarrier_arrive_shared(ptr addrspace(3) %0, i32 %1) {
-  // CHECK-NEXT: %3 = call i64 @llvm.nvvm.mbarrier.arrive.scope.cta.space.cta(ptr addrspace(3) %0, i32 1)
+  // CHECK-NEXT: %3 = call i64 @llvm.nvvm.mbarrier.arrive.shared(ptr addrspace(3) %0)
   // CHECK-NEXT: %4 = call i64 @llvm.nvvm.mbarrier.arrive.scope.cta.space.cta(ptr addrspace(3) %0, i32 %1)
   // CHECK-NEXT: %5 = call i64 @llvm.nvvm.mbarrier.arrive.scope.cta.space.cta(ptr addrspace(3) %0, i32 %1)
   // CHECK-NEXT: %6 = call i64 @llvm.nvvm.mbarrier.arrive.scope.cluster.space.cta(ptr addrspace(3) %0, i32 %1)