* Added support for additional rounding modes

Signed-off-by: Gideon Kassa <gkassa@nvidia.com>

Author: gideon116

changelog.d/float-rounding-modes.md

@@ -0,0 +1 @@
+Added support for specifying IEEE rounding modes in float to float `astype` conversions.

experimental/cuda-lang/test/passes/test_flatten_cfg.py

@@ -23,7 +23,7 @@ def test_kernel(A):
             A[0] = 0
 
     # BEFORE: $[[ITEM:[0-9]+]]: int32 = load_pointer
-    # BEFORE: $[[ITEM_CASTED:[0-9]+]]: bool_ = tile_astype(x=$[[ITEM]])
+    # BEFORE: $[[ITEM_CASTED:[0-9]+]]: bool_ = tile_astype(x=$[[ITEM]], rounding_mode=None)
     # BEFORE: if(cond=$[[ITEM_CASTED]])
     # BEFORE: then
     # BEFORE:   store_pointer
@@ -37,7 +37,7 @@ def test_kernel(A):
 
     # AFTER: ^entry({{.+}}):
     # AFTER:   $[[ITEM:[0-9]+]]: int32 = load_pointer
-    # AFTER:   $[[ITEM_CASTED:[0-9]+]]: bool_ = tile_astype(x=$[[ITEM]])
+    # AFTER:   $[[ITEM_CASTED:[0-9]+]]: bool_ = tile_astype(x=$[[ITEM]], rounding_mode=None)
     # AFTER:   cond_br $[[ITEM_CASTED]]: bool_ ^then() ^else()
     # AFTER: ^then():
     # AFTER:   store_pointer
@@ -87,7 +87,7 @@ def test_kernel(cond1, cond2):
 def test_flatten_ifelse_phi_merge():
     def test_kernel(A):
         # CHECK: $[[ITEM:[0-9]+]]: int32 = load_pointer
-        # CHECK: $[[ITEM_BOOL:[0-9]+]]: bool_ = tile_astype(x=$[[ITEM]])
+        # CHECK: $[[ITEM_BOOL:[0-9]+]]: bool_ = tile_astype(x=$[[ITEM]], rounding_mode=None)
         # CHECK: cond_br $[[ITEM_BOOL]]: bool_ ^then() ^else()
         if A[0]:
             # CHECK: $[[ITEM_1:[0-9]+]]: int32 = load_pointer

src/cuda/tile/_ir/arithmetic_ops.py

@@ -127,24 +127,28 @@ def broadcast_to(x: Var[TensorLikeTy], shape: Sequence[int]) -> Var[TensorLikeTy
 @dataclass(eq=False)
 class TileAsType(Operation, opcode="tile_astype"):
     x: Var = operand()
+    rounding_mode: RoundingMode | None = attribute(default=None)
 
     @override
     def generate_bytecode(self, ctx: BytecodeContext) -> bc.Value:
         value = ctx.get_value(self.x)
-        return convert_dtype(ctx, value, ctx.typeof(self.x), ctx.typeof(self.result_var))
+
+        return convert_dtype(ctx, value, ctx.typeof(self.x), ctx.typeof(self.result_var),
+                             rounding_mode=self.rounding_mode)
 
 
-def astype(x: Var[TensorLikeTy], dtype: DType) -> Var[TensorLikeTy]:
+def astype(x: Var[TensorLikeTy], dtype: DType, *,
+           rounding_mode: RoundingMode | None = None) -> Var[TensorLikeTy]:
     x_ty = x.get_type()
     if x_ty.tensor_dtype() == dtype:
         return x
 
-    if x.is_constant():
+    if x.is_constant() and rounding_mode in (None, RoundingMode.RN):
         val = numeric_dtype_category(dtype).pytype(x.get_constant())
         return strictly_typed_const(val, x.ctx.typing_hooks.get_tensor_like_type(dtype, ()))
 
     result_ty = x.ctx.typing_hooks.get_tensor_like_type(dtype, x_ty.tensor_shape())
-    return add_operation(TileAsType, result_ty, x=x)
+    return add_operation(TileAsType, result_ty, x=x, rounding_mode=rounding_mode)
 
 
 def dtype_constructor(new_dtype: DType, x: Var) -> Var[TensorLikeTy]:

src/cuda/tile/_ir/ops.py

@@ -21,6 +21,7 @@
     make_aggregate, MemoryEffect, attribute, operand,
     BlockRestriction, add_operation_variadic,
 )
+from cuda.tile._ir.ops_utils import get_ftof_rounding_min_version
 from .aggregate_support import unflatten_aggregates
 from .arithmetic_ops import reshape, broadcast_to, astype, compare_tensorlike, \
     binary_bitwise_tensorlike, bitwise_shift_tensorlike, binary_arithmetic_tensorlike, \
@@ -2810,10 +2811,29 @@ def assert_impl(cond: Var, message: Var) -> None:
 
 
 @impl(ct.astype)
-def astype_impl(x: Var, dtype: Var) -> Var:
-    require_tile_type(x)
+def astype_impl(x: Var, dtype: Var, rounding_mode: Var) -> Var:
+    x_ty = require_tile_type(x)
     dtype = require_dtype_spec(dtype)
-    return astype(x, dtype)
+    rounding_mode = require_optional_constant_enum(rounding_mode, RoundingMode)
+
+    is_ftof = datatype.is_float(x_ty.tensor_dtype()) and datatype.is_float(dtype)
+    if not is_ftof and rounding_mode is not None:
+        raise TileTypeError("rounding_mode is only valid for float to float conversions")
+
+    if x_ty.tensor_dtype() == dtype:
+        return x
+
+    if is_ftof:
+        rounding_mode, min_bc_version = get_ftof_rounding_min_version(x_ty.tensor_dtype(),
+                                                                      dtype, rounding_mode)
+        cur_bc_version = Builder.get_current().ir_ctx.tileiras_version
+        if min_bc_version is not None and cur_bc_version < min_bc_version:
+            raise TileUnsupportedFeatureError(
+                f"The requested conversion and rounding_mode require tileiras "
+                f"{min_bc_version.as_string()} or later. Current version is "
+                f"{cur_bc_version.as_string()}.")
+
+    return astype(x, dtype, rounding_mode=rounding_mode)
 
 
 @dataclass(eq=False)

src/cuda/tile/_ir/ops_utils.py

@@ -4,6 +4,7 @@
 import itertools
 import math
 from contextlib import contextmanager
+from collections import defaultdict
 
 from dataclasses import dataclass, field
 from typing import Optional, Tuple, Dict, Any, Sequence, Literal
@@ -123,6 +124,7 @@ def get_default_rounding_mode(opname: Optional[str] = None):
     RoundingMode.RZ: bc.RoundingMode.ZERO,
     RoundingMode.RM: bc.RoundingMode.NEGATIVE_INF,
     RoundingMode.RP: bc.RoundingMode.POSITIVE_INF,
+    RoundingMode.RA: bc.RoundingMode.NEAREST_AWAY,
     RoundingMode.FULL: bc.RoundingMode.FULL,
     RoundingMode.APPROX: bc.RoundingMode.APPROX,
     RoundingMode.RZI: bc.RoundingMode.NEAREST_INT_TO_ZERO
@@ -225,6 +227,79 @@ def check_shapes_eq(a: TileTy, b: TileTy,
                             f"got {a.shape} and {b.shape}", loc)
 
 
+F64 = datatype.float64
+F32 = datatype.float32
+TF32 = datatype.tfloat32
+F16 = datatype.float16
+BF16 = datatype.bfloat16
+F8E5M2 = datatype.float8_e5m2
+F8E8M0FNU = datatype.float8_e8m0fnu
+F8E4M3FN = datatype.float8_e4m3fn
+F4E2M1FN = datatype.float4_e2m1fn
+ALL = (F64, F32, TF32, F16, BF16, F8E5M2, F8E8M0FNU, F8E4M3FN, F4E2M1FN)
+B133 = BytecodeVersion.V_13_3
+B134 = BytecodeVersion.V_13_4
+
+_FTOF_ROUNDING_ROWS = (
+        {(i, F64): (RoundingMode.RN, None) for i in ALL},
+        {(i, F32): (RoundingMode.RN, None) for i in ALL},
+        {(i, TF32): (RoundingMode.RN, None) for i in ALL},
+        {(i, BF16): (RoundingMode.RN, None) for i in ALL},
+        {(i, F16): (RoundingMode.RN, None) for i in ALL},
+        {(i, F8E4M3FN): (RoundingMode.RN, None) for i in ALL},
+        {(i, F8E5M2): (RoundingMode.RN, None) for i in ALL},
+        {(i, F4E2M1FN): (RoundingMode.RN, None) for i in ALL},
+
+        {(i, F64): (RoundingMode.RZ, B134) for i in ALL},
+        {(i, F32): (RoundingMode.RZ, B134) for i in ALL},
+        {(i, TF32): (RoundingMode.RZ, B134) for i in ALL},
+        {(i, F16): (RoundingMode.RZ, B134) for i in ALL},
+        {(i, BF16): (RoundingMode.RZ, B134) for i in ALL},
+        {(i, F8E8M0FNU): (RoundingMode.RZ, B133) for i in ALL if i not in {F64, F8E5M2, F8E4M3FN}},
+        {(i, F8E8M0FNU): (RoundingMode.RZ, B134) for i in (F64, F8E5M2, F8E4M3FN)},
+
+        {(i, F64): (RoundingMode.RM, B134) for i in ALL},
+        {(i, F32): (RoundingMode.RM, B134) for i in ALL if i not in {F8E8M0FNU}},
+        {(i, F16): (RoundingMode.RM, B134) for i in ALL if i not in {F64, F32, TF32, F8E8M0FNU}},
+        {(i, F64): (RoundingMode.RP, B134) for i in ALL},
+        {(i, F32): (RoundingMode.RP, B134) for i in ALL if i not in {F8E8M0FNU}},
+        {(i, F16): (RoundingMode.RP, B134) for i in ALL if i not in {F64, F32, TF32, F8E8M0FNU}},
+        {(i, F8E8M0FNU): (RoundingMode.RP, B133) for i in ALL if i not in {F64, F8E5M2, F8E4M3FN}},
+        {(i, F8E8M0FNU): (RoundingMode.RP, B134) for i in (F64, F8E5M2, F8E4M3FN)},
+
+        {(i, F64): (RoundingMode.RA, B134) for i in ALL},
+        {(i, F32): (RoundingMode.RA, B134) for i in ALL if i not in {F64, F8E8M0FNU}},
+        {(i, TF32): (RoundingMode.RA, B134) for i in ALL if i not in {F64, F8E8M0FNU}},
+        {(i, F16): (RoundingMode.RA, B134) for i in ALL if i not in {F64, F32, TF32, F8E8M0FNU}}
+)
+
+# {(from, to): {RoundingMode_1: BC_Version, RoundingMode_2: BC_Version}}
+FTOF_ROUNDING_REGISTRY = defaultdict(dict)
+for row in _FTOF_ROUNDING_ROWS:
+    for from_to, (mode, version) in row.items():
+        FTOF_ROUNDING_REGISTRY[from_to][mode] = version
+
+
+def get_ftof_rounding_min_version(from_dtype: datatype.DType, to_dtype: datatype.DType,
+                                  rounding_mode: RoundingMode | None
+                                  ) -> tuple[RoundingMode, BytecodeVersion | None]:
+
+    conversion_pair = (from_dtype, to_dtype)
+    supported = FTOF_ROUNDING_REGISTRY.get(conversion_pair, None)
+    if supported is None:
+        raise TileTypeError(f"float conversion from {from_dtype} to {to_dtype} "
+                            "is not supported")
+
+    rounding_mode = RoundingMode.RN if rounding_mode is None else rounding_mode
+    if rounding_mode not in supported:
+        raise TileTypeError(
+            f"rounding_mode={rounding_mode} is not supported "
+            f"for conversion from {from_dtype} to {to_dtype}, "
+            f"supported rounding modes for this conversion are {tuple(supported.keys())}")
+
+    return (rounding_mode, supported[rounding_mode])
+
+
 class CompareOrdering(Enum):
     ORDERED = "ordered"
     UNORDERED = "unordered"

src/cuda/tile/_ir2bytecode.py

@@ -18,6 +18,7 @@
 import cuda.tile._bytecode as bc
 from cuda.tile._compiler_options import CompilerOptions
 from cuda.tile._exception import TileInternalError, TileError, FunctionDesc
+from cuda.tile._numeric_semantics import RoundingMode
 from cuda.tile._ir.ir import Block, Loc, Var, IRContext
 from cuda.tile._ir.ops_utils import (
     padding_mode_to_bytecode, rounding_mode_to_bytecode,
@@ -130,7 +131,8 @@ def _flatten_bools(value) -> tuple[bool]:
     return sum((_flatten_bools(v) for v in value), start=())
 
 
-def _get_type_conversion_encoder(from_dtype: Type, to_dtype: Type):
+def _get_type_conversion_encoder(from_dtype: Type, to_dtype: Type, *,
+                                 rounding_mode: RoundingMode | None = None):
 
     def kind(t):
         if datatype.is_float(t):
@@ -140,11 +142,16 @@ def kind(t):
         raise TileInternalError(f'Unsupported dtype: {t}')
 
     from_kind, to_kind = kind(from_dtype), kind(to_dtype)
+
+    if rounding_mode is not None:
+        rounding_mode = rounding_mode_to_bytecode[rounding_mode]
+    else:
+        rounding_mode = bc.RoundingMode.NEAREST_EVEN
+
     round_to_float = rounding_mode_to_bytecode[get_default_rounding_mode()]
     partial = functools.partial
     match from_kind, to_kind:
-        case 'f', 'f': return partial(bc.encode_FToFOp,
-                                      rounding_mode=bc.RoundingMode.NEAREST_EVEN)
+        case 'f', 'f': return partial(bc.encode_FToFOp, rounding_mode=rounding_mode)
         case 'f', 'si': return partial(bc.encode_FToIOp,
                                        signedness=bc.Signedness.Signed,
                                        rounding_mode=bc.RoundingMode.NEAREST_INT_TO_ZERO,
@@ -171,8 +178,8 @@ def kind(t):
     raise NotImplementedError(f"Type coversion from {from_dtype} to {to_dtype} not implemented")
 
 
-def convert_dtype(ctx: "BytecodeContext", val: bc.Value,
-                  fromty: Type, toty: Type) -> bc.Value:
+def convert_dtype(ctx: "BytecodeContext", val: bc.Value, fromty: Type,
+                  toty: Type, *, rounding_mode: RoundingMode | None = None) -> bc.Value:
     from_dtype = fromty.dtype if isinstance(fromty, TileTy) else fromty
     to_dtype = toty.dtype if isinstance(toty, TileTy) else toty
     toty_id = typeid(ctx.type_table, toty)
@@ -188,7 +195,7 @@ def convert_dtype(ctx: "BytecodeContext", val: bc.Value,
             comparison_predicate=bc.ComparisonPredicate.NOT_EQUAL,
             signedness=datatype.get_signedness(from_dtype))
     else:
-        encoder = _get_type_conversion_encoder(from_dtype, to_dtype)
+        encoder = _get_type_conversion_encoder(from_dtype, to_dtype, rounding_mode=rounding_mode)
         return encoder(ctx.builder, toty_id, val)
 
 

src/cuda/tile/_numeric_semantics.py

@@ -11,7 +11,7 @@ class RoundingMode(Enum):
     """
 
     RN = "nearest_even"
-    """Rounds the nearest (ties to even)."""
+    """Round to nearest (ties to even)."""
 
     RZ = "zero"
     """Round towards zero (truncate)."""
@@ -22,6 +22,9 @@ class RoundingMode(Enum):
     RP = "positive_inf"
     """Round towards positive infinity."""
 
+    RA = "nearest_away"
+    """Round to nearest (ties away from zero)."""
+
     FULL = "full"
     """Full precision rounding mode."""
 

src/cuda/tile/_stub.py

@@ -605,9 +605,9 @@ def transpose(self, axis0=None, axis1=None) -> "Tile":
         """See :py:func:`transpose`."""
         return transpose(self, axis0, axis1)
 
-    def astype(self, dtype) -> "Tile":
+    def astype(self, dtype, *, rounding_mode: Optional[RoundingMode] = None) -> "Tile":
         """See :py:func:`astype`."""
-        return astype(self, dtype)
+        return astype(self, dtype, rounding_mode=rounding_mode)
 
     @function
     def __index__(self) -> int:
@@ -2559,12 +2559,15 @@ def transpose(x, /, axis0=None, axis1=None) -> Tile:
 
 
 @stub
-def astype(x, dtype, /) -> Tile:
+def astype(x, dtype, /, *, rounding_mode: Optional[RoundingMode] = None) -> Tile:
     """Converts a tile to the specified data type.
 
     Args:
         x (Tile): input tile.
         dtype (DType): target data type.
+        rounding_mode (RoundingMode): optional rounding mode for explicit float to
+                                      float conversions. Default is round to nearest
+                                      with ties to even.
 
     Returns:
         Tile: