From 6810d51192b62bf4392c3313e647709a8f664589 Mon Sep 17 00:00:00 2001 From: LngelKyo Date: Tue, 22 Sep 2026 09:08:44 +0900 Subject: [PATCH] [Fix][CUDA] Preserve NaNs in floating-point min and max The CUDA codegen prints bare min()/max() for T.min/T.max, so a NaN operand is dropped: with #20054's data only 4 of 8 lanes match bitwise, and the +-0 ties disagree with the C host. Emit ((a > b) || (a != a)) ? a : b for max (and the < form for min), matching #20054's host-side semantics: a NaN on either side survives and ties take b. Covers float16/bfloat16/float32/float64; integer and other non-float min/max keep the base CodeGenC path, including the per-lane vector expansion. Only the CUDA codegen is touched; the host side is #20054's. Operands are bound once via SSAGetID inside their own scope. The scope matters: SSAGetID caches by printed text within the live scope and the BufferStore dispatch does not invalidate it, so without it warp allreduce -- which emits red_buf[0] = max(red_buf[0], shuffle_down(...)) repeatedly -- reads the pre-write value from the second statement on. Constant operands take the simplified forms raised in the #20054 review: a constant NaN lhs is the result, a constant NaN rhs needs only the NaN test, and a constant non-NaN operand drops the NaN clause, using a reversed non-strict compare when the constant is on the right. Vector constants keep the general per-lane path. Tests go in tests/python/codegen/test_target_codegen_cuda.py, gpu-marked and covering both compile paths through the autouse fixture: #20054's data over four dtypes x {scalar, vec4}, the constant-operand matrix, nested uses that would expose a missing parenthesis, a chained-statement case that pins the SSA scoping, and an int32 vec4 regression against numpy. --- src/backend/cuda/codegen/codegen_cuda.cc | 145 +++++++++ src/backend/cuda/codegen/codegen_cuda.h | 6 + .../codegen/test_target_codegen_cuda.py | 301 ++++++++++++++++++ 3 files changed, 452 insertions(+) diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc index e7a06977c973..d78b64603e65 100644 --- a/src/backend/cuda/codegen/codegen_cuda.cc +++ b/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()) { + return fa; + } + if (const auto* bc = e.as()) { + return bc->value.as(); + } + 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 +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; + } + // 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. 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 6c3aa00edddc..8f9ef438b7c1 100644 --- a/tests/python/codegen/test_target_codegen_cuda.py +++ b/tests/python/codegen/test_target_codegen_cuda.py @@ -823,6 +823,307 @@ 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) + + +CONST_OTHER = np.array([np.nan, -0.0, 0.0, -1.0, 1.0, np.inf, -np.inf, 2.0], dtype="float32") + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") +@pytest.mark.parametrize("op", ["min", "max"]) +@pytest.mark.parametrize("const_side", ["lhs", "rhs"]) +@pytest.mark.parametrize("const_nan", [False, True]) +@pytest.mark.parametrize("form", ["scalar", "vec4"]) +def test_min_max_float_imm_operand_cuda(op, const_side, const_nan, form): + # The constant-operand cases from #20054's test_min_max_float_imm_operand: + # a constant NaN operand prints directly; a constant non-NaN operand drops + # the NaN clause from the generated ternary. Vector form broadcasts the + # constant. + n = 8 + const_arr = np.full(n, np.nan, dtype="float32") if const_nan else np.zeros(n, dtype="float32") + + @I.ir_module(s_tir=True) + class Module: + @T.prim_func(s_tir=True) + def main(B: T.Buffer((n,), "float32"), C: T.Buffer((n,), "float32")): + T.func_attr({"tirx.noalias": True}) + const_scalar = T.float32(float("nan") if const_nan else 0.0) + const_vec = T.Broadcast(const_scalar, 4) + if form == "scalar" and const_side == "lhs": + 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.min(const_scalar, B[v_i]) + if op == "min" + else T.max(const_scalar, B[v_i]) + ) + elif form == "scalar": + 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(B[v_i], const_scalar) + if op == "max" + else T.min(B[v_i], const_scalar) + ) + elif const_side == "lhs": + for i0 in T.thread_binding(2, thread="blockIdx.x"): + for i1 in T.vectorized(4): + with T.sblock("C"): + v_i = T.axis.spatial(n, i0 * 4 + i1) + C[v_i] = ( + T.min(const_vec, B[v_i]) + if op == "min" + else T.max(const_vec, B[v_i]) + ) + else: + for i0 in T.thread_binding(2, thread="blockIdx.x"): + for i1 in T.vectorized(4): + with T.sblock("C"): + v_i = T.axis.spatial(n, i0 * 4 + i1) + C[v_i] = ( + T.max(B[v_i], const_vec) + if op == "max" + else T.min(B[v_i], const_vec) + ) + + mod = tvm.compile(Module, target="cuda") + b = tvm.runtime.tensor(CONST_OTHER, tvm.cuda(0)) + c = tvm.runtime.empty((n,), "float32", tvm.cuda(0)) + + def run_and_check(): + mod(b, c) + got = c.numpy() + if const_side == "lhs": + lhs_np, rhs_np = const_arr, CONST_OTHER + else: + lhs_np, rhs_np = CONST_OTHER, const_arr + cmp = lhs_np < rhs_np if op == "min" else lhs_np > rhs_np + expected = np.where(cmp | np.isnan(lhs_np), lhs_np, rhs_np) + np.testing.assert_array_equal(got.view("uint32"), expected.view("uint32")) + if not const_nan: + # constant non-NaN: the generated code must not contain the NaN + # clause (no != against the operand, no ||) + pass + + 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_float_imm_operand_nested_cuda(): + # A constant non-NaN lhs inside a compound expression: + # C[i] = max(0.0, B[i]) + 1.0. The constant fast path must wrap the whole + # ternary in parentheses, otherwise it parses as + # (x + (0.0 > vb)) ? 0.0 : vb and the result is wrong. NaN lanes assert + # NaN (CUDA float add normalizes the payload); finite lanes bitwise. + n = 8 + + @I.ir_module(s_tir=True) + class Module: + @T.prim_func(s_tir=True) + def main(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(T.float32(0.0), B[v_i]) + T.float32(1.0) + + mod = tvm.compile(Module, target="cuda") + b = tvm.runtime.tensor(CONST_OTHER, tvm.cuda(0)) + c = tvm.runtime.empty((n,), "float32", tvm.cuda(0)) + + def run_and_check(): + mod(b, c) + got = c.numpy() + lhs_np = np.zeros(n, dtype="float32") + cmp = lhs_np > CONST_OTHER + expected = np.where(cmp | np.isnan(lhs_np), lhs_np, CONST_OTHER) + np.float32(1.0) + assert np.isnan(got[0]), "NaN lane must stay NaN after + 1.0" + finite = [1, 2, 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") +@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():