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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-arith

Author: Prakhar Dixit (Prakhar-Dixit)

<details>
<summary>Changes</summary>

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)

---
Full diff: https://github.com/llvm/llvm-project/pull/233234.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td (+2-1) 
- (modified) mlir/test/Dialect/Arith/canonicalize.mlir (+22) 


``````````diff
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

``````````

</details>


https://github.com/llvm/llvm-project/pull/233234