From 6810d51192b62bf4392c3313e647709a8f664589 Mon Sep 17 00:00:00 2001
From: LngelKyo <lngel.kyo@a2o-labs.com>
Date: Tue, 22 Sep 2026 09:08:44 +0900
Subject: [PATCH 1/8] [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<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.
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 <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;
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():

From bb6e5ced3346c16ff73d2d5cc91ce6da75d9a472 Mon Sep 17 00:00:00 2001
From: LngelKyo <lngel.kyo@a2o-labs.com>
Date: Wed, 23 Sep 2026 10:26:17 +0900
Subject: [PATCH 2/8] [Fix][CUDA] Inline the min/max constant-operand test
 constants

The float_imm tests wrote `const_scalar = T.float32(...)` at the top of the
prim_func. TVMScript parses that assignment as a one-element host-side local
buffer plus a store, not as a constant: under CUDA the buffer is lowered to
device workspace (cudaMalloc) and the host then writes the device pointer,
which segfaults with no diagnostic. The kernel also read the operand back
through a load, so the constant fast path was never exercised. The constants
are now inlined into the min/max call, the vec4 form uses scalar indices, and
the NaN lanes are checked with isnan instead of bitwise equality because CUDA
emits a constant NaN as CUDART_NAN_F (0x7fffffff) while numpy uses 0x7fc00000.
Source assertions now prove which path each form took.

AsFloatImm narrows to scalar FloatImm: its Broadcast branch sat behind the
lanes() == 1 guard and was unreachable; vector constants take the general
per-lane path.
---
 src/backend/cuda/codegen/codegen_cuda.cc      | 15 +--
 .../codegen/test_target_codegen_cuda.py       | 96 ++++++++++++++-----
 2 files changed, 74 insertions(+), 37 deletions(-)

diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc
index d78b64603e65..90c4168e2a14 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -80,17 +80,10 @@ TVM_FFI_INLINE bool IsPackedFloat(const PrimType& 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;
-}
+// (an ordered compare against a constant is never NaN). Scalar constants
+// only: a vector constant (Broadcast of a FloatImm) takes the general
+// per-lane path.
+TVM_FFI_INLINE const FloatImmNode* AsFloatImm(const PrimExpr& e) { return e.as<FloatImmNode>(); }
 
 // 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
diff --git a/tests/python/codegen/test_target_codegen_cuda.py b/tests/python/codegen/test_target_codegen_cuda.py
index 8f9ef438b7c1..870e163ed9e2 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -988,37 +988,45 @@ def run_and_check():
 @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.
+    # The constant-operand cases from #20054's test_min_max_float_imm_operand.
+    # The constant is written inline as T.float32(const_val) so that it reaches
+    # codegen as a FloatImm; a statement `c = T.float32(...)` in the prim_func
+    # body would instead become a one-element local buffer read through a load.
+    # scalar: lanes == 1, so codegen takes the constant fast path, none of
+    #   whose forms emits the "||" of the general ternary.
+    # vec4: VectorizeLoop turns the literal into Broadcast(FloatImm, 4); the fast
+    #   path is scalar-only, so this form checks a broadcast constant on the
+    #   per-lane general path.
     n = 8
-    const_arr = np.full(n, np.nan, dtype="float32") if const_nan else np.zeros(n, dtype="float32")
+    const_val = float("nan") if const_nan else 0.0
+    # B's NaN gets its own payload, so a lane that must return B's NaN cannot
+    # pass by returning the constant NaN, and vice versa.
+    b_nan_bits = 0x7FC00033
+    b_np = CONST_OTHER.copy()
+    b_np.view("uint32")[0] = b_nan_bits
 
     @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])
+                            T.min(T.float32(const_val), B[v_i])
                             if op == "min"
-                            else T.max(const_scalar, B[v_i])
+                            else T.max(T.float32(const_val), 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)
+                            T.min(B[v_i], T.float32(const_val))
+                            if op == "min"
+                            else T.max(B[v_i], T.float32(const_val))
                         )
             elif const_side == "lhs":
                 for i0 in T.thread_binding(2, thread="blockIdx.x"):
