[MLIR][NVVM] Declare InferIntRangeInterface for SpecialRangeableRegisterOp

`SpecialRangeableRegister` NVVM operations may have a `range` attribute set. When this attribute is present, it becomes possible to determine their range.

This PR declares the `InferIntRangeInterface` for all `SpecialRangeableRegister` operations and extracts range information from the `range` attribute.

Author: grypp

mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h

@@ -19,6 +19,7 @@
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "llvm/IR/IntrinsicsNVPTX.h"
 

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

@@ -18,6 +18,7 @@ include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.td"
+include "mlir/Interfaces/InferIntRangeInterface.td"
 
 def LLVM_PointerGeneric : LLVM_PointerInAddressSpace<0>;
 def LLVM_PointerGlobal : LLVM_PointerInAddressSpace<1>;
@@ -134,8 +135,8 @@ class NVVM_SpecialRegisterOp<string mnemonic, list<Trait> traits = []> :
   let assemblyFormat = "attr-dict `:` type($res)";
 }
 
-class NVVM_SpecialRangeableRegisterOp<string mnemonic, list<Trait> traits = []> :
-  NVVM_SpecialRegisterOp<mnemonic, traits> {
+class NVVM_SpecialRangeableRegisterOp<string mnemonic> :
+  NVVM_SpecialRegisterOp<mnemonic, [DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>]> {
   let arguments = (ins OptionalAttr<LLVM_ConstantRangeAttr>:$range);
   let assemblyFormat = "(`range` $range^)? attr-dict `:` type($res)";
   let llvmBuilder = baseLlvmBuilder # setRangeRetAttrCode # baseLlvmBuilderCoda;
@@ -147,6 +148,21 @@ class NVVM_SpecialRangeableRegisterOp<string mnemonic, list<Trait> traits = []>
       build($_builder, $_state, resultType, ::mlir::LLVM::ConstantRangeAttr{});
     }]>
   ];
+
+  // Define this method for the InferIntRangeInterface.
+  let extraClassDefinition = [{
+    // Infer the result ranges based on the range attribute.
+    void $cppClass::inferResultRanges(
+        ArrayRef<::mlir::ConstantIntRanges> argRanges,
+        SetIntRangeFn setResultRanges) {
+        if (auto rangeAttr = getOperation()->getAttrOfType<LLVM::ConstantRangeAttr>("range")) {
+          setResultRanges(getResult(), 
+                          {rangeAttr.getLower(), rangeAttr.getUpper(),
+                          rangeAttr.getLower(), rangeAttr.getUpper()});
+        }
+    }
+  }];
+
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/LLVMIR/nvvm-test-range.mlir

@@ -0,0 +1,28 @@
+// RUN: mlir-opt -int-range-optimizations -canonicalize %s | FileCheck %s
+gpu.module @module{
+    gpu.func @kernel_1() kernel {
+        %tidx = nvvm.read.ptx.sreg.tid.x range <i32, 0, 32> : i32
+        %tidy = nvvm.read.ptx.sreg.tid.y range <i32, 0, 128> : i32
+        %tidz = nvvm.read.ptx.sreg.tid.z range <i32, 0, 4> : i32
+        %c64 = arith.constant 64 : i32
+        
+        %1 = arith.cmpi sgt, %tidx, %c64 : i32
+        scf.if %1 {
+            gpu.printf "threadidx"
+        }
+        %2 = arith.cmpi sgt, %tidy, %c64 : i32
+        scf.if %2 {
+            gpu.printf "threadidy"
+        }
+        %3 = arith.cmpi sgt, %tidz, %c64 : i32
+        scf.if %3 {
+            gpu.printf "threadidz"
+        }
+        gpu.return
+    }
+}
+
+// CHECK-LABEL: gpu.func @kernel_1
+// CHECK-NOT: gpu.printf "threadidx"
+// CHECK: gpu.printf "threadidy"
+// CHECK-NOT: gpu.printf "threadidz"