Skip to main content

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 to DType.int64.

Args:

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?