@@ -1026,9 +1034,9 @@ def main(B: T.Buffer((n,), "float32"), C: T.Buffer((n,), "float32")):
                         with T.sblock("C"):
                             v_i = T.axis.spatial(n, i0 * 4 + i1)
                             C[v_i] = (
-                                T.min(const_vec, B[v_i])
+                                T.min(T.float32(const_val), B[v_i])
                                 if op == "min"
-                                else T.max(const_vec, B[v_i])
+                                else T.max(T.float32(const_val), B[v_i])
                             )
             else:
                 for i0 in T.thread_binding(2, thread="blockIdx.x"):
@@ -1036,29 +1044,61 @@ def main(B: T.Buffer((n,), "float32"), C: T.Buffer((n,), "float32")):
                         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)
+                                T.min(B[v_i], T.float32(const_val))
+                                if op == "min"
+                                else T.max(B[v_i], T.float32(const_val))
                             )
 
     mod = tvm.compile(Module, target="cuda")
-    b = tvm.runtime.tensor(CONST_OTHER, tvm.cuda(0))
+
+    # The CUDA header emitted before the kernel contains "||" in preprocessor
+    # lines, so only the text from the kernel declaration on is checked.
+    src = mod.mod.imports[0].inspect_source()
+    kernel = src[src.index("__global__") :]
+    if form == "scalar":
+        assert "||" not in kernel, kernel
+        if const_nan and const_side == "lhs":
+            # the constant NaN is the result: no select at all
+            assert "?" not in kernel, kernel
+            assert "CUDART_NAN_F" in kernel, kernel
+        elif const_nan:
+            # rhs NaN: only the (a != a) test is left
+            assert "!=" in kernel, kernel
+        else:
+            # non-NaN constant: an ordered compare only, no NaN test
+            assert "!=" not in kernel, kernel
+            if const_side == "lhs":
+                assert (" < " if op == "min" else " > ") in kernel, kernel
+            else:
+                assert (" >= " if op == "min" else " <= ") in kernel, kernel
+    else:
+        assert "||" in kernel, kernel
+
+    b = tvm.runtime.tensor(b_np, tvm.cuda(0))
     c = tvm.runtime.empty((n,), "float32", tvm.cuda(0))
 
     def run_and_check():
         mod(b, c)
         got = c.numpy()
+        const_np = np.full(n, const_val, dtype="float32")
         if const_side == "lhs":
-            lhs_np, rhs_np = const_arr, CONST_OTHER
+            lhs_np, rhs_np = const_np, b_np
         else:
-            lhs_np, rhs_np = CONST_OTHER, const_arr
+            lhs_np, rhs_np = b_np, const_np
         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
+        take_lhs = cmp | np.isnan(lhs_np)
+        expected = np.where(take_lhs, lhs_np, rhs_np)
+        takes_const = take_lhs if const_side == "lhs" else ~take_lhs
+        if const_nan:
+            # A constant NaN is emitted as CUDART_NAN_F, whose bits differ from
+            # numpy's NaN: lanes that return it are checked for NaN, and for not
+            # being B's NaN.
+            assert np.isnan(got[takes_const]).all(), got
+            assert (got.view("uint32")[takes_const] != b_nan_bits).all(), got
+            exact = ~takes_const
+        else:
+            exact = np.ones(n, dtype=bool)
+        np.testing.assert_array_equal(got.view("uint32")[exact], expected.view("uint32")[exact])
 
     tvm.testing.run_with_gpu_lock(run_and_check)
 
@@ -1084,6 +1124,10 @@ def main(B: T.Buffer((n,), "float32"), C: T.Buffer((n,), "float32")):
                     C[v_i] = T.max(T.float32(0.0), B[v_i]) + T.float32(1.0)
 
     mod = tvm.compile(Module, target="cuda")
+    src = mod.mod.imports[0].inspect_source()
+    kernel = src[src.index("__global__") :]
+    # the inline 0.0 reaches codegen as a FloatImm: fast path, no NaN clause
+    assert "||" not in kernel and "!=" not in kernel, kernel
     b = tvm.runtime.tensor(CONST_OTHER, tvm.cuda(0))
     c = tvm.runtime.empty((n,), "float32", tvm.cuda(0))
 

