[Fix][Codegen] Preserve NaNs in floating-point min and max

src/arith/const_fold.h

@@ -356,7 +356,7 @@ inline ffi::Optional<PrimExpr> TryConstFold<tirx::Min>(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<PrimExpr> TryConstFold<tirx::Max>(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;

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) {

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 <tvm/tirx/expr.h>
+
+#include <cmath>
+
+namespace tvm {
+namespace codegen {
+
+enum class ConstFloatKind {
+  kNotConst,
+  kNonNaN,
+  kNaN,
+};
+
+inline ConstFloatKind GetConstFloatKind(const PrimExpr& expr) {
+  const FloatImmNode* value = expr.as<FloatImmNode>();
+  if (const auto* broadcast = expr.as<tirx::BroadcastNode>()) {
+    // MakeConst represents vector-valued constants as a broadcast of a
+    // scalar immediate, including fixed-length and scalable vectors.
+    value = broadcast->value.as<FloatImmNode>();
+  }
+  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_

src/target/source/codegen_c_host.cc

@@ -32,6 +32,8 @@
 #include <utility>
 #include <vector>
 
+#include "../min_max_utils.h"
+
 namespace tvm {
 namespace codegen {
 
@@ -347,25 +349,56 @@ 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 <typename T>
 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<PrimType>();
+  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());
   std::ostringstream temp_b;
   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) {

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 <typename T>
-  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(*)
 };
 

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()