You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/10/22 23:17:59 UTC
(tvm-tirx) branch main updated: [FIX] Fold constants only when results fit their types (#26)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 42abbce7 [FIX] Fold constants only when results fit their types (#26)
42abbce7 is described below
commit 42abbce7d84759563f328fa6b7baced97fba4f33
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Thu Oct 22 19:17:53 2026 -0400
[FIX] Fold constants only when results fit their types (#26)
---
include/tvm/ir/prim/const_fold.h | 135 +++++++++++++++------------
include/tvm/ir/prim/op.h | 3 +-
python/tvm_tirx/ir/_constant.py | 6 ++
python/tvm_tirx/ir/type.py | 6 +-
src/ir/prim/op.cc | 106 ++++++++++++++++-----
src/support/limits.h | 2 +-
src/sym/analyzer.cc | 3 +-
src/sym/canonical_simplify.cc | 6 ++
src/sym/const_fold.h | 1 -
src/sym/rewrite_simplify.cc | 12 +--
tests/python/tirx/base/test_imm_values.py | 89 ++++++++++++------
tests/python/tirx/transform/test_simplify.py | 4 +-
12 files changed, 243 insertions(+), 130 deletions(-)
diff --git a/include/tvm/ir/prim/const_fold.h b/include/tvm/ir/prim/const_fold.h
index 8fcb7c17..83398094 100644
--- a/include/tvm/ir/prim/const_fold.h
+++ b/include/tvm/ir/prim/const_fold.h
@@ -86,45 +86,39 @@ inline bool IsIndexTypedExpr(const ExprNode* expr) {
inline bool IsIndexTypedExpr(const PrimExpr& expr) { return IsIndexTypedExpr(expr.get()); }
-/*! \brief Normalize exact arithmetic to the result's integer width. */
-inline ffi::BigInt GetFoldResult(ffi::BigInt value, const PrimType& dtype) {
- bool is_signed = dtype.MatchesCode(DLDataTypeCode::kDLInt);
- if (auto small = value.as<int64_t>(); small.has_value()) {
- if (dtype.bits() >= 64 && (is_signed || *small >= 0)) return value;
- if (dtype.bits() < 64) {
- uint64_t bits = static_cast<uint64_t>(*small) & ((uint64_t{1} << dtype.bits()) - 1);
- int64_t result = static_cast<int64_t>(bits);
- if (is_signed) {
- int64_t sign_bit = int64_t{1} << (dtype.bits() - 1);
- result = (result ^ sign_bit) - sign_bit;
- }
- return ffi::BigInt(result);
+/*! \brief Whether an exact integer is a valid literal of the given scalar type. */
+inline bool IsRepresentable(const ffi::BigInt& value, const PrimType& dtype) {
+ if (!dtype.IsScalar() || dtype.bits() <= 0) return false;
+ if (auto small = value.as<int64_t>()) {
+ if (dtype.MatchesCode(DLDataTypeCode::kDLUInt)) {
+ return *small >= 0 &&
+ (dtype.bits() >= 63 || static_cast<uint64_t>(*small) < (uint64_t{1} << dtype.bits()));
+ }
+ if (dtype.MatchesCode(DLDataTypeCode::kDLBool) ||
+ (dtype.MatchesCode(DLDataTypeCode::kDLInt) && dtype.bits() == 1)) {
+ return *small == 0 || *small == 1;
}
+ if (dtype.MatchesCode(DLDataTypeCode::kDLInt)) {
+ if (dtype.bits() >= 64) return true;
+ int64_t limit = int64_t{1} << (dtype.bits() - 1);
+ return *small >= -limit && *small < limit;
+ }
+ return false;
}
- ffi::BigInt modulus = ffi::BigInt(1) << dtype.bits();
- value &= modulus - 1;
- if (is_signed && value >= (modulus >> 1)) value -= modulus;
- return value;
-}
-
-/*! \brief Helper to get fp32 const folding result repr in double. */
-inline double GetFoldResultDoubleRepr(float x) {
- double res = static_cast<double>(x);
- if (std::isinf(res) || std::isnan(res)) {
- return res;
+ if (dtype.MatchesCode(DLDataTypeCode::kDLUInt)) {
+ return value >= 0 && value < (ffi::BigInt(1) << dtype.bits());
}
- // certain platform (eg, on gcc7-i386) do the folding arithmetic
- // on float and write back to double is optimized to double
- // precision arithmetic, this is legal and we check the output
- // range thus to ensure consistency when the float result is inf.
- if (res < std::numeric_limits<float>::lowest()) {
- LOG(WARNING) << "underlying float value overflow";
- return -std::numeric_limits<double>::infinity();
- } else if (res > std::numeric_limits<float>::max()) {
- LOG(WARNING) << "underlying float value overflow";
- return std::numeric_limits<double>::infinity();
+ if (dtype.MatchesCode(DLDataTypeCode::kDLInt) && dtype.bits() > 1) {
+ ffi::BigInt limit = ffi::BigInt(1) << (dtype.bits() - 1);
+ return value >= -limit && value < limit;
}
- return res;
+ return false;
+}
+
+/*! \brief Check float32 range, allowing rounding and supported special values. */
+inline bool IsInFloat32Range(double value) {
+ return !std::isfinite(value) || (value >= std::numeric_limits<float>::lowest() &&
+ value <= std::numeric_limits<float>::max());
}
} // namespace details
@@ -151,14 +145,18 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Add>(PrimExpr a, PrimExpr b) {
PrimType result_ty = a.ty();
if (pa && pb) {
ffi::BigInt res = pa->value + pb->value;
- return IntImm(result_ty, details::GetFoldResult(res, result_ty));
+ if (!details::IsRepresentable(res, result_ty)) return std::nullopt;
+ return IntImm(result_ty, std::move(res));
}
if (pa && pa->value == 0) return b;
if (pb && pb->value == 0) return a;
if (fa && fb) {
if (result_ty.bits() == 32) {
- return FloatImm(result_ty, details::GetFoldResultDoubleRepr(static_cast<float>(fa->value) +
- static_cast<float>(fb->value)));
+ float lhs = static_cast<float>(fa->value);
+ float rhs = static_cast<float>(fb->value);
+ // Check before rounding can turn a finite overflow into infinity.
+ if (!details::IsInFloat32Range(static_cast<double>(lhs) + rhs)) return std::nullopt;
+ return FloatImm(result_ty, static_cast<float>(lhs + rhs));
} else if (result_ty.bits() == 64) {
return FloatImm(result_ty, fa->value + fb->value);
}
@@ -172,22 +170,20 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Add>(PrimExpr a, PrimExpr b) {
template <>
inline ffi::Optional<PrimExpr> TryConstFold<prim::Sub>(PrimExpr a, PrimExpr b) {
TVM_PRIM_CONST_PROPAGATION({
- TVM_FFI_ICHECK(!((pa && pa->ty.as_or_throw<PrimType>().MatchesCode(DLDataTypeCode::kDLUInt) &&
- pa->value == 0U) &&
- (pb && pb->ty.as_or_throw<PrimType>().MatchesCode(DLDataTypeCode::kDLUInt) &&
- pb->value > 0U)))
- << "Checked failed. Minuend 's value is 0U and it's dtype is uint "
- << "while Subtrahend's dtype is uint; which will cause a negative uint";
PrimType result_ty = a.ty();
if (pa && pb) {
ffi::BigInt res = pa->value - pb->value;
- return IntImm(result_ty, details::GetFoldResult(res, result_ty));
+ if (!details::IsRepresentable(res, result_ty)) return std::nullopt;
+ return IntImm(result_ty, std::move(res));
}
if (pb && pb->value == 0) return a;
if (fa && fb) {
if (result_ty.bits() == 32) {
- return FloatImm(result_ty, details::GetFoldResultDoubleRepr(static_cast<float>(fa->value) -
- static_cast<float>(fb->value)));
+ float lhs = static_cast<float>(fa->value);
+ float rhs = static_cast<float>(fb->value);
+ // Check before rounding can turn a finite overflow into infinity.
+ if (!details::IsInFloat32Range(static_cast<double>(lhs) - rhs)) return std::nullopt;
+ return FloatImm(result_ty, static_cast<float>(lhs - rhs));
} else if (result_ty.bits() == 64) {
return FloatImm(result_ty, fa->value - fb->value);
}
@@ -203,7 +199,8 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Mul>(PrimExpr a, PrimExpr b) {
PrimType result_ty = a.ty();
if (pa && pb) {
ffi::BigInt res = pa->value * pb->value;
- return IntImm(result_ty, details::GetFoldResult(res, result_ty));
+ if (!details::IsRepresentable(res, result_ty)) return std::nullopt;
+ return IntImm(result_ty, std::move(res));
}
if (pa) {
if (pa->value == 1) return b;
@@ -215,8 +212,11 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Mul>(PrimExpr a, PrimExpr b) {
}
if (fa && fb) {
if (result_ty.bits() == 32) {
- return FloatImm(result_ty, details::GetFoldResultDoubleRepr(static_cast<float>(fa->value) *
- static_cast<float>(fb->value)));
+ float lhs = static_cast<float>(fa->value);
+ float rhs = static_cast<float>(fb->value);
+ // Check before rounding can turn a finite overflow into infinity.
+ if (!details::IsInFloat32Range(static_cast<double>(lhs) * rhs)) return std::nullopt;
+ return FloatImm(result_ty, static_cast<float>(lhs * rhs));
} else if (result_ty.bits() == 64) {
return FloatImm(result_ty, fa->value * fb->value);
}
@@ -242,7 +242,8 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Div>(PrimExpr a, PrimExpr b) {
// NOTE: this will assumes truc div.
TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
ffi::BigInt res = pa->value / pb->value;
- return IntImm(result_ty, details::GetFoldResult(res, result_ty));
+ if (!details::IsRepresentable(res, result_ty)) return std::nullopt;
+ return IntImm(result_ty, std::move(res));
}
if (pa) {
if (pa->value == 0) return a;
@@ -254,8 +255,11 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Div>(PrimExpr a, PrimExpr b) {
if (fa && fb) {
TVM_FFI_ICHECK_NE(fb->value, 0) << "Divide by zero";
if (result_ty.bits() == 32) {
- return FloatImm(result_ty, details::GetFoldResultDoubleRepr(static_cast<float>(fa->value) /
- static_cast<float>(fb->value)));
+ float lhs = static_cast<float>(fa->value);
+ float rhs = static_cast<float>(fb->value);
+ // Check before rounding can turn a finite overflow into infinity.
+ if (!details::IsInFloat32Range(static_cast<double>(lhs) / rhs)) return std::nullopt;
+ return FloatImm(result_ty, static_cast<float>(lhs / rhs));
} else if (result_ty.bits() == 64) {
return FloatImm(result_ty, fa->value / fb->value);
}
@@ -276,7 +280,8 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Mod>(PrimExpr a, PrimExpr b) {
if (pa && pb) {
TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
ffi::BigInt res = pa->value % pb->value;
- return IntImm(result_ty, details::GetFoldResult(res, result_ty));
+ if (!details::IsRepresentable(res, result_ty)) return std::nullopt;
+ return IntImm(result_ty, std::move(res));
}
if (pa) {
if (pa->value == 0) return a;
@@ -289,7 +294,9 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::Mod>(PrimExpr a, PrimExpr b) {
// Fold wide literals without enabling symbolic index rewrites for their type.
if (pa && pb && a.ty().bits() >= 64) {
TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
- return IntImm(a.ty(), details::GetFoldResult(pa->value % pb->value, a.ty()));
+ ffi::BigInt res = pa->value % pb->value;
+ if (!details::IsRepresentable(res, a.ty())) return std::nullopt;
+ return IntImm(a.ty(), std::move(res));
}
return std::nullopt;
}
@@ -301,7 +308,8 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::FloorDiv>(PrimExpr a, PrimExpr
if (pa && pb) {
TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
ffi::BigInt res = ffi::floordiv(pa->value, pb->value);
- return IntImm(result_ty, details::GetFoldResult(res, result_ty));
+ if (!details::IsRepresentable(res, result_ty)) return std::nullopt;
+ return IntImm(result_ty, std::move(res));
}
if (pa) {
if (pa->value == 0) return a;
@@ -312,9 +320,11 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::FloorDiv>(PrimExpr a, PrimExpr
}
if (fa && fb && fb->value != 0) {
if (result_ty.bits() == 32) {
- return FloatImm(
- result_ty, details::GetFoldResultDoubleRepr(std::floor(static_cast<float>(fa->value) /
- static_cast<float>(fb->value))));
+ float lhs = static_cast<float>(fa->value);
+ float rhs = static_cast<float>(fb->value);
+ if (!details::IsInFloat32Range(static_cast<double>(lhs) / rhs)) return std::nullopt;
+ float quotient = lhs / rhs;
+ return FloatImm(result_ty, std::floor(quotient));
} else if (result_ty.bits() == 64) {
return FloatImm(result_ty, std::floor(fa->value / fb->value));
} else {
@@ -337,7 +347,8 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::FloorMod>(PrimExpr a, PrimExpr
if (pa && pb) {
TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
ffi::BigInt res = ffi::floormod(pa->value, pb->value);
- return IntImm(result_ty, details::GetFoldResult(res, result_ty));
+ if (!details::IsRepresentable(res, result_ty)) return std::nullopt;
+ return IntImm(result_ty, std::move(res));
}
if (pa) {
if (pa->value == 0) return a;
@@ -350,7 +361,9 @@ inline ffi::Optional<PrimExpr> TryConstFold<prim::FloorMod>(PrimExpr a, PrimExpr
// Fold wide literals without enabling symbolic index rewrites for their type.
if (pa && pb && (a.ty().bits() >= 64 || a.ty().MatchesCode(DLDataTypeCode::kDLUInt))) {
TVM_FFI_ICHECK_NE(pb->value, 0) << "Divide by zero";
- return IntImm(a.ty(), details::GetFoldResult(ffi::floormod(pa->value, pb->value), a.ty()));
+ ffi::BigInt res = ffi::floormod(pa->value, pb->value);
+ if (!details::IsRepresentable(res, a.ty())) return std::nullopt;
+ return IntImm(a.ty(), std::move(res));
}
return std::nullopt;
}
diff --git a/include/tvm/ir/prim/op.h b/include/tvm/ir/prim/op.h
index ff8563f7..2f976c7b 100644
--- a/include/tvm/ir/prim/op.h
+++ b/include/tvm/ir/prim/op.h
@@ -751,7 +751,8 @@ TVM_TIRX_DLL PrimExpr max_value(PrimType dtype, Location loc = UnknownLoc());
* \param loc The location of this operation in the source.
* \return The result expression.
* \note This function may return value if the type is the same. Constant folding
- * uses MakeConst for scalar and vector constants.
+ * is limited to supported conversions whose results fit the destination type.
+ * Use MakeConst to construct typed literals directly.
*/
TVM_TIRX_DLL PrimExpr cast(PrimExpr value, PrimType t, Location loc = UnknownLoc());
/*! \brief Construct integer absolute value; floating absolute value belongs to TIRX. */
diff --git a/python/tvm_tirx/ir/_constant.py b/python/tvm_tirx/ir/_constant.py
index e38e57e4..246f5754 100644
--- a/python/tvm_tirx/ir/_constant.py
+++ b/python/tvm_tirx/ir/_constant.py
@@ -84,4 +84,10 @@ def const(value, ty=None, loc: Location = UNKNOWN_LOC, *, dtype=None) -> "Expr":
ty = "float32"
else:
raise TypeError(f"Cannot infer a constant type from {type(value).__name__}")
+ if isinstance(value, np.bool_):
+ value = bool(value)
+ elif isinstance(value, np.integer):
+ value = int(value)
+ elif isinstance(value, np.floating):
+ value = float(value)
return _ffi_api._const(value, ty, loc)
diff --git a/python/tvm_tirx/ir/type.py b/python/tvm_tirx/ir/type.py
index f81555e4..e742d907 100644
--- a/python/tvm_tirx/ir/type.py
+++ b/python/tvm_tirx/ir/type.py
@@ -22,6 +22,7 @@ from collections.abc import Callable
from types import ModuleType
from typing import ClassVar
+import numpy as np
import tvm_ffi
from tvm_tirx.ir.location import UNKNOWN_LOC, Location
@@ -122,7 +123,8 @@ class PrimType(Type):
self.__init_handle_by_constructor__(_ffi_api.PrimType, dtype)
def __call__(self, expr=None):
- """Create a fresh variable, or convert a primitive expression to this type."""
+ """Create a variable, construct a typed literal, or cast an expression."""
+ from ._constant import const
from .expr import Var
from .prim import _ffi_api as prim_ffi_api
@@ -130,6 +132,8 @@ class PrimType(Type):
return Var("", self)
if isinstance(expr, str) and str(self.dtype) not in ("bool", "void", ""):
expr = float(expr)
+ if isinstance(expr, int | float | np.integer | np.floating | np.bool_):
+ return const(expr, self, UNKNOWN_LOC)
return prim_ffi_api._cast(expr, self, UNKNOWN_LOC)
def __getitem__(self, shape):
diff --git a/src/ir/prim/op.cc b/src/ir/prim/op.cc
index 8cd61e1b..226882e4 100644
--- a/src/ir/prim/op.cc
+++ b/src/ir/prim/op.cc
@@ -22,6 +22,7 @@
#include <cmath>
+#include "../../support/limits.h"
#include "op_utils.h"
namespace tvm {
@@ -214,7 +215,7 @@ PrimExpr max_value(PrimType value_ty, Location loc) {
return FloatImm(value_ty, 65504.0, loc);
}
} else if (IsBFloat16Type(dtype)) {
- return FloatImm(value_ty, std::numeric_limits<float>::max(), loc);
+ return FloatImm(value_ty, support::kMaxBFloat16, loc);
} else if (IsFloat8Type(dtype)) {
// according to https://arxiv.org/pdf/2209.05433.pdf
if (dtype.code() == DLDataTypeCode::kDLFloat8_e5m2) {
@@ -268,7 +269,7 @@ PrimExpr min_value(PrimType value_ty, Location loc) {
return FloatImm(value_ty, -65504.0, loc);
}
} else if (IsBFloat16Type(dtype)) {
- return FloatImm(value_ty, std::numeric_limits<float>::lowest(), loc);
+ return FloatImm(value_ty, -support::kMaxBFloat16, loc);
} else if (IsFloat8Type(dtype)) {
// according to https://arxiv.org/pdf/2209.05433.pdf
if (dtype.code() == DLDataTypeCode::kDLFloat8_e5m2) {
@@ -309,24 +310,69 @@ PrimExpr cast(PrimExpr value, PrimType t, Location loc) {
// const fold IntImm as they are used in index computations
if (dtype.IsScalar()) {
if (const IntImmNode* op = value.as<IntImmNode>()) {
- return MakeConst(dtype, op->value, op->loc);
+ if (dtype.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt,
+ DLDataTypeCode::kDLBool)) {
+ if (details::IsRepresentable(op->value, dtype)) return IntImm(dtype, op->value, op->loc);
+ } else if (dtype.MatchesElementType(DLDataTypeCode::kDLFloat, 32)) {
+ if (auto small = op->value.as<int64_t>()) {
+ return FloatImm(dtype, static_cast<float>(*small), op->loc);
+ }
+ } else if (dtype.MatchesElementType(DLDataTypeCode::kDLFloat, 16) ||
+ dtype.MatchesElementType(DLDataTypeCode::kDLBfloat, 16)) {
+ double source = static_cast<double>(op->value);
+ double limit = dtype.MatchesCode(DLDataTypeCode::kDLBfloat) ? support::kMaxBFloat16
+ : support::kMaxFloat16;
+ // Keep the integer exact until the destination performs its rounding.
+ if (source >= -limit && source <= limit && ffi::BigInt(source) == op->value) {
+ return FloatImm(dtype, source, op->loc);
+ }
+ }
} else if (const FloatImmNode* op = value.as<FloatImmNode>()) {
- return MakeConst(dtype, op->value, op->loc);
+ if (std::isfinite(op->value) &&
+ (value.ty().MatchesElementType(DLDataTypeCode::kDLFloat, 16) ||
+ value.ty().MatchesElementType(DLDataTypeCode::kDLBfloat, 16))) {
+ // FloatImm stores an unrounded double. Preserve the cast when using it
+ // directly would skip rounding in the source's 16-bit floating type.
+ bool is_bfloat = value.ty().MatchesCode(DLDataTypeCode::kDLBfloat);
+ int exponent;
+ std::frexp(op->value, &exponent);
+ int step = std::max(exponent - (is_bfloat ? 8 : 11), is_bfloat ? -133 : -24);
+ double scaled = std::ldexp(op->value, -step);
+ if (std::trunc(scaled) != scaled) return prim::Cast(std::move(t), value, loc);
+ }
+ double source = value.ty().MatchesElementType(DLDataTypeCode::kDLFloat, 32)
+ ? static_cast<float>(op->value)
+ : op->value;
+ if (dtype.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt,
+ DLDataTypeCode::kDLBool)) {
+ // Non-finite values cannot be converted to BigInt. Boolean casts also
+ // need to preserve nonzero testing rather than truncate a fraction.
+ if (std::isfinite(source) &&
+ (!dtype.MatchesCode(DLDataTypeCode::kDLBool) || source == 0 || source == 1)) {
+ ffi::BigInt result(source);
+ if (details::IsRepresentable(result, dtype)) {
+ return IntImm(dtype, std::move(result), op->loc);
+ }
+ }
+ } else if (dtype.MatchesElementType(DLDataTypeCode::kDLFloat, 32)) {
+ if (details::IsInFloat32Range(source)) {
+ return FloatImm(dtype, static_cast<float>(source), op->loc);
+ }
+ } else if (dtype.MatchesElementType(DLDataTypeCode::kDLFloat, 16) ||
+ dtype.MatchesElementType(DLDataTypeCode::kDLBfloat, 16)) {
+ double limit = dtype.MatchesCode(DLDataTypeCode::kDLBfloat) ? support::kMaxBFloat16
+ : support::kMaxFloat16;
+ if (!std::isfinite(source) || (source >= -limit && source <= limit)) {
+ return FloatImm(dtype, source, op->loc);
+ }
+ }
}
return prim::Cast(std::move(t), value, loc);
} else {
PrimType elem_ty = dtype.WithLanes(1);
if (!value.ty().IsScalableVector() && !value.ty().IsFixedLengthVector()) {
// manually unroll cast
- if (value.ty() != elem_ty) {
- if (const IntImmNode* op = value.as<IntImmNode>()) {
- value = MakeConst(elem_ty, op->value, op->loc);
- } else if (const FloatImmNode* op = value.as<FloatImmNode>()) {
- value = MakeConst(elem_ty, op->value, op->loc);
- } else {
- value = prim::Cast(elem_ty, value, loc);
- }
- }
+ value = cast(value, elem_ty, loc);
if (dtype.IsScalableVector()) {
return prim::Broadcast(
value,
@@ -349,10 +395,21 @@ PrimExpr cast(PrimExpr value, PrimType t, Location loc) {
if (const auto* broadcast = value.as<prim::BroadcastNode>()) {
return prim::Broadcast(cast(broadcast->value, elem_ty, loc), broadcast->lanes, loc);
} else if (const auto* ramp = value.as<prim::RampNode>()) {
- if (dtype.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
- // only cast to index data type can be folded to ramp
- return prim::Ramp(cast(ramp->base, elem_ty, loc), cast(ramp->stride, elem_ty, loc),
- ramp->lanes, loc);
+ // Casting a ramp's base and stride is valid only when every lane
+ // remains the same integer in both types. In particular, widening
+ // must not remove source overflow, and floating casts need not be affine.
+ const auto* base = ramp->base.as<IntImmNode>();
+ const auto* stride = ramp->stride.as<IntImmNode>();
+ if (dtype.IsFixedLengthVector() && base && stride &&
+ dtype.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
+ ffi::BigInt last = base->value + stride->value * (dtype.lanes() - 1);
+ if (details::IsRepresentable(base->value, elem_ty) &&
+ details::IsRepresentable(stride->value, elem_ty) &&
+ details::IsRepresentable(last, elem_ty) &&
+ details::IsRepresentable(last, ramp->base.ty())) {
+ return prim::Ramp(cast(ramp->base, elem_ty, loc), cast(ramp->stride, elem_ty, loc),
+ ramp->lanes, loc);
+ }
}
}
return prim::Cast(std::move(t), value, loc);
@@ -400,10 +457,9 @@ PrimExpr neg(PrimExpr a, Location loc) {
const FloatImmNode* fa = a.as<FloatImmNode>();
if (pa) {
ffi::BigInt value = -pa->value;
- if (a.ty().MatchesCode(DLDataTypeCode::kDLInt) && a.ty().bits() >= 64) {
- value = prim::details::GetFoldResult(std::move(value), a.ty());
+ if (prim::details::IsRepresentable(value, a.ty())) {
+ return IntImm(a.ty(), std::move(value), loc);
}
- return IntImm(a.ty(), std::move(value), loc);
}
if (fa) return FloatImm(a.ty(), -fa->value, loc);
return MakeConst(a.ty(), 0, loc) - a;
@@ -651,8 +707,9 @@ PrimExpr left_shift(PrimExpr a, PrimExpr b, Location loc) {
<< result_ty;
if (pa && pb) {
ffi::BigInt value = pa->value << pb->value;
- if (result_ty.bits() >= 64) value = prim::details::GetFoldResult(std::move(value), result_ty);
- return IntImm(result_ty, std::move(value), loc);
+ if (prim::details::IsRepresentable(value, result_ty)) {
+ return IntImm(result_ty, std::move(value), loc);
+ }
}
if (pb) {
if (pb->value == 0) return a;
@@ -733,8 +790,9 @@ PrimExpr prim::IntegerAbs(PrimExpr x, Location loc) {
if (x.ty().MatchesCode(DLDataTypeCode::kDLInt)) {
if (const IntImmNode* px = x.as<IntImmNode>()) {
ffi::BigInt value = px->value < 0 ? -px->value : px->value;
- if (x.ty().bits() >= 64) value = prim::details::GetFoldResult(std::move(value), x.ty());
- return IntImm(x.ty(), std::move(value), px->loc);
+ if (prim::details::IsRepresentable(value, x.ty())) {
+ return IntImm(x.ty(), std::move(value), px->loc);
+ }
}
return Select(x >= MakeConst(x.ty(), 0), x, -x, loc);
}
diff --git a/src/support/limits.h b/src/support/limits.h
index eb886432..09e71c2d 100644
--- a/src/support/limits.h
+++ b/src/support/limits.h
@@ -33,7 +33,7 @@ constexpr double kMaxFloat16 = 65504.0;
// 2^127 * (1 + 127/128)
// See https://en.wikipedia.org/wiki/Bfloat16_floating-point_format
-constexpr double kMaxBFloat16 = 3.895313892515354759047080037148786688e38;
+constexpr double kMaxBFloat16 = 0x1.fep127;
// 2^15 * (1 + 3/4)
// See https://arxiv.org/pdf/2209.05433.pdf
diff --git a/src/sym/analyzer.cc b/src/sym/analyzer.cc
index 6183fb72..51e18649 100644
--- a/src/sym/analyzer.cc
+++ b/src/sym/analyzer.cc
@@ -187,7 +187,8 @@ bool AnalyzerObj::CanProveEqual(const PrimExpr& lhs, const PrimExpr& rhs) {
prim::is_neg_inf(rhs)) {
return lhs.same_as(rhs);
}
- return CanProve(lhs - rhs == 0);
+ PrimExpr difference = lhs - rhs;
+ return CanProve(difference == prim::MakeConst(difference.ty(), 0));
}
bool AnalyzerObj::CanProveLessEqualThanSymbolicShapeValue(const PrimExpr& lhs,
diff --git a/src/sym/canonical_simplify.cc b/src/sym/canonical_simplify.cc
index afbb1538..948e9495 100644
--- a/src/sym/canonical_simplify.cc
+++ b/src/sym/canonical_simplify.cc
@@ -774,6 +774,8 @@ UnchangedOr<PrimExpr> CanonicalSimplifier::Impl::Mutate_(const prim::AddNode* op
// const folding
if (auto const_res = TryConstFold<prim::Add>(a, b)) return *std::move(const_res);
+ // A declined literal fold must not be materialized through canonical coefficients.
+ if (a.as<IntImmNode>() && b.as<IntImmNode>()) return prim::Add(a, b);
// canonical form simplification.
SumExpr ret = ToSumExpr(std::move(a));
@@ -799,6 +801,8 @@ UnchangedOr<PrimExpr> CanonicalSimplifier::Impl::Mutate_(const prim::SubNode* op
// const folding
if (auto const_res = TryConstFold<prim::Sub>(a, b)) return *std::move(const_res);
+ // A declined literal fold must not be materialized through canonical coefficients.
+ if (a.as<IntImmNode>() && b.as<IntImmNode>()) return prim::Sub(a, b);
// canonical form simplification.
SumExpr ret = ToSumExpr(std::move(a));
@@ -824,6 +828,8 @@ UnchangedOr<PrimExpr> CanonicalSimplifier::Impl::Mutate_(const prim::MulNode* op
// const folding
if (auto const_res = TryConstFold<prim::Mul>(a, b)) return *std::move(const_res);
+ // A declined literal fold must not be materialized through canonical coefficients.
+ if (a.as<IntImmNode>() && b.as<IntImmNode>()) return prim::Mul(a, b);
// x * c
if (a.as<IntImmNode>()) {
diff --git a/src/sym/const_fold.h b/src/sym/const_fold.h
index 6038ea3e..cfe9c0ab 100644
--- a/src/sym/const_fold.h
+++ b/src/sym/const_fold.h
@@ -32,7 +32,6 @@ namespace tvm {
namespace sym {
using prim::TryConstFold;
-using prim::details::GetFoldResultDoubleRepr;
using prim::details::IsIndexType;
using prim::details::IsIndexTypedExpr;
diff --git a/src/sym/rewrite_simplify.cc b/src/sym/rewrite_simplify.cc
index f72a48d9..9e4e5b40 100644
--- a/src/sym/rewrite_simplify.cc
+++ b/src/sym/rewrite_simplify.cc
@@ -669,15 +669,9 @@ UnchangedOr<PrimExpr> RewriteSimplifier::Impl::Mutate_(const prim::DivNode* op,
// We adopt the default C division uses truncation instead of floordiv.
// This means most rules need to check non-negativeness of the operands.
- // TryConstFold doesn't work for negative cases because it is also used by legacy
- // parts of tvm which still assume euclidean div. In this simplifier we assume that the division
- // is truncated, so perform const folding again.
- // NOTE: trunc div required
- if (truncdiv(c1, c2).Match(ret)) {
- ffi::BigInt c1val = c1.Eval()->value;
- ffi::BigInt c2val = c2.Eval()->value;
- return IntImm(op->ty.as_or_throw<PrimType>(), (c1val / c2val));
- }
+ // TryConstFold already handles truncation for both signs. A declined
+ // literal division has an unrepresentable result and must remain typed.
+ if (truncdiv(c1, c2).Match(ret)) return ret;
// while it is always true for trunc div
// restrict to common case(positive div)
diff --git a/tests/python/tirx/base/test_imm_values.py b/tests/python/tirx/base/test_imm_values.py
index 75d701ae..1f12b304 100644
--- a/tests/python/tirx/base/test_imm_values.py
+++ b/tests/python/tirx/base/test_imm_values.py
@@ -86,14 +86,39 @@ def test_tir_large_py_int_literals(dtype, literals):
def test_tir_intimm_overflow():
- assert int(tirx.const(255, "uint8") + tirx.const(1, "uint8")) == 0
- assert int(tirx.const(2**31 - 1, "int32") + tirx.const(1, "int32")) == -(2**31)
- assert int(tirx.const(2**32 - 1, "uint32") + tirx.const(1, "uint32")) == 0
- assert int(tirx.const(2**63 - 1, "int64") + tirx.const(1, "int64")) == -(2**63)
- assert int(tirx.const(2**32, "uint64") * tirx.const(2**32, "uint64")) == 0
- # customized int types
- assert int(tirx.const(7, "int4") + tirx.const(1, "int4")) == -8
- assert int(tirx.const(2**39 - 1, "int40") + tirx.const(1, "int40")) == -(2**39)
+ assert isinstance(tirx.const(255, "uint8") + tirx.const(1, "uint8"), tirx.Add)
+ assert isinstance(tirx.const(2**31 - 1, "int32") + tirx.const(1, "int32"), tirx.Add)
+ assert isinstance(tirx.const(2**32 - 1, "uint32") + tirx.const(1, "uint32"), tirx.Add)
+ assert isinstance(tirx.const(2**63 - 1, "int64") + tirx.const(1, "int64"), tirx.Add)
+ assert isinstance(tirx.const(2**32, "uint64") * tirx.const(2**32, "uint64"), tirx.Mul)
+ # Customized integer widths follow the same conservative policy.
+ assert isinstance(tirx.const(7, "int4") + tirx.const(1, "int4"), tirx.Add)
+ assert isinstance(tirx.const(2**39 - 1, "int40") + tirx.const(1, "int40"), tirx.Add)
+
+
+def test_tir_integer_cast_folding():
+ from tvm_tirx.ir.prim import Broadcast
+
+ small = tirx.const(127, "int32")
+ large = tirx.const(128, "int32")
+ assert small.astype("int8").value == 127
+ narrowed = large.astype("int8")
+ assert isinstance(narrowed, tirx.Cast) and narrowed.value.same_as(large)
+ assert narrowed.ty == tvm_tirx.ir.PrimType("int8")
+ assert isinstance(tirx.const(-1, "int32").astype("uint32"), tirx.Cast)
+ assert isinstance(tirx.const(2**64 - 1, "uint64").astype("int64"), tirx.Cast)
+ assert tirx.const(2**63, "uint64").astype("int128").value == 2**63
+
+ folded = small.astype("int8x4")
+ assert isinstance(folded, Broadcast) and folded.value.value == 127
+ deferred = large.astype("int8x4")
+ assert isinstance(deferred, Broadcast) and isinstance(deferred.value, tirx.Cast)
+ assert deferred.ty == tvm_tirx.ir.PrimType("int8x4")
+ vector = Broadcast(large, 4).astype("int8x4")
+ tvm_tirx.ir.assert_structural_equal(vector, deferred)
+ scalable = large.astype("int8xvscalex4")
+ assert isinstance(scalable, Broadcast) and isinstance(scalable.value, tirx.Cast)
+ assert scalable.ty == tvm_tirx.ir.PrimType("int8xvscalex4")
def compare_float_value(value, expect, msg):
@@ -109,7 +134,7 @@ def compare_float_value(value, expect, msg):
"dtype, literals",
[
["float16", [-65504.0, 3.14, 65504.0, np.inf, np.nan]],
- ["bfloat16", [-3.38953139e38, 3.38953139e38, 3.14]],
+ ["bfloat16", [-float.fromhex("0x1.fep127"), float.fromhex("0x1.fep127"), 3.14]],
["float32", [np.finfo("float32").min, 3.14, np.finfo("float32").max, np.inf, np.nan]],
["float64", [np.finfo("float64").min, 3.14, np.finfo("float64").max, np.inf, np.nan]],
],
@@ -229,6 +254,8 @@ def check_tir_const_fold(
else:
y = np.random.uniform(y_range[0], y_range[1])
+ if dtype.startswith(("int", "uint")):
+ x, y = int(x), int(y)
if skip_overflow:
py_res = foldf(x, y)
if isinstance(py_res, tvm_tirx.ir.IntImm | tvm_tirx.ir.FloatImm):
@@ -247,6 +274,13 @@ def check_tir_const_fold(
+ f"if it fails please report it in GitHub issue together with this seed {seed}\n"
)
if dtype.startswith("float"):
+ if not isinstance(fold_res, tvm_tirx.ir.FloatImm):
+ exact_res = foldf(float(np.float32(x)), float(np.float32(y)))
+ assert math.isfinite(exact_res) and abs(exact_res) > float(ninfo.max), flaky_msg
+ assert fold_res.ty == tvm_tirx.ir.PrimType(dtype)
+ if expect:
+ compare_float_value(expect, calc_res, flaky_msg)
+ return
compare_float_value(calc_res, fold_res.value, flaky_msg)
if expect:
compare_float_value(expect, calc_res, flaky_msg)
@@ -344,9 +378,9 @@ def test_tir_int8_const_fold():
ffloordiv = tvm_tirx.tirx.compile(imm_floordiv, target="llvm")
ftruncdiv = tvm_tirx.tirx.compile(imm_truncdiv, target="llvm")
- # overflow
- check_tir_const_fold("int8", lambda x, y: x + y, fadd, 127, 1, -128)
- check_tir_const_fold("int8", lambda x, y: x * y, fmul, 127, 127, 1)
+ # Overflow remains a typed operation instead of manufacturing a wrapped literal.
+ assert isinstance(tirx.const(127, "int8") + tirx.const(1, "int8"), tirx.Add)
+ assert isinstance(tirx.const(127, "int8") * tirx.const(127, "int8"), tirx.Mul)
# divide by zero
with pytest.raises(RuntimeError):
@@ -363,9 +397,9 @@ def test_tir_int8_const_fold():
)
# randomized check
- check_tir_const_fold("int8", lambda x, y: x * y, fmul)
- check_tir_const_fold("int8", lambda x, y: x + y, fadd)
- check_tir_const_fold("int8", lambda x, y: x - y, fsub)
+ check_tir_const_fold("int8", lambda x, y: x * y, fmul, skip_overflow=True)
+ check_tir_const_fold("int8", lambda x, y: x + y, fadd, skip_overflow=True)
+ check_tir_const_fold("int8", lambda x, y: x - y, fsub, skip_overflow=True)
check_tir_const_fold(
"int8", lambda x, y: tirx.floordiv(x, y), ffloordiv, y_range=(1, np.iinfo("int8").max)
)
@@ -404,12 +438,9 @@ def test_tir_uint8_const_fold():
ffloordiv = tvm_tirx.tirx.compile(imm_floordiv, target="llvm")
ftruncdiv = tvm_tirx.tirx.compile(imm_truncdiv, target="llvm")
- # overflow
- check_tir_const_fold("uint8", lambda x, y: x + y, fadd, 255, 1, 0)
-
- # zero sub
- with pytest.raises(RuntimeError):
- check_tir_const_fold("uint8", lambda x, y: x - y, fsub, 0, 10)
+ # Both overflow and negative unsigned results decline folding.
+ assert isinstance(tirx.const(255, "uint8") + tirx.const(1, "uint8"), tirx.Add)
+ assert isinstance(tirx.const(0, "uint8") - tirx.const(10, "uint8"), tirx.Sub)
# divide by zero
with pytest.raises(RuntimeError):
@@ -428,9 +459,9 @@ def test_tir_uint8_const_fold():
)
# randomized check
- check_tir_const_fold("uint8", lambda x, y: x * y, fmul)
- check_tir_const_fold("uint8", lambda x, y: x + y, fadd)
- check_tir_const_fold("uint8", lambda x, y: x - y, fsub)
+ check_tir_const_fold("uint8", lambda x, y: x * y, fmul, skip_overflow=True)
+ check_tir_const_fold("uint8", lambda x, y: x + y, fadd, skip_overflow=True)
+ check_tir_const_fold("uint8", lambda x, y: x - y, fsub, skip_overflow=True)
check_tir_const_fold(
"uint8", lambda x, y: tirx.floordiv(x, y), ffloordiv, y_range=(1, np.iinfo("uint8").max)
)
@@ -479,9 +510,9 @@ def test_tir_int32_const_fold():
ftruncdiv = tvm_tirx.tirx.compile(imm_truncdiv, target="llvm")
ftruncmod = tvm_tirx.tirx.compile(imm_truncmod, target="llvm")
- # i32 overflow is not specified, only check for range
- assert -(2**31) <= int(tirx.const(2**31 - 1, "int32") + tirx.const(1, "int32")) < 2**31
- assert -(2**31) <= int(tirx.const(-(2**31), "int32") - tirx.const(1, "int32")) < 2**31
+ # Out-of-range exact results remain typed operations.
+ assert isinstance(tirx.const(2**31 - 1, "int32") + tirx.const(1, "int32"), tirx.Add)
+ assert isinstance(tirx.const(-(2**31), "int32") - tirx.const(1, "int32"), tirx.Sub)
# divide by zero
with pytest.raises(RuntimeError):
@@ -557,8 +588,8 @@ def test_tir_uint32_const_fold():
ffloordiv = tvm_tirx.tirx.compile(imm_floordiv, target="llvm")
ftruncdiv = tvm_tirx.tirx.compile(imm_truncdiv, target="llvm")
- # u32 overflow is not specified, only check for range
- assert 0 <= int(tirx.const(2**32 - 1, "uint32") + tirx.const(1, "uint32")) < 2**32
+ # Out-of-range exact results remain typed operations.
+ assert isinstance(tirx.const(2**32 - 1, "uint32") + tirx.const(1, "uint32"), tirx.Add)
# divide by zero
with pytest.raises(RuntimeError):
diff --git a/tests/python/tirx/transform/test_simplify.py b/tests/python/tirx/transform/test_simplify.py
index 4fdecdac..5bf482f3 100644
--- a/tests/python/tirx/transform/test_simplify.py
+++ b/tests/python/tirx/transform/test_simplify.py
@@ -497,7 +497,7 @@ def test_if_then_else_expr():
def test_ceil_log2_int():
- """Simplify expressions resulting from topi.math.ceil_log2"""
+ """Preserve ceil_log2 when its integer-to-float64 cast does not fold."""
@T.function(private=True)
def before(A: T.Tensor(1, "int32")):
@@ -505,7 +505,7 @@ def test_ceil_log2_int():
@T.function(private=True)
def expected(A: T.Tensor(1, "int32")):
- A[0] = 4
+ A[0] = T.Cast("int32", T.ceil(T.log2(T.Cast("float64", 14))))
after = _apply_simplify(before)
tvm_tirx.ir.assert_structural_equal(after, expected)