From 525347dd8eb25672b4cf8c6c3f698fb03bc7acb3 Mon Sep 17 00:00:00 2001
From: LngelKyo <lngel.kyo@a2o-labs.com>
Date: Thu, 24 Sep 2026 18:04:24 +0900
Subject: [PATCH 3/8] [Fix][CUDA] Check the NaN lanes of the constant-operand
 test with isnan

test_min_max_float_imm_operand_cuda compared every lane bitwise, and gave B's
NaN lane its own payload (0x7FC00033) so that "which operand was selected"
could be read off the payload. On sm_89 the nvcc-compiled min(c, B) and
max(B, c) forms return the canonical NaN (0x7FFFFFFF) for that lane: nvcc
lowers the compare-select to a min/max instruction that quiets the payload,
which is allowed. The contract of the NaN-preserving min/max is that the
result is NaN whenever an operand is NaN, not which NaN.

The check now asserts isnan on every lane whose expected value is NaN, no NaN
on the other lanes, and bitwise equality (sign of zero included) on the finite
lanes. Which operand was selected is still visible from the finite lanes.
---
 .../codegen/test_target_codegen_cuda.py       | 29 +++++++++----------
 1 file changed, 14 insertions(+), 15 deletions(-)

diff --git a/tests/python/codegen/test_target_codegen_cuda.py b/tests/python/codegen/test_target_codegen_cuda.py
index 870e163ed9e2..1c77a447d728 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -999,11 +999,9 @@ def test_min_max_float_imm_operand_cuda(op, const_side, const_nan, form):
     #   per-lane general path.
     n = 8
     const_val = float("nan") if const_nan else 0.0
-    # B's NaN gets its own payload, so a lane that must return B's NaN cannot
-    # pass by returning the constant NaN, and vice versa.
-    b_nan_bits = 0x7FC00033
+    # B carries a NaN lane so the (NaN B, finite constant) branch is exercised.
     b_np = CONST_OTHER.copy()
-    b_np.view("uint32")[0] = b_nan_bits
+    b_np[0] = np.float32("nan")
 
     @I.ir_module(s_tir=True)
     class Module:
@@ -1088,17 +1086,18 @@ def run_and_check():
         cmp = lhs_np < rhs_np if op == "min" else lhs_np > rhs_np
         take_lhs = cmp | np.isnan(lhs_np)
         expected = np.where(take_lhs, lhs_np, rhs_np)
-        takes_const = take_lhs if const_side == "lhs" else ~take_lhs
-        if const_nan:
-            # A constant NaN is emitted as CUDART_NAN_F, whose bits differ from
-            # numpy's NaN: lanes that return it are checked for NaN, and for not
-            # being B's NaN.
-            assert np.isnan(got[takes_const]).all(), got
-            assert (got.view("uint32")[takes_const] != b_nan_bits).all(), got
-            exact = ~takes_const
-        else:
-            exact = np.ones(n, dtype=bool)
-        np.testing.assert_array_equal(got.view("uint32")[exact], expected.view("uint32")[exact])
+        # The contract is NaN propagation, not the NaN payload: a constant NaN is
+        # emitted as CUDART_NAN_F, and nvcc may lower the compare-select to a
+        # min/max instruction that returns the canonical NaN (seen on sm_89 for
+        # min(c, B) and max(B, c): 0x7FFFFFFF where B held 0x7FC00033). NaN lanes
+        # are checked with isnan; every finite lane, including the sign of zero,
+        # bitwise.
+        nan_lanes = np.isnan(expected)
+        assert np.isnan(got[nan_lanes]).all(), got
+        assert not np.isnan(got[~nan_lanes]).any(), got
+        np.testing.assert_array_equal(
+            got.view("uint32")[~nan_lanes], expected.view("uint32")[~nan_lanes]
+        )
 
     tvm.testing.run_with_gpu_lock(run_and_check)
 

From ac9563e0dcf7e79fbfc3a9fcdd7e139ee7a444e9 Mon Sep 17 00:00:00 2001
From: LngelKyo <lngel.kyo@a2o-labs.com>
Date: Sun, 27 Sep 2026 07:33:41 +0900
Subject: [PATCH 4/8] [Fix][CUDA] Write the min/max tests as plain TIRx with
 T.constexpr selections

