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

src/target/llvm/codegen_llvm.cc

@@ -1650,13 +1650,29 @@ 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) {
   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) {

src/target/source/codegen_c_host.cc

@@ -364,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<PrimType>();
+  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 << "))";
+  } else {
+    os << "((" << a_id << ") " << compare << " (" << b_id << ") "
+       << "? (" << a_id << ") : (" << b_id << "))";
+  }
 }
 
 ffi::Module BuildCHost(IRModule mod, Target target) {

tests/python/codegen/test_target_codegen.py

@@ -162,5 +162,59 @@ 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")
+
+    @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]] = 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)
+    func = min_func if operation == "min" else max_func
+    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))
+
+
 if __name__ == "__main__":
     tvm.testing.main()