astile

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

Author: blinxt

changelog.d/astile.md

@@ -0,0 +1,4 @@
+<!--- SPDX-FileCopyrightText: Copyright (c) <2026> NVIDIA CORPORATION & AFFILIATES. All rights reserved. -->
+<!--- SPDX-License-Identifier: Apache-2.0 -->
+
+- Added `ct.astile` for creating a tile from a scalar (yielding a 0-d tile) or a (possibly nested) tuple of scalars whose nesting determines the tile's shape.

docs/source/operations.rst

@@ -37,10 +37,12 @@ Factory
    :nosignatures:
 
    arange
+   astile
    full
    ones
    zeros
 
+
 .. _operations-shape-dtype:
 
 Shape & DType

src/cuda/tile/__init__.py

@@ -74,6 +74,7 @@
     argmax,
     argmin,
     assert_,
+    astile,
     astype,
     atan2,
     atomic_add,
@@ -235,6 +236,7 @@
     "argmax",
     "argmin",
     "assert_",
+    "astile",
     "astype",
     "atan2",
     "atomic_add",

src/cuda/tile/_ir/op_impl.py

@@ -227,15 +227,15 @@ class _CurrentStub(threading.local):
 _current_stub = _CurrentStub()
 
 
-def _is_0d_tile(ty: Type, dtype_predicate: Callable[[DType], bool] = lambda _: True) -> bool:
+def is_0d_tile(ty: Type, dtype_predicate: Callable[[DType], bool] = lambda _: True) -> bool:
     return isinstance(ty, TileTy) and ty.ndim == 0 and dtype_predicate(ty.dtype)
 
 
 def require_constant_int(var: Var) -> int:
     if not var.is_constant():
         raise _make_type_error("Expected an integer constant, but given value is not constant", var)
     ty = var.get_type()
-    if not _is_0d_tile(ty, is_integral):
+    if not is_0d_tile(ty, is_integral):
         raise _make_type_error(f"Expected an integer constant, but given value has type {ty}",
                                var)
     return var.get_constant()
@@ -251,14 +251,14 @@ def require_constant_bool(var: Var) -> bool:
     if not var.is_constant():
         raise _make_type_error("Expected a boolean constant, but given value is not constant", var)
     ty = var.get_type()
-    if not _is_0d_tile(ty, is_boolean):
+    if not is_0d_tile(ty, is_boolean):
         raise _make_type_error(f"Expected a boolean constant, but given value has type {ty}", var)
     return var.get_constant()
 
 
 def require_constant_scalar(var: Var) -> bool | int | float:
     ty = var.get_type()
-    if not _is_0d_tile(ty):
+    if not is_0d_tile(ty):
         raise _make_type_error(f"Expected a scalar constant, but given value has type {ty}", var)
     if not var.is_constant():
         raise _make_type_error(f"Expected a constant, but given value has non-constant type {ty}",
@@ -274,7 +274,7 @@ def require_constant_scalar_tuple(var: Var) -> tuple[bool | int | float, ...]:
     tuple_val = var.get_aggregate()
     assert isinstance(tuple_val, TupleValue)
     for i, (item_ty, item) in enumerate(zip(ty.value_types, tuple_val.items, strict=True)):
-        if not _is_0d_tile(item_ty):
+        if not is_0d_tile(item_ty):
             raise _make_type_error(f"Expected a tuple of scalar constants,"
                                    f" but item at position #{i} has type {item_ty}", var)
         if not item.is_constant():
@@ -323,6 +323,12 @@ 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_optional_constant_enum(var: Var, enum: EnumMeta):
     if var.is_constant() and var.get_constant() is None:
         return None
@@ -363,14 +369,14 @@ def require_constant_int_tuple(var: Var, allow_single_int: bool = False) -> Tupl
                                " but given value is not constant", var)
 
     ty = var.get_type()
-    if allow_single_int and _is_0d_tile(ty):
+    if allow_single_int and is_0d_tile(ty):
         return require_constant_int(var),
 
     if not isinstance(ty, TupleTy):
         raise _make_type_error(f"Expected a tuple, but given value has type {ty}", var)
 
     for i, item_ty in enumerate(ty.value_types):