Main now builds the codegen tests as plain TIRx (#20437) and selects a Python
branch in a prim_func body only when the condition is marked with T.constexpr
(#20423). Write the min/max tests the same way: no S-TIR blocks, direct
indexing inside the thread-bound and vectorized loops, and T.constexpr for the
op, constant-side and form selections. What the tests check is unchanged.

The float-imm test's comment no longer describes what a body assignment turns
into; only the reason for writing the constant inline is kept.
---
 .../codegen/test_target_codegen_cuda.py       | 119 ++++++++----------
 1 file changed, 51 insertions(+), 68 deletions(-)

diff --git a/tests/python/codegen/test_target_codegen_cuda.py b/tests/python/codegen/test_target_codegen_cuda.py
index 48da979a7c5f..38c1dcb5b5dc 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -812,35 +812,31 @@ def run_and_check():
 
 
 def _min_max_nan_module(op, dt, n=8, vectorize=False, composite=False):
-    @I.ir_module(s_tir=True)
+    @I.ir_module
     class Module:
-        @T.prim_func(s_tir=True)
+        @T.prim_func
         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:
+                if T.constexpr(composite):
+                    C[i0 * 4 + 0] = T.max(A[i0 * 4 + 0], B[i0 * 4 + 0]) + T.float32(1.0)
+                    C[i0 * 4 + 1] = T.max(A[i0 * 4 + 1], B[i0 * 4 + 1]) + T.float32(1.0)
+                    C[i0 * 4 + 2] = T.max(A[i0 * 4 + 2], B[i0 * 4 + 2]) + T.float32(1.0)
+                    C[i0 * 4 + 3] = T.max(A[i0 * 4 + 3], B[i0 * 4 + 3]) + T.float32(1.0)
+                elif T.constexpr(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])
+                        C[i0 * 4 + i1] = (
+                            T.max(A[i0 * 4 + i1], B[i0 * 4 + i1])
+                            if T.constexpr(op == "max")
+                            else T.min(A[i0 * 4 + i1], B[i0 * 4 + i1])
+                        )
                 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])
+                        C[i0 * 4 + i1] = (
+                            T.max(A[i0 * 4 + i1], B[i0 * 4 + i1])
+                            if T.constexpr(op == "max")
+                            else T.min(A[i0 * 4 + i1], B[i0 * 4 + i1])
+                        )
 
     return Module
 
@@ -938,9 +934,9 @@ def test_min_max_chained_statements_cuda():
     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)
+    @I.ir_module
     class Module:
-        @T.prim_func(s_tir=True)
+        @T.prim_func
         def main(
             A: T.Buffer((n,), "float32"),
             B: T.Buffer((n,), "float32"),
@@ -948,10 +944,8 @@ def main(
         ):
             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])
+                C[i] = T.max(C[i], A[i])
+                C[i] = T.max(C[i], B[i])
 
     mod = tvm.compile(Module, target="cuda")
     a = tvm.runtime.tensor(a_np, tvm.cuda(0))
@@ -978,8 +972,7 @@ def run_and_check():
 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.
     # The constant is written inline as T.float32(const_val) so that it reaches
-    # codegen as a FloatImm; a statement `c = T.float32(...)` in the prim_func
-    # body would instead become a one-element local buffer read through a load.
+    # codegen as a FloatImm.
     # scalar: lanes == 1, so codegen takes the constant fast path, none of
     #   whose forms emits the "||" of the general ternary.
     # vec4: VectorizeLoop turns the literal into Broadcast(FloatImm, 4); the fast
@@ -991,49 +984,41 @@ def test_min_max_float_imm_operand_cuda(op, const_side, const_nan, form):
     b_np = CONST_OTHER.copy()
     b_np[0] = np.float32("nan")
 
-    @I.ir_module(s_tir=True)
+    @I.ir_module
     class Module:
-        @T.prim_func(s_tir=True)
+        @T.prim_func
         def main(B: T.Buffer((n,), "float32"), C: T.Buffer((n,), "float32")):
             T.func_attr({"tirx.noalias": True})
