[MLIR][NVVM] Add exit (#120251)

PR adds `exit` instruction to nvvm dialect.

Author: grypp

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
 //===----------------------------------------------------------------------===//

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