[Mlir-commits] [mlir] [MLIR][NVVM] Add nanosleep (PR #154697)

[Guray Ozen]

https://github.com/grypp updated https://github.com/llvm/llvm-project/pull/154697

>From da793ae13a29d5b886122491597425ea525871de Mon Sep 17 00:00:00 2001
From: Guray Ozen <guray.ozen at gmail.com>
Date: Thu, 21 Aug 2025 08:49:42 +0000
Subject: [PATCH 1/2] [MLIR][NVVM] Add nanosleep

---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 30 +++++++++++++++++++++
 mlir/test/Target/LLVMIR/nvvmir-invalid.mlir |  8 ++++++
 mlir/test/Target/LLVMIR/nvvmir.mlir         |  9 +++++++
 3 files changed, 47 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index f9cd58de8915f..2205a77a3bd0c 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -401,6 +401,36 @@ def NVVM_ReduxOp :
    }];   
 }
 
+//===----------------------------------------------------------------------===//
+// NVVM nanosleep
+//===----------------------------------------------------------------------===//
+
+def NVVM_NanosleepOp : NVVM_Op<"nanosleep">,
+  Arguments<(ins 
+  ConfinedAttr<I32Attr, [IntMinValue<1>, IntMaxValue<1000000>]>:$duration)> 
+{
+  let summary = "Suspends the thread for a sleep.";
+
+  let description = [{
+    The Op suspends the thread for a sleep duration approximately close to the 
+    delay `$duration`, specified in nanoseconds. 
+
+    The sleep duration is approximated, but guaranteed to be in the 
+    interval [0, 2*t]. The maximum sleep duration is 1 millisecond. 
+    The implementation may reduce the sleep duration for individual threads 
+    within a warp such that all sleeping threads in the warp wake up together.
+
+    [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#miscellaneous-instructions-nanosleep)
+  }];
+  
+  string llvmBuilder = [{
+      createIntrinsicCall(builder, 
+                          llvm::Intrinsic::nvvm_nanosleep, 
+                          {builder.getInt32($duration)});
+  }];
+  let assemblyFormat = "attr-dict $duration";
+}
+
 //===----------------------------------------------------------------------===//
 // NVVM Performance Monitor events
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir b/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir
index 33398cfb92429..703b8b38b3d87 100644
--- a/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir
@@ -495,3 +495,11 @@ llvm.func @ld_matrix(%arg0: !llvm.ptr<3>) {
   %l = nvvm.ldmatrix %arg0 {num = 1 : i32, layout = #nvvm.mma_layout<col>, shape = #nvvm.ld_st_matrix_shape<m = 16, n = 16>, eltType  = #nvvm.ld_st_matrix_elt_type<b8>} : (!llvm.ptr<3>) -> i32
   llvm.return
 }
+
+// -----
+
+llvm.func @nanosleep() {
+  // expected-error at +1 {{integer constant out of range for attribute}}
+  nvvm.nanosleep 100000000000000
+  llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index c8ba91efbff4d..16191d925959b 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -959,3 +959,12 @@ llvm.func @nvvm_pmevent() {
   nvvm.pmevent mask = 4
   llvm.return
 }
+
+// -----
+
+// CHECK-LABEL: @nanosleep
+llvm.func @nanosleep() {
+  // CHECK: call void @llvm.nvvm.nanosleep(i32 4000)
+  nvvm.nanosleep 4000
+  llvm.return
+}

>From 936d4d1b5288ad6e831dd16215fbb5c051834526 Mon Sep 17 00:00:00 2001
From: Guray Ozen <guray.ozen at gmail.com>
Date: Thu, 21 Aug 2025 11:04:16 +0200
Subject: [PATCH 2/2] Update mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

Co-authored-by: Copilot <175728472+Copilot at users.noreply.github.com>
---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 2205a77a3bd0c..589a3c4221358 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -409,7 +409,7 @@ def NVVM_NanosleepOp : NVVM_Op<"nanosleep">,
   Arguments<(ins 
   ConfinedAttr<I32Attr, [IntMinValue<1>, IntMaxValue<1000000>]>:$duration)> 
 {
-  let summary = "Suspends the thread for a sleep.";
+  let summary = "Suspends the thread for a specified duration.";
 
   let description = [{
     The Op suspends the thread for a sleep duration approximately close to the