-        if not _is_0d_tile(item_ty, is_integral):
+        if not is_0d_tile(item_ty, is_integral):
             raise _make_type_error(f"Expected a tuple of integers,"
                                    f" but element #{i} has type {item_ty}", var)
 
@@ -472,7 +478,7 @@ def require_signed_integer_0d_tile_type(var: Var) -> TileTy:
 
 def require_bool(var: Var) -> TileTy:
     ty = var.get_type()
-    if not _is_0d_tile(ty, is_boolean):
+    if not is_0d_tile(ty, is_boolean):
         raise _make_type_error(f"Expected a bool, but given value has type {ty}", var)
     return ty
 

src/cuda/tile/_ir/ops.py

@@ -4,6 +4,7 @@
 import builtins
 import dataclasses
 import enum
+import functools
 import math
 import operator
 from contextlib import contextmanager
@@ -19,7 +20,7 @@
 from cuda.tile import _datatype as datatype
 from cuda.tile import RoundingMode, MemoryOrder, MemoryScope
 from cuda.tile._mutex import tile_mutex
-from cuda.tile._exception import TileTypeError, TileSyntaxError, TileError, \
+from cuda.tile._exception import TileInternalError, TileTypeError, TileSyntaxError, TileError, \
     TileStaticAssertionError, TileStaticEvalError, TileValueError, TileUnsupportedFeatureError
 from cuda.tile._ir.ir import (
     Operation, Var, Loc, Block, add_operation, Builder, enter_nested_block, nested_block,
@@ -34,7 +35,8 @@
 from . import hir, hir_stubs
 from .hir import ResolvedName
 from .op_impl import (
-    ImplRegistry, require_constant_int, require_constant_int_tuple,
+    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,
@@ -3283,14 +3285,14 @@ def num_tiles_impl(array: Var, axis: Var, shape: Var, order: Var) -> Var:
     return space_shape[axis]
 
 
-def full_const(shape: Sequence[int], fill_value: int | float, dtype: DType) -> Var:
+def _const(shape: Sequence[int], value: int | float | bool | tuple, dtype: DType) -> Var:
     res_ty = make_tile_ty(dtype, shape)
-    return strictly_typed_const(fill_value, res_ty)
+    return strictly_typed_const(value, res_ty)
 
 
 def full(shape: Sequence[int], fill_value: Var, dtype: DType) -> Var:
     if fill_value.is_constant():
-        return full_const(shape, fill_value.get_constant(), dtype)
+        return _const(shape, fill_value.get_constant(), dtype)
     fill_value = astype(fill_value, dtype)
     return broadcast_to(fill_value, shape)
 
@@ -3307,14 +3309,93 @@ def full_impl(shape: Var, fill_value: Var, dtype: Var) -> Var:
 def ones_impl(shape: Var, dtype: Var) -> Var:
     shape = require_constant_shape(shape, allow_single_int=True)
     dtype = require_dtype_spec(dtype)
-    return full_const(shape, 1, dtype)
+    return _const(shape, 1, dtype)
 
 
 @impl(ct.zeros)
 def zeros_impl(shape: Var, dtype: Var) -> Var:
     shape = require_constant_shape(shape, allow_single_int=True)
     dtype = require_dtype_spec(dtype)
-    return full_const(shape, 0, dtype)
+    return _const(shape, 0, dtype)
+
+
+def _path_str(path: tuple[int, ...]) -> str:
+    return "value" + "".join(f"[{i}]" for i in path)
+
+
+def _tuple_shape(ty: Type, path: tuple[int, ...]) -> tuple[int, ...]:
+    path_str = _path_str(path)
+    if not isinstance(ty, TupleTy):
+        if not isinstance(ty, TileTy):
+            raise TileTypeError(
+                f"Expected scalar elements at {path_str}; "
+                f"got element of type {ty}")
+
+        if ty.ndim != 0:
+            raise TileTypeError(
+                f"Expected scalar elements at {path_str}; "
+                f"got a tile of shape {ty.shape}")
+
+        assert is_0d_tile(ty)
+        return ()
+
+    n = len(ty)
+    if not _is_power_of_2(n):
+        raise TileTypeError(f"Tuple length {n} at {path_str} is not a power of 2")
+
+    inner_shapes = {_tuple_shape(t, path + (i,)) for i, t in enumerate(ty.value_types)}
+    if len(inner_shapes) != 1:
+        raise TileTypeError(f"Tuple has non-uniform inner shapes at {path_str}")
+
+    return (n,) + inner_shapes.pop()
+
+
+def _flatten_tuple(value: Var) -> tuple[Var, ...]:
+    value_ty = value.get_type()
+    if not isinstance(value_ty, TupleTy):
+        return (value,)
+    return sum((_flatten_tuple(i) for i in value.get_aggregate().items), start=())
+
+
+def _cat_tuple(tiles: tuple[Var, ...]) -> Var:
+    if len(tiles) == 0:
+        raise TileInternalError("Expected non-empty tile tuple")
+
+    if len(tiles) == 1:
+        require_0d_tile_type(tiles[0])
+        return reshape(tiles[0], (1,))
+
+    assert len(tiles) % 2 == 0
+    mid = len(tiles) // 2
+    left = _cat_tuple(tiles[:mid])
+    right = _cat_tuple(tiles[mid:])
+    return cat((left, right), axis=0)
+
+
+@impl(ct.astile)
+def astile_impl(value: Var, dtype: Var) -> Var:
+    dtype: Optional[DType] = require_optional_dtype_spec(dtype)
+    value_ty = value.get_type()
+    if is_0d_tile(value_ty):
+        return value if dtype is None else astype(value, dtype)
+
+    if not isinstance(value_ty, TupleTy):
+        raise TileTypeError(
+                f"Expected a scalar or (possibly nested) tuple of scalars; "
+                f"got value of type {value_ty}")
+
+    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)
+
+    tiles = tuple(astype(t, dtype) for t in tiles)
+    flat = _cat_tuple(tiles)
+    return reshape(flat, shape)
 
 
 _TileShape = Tuple[int, ...]
@@ -4116,29 +4197,24 @@ def generate_bytecode(self, ctx: BytecodeContext) -> bc.Value:
         return bc.encode_CatOp(ctx.builder, return_type_id, x_value, y_value, self.axis)
 
 
-def cat(tiles: Var, axis: int) -> Var:
-    tuple_ty = require_tuple_type(tiles)
-    items = tiles.get_aggregate().items
-    if len(items) == 1:
-        return items[0]
-
-    if len(tuple_ty) == 0:
+def cat(tiles: tuple[Var, ...], axis: int) -> Var:
+    if len(tiles) == 0:
         raise TileTypeError("cat() received an empty tuple")
-    elif len(items) == 1:
-        return items[0]
-    elif len(tuple_ty) > 2:
-        raise TileTypeError(f"cat() supports at most 2 tiles, got {len(tuple_ty)}")
+    if len(tiles) == 1:
+        return tiles[0]
+    if len(tiles) > 2:
+        raise TileTypeError(f"cat() supports at most 2 tiles, got {len(tiles)}")
 
-    x_tile, y_tile = items
+    x_tile, y_tile = tiles
 
-    if not isinstance(first_tile := tuple_ty.value_types[0], TileTy):
-        raise TileTypeError(f"Expected tuple of Tile, got a {first_tile}")
+    if not isinstance(first_tile_ty := tiles[0].get_type(), TileTy):
+        raise TileTypeError(f"Expected tuple of Tile, got a {first_tile_ty}")
 
-    dtype = first_tile.dtype
-    rank = first_tile.ndim
-    shape_value = list(first_tile.shape)
+    dtype = first_tile_ty.dtype
+    rank = first_tile_ty.ndim
+    shape_value = list(first_tile_ty.shape)
     axis = normalize_axis(axis, rank)
-    for tile_ty in tuple_ty.value_types[1:]:
+    for tile_ty in (t.get_type() for t in tiles[1:]):
         if not isinstance(tile_ty, TileTy):
             raise TileTypeError(f"Expected tuple of Tile, got a {tile_ty}")
         if tile_ty.ndim != rank:
@@ -4167,8 +4243,9 @@ def _is_power_of_2(x: int):
 
 @impl(ct.cat)
 def cat_impl(tiles: Var, axis: Var) -> Var:
+    require_tuple_type(tiles)
     const_axis = require_constant_int(axis)
-    return cat(tiles, const_axis)
+    return cat(tiles.get_aggregate().items, const_axis)
 
 
 # Does not support broadcasting or type promotion

src/cuda/tile/_ir/typing_support.py

@@ -190,7 +190,7 @@ def get_constant_value(val: Any) -> Any:
         return val
     if is_dtype(val):
         return to_dtype(val)
-    if isinstance(val, tuple) and not any(isinstance(x, tuple) for x in val):
+    if isinstance(val, tuple):
         return tuple(get_constant_value(x) for x in val)
     typ = type(val)
     prefix = "" if typ.__module__ == "builtins" else f"{typ.__module__}."

src/cuda/tile/_ir2bytecode.py

@@ -24,7 +24,7 @@
 )
 from cuda.tile._ir.type import (
     PartitionViewTy, StridedViewTy, GatherScatterViewTy, Type, TileTy, PointerTy, TokenTy,
-    TupleTy, ArrayTy, size_to_bytecode,
+    ArrayTy, size_to_bytecode,
 )
 
 
