[Mlir-commits] [mlir] [NVVM][MLIR] Remove Pure trait from clock, clock64, globaltimer Ops (PR #147608)

[Pradeep Kumar]

https://github.com/schwarzschild-radius updated https://github.com/llvm/llvm-project/pull/147608

>From ddab1159579d81f3c9bcb55303347b974de4af9c Mon Sep 17 00:00:00 2001
From: pradeepku <pradeepku at nvidia.com>
Date: Mon, 7 Jul 2025 14:05:00 +0530
Subject: [PATCH] [NVVM][MLIR] Remove Pure trait from clock, clock64,
 globaltimer Ops

This commit removes Pure trait from clock, clock64 and globaltimer Ops by creating NVVM_NCSpecialRegisterOp class to represent Ops which return non-constant values. This prevents CSE pass from optimizing away redundant uses of them
---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 25 +++++++++-----
 mlir/test/Dialect/LLVMIR/cse-nvvm.mlir      | 37 +++++++++++++++++++++
 2 files changed, 54 insertions(+), 8 deletions(-)
 create mode 100644 mlir/test/Dialect/LLVMIR/cse-nvvm.mlir

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 6895e946b8a45..f57a89c31b283 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -153,14 +153,23 @@ class NVVM_IntrOp<string mnem, list<Trait> traits = [],
 // NVVM special register op definitions
 //===----------------------------------------------------------------------===//
 
-class NVVM_SpecialRegisterOp<string mnemonic, list<Trait> traits = []> :
+// NVVM_PureSpecialRegisterOp represents special register ops that can
+// speculated and does not touch memory. These operations are always
+// legal to hoist or sink.
+class NVVM_PureSpecialRegisterOp<string mnemonic, list<Trait> traits = []> :
   NVVM_IntrOp<mnemonic, !listconcat(traits, [Pure]), 1> {
   let arguments = (ins);
   let assemblyFormat = "attr-dict `:` type($res)";
 }
 
+class NVVM_SpecialRegisterOp<string mnemonic, list<Trait> traits = []> :
+  NVVM_IntrOp<mnemonic, traits, 1> {
+  let arguments = (ins);
+  let assemblyFormat = "attr-dict `:` type($res)";
+}
+
 class NVVM_SpecialRangeableRegisterOp<string mnemonic, list<Trait> traits = []> :
-  NVVM_SpecialRegisterOp<mnemonic,
+  NVVM_PureSpecialRegisterOp<mnemonic,
     !listconcat(traits,
       [DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>])> {
   let arguments = (ins OptionalAttr<LLVM_ConstantRangeAttr>:$range);
@@ -199,11 +208,11 @@ def NVVM_GridIdOp : NVVM_SpecialRangeableRegisterOp<"read.ptx.sreg.gridid">;
 
 //===----------------------------------------------------------------------===//
 // Lane Mask Comparison Ops
-def NVVM_LaneMaskEqOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.lanemask.eq">;
-def NVVM_LaneMaskLeOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.lanemask.le">;
-def NVVM_LaneMaskLtOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.lanemask.lt">;
-def NVVM_LaneMaskGeOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.lanemask.ge">;
-def NVVM_LaneMaskGtOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.lanemask.gt">;
+def NVVM_LaneMaskEqOp : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.lanemask.eq">;
+def NVVM_LaneMaskLeOp : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.lanemask.le">;
+def NVVM_LaneMaskLtOp : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.lanemask.lt">;
+def NVVM_LaneMaskGeOp : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.lanemask.ge">;
+def NVVM_LaneMaskGtOp : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.lanemask.gt">;
 
 //===----------------------------------------------------------------------===//
 // Thread index and range
@@ -256,7 +265,7 @@ def NVVM_GlobalTimerOp : NVVM_SpecialRegisterOp<"read.ptx.sreg.globaltimer">;
 //===----------------------------------------------------------------------===//
 // envreg registers
 foreach index = !range(0, 32) in {
-  def NVVM_EnvReg # index # Op : NVVM_SpecialRegisterOp<"read.ptx.sreg.envreg" # index>;
+  def NVVM_EnvReg # index # Op : NVVM_PureSpecialRegisterOp<"read.ptx.sreg.envreg" # index>;
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/LLVMIR/cse-nvvm.mlir b/mlir/test/Dialect/LLVMIR/cse-nvvm.mlir
new file mode 100644
index 0000000000000..8d24c3846f178
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/cse-nvvm.mlir
@@ -0,0 +1,37 @@
+// RUN: mlir-opt %s -cse -split-input-file -verify-diagnostics | FileCheck %s
+
+// CHECK-LABEL: @nvvm_special_regs_clock
+llvm.func @nvvm_special_regs_clock() -> !llvm.struct<(i32, i32)> {
+  %0 = llvm.mlir.zero: !llvm.struct<(i32, i32)>
+  // CHECK:  {{.*}} = nvvm.read.ptx.sreg.clock
+  %1 = nvvm.read.ptx.sreg.clock : i32
+  // CHECK:  {{.*}} = nvvm.read.ptx.sreg.clock
+  %2 = nvvm.read.ptx.sreg.clock : i32
+  %4 = llvm.insertvalue %1, %0[0]: !llvm.struct<(i32, i32)>
+  %5 = llvm.insertvalue %2, %4[1]: !llvm.struct<(i32, i32)>
+  llvm.return %5: !llvm.struct<(i32, i32)>
+}
+
+// CHECK-LABEL: @nvvm_special_regs_clock64
+llvm.func @nvvm_special_regs_clock64() -> !llvm.struct<(i64, i64)> {
+  %0 = llvm.mlir.zero: !llvm.struct<(i64, i64)>
+  // CHECK:  {{.*}} = nvvm.read.ptx.sreg.clock64
+  %1 = nvvm.read.ptx.sreg.clock64 : i64
+  // CHECK:  {{.*}} = nvvm.read.ptx.sreg.clock64
+  %2 = nvvm.read.ptx.sreg.clock64 : i64
+  %4 = llvm.insertvalue %1, %0[0]: !llvm.struct<(i64, i64)>
+  %5 = llvm.insertvalue %2, %4[1]: !llvm.struct<(i64, i64)>
+  llvm.return %5: !llvm.struct<(i64, i64)>
+}
+
+// CHECK-LABEL: @nvvm_special_regs_globaltimer
+llvm.func @nvvm_special_regs_globaltimer() -> !llvm.struct<(i64, i64)> {
+  %0 = llvm.mlir.zero: !llvm.struct<(i64, i64)>
+  // CHECK:  {{.*}} = nvvm.read.ptx.sreg.globaltimer
+  %1 = nvvm.read.ptx.sreg.globaltimer : i64
+  // CHECK:  {{.*}} = nvvm.read.ptx.sreg.globaltimer
+  %2 = nvvm.read.ptx.sreg.globaltimer : i64
+  %4 = llvm.insertvalue %1, %0[0]: !llvm.struct<(i64, i64)>
+  %5 = llvm.insertvalue %2, %4[1]: !llvm.struct<(i64, i64)>
+  llvm.return %5: !llvm.struct<(i64, i64)>
+}