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 effeeca8 [FIX] Preserve NaNs in host floating-point min and max (#31)
effeeca8 is described below
commit effeeca8e20101ed705c43908bb044389116ec80
Author: Shushi Hong <[email protected]>
AuthorDate: Fri Oct 23 15:22:32 2026 -0400
[FIX] Preserve NaNs in host floating-point min and max (#31)
This PR fixes floating-point min/max NaN propagation in TVM-TIRx’s
constant folding, C host codegen, and LLVM codegen. C host codegen emits
NaN-preserving conditional expressions, simplifies cases involving
floating-point constants, and prints valid NaN/Inf literals. LLVM
codegen uses `llvm.minimum`/`llvm.maximum` for floating-point scalars
and vectors, adopting their signed-zero rules. Tests cover NaNs in
either operand, scalar and broadcast constants, infinities, signed zero,
C host float32/64, LLVM float16/32/64 and legalized BF16 scalars and
fixed-length vectors, scalable-vector codegen, and constant folding.
---
include/tvm/ir/prim/const_fold.h | 4 +-
src/codegen/c/codegen_c_host.cc | 76 +++++++++++++++++++++-----
src/codegen/c/codegen_c_host.h | 13 ++---
src/codegen/llvm/codegen_llvm.cc | 24 ++++++---
src/codegen/llvm/codegen_llvm.h | 1 +
src/codegen/min_max_utils.h | 70 ++++++++++++++++++++++++
tests/python/codegen/c/test_basic.py | 58 ++++++++++++++++++++
tests/python/codegen/llvm/test_basic.py | 89 +++++++++++++++++++++++++++++++
tests/python/tirx/base/test_imm_values.py | 20 +++++++
9 files changed, 327 insertions(+), 28 deletions(-)
diff --git a/include/tvm/ir/prim/const_fold.h b/include/tvm/ir/prim/const_fold.h
index fc435eda..6be3adb2 100644
--- a/include/tvm/ir/prim/const_fold.h
+++ b/include/tvm/ir/prim/const_fold.h
@@ -373,7 +373,7 @@ inline ffi::Optional<PrimExpr>
TryConstFold<prim::Min>(PrimExpr a, PrimExpr b) {
TVM_PRIM_CONST_PROPAGATION({
PrimType result_ty = a.ty();
if (pa && pb) return IntImm(result_ty, std::min(pa->value, pb->value));
- if (fa && fb) return FloatImm(result_ty, std::min(fa->value, fb->value));
+ if (fa && fb) return std::isnan(fa->value) || fa->value < fb->value ? a :
b;
});
if (a.same_as(b)) return a;
return std::nullopt;
@@ -384,7 +384,7 @@ inline ffi::Optional<PrimExpr>
TryConstFold<prim::Max>(PrimExpr a, PrimExpr b) {
TVM_PRIM_CONST_PROPAGATION({
PrimType result_ty = a.ty();
if (pa && pb) return IntImm(result_ty, std::max(pa->value, pb->value));
- if (fa && fb) return FloatImm(result_ty, std::max(fa->value, fb->value));
+ if (fa && fb) return std::isnan(fa->value) || fa->value > fb->value ? a :
b;
});
if (a.same_as(b)) return a;
return std::nullopt;
diff --git a/src/codegen/c/codegen_c_host.cc b/src/codegen/c/codegen_c_host.cc
index 4e56650f..92b367fd 100644
--- a/src/codegen/c/codegen_c_host.cc
+++ b/src/codegen/c/codegen_c_host.cc
@@ -31,10 +31,13 @@
#include <tvm/tirx/op/abi.h>
#include <algorithm>
+#include <cmath>
#include <string>
#include <utility>
#include <vector>
+#include "../min_max_utils.h"
+
namespace tvm {
namespace codegen {
@@ -350,26 +353,71 @@ void CodeGenCHost::Dispatch_(const AssertStmtObj* op) {
// NOLINT(*)
}
}
+void CodeGenCHost::Dispatch_(const FloatImmObj* op, std::ostream& os) { //
NOLINT(*)
+ if (std::isfinite(op->value)) {
+ CodeGenC::Dispatch_(op, os);
+ return;
+ }
+
+ const char* literal;
+ if (std::isnan(op->value)) {
+ literal = "NAN";
+ } else if (op->value < 0) {
+ literal = "-INFINITY";
+ } else {
+ literal = "INFINITY";
+ }
+ os << "((";
+ PrintType(op->ty.as_or_throw<PrimType>(), os);
+ os << ")" << literal << ")";
+}
+
void CodeGenCHost::Dispatch_(const prim::MinObj* op, std::ostream& os) { //
NOLINT(*)
- PrintTernaryCondExpr(op, "<", os);
+ PrintMinMax(op->a, op->b, /*is_min=*/true, os);
}
void CodeGenCHost::Dispatch_(const prim::MaxObj* op, std::ostream& os) { //
NOLINT(*)
- PrintTernaryCondExpr(op, ">", os);
+ PrintMinMax(op->a, op->b, /*is_min=*/false, os);
}
-template <typename T>
-inline void CodeGenCHost::PrintTernaryCondExpr(const T* op, const char*
compare,
- std::ostream& os) { //
NOLINT(*)
- std::ostringstream temp_a;
- Dispatch(op->a, temp_a);
- std::string a_id = SSAGetID(temp_a.str(), op->a.ty());
- std::ostringstream temp_b;
- Dispatch(op->b, temp_b);
- std::string b_id = SSAGetID(temp_b.str(), op->b.ty());
-
- os << "((" << a_id << ") " << compare << " (" << b_id << ") "
- << "? (" << a_id << ") : (" << b_id << "))";
+void CodeGenCHost::PrintMinMax(const PrimExpr& lhs, const PrimExpr& rhs, bool
is_min,
+ std::ostream& os) { // NOLINT(*)
+ std::ostringstream temp_lhs;
+ Dispatch(lhs, temp_lhs);
+ std::string lhs_value = "(" + SSAGetID(temp_lhs.str(), lhs.ty()) + ")";
+ std::ostringstream temp_rhs;
+ Dispatch(rhs, temp_rhs);
+ std::string rhs_value = "(" + SSAGetID(temp_rhs.str(), rhs.ty()) + ")";
+ PrimType dtype = lhs.ty().as_or_throw<PrimType>();
+ FloatNaNStatus lhs_nan_status = GetFloatNaNStatus(lhs);
+ FloatNaNStatus rhs_nan_status = GetFloatNaNStatus(rhs);
+ const char* compare = is_min ? "<" : ">";
+ const char* take_rhs_compare = is_min ? ">=" : "<=";
+
+ std::string condition;
+ std::string true_value = lhs_value;
+ std::string false_value = rhs_value;
+ if (!dtype.MatchesCode(DLDataTypeCode::kDLFloat) ||
+ lhs_nan_status == FloatNaNStatus::kKnownNotNaN) {
+ // Integers and a known non-NaN lhs need only the strict comparison.
+ condition = lhs_value + " " + compare + " " + rhs_value;
+ } else if (lhs_nan_status == FloatNaNStatus::kKnownNaN) {
+ os << lhs_value;
+ return;
+ } else if (rhs_nan_status == FloatNaNStatus::kKnownNaN) {
+ // isnan(lhs) ? lhs : rhs: select lhs only if it is also NaN.
+ condition = lhs_value + " != " + lhs_value;
+ } else if (rhs_nan_status == FloatNaNStatus::kKnownNotNaN) {
+ // For min: lhs >= rhs ? rhs : lhs. A NaN lhs makes >= false and is kept.
+ condition = lhs_value + " " + take_rhs_compare + " " + rhs_value;
+ true_value = rhs_value;
+ false_value = lhs_value;
+ } else {
+ // For min: (lhs < rhs || isnan(lhs)) ? lhs : rhs. Max uses > instead.
+ condition = "(" + lhs_value + " " + compare + " " + rhs_value + ") || (" +
lhs_value +
+ " != " + lhs_value + ")";
+ }
+ os << "((" << condition << ") ? " << true_value << " : " << false_value <<
")";
}
ffi::Module BuildCHost(IRModule mod, Target target) {
diff --git a/src/codegen/c/codegen_c_host.h b/src/codegen/c/codegen_c_host.h
index 3c296694..abee26ad 100644
--- a/src/codegen/c/codegen_c_host.h
+++ b/src/codegen/c/codegen_c_host.h
@@ -62,6 +62,7 @@ class CodeGenCHost : public CodeGenC {
// overload visitor functions
void Dispatch_(const prim::BroadcastObj* op, std::ostream& os) final; //
NOLINT(*)
void Dispatch_(const CallObj* op, std::ostream& os) override; //
NOLINT(*)
+ void Dispatch_(const FloatImmObj* op, std::ostream& os) final; //
NOLINT(*)
// overload min and max to use the ternary operator, so we don't rely on the
// standard library implementations
void Dispatch_(const prim::MinObj* op, std::ostream& os) final; // NOLINT(*)
@@ -96,15 +97,15 @@ class CodeGenCHost : public CodeGenC {
void PrintGetFuncFromBackend(const std::string& func_name, const
std::string& packed_func_name);
void PrintCallPacked(const CallObj* op);
/*!
- * \brief Print ternary conditional operator implementing binary `op`
+ * \brief Print min/max, propagating floating-point NaNs and selecting rhs
on ties.
* Forces the operands to be in SSA form.
- * \param op binary operator being expressed
- * \param compare string representation of comparison operator
+ * \param lhs Left operand.
+ * \param rhs Right operand.
+ * \param is_min Whether to print min (true) or max (false).
* \param os stream reference to print into
*/
- template <typename T>
- inline void PrintTernaryCondExpr(const T* op, const char* compare,
- std::ostream& os); // NOLINT(*)
+ void PrintMinMax(const PrimExpr& lhs, const PrimExpr& rhs, bool is_min,
+ std::ostream& os); // NOLINT(*)
};
} // namespace codegen
diff --git a/src/codegen/llvm/codegen_llvm.cc b/src/codegen/llvm/codegen_llvm.cc
index b72c37ba..85093f12 100644
--- a/src/codegen/llvm/codegen_llvm.cc
+++ b/src/codegen/llvm/codegen_llvm.cc
@@ -1640,15 +1640,27 @@ llvm::Value* CodeGenLLVM::Dispatch_(const prim::ModObj*
op) {
}
llvm::Value* CodeGenLLVM::Dispatch_(const prim::MinObj* op) {
- llvm::Value* a = MakeValue(op->a);
- llvm::Value* b = MakeValue(op->b);
- return builder_->CreateSelect(CreateLT(PrimType(op->a.ty()->dtype), a, b),
a, b);
+ return CreateMinMax(op->a, op->b, /*is_min=*/true);
}
llvm::Value* CodeGenLLVM::Dispatch_(const prim::MaxObj* op) {
- llvm::Value* a = MakeValue(op->a);
- llvm::Value* b = MakeValue(op->b);
- return builder_->CreateSelect(CreateGT(PrimType(op->a.ty()->dtype), a, b),
a, b);
+ return CreateMinMax(op->a, op->b, /*is_min=*/false);
+}
+
+llvm::Value* CodeGenLLVM::CreateMinMax(const PrimExpr& lhs, const PrimExpr&
rhs, bool is_min) {
+ llvm::Value* lhs_value = MakeValue(lhs);
+ llvm::Value* rhs_value = MakeValue(rhs);
+ if (lhs_value->getType()->isFPOrFPVectorTy()) {
+ // Propagate NaNs lane-wise for scalars, fixed-length vectors, and
scalable vectors.
+ // On a signed-zero tie, minimum selects -0 and maximum selects +0.
+ return is_min ? builder_->CreateMinimum(lhs_value, rhs_value)
+ : builder_->CreateMaximum(lhs_value, rhs_value);
+ }
+
+ PrimType dtype(lhs.ty()->dtype);
+ llvm::Value* take_lhs =
+ is_min ? CreateLT(dtype, lhs_value, rhs_value) : CreateGT(dtype,
lhs_value, rhs_value);
+ return builder_->CreateSelect(take_lhs, lhs_value, rhs_value);
}
llvm::Value* CodeGenLLVM::Dispatch_(const prim::EQObj* op) {
diff --git a/src/codegen/llvm/codegen_llvm.h b/src/codegen/llvm/codegen_llvm.h
index 9e2355ba..82a44cc1 100644
--- a/src/codegen/llvm/codegen_llvm.h
+++ b/src/codegen/llvm/codegen_llvm.h
@@ -492,6 +492,7 @@ class CodeGenLLVM : public
tvm::ExprFunctor<llvm::Value*(const Expr&)>,
llvm::Value* CreateLE(PrimType t, llvm::Value* a, llvm::Value* b);
llvm::Value* CreateGT(PrimType t, llvm::Value* a, llvm::Value* b);
llvm::Value* CreateGE(PrimType t, llvm::Value* a, llvm::Value* b);
+ llvm::Value* CreateMinMax(const PrimExpr& lhs, const PrimExpr& rhs, bool
is_min);
llvm::Value* CreateAdd(PrimType t, llvm::Value* a, llvm::Value* b);
llvm::Value* CreateSub(PrimType t, llvm::Value* a, llvm::Value* b);
llvm::Value* CreateMul(PrimType t, llvm::Value* a, llvm::Value* b);
diff --git a/src/codegen/min_max_utils.h b/src/codegen/min_max_utils.h
new file mode 100644
index 00000000..b78c1e1d
--- /dev/null
+++ b/src/codegen/min_max_utils.h
@@ -0,0 +1,70 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file min_max_utils.h
+ * \brief Common utilities for lowering floating-point min and max.
+ */
+#ifndef TVM_CODEGEN_MIN_MAX_UTILS_H_
+#define TVM_CODEGEN_MIN_MAX_UTILS_H_
+
+#include <tvm/ir/prim/expr.h>
+
+#include <cmath>
+
+namespace tvm {
+namespace codegen {
+
+/*! \brief Compile-time NaN information about an expression. */
+enum class FloatNaNStatus {
+ /*! \brief Whether the expression is NaN has not been determined. */
+ kUnknown,
+ /*! \brief The scalar value or every vector lane is known to be NaN. */
+ kKnownNaN,
+ /*! \brief The scalar value or every vector lane is known not to be NaN. */
+ kKnownNotNaN,
+};
+
+/*!
+ * \brief Determine an expression's NaN status from floating-point constants.
+ *
+ * Recognizes a scalar FloatImm or a Broadcast of a FloatImm, including
+ * fixed-length and scalable broadcasts. Non-NaN constants include infinities.
+ * All other expression forms have unknown NaN status.
+ *
+ * \param expr The expression to classify.
+ * \return FloatNaNStatus::kKnownNaN or FloatNaNStatus::kKnownNotNaN for a
recognized
+ * constant; FloatNaNStatus::kUnknown otherwise.
+ */
+inline FloatNaNStatus GetFloatNaNStatus(const PrimExpr& expr) {
+ const FloatImmObj* value = expr.as<FloatImmObj>();
+ if (const auto* broadcast = expr.as<prim::BroadcastObj>()) {
+ // Recognize uniform vector constants, including scalable broadcasts.
+ value = broadcast->value.as<FloatImmObj>();
+ }
+ if (value == nullptr) {
+ return FloatNaNStatus::kUnknown;
+ }
+ return std::isnan(value->value) ? FloatNaNStatus::kKnownNaN :
FloatNaNStatus::kKnownNotNaN;
+}
+
+} // namespace codegen
+} // namespace tvm
+
+#endif // TVM_CODEGEN_MIN_MAX_UTILS_H_
diff --git a/tests/python/codegen/c/test_basic.py
b/tests/python/codegen/c/test_basic.py
index 52d5bde2..2ab26b80 100644
--- a/tests/python/codegen/c/test_basic.py
+++ b/tests/python/codegen/c/test_basic.py
@@ -16,6 +16,7 @@
# under the License.
import numpy as np
+import pytest
import tvm_tirx
import tvm_tirx.testing
@@ -299,5 +300,62 @@ def test_if_then_else_avoids_extraneous_parentheses():
)
[email protected]("dtype", ["float32", "float64"])
+def test_min_max_nan_preserving(dtype):
+ n = 12
+ nan = float("nan")
+ infinity = float("inf")
+
+ @T.function
+ def main(A: T.Tensor((n,), dtype), B: T.Tensor((n,), dtype), C:
T.Tensor((14, n), dtype)):
+ for i in range(n):
+ # Rows 0-4: min(a, b), min(0, a), min(a, 0), min(NaN, a), min(a,
NaN).
+ C[0, i] = T.min(A[i], B[i])
+ C[1, i] = T.min(T.cast(0, dtype), A[i])
+ C[2, i] = T.min(A[i], T.cast(0, dtype))
+ C[3, i] = T.min(T.cast(nan, dtype), A[i])
+ C[4, i] = T.min(A[i], T.cast(nan, dtype))
+ # Rows 5-9: max with the same operand pairs.
+ C[5, i] = T.max(A[i], B[i])
+ C[6, i] = T.max(T.cast(0, dtype), A[i])
+ C[7, i] = T.max(A[i], T.cast(0, dtype))
+ C[8, i] = T.max(T.cast(nan, dtype), A[i])
+ C[9, i] = T.max(A[i], T.cast(nan, dtype))
+ # Rows 10-11: parentheses for general and constant max inside
arithmetic.
+ C[10, i] = T.max(A[i], B[i]) + T.cast(1, dtype)
+ C[11, i] = T.max(T.cast(0, dtype), A[i]) + T.cast(1, dtype)
+ # Rows 12-13: positive and negative infinity constants.
+ C[12, i] = T.min(A[i], T.cast(infinity, dtype))
+ C[13, i] = T.max(A[i], T.cast(-infinity, dtype))
+
+ compiled = tvm_tirx.tirx.compile(main, target="c")
+ a = np.array([np.nan, 1, np.nan, 0.0, -0.0, 3, 2, -5, np.inf, -np.inf, 1,
-1], dtype=dtype)
+ b = np.array([1, np.nan, np.nan, -0.0, 0.0, 2, 2, -4, 1, 1, np.inf,
-np.inf], dtype=dtype)
+ out = tvm_tirx.runtime.empty((14, n), dtype)
+ compiled(tvm_tirx.runtime.tensor(a), tvm_tirx.runtime.tensor(b), out)
+
+ zero, nan_array = np.zeros_like(a), np.full_like(a, np.nan)
+ operand_pairs = ((a, b), (zero, a), (a, zero), (nan_array, a), (a,
nan_array))
+ expected_by_op = {}
+ for name, compare in (("min", np.less), ("max", np.greater)):
+ expected_by_op[name] = [
+ np.where(compare(lhs, rhs) | np.isnan(lhs), lhs, rhs) for lhs, rhs
in operand_pairs
+ ]
+ expected = expected_by_op["min"] + expected_by_op["max"]
+ max_ab, max_zero_a = expected_by_op["max"][:2]
+ expected.extend([max_ab + np.ones_like(a), max_zero_a + np.ones_like(a)])
+ # min(a, +Inf) and max(a, -Inf) both preserve a, including NaNs and signed
zero.
+ expected.extend([a, a])
+ expected = np.asarray(expected, dtype=dtype)
+ got = out.numpy()
+ nan_mask = np.isnan(expected)
+ np.testing.assert_array_equal(np.isnan(got), nan_mask)
+ # NaN payloads may be canonicalized; preserve exact non-NaN values and
signed zero.
+ uint_dtype = f"uint{a.dtype.itemsize * 8}"
+ np.testing.assert_array_equal(
+ got.view(uint_dtype)[~nan_mask], expected.view(uint_dtype)[~nan_mask]
+ )
+
+
if __name__ == "__main__":
tvm_tirx.testing.main()
diff --git a/tests/python/codegen/llvm/test_basic.py
b/tests/python/codegen/llvm/test_basic.py
index 36b7fe9a..533a3c55 100644
--- a/tests/python/codegen/llvm/test_basic.py
+++ b/tests/python/codegen/llvm/test_basic.py
@@ -1314,5 +1314,94 @@ def test_llvm_module_properties_and_formats(tmp_path):
np.testing.assert_equal(data.numpy(), np.arange(10))
[email protected](not env.has_llvm(), reason="need llvm")
[email protected](
+ "dtype,lanes",
+ [(dtype, lanes) for dtype in ("float16", "float32", "float64", "bfloat16")
for lanes in (1, 4)]
+ + [("float32", "scalable")],
+)
+def test_llvm_min_max_nan_preserving(dtype, lanes):
+ scalable = lanes == "scalable"
+ extent = 4 * tvm_tirx.tirx.vscale() if scalable else lanes
+ n = 128 if scalable else 12
+ nan = float("nan")
+ infinity = float("inf")
+
+ @T.function
+ def main(A: T.Tensor((n,), dtype), B: T.Tensor((n,), dtype), C:
T.Tensor((14, n), dtype)):
+ for outer in range(n // extent):
+ i = outer if T.constexpr(lanes == 1) else T.ramp(outer * extent,
1, extent)
+ # Rows 0-4: min(a, b), min(0, a), min(a, 0), min(NaN, a), min(a,
NaN).
+ C[0, i] = T.min(A[i], B[i])
+ C[1, i] = T.min(T.cast(0, dtype), A[i])
+ C[2, i] = T.min(A[i], T.cast(0, dtype))
+ C[3, i] = T.min(T.cast(nan, dtype), A[i])
+ C[4, i] = T.min(A[i], T.cast(nan, dtype))
+ # Rows 5-9: max with the same operand pairs.
+ C[5, i] = T.max(A[i], B[i])
+ C[6, i] = T.max(T.cast(0, dtype), A[i])
+ C[7, i] = T.max(A[i], T.cast(0, dtype))
+ C[8, i] = T.max(T.cast(nan, dtype), A[i])
+ C[9, i] = T.max(A[i], T.cast(nan, dtype))
+ # Rows 10-11: use max results inside arithmetic.
+ C[10, i] = T.max(A[i], B[i]) + T.cast(1, dtype)
+ C[11, i] = T.max(T.cast(0, dtype), A[i]) + T.cast(1, dtype)
+ # Rows 12-13: positive and negative infinity constants.
+ C[12, i] = T.min(A[i], T.cast(infinity, dtype))
+ C[13, i] = T.max(A[i], T.cast(-infinity, dtype))
+
+ target = {"kind": "llvm", "opt-level": 0}
+ if scalable:
+ target.update(mtriple="aarch64-linux-gnu", mattr=["+sve"])
+ compiled = tvm_tirx.tirx.compile(main, target=target)
+ # BF16 arithmetic is legalized to float32 before LLVM codegen.
+ llvm_type = {"float16": "half", "float32": "float", "float64": "double",
"bfloat16": "float"}[
+ dtype
+ ]
+ if scalable:
+ llvm_type = f"<vscale x 4 x {llvm_type}>"
+ elif lanes != 1:
+ llvm_type = f"<{lanes} x {llvm_type}>"
+ llvm_ir = compiled.mod.inspect_source("ll")
+ for intrinsic in ("minimum", "maximum"):
+ assert f"call {llvm_type} @llvm.{intrinsic}." in llvm_ir
+ if scalable:
+ return # Cross-compile SVE; execution requires SVE hardware.
+
+ np_dtype = pytest.importorskip("ml_dtypes").bfloat16 if dtype ==
"bfloat16" else dtype
+ a = np.array([np.nan, 1, np.nan, 0.0, -0.0, 3, 2, -5, np.inf, -np.inf, 1,
-1], dtype=np_dtype)
+ b = np.array([1, np.nan, np.nan, -0.0, 0.0, 2, 2, -4, 1, 1, np.inf,
-np.inf], dtype=np_dtype)
+ out = tvm_tirx.runtime.empty((14, n), dtype)
+ compiled(tvm_tirx.runtime.tensor(a), tvm_tirx.runtime.tensor(b), out)
+
+ zero, nan_array = np.zeros_like(a), np.full_like(a, np.nan)
+ operand_pairs = ((a, b), (zero, a), (a, zero), (nan_array, a), (a,
nan_array))
+ expected_by_op = {}
+ for name, compare in (("min", np.less), ("max", np.greater)):
+ values = []
+ for lhs, rhs in operand_pairs:
+ take_lhs = compare(lhs, rhs) | np.isnan(lhs)
+ # IEEE minimum selects -0 and maximum selects +0 on a zero tie.
+ both_zero = (lhs == 0) & (rhs == 0)
+ preferred_zero = np.signbit(lhs) if name == "min" else
~np.signbit(lhs)
+ take_lhs |= both_zero & preferred_zero
+ values.append(np.where(take_lhs, lhs, rhs))
+ expected_by_op[name] = values
+ expected = expected_by_op["min"] + expected_by_op["max"]
+ max_ab, max_zero_a = expected_by_op["max"][:2]
+ expected.extend([max_ab + np.ones_like(a), max_zero_a + np.ones_like(a)])
+ # min(a, +Inf) and max(a, -Inf) both preserve a, including NaNs and signed
zero.
+ expected.extend([a, a])
+ expected = np.asarray(expected, dtype=np_dtype)
+ got = out.numpy()
+ nan_mask = np.isnan(expected)
+ np.testing.assert_array_equal(np.isnan(got), nan_mask)
+ # NaN payloads may be canonicalized; preserve exact non-NaN values and
signed zero.
+ uint_dtype = f"uint{a.dtype.itemsize * 8}"
+ np.testing.assert_array_equal(
+ got.view(uint_dtype)[~nan_mask], expected.view(uint_dtype)[~nan_mask]
+ )
+
+
if __name__ == "__main__":
tvm_tirx.testing.main()
diff --git a/tests/python/tirx/base/test_imm_values.py
b/tests/python/tirx/base/test_imm_values.py
index 1f12b304..4c8904d8 100644
--- a/tests/python/tirx/base/test_imm_values.py
+++ b/tests/python/tirx/base/test_imm_values.py
@@ -121,6 +121,26 @@ def test_tir_integer_cast_folding():
assert scalable.ty == tvm_tirx.ir.PrimType("int8xvscalex4")
[email protected]("dtype", ["float16", "bfloat16", "float32",
"float64"])
+def test_min_max_floatimm_const_fold(dtype):
+ for operation in (tirx.min, tirx.max):
+ for lhs, rhs, expected_side in (
+ (np.nan, 1.0, "lhs"),
+ (1.0, np.nan, "rhs"),
+ (np.nan, np.nan, "lhs"),
+ (0.0, -0.0, "rhs"),
+ (-0.0, 0.0, "rhs"),
+ ):
+ a, b = tirx.const(lhs, dtype), tirx.const(rhs, dtype)
+ result = operation(a, b)
+ expected = a if expected_side == "lhs" else b
+ if math.isnan(expected.value):
+ assert math.isnan(result.value)
+ else:
+ assert result.value == expected.value
+ assert np.signbit(result.value) == np.signbit(expected.value)
+
+
def compare_float_value(value, expect, msg):
if math.isfinite(value):
assert np.abs(value - expect) < 1e-5, f"{value} vs {expect}, {msg}"
