[llvm] [clang] [NVPTX] Add builtin support for 'nanosleep' PTX instrunction (PR #79888)

[Joseph Huber via cfe-commits]

https://github.com/jhuber6 created https://github.com/llvm/llvm-project/pull/79888

Summary:
This patch adds a builtin for the `nanosleep` PTX function. It takes
either an immediate or a register and sleeps for [0, 2t] nanoseconds
given t. More information at the documentation:
https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#miscellaneous-instructions-nanosleep


>From 44b134ae71e0accab720071b4ced9ccbe74e8078 Mon Sep 17 00:00:00 2001
From: Joseph Huber <huberjn at outlook.com>
Date: Mon, 29 Jan 2024 13:53:56 -0600
Subject: [PATCH] [NVPTX] Add builtin support for 'nanosleep' PTX instrunction

Summary:
This patch adds a builtin for the `nanosleep` PTX function. It takes
either an immediate or a register and sleeps for [0, 2t] nanoseconds
given t. More information at the documentation:
https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#miscellaneous-instructions-nanosleep
---
 clang/include/clang/Basic/BuiltinsNVPTX.def |  1 +
 clang/test/CodeGen/builtins-nvptx.c         | 11 +++++++++++
 llvm/include/llvm/IR/IntrinsicsNVVM.td      |  4 ++++
 llvm/lib/Target/NVPTX/NVPTXIntrinsics.td    |  6 ++++++
 llvm/test/CodeGen/NVPTX/nanosleep.ll        | 20 ++++++++++++++++++++
 5 files changed, 42 insertions(+)
 create mode 100644 llvm/test/CodeGen/NVPTX/nanosleep.ll

diff --git a/clang/include/clang/Basic/BuiltinsNVPTX.def b/clang/include/clang/Basic/BuiltinsNVPTX.def
index 0f2e8260143be78..ef3a37c8753d162 100644
--- a/clang/include/clang/Basic/BuiltinsNVPTX.def
+++ b/clang/include/clang/Basic/BuiltinsNVPTX.def
@@ -155,6 +155,7 @@ BUILTIN(__nvvm_read_ptx_sreg_pm3, "i", "n")
 // MISC
 
 BUILTIN(__nvvm_prmt, "UiUiUiUi", "")
+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 353f3ebb608c2b1..b209e2fbad98fb0 100644
--- a/clang/test/CodeGen/builtins-nvptx.c
+++ b/clang/test/CodeGen/builtins-nvptx.c
@@ -810,6 +810,17 @@ __device__ void nvvm_vote(int pred) {
   // CHECK: ret void
 }
 
+// CHECK-LABEL: nvvm_nanosleep
+__device__ void nvvm_nanosleep(int d) {
+#if __CUDA_ARCH__ >= 700
+  // CHECK_PTX70_SM80: call void @llvm.nvvm.nanosleep
+  __nvvm_nanosleep(d);
+
+  // CHECK_PTX70_SM80: call void @llvm.nvvm.nanosleep
+  __nvvm_nanosleep(1);
+#endif
+}
+
 // CHECK-LABEL: nvvm_mbarrier
 __device__ void nvvm_mbarrier(long long* addr, __attribute__((address_space(3))) long long* sharedAddr, int count, long long state) {
   #if __CUDA_ARCH__ >= 800
diff --git a/llvm/include/llvm/IR/IntrinsicsNVVM.td b/llvm/include/llvm/IR/IntrinsicsNVVM.td
index 5a5ba2592e1467e..5d863b283d0466e 100644
--- a/llvm/include/llvm/IR/IntrinsicsNVVM.td
+++ b/llvm/include/llvm/IR/IntrinsicsNVVM.td
@@ -557,6 +557,10 @@ let TargetPrefix = "nvvm" in {
       DefaultAttrsIntrinsic<[llvm_i32_ty], [llvm_i32_ty, llvm_i32_ty, llvm_i32_ty],
         [IntrNoMem, IntrSpeculatable]>;
 
+  def int_nvvm_nanosleep : ClangBuiltin<"__nvvm_nanosleep">,
+      DefaultAttrsIntrinsic<[], [llvm_i32_ty],
+                            [IntrConvergent, IntrNoMem, IntrHasSideEffects]>;
+
 //
 // Min Max
 //
diff --git a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
index 33f1e4a43e072af..133514f4f48024e 100644
--- a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
+++ b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
@@ -634,6 +634,12 @@ class F_MATH_3<string OpcStr, NVPTXRegClass t_regclass,
 def INT_NVVM_PRMT : F_MATH_3<"prmt.b32 \t$dst, $src0, $src1, $src2;", Int32Regs,
   Int32Regs, Int32Regs, Int32Regs, int_nvvm_prmt>;
 
+def INT_NVVM_NANOSLEEP_I : NVPTXInst<(outs), (ins i32imm:$i), "nanosleep.u32 \t$i;",
+                             [(int_nvvm_nanosleep imm:$i)]>,
+        Requires<[hasPTX<63>, hasSM<70>]>;
+def INT_NVVM_NANOSLEEP_R : NVPTXInst<(outs), (ins Int32Regs:$i), "nanosleep.u32 \t$i;",
+                             [(int_nvvm_nanosleep Int32Regs:$i)]>,
+        Requires<[hasPTX<63>, hasSM<70>]>;
 //
 // Min Max
 //
diff --git a/llvm/test/CodeGen/NVPTX/nanosleep.ll b/llvm/test/CodeGen/NVPTX/nanosleep.ll
new file mode 100644
index 000000000000000..1b2a7bf9476cf5f
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/nanosleep.ll
@@ -0,0 +1,20 @@
+; RUN: llc < %s -march=nvptx64 -O2 -mcpu=sm_70 -mattr=+ptx63 | FileCheck %s
+; RUN: %if ptxas %{ llc < %s -march=nvptx64 -mcpu=sm_70 -mattr=+ptx63 | %ptxas-verify %}
+
+declare void @llvm.nvvm.nanosleep(i32)
+
+; CHECK-LABEL: test_nanosleep_r
+define void @test_nanosleep_r(i32 noundef %d) {
+entry:
+; CHECK: nanosleep.u32   %[[REG:.+]];
+  call void @llvm.nvvm.nanosleep(i32 %d)
+  ret void
+}
+
+; CHECK-LABEL: test_nanosleep_i
+define void @test_nanosleep_i() {
+entry:
+; CHECK: nanosleep.u32   42;
+  call void @llvm.nvvm.nanosleep(i32 42)
+  ret void
+}