From d1e796891c283e25468fa6b1842ccbafdc34a375 Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Mon, 27 Jul 2026 02:25:43 -0400 Subject: [PATCH 1/4] [FIX][CODEGEN] Preserve NaNs in floating-point max --- src/target/llvm/codegen_llvm.cc | 10 ++++- src/target/source/codegen_c_host.cc | 18 ++++++++- tests/python/codegen/test_target_codegen.py | 41 +++++++++++++++++++++ 3 files changed, 67 insertions(+), 2 deletions(-) diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc index 97bd1b0f2644..e87ae8f22a3a 100644 --- a/src/target/llvm/codegen_llvm.cc +++ b/src/target/llvm/codegen_llvm.cc @@ -1656,7 +1656,15 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const MinNode* op) { llvm::Value* CodeGenLLVM::VisitExpr_(const MaxNode* 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); + PrimType dtype(op->a.ty()->dtype); + llvm::Value* take_a = CreateGT(dtype, a, b); + if (dtype.MatchesCode(DLDataTypeCode::kDLFloat)) { + // Keep the ordered comparison so a NaN in b selects b, then explicitly + // select a when a is NaN. This also retains the existing second-operand + // tie behavior, including for signed zero. + take_a = builder_->CreateOr(take_a, builder_->CreateFCmpUNO(a, a)); + } + return builder_->CreateSelect(take_a, a, b); } llvm::Value* CodeGenLLVM::VisitExpr_(const EQNode* op) { diff --git a/src/target/source/codegen_c_host.cc b/src/target/source/codegen_c_host.cc index 709b3ec4a9e8..a0c93ea178bd 100644 --- a/src/target/source/codegen_c_host.cc +++ b/src/target/source/codegen_c_host.cc @@ -351,7 +351,23 @@ void CodeGenCHost::VisitExpr_(const MinNode* op, std::ostream& os) { // NOLINT( } void CodeGenCHost::VisitExpr_(const MaxNode* op, std::ostream& os) { // NOLINT(*) - PrintTernaryCondExpr(op, ">", os); + PrimType dtype = op->ty.as_or_throw(); + if (!dtype.MatchesCode(DLDataTypeCode::kDLFloat)) { + PrintTernaryCondExpr(op, ">", os); + return; + } + + std::ostringstream temp_a; + VisitExpr(op->a, temp_a); + std::string a_id = SSAGetID(temp_a.str(), op->a.ty()); + std::ostringstream temp_b; + VisitExpr(op->b, temp_b); + std::string b_id = SSAGetID(temp_b.str(), op->b.ty()); + + // Preserve NaNs from either operand while retaining the existing behavior + // of selecting the second operand when both operands compare equal. + os << "((" << a_id << ") > (" << b_id << ") ? (" << a_id << ") : ((" << a_id << ") == (" << a_id + << ") ? (" << b_id << ") : (" << a_id << ")))"; } template diff --git a/tests/python/codegen/test_target_codegen.py b/tests/python/codegen/test_target_codegen.py index 7157ae0f69bf..579b100647a8 100644 --- a/tests/python/codegen/test_target_codegen.py +++ b/tests/python/codegen/test_target_codegen.py @@ -162,5 +162,46 @@ def test_loop_step( assert c_result[i] == 0.0 +@pytest.mark.parametrize( + "target,dtype,uint_dtype,nan_a,nan_b", + [ + ("llvm", "float16", "uint16", 0x7E11, 0x7E22), + ("c", "float32", "uint32", 0x7FC00011, 0x7FC00022), + ("llvm", "float32", "uint32", 0x7FC00011, 0x7FC00022), + ("c", "float64", "uint64", 0x7FF8000000000011, 0x7FF8000000000022), + ("llvm", "float64", "uint64", 0x7FF8000000000011, 0x7FF8000000000022), + ], +) +def test_max_nan_preserving(target, dtype, uint_dtype, nan_a, nan_b): + if target != "c" and not tvm.testing.device_enabled(target): + pytest.skip(f"{target} not enabled") + + @T.prim_func(s_tir=True) + def max_func( + A: T.Buffer((8,), dtype), + B: T.Buffer((8,), dtype), + C: T.Buffer((8,), dtype), + ): + T.func_attr({"tirx.noalias": True}) + for i in range(8): + C[i] = T.max(A[i], B[i]) + + a_np = np.array([0.0, 1.0, 0.0, 0.0, -0.0, 3.0, 2.0, -5.0], dtype=dtype) + b_np = np.array([1.0, 0.0, 0.0, -0.0, 0.0, 2.0, 2.0, -4.0], dtype=dtype) + a_bits = a_np.view(uint_dtype) + b_bits = b_np.view(uint_dtype) + a_bits[[0, 2]] = nan_a + b_bits[[1, 2]] = nan_b + + dev = tvm.cpu() + a = tvm.runtime.tensor(a_np, dev) + b = tvm.runtime.tensor(b_np, dev) + c = tvm.runtime.empty((8,), dtype, dev) + tvm.compile(max_func, target=target)(a, b, c) + + expected = np.where((a_np > b_np) | np.isnan(a_np), a_np, b_np) + np.testing.assert_array_equal(c.numpy().view(uint_dtype), expected.view(uint_dtype)) + + if __name__ == "__main__": tvm.testing.main() From 8482783e117c67dec9c0a50576fbd36127dfe4ad Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Mon, 27 Jul 2026 02:53:25 -0400 Subject: [PATCH 2/4] [FIX][CODEGEN] Preserve NaNs in floating-point min --- src/target/llvm/codegen_llvm.cc | 10 ++++++- src/target/source/codegen_c_host.cc | 30 ++++++++------------- tests/python/codegen/test_target_codegen.py | 19 ++++++++++--- 3 files changed, 36 insertions(+), 23 deletions(-) diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc index e87ae8f22a3a..76b80b9eb7ca 100644 --- a/src/target/llvm/codegen_llvm.cc +++ b/src/target/llvm/codegen_llvm.cc @@ -1650,7 +1650,15 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const ModNode* op) { llvm::Value* CodeGenLLVM::VisitExpr_(const MinNode* 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); + PrimType dtype(op->a.ty()->dtype); + llvm::Value* take_a = CreateLT(dtype, a, b); + if (dtype.MatchesCode(DLDataTypeCode::kDLFloat)) { + // Keep the ordered comparison so a NaN in b selects b, then explicitly + // select a when a is NaN. This also retains the existing second-operand + // tie behavior, including for signed zero. + take_a = builder_->CreateOr(take_a, builder_->CreateFCmpUNO(a, a)); + } + return builder_->CreateSelect(take_a, a, b); } llvm::Value* CodeGenLLVM::VisitExpr_(const MaxNode* op) { diff --git a/src/target/source/codegen_c_host.cc b/src/target/source/codegen_c_host.cc index a0c93ea178bd..daca461ba333 100644 --- a/src/target/source/codegen_c_host.cc +++ b/src/target/source/codegen_c_host.cc @@ -351,23 +351,7 @@ void CodeGenCHost::VisitExpr_(const MinNode* op, std::ostream& os) { // NOLINT( } void CodeGenCHost::VisitExpr_(const MaxNode* op, std::ostream& os) { // NOLINT(*) - PrimType dtype = op->ty.as_or_throw(); - if (!dtype.MatchesCode(DLDataTypeCode::kDLFloat)) { - PrintTernaryCondExpr(op, ">", os); - return; - } - - std::ostringstream temp_a; - VisitExpr(op->a, temp_a); - std::string a_id = SSAGetID(temp_a.str(), op->a.ty()); - std::ostringstream temp_b; - VisitExpr(op->b, temp_b); - std::string b_id = SSAGetID(temp_b.str(), op->b.ty()); - - // Preserve NaNs from either operand while retaining the existing behavior - // of selecting the second operand when both operands compare equal. - os << "((" << a_id << ") > (" << b_id << ") ? (" << a_id << ") : ((" << a_id << ") == (" << a_id - << ") ? (" << b_id << ") : (" << a_id << ")))"; + PrintTernaryCondExpr(op, ">", os); } template @@ -380,8 +364,16 @@ inline void CodeGenCHost::PrintTernaryCondExpr(const T* op, const char* compare, VisitExpr(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 << "))"; + PrimType dtype = op->ty.template as_or_throw(); + if (dtype.MatchesCode(DLDataTypeCode::kDLFloat)) { + // Preserve NaNs from either operand while retaining the existing behavior + // of selecting the second operand when both operands compare equal. + os << "((" << a_id << ") " << compare << " (" << b_id << ") ? (" << a_id << ") : ((" << a_id + << ") == (" << a_id << ") ? (" << b_id << ") : (" << a_id << ")))"; + } else { + os << "((" << a_id << ") " << compare << " (" << b_id << ") " + << "? (" << a_id << ") : (" << b_id << "))"; + } } ffi::Module BuildCHost(IRModule mod, Target target) { diff --git a/tests/python/codegen/test_target_codegen.py b/tests/python/codegen/test_target_codegen.py index 579b100647a8..a3335f3d9738 100644 --- a/tests/python/codegen/test_target_codegen.py +++ b/tests/python/codegen/test_target_codegen.py @@ -172,7 +172,8 @@ def test_loop_step( ("llvm", "float64", "uint64", 0x7FF8000000000011, 0x7FF8000000000022), ], ) -def test_max_nan_preserving(target, dtype, uint_dtype, nan_a, nan_b): +@pytest.mark.parametrize("operation", ["min", "max"]) +def test_min_max_nan_preserving(target, dtype, uint_dtype, nan_a, nan_b, operation): if target != "c" and not tvm.testing.device_enabled(target): pytest.skip(f"{target} not enabled") @@ -186,6 +187,16 @@ def max_func( for i in range(8): C[i] = T.max(A[i], B[i]) + @T.prim_func(s_tir=True) + def min_func( + A: T.Buffer((8,), dtype), + B: T.Buffer((8,), dtype), + C: T.Buffer((8,), dtype), + ): + T.func_attr({"tirx.noalias": True}) + for i in range(8): + C[i] = T.min(A[i], B[i]) + a_np = np.array([0.0, 1.0, 0.0, 0.0, -0.0, 3.0, 2.0, -5.0], dtype=dtype) b_np = np.array([1.0, 0.0, 0.0, -0.0, 0.0, 2.0, 2.0, -4.0], dtype=dtype) a_bits = a_np.view(uint_dtype) @@ -197,9 +208,11 @@ def max_func( a = tvm.runtime.tensor(a_np, dev) b = tvm.runtime.tensor(b_np, dev) c = tvm.runtime.empty((8,), dtype, dev) - tvm.compile(max_func, target=target)(a, b, c) + func = min_func if operation == "min" else max_func + tvm.compile(func, target=target)(a, b, c) - expected = np.where((a_np > b_np) | np.isnan(a_np), a_np, b_np) + compare = a_np < b_np if operation == "min" else a_np > b_np + expected = np.where(compare | np.isnan(a_np), a_np, b_np) np.testing.assert_array_equal(c.numpy().view(uint_dtype), expected.view(uint_dtype)) From ba4784e9a640f8ff77a7d2ad92c162cf0d92568b Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Mon, 27 Jul 2026 03:24:24 -0400 Subject: [PATCH 3/4] [REFACTOR][CODEGEN] Simplify NaN-preserving C min/max --- src/target/source/codegen_c_host.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/target/source/codegen_c_host.cc b/src/target/source/codegen_c_host.cc index daca461ba333..a468f0163997 100644 --- a/src/target/source/codegen_c_host.cc +++ b/src/target/source/codegen_c_host.cc @@ -368,8 +368,8 @@ inline void CodeGenCHost::PrintTernaryCondExpr(const T* op, const char* compare, if (dtype.MatchesCode(DLDataTypeCode::kDLFloat)) { // Preserve NaNs from either operand while retaining the existing behavior // of selecting the second operand when both operands compare equal. - os << "((" << a_id << ") " << compare << " (" << b_id << ") ? (" << a_id << ") : ((" << a_id - << ") == (" << a_id << ") ? (" << b_id << ") : (" << a_id << ")))"; + os << "(((" << a_id << ") " << compare << " (" << b_id << ") || (" << a_id << ") != (" << a_id + << ")) ? (" << a_id << ") : (" << b_id << "))"; } else { os << "((" << a_id << ") " << compare << " (" << b_id << ") " << "? (" << a_id << ") : (" << b_id << "))"; From c8f0e740e312bee6da68fc1542c577e91161fd82 Mon Sep 17 00:00:00 2001 From: tlopex <820958424@qq.com> Date: Mon, 3 Aug 2026 18:53:43 -0400 Subject: [PATCH 4/4] [FIX][CODEGEN] Simplify min/max with constant operands --- src/arith/const_fold.h | 4 +- src/target/llvm/codegen_llvm.cc | 59 ++++++-- src/target/min_max_utils.h | 56 +++++++ src/target/source/codegen_c_host.cc | 39 ++++- src/target/source/codegen_c_host.h | 5 +- tests/python/codegen/test_target_codegen.py | 137 +++++++++++++++--- tests/python/tirx-base/test_tir_imm_values.py | 22 +++ 7 files changed, 277 insertions(+), 45 deletions(-) create mode 100644 src/target/min_max_utils.h diff --git a/src/arith/const_fold.h b/src/arith/const_fold.h index f7f46fae78a4..bb2ff7ca150b 100644 --- a/src/arith/const_fold.h +++ b/src/arith/const_fold.h @@ -356,7 +356,7 @@ inline ffi::Optional TryConstFold(PrimExpr a, PrimExpr b) { TVM_ARITH_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; @@ -367,7 +367,7 @@ inline ffi::Optional TryConstFold(PrimExpr a, PrimExpr b) { TVM_ARITH_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/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc index 76b80b9eb7ca..5ba1f6f9f461 100644 --- a/src/target/llvm/codegen_llvm.cc +++ b/src/target/llvm/codegen_llvm.cc @@ -93,6 +93,7 @@ #include "../../arith/pattern_match.h" #include "../build_common.h" +#include "../min_max_utils.h" #include "codegen_params.h" #include "llvm_instance.h" @@ -1651,12 +1652,29 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const MinNode* op) { llvm::Value* a = MakeValue(op->a); llvm::Value* b = MakeValue(op->b); PrimType dtype(op->a.ty()->dtype); - llvm::Value* take_a = CreateLT(dtype, a, b); - if (dtype.MatchesCode(DLDataTypeCode::kDLFloat)) { - // Keep the ordered comparison so a NaN in b selects b, then explicitly - // select a when a is NaN. This also retains the existing second-operand - // tie behavior, including for signed zero. - take_a = builder_->CreateOr(take_a, builder_->CreateFCmpUNO(a, a)); + llvm::Value* take_a; + if (!dtype.MatchesCode(DLDataTypeCode::kDLFloat)) { + take_a = CreateLT(dtype, a, b); + } else { + ConstFloatKind a_kind = GetConstFloatKind(op->a); + ConstFloatKind b_kind = GetConstFloatKind(op->b); + if (a_kind == ConstFloatKind::kNaN) { + return a; + } else if (a_kind == ConstFloatKind::kNonNaN) { + // The ordered comparison already selects b if b is NaN. + take_a = CreateLT(dtype, a, b); + } else if (b_kind == ConstFloatKind::kNaN) { + take_a = builder_->CreateFCmpUNO(a, a); + } else if (b_kind == ConstFloatKind::kNonNaN) { + // With a known non-NaN rhs, an unordered comparison is true exactly + // when a < b or a is NaN. + take_a = builder_->CreateFCmpULT(a, b); + } else { + // Keep the ordered comparison so a NaN in b selects b, then explicitly + // select a when a is NaN. This also retains the existing second-operand + // tie behavior, including for signed zero. + take_a = builder_->CreateOr(CreateLT(dtype, a, b), builder_->CreateFCmpUNO(a, a)); + } } return builder_->CreateSelect(take_a, a, b); } @@ -1665,12 +1683,29 @@ llvm::Value* CodeGenLLVM::VisitExpr_(const MaxNode* op) { llvm::Value* a = MakeValue(op->a); llvm::Value* b = MakeValue(op->b); PrimType dtype(op->a.ty()->dtype); - llvm::Value* take_a = CreateGT(dtype, a, b); - if (dtype.MatchesCode(DLDataTypeCode::kDLFloat)) { - // Keep the ordered comparison so a NaN in b selects b, then explicitly - // select a when a is NaN. This also retains the existing second-operand - // tie behavior, including for signed zero. - take_a = builder_->CreateOr(take_a, builder_->CreateFCmpUNO(a, a)); + llvm::Value* take_a; + if (!dtype.MatchesCode(DLDataTypeCode::kDLFloat)) { + take_a = CreateGT(dtype, a, b); + } else { + ConstFloatKind a_kind = GetConstFloatKind(op->a); + ConstFloatKind b_kind = GetConstFloatKind(op->b); + if (a_kind == ConstFloatKind::kNaN) { + return a; + } else if (a_kind == ConstFloatKind::kNonNaN) { + // The ordered comparison already selects b if b is NaN. + take_a = CreateGT(dtype, a, b); + } else if (b_kind == ConstFloatKind::kNaN) { + take_a = builder_->CreateFCmpUNO(a, a); + } else if (b_kind == ConstFloatKind::kNonNaN) { + // With a known non-NaN rhs, an unordered comparison is true exactly + // when a > b or a is NaN. + take_a = builder_->CreateFCmpUGT(a, b); + } else { + // Keep the ordered comparison so a NaN in b selects b, then explicitly + // select a when a is NaN. This also retains the existing second-operand + // tie behavior, including for signed zero. + take_a = builder_->CreateOr(CreateGT(dtype, a, b), builder_->CreateFCmpUNO(a, a)); + } } return builder_->CreateSelect(take_a, a, b); } diff --git a/src/target/min_max_utils.h b/src/target/min_max_utils.h new file mode 100644 index 000000000000..d1beab3ece52 --- /dev/null +++ b/src/target/min_max_utils.h @@ -0,0 +1,56 @@ +/* + * 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_TARGET_MIN_MAX_UTILS_H_ +#define TVM_TARGET_MIN_MAX_UTILS_H_ + +#include + +#include + +namespace tvm { +namespace codegen { + +enum class ConstFloatKind { + kNotConst, + kNonNaN, + kNaN, +}; + +inline ConstFloatKind GetConstFloatKind(const PrimExpr& expr) { + const FloatImmNode* value = expr.as(); + if (const auto* broadcast = expr.as()) { + // MakeConst represents vector-valued constants as a broadcast of a + // scalar immediate, including fixed-length and scalable vectors. + value = broadcast->value.as(); + } + if (value == nullptr) { + return ConstFloatKind::kNotConst; + } + return std::isnan(value->value) ? ConstFloatKind::kNaN : ConstFloatKind::kNonNaN; +} + +} // namespace codegen +} // namespace tvm + +#endif // TVM_TARGET_MIN_MAX_UTILS_H_ diff --git a/src/target/source/codegen_c_host.cc b/src/target/source/codegen_c_host.cc index a468f0163997..ac1a707b2cc4 100644 --- a/src/target/source/codegen_c_host.cc +++ b/src/target/source/codegen_c_host.cc @@ -32,6 +32,8 @@ #include #include +#include "../min_max_utils.h" + namespace tvm { namespace codegen { @@ -347,16 +349,25 @@ void CodeGenCHost::VisitStmt_(const AssertStmtNode* op) { // NOLINT(*) } void CodeGenCHost::VisitExpr_(const MinNode* op, std::ostream& os) { // NOLINT(*) - PrintTernaryCondExpr(op, "<", os); + PrintTernaryCondExpr(op, "<", ">=", os); } void CodeGenCHost::VisitExpr_(const MaxNode* op, std::ostream& os) { // NOLINT(*) - PrintTernaryCondExpr(op, ">", os); + PrintTernaryCondExpr(op, ">", "<=", os); } template inline void CodeGenCHost::PrintTernaryCondExpr(const T* op, const char* compare, + const char* reverse_compare, std::ostream& os) { // NOLINT(*) + PrimType dtype = op->ty.template as_or_throw(); + ConstFloatKind a_kind = ConstFloatKind::kNotConst; + ConstFloatKind b_kind = ConstFloatKind::kNotConst; + if (dtype.MatchesCode(DLDataTypeCode::kDLFloat)) { + a_kind = GetConstFloatKind(op->a); + b_kind = GetConstFloatKind(op->b); + } + std::ostringstream temp_a; VisitExpr(op->a, temp_a); std::string a_id = SSAGetID(temp_a.str(), op->a.ty()); @@ -364,12 +375,26 @@ inline void CodeGenCHost::PrintTernaryCondExpr(const T* op, const char* compare, VisitExpr(op->b, temp_b); std::string b_id = SSAGetID(temp_b.str(), op->b.ty()); - PrimType dtype = op->ty.template as_or_throw(); if (dtype.MatchesCode(DLDataTypeCode::kDLFloat)) { - // Preserve NaNs from either operand while retaining the existing behavior - // of selecting the second operand when both operands compare equal. - os << "(((" << a_id << ") " << compare << " (" << b_id << ") || (" << a_id << ") != (" << a_id - << ")) ? (" << a_id << ") : (" << b_id << "))"; + if (a_kind == ConstFloatKind::kNaN) { + os << "(" << a_id << ")"; + } else if (a_kind == ConstFloatKind::kNonNaN) { + os << "((" << a_id << ") " << compare << " (" << b_id << ") ? (" << a_id << ") : (" << b_id + << "))"; + } else if (b_kind == ConstFloatKind::kNaN) { + os << "((" << a_id << ") != (" << a_id << ") ? (" << a_id << ") : (" << b_id << "))"; + } else if (b_kind == ConstFloatKind::kNonNaN) { + // Reversing the select avoids a separate NaN test: if a is NaN, the + // ordered comparison is false and a is selected. Equality still + // selects b, preserving the existing signed-zero behavior. + os << "((" << a_id << ") " << reverse_compare << " (" << b_id << ") ? (" << b_id << ") : (" + << a_id << "))"; + } else { + // Preserve NaNs from either operand while retaining the existing behavior + // of selecting the second operand when both operands compare equal. + os << "(((" << a_id << ") " << compare << " (" << b_id << ") || (" << a_id << ") != (" << a_id + << ")) ? (" << a_id << ") : (" << b_id << "))"; + } } else { os << "((" << a_id << ") " << compare << " (" << b_id << ") " << "? (" << a_id << ") : (" << b_id << "))"; diff --git a/src/target/source/codegen_c_host.h b/src/target/source/codegen_c_host.h index e9b89e6e3446..4f7734028d95 100644 --- a/src/target/source/codegen_c_host.h +++ b/src/target/source/codegen_c_host.h @@ -95,11 +95,12 @@ class CodeGenCHost : public CodeGenC { * \brief Print ternary conditional operator implementing binary `op` * Forces the operands to be in SSA form. * \param op binary operator being expressed - * \param compare string representation of comparison operator + * \param compare string representation of the strict comparison operator + * \param reverse_compare string representation of the reverse non-strict comparison operator * \param os stream reference to print into */ template - inline void PrintTernaryCondExpr(const T* op, const char* compare, + inline void PrintTernaryCondExpr(const T* op, const char* compare, const char* reverse_compare, std::ostream& os); // NOLINT(*) }; diff --git a/tests/python/codegen/test_target_codegen.py b/tests/python/codegen/test_target_codegen.py index a3335f3d9738..b3aa532425be 100644 --- a/tests/python/codegen/test_target_codegen.py +++ b/tests/python/codegen/test_target_codegen.py @@ -162,20 +162,9 @@ def test_loop_step( assert c_result[i] == 0.0 -@pytest.mark.parametrize( - "target,dtype,uint_dtype,nan_a,nan_b", - [ - ("llvm", "float16", "uint16", 0x7E11, 0x7E22), - ("c", "float32", "uint32", 0x7FC00011, 0x7FC00022), - ("llvm", "float32", "uint32", 0x7FC00011, 0x7FC00022), - ("c", "float64", "uint64", 0x7FF8000000000011, 0x7FF8000000000022), - ("llvm", "float64", "uint64", 0x7FF8000000000011, 0x7FF8000000000022), - ], -) -@pytest.mark.parametrize("operation", ["min", "max"]) -def test_min_max_nan_preserving(target, dtype, uint_dtype, nan_a, nan_b, operation): - if target != "c" and not tvm.testing.device_enabled(target): - pytest.skip(f"{target} not enabled") +def test_min_max_nan_preserving(): + dtype = "float32" + uint_dtype = "uint32" @T.prim_func(s_tir=True) def max_func( @@ -201,19 +190,123 @@ def min_func( b_np = np.array([1.0, 0.0, 0.0, -0.0, 0.0, 2.0, 2.0, -4.0], dtype=dtype) a_bits = a_np.view(uint_dtype) b_bits = b_np.view(uint_dtype) - a_bits[[0, 2]] = nan_a - b_bits[[1, 2]] = nan_b + a_bits[[0, 2]] = 0x7FC00011 + b_bits[[1, 2]] = 0x7FC00022 dev = tvm.cpu() a = tvm.runtime.tensor(a_np, dev) b = tvm.runtime.tensor(b_np, dev) - c = tvm.runtime.empty((8,), dtype, dev) - func = min_func if operation == "min" else max_func - tvm.compile(func, target=target)(a, b, c) + targets = ["c"] + if tvm.testing.device_enabled("llvm"): + targets.append("llvm") + + for target in targets: + for operation, func in [("min", min_func), ("max", max_func)]: + c = tvm.runtime.empty((8,), dtype, dev) + tvm.compile(func, target=target)(a, b, c) + compare = a_np < b_np if operation == "min" else a_np > b_np + expected = np.where(compare | np.isnan(a_np), a_np, b_np) + np.testing.assert_array_equal(c.numpy().view(uint_dtype), expected.view(uint_dtype)) + + +def _make_min_max_func(operation, const_side, dtype="float32", extent=1, const_value=0): + a_buffer = tvm.tirx.decl_buffer((extent,), dtype, name="A") + c_buffer = tvm.tirx.decl_buffer((extent,), dtype, name="C") + index = tvm.tirx.Var("i", "int32") + constant = tvm.tirx.const(const_value, dtype) + dynamic = tvm.tirx.BufferLoad(a_buffer, [index]) + lhs, rhs = (constant, dynamic) if const_side == "lhs" else (dynamic, constant) + result = {"min": tvm.tirx.min, "max": tvm.tirx.max}[operation](lhs, rhs) + body = tvm.tirx.For( + index, + 0, + extent, + tvm.tirx.ForKind.SERIAL, + tvm.tirx.BufferStore(c_buffer, result, [index]), + ) + return tvm.tirx.PrimFunc([a_buffer, c_buffer], body).with_attr("global_symbol", "main") + + +@pytest.mark.parametrize( + "target,operation,const_side", + [ + ("c", "min", "lhs"), + ("c", "max", "rhs"), + ("llvm", "min", "rhs"), + ("llvm", "max", "lhs"), + ], +) +def test_min_max_float_imm_operand(target, operation, const_side): + if target != "c" and not tvm.testing.device_enabled(target): + pytest.skip(f"{target} not enabled") - compare = a_np < b_np if operation == "min" else a_np > b_np - expected = np.where(compare | np.isnan(a_np), a_np, b_np) - np.testing.assert_array_equal(c.numpy().view(uint_dtype), expected.view(uint_dtype)) + func = _make_min_max_func(operation, const_side, extent=7) + compile_target = {"kind": "llvm", "opt-level": 0} if target == "llvm" else target + compiled = tvm.compile(func, target=compile_target) + + a_np = np.array([np.nan, -0.0, 0.0, -1.0, 1.0, np.inf, -np.inf], dtype="float32") + c = tvm.runtime.empty(a_np.shape, "float32", tvm.cpu()) + compiled(tvm.runtime.tensor(a_np), c) + + zero = np.zeros_like(a_np) + lhs, rhs = (zero, a_np) if const_side == "lhs" else (a_np, zero) + compare = lhs < rhs if operation == "min" else lhs > rhs + expected = np.where(compare | np.isnan(lhs), lhs, rhs) + np.testing.assert_array_equal(c.numpy().view("uint32"), expected.view("uint32")) + + if target == "llvm": + predicate = { + ("min", "lhs"): "olt", + ("min", "rhs"): "ult", + ("max", "lhs"): "ogt", + ("max", "rhs"): "ugt", + }[operation, const_side] + llvm_ir = compiled.mod.inspect_source("ll") + assert f"fcmp {predicate}" in llvm_ir + assert "fcmp uno" not in llvm_ir + else: + predicate = { + ("min", "lhs"): " < ", + ("min", "rhs"): " >= ", + ("max", "lhs"): " > ", + ("max", "rhs"): " <= ", + }[operation, const_side] + result_lines = [ + line + for line in compiled.mod.inspect_source().splitlines() + if " = " in line and " ? " in line + ] + assert len(result_lines) == 1 + assert predicate in result_lines[0] + assert "||" not in result_lines[0] + assert "!=" not in result_lines[0] + + +def test_llvm_min_max_broadcast_float_imm_operand(): + if not tvm.testing.device_enabled("llvm"): + pytest.skip("llvm not enabled") + + func = _make_min_max_func("min", "rhs", dtype="float32x4") + llvm_ir = tvm.compile(func, target={"kind": "llvm", "opt-level": 0}).mod.inspect_source("ll") + assert "fcmp ult <4 x float>" in llvm_ir + assert "fcmp uno" not in llvm_ir + + +def test_llvm_min_max_nan_float_imm_operand(): + if not tvm.testing.device_enabled("llvm"): + pytest.skip("llvm not enabled") + + for operation, const_side in [("min", "lhs"), ("max", "rhs")]: + func = _make_min_max_func(operation, const_side, const_value=np.nan) + llvm_ir = tvm.compile(func, target={"kind": "llvm", "opt-level": 0}).mod.inspect_source( + "ll" + ) + fcmp_lines = [line for line in llvm_ir.splitlines() if " fcmp " in line] + if const_side == "lhs": + assert not fcmp_lines + else: + assert len(fcmp_lines) == 1 + assert " fcmp uno " in fcmp_lines[0] if __name__ == "__main__": diff --git a/tests/python/tirx-base/test_tir_imm_values.py b/tests/python/tirx-base/test_tir_imm_values.py index 2d9048f7d1d5..2a85dd289cf0 100644 --- a/tests/python/tirx-base/test_tir_imm_values.py +++ b/tests/python/tirx-base/test_tir_imm_values.py @@ -147,6 +147,28 @@ def test_tir_special_floatimms(dtype, literal): compare_float_value(x.value, literal, "imm value should match feed value") +def test_tir_min_max_floatimm_const_fold(): + dtype = "float32" + uint_dtype = "uint32" + lhs_nan, rhs_nan = np.array([0x7FC00011, 0x7FC00022], dtype=uint_dtype).view(dtype) + cases = { + "lhs_nan": (lhs_nan, 1.0, "lhs"), + "rhs_nan": (1.0, rhs_nan, "rhs"), + "both_nan": (lhs_nan, rhs_nan, "lhs"), + "signed_zero_tie": (0.0, -0.0, "rhs"), + } + + for operation_name, operation in [("min", tirx.min), ("max", tirx.max)]: + for case, (lhs_value, rhs_value, expected_side) in cases.items(): + lhs = tirx.const(lhs_value, dtype) + rhs = tirx.const(rhs_value, dtype) + result = operation(lhs, rhs) + expected = lhs if expected_side == "lhs" else rhs + result_bits = np.asarray(result.value, dtype=dtype).view(uint_dtype).item() + expected_bits = np.asarray(expected.value, dtype=dtype).view(uint_dtype).item() + assert result_bits == expected_bits, f"{operation_name}: {case}" + + @pytest.mark.skipif(not env.has_llvm(), reason="need llvm") def test_tir_too_large_literal_f64(): # Behavior check: if literal f64 value is out of dtype range, the