Mojo function
lt_to_tt_idx
lt_to_tt_idx[dtype: DType, lt_layout: Layout, //, ResultLayout: TensorLayout = Layout[*?, *?], linear_idx_type: DType = DType.int64](lt: LayoutTensor[dtype, lt_layout, address_space=lt.address_space, element_layout=lt.element_layout, layout_int_type=lt.layout_int_type, linear_idx_type=lt.linear_idx_type, masked=lt.masked, alignment=lt.alignment]) -> TileTensor[dtype, Layout[ResultLayout._shape_types, ResultLayout._stride_types], lt.origin, address_space=lt.address_space, linear_idx_type=linear_idx_type]
Like lt_to_tt but with an explicit linear_idx_type override.
Use this variant when the default TileTensor index-type heuristic
(int32 for SHARED/CONSTANT or small static cosize, int64
otherwise) does not match what callers downstream want. The most
common reason is preserving int64 indexing for DRAM tiles derived
from a tensor with runtime dimensions: the parent LayoutTensor
uses int64 for its address arithmetic, but a tile coming out of
lt_to_tt (whose own layout is fully static post-tiling) defaults
to int32 -- forcing _distribute() to do narrow-then-widen index
arithmetic for every offset, which is measurable in tight inner
loops.
Parameters:
- dtype (
DType): Element type of the tensor. - lt_layout (
Layout): The legacy Layout of the LayoutTensor. - ResultLayout (
TensorLayout): The target TileTensor layout type. - linear_idx_type (
DType): Integer type used for the result's offset arithmetic. Defaults toDType.int64.
Args:
- lt (
LayoutTensor[dtype, lt_layout, address_space=lt.address_space, element_layout=lt.element_layout, layout_int_type=lt.layout_int_type, linear_idx_type=lt.linear_idx_type, masked=lt.masked, alignment=lt.alignment]): The LayoutTensor to convert.
Returns:
TileTensor[dtype, Layout[ResultLayout._shape_types, ResultLayout._stride_types], lt.origin, address_space=lt.address_space, linear_idx_type=linear_idx_type]: A TileTensor with the same data and the requested
linear_idx_type.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!