[Mlir-commits] [mlir] [mlir][arith] Support hybrid optimization for indexcast indexcastui chains (PR #227199)

[Pranav Alavandi]

https://github.com/pranavalavandi updated https://github.com/llvm/llvm-project/pull/227199

>From 304ab97f95ba36e9d4e3d955c0ab61012af8dbe2 Mon Sep 17 00:00:00 2001
From: pranavalavandi <pranav.alavandi at q-ctrl.com>
Date: Tue, 29 Sep 2026 16:30:38 +1000
Subject: [PATCH] [mlir][arith] Support hybrid optimization for indexcast
 indexcastui

---
 .../Transforms/IntRangeOptimizations.cpp      | 26 ++++++++++++-------
 .../Dialect/Arith/int-range-narrowing.mlir    | 22 ++++++++++++++++
 2 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
index 8d2da09969afe..dd0e5facba1a8 100644
--- a/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/IntRangeOptimizations.cpp
@@ -503,17 +503,22 @@ struct NarrowCmpI final : OpRewritePattern<arith::CmpIOp> {
   SmallVector<unsigned, 4> targetBitwidths;
 };
 
-/// Fold index_cast(index_cast(%arg: i8, index), i8) -> %arg
+/// Fold:
+///  - index_cast(index_cast(%arg: i8, index), i8) -> %arg
+///  - index_castui(index_castui(%arg: i8, index), i8) -> %arg
+///  - index_cast(index_castui(%arg: i8, index), i8) -> %arg
+///  - index_castui(index_cast(%arg: i8, index), i8) -> %arg
 /// This pattern assumes all passed `targetBitwidths` are not wider than index
 /// type.
-template <typename CastOp>
-struct FoldIndexCastChain final : OpRewritePattern<CastOp> {
+
+template <typename OuterCastOp, typename InnerCastOp>
+struct FoldIndexCastChain final : OpRewritePattern<OuterCastOp> {
   FoldIndexCastChain(MLIRContext *context, ArrayRef<unsigned> target)
-      : OpRewritePattern<CastOp>(context), targetBitwidths(target) {}
+      : OpRewritePattern<OuterCastOp>(context), targetBitwidths(target) {}
 
-  LogicalResult matchAndRewrite(CastOp op,
+  LogicalResult matchAndRewrite(OuterCastOp op,
                                 PatternRewriter &rewriter) const override {
-    auto srcOp = op.getIn().template getDefiningOp<CastOp>();
+    auto srcOp = op.getIn().template getDefiningOp<InnerCastOp>();
     if (!srcOp)
       return rewriter.notifyMatchFailure(op, "doesn't come from an index cast");
 
@@ -794,9 +799,12 @@ void mlir::arith::populateIntRangeNarrowingPatterns(
     ArrayRef<unsigned> bitwidthsSupported) {
   patterns.add<NarrowElementwise, NarrowCmpI>(patterns.getContext(), solver,
                                               bitwidthsSupported);
-  patterns.add<FoldIndexCastChain<arith::IndexCastUIOp>,
-               FoldIndexCastChain<arith::IndexCastOp>>(patterns.getContext(),
-                                                       bitwidthsSupported);
+
+  patterns.add<FoldIndexCastChain<arith::IndexCastOp, arith::IndexCastOp>,
+               FoldIndexCastChain<arith::IndexCastOp, arith::IndexCastUIOp>,
+               FoldIndexCastChain<arith::IndexCastUIOp, arith::IndexCastOp>,
+               FoldIndexCastChain<arith::IndexCastUIOp, arith::IndexCastUIOp>>(
+      patterns.getContext(), bitwidthsSupported);
 }
 
 void mlir::arith::populateControlFlowValuesNarrowingPatterns(
diff --git a/mlir/test/Dialect/Arith/int-range-narrowing.mlir b/mlir/test/Dialect/Arith/int-range-narrowing.mlir
index 780e4a918d159..71dcfc985011d 100644
--- a/mlir/test/Dialect/Arith/int-range-narrowing.mlir
+++ b/mlir/test/Dialect/Arith/int-range-narrowing.mlir
@@ -177,6 +177,28 @@ func.func @test_add_cmpi_i64() -> i1 {
   return %5 : i1
 }
 
+//===----------------------------------------------------------------------===//
+// Test folding index cast chains
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func.func @fold_index_cast_castui_chain
+// CHECK-SAME:    (%[[ARG0:.+]]: i8)
+// CHECK-NEXT:    return %[[ARG0]] : i8
+func.func @fold_index_cast_castui_chain(%arg0: i8) -> i8 {
+  %0 = arith.index_cast %arg0 : i8 to index
+  %1 = arith.index_castui %0 : index to i8
+  return %1 : i8
+}
+
+// CHECK-LABEL: func.func @fold_index_castui_cast_chain
+// CHECK-SAME:    (%[[ARG0:.+]]: i8)
+// CHECK-NEXT:    return %[[ARG0]] : i8
+func.func @fold_index_castui_cast_chain(%arg0: i8) -> i8 {
+  %0 = arith.index_castui %arg0 : i8 to index
+  %1 = arith.index_cast %0 : index to i8
+  return %1 : i8
+}
+
 //===----------------------------------------------------------------------===//
 // arith.addi
 //===----------------------------------------------------------------------===//