[mlir][AffineExpr] Fix UBSan signed-integer overflow in subtraction operators#186155
[mlir][AffineExpr] Fix UBSan signed-integer overflow in subtraction operators#186155joker-eph wants to merge 2 commits into
Conversation
|
@llvm/pr-subscribers-mlir-core @llvm/pr-subscribers-mlir Author: Mehdi Amini (joker-eph) ChangesAffineConstantExpr can hold INT64_MIN. Two sites in AsmPrinter.cpp negated Fix the two printer sites by casting to uint64_t before negating: Fix the parser to accept 9223372036854775808 specifically: the existing Add two test cases to mlir/test/IR/affine-map.mlir:
Fixes #162774 Assisted-by: Claude Code Full diff: https://github.com/llvm/llvm-project/pull/186155.diff 3 Files Affected:
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} : () -> ()
|
2f251d1 to
83a185f
Compare
83a185f to
3dee744
Compare
…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
3dee744 to
7666330
Compare
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 can this be merged? |
AffineExpr::operator-(int64_t)was negating the value with unary-,which is undefined behavior for INT64_MIN. Similarly,
AffineExpr::operator-(AffineExpr)called unaryoperator-()on aconstant AffineExpr, which multiplies by -1 — also overflowing for
INT64_MIN.
Fix both by using unsigned arithmetic for the negation:
operator-(int64_t)now casts touint64_tbefore negating.operator-(AffineExpr)short-circuits constant RHS throughoperator-(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 subtracta constant expression equal to INT64_MIN.
Fixes #162774
Assisted-by: Claude Code