-            if form == "scalar" and const_side == "lhs":
+            if T.constexpr(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(T.float32(const_val), B[v_i])
-                            if op == "min"
-                            else T.max(T.float32(const_val), B[v_i])
-                        )
-            elif form == "scalar":
+                    C[i] = (
+                        T.min(T.float32(const_val), B[i])
+                        if T.constexpr(op == "min")
+                        else T.max(T.float32(const_val), B[i])
+                    )
+            elif T.constexpr(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.min(B[v_i], T.float32(const_val))
-                            if op == "min"
-                            else T.max(B[v_i], T.float32(const_val))
-                        )
-            elif const_side == "lhs":
+                    C[i] = (
+                        T.min(B[i], T.float32(const_val))
+                        if T.constexpr(op == "min")
+                        else T.max(B[i], T.float32(const_val))
+                    )
+            elif T.constexpr(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(T.float32(const_val), B[v_i])
-                                if op == "min"
-                                else T.max(T.float32(const_val), B[v_i])
-                            )
+                        C[i0 * 4 + i1] = (
+                            T.min(T.float32(const_val), B[i0 * 4 + i1])
+                            if T.constexpr(op == "min")
+                            else T.max(T.float32(const_val), B[i0 * 4 + i1])
+                        )
             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.min(B[v_i], T.float32(const_val))
-                                if op == "min"
-                                else T.max(B[v_i], T.float32(const_val))
-                            )
+                        C[i0 * 4 + i1] = (
+                            T.min(B[i0 * 4 + i1], T.float32(const_val))
+                            if T.constexpr(op == "min")
+                            else T.max(B[i0 * 4 + i1], T.float32(const_val))
+                        )
 
     mod = tvm.compile(Module, target="cuda")
 
@@ -1100,15 +1085,13 @@ def test_min_max_float_imm_operand_nested_cuda():
     # NaN (CUDA float add normalizes the payload); finite lanes bitwise.
     n = 8
 
-    @I.ir_module(s_tir=True)
+    @I.ir_module
     class Module:
-        @T.prim_func(s_tir=True)
+        @T.prim_func
         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)
+                C[i] = T.max(T.float32(0.0), B[i]) + T.float32(1.0)
 
     mod = tvm.compile(Module, target="cuda")
     src = mod.mod.imports[0].inspect_source()

From 7ed01f5aa9bce08f3601c5e785cbfe356e00d64d Mon Sep 17 00:00:00 2001
From: LngelKyo <lngel.kyo@a2o-labs.com>
Date: Sat, 3 Oct 2026 15:06:55 +0900
Subject: [PATCH 5/8] [Test] Spell the two new min/max tests with T.Tensor
 after #20519

---
 tests/python/codegen/test_target_codegen_cuda.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tests/python/codegen/test_target_codegen_cuda.py b/tests/python/codegen/test_target_codegen_cuda.py
index a92e5320ebe5..d23cd12e969a 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -952,7 +952,7 @@ def _min_max_nan_module(op, dt, n=8, vectorize=False, composite=False):
     @I.ir_module
     class Module:
         @T.prim_func
-        def main(A: T.Buffer((n,), dt), B: T.Buffer((n,), dt), C: T.Buffer((n,), dt)):
+        def main(A: T.Tensor((n,), dt), B: T.Tensor((n,), dt), C: T.Tensor((n,), dt)):
             T.func_attr({"tirx.noalias": True})
             for i0 in T.thread_binding(2, thread="blockIdx.x"):
                 if T.constexpr(composite):
@@ -1075,9 +1075,9 @@ def test_min_max_chained_statements_cuda():
     class Module:
         @T.prim_func
         def main(
-            A: T.Buffer((n,), "float32"),
-            B: T.Buffer((n,), "float32"),
-            C: T.Buffer((n,), "float32"),
+            A: T.Tensor((n,), "float32"),
+            B: T.Tensor((n,), "float32"),
+            C: T.Tensor((n,), "float32"),
         ):
             T.func_attr({"tirx.noalias": True})
             for i in T.thread_binding(n, thread="threadIdx.x"):
