diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc index 55e61927b156..07605a9119dd 100644 --- a/src/backend/cuda/codegen/codegen_cuda.cc +++ b/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 +void CodeGenCUDA::PrintMinMaxNanPreservingImpl(const T* op, const char* opstr, + std::ostream& os) { // NOLINT(*) + PrimType op_ty = op->ty.template as_or_throw(); + 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. diff --git a/src/backend/cuda/codegen/codegen_cuda.h b/src/backend/cuda/codegen/codegen_cuda.h index c7703f9e0b9c..2f051530f5e9 100644 --- a/src/backend/cuda/codegen/codegen_cuda.h +++ b/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 + 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; diff --git a/tests/python/codegen/test_target_codegen_cuda.py b/tests/python/codegen/test_target_codegen_cuda.py index d0459f3e958f..6943b089dfd2 100644 --- a/tests/python/codegen/test_target_codegen_cuda.py +++ b/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():