[Mlir-commits] [mlir] [MLIR][NVVM] Add exit (PR #120251)

[Guray Ozen]

https://github.com/grypp created https://github.com/llvm/llvm-project/pull/120251

PR adds `exit` instruction to nvvm dialect.

>From 2c9468a11c8203c5139099373fa1e7e7d0b1f7ad Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Tue, 17 Dec 2024 16:52:53 +0100
Subject: [PATCH] [MLIR][NVVM] Add exit

PR adds `exit` instruction to nvvm dialect.
---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 14 ++++++++++++++
 mlir/test/Target/LLVMIR/nvvmir.mlir         | 10 ++++++++++
 2 files changed, 24 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 5d8772d9d5c5f5..530135b912b9e6 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -2316,6 +2316,20 @@ def NVVM_WgmmaMmaAsyncOp : NVVM_Op<"wgmma.mma_async",
   }];
 }
 
+def NVVM_Exit : NVVM_Op<"exit"> {
+  let summary = "Exit Op";
+  let description = [{
+    Ends execution of a thread.
+    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#control-flow-instructions-exit)
+  }];
+  string llvmBuilder = [{
+    createIntrinsicCall(builder, llvm::Intrinsic::nvvm_exit);
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
+
 //===----------------------------------------------------------------------===//
 // NVVM breakpoint Op
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index 2749e42c40bc34..6a32190694b470 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -696,6 +696,16 @@ llvm.func @nvvm_fence_proxy_tensormap_generic_acquire(%addr : !llvm.ptr) {
   nvvm.fence.proxy.acquire #nvvm.mem_scope<sys> %addr, %c128
   llvm.return
 }
+// -----
+
+// CHECK-LABEL: @nvvm_exit
+llvm.func @nvvm_exit() {
+  // CHECK: call void @llvm.nvvm.exit()
+  nvvm.exit
+  llvm.return
+}
+
+
 
 // -----
 // CHECK-LABEL: @nvvm_breakpoint