@@ -1124,7 +1124,7 @@ def test_min_max_float_imm_operand_cuda(op, const_side, const_nan, form):
     @I.ir_module
     class Module:
         @T.prim_func
-        def main(B: T.Buffer((n,), "float32"), C: T.Buffer((n,), "float32")):
+        def main(B: T.Tensor((n,), "float32"), C: T.Tensor((n,), "float32")):
             T.func_attr({"tirx.noalias": True})
             if T.constexpr(form == "scalar" and const_side == "lhs"):
                 for i in T.thread_binding(n, thread="threadIdx.x"):
@@ -1225,7 +1225,7 @@ def test_min_max_float_imm_operand_nested_cuda():
     @I.ir_module
     class Module:
         @T.prim_func
-        def main(B: T.Buffer((n,), "float32"), C: T.Buffer((n,), "float32")):
+        def main(B: T.Tensor((n,), "float32"), C: T.Tensor((n,), "float32")):
             T.func_attr({"tirx.noalias": True})
             for i in T.thread_binding(n, thread="threadIdx.x"):
                 C[i] = T.max(T.float32(0.0), B[i]) + T.float32(1.0)

From 09629da38e6a53d3a94f3300d7b6501666e58f57 Mon Sep 17 00:00:00 2001
From: LngelKyo <lngel.kyo@a2o-labs.com>
Date: Wed, 14 Oct 2026 19:59:49 +0900
Subject: [PATCH 6/8] [Test] Spell the min/max tests with T.function and
 T.launch_thread after #20586 / #20660

---
 .../codegen/test_target_codegen_cuda.py       | 24 +++++++++----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/tests/python/codegen/test_target_codegen_cuda.py b/tests/python/codegen/test_target_codegen_cuda.py
index 36e0ca62f69b..98b9ba5d3cff 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -974,10 +974,10 @@ def run_and_check():
 def _min_max_nan_module(op, dt, n=8, vectorize=False, composite=False):
     @I.ir_module
     class Module:
-        @T.prim_func
+        @T.function
         def main(A: T.Tensor((n,), dt), B: T.Tensor((n,), dt), C: T.Tensor((n,), dt)):
             T.func_attr({"tirx.noalias": True})
-            for i0 in T.thread_binding(2, thread="blockIdx.x"):
+            with T.launch_thread("blockIdx.x", 2) as i0:
                 if T.constexpr(composite):
                     C[i0 * 4 + 0] = T.max(A[i0 * 4 + 0], B[i0 * 4 + 0]) + T.float32(1.0)
                     C[i0 * 4 + 1] = T.max(A[i0 * 4 + 1], B[i0 * 4 + 1]) + T.float32(1.0)
@@ -991,7 +991,7 @@ def main(A: T.Tensor((n,), dt), B: T.Tensor((n,), dt), C: T.Tensor((n,), dt)):
                             else T.min(A[i0 * 4 + i1], B[i0 * 4 + i1])
                         )
                 else:
