astile requires dtype

Signed-off-by: Boyan Li <boyanl@nvidia.com>

Author: blinxt

src/cuda/tile/_ir/op_impl.py

@@ -323,12 +323,6 @@ def require_dtype_spec(var: Var) -> DType:
     return ty.dtype
 
 
-def require_optional_dtype_spec(var: Var) -> DType | None:
-    if var.is_constant() and var.get_constant() is None:
-        return None
-    return require_dtype_spec(var)
-
-
 def require_constant_pointer_info(var: Var) -> PointerInfo:
     ty = var.get_type()
     if not isinstance(ty, PointerInfoTy):

src/cuda/tile/_ir/ops.py

@@ -36,7 +36,6 @@
 from .hir import ResolvedName
 from .op_impl import (
     ImplRegistry, is_0d_tile, require_constant_int, require_constant_int_tuple,
-    require_optional_dtype_spec,
     require_signed_integer_0d_tile_type,
     require_tile_type, normalize_axis, require_dtype_spec,
     require_constant_bool, require_optional_constant_enum,
@@ -3413,10 +3412,10 @@ def _cat_tuple(tiles: tuple[Var, ...]) -> Var:
 
 @impl(ct.astile)
 def astile_impl(value: Var, dtype: Var) -> Var:
-    dtype: Optional[DType] = require_optional_dtype_spec(dtype)
+    dtype = require_dtype_spec(dtype)
     value_ty = value.get_type()
     if is_0d_tile(value_ty):
-        return value if dtype is None else astype(value, dtype)
+        return astype(value, dtype)
 
     if not isinstance(value_ty, TupleTy):
         raise TileTypeError(
@@ -3425,9 +3424,6 @@ def astile_impl(value: Var, dtype: Var) -> Var:
 
     shape = _tuple_shape(value_ty, path=())
     tiles = _flatten_tuple(value)
-    dtype = (functools.reduce(promote_dtypes, (require_0d_tile_type(t).dtype for t in tiles))
-             if dtype is None
-             else dtype)
 
     if value.is_constant():
         return _const(shape, value.get_constant(), dtype)

src/cuda/tile/_stub.py

@@ -1962,20 +1962,18 @@ def zeros(shape, dtype) -> Tile:
 
 
 @stub
-def astile(value, /, *, dtype: Optional[DType] = None) -> Tile:
+def astile(value, /, *, dtype: DType) -> Tile:
     """Creates a tile from a value.
 
     Args:
         value (scalar | (nested) tuple of scalar): A scalar (yielding a 0-d tile),
             or a (possibly nested) tuple of scalars whose nesting determines the 
             tile's shape. Every tuple's length must be a power of two, and sibling tuples
             at each level must have uniform length.
-        dtype (DType, optional): The |Data type| of the tile. If ``None``, the
-            dtype is inferred from ``value``.
+        dtype (DType): The |Data type| of the tile.
 
     Returns:
-        Tile: A tile shaped from ``value``, with elements cast to ``dtype`` if
-        given.
+        Tile: A tile shaped from ``value``, with elements cast to ``dtype``.
 
     Examples:
 

test/test_astile.py

@@ -129,24 +129,6 @@ def kernel(X):
     assert_equal(x, ref)
 
 
-@pytest.mark.parametrize("value,expected_dtype", [
-    ((1,), ct.int32),
-    ((-2**42,), ct.int64),
-    ((2**63,), ct.uint64),
-    ((2.5,), ct.float32),
-    ((True,), ct.bool_),
-    ((1, True), ct.int32),
-    ((1, 2.5, True, 2**63), ct.float32),
-])
-def test_astile_dtype_infer_const(value, expected_dtype):
-    @ct.kernel
-    def kernel():
-        t = ct.astile(value)
-        ct.static_assert(t.dtype == expected_dtype)
-
-    ct.launch(torch.cuda.current_stream(), (1,), kernel, ())
-
-
 def test_astile_scalar_runtime():
     @ct.kernel
     def kernel(X, a: float):
@@ -182,24 +164,6 @@ def kernel(X, a: int, b: int, c: float, d: bool):
     assert_equal(x, ref)
 
 
-@pytest.mark.parametrize("ann1,ann2,val1,val2,expected_dtype", [
-    (int, int, 1, 2, ct.int32),
-    (int, ct.ScalarInt64, 1, 2, ct.int64),
-    (float, float, 1.5, 2.5, ct.float32),
-    (bool, bool, True, False, ct.bool_),
-    (int, float, 1, 2.5, ct.float32),
-    (int, bool, 1, True, ct.int32),
-    (float, bool, 2.5, True, ct.float32),
-])
-def test_astile_dtype_infer_runtime(ann1, ann2, val1, val2, expected_dtype):
-    @ct.kernel
-    def kernel(a: ann1, b: ann2):
-        t = ct.astile((a, b))
-        ct.static_assert(t.dtype == expected_dtype)
-
-    ct.launch(torch.cuda.current_stream(), (1,), kernel, (val1, val2))
-
-
 def test_astile_3d_mixed():
     @ct.kernel
     def kernel(X, a: int, b: int, c: float, d: bool):
@@ -263,7 +227,7 @@ def kernel(X):
 def test_astile_top_level_not_supported():
     @ct.kernel
     def kernel():
-        ct.astile(ct.full((4,), 1, dtype=ct.int32))
+        ct.astile(ct.full((4,), 1, dtype=ct.int32), dtype=ct.int32)
     with pytest.raises(TileTypeError,
                        match=r"Expected a scalar or \(possibly nested\) tuple of scalars"):
         ct.launch(torch.cuda.current_stream(), (1,), kernel, ())