[Mlir-commits] [mlir] [mlir][rocdl] Add `InferIntRangeInterface` to readlane operations. (PR #183610)

[Erick Ochoa Lopez]

https://github.com/amd-eochoalo created https://github.com/llvm/llvm-project/pull/183610

Will allow integer range analysis to work for these operations.

>From b2ddf8d20566466c850e1ea9275ca381338c17ba Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Thu, 26 Feb 2026 14:43:42 -0500
Subject: [PATCH 1/2] [mlir][rocdl] Add InferIntRangeInterface to readlane
 operations.

---
 .../mlir/Dialect/LLVMIR/ROCDLDialect.h        |  1 +
 mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td  | 27 +++++++++++++++++--
 mlir/lib/Dialect/LLVMIR/CMakeLists.txt        |  1 +
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h b/mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h
index ce1fe5a03c494..1474b612afd3a 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h
@@ -26,6 +26,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"
 
 ///// Ops /////
diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 1f2d472611120..a46118f30ae97 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -15,6 +15,7 @@
 
 include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
+include "mlir/Interfaces/InferIntRangeInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 //===----------------------------------------------------------------------===//
@@ -381,7 +382,10 @@ def ROCDL_BallotOp :
   let assemblyFormat = "$pred attr-dict `:` type($res)";
 }
 
-def ROCDL_ReadfirstlaneOp : ROCDL_IntrOp<"readfirstlane", [], [0], [AllTypesMatch<["res", "src"]>], 1>,
+def ROCDL_ReadfirstlaneOp : ROCDL_IntrOp<"readfirstlane", [], [0],
+  [AllTypesMatch<["res", "src"]>,
+   DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>],
+  1>,
   Arguments<(ins LLVM_Type:$src)> {
   let results = (outs LLVM_Type:$res);
   let summary = "Get the value in first active lane.";
@@ -402,9 +406,20 @@ def ROCDL_ReadfirstlaneOp : ROCDL_IntrOp<"readfirstlane", [], [0], [AllTypesMatc
   let assemblyFormat = [{
     $src attr-dict `:` type($res)
   }];
+
+  let extraClassDefinition = [{
+    void $cppClass::inferResultRanges(
+        ArrayRef<::mlir::ConstantIntRanges> argRanges,
+        SetIntRangeFn setResultRanges) {
+      setResultRanges(getResult(), argRanges[0]);
+    }
+  }];
 }
 
-def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res", "src0"]>], 1>,
+def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0],
+  [AllTypesMatch<["res", "src0"]>,
+   DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>],
+  1>,
   Arguments<(ins LLVM_Type:$src0,
                  I32:$src1)> {
   let results = (outs LLVM_Type:$res);
@@ -426,6 +441,14 @@ def ROCDL_ReadlaneOp : ROCDL_IntrOp<"readlane", [], [0], [AllTypesMatch<["res",
   let assemblyFormat = [{
     $src0 `,` $src1  attr-dict `:` `(` type($src0) `,` type($src1) `)` `->` type($res)
   }];
+
+  let extraClassDefinition = [{
+    void $cppClass::inferResultRanges(
+        ArrayRef<::mlir::ConstantIntRanges> argRanges,
+        SetIntRangeFn setResultRanges) {
+      setResultRanges(getResult(), argRanges[0]);
+    }
+  }];
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/LLVMIR/CMakeLists.txt b/mlir/lib/Dialect/LLVMIR/CMakeLists.txt
index 21b60ae747a71..9aff8bbc15c76 100644
--- a/mlir/lib/Dialect/LLVMIR/CMakeLists.txt
+++ b/mlir/lib/Dialect/LLVMIR/CMakeLists.txt
@@ -86,6 +86,7 @@ add_mlir_dialect_library(MLIRROCDLDialect
   Core
 
   LINK_LIBS PUBLIC
+  MLIRInferIntRangeInterface
   MLIRIR
   MLIRLLVMDialect
   MLIRSideEffectInterfaces

>From ae4ebb2835ea6b7c268219f81c4b20cfaf7cade2 Mon Sep 17 00:00:00 2001
From: Erick Ochoa <erick.ochoalopez at amd.com>
Date: Thu, 26 Feb 2026 14:49:55 -0500
Subject: [PATCH 2/2] Add test

---
 .../LLVMIR/rocdl-int-range-interface.mlir     | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)
 create mode 100644 mlir/test/Dialect/LLVMIR/rocdl-int-range-interface.mlir

diff --git a/mlir/test/Dialect/LLVMIR/rocdl-int-range-interface.mlir b/mlir/test/Dialect/LLVMIR/rocdl-int-range-interface.mlir
new file mode 100644
index 0000000000000..366df5656ec11
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/rocdl-int-range-interface.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-opt -int-range-optimizations -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: func @readfirstlane
+// CHECK: test.reflect_bounds {smax = 10 : si32, smin = 0 : si32, umax = 10 : ui32, umin = 0 : ui32}
+func.func @readfirstlane() -> i32 {
+  %0 = test.with_bounds { umin = 0 : ui32, umax = 10 : ui32, smin = 0 : si32, smax = 10 : si32 } : i32
+  %1 = rocdl.readfirstlane %0 : i32
+  %2 = test.reflect_bounds %1 : i32
+  return %2 : i32
+}
+
+// -----
+
+// CHECK-LABEL: func @readlane
+// CHECK: test.reflect_bounds {smax = 10 : si32, smin = 0 : si32, umax = 10 : ui32, umin = 0 : ui32}
+func.func @readlane(%idx: i32) -> i32 {
+  %0 = test.with_bounds { umin = 0 : ui32, umax = 10 : ui32, smin = 0 : si32, smax = 10 : si32 } : i32
+  %1 = rocdl.readlane %0, %idx : (i32, i32) -> i32
+  %2 = test.reflect_bounds %1 : i32
+  return %2 : i32
+}