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

[llvmlistbot at llvm.org]

Author: Durgadoss R
Date: 2025-12-17T10:05:29+05:30
New Revision: 9d76b2cfce0b8f5aaf5357137a1aa8ff7279d593

URL: https://github.com/llvm/llvm-project/commit/9d76b2cfce0b8f5aaf5357137a1aa8ff7279d593
DIFF: https://github.com/llvm/llvm-project/commit/9d76b2cfce0b8f5aaf5357137a1aa8ff7279d593.diff

LOG: [MLIR][NVVM] Fix the lowering of legacy mbar.arrive (#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.

Signed-off-by: Durgadoss R <durgadossr at nvidia.com>

Added: 
    

Modified: 
    mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
    mlir/test/Target/LLVMIR/nvvm/mbar_arrive.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index a3ff904146b92..6e22bdc3ab135 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3241,13 +3241,19 @@ mlir::NVVM::IDArgPair MBarrierArriveOp::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)