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

[Dhairyashil R G]

https://github.com/dhairyashilRG updated https://github.com/llvm/llvm-project/pull/222916

>From f01ce51a67bdda11ff1ac900197df1573675c8c2 Mon Sep 17 00:00:00 2001
From: Dhairyashil R G <dhairyashil25 at gmail.com>
Date: Fri, 11 Sep 2026 18:12:06 +0530
Subject: [PATCH] [MLIR][NVVM] Fix the lowering of legacy mbar.arrive_drop

We have the most basic mbarrier.arrive_drop 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.
Previously every configuration lowered to one of the
mbarrier.arrive.drop.{,relaxed.}scope.*.space.* intrinsics, which
require [PTX80, SM90] or [PTX86, SM90], so the Op could not be
selected below sm_90 at all.

This is the same fix as #172476, which was not extended to
arrive_drop at the time.

Assisted-by: Claude
---
 mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp      | 10 ++++++++--
 .../Target/LLVMIR/nvvm/mbar_arrive_drop.mlir    | 17 +++++++++++++++--
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 688b3b0cd468a..9acb3cfd72b70 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -4255,13 +4255,19 @@ mlir::NVVM::IDArgPair MBarrierArriveDropOp::getIntrinsicIDAndArgs(
   if (needCast)
     mbar = castPtrToAddrSpace(builder, mbar, NVVMMemorySpace::Shared);
 
+  // We have the most basic mbarrier.arrive_drop 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_drop_scope_cta_space_cta))
+    return {llvm::Intrinsic::nvvm_mbarrier_arrive_drop_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_drop.mlir b/mlir/test/Target/LLVMIR/nvvm/mbar_arrive_drop.mlir
index 3e06d3825fd1b..255367b5bb051 100644
--- a/mlir/test/Target/LLVMIR/nvvm/mbar_arrive_drop.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/mbar_arrive_drop.mlir
@@ -3,7 +3,7 @@
 llvm.func @mbarrier_arrive_drop_generic(%barrier: !llvm.ptr, %count : i32) {
   // CHECK-LABEL: define void @mbarrier_arrive_drop_generic(ptr %0, i32 %1) {
   // CHECK-NEXT: %3 = addrspacecast ptr %0 to ptr addrspace(3)
-  // CHECK-NEXT: %4 = call i64 @llvm.nvvm.mbarrier.arrive.drop.scope.cta.space.cta(ptr addrspace(3) %3, i32 1)
+  // CHECK-NEXT: %4 = call i64 @llvm.nvvm.mbarrier.arrive.drop.shared(ptr addrspace(3) %3)
   // CHECK-NEXT: %5 = addrspacecast ptr %0 to ptr addrspace(3)
   // CHECK-NEXT: %6 = call i64 @llvm.nvvm.mbarrier.arrive.drop.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_drop_generic(%barrier: !llvm.ptr, %count : i32) {
 
 llvm.func @mbarrier_arrive_drop_shared(%barrier: !llvm.ptr<3>, %count : i32) {
   // CHECK-LABEL: define void @mbarrier_arrive_drop_shared(ptr addrspace(3) %0, i32 %1) {
-  // CHECK-NEXT: %3 = call i64 @llvm.nvvm.mbarrier.arrive.drop.scope.cta.space.cta(ptr addrspace(3) %0, i32 1)
+  // CHECK-NEXT: %3 = call i64 @llvm.nvvm.mbarrier.arrive.drop.shared(ptr addrspace(3) %0)
   // CHECK-NEXT: %4 = call i64 @llvm.nvvm.mbarrier.arrive.drop.scope.cta.space.cta(ptr addrspace(3) %0, i32 %1)
   // CHECK-NEXT: %5 = call i64 @llvm.nvvm.mbarrier.arrive.drop.scope.cta.space.cta(ptr addrspace(3) %0, i32 %1)
   // CHECK-NEXT: %6 = call i64 @llvm.nvvm.mbarrier.arrive.drop.scope.cluster.space.cta(ptr addrspace(3) %0, i32 %1)
@@ -101,3 +101,16 @@ llvm.func @mbarrier_arrive_drop_nocomplete_shared(%barrier: !llvm.ptr<3>) {
   %0 = nvvm.mbarrier.arrive_drop.nocomplete %barrier, %count : !llvm.ptr<3>, i32  -> i64
   llvm.return
 }
+
+llvm.func @mbarrier_arrive_drop_ignore_retval(%barrier: !llvm.ptr<3>) {
+  // CHECK-LABEL: define void @mbarrier_arrive_drop_ignore_retval(ptr addrspace(3) %0) {
+  // CHECK-NEXT: %2 = call i64 @llvm.nvvm.mbarrier.arrive.drop.shared(ptr addrspace(3) %0)
+  // CHECK-NEXT: %3 = call i64 @llvm.nvvm.mbarrier.arrive.drop.shared(ptr addrspace(3) %0)
+  // CHECK-NEXT: ret void
+  // CHECK-NEXT: }
+  // Result silently discarded (backward compatible form)
+  nvvm.mbarrier.arrive_drop %barrier : !llvm.ptr<3>
+  // Result explicitly captured
+  %0 = nvvm.mbarrier.arrive_drop %barrier : !llvm.ptr<3> -> i64
+  llvm.return
+}