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

src/backend/cuda/codegen/codegen_cuda.cc

@@ -77,6 +77,16 @@ TVM_FFI_INLINE bool IsPackedFloat(const PrimType& ty) {
   return IsFloat8(ty) || IsFloat6(ty) || IsFloat4(ty);
 }
 
+// NaN-preserving min/max is emitted for float, half and bfloat16: the C
+// ternary below keeps the NaN operand instead of discarding it (apache/tvm
+// PR #20054). Integer min/max stays as it is.
+TVM_FFI_INLINE bool IsFloatMinMaxNanPreserving(const PrimType& ty) {
+  bool is_fp = ty.MatchesCode(DLDataTypeCode::kDLFloat) &&
+               (ty.bits() == 16 || ty.bits() == 32 || ty.bits() == 64);
+  bool is_bf16 = ty.MatchesCode(DLDataTypeCode::kDLBfloat) && ty.bits() == 16;
+  return is_fp || is_bf16;
+}
+
 }  // namespace
 
 std::string GetFP8Type(const PrimType& type_ty) {
@@ -686,6 +696,78 @@ void CodeGenCUDA::PrintVecConstructor(const PrimType& t, std::ostream& os) {
   PrintType(t, os);
 }
 
+// min/max for CUDA, keeping a NaN operand (apache/tvm PR #20054), for
+// float16/bfloat16/float32/float64. For a NaN-free pair this is min()/max();
+// the lhs-NaN clause fires on the first operand and the rhs-NaN case is
+// covered because the ordering compare is false when either operand is NaN.
+template <typename T>
+void CodeGenCUDA::PrintMinMaxNanPreservingImpl(const T* op, const char* opstr,
+                                               std::ostream& os) {  // NOLINT(*)
+  PrimType op_ty = op->ty.template as_or_throw<PrimType>();
+  const char* cmp = (opstr[0] == 'm' && opstr[1] == 'i') ? "<" : ">";
+  if (!IsFloatMinMaxNanPreserving(op_ty)) {
+    // Integer / non-float min/max keeps the base-codegen path, including the
+    // vectorized per-lane expansion for lanes > 1.
+    this->CodeGenC::Dispatch_(op, os);
+    return;
+  }
+  if (op_ty.lanes() == 1) {
+    // Bind both operands once (SSA), then reference the temporaries: the
+    // expressions must not be re-evaluated per clause. The bindings live in
+    // their own scope so a later statement that prints the same text (e.g.
+    // `red_buf[0] = max(red_buf[0], shuffle_down(...))` repeated by warp
+    // reduction) does not hit the cache and read the pre-write value.
+    int ssa_scope = BeginScope();
+    std::string va = this->SSAGetID(this->PrintExpr(op->a), op->a.ty());
+    std::string vb = this->SSAGetID(this->PrintExpr(op->b), op->b.ty());
+    os << "(((" << va << ' ' << cmp << ' ' << vb << ") || (" << va << " != " << va << ")) ? " << va
+       << " : " << vb << ")";
+    EndScope(ssa_scope);
+  } else {
+    this->PrintVecBinaryOpNanPreserving(std::string(opstr), op_ty, op->a, op->b, cmp, os);
+  }
+}
+
+void CodeGenCUDA::Dispatch_(const prim::MinNode* op, std::ostream& os) {  // NOLINT(*)
+  this->PrintMinMaxNanPreservingImpl(op, "min", os);
+}
+
+void CodeGenCUDA::Dispatch_(const prim::MaxNode* op, std::ostream& os) {  // NOLINT(*)
+  this->PrintMinMaxNanPreservingImpl(op, "max", os);
+}
+
+void CodeGenCUDA::PrintVecBinaryOpNanPreserving(const std::string& op, const PrimType& t,
+                                                PrimExpr lhs, PrimExpr rhs, const char* cmp,
+                                                std::ostream& os) {  // NOLINT(*)
+  std::string sret = name_supply_->FreshName("_");
+  this->PrintIndent();
+  this->PrintType(t, stream);
+  stream << ' ' << sret << ";\n";
+  int ssa_scope = BeginScope();
+  {
+    std::string vlhs = SSAGetID(PrintExpr(lhs), lhs.ty());
+    std::string vrhs = SSAGetID(PrintExpr(rhs), rhs.ty());
+    for (int i = 0, lanes = t.lanes(); i < lanes; ++i) {
+      std::ostringstream value_temp;
+      value_temp << "((";
+      PrintVecElemLoad(vlhs, lhs.ty(), i, value_temp);
+      value_temp << ' ' << cmp << ' ';
+      PrintVecElemLoad(vrhs, rhs.ty(), i, value_temp);
+      value_temp << ") || (";
+      PrintVecElemLoad(vlhs, lhs.ty(), i, value_temp);
+      value_temp << " != ";
+      PrintVecElemLoad(vlhs, lhs.ty(), i, value_temp);
+      value_temp << ")) ? ";
+      PrintVecElemLoad(vlhs, lhs.ty(), i, value_temp);
+      value_temp << " : ";
+      PrintVecElemLoad(vrhs, rhs.ty(), i, value_temp);
+      PrintVecElemStore(sret, t, i, value_temp.str());
+    }
+  }
+  EndScope(ssa_scope);
+  os << sret;
+}
+
 void CodeGenCUDA::PrintVecBinaryOp(const std::string& op, const PrimType& t, PrimExpr lhs,
                                    PrimExpr rhs, std::ostream& os) {  // NOLINT(*)
   // Declare the result.

src/backend/cuda/codegen/codegen_cuda.h

@@ -77,6 +77,12 @@ class CodeGenCUDA final : public CodeGenC {
   void Dispatch_(const FloatImmNode* op, std::ostream& os) final;
   void Dispatch_(const CallNode* op, std::ostream& os) final;
   void Dispatch_(const prim::CastNode* op, std::ostream& os) final;
+  void Dispatch_(const prim::MinNode* op, std::ostream& os) final;  // NOLINT(*)
+  void Dispatch_(const prim::MaxNode* op, std::ostream& os) final;  // NOLINT(*)
+  template <typename T>
+  void PrintMinMaxNanPreservingImpl(const T* op, const char* opstr, std::ostream& os);
+  void PrintVecBinaryOpNanPreserving(const std::string& op, const PrimType& t, PrimExpr lhs,
+                                     PrimExpr rhs, const char* cmp, std::ostream& os);
   void Dispatch_(const EvaluateNode* op) final;
   void Dispatch_(const ReturnNode* op) final;
   void Dispatch_(const AllocBufferNode* op) final;

tests/python/codegen/test_target_codegen_cuda.py

@@ -782,6 +782,185 @@ def run_and_check():
         run_test(*func, "float16")
 
 
+def _min_max_nan_module(op, dt, n=8, vectorize=False, composite=False):
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(A: T.Buffer((n,), dt), B: T.Buffer((n,), dt), C: T.Buffer((n,), dt)):
+            T.func_attr({"tirx.noalias": True})
+            for i0 in T.thread_binding(2, thread="blockIdx.x"):
+                if composite:
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(n, i0 * 4 + 0)
+                        C[v_i] = T.max(A[v_i], B[v_i]) + T.float32(1.0)
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(n, i0 * 4 + 1)
+                        C[v_i] = T.max(A[v_i], B[v_i]) + T.float32(1.0)
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(n, i0 * 4 + 2)
+                        C[v_i] = T.max(A[v_i], B[v_i]) + T.float32(1.0)
+                    with T.sblock("C"):
+                        v_i = T.axis.spatial(n, i0 * 4 + 3)
+                        C[v_i] = T.max(A[v_i], B[v_i]) + T.float32(1.0)
+                elif vectorize:
+                    for i1 in T.vectorized(4):
+                        with T.sblock("C"):
+                            v_i = T.axis.spatial(n, i0 * 4 + i1)
+                            C[v_i] = T.max(A[v_i], B[v_i]) if op == "max" else T.min(A[v_i], B[v_i])
+                else:
+                    for i1 in T.thread_binding(4, thread="threadIdx.x"):
+                        with T.sblock("C"):
+                            v_i = T.axis.spatial(n, i0 * 4 + i1)
+                            C[v_i] = T.max(A[v_i], B[v_i]) if op == "max" else T.min(A[v_i], B[v_i])
+
+    return Module
+
+
+# Same data as #20054. The NaN payloads are written into the float32 bit
+# pattern; the fp16/bf16 arms use explicit NaN codes for the NaN lanes and the
+# same finite values narrowed losslessly.
+A_F32 = np.array([0.0, 1.0, 0.0, 0.0, -0.0, 3.0, 2.0, -5.0], dtype="float32")
+B_F32 = np.array([1.0, 0.0, 0.0, -0.0, 0.0, 2.0, 2.0, -4.0], dtype="float32")
+A_F32.view("uint32")[[0, 2]] = 0x7FC00011
+B_F32.view("uint32")[[1, 2]] = 0x7FC00022
+
+
+def _nan_preserving_expected(a, b, op):
+    cmp = a < b if op == "min" else a > b
+    return np.where(cmp | np.isnan(a), a, b)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("op", ["min", "max"])
+@pytest.mark.parametrize("form", ["scalar", "vec4"])
+@pytest.mark.parametrize("dt", ["float32", "float64", "float16", "bfloat16"])
+def test_min_max_nan_preserving_cuda(op, dt, form):
+    n = 8
+    vectorize = form == "vec4"
+    if dt == "float32":
+        a_np, b_np = A_F32, B_F32
+    else:
+        a_np = np.array([np.nan, 1.0, np.nan, 0.0, -0.0, 3.0, 2.0, -5.0], dtype=dt)
+        b_np = np.array([1.0, np.nan, np.nan, -0.0, 0.0, 2.0, 2.0, -4.0], dtype=dt)
+
+    mod = tvm.compile(_min_max_nan_module(op, dt, n, vectorize=vectorize), target="cuda")
+    a = tvm.runtime.tensor(a_np, tvm.cuda(0))
+    b = tvm.runtime.tensor(b_np, tvm.cuda(0))
+    c = tvm.runtime.empty((n,), dt, tvm.cuda(0))
+
+    def run_and_check():
+        mod(a, b, c)
+        got = c.numpy()
+        expected = _nan_preserving_expected(a_np, b_np, op)
+        if dt == "float64":
+            np.testing.assert_array_equal(got.view("uint64"), expected.view("uint64"))
+        elif dt == "float32":
+            np.testing.assert_array_equal(got.view("uint32"), expected.view("uint32"))
+        else:
+            np.testing.assert_array_equal(got.view("uint16"), expected.view("uint16"))
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_min_max_nan_preserving_composite_cuda():
+    # The scalar ternary must survive nesting inside a compound expression:
+    # C[i] = max(A[i], B[i]) + 1.0 without being parsed as (x + cond) ? va : vb.
+    # NaN lanes cannot be compared bitwise (CUDA float add normalizes the NaN
+    # payload), so lanes 0-2 assert NaN and the finite lanes assert bitwise.
+    n = 8
+    a_np = A_F32
+    b_np = B_F32
+    mod = tvm.compile(
+        _min_max_nan_module("max", "float32", n, vectorize=False, composite=True),
+        target="cuda",
+    )
+    a = tvm.runtime.tensor(a_np, tvm.cuda(0))
+    b = tvm.runtime.tensor(b_np, tvm.cuda(0))
+    c = tvm.runtime.empty((n,), "float32", tvm.cuda(0))
+
+    def run_and_check():
+        mod(a, b, c)
+        got = c.numpy()
+        expected = _nan_preserving_expected(a_np, b_np, "max") + np.float32(1.0)
+        assert np.isnan(got[[0, 1, 2]]).all(), "NaN lanes must stay NaN after + 1.0"
+        finite = [3, 4, 5, 6, 7]
+        np.testing.assert_array_equal(got.view("uint32")[finite], expected.view("uint32")[finite])
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+def test_min_max_chained_statements_cuda():
+    # Two consecutive min/max statements that both read C, with a write in
+    # between: C[v] = max(C[v], A[v]); C[v] = max(C[v], B[v]). The scalar SSA
+    # binding must not cache the read across statements — SSAGetID caches by
+    # printed text, and the first statement's binding of the read of C must
+    # not be hit by the second statement after the first write. This broke
+    # warp allreduce, where red_buf[0] = max(red_buf[0], shuffle_down(...)) is
+    # emitted repeatedly. C is pre-filled with c0 so a cache hit (if present)
+    # reads a known wrong value; with c0 = 0 and a > max(c0, b), the pre-fix
+    # code computes max(c0, b) on at least one lane and fails.
+    n = 8
+    c0_np = np.zeros(n, dtype="float32")
+    a_np = np.array([5.0, 9.0, 3.0, 7.0, 1.0, 8.0, 4.0, 6.0], dtype="float32")
+    b_np = np.array([1.0, 2.0, 2.0, 1.0, 0.0, 3.0, 2.0, 5.0], dtype="float32")
+
+    @I.ir_module(s_tir=True)
+    class Module:
+        @T.prim_func(s_tir=True)
+        def main(
+            A: T.Buffer((n,), "float32"),
+            B: T.Buffer((n,), "float32"),
+            C: T.Buffer((n,), "float32"),
+        ):
+            T.func_attr({"tirx.noalias": True})
+            for i in T.thread_binding(n, thread="threadIdx.x"):
+                with T.sblock("C"):
+                    v_i = T.axis.spatial(n, i)
+                    C[v_i] = T.max(C[v_i], A[v_i])
+                    C[v_i] = T.max(C[v_i], B[v_i])
+
+    mod = tvm.compile(Module, target="cuda")
+    a = tvm.runtime.tensor(a_np, tvm.cuda(0))
+    b = tvm.runtime.tensor(b_np, tvm.cuda(0))
+    c = tvm.runtime.tensor(c0_np, tvm.cuda(0))
+
+    def run_and_check():
+        mod(a, b, c)
+        expected = np.maximum(np.maximum(c0_np, a_np), b_np)
+        np.testing.assert_array_equal(c.numpy(), expected)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
+@pytest.mark.parametrize("op", ["min", "max"])
+def test_min_max_int_vector_cuda(op):
+    # Integer min/max keeps the base codegen path, including the per-lane
+    # vector expansion: lanes=4 must compile and match numpy bitwise.
+    n = 8
+    dt = "int32"
+    a_np = np.array([3, -5, 7, 0, -1, 9, 2, -8], dtype=dt)
+    b_np = np.array([1, 9, -2, 4, -3, 9, 5, -8], dtype=dt)
+
+    mod = tvm.compile(_min_max_nan_module(op, dt, n, vectorize=True), target="cuda")
+    a = tvm.runtime.tensor(a_np, tvm.cuda(0))
+    b = tvm.runtime.tensor(b_np, tvm.cuda(0))
+    c = tvm.runtime.empty((n,), dt, tvm.cuda(0))
+
+    def run_and_check():
+        mod(a, b, c)
+        expected = np.minimum(a_np, b_np) if op == "min" else np.maximum(a_np, b_np)
+        np.testing.assert_array_equal(c.numpy(), expected)
+
+    tvm.testing.run_with_gpu_lock(run_and_check)
+
+
 @pytest.mark.gpu
 @pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
 def test_round_ties_to_even():