[MLIR][NVVM] Add nanosleep

[grypp]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Guray Ozen (grypp)

Changes

---

Full diff: https://github.com/llvm/llvm-project/pull/154697.diff

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td (+30)
• (modified) mlir/test/Target/LLVMIR/nvvmir-invalid.mlir (+8)
• (modified) mlir/test/Target/LLVMIR/nvvmir.mlir (+9)

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@+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
+}

[Copilot]

Pull Request Overview

This PR adds support for the NVVM nanosleep instruction in MLIR's NVVM dialect. This operation allows threads to suspend execution for a specified duration in nanoseconds.

• Adds a new nvvm.nanosleep operation with constrained duration parameter (1-1,000,000 nanoseconds)
• Includes test coverage for both valid usage and invalid parameter validation
• Implements the operation to map to the corresponding LLVM intrinsic

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td	Defines the new NVVM_NanosleepOp with parameter constraints and LLVM intrinsic mapping
mlir/test/Target/LLVMIR/nvvmir.mlir	Adds positive test case for nanosleep operation
mlir/test/Target/LLVMIR/nvvmir-invalid.mlir	Adds negative test case for invalid duration parameter

[durga4github]

Clean, LGTM!

[schwarzschild-radius]

@grypp Should this be an SSA value instead of an attribute? PTX spec says "t may be a register or an immediate value."

[grypp]

You're right. I added this op just by looking at PTX instruction. My bad. Let me change this SSA value.