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 336af7322 [FIX] Fold float64 literals and require LLVM 17 (#95)
336af7322 is described below
commit 336af73223991bd5586efe528bd59f888c0efc46
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 27 15:45:55 2026 -0400
[FIX] Fold float64 literals and require LLVM 17 (#95)
Keep C++ and Rust float64 literal folding consistent and restore the
folded ceil-log2 expectation. Require LLVM 17 or newer, retaining the
declared packaged LLVM toolchain for CPU and GPU builds.
Remove the ARM popcount instruction-count test.
---
cmake/modules/LLVM.cmake | 4 ++--
cmake/utils/FindLLVM.cmake | 6 ++---
rust/src/ir/prim/op.rs | 11 ++++++++++
src/ir/prim/op.cc | 11 ++++++++++
tests/python/codegen/llvm/test_arm.py | 33 ----------------------------
tests/python/tirx/transform/test_simplify.py | 4 ++--
6 files changed, 29 insertions(+), 40 deletions(-)
diff --git a/cmake/modules/LLVM.cmake b/cmake/modules/LLVM.cmake
index 2475cff3f..e9cb0f412 100644
--- a/cmake/modules/LLVM.cmake
+++ b/cmake/modules/LLVM.cmake
@@ -40,8 +40,8 @@ if(NOT ${TVM_TIRX_USE_LLVM} MATCHES ${IS_FALSE_PATTERN})
set(TVM_LLVM_BACKEND_LIB tvm_tirx_llvm_backend_native)
endif()
find_llvm(${TVM_TIRX_USE_LLVM})
- if (${TVM_LLVM_VERSION} LESS 150)
- message(FATAL_ERROR "LLVM version 15.0 or greater is required.")
+ if (${TVM_LLVM_VERSION} LESS 170)
+ message(FATAL_ERROR "LLVM version 17.0 or greater is required.")
endif()
include_directories(SYSTEM ${LLVM_INCLUDE_DIRS})
add_definitions(${LLVM_DEFINITIONS})
diff --git a/cmake/utils/FindLLVM.cmake b/cmake/utils/FindLLVM.cmake
index a1e3d964d..93c5f9815 100644
--- a/cmake/utils/FindLLVM.cmake
+++ b/cmake/utils/FindLLVM.cmake
@@ -106,7 +106,7 @@ macro(find_llvm use_llvm)
message(FATAL_ERROR
"LLVM development files were not found. From the source root, run "
"`uv sync --only-group build` and configure with `uv run --no-sync
cmake`, "
- "or install LLVM 15 or newer, including "
+ "or install LLVM 17 or newer, including "
"its development headers and libraries (for example, llvm-dev on
Debian/Ubuntu "
"or llvm-devel on Fedora). Set LLVM_DIR to the directory containing "
"LLVMConfig.cmake, add the LLVM installation to CMAKE_PREFIX_PATH, or
set "
@@ -285,8 +285,8 @@ macro(find_llvm use_llvm)
message(STATUS "Found LLVM_DEFINITIONS=" "${LLVM_DEFINITIONS}")
message(STATUS "Found LLVM_LIBS=" "${LLVM_LIBS}")
message(STATUS "Found TVM_LLVM_VERSION=" ${TVM_LLVM_VERSION})
- if (${TVM_LLVM_VERSION} LESS 150)
- message(FATAL_ERROR "TVM requires LLVM 15.0 or higher.")
+ if (${TVM_LLVM_VERSION} LESS 170)
+ message(FATAL_ERROR "TVM requires LLVM 17.0 or higher.")
endif()
message(STATUS "Found TVM_LLVM_HAS_AARCH64_TARGET="
${TVM_LLVM_HAS_AARCH64_TARGET})
diff --git a/rust/src/ir/prim/op.rs b/rust/src/ir/prim/op.rs
index 0aa5f9da2..ec62cf30e 100644
--- a/rust/src/ir/prim/op.rs
+++ b/rust/src/ir/prim/op.rs
@@ -561,6 +561,7 @@ fn is_rounded_to_16_bits(value: f64, bfloat: bool) -> bool {
fn fold_cast(value: &Expr, to: DLDataType) -> Result<Option<Expr>> {
let integer = matches!(to.code, support::INT | support::UINT |
support::BOOL);
let float32 = to.code == support::FLOAT && to.bits == 32;
+ let float64 = to.code == support::FLOAT && to.bits == 64;
let limit_16 = match (to.code, to.bits) {
(support::FLOAT, 16) => Some(support::MAX_FLOAT16),
(support::BFLOAT, 16) => Some(support::MAX_BFLOAT16),
@@ -581,6 +582,11 @@ fn fold_cast(value: &Expr, to: DLDataType) ->
Result<Option<Expr>> {
if let Some(small) = imm.value.to_i64() {
return float(f64::from(small as f32), &imm.loc);
}
+ } else if float64 {
+ let limit = BigInt::try_from(f64::MAX)?;
+ if (-&limit..=limit).contains(&imm.value) {
+ return float(imm.value.to_f64()?, &imm.loc);
+ }
} else if let Some(limit) = limit_16 {
let source = imm.value.to_f64()?;
// Keep the integer exact until the destination performs its
rounding.
@@ -618,6 +624,11 @@ fn fold_cast(value: &Expr, to: DLDataType) ->
Result<Option<Expr>> {
if const_fold::is_in_float32_range(source) {
return float(f64::from(source as f32), &imm.loc);
}
+ } else if float64
+ && ((from.code == support::FLOAT && matches!(from.bits, 16 | 32))
|| is_bfloat16(from))
+ {
+ // Source rounding was checked above; widening preserves its value
exactly.
+ return float(source, &imm.loc);
} else if let Some(limit) = limit_16 {
if !source.is_finite() || (-limit..=limit).contains(&source) {
return float(source, &imm.loc);
diff --git a/src/ir/prim/op.cc b/src/ir/prim/op.cc
index d53568f6b..46777c54b 100644
--- a/src/ir/prim/op.cc
+++ b/src/ir/prim/op.cc
@@ -317,6 +317,11 @@ PrimExpr cast(PrimExpr value, PrimType t, Location loc) {
if (auto small = op->value.as<int64_t>()) {
return FloatImm(dtype, static_cast<float>(*small), op->loc);
}
+ } else if (dtype.MatchesElementType(DLDataTypeCode::kDLFloat, 64)) {
+ ffi::BigInt limit(std::numeric_limits<double>::max());
+ if (op->value >= -limit && op->value <= limit) {
+ return FloatImm(dtype, static_cast<double>(op->value), op->loc);
+ }
} else if (dtype.MatchesElementType(DLDataTypeCode::kDLFloat, 16) ||
dtype.MatchesElementType(DLDataTypeCode::kDLBfloat, 16)) {
double source = static_cast<double>(op->value);
@@ -358,6 +363,12 @@ PrimExpr cast(PrimExpr value, PrimType t, Location loc) {
if (detail::IsInFloat32Range(source)) {
return FloatImm(dtype, static_cast<float>(source), op->loc);
}
+ } else if (dtype.MatchesElementType(DLDataTypeCode::kDLFloat, 64) &&
+ (value.ty().MatchesElementType(DLDataTypeCode::kDLFloat, 32)
||
+ value.ty().MatchesElementType(DLDataTypeCode::kDLFloat, 16)
||
+ value.ty().MatchesElementType(DLDataTypeCode::kDLBfloat,
16))) {
+ // Source rounding was checked above; widening preserves its value
exactly.
+ return FloatImm(dtype, source, op->loc);
} else if (dtype.MatchesElementType(DLDataTypeCode::kDLFloat, 16) ||
dtype.MatchesElementType(DLDataTypeCode::kDLBfloat, 16)) {
double limit = dtype.MatchesCode(DLDataTypeCode::kDLBfloat) ?
support::kMaxBFloat16
diff --git a/tests/python/codegen/llvm/test_arm.py
b/tests/python/codegen/llvm/test_arm.py
index 78c13870e..c9dc3c32c 100644
--- a/tests/python/codegen/llvm/test_arm.py
+++ b/tests/python/codegen/llvm/test_arm.py
@@ -21,38 +21,6 @@ from tvm_tirx.script import ir as I
from tvm_tirx.script import tirx as T
-def test_popcount():
- target = {
- "kind": "llvm",
- "mtriple": "armv7l-none-linux-gnueabihf",
- "mcpu": "cortex-a53",
- "mattr": ["+neon"],
- }
-
- def check_correct_assembly(type, elements, counts):
- @I.ir_module
- class Module:
- @T.function
- def main(A: T.Tensor((elements,), type), B: T.Tensor((elements,),
type)):
- T.func_attr({"tirx.noalias": True})
- for i in T.vectorized(elements):
- B[i] = T.popcount(A[i])
-
- f = tvm_tirx.tirx.compile(Module, target=target).mod
- # Verify we see the correct number of vpaddl and vcnt instructions in
the assembly
- assembly = f.inspect_source("asm")
- matches = re.findall("vpaddl", assembly)
- assert len(matches) == counts
- matches = re.findall("vcnt", assembly)
- assert len(matches) == 1
-
- check_correct_assembly("uint16", 8, 1)
- check_correct_assembly("uint16", 4, 1)
- check_correct_assembly("uint32", 4, 2)
- check_correct_assembly("uint32", 2, 2)
- check_correct_assembly("uint64", 2, 3)
-
-
def test_vmlal_s16():
target = {
"kind": "llvm",
@@ -121,5 +89,4 @@ def test_vmlal_s16():
if __name__ == "__main__":
- test_popcount()
test_vmlal_s16()
diff --git a/tests/python/tirx/transform/test_simplify.py
b/tests/python/tirx/transform/test_simplify.py
index d60942bf6..dad1daae2 100644
--- a/tests/python/tirx/transform/test_simplify.py
+++ b/tests/python/tirx/transform/test_simplify.py
@@ -500,7 +500,7 @@ def test_if_then_else_expr():
def test_ceil_log2_int():
- """Preserve ceil_log2 when its integer-to-float64 cast does not fold."""
+ """Fold ceil_log2 of an integer constant through its float64 cast."""
@T.function(private=True)
def before(A: T.Tensor(1, "int32")):
@@ -508,7 +508,7 @@ def test_ceil_log2_int():
@T.function(private=True)
def expected(A: T.Tensor(1, "int32")):
- A[0] = T.Cast("int32", T.ceil(T.log2(T.Cast("float64", 14))))
+ A[0] = 4
after = _apply_simplify(before)
tvm_tirx.ir.assert_structural_equal(after, expected)
