[Mlir-commits] [mlir] [MLIR][Index] Restrict subtraction comparison canonicalization (PR #216831)

[llvmlistbot at llvm.org]

https://github.com/at0m741 created https://github.com/llvm/llvm-project/pull/216831

Fix the unsafe canonicalization of ordered comparisons involving `index.sub`.

Ordered comparisons may change semantics when the subtraction wraps, so the
rewrite is now restricted to `eq` and `ne`.

Regression tests cover ordered comparisons while preserving the valid `eq` and
`ne` canonicalizations.

Fixes #216042 

>From 420f71a7ecb866f0f1cd643c41768bd164feca80 Mon Sep 17 00:00:00 2001
From: at0m741 <louistouz at icloud.com>
Date: Mon, 17 Aug 2026 22:33:51 +0200
Subject: [PATCH] [MLIR][Index] Restrict subtraction cmp canonicalization

---
 mlir/lib/Dialect/Index/IR/IndexOps.cpp        | 16 ++++++--
 .../Dialect/Index/index-canonicalize.mlir     | 38 +++++++++++++++++--
 2 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Dialect/Index/IR/IndexOps.cpp b/mlir/lib/Dialect/Index/IR/IndexOps.cpp
index 2b1baa8b44643..b6863befb3fbd 100644
--- a/mlir/lib/Dialect/Index/IR/IndexOps.cpp
+++ b/mlir/lib/Dialect/Index/IR/IndexOps.cpp
@@ -693,9 +693,11 @@ OpFoldResult CmpOp::fold(FoldAdaptor adaptor) {
   return {};
 }
 
-/// Canonicalize
-/// `x - y cmp 0` to `x cmp y`. or `x - y cmp 0` to `x cmp y`.
-/// `0 cmp x - y` to `y cmp x`. or `0 cmp x - y` to `y cmp x`.
+/// Canonicalize `(x - y) == 0` to `x == y` and `(x - y) != 0` to `x != y`.
+/// Likewise, canonicalize `0 == (x - y)` and `0 != (x - y)`.
+///
+/// Ordered comparisons cannot be canonicalized this way because subtraction
+/// may wrap at the target index bitwidth.
 LogicalResult CmpOp::canonicalize(CmpOp op, PatternRewriter &rewriter) {
   IntegerAttr cmpRhs;
   IntegerAttr cmpLhs;
@@ -707,12 +709,19 @@ LogicalResult CmpOp::canonicalize(CmpOp op, PatternRewriter &rewriter) {
   if (!rhsIsZero && !lhsIsZero)
     return rewriter.notifyMatchFailure(op.getLoc(),
                                        "cmp is not comparing something with 0");
+
   SubOp subOp = rhsIsZero ? op.getLhs().getDefiningOp<index::SubOp>()
                           : op.getRhs().getDefiningOp<index::SubOp>();
   if (!subOp)
     return rewriter.notifyMatchFailure(
         op.getLoc(), "non-zero operand is not a result of subtraction");
 
+  if (op.getPred() != IndexCmpPredicate::EQ &&
+      op.getPred() != IndexCmpPredicate::NE)
+    return rewriter.notifyMatchFailure(
+        op.getLoc(),
+        "only eq and ne comparisons can be canonicalized through subtraction");
+
   index::CmpOp newCmp;
   if (rhsIsZero)
     newCmp = index::CmpOp::create(rewriter, op.getLoc(), op.getPred(),
@@ -720,6 +729,7 @@ LogicalResult CmpOp::canonicalize(CmpOp op, PatternRewriter &rewriter) {
   else
     newCmp = index::CmpOp::create(rewriter, op.getLoc(), op.getPred(),
                                   subOp.getRhs(), subOp.getLhs());
+
   rewriter.replaceOp(op, newCmp);
   return success();
 }
diff --git a/mlir/test/Dialect/Index/index-canonicalize.mlir b/mlir/test/Dialect/Index/index-canonicalize.mlir
index 45da6ea57d796..06fc9abba8bcc 100644
--- a/mlir/test/Dialect/Index/index-canonicalize.mlir
+++ b/mlir/test/Dialect/Index/index-canonicalize.mlir
@@ -609,13 +609,43 @@ func.func @cmp(%arg0: index) -> (i1, i1, i1, i1, i1, i1) {
 
   // CHECK-DAG: %[[TRUE:.*]] = index.bool.constant true
   // CHECK-DAG: %[[FALSE:.*]] = index.bool.constant false
-  // CHECK-DAG: [[IDX0:%.*]] = index.constant 0
-  // CHECK-DAG: [[V4:%.*]] = index.cmp sgt([[IDX0]], %arg0)
-  // CHECK-DAG: [[V5:%.*]] = index.cmp sgt(%arg0, [[IDX0]])
-  // CHECK: return %[[FALSE]], %[[TRUE]], %[[TRUE]], %[[FALSE]]
+  // CHECK-DAG: %[[IDX0:.*]] = index.constant 0
+  // CHECK: %[[SUB0:.*]] = index.sub %[[IDX0]], %arg0
+  // CHECK-NEXT: %[[V4:.*]] = index.cmp sgt(%[[SUB0]], %[[IDX0]])
+  // CHECK-NEXT: %[[SUB1:.*]] = index.sub %[[IDX0]], %arg0
+  // CHECK-NEXT: %[[V5:.*]] = index.cmp sgt(%[[IDX0]], %[[SUB1]])
+  // CHECK-NEXT: return %[[FALSE]], %[[TRUE]], %[[TRUE]], %[[FALSE]], %[[V4]], %[[V5]]
   return %0, %1, %2, %3, %5, %7 : i1, i1, i1, i1, i1, i1
 }
 
+// CHECK-LABEL: @cmp_sub_ordered
+func.func @cmp_sub_ordered(%arg0: index) -> i1 {
+  %zero = index.constant 0
+  %two = index.constant 2
+  %sub = index.sub %arg0, %two
+  %cmp = index.cmp ule(%sub, %zero)
+
+  // CHECK-DAG: %[[ZERO:.*]] = index.constant 0
+  // CHECK-DAG: %[[TWO:.*]] = index.constant 2
+  // CHECK: %[[SUB:.*]] = index.sub %arg0, %[[TWO]]
+  // CHECK-NEXT: %[[CMP:.*]] = index.cmp ule(%[[SUB]], %[[ZERO]])
+  // CHECK-NEXT: return %[[CMP]]
+  return %cmp : i1
+}
+
+// CHECK-LABEL: @cmp_sub_eq_ne
+func.func @cmp_sub_eq_ne(%arg0: index, %arg1: index) -> (i1, i1) {
+  %zero = index.constant 0
+  %sub = index.sub %arg0, %arg1
+  %eq = index.cmp eq(%sub, %zero)
+  %ne = index.cmp ne(%zero, %sub)
+
+  // CHECK-DAG: %[[EQ:.*]] = index.cmp eq(%arg0, %arg1)
+  // CHECK-DAG: %[[NE:.*]] = index.cmp ne(%arg1, %arg0)
+  // CHECK: return %[[EQ]], %[[NE]]
+  return %eq, %ne : i1, i1
+}
+
 // CHECK-LABEL: @cmp_same_args
 func.func @cmp_same_args(%a: index) -> (i1, i1, i1, i1, i1, i1, i1, i1, i1, i1) {
   %0 = index.cmp eq(%a, %a)