@@ -103,6 +103,31 @@ def _constant_to_bytes(value: int | float, dtype: DType) -> bytes:
         raise TypeError(f"Cannot make a constant out of {dtype}")
 
 
+# Encode a potentially nested constant tuple as the raw row-major byte buffer according to MLIR's
+# "DenseElementsAttr" non-splat format.
+def _constant_tuple_to_bytes(value, dtype: DType, shape: tuple[int, ...]) -> bytes:
+    # Note that MLIR requires bit packing for non-splat DenseElementsAttr<i1>
+    if dtype == datatype.bool_ and isinstance(value, tuple):
+        flat_bools = _flatten_bools(value)
+        bits = 0
+        for i, v in enumerate(flat_bools):
+            if v:
+                bits |= 1 << i
+        return bits.to_bytes((len(flat_bools) + 7) // 8, "little")
+
+    if len(shape) == 0:
+        return _constant_to_bytes(value, dtype)
+
+    assert len(value) == shape[0]
+    return b"".join(_constant_tuple_to_bytes(c, dtype, shape[1:]) for c in value)
+
+
+def _flatten_bools(value) -> tuple[bool]:
+    if not isinstance(value, tuple):
+        return (bool(value),)
+    return sum((_flatten_bools(v) for v in value), start=())
+
+
 def _get_type_conversion_encoder(from_dtype: Type, to_dtype: Type):
 
     def kind(t):
@@ -401,20 +426,23 @@ def bitcast(self, value: bc.Value, fromty: Type, toty: Type) -> bc.Value:
             value = bc.encode_BitcastOp(self.builder, typeid(self.type_table, toty), value)
         return value
 
-    def constant(self, value: int | float, ty: Type) -> bc.Value:
-        if isinstance(ty, TileTy):
-            dtype = ty.dtype
-        else:
-            raise TypeError(f"Cannot make a constant tuple out of {ty}")
+    def constant(self, value, ty: Type) -> bc.Value:
+        if not isinstance(ty, TileTy):
+            raise TypeError(f"Cannot encode a constant of type {ty}; expected a TileTy")
 
-        data = _constant_to_bytes(value, dtype)
-        return bc.encode_ConstantOp(self.builder, typeid(self.type_table, ty), data)
+        def get_numel(v):
+            if not isinstance(v, tuple):
+                return 1
+            return sum(get_numel(i) for i in v)
 
-    def constant_tuple(self, value, ty: Type) -> Tuple[bc.Value, ...]:
-        if isinstance(ty, TupleTy):
-            return sum((self.constant_tuple(item_val, item_ty)
-                        for item_ty, item_val in zip(ty.value_types, value, strict=True)), ())
-        return self.constant(value, ty),
+        if get_numel(value) == 1:
+            while isinstance(value, tuple):
+                value = value[0]
+            data = _constant_to_bytes(value, ty.dtype)
+        else:
+            assert isinstance(value, tuple)
+            data = _constant_tuple_to_bytes(value, ty.dtype, ty.shape)
+        return bc.encode_ConstantOp(self.builder, typeid(self.type_table, ty), data)
 
     def index_tuple(self,
                     index: tuple[Var, ...], *, keep_i64: bool = False) -> Tuple[bc.Value, ...]: