Expose propagate_nan and fix argmin/argmax NaN handling to match min/max

Signed-off-by: Qiqi Xiao <qiqix@nvidia.com>

Author: xiaoqiqi177

changelog.d/propagate-nan.md

@@ -0,0 +1,7 @@
+- Added a ``propagate_nan`` option to ``ct.min()``, ``ct.max()``,
+  ``ct.argmin()``, ``ct.argmax()``, ``ct.minimum()`` and ``ct.maximum()``. By
+  default ``NaN`` values are ignored; with ``propagate_nan=True`` a ``NaN``
+  propagates -- ``min``/``max``/``minimum``/``maximum`` return ``NaN`` and
+  ``argmin``/``argmax`` return the index of the first ``NaN``.
+- Fixed ``ct.argmin()`` and ``ct.argmax()`` to ignore ``NaN`` values under the
+  default ``propagate_nan=False``, consistent with ``ct.min()`` and ``ct.max()``.

experimental/cuda-lang/src/cuda/lang/_passes/flatten_cfg.py

@@ -91,8 +91,6 @@ def flatten_continue(self, op: Continue, current: Block) -> Block:
                     fn="add",
                     lhs=self.loop_iv,
                     rhs=self.loop_step,
-                    rounding_mode=None,
-                    flush_to_zero=False,
                 )
             )
             args = (next_iv,) + args

src/cuda/tile/_ir/arithmetic_ops.py

@@ -357,8 +357,9 @@ def _bitwise_shift_tensorlike_impl(fn: str, x: Var[TensorLikeTy], y: Var[TensorL
 @dataclass(eq=False)
 class RawBinaryArithmeticOperation(Operation, opcode="raw_binary_arith"):
     fn: str = attribute()
-    rounding_mode: RoundingMode | None = attribute()
-    flush_to_zero: bool = attribute()
+    rounding_mode: RoundingMode | None = attribute(default=None)
+    flush_to_zero: bool = attribute(default=False)
+    propagate_nan: bool = attribute(default=False)
     lhs: Var[TensorLikeTy] = operand()
     rhs: Var[TensorLikeTy] = operand()
 
@@ -421,14 +422,14 @@ def generate_bytecode(self, ctx: BytecodeContext) -> bc.Value:
                                         signedness=get_signedness(dtype))
             case "min", "float":
                 return bc.encode_MinFOp(ctx.builder, res_typeid, lhs, rhs,
-                                        propagate_nan=False,
+                                        propagate_nan=self.propagate_nan,
                                         flush_to_zero=self.flush_to_zero)
             case "max", "int":
                 return bc.encode_MaxIOp(ctx.builder, res_typeid, lhs, rhs,
                                         signedness=get_signedness(dtype))
             case "max", "float":
                 return bc.encode_MaxFOp(ctx.builder, res_typeid, lhs, rhs,
-                                        propagate_nan=False,
+                                        propagate_nan=self.propagate_nan,
                                         flush_to_zero=self.flush_to_zero)
             case "c_mod", "float":
                 # C-style modulo
@@ -444,18 +445,21 @@ def generate_bytecode(self, ctx: BytecodeContext) -> bc.Value:
 
 def binary_arithmetic_tensorlike_raw(fn: str, x: Var[TensorLikeTy], y: Var[TensorLikeTy],
                                      rounding_mode: RoundingMode | None = None,
-                                     flush_to_zero: bool = False) -> Var[TensorLikeTy]:
+                                     flush_to_zero: bool = False,
+                                     propagate_nan: bool = False) -> Var[TensorLikeTy]:
     ty = x.get_type()
     assert ty == y.get_type(), f"{ty} != {y.get_type()}"
     # FIXME: remove cutile-specific check
     check_rd_and_ftz(fn, rounding_mode, flush_to_zero, ty.tensor_dtype())
     return add_operation(RawBinaryArithmeticOperation, ty, fn=fn, lhs=x, rhs=y,
-                         rounding_mode=rounding_mode, flush_to_zero=flush_to_zero)
+                         rounding_mode=rounding_mode, flush_to_zero=flush_to_zero,
+                         propagate_nan=propagate_nan)
 
 
 def binary_arithmetic_tensorlike(fn: str, x: Var[TensorLikeTy], y: Var[TensorLikeTy],
                                  rounding_mode: RoundingMode | None = None,
-                                 flush_to_zero: bool = False) -> Var[TensorLikeTy]:
+                                 flush_to_zero: bool = False,
+                                 propagate_nan: bool = False) -> Var[TensorLikeTy]:
     x_ty = x.get_loose_type()
     y_ty = y.get_loose_type()
 
@@ -481,7 +485,8 @@ def binary_arithmetic_tensorlike(fn: str, x: Var[TensorLikeTy], y: Var[TensorLik
     if x.is_constant() and y.is_constant():
         return binop_propagate_constant(fn, x.get_constant(), y.get_constant(), common_ty)
 
-    return binary_arithmetic_tensorlike_raw(fn, x, y, rounding_mode, flush_to_zero)
+    return binary_arithmetic_tensorlike_raw(fn, x, y, rounding_mode, flush_to_zero,
+                                            propagate_nan=propagate_nan)
 
 
 @impl(operator.add, fixed_args=["add"], overload=(TensorLikeTy, TensorLikeTy))
@@ -590,8 +595,8 @@ def where(cond: Var[TensorLikeTy],
 @dataclass(eq=False)
 class Unary(Operation, opcode="unaryop"):
     fn: str = attribute()
-    rounding_mode: RoundingMode | None = attribute()
-    flush_to_zero: bool = attribute()
+    rounding_mode: RoundingMode | None = attribute(default=None)
+    flush_to_zero: bool = attribute(default=False)
     operand: Var = operand()
 
     @override
@@ -725,8 +730,7 @@ def logical_not_impl(x: Var[TensorLikeTy]) -> Var[TensorLikeTy]:
     if x.is_constant():
         return strictly_typed_const(not x.get_constant(), x.get_type())
 
-    return add_operation(Unary, x.get_type(), fn="invert", operand=x,
-                         rounding_mode=None, flush_to_zero=False)
+    return add_operation(Unary, x.get_type(), fn="invert", operand=x)
 
 
 @impl(operator.pos, overload=(TensorLikeTy,))

src/cuda/tile/_ir/ops.py

@@ -166,10 +166,11 @@ def atan2_impl(x1: Var, x2: Var) -> Var:
 @impl(ct.minimum, fixed_args=["min"])
 @impl(ct.maximum, fixed_args=["max"])
 def tile_binary_arithmetic_function_impl_with_ftz(fn: str, x: Var, y: Var,
-                                                  flush_to_zero: Var) -> Var:
+                                                  flush_to_zero: Var, propagate_nan: Var) -> Var:
     flush_to_zero = require_constant_bool(flush_to_zero)
+    propagate_nan = require_constant_bool(propagate_nan)
     return binary_arithmetic_tensorlike(fn, ensure_tile(x), ensure_tile(y),
-                                        flush_to_zero=flush_to_zero)
+                                        flush_to_zero=flush_to_zero, propagate_nan=propagate_nan)
 
 
 @impl(ct.add, fixed_args=["add"])
@@ -2331,7 +2332,7 @@ def _get_reduction_shape(shape: Tuple[int, ...],
 
 async def reduce_simple(fn: str, x: Var, axis: int | None | tuple[int, ...], keepdims: bool,
                         rounding_mode: Optional[RoundingMode] = None,
-                        flush_to_zero: bool = False) -> Var:
+                        flush_to_zero: bool = False, propagate_nan: bool = False) -> Var:
     x_type = require_tile_type(x)
     if not datatype.is_arithmetic(x_type.dtype):
         raise TileTypeError(f"Non-arithmetic dtype {x_type.dtype} is unsupported for reduction")
@@ -2351,7 +2352,8 @@ async def reduce_simple(fn: str, x: Var, axis: int | None | tuple[int, ...], kee
     async def body(lhs: tuple[Var], rhs: tuple[Var]) -> tuple[Var]:
         [lhs], [rhs] = lhs, rhs
         ret = binary_arithmetic_tensorlike(fn, lhs, rhs,
-                                           rounding_mode=rounding_mode, flush_to_zero=flush_to_zero)
+                                           rounding_mode=rounding_mode, flush_to_zero=flush_to_zero,
+                                           propagate_nan=propagate_nan)
         return (ret,)
 
     [ret] = await reduce((x,), (id_val,), axis, keepdims, body)
@@ -2399,14 +2401,17 @@ async def reduce_impl_with_rd_and_ftz(fn: str, x: Var, axis: Var, keepdims: Var,
 @impl(ct.max, fixed_args=["max"])
 @impl(ct.min, fixed_args=["min"])
 async def reduce_impl_with_ftz(fn: str, x: Var, axis: Var, keepdims: Var,
-                               flush_to_zero: Var) -> Var:
+                               flush_to_zero: Var, propagate_nan: Var) -> Var:
     axis = _parse_reduce_axis(axis)
     keepdims = require_constant_bool(keepdims)
     flush_to_zero = require_constant_bool(flush_to_zero)
-    return await reduce_simple(fn, x, axis, keepdims, flush_to_zero=flush_to_zero)
+    propagate_nan = require_constant_bool(propagate_nan)
+    return await reduce_simple(fn, x, axis, keepdims, flush_to_zero=flush_to_zero,
+                               propagate_nan=propagate_nan)
 
 
-async def argmax_argmin(fn: str, x: Var, axis: Optional[int], keepdims: bool) -> Var:
+async def argmax_argmin(fn: str, x: Var, axis: Optional[int], keepdims: bool,
+                        propagate_nan: bool = False) -> Var:
     require_tile_type(x)
     final_shape = None
     if axis is None:
@@ -2434,14 +2439,34 @@ async def argmax_argmin(fn: str, x: Var, axis: Optional[int], keepdims: bool) ->
             cmp = "gt"
         case _: assert False
 
+    is_float_dtype = datatype.is_float(x_type.dtype)
+
     async def body(lhs: tuple[Var, Var], rhs: tuple[Var, Var]) -> tuple[Var, Var]:
         lhs_val, lhs_idx = lhs
         rhs_val, rhs_idx = rhs
-        val_strict = compare_tensorlike_raw(cmp, lhs_val, rhs_val)
+        lhs_win = compare_tensorlike_raw(cmp, lhs_val, rhs_val)
         val_equal = compare_tensorlike_raw("eq", lhs_val, rhs_val)
+        if is_float_dtype:
+            lhs_is_nan = compare_tensorlike_raw("ne", lhs_val, lhs_val)
+            rhs_is_nan = compare_tensorlike_raw("ne", rhs_val, rhs_val)
+            if propagate_nan:
+                # Mirror min/max's propagate_nan=True semantics by
+                # treating NaN as the best possible value.
+                rhs_not_nan = compare_tensorlike_raw("eq", rhs_val, rhs_val)
+                lhs_nan_rhs_finite = binary_bitwise_tensorlike_raw("and_", lhs_is_nan, rhs_not_nan)
+                lhs_win = binary_bitwise_tensorlike_raw("or_", lhs_win, lhs_nan_rhs_finite)
+            else:
+                # Mirror min/max's propagate_nan=False semantics by
+                # treating NaN as the worst possible value.
+                lhs_not_nan = compare_tensorlike_raw("eq", lhs_val, lhs_val)
+                lhs_finite_rhs_nan = binary_bitwise_tensorlike_raw("and_", lhs_not_nan, rhs_is_nan)
+                lhs_win = binary_bitwise_tensorlike_raw("or_", lhs_win, lhs_finite_rhs_nan)
+            # two NaNs count as "equal" so the index tiebreak (smallest index) decides.
+            both_nan = binary_bitwise_tensorlike_raw("and_", lhs_is_nan, rhs_is_nan)
+            val_equal = binary_bitwise_tensorlike_raw("or_", val_equal, both_nan)
         index_lt = compare_tensorlike_raw("lt", lhs_idx, rhs_idx)
         val_equal_and_index_lt = binary_bitwise_tensorlike_raw("and_", val_equal, index_lt)
-        cond = binary_bitwise_tensorlike_raw("or_", val_strict, val_equal_and_index_lt)
+        cond = binary_bitwise_tensorlike_raw("or_", lhs_win, val_equal_and_index_lt)
         res = where_raw(cond, lhs_val, rhs_val)
         idx = where_raw(cond, lhs_idx, rhs_idx)
         return res, idx
@@ -2456,10 +2481,11 @@ async def body(lhs: tuple[Var, Var], rhs: tuple[Var, Var]) -> tuple[Var, Var]:
 
 @impl(ct.argmax, fixed_args=["argmax"])
 @impl(ct.argmin, fixed_args=["argmin"])
-async def argmax_argmin_impl(fn: str, x: Var, axis: Var, keepdims: Var) -> Var:
+async def argmax_argmin_impl(fn: str, x: Var, axis: Var, keepdims: Var, propagate_nan: Var) -> Var:
     axis = require_optional_constant_int(axis)
     keepdims = require_constant_bool(keepdims)
-    return await argmax_argmin(fn, x, axis, keepdims)
+    propagate_nan = require_constant_bool(propagate_nan)
+    return await argmax_argmin(fn, x, axis, keepdims, propagate_nan=propagate_nan)
 
 
 @dataclass(eq=False)

src/cuda/tile/_passes/loop_split.py

@@ -100,20 +100,17 @@ def _split_loop(loop: Loop, cond: _Condition, if_ops_to_flatten: Set[IfElse], ne
         one_var.set_type(range_dtype)
         plus_one_var = new_block.make_temp_var(loc)
         new_block.append(RawBinaryArithmeticOperation(fn="add", lhs=split_value, rhs=one_var,
-                                                      rounding_mode=None, flush_to_zero=False,
                                                       result_vars=(plus_one_var,),
                                                       loc=loc))
         plus_one_var.set_type(range_dtype)
         split_value = plus_one_var
 
     first_loop_stop = new_block.make_temp_var(loc)
     new_block.append(RawBinaryArithmeticOperation(fn="min", lhs=loop.stop, rhs=split_value,
-                                                  rounding_mode=None, flush_to_zero=False,
                                                   result_vars=(first_loop_stop,), loc=loc))
 
     second_loop_start = new_block.make_temp_var(loc)
     new_block.append(RawBinaryArithmeticOperation(fn="max", lhs=loop.start, rhs=split_value,
-                                                  rounding_mode=None, flush_to_zero=False,
                                                   result_vars=(second_loop_start,), loc=loc))
 
     for var in first_loop_stop, second_loop_start:

src/cuda/tile/_passes/rewrite_patterns.py

@@ -120,7 +120,7 @@ def fuse_mul_addsub(op: RawBinaryArithmeticOperation, ctx: MatchContext):
         neg_target = mul_op.lhs if rhs_is_mul else acc
         negated = ctx.make_temp_var(op.loc)
         ctx.set_type(negated, ctx.typeof(neg_target))
-        new_ops.append(Unary(fn="neg", operand=neg_target, rounding_mode=None, flush_to_zero=False,
+        new_ops.append(Unary(fn="neg", operand=neg_target,
                              result_vars=(negated,), loc=op.loc))
         if rhs_is_mul:
             mul_lhs = negated

src/cuda/tile/_stub.py

@@ -2638,6 +2638,12 @@ def _math_op_extra_block(f, indent):
                 f"{name} (const bool): If True, flushes subnormal inputs and results to "
                 "sign-preserving zero, default is False."
             )
+        elif name == "propagate_nan":
+            extra.append(
+                f"{name} (const bool): Float types only. If False (the default), NaN inputs "
+                "are ignored. If True, NaN propagates: a NaN makes the result NaN for min/max, "
+                "and argmin/argmax return the index of the first NaN."
+            )
     return ("\n" + textwrap.indent("\n".join(extra), indent)) if extra else ""
 
 
@@ -2725,7 +2731,8 @@ def sum(x, /, axis=None, *, keepdims=False, rounding_mode: Optional[RoundingMode
 
 @_doc_reduce_op
 @stub
-def max(x, /, axis=None, *, keepdims=False, flush_to_zero: bool = False) -> Tile:
+def max(x, /, axis=None, *, keepdims=False, flush_to_zero: bool = False,
+        propagate_nan: bool = False) -> Tile:
     """
     Examples:
 
@@ -2777,7 +2784,8 @@ def max(x, /, axis=None, *, keepdims=False, flush_to_zero: bool = False) -> Tile
 
 @_doc_reduce_op
 @stub
-def min(x, /, axis=None, *, keepdims=False, flush_to_zero: bool = False) -> Tile:
+def min(x, /, axis=None, *, keepdims=False, flush_to_zero: bool = False,
+        propagate_nan: bool = False) -> Tile:
     """
     Examples:
 
@@ -2865,7 +2873,7 @@ def prod(x, /, axis=None, *, keepdims=False, rounding_mode: Optional[RoundingMod
 
 @_doc_reduce_op
 @stub
-def argmax(x, /, axis=None, *, keepdims=False) -> Tile:
+def argmax(x, /, axis=None, *, keepdims=False, propagate_nan: bool = False) -> Tile:
     """
     Examples:
 
@@ -2919,7 +2927,7 @@ def argmax(x, /, axis=None, *, keepdims=False) -> Tile:
 
 @_doc_reduce_op
 @stub
-def argmin(x, /, axis=None, *, keepdims=False) -> Tile:
+def argmin(x, /, axis=None, *, keepdims=False, propagate_nan: bool = False) -> Tile:
     """
     Examples:
 
@@ -3462,7 +3470,7 @@ def bitwise_not(x, /) -> TileOrScalar:
 
 @_doc_binary_op('min')
 @stub
-def minimum(x, y, /, *, flush_to_zero: bool = False) -> TileOrScalar:
+def minimum(x, y, /, *, flush_to_zero: bool = False, propagate_nan: bool = False) -> TileOrScalar:
     """
     Examples:
 
@@ -3482,7 +3490,7 @@ def minimum(x, y, /, *, flush_to_zero: bool = False) -> TileOrScalar:
 
 @_doc_binary_op('max')
 @stub
-def maximum(x, y, /, *, flush_to_zero: bool = False) -> TileOrScalar:
+def maximum(x, y, /, *, flush_to_zero: bool = False, propagate_nan: bool = False) -> TileOrScalar:
     """
     Examples:
 

test/test_binary_elementwise.py

@@ -405,6 +405,28 @@ def test_array_maxmin_flush_to_zero(shape, tile, dtype, op_func, tile_op, flush_
         launch_binary(kernel, x, y, z, tile)
 
 
+@pytest.mark.parametrize("propagate_nan", [False, True])
+@pytest.mark.parametrize("op_func, torch_op", [
+    pytest.param("ct.maximum", torch.maximum, id="ct.maximum"),
+    pytest.param("ct.minimum", torch.minimum, id="ct.minimum"),
+])
+@pytest.mark.parametrize("dtype", float_dtypes, ids=dtype_id)
+def test_array_maxmin_nan(shape, tile, dtype, op_func, torch_op, propagate_nan, tmp_path):
+    x = make_tensor(shape, dtype=dtype, device='cuda')
+    y = make_tensor(shape, dtype=dtype, device='cuda')
+    x[1::3] = float("nan")
+    y[0::3] = float("nan")
+    z = torch.zeros_like(x)
+    kernel = array_kernel("maxmin_nan",
+                          f"tz = {op_func}(tx, ty, propagate_nan={propagate_nan})", tmp_path)
+    launch_binary(kernel, x, y, z, tile)
+    if propagate_nan:
+        ref = torch_op(x, y)
+    else:
+        ref = torch.where(torch.isnan(x), y, torch.where(torch.isnan(y), x, torch_op(x, y)))
+    assert_equal(z, ref)
+
+
 @pytest.mark.parametrize("is_constant", [False, True])
 @pytest.mark.parametrize("x", [100, 100.0])
 @pytest.mark.parametrize("y", [23, 2.3])