[Mlir-commits] [mlir] [mlir][nvvm] Add clock and clock64 special registers (PR #77088)

[Guray Ozen]

https://github.com/grypp created https://github.com/llvm/llvm-project/pull/77088

Tihs PR adds `clock` and `clock64` special registers to NVVM dialect.

>From fba3852b98672e0d73d78f20253ff4fce3432ebe Mon Sep 17 00:00:00 2001
From: Guray Ozen <guray.ozen at gmail.com>
Date: Fri, 5 Jan 2024 13:35:46 +0100
Subject: [PATCH 1/2] [mlir][nvvm] Add clock and clock64 special registers

Tihs PR adds `clock` and `clock64` special registers to NVVM dialect.
---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 5 +++++
 mlir/test/Target/LLVMIR/nvvmir.mlir         | 4 ++++
 2 files changed, 9 insertions(+)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 57986f291de745..52857164ffafb8 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -156,6 +156,11 @@ def NVVM_GridInClusterDimZOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.cluster.nct
 def NVVM_ClusterId : NVVM_SpecialRegisterOp<"read.ptx.sreg.cluster.ctarank">;
 def NVVM_ClusterDim : NVVM_SpecialRegisterOp<"read.ptx.sreg.cluster.nctarank">;
 
+//===----------------------------------------------------------------------===//
+// Clock registers
+def NVVM_ClockOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.clock">;
+def NVVM_Clock64Op : NVVM_SpecialRegisterOp<"read.ptx.sreg.clock64">;
+
 //===----------------------------------------------------------------------===//
 // NVVM approximate op definitions
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index 3fed2c24b314fe..d08f12943ba441 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -58,6 +58,10 @@ llvm.func @nvvm_special_regs() -> i32 {
   %27 = nvvm.read.ptx.sreg.cluster.ctarank : i32
   // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.cluster.nctarank
   %28 = nvvm.read.ptx.sreg.cluster.nctarank : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.clock
+  %29 = nvvm.read.ptx.sreg.clock : i32
+  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.clock64
+  %30 = nvvm.read.ptx.sreg.clock64 : i64
   
   llvm.return %1 : i32
 }

>From 8b963a252e7a9d94f3126ec04262c2ec2b4b4364 Mon Sep 17 00:00:00 2001
From: Guray Ozen <guray.ozen at gmail.com>
Date: Fri, 5 Jan 2024 13:37:17 +0100
Subject: [PATCH 2/2] fix

---
 mlir/test/Target/LLVMIR/nvvmir.mlir | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index d08f12943ba441..6076fce598fbd2 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -60,7 +60,7 @@ llvm.func @nvvm_special_regs() -> i32 {
   %28 = nvvm.read.ptx.sreg.cluster.nctarank : i32
   // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.clock
   %29 = nvvm.read.ptx.sreg.clock : i32
-  // CHECK: call i32 @llvm.nvvm.read.ptx.sreg.clock64
+  // CHECK: call i64 @llvm.nvvm.read.ptx.sreg.clock64
   %30 = nvvm.read.ptx.sreg.clock64 : i64
   
   llvm.return %1 : i32