-                    for i1 in T.thread_binding(4, thread="threadIdx.x"):
+                    with T.launch_thread("threadIdx.x", 4) as i1:
                         C[i0 * 4 + i1] = (
                             T.max(A[i0 * 4 + i1], B[i0 * 4 + i1])
                             if T.constexpr(op == "max")
@@ -1096,14 +1096,14 @@ def test_min_max_chained_statements_cuda():
 
     @I.ir_module
     class Module:
-        @T.prim_func
+        @T.function
         def main(
             A: T.Tensor((n,), "float32"),
             B: T.Tensor((n,), "float32"),
             C: T.Tensor((n,), "float32"),
         ):
             T.func_attr({"tirx.noalias": True})
-            for i in T.thread_binding(n, thread="threadIdx.x"):
+            with T.launch_thread("threadIdx.x", n) as i:
                 C[i] = T.max(C[i], A[i])
                 C[i] = T.max(C[i], B[i])
 
@@ -1146,25 +1146,25 @@ def test_min_max_float_imm_operand_cuda(op, const_side, const_nan, form):
 
     @I.ir_module
     class Module:
-        @T.prim_func
+        @T.function
         def main(B: T.Tensor((n,), "float32"), C: T.Tensor((n,), "float32")):
             T.func_attr({"tirx.noalias": True})
             if T.constexpr(form == "scalar" and const_side == "lhs"):
-                for i in T.thread_binding(n, thread="threadIdx.x"):
+                with T.launch_thread("threadIdx.x", n) as i:
                     C[i] = (
                         T.min(T.float32(const_val), B[i])
                         if T.constexpr(op == "min")
                         else T.max(T.float32(const_val), B[i])
                     )
             elif T.constexpr(form == "scalar"):
-                for i in T.thread_binding(n, thread="threadIdx.x"):
+                with T.launch_thread("threadIdx.x", n) as i:
                     C[i] = (
                         T.min(B[i], T.float32(const_val))
                         if T.constexpr(op == "min")
                         else T.max(B[i], T.float32(const_val))
                     )
             elif T.constexpr(const_side == "lhs"):
-                for i0 in T.thread_binding(2, thread="blockIdx.x"):
+                with T.launch_thread("blockIdx.x", 2) as i0:
                     for i1 in T.vectorized(4):
                         C[i0 * 4 + i1] = (
                             T.min(T.float32(const_val), B[i0 * 4 + i1])
@@ -1172,7 +1172,7 @@ def main(B: T.Tensor((n,), "float32"), C: T.Tensor((n,), "float32")):
                             else T.max(T.float32(const_val), B[i0 * 4 + i1])
                         )
             else:
-                for i0 in T.thread_binding(2, thread="blockIdx.x"):
+                with T.launch_thread("blockIdx.x", 2) as i0:
                     for i1 in T.vectorized(4):
                         C[i0 * 4 + i1] = (
                             T.min(B[i0 * 4 + i1], T.float32(const_val))
@@ -1247,10 +1247,10 @@ def test_min_max_float_imm_operand_nested_cuda():
 
     @I.ir_module
     class Module:
-        @T.prim_func
+        @T.function
         def main(B: T.Tensor((n,), "float32"), C: T.Tensor((n,), "float32")):
             T.func_attr({"tirx.noalias": True})
-            for i in T.thread_binding(n, thread="threadIdx.x"):
+            with T.launch_thread("threadIdx.x", n) as i:
                 C[i] = T.max(T.float32(0.0), B[i]) + T.float32(1.0)
 
     mod = tvm.compile(Module, target="cuda")

From c3929f9af3bd98593440c2bc3c920d26c9061b59 Mon Sep 17 00:00:00 2001
From: LngelKyo <lngel.kyo@a2o-labs.com>
Date: Wed, 14 Oct 2026 20:34:41 +0900
Subject: [PATCH 7/8] [Test] Run the new min/max tests on both CUDA compiler
 paths via the backend_config fixture

---
 .../codegen/test_target_codegen_cuda.py       | 24 +++++++++----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/tests/python/codegen/test_target_codegen_cuda.py b/tests/python/codegen/test_target_codegen_cuda.py
index 98b9ba5d3cff..6328eb156e4c 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -1020,7 +1020,7 @@ def _nan_preserving_expected(a, b, op):
 @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):
+def test_min_max_nan_preserving_cuda(op, dt, form, *, backend_config):
     n = 8
     vectorize = form == "vec4"
     if dt == "float32":
@@ -1029,7 +1029,7 @@ def test_min_max_nan_preserving_cuda(op, dt, form):
         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")
+    mod = tvm.compile(_min_max_nan_module(op, dt, n, vectorize=vectorize), target={"kind": "cuda", "backend_config": backend_config})
     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))
@@ -1050,7 +1050,7 @@ def run_and_check():
 
 @pytest.mark.gpu
 @pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_min_max_nan_preserving_composite_cuda():
+def test_min_max_nan_preserving_composite_cuda(*, backend_config):
     # 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
@@ -1060,7 +1060,7 @@ def test_min_max_nan_preserving_composite_cuda():
     b_np = B_F32
     mod = tvm.compile(
         _min_max_nan_module("max", "float32", n, vectorize=False, composite=True),
-        target="cuda",
+        target={"kind": "cuda", "backend_config": backend_config},
     )
     a = tvm.runtime.tensor(a_np, tvm.cuda(0))
     b = tvm.runtime.tensor(b_np, tvm.cuda(0))
