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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-nvptx

Author: Dhairyashil R G (dhairyashilRG)

<details>
<summary>Changes</summary>

`nvvm.mbarrier.arrive_drop` cannot be compiled for sm_80 today, even though the PTX instruction behind it has existed since sm_80. Using the Op with an sm_80 target fails in the backend:

```
LLVM ERROR: Cannot select: intrinsic %llvm.nvvm.mbarrier.arrive.drop.scope.cta.space.cta
```

The Op always lowers to one of the newer `mbarrier.arrive.drop.*` intrinsics, and all of those are `Requires<[PTX80, SM90]>` or `[PTX86, SM90]` in the backend. The older `llvm.nvvm.mbarrier.arrive.drop{,.shared}` intrinsics, which are available from sm_80, are never used.

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 use the legacy intrinsic when all these conditions are met - the same fix made for `mbarrier.arrive`, which was not extended to `arrive_drop` at the time.

After the patch:

```
mlir-translate -mlir-to-llvmir   ->  call i64 @<!-- -->llvm.nvvm.mbarrier.arrive.drop.shared(ptr addrspace(3) %0)
llc -mtriple=nvptx64 -mcpu=sm_80 ->  mbarrier.arrive_drop.shared.b64 %rd2, [%rd1];
```

**Why only that one configuration:** the `count` operand is what raises the requirement. `mbarrier.arrive_drop.shared.b64 state, [addr]` is PTX ISA 7.0 / sm_80, while the forms that take an explicit `count` are PTX ISA 8.0 / sm_90 (8.6 for `.relaxed`). So the other configurations correctly keep needing sm_90, and only the no-count form has an sm_80 instruction to fall back to.

Both intrinsics produce the same instruction, and neither spells `.release.cta` because those are the defaults. `llvm/test/CodeGen/NVPTX/mbarrier.ll` already compiles `mbarrier.arrive_drop.shared.b64` at `-mcpu=sm_80` and runs `%ptxas-verify -arch=sm_80` over it.

Two notes for reviewers:

- The blank line before the final `return` is removed so that this function and `MBarrierArriveOp::getIntrinsicIDAndArgs` are identical apart from the intrinsic names.
- Which also means the two are now exact duplicates, and `MBarrierArriveExpectTxOp` / `MBarrierArriveDropExpectTxOp` share the same index computation. **I will send a follow-up NFC** folding the four into a common helper, kept out of this patch so the functional change stays small.

## Testing

`mlir/test/Target/LLVMIR/nvvm/mbar_arrive_drop.mlir`: the two default-form calls now check the legacy intrinsic - the same two-line update #<!-- -->172476 made to `mbar_arrive.mlir`. The file's other 22 calls are unchanged and act as the negative half of the test: `shared_cluster`, `scope = cluster`, `relaxed` and explicit `count` must keep using the newer intrinsics.

Added `@<!-- -->mbarrier_arrive_drop_ignore_retval`, mirroring the function of the same name in `mbar_arrive.mlir`. The result-less `shared::cta` form reaches the new branch and had no coverage in either file.

Assisted-by: Claude

---
Full diff: https://github.com/llvm/llvm-project/pull/222916.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp (+8-2) 
- (modified) mlir/test/Target/LLVMIR/nvvm/mbar_arrive_drop.mlir (+15-2) 


``````````diff
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
+}

``````````

</details>


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