Skip to content

[mlir][AffineExpr] Fix UBSan signed-integer overflow in subtraction operators#186155

Open
joker-eph wants to merge 2 commits into
llvm:mainfrom
joker-eph:fix/issue-162774
Open

[mlir][AffineExpr] Fix UBSan signed-integer overflow in subtraction operators#186155
joker-eph wants to merge 2 commits into
llvm:mainfrom
joker-eph:fix/issue-162774

Conversation

@joker-eph

@joker-eph joker-eph commented Mar 12, 2026

Copy link
Copy Markdown
Contributor

AffineExpr::operator-(int64_t) was negating the value with unary -,
which is undefined behavior for INT64_MIN. Similarly,
AffineExpr::operator-(AffineExpr) called unary operator-() on a
constant AffineExpr, which multiplies by -1 — also overflowing for
INT64_MIN.

Fix both by using unsigned arithmetic for the negation:

  • operator-(int64_t) now casts to uint64_t before negating.
  • operator-(AffineExpr) short-circuits constant RHS through
    operator-(int64_t) instead of the multiply-by-(-1) path.

This correctly round-trips affine maps like d0 - 9223372036854775808
(printed for d0 + INT64_MIN) and fixes any other callers that subtract
a constant expression equal to INT64_MIN.

Fixes #162774

Assisted-by: Claude Code

@llvmbot llvmbot added mlir:core MLIR Core Infrastructure mlir labels Mar 12, 2026
@llvmbot

llvmbot commented Mar 12, 2026

Copy link
Copy Markdown
Member

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

Changes

AffineConstantExpr can hold INT64_MIN. Two sites in AsmPrinter.cpp negated
the stored value with unary - on a signed int64_t, which is undefined
behaviour (signed overflow). A third site in AffineParser.cpp rejected the
literal 9223372036854775808 (2^63) during parsing even though the printer
emits exactly that magnitude when representing INT64_MIN as "d0 - 9223372036854775808",
breaking the round-trip.

Fix the two printer sites by casting to uint64_t before negating:
-(uint64_t)val

Fix the parser to accept 9223372036854775808 specifically: the existing
guard (int64_t)*val < 0 is relaxed to also allow the single value 2^63,
whose bit pattern is INT64_MIN. When this value is consumed, the sub-expr
path in getAffineBinaryOpExpr is updated to compute lhs + (-rhs) via
unsigned negation, avoiding a second UB site.

Add two test cases to mlir/test/IR/affine-map.mlir:

  • A map whose constant folds to INT64_MIN (sum of two -2^62 constants),
    verifying the printer no longer invokes UB.
  • A direct round-trip parse of affine_map<(d0) -> (d0 - 9223372036854775808)>.

Fixes #162774

Assisted-by: Claude Code


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

3 Files Affected:

  • (modified) mlir/lib/AsmParser/AffineParser.cpp (+17-1)
  • (modified) mlir/lib/IR/AsmPrinter.cpp (+5-2)
  • (modified) mlir/test/IR/affine-map.mlir (+13)
diff --git a/mlir/lib/AsmParser/AffineParser.cpp b/mlir/lib/AsmParser/AffineParser.cpp
index 1797611858c06..8d6402cc86c01 100644
--- a/mlir/lib/AsmParser/AffineParser.cpp
+++ b/mlir/lib/AsmParser/AffineParser.cpp
@@ -25,6 +25,7 @@
 #include "llvm/Support/raw_ostream.h"
 #include <cassert>
 #include <cstdint>
+#include <limits>
 #include <utility>
 
 using namespace mlir;
