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

src/backend/cuda/codegen/codegen_cuda.cc

@@ -77,6 +77,31 @@ TVM_FFI_INLINE bool IsPackedFloat(const PrimType& ty) {
   return IsFloat8(ty) || IsFloat6(ty) || IsFloat4(ty);
 }
 
+// Constant-operand handling for min/max, same semantics as apache/tvm
+// #20054's min_max_utils.h (host side). A constant NaN operand is printed
+// directly; a constant non-NaN operand lets the ternary skip its NaN clause
+// (an ordered compare against a constant is never NaN). The two
+// implementations can be merged once the host/CUDA split is revisited.
+TVM_FFI_INLINE const FloatImmNode* AsFloatImm(const PrimExpr& e) {
+  if (const auto* fa = e.as<FloatImmNode>()) {
+    return fa;
+  }
+  if (const auto* bc = e.as<BroadcastNode>()) {
+    return bc->value.as<FloatImmNode>();
+  }
+  return nullptr;
+}
+
+// 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 +711,126 @@ 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;
+  }
+  // Constant-operand fast paths, matching #20054's min_max_utils.h. A
+  // constant NaN operand is printed as-is; a constant non-NaN operand lets
+  // the ternary drop its NaN clause. The constant side is not SSA-bound.
+  const FloatImmNode* ca = AsFloatImm(op->a);
+  const FloatImmNode* cb = AsFloatImm(op->b);
+  if (op_ty.lanes() == 1 && (ca != nullptr || cb != nullptr)) {
+    if (ca != nullptr && std::isnan(ca->value)) {
+      // a constant NaN lhs is the result: print it directly.
+      this->PrintExpr(op->a, os);
+      return;
+    }
+    int ssa_scope = BeginScope();
+    if (cb != nullptr && std::isnan(cb->value)) {
+      // rhs is a constant NaN: keep a on any comparison outcome. a is bound
+      // once and referenced in both clauses.
+      std::string va = SSAGetID(PrintExpr(op->a), op->a.ty());
+      os << "((" << va << " != " << va << ") ? " << va << " : ";
+      PrintConst(cb, os, this);
+      os << ")";
+      EndScope(ssa_scope);
+      return;
+    }
+    if (ca != nullptr) {
+      // lhs is a constant non-NaN: ordered compare, b's NaN makes the compare
+      // false and b is kept. b bound once.
+      std::string vb = SSAGetID(PrintExpr(op->b), op->b.ty());
+      os << "((";
+      PrintConst(ca, os, this);
+      os << ' ' << cmp << ' ' << vb << ") ? ";
+      PrintConst(ca, os, this);
+      os << " : " << vb << ")";
+      EndScope(ssa_scope);
+      return;
+    }
+    {
+      // rhs is a constant non-NaN: reverse to a non-strict compare so a's NaN
+      // makes the compare false and a is kept; ties take b. a bound once.
+      std::string va = SSAGetID(PrintExpr(op->a), op->a.ty());
+      const char* rcmp = (cmp[0] == '<') ? ">=" : "<=";
+      os << "((" << va << ' ' << rcmp << ' ';
+      PrintConst(cb, os, this);
+      os << ") ? ";
+      PrintConst(cb, os, this);
+      os << " : " << va << ')';
+      EndScope(ssa_scope);
+      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;