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 97bd1b0f2644..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" @@ -1650,13 +1651,63 @@ 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; + 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); } 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; + 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); } llvm::Value* CodeGenLLVM::VisitExpr_(const EQNode* op) { 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 709b3ec4a9e8..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,8 +375,30 @@ 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 << "))"; + if (dtype.MatchesCode(DLDataTypeCode::kDLFloat)) { + 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 << "))"; + } } ffi::Module BuildCHost(IRModule mod, Target target) { 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 7157ae0f69bf..b3aa532425be 100644 --- a/tests/python/codegen/test_target_codegen.py +++ b/tests/python/codegen/test_target_codegen.py @@ -162,5 +162,152 @@ def test_loop_step( assert c_result[i] == 0.0 +def test_min_max_nan_preserving(): + dtype = "float32" + uint_dtype = "uint32" + + @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]) + + @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) + b_bits = b_np.view(uint_dtype) + 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) + 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") + + 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__": tvm.testing.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