[Mlir-commits] [mlir] [mlir][arith] Restrict mulsi_extended(x, 1) rewrite to integer types (PR #233234)

[Prakhar Dixit]

https://github.com/Prakhar-Dixit created https://github.com/llvm/llvm-project/pull/233234

mlir-opt --canonicalize crashes on arith.mulsi_extended with index operands and a constant 1.

MulSIExtendedRHSOne builds arith.extsi with the op's type, so with index it creates `extsi : i1 to index`, which the verifier rejects. The greedy driver folds new ops before verification, so ExtSIOp::fold hits cast<IntegerType>.

Fix: restrict the pattern to SignlessFixedWidthIntegerLike, the type constraint arith.extsi requires for its result. Index operands are left unchanged.

Add scalar and vector index tests.

Fixes [#231478](https://github.com/llvm/llvm-project/issues/231478)

>From 371c4793a9f00d43dca52348617f2b2a9186777c Mon Sep 17 00:00:00 2001
From: Prakhar Dixit <dixitprakhar11 at gmail.com>
Date: Sun, 18 Oct 2026 19:05:42 +0530
Subject: [PATCH] [mlir][arith] Restrict mulsi_extended(x, 1) rewrite to
 integer types

mlir-opt --canonicalize crashes on arith.mulsi_extended with index operands and a constant 1.

MulSIExtendedRHSOne builds arith.extsi with the op's type, so with index
it creates `extsi : i1 to index`, which the verifier rejects. The greedy driver
folds new ops before verification, so ExtSIOp::fold hits cast<IntegerType>.

Fix: restrict the pattern to SignlessFixedWidthIntegerLike, the type constraint
arith.extsi requires for its result. Index operands are left unchanged.

Add scalar and vector index tests.

Fixes #231478
---
 .../Dialect/Arith/IR/ArithCanonicalization.td |  3 ++-
 mlir/test/Dialect/Arith/canonicalize.mlir     | 22 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
index 2615bcca23e35c..476060a3152702 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
+++ b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
@@ -225,7 +225,8 @@ def IsScalarOrSplatOne :
 
 // mulsi_extended(x, 1) -> [x, extsi(cmpi slt, x, 0)]
 def MulSIExtendedRHSOne :
-    Pattern<(Arith_MulSIExtendedOp $x, (ConstantLikeMatcher AnyAttr:$c1)),
+    Pattern<(Arith_MulSIExtendedOp SignlessFixedWidthIntegerLike:$x,
+                                   (ConstantLikeMatcher AnyAttr:$c1)),
             [(replaceWithValue $x),
              (Arith_ExtSIOp(Arith_CmpIOp
                 ConstantEnumCase<Arith_CmpIPredicateAttr, "slt">,
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 0da65a4a0f4028..40c5348062c1b7 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -2621,6 +2621,28 @@ func.func @mulsiExtendedOneRhsSplatI1(%arg0: vector<3xi1>) -> (vector<3xi1>, vec
   return %low, %high : vector<3xi1>, vector<3xi1>
 }
 
+// CHECK-LABEL: @mulsiExtendedOneRhsIndex
+//  CHECK-SAME:   (%[[ARG:.+]]: index) -> (index, index)
+//  CHECK-NEXT:   %[[ONE:.+]] = arith.constant 1 : index
+//  CHECK-NEXT:   %[[LOW:.+]], %[[HIGH:.+]] = arith.mulsi_extended %[[ARG]], %[[ONE]] : index
+//  CHECK-NEXT:   return %[[LOW]], %[[HIGH]] : index, index
+func.func @mulsiExtendedOneRhsIndex(%arg0: index) -> (index, index) {
+  %one = arith.constant 1 : index
+  %low, %high = arith.mulsi_extended %arg0, %one: index
+  return %low, %high : index, index
+}
+
+// CHECK-LABEL: @mulsiExtendedOneRhsSplatIndex
+//  CHECK-SAME:   (%[[ARG:.+]]: vector<3xindex>) -> (vector<3xindex>, vector<3xindex>)
+//  CHECK-NEXT:   %[[ONE:.+]] = arith.constant dense<1> : vector<3xindex>
+//  CHECK-NEXT:   %[[LOW:.+]], %[[HIGH:.+]] = arith.mulsi_extended %[[ARG]], %[[ONE]] : vector<3xindex>
+//  CHECK-NEXT:   return %[[LOW]], %[[HIGH]] : vector<3xindex>, vector<3xindex>
+func.func @mulsiExtendedOneRhsSplatIndex(%arg0: vector<3xindex>) -> (vector<3xindex>, vector<3xindex>) {
+  %one = arith.constant dense<1> : vector<3xindex>
+  %low, %high = arith.mulsi_extended %arg0, %one: vector<3xindex>
+  return %low, %high : vector<3xindex>, vector<3xindex>
+}
+
 // CHECK-LABEL: @mulsiExtendedUnusedHigh
 //  CHECK-SAME:   (%[[ARG:.+]]: i32) -> i32
 //  CHECK-NEXT:   %[[RES:.+]] = arith.muli %[[ARG]], %[[ARG]] : i32