[llvm] [clang] [NVPTX] Add builtin for 'exit' handling (PR #79777)

[Joseph Huber via llvm-commits]

https://github.com/jhuber6 updated https://github.com/llvm/llvm-project/pull/79777

>From ea3b32593dd0f2035020313176c6e1a131ef8eb4 Mon Sep 17 00:00:00 2001
From: Joseph Huber <huberjn at outlook.com>
Date: Sun, 28 Jan 2024 21:27:37 -0600
Subject: [PATCH] [NVPTX] Add builtin for 'exit' handling

Summary:
The PTX ISA has always supported the 'exit' instruction to terminate
individual threads. This patch adds a builtin to handle it. See the PTX
documentation for further details.
https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#control-flow-instructions-exit
---
 clang/include/clang/Basic/BuiltinsNVPTX.def | 1 +
 clang/test/CodeGen/builtins-nvptx.c         | 7 +++++++
 llvm/include/llvm/IR/IntrinsicsNVVM.td      | 4 ++++
 llvm/lib/Target/NVPTX/NVPTXIntrinsics.td    | 3 +++
 llvm/test/CodeGen/NVPTX/intrinsics.ll       | 8 ++++++++
 5 files changed, 23 insertions(+)

diff --git a/clang/include/clang/Basic/BuiltinsNVPTX.def b/clang/include/clang/Basic/BuiltinsNVPTX.def
index 4ce8cb111b5cb89..ed67f0877aee37d 100644
--- a/clang/include/clang/Basic/BuiltinsNVPTX.def
+++ b/clang/include/clang/Basic/BuiltinsNVPTX.def
@@ -157,6 +157,7 @@ BUILTIN(__nvvm_read_ptx_sreg_pm3, "i", "n")
 // MISC
 
 BUILTIN(__nvvm_prmt, "UiUiUiUi", "")
+BUILTIN(__nvvm_exit, "v", "r")
 TARGET_BUILTIN(__nvvm_nanosleep, "vi", "n", AND(SM_70, PTX63))
 
 // Min Max
diff --git a/clang/test/CodeGen/builtins-nvptx.c b/clang/test/CodeGen/builtins-nvptx.c
index 6649e23fa9c4a7c..4c2cca2f5af4ca7 100644
--- a/clang/test/CodeGen/builtins-nvptx.c
+++ b/clang/test/CodeGen/builtins-nvptx.c
@@ -173,6 +173,13 @@ __device__ void activemask() {
 
 }
 
+__device__ void exit() {
+
+// CHECK: call void @llvm.nvvm.exit()
+
+  __nvvm_exit();
+
+}
 
 // NVVM intrinsics
 
diff --git a/llvm/include/llvm/IR/IntrinsicsNVVM.td b/llvm/include/llvm/IR/IntrinsicsNVVM.td
index f81fe6d6e74ba84..e432f43f98a305f 100644
--- a/llvm/include/llvm/IR/IntrinsicsNVVM.td
+++ b/llvm/include/llvm/IR/IntrinsicsNVVM.td
@@ -4813,4 +4813,8 @@ def int_nvvm_setmaxnreg_dec_sync_aligned_u32
               [IntrConvergent, IntrNoMem, IntrHasSideEffects, ImmArg<ArgIndex<0>>],
               "llvm.nvvm.setmaxnreg.dec.sync.aligned.u32">;
 
+// Exit
+def int_nvvm_exit : ClangBuiltin<"__nvvm_exit">,
+    Intrinsic<[], [], [IntrConvergent, IntrInaccessibleMemOnly, IntrNoReturn]>;
+
 } // let TargetPrefix = "nvvm"
diff --git a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
index c6f89f1e7829921..87ae1ef6d738bf8 100644
--- a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
+++ b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
@@ -6844,4 +6844,7 @@ multiclass SET_MAXNREG<string Action, Intrinsic Intr> {
 
 defm INT_SET_MAXNREG_INC : SET_MAXNREG<"inc", int_nvvm_setmaxnreg_inc_sync_aligned_u32>;
 defm INT_SET_MAXNREG_DEC : SET_MAXNREG<"dec", int_nvvm_setmaxnreg_dec_sync_aligned_u32>;
+
 } // isConvergent
+
+def INT_EXIT : NVPTXInst<(outs), (ins), "exit;", [(int_nvvm_exit)]>;
diff --git a/llvm/test/CodeGen/NVPTX/intrinsics.ll b/llvm/test/CodeGen/NVPTX/intrinsics.ll
index c09c7a72fd10181..4b7d5c8f2390769 100644
--- a/llvm/test/CodeGen/NVPTX/intrinsics.ll
+++ b/llvm/test/CodeGen/NVPTX/intrinsics.ll
@@ -133,6 +133,13 @@ define i64 @test_clock64() {
   ret i64 %ret
 }
 
+; CHECK-LABEL: test_exit
+define void @test_exit() {
+; CHECK: exit;
+  call void @llvm.nvvm.exit()
+  ret void
+}
+
 declare float @llvm.fabs.f32(float)
 declare double @llvm.fabs.f64(double)
 declare float @llvm.nvvm.sqrt.f(float)
@@ -146,3 +153,4 @@ declare i64 @llvm.ctpop.i64(i64)
 declare i32 @llvm.nvvm.read.ptx.sreg.tid.x()
 declare i32 @llvm.nvvm.read.ptx.sreg.clock()
 declare i64 @llvm.nvvm.read.ptx.sreg.clock64()
+declare void @llvm.nvvm.exit()