@@ -163,6 +164,14 @@ AffineExpr AffineParser::getAffineBinaryOpExpr(AffineLowPrecOp op,
   case AffineLowPrecOp::Add:
     return lhs + rhs;
   case AffineLowPrecOp::Sub:
+    // When rhs is a constant, compute lhs + (-rhs) directly using unsigned
+    // negation to avoid signed integer overflow for INT64_MIN. This is needed
+    // to correctly round-trip expressions like "d0 - 9223372036854775808",
+    // where the printer emits that form for "d0 + INT64_MIN".
+    if (auto constRhs = dyn_cast<AffineConstantExpr>(rhs)) {
+      int64_t negated = static_cast<int64_t>(-(uint64_t)constRhs.getValue());
+      return lhs + builder.getAffineConstantExpr(negated);
+    }
     return lhs - rhs;
   case AffineLowPrecOp::LNoOp:
     llvm_unreachable("can't create affine expression for null low prec op");
@@ -345,7 +354,14 @@ AffineExpr AffineParser::parseSymbolSSAIdExpr() {
 ///   affine-expr ::= integer-literal
 AffineExpr AffineParser::parseIntegerExpr() {
   auto val = getToken().getUInt64IntegerValue();
-  if (!val.has_value() || (int64_t)*val < 0)
+  // Allow 9223372036854775808 (= 2^63 = |INT64_MIN|) because the printer
+  // emits it as the magnitude in "... - 9223372036854775808" to represent
+  // affine expressions containing INT64_MIN (e.g. "d0 + INT64_MIN" is
+  // printed as "d0 - 9223372036854775808"). The cast to int64_t yields
+  // INT64_MIN, which is the correct internal representation.
+  if (!val.has_value() ||
+      ((int64_t)*val < 0 &&
+       *val != (uint64_t)std::numeric_limits<int64_t>::min()))
     return emitError("constant too large for index"), nullptr;
 
   consumeToken(Token::integer);
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index b3242f838fc1d..49dc8c0cc9388 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -3234,7 +3234,9 @@ void AsmPrinter::Impl::printAffineExprInternal(
           os << " - ";
           printAffineExprInternal(rhs.getLHS(), BindingStrength::Strong,
                                   printValueName);
-          os << " * " << -rrhs.getValue();
+          // Use unsigned negation to avoid signed integer overflow for
+          // INT64_MIN.
+          os << " * " << -(uint64_t)rrhs.getValue();
           if (enclosingTightness == BindingStrength::Strong)
             os << ')';
           return;
@@ -3247,7 +3249,8 @@ void AsmPrinter::Impl::printAffineExprInternal(
   if (auto rhsConst = dyn_cast<AffineConstantExpr>(rhsExpr)) {
     if (rhsConst.getValue() < 0) {
       printAffineExprInternal(lhsExpr, BindingStrength::Weak, printValueName);
-      os << " - " << -rhsConst.getValue();
+      // Use unsigned negation to avoid signed integer overflow for INT64_MIN.
+      os << " - " << -(uint64_t)rhsConst.getValue();
       if (enclosingTightness == BindingStrength::Strong)
         os << ')';
       return;
diff --git a/mlir/test/IR/affine-map.mlir b/mlir/test/IR/affine-map.mlir
index 86bdaafd79f32..5a3f30536d820 100644
--- a/mlir/test/IR/affine-map.mlir
+++ b/mlir/test/IR/affine-map.mlir
@@ -448,3 +448,16 @@ func.func private @f56(memref<1x1xi8, #map56>)
 
 // CHECK: "f69"() {map = #map{{[0-9]*}}} : () -> ()
 "f69"() {map = #map69} : () -> ()
+
+// Test that affine expressions with INT64_MIN as a constant are printed
+// without signed integer overflow (printed as "d0 - 9223372036854775808")
+// and can be parsed back (round-trip). The value INT64_MIN arises when
+// affine simplification sums two large negative constants.
+#map_int64min = affine_map<(d0) -> (d0 + (-4611686018427387904) + (-4611686018427387904))>
+// CHECK: "f_int64min"() {map = #map{{[0-9]*}}} : () -> ()
+"f_int64min"() {map = #map_int64min} : () -> ()
+
+// Round-trip: parse a map already containing "d0 - 9223372036854775808".
+#map_int64min_rt = affine_map<(d0) -> (d0 - 9223372036854775808)>
+// CHECK: "f_int64min_rt"() {map = #map{{[0-9]*}}} : () -> ()
+"f_int64min_rt"() {map = #map_int64min_rt} : () -> ()

Comment thread mlir/lib/AsmParser/AffineParser.cpp Outdated
@joker-eph joker-eph changed the title [mlir] Fix UBSan signed overflow when printing/parsing affine INT64_MIN constants [mlir][AffineExpr] Fix UBSan signed-integer overflow in subtraction operators Mar 16, 2026
@joker-eph joker-eph requested a review from jsetoain March 16, 2026 12:06
…perators

`AffineExpr::operator-(int64_t)` was negating the value with unary `-`,
which is undefined behavior for INT64_MIN. Similarly,
`AffineExpr::operator-(AffineExpr)` called unary `operator-()` on a
constant AffineExpr, which multiplies by -1 — also overflowing for
INT64_MIN.

Fix both by using unsigned arithmetic for the negation:
- `operator-(int64_t)` now casts to `uint64_t` before negating.
- `operator-(AffineExpr)` short-circuits constant RHS through
  `operator-(int64_t)` instead of the multiply-by-(-1) path.

This correctly round-trips affine maps like `d0 - 9223372036854775808`
(printed for `d0 + INT64_MIN`) and fixes any other callers that subtract
a constant expression equal to INT64_MIN.

Fixes llvm#162774

Assisted-by: Claude Code
Comment thread mlir/test/IR/affine-map.mlir Outdated
Address reviewer feedback: add a CHECK line that verifies the map
expression content (`(d0) -> (d0 - 9223372036854775808)`) rather
than just matching the generic `#map{{[0-9]*}}` reference.  Also note
that the two equivalent maps are deduplicated.

Assisted-by: Claude Code
@joker-eph joker-eph requested a review from bondhugula March 24, 2026 10:26

@bondhugula bondhugula left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM.

Comment thread mlir/test/IR/affine-map.mlir Outdated

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Likewise.

@ayokunle321

Copy link
Copy Markdown
Contributor

@joker-eph can this be merged?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

mlir:core MLIR Core Infrastructure mlir

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[mlir] UBSan alert mlir/lib/IR/AsmPrinter.cpp

5 participants