@@ -1079,7 +1079,7 @@ def run_and_check():
 
 @pytest.mark.gpu
 @pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
-def test_min_max_chained_statements_cuda():
+def test_min_max_chained_statements_cuda(*, backend_config):
     # 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
@@ -1107,7 +1107,7 @@ def main(
                 C[i] = T.max(C[i], A[i])
                 C[i] = T.max(C[i], B[i])
 
-    mod = tvm.compile(Module, target="cuda")
+    mod = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
     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))
@@ -1129,7 +1129,7 @@ def run_and_check():
 @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):
+def test_min_max_float_imm_operand_cuda(op, const_side, const_nan, form, *, backend_config):
     # The constant-operand cases from #20054's test_min_max_float_imm_operand.
     # The constant is written inline as T.float32(const_val) so that it reaches
     # codegen as a FloatImm.
@@ -1180,7 +1180,7 @@ def main(B: T.Tensor((n,), "float32"), C: T.Tensor((n,), "float32")):
                             else T.max(B[i0 * 4 + i1], T.float32(const_val))
                         )
 
-    mod = tvm.compile(Module, target="cuda")
+    mod = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
 
     # The CUDA header emitted before the kernel contains "||" in preprocessor
     # lines, so only the text from the kernel declaration on is checked.
@@ -1237,7 +1237,7 @@ def 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():
+def test_min_max_float_imm_operand_nested_cuda(*, backend_config):
     # 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
@@ -1253,7 +1253,7 @@ def main(B: T.Tensor((n,), "float32"), C: T.Tensor((n,), "float32")):
             with T.launch_thread("threadIdx.x", n) as i:
                 C[i] = T.max(T.float32(0.0), B[i]) + T.float32(1.0)
 
-    mod = tvm.compile(Module, target="cuda")
+    mod = tvm.compile(Module, target={"kind": "cuda", "backend_config": backend_config})
     src = mod.mod.imports[0].inspect_source()
     kernel = src[src.index("__global__") :]
     # the inline 0.0 reaches codegen as a FloatImm: fast path, no NaN clause
@@ -1277,7 +1277,7 @@ def 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):
+def test_min_max_int_vector_cuda(op, *, backend_config):
     # Integer min/max keeps the base codegen path, including the per-lane
     # vector expansion: lanes=4 must compile and match numpy bitwise.
     n = 8
@@ -1285,7 +1285,7 @@ def test_min_max_int_vector_cuda(op):
     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")
+    mod = tvm.compile(_min_max_nan_module(op, dt, n, vectorize=True), target={"kind": "cuda", "backend_config": backend_config})
     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))

From e832b003d6b36775931a58fb740161361bcd12da Mon Sep 17 00:00:00 2001
From: LngelKyo <lngel.kyo@a2o-labs.com>
Date: Wed, 14 Oct 2026 22:13:07 +0900
Subject: [PATCH 8/8] [Test] Wrap the two long tvm.compile calls in the min/max
 tests for ruff format

The backend_config fixture change left two tvm.compile(...) calls over
100 columns. This is the output of ruff format (pre-commit, ruff 0.12.3);
no functional change.
---
 tests/python/codegen/test_target_codegen_cuda.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/tests/python/codegen/test_target_codegen_cuda.py b/tests/python/codegen/test_target_codegen_cuda.py
index 6328eb156e4c..2abeba985f0c 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -1029,7 +1029,10 @@ def test_min_max_nan_preserving_cuda(op, dt, form, *, backend_config):
         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={"kind": "cuda", "backend_config": backend_config})
+    mod = tvm.compile(
+        _min_max_nan_module(op, dt, n, vectorize=vectorize),
+        target={"kind": "cuda", "backend_config": backend_config},
+    )
     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))
@@ -1285,7 +1288,10 @@ def test_min_max_int_vector_cuda(op, *, backend_config):
     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={"kind": "cuda", "backend_config": backend_config})
+    mod = tvm.compile(
+        _min_max_nan_module(op, dt, n, vectorize=True),
+        target={"kind": "cuda", "backend_config": backend_config},
+    )
     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))