IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo trait

TensorEngine

Defines how tile tensor operates on borrowed storage.

A conforming engine names a concrete, register-passable StorageType handle for storage that is owned elsewhere, and supplies the static operations that act on it: offset, distance, reinterpretation, and raw pointer access. Element loads and stores go through the raw pointer unsafe_ptr returns. An engine never owns the underlying memory; the handle's origin parameter tracks the lifetime and mutability of the borrowed storage.

Implemented traits​

AnyType

comptime members​

element_size​

comptime element_size = Int(1)

OffsetResultType​

comptime OffsetResultType[offset_types: TypeList[values]]

The storage type produced by offsetting with a given coordinate.

Parameters​

  • ​offset_types (TypeList[values]): The coordinate element types of the applied offset.

StorageType​

comptime StorageType[mut: Bool, //, dtype: DType, origin: Origin[mut=mut], address_space: AddressSpace]

The concrete, register-passable handle to the borrowed storage.

Every operation in this trait acts on values of this type. It is parameterized on the element dtype, the origin that tracks the lifetime and mutability of the borrowed storage, and the address_space the storage resides in, so a single conforming type describes a whole family of handles.

Parameters​

  • ​mut (Bool): The mutability of the borrowed storage, inferred from origin.
  • ​dtype (DType): The element data type of the borrowed storage.
  • ​origin (Origin[mut=mut]): The origin tracking the lifetime of the borrowed storage.
  • ​address_space (AddressSpace): The address space the borrowed storage resides in.

Required methods​

unsafe_cast​

static def unsafe_cast[to_mut: Bool, //, to_dtype: DType, to_origin: Origin[mut=to_mut], to_address_space: AddressSpace](storage: Self.StorageType[origin.mut, origin_of(origin), dtype, origin, address_space]) -> Self.StorageType[to_mut, origin_of(to_origin), to_dtype, to_origin, to_address_space]

Reinterprets a storage handle with new type parameters.

This performs an unchecked reinterpretation of the underlying reference; no conversion of the stored elements takes place. The caller is responsible for ensuring the new dtype, origin, and address_space are valid for the referenced storage.

Parameters:

  • ​to_mut (Bool): The mutability to reinterpret the storage as.
  • ​to_dtype (DType): The element data type to reinterpret the storage as.
  • ​to_origin (Origin[mut=to_mut]): The origin to reinterpret the storage as.
  • ​to_address_space (AddressSpace): The address space to reinterpret the storage as.

Args:

  • ​storage (_Self.StorageType[origin.mut, origin_of(origin), dtype, origin, address_space]): The storage to reinterpret.

Returns:

_Self.StorageType[to_mut, origin_of(to_origin), to_dtype, to_origin, to_address_space]: A handle referring to the same storage, viewed with the new type parameters.

offset​

static def offset[offset_mut: Bool, offset_types: TypeList[offset_types.values], //, offset_dtype: DType, offset_origin: Origin[mut=offset_mut], offset_address_space: AddressSpace](var storage: Self.StorageType[offset_mut, origin_of(offset_origin), offset_dtype, offset_origin, offset_address_space], var offset_coord: Coord[offset_types]) -> Self.OffsetResultType[offset_types.values, offset_types].StorageType[offset_mut, origin_of(offset_origin), offset_dtype, offset_origin, offset_address_space]

Returns a storage handle offset by a number of scalar elements.

Parameters:

  • ​offset_mut (Bool): The mutability of the storage, inferred from offset_origin.
  • ​offset_types (TypeList[offset_types.values]): The coordinate element types of offset_coord.
  • ​offset_dtype (DType): The element data type of the storage.
  • ​offset_origin (Origin[mut=offset_mut]): The origin tracking the lifetime of the storage.
  • ​offset_address_space (AddressSpace): The address space the storage resides in.

Args:

  • ​storage (_Self.StorageType[offset_mut, origin_of(offset_origin), offset_dtype, offset_origin, offset_address_space]): The storage to offset from.
  • ​offset_coord (Coord[offset_types]): A flat coordinate whose components sum to the number of scalar elements to advance the handle by.

Returns:

_Self.OffsetResultType[offset_types.values, offset_types].StorageType[offset_mut, origin_of(offset_origin), offset_dtype, offset_origin, offset_address_space]: A handle of the same type starting the given number of scalar elements into the referenced storage.

copy_from​

static def copy_from[SelfLayoutType: TensorLayout, self_origin: MutOrigin, self_address_space: AddressSpace, OtherLayoutType: TensorLayout, other_mut: Bool, other_origin: Origin[mut=other_mut], other_address_space: AddressSpace, //, dst_dtype: DType, src_dtype: DType, OtherEngine: TensorEngine](storage: Tuple[Self.StorageType[True, origin_of(self_origin), dst_dtype, self_origin, self_address_space], SelfLayoutType], other: Tuple[OtherEngine.StorageType[other_mut, origin_of(other_origin), src_dtype, other_origin, other_address_space], OtherLayoutType])

Copies the elements of other into storage, in place.

Performs an element-by-element copy from other into storage, respecting the layouts of both operands. Each logical element is loaded from other using its layout and stored into storage using its own layout, so the copy works correctly even when the two sides have different shapes or strides (as long as they agree on total element count). When both operands have fully static, row-major layouts and a scalar logical element, the copy widens to SIMD load + cast + SIMD store using the narrower of the two dtypes' native SIMD widths.

TileTensor.copy_from reaches this through the source engine's copy_to, whose default forwards here; the source is read through OtherEngine.unsafe_ptr.

  • Both operands must have statically known shapes with matching total element count.
  • Both operands must have the same logical element size.
  • Source and destination dtypes may differ; each logical element is cast to the destination dtype.

Parameters:

  • ​SelfLayoutType (TensorLayout): The layout type of the destination storage.
  • ​self_origin (MutOrigin): The origin of the destination storage.
  • ​self_address_space (AddressSpace): The address space of the destination storage.
  • ​OtherLayoutType (TensorLayout): The layout type of the source storage.
  • ​other_mut (Bool): The mutability of the source storage.
  • ​other_origin (Origin[mut=other_mut]): The origin of the source storage.
  • ​other_address_space (AddressSpace): The address space of the source storage.
  • ​dst_dtype (DType): The element data type of the destination storage.
  • ​src_dtype (DType): The element data type of the source storage.
  • ​OtherEngine (TensorEngine): The engine of the source. May differ from Self as long as the two engines are copy-compatible (same logical element size).

Args:

distance​

static def distance[dtype: DType, address_space: AddressSpace, //](storage: Self.StorageType[False, origin_of(origin), dtype, origin, address_space], other: Self.StorageType[False, origin_of(origin), dtype, origin, address_space]) -> Int

Returns the scalar-element distance from other to storage.

Parameters:

  • ​dtype (DType): The storages' DType.
  • ​address_space (AddressSpace): The storages' AddressSpace.

Args:

  • ​storage (_Self.StorageType[False, origin_of(origin), dtype, origin, address_space]): The storage to measure the distance to.
  • ​other (_Self.StorageType[False, origin_of(origin), dtype, origin, address_space]): The storage to measure the distance from.

Returns:

Int: The number of scalar elements separating the two handles. The value is positive when storage is ahead of other and negative when it precedes other.

Provided methods​

write_type_name_to​

static def write_type_name_to(mut writer: T)

Writes the engine type name representation to the writer.

Args:

  • ​writer (T): The Writer to output to.

copy_to​

static def copy_to[SelfLayoutType: TensorLayout, self_mut: Bool, self_origin: Origin[mut=self_mut], self_address_space: AddressSpace, OtherLayoutType: TensorLayout, other_origin: MutOrigin, other_address_space: AddressSpace, //, src_dtype: DType, dst_dtype: DType, OtherEngine: TensorEngine](storage: Tuple[Self.StorageType[self_mut, origin_of(self_origin), src_dtype, self_origin, self_address_space], SelfLayoutType], other: Tuple[OtherEngine.StorageType[True, origin_of(other_origin), dst_dtype, other_origin, other_address_space], OtherLayoutType])

Copies the elements of storage into other, in place.

This is the source side of a copy. TileTensor.copy_from calls it on the source tensor's engine, and the default forwards to OtherEngine.copy_from, so an engine only overrides it when reading its storage takes more than unsafe_ptr: an engine with no pointer representation runs its own load loop here and stores through the destination's pointer.

Parameters:

  • ​SelfLayoutType (TensorLayout): The layout type of the source storage.
  • ​self_mut (Bool): The mutability of the source storage.
  • ​self_origin (Origin[mut=self_mut]): The origin of the source storage.
  • ​self_address_space (AddressSpace): The address space of the source storage.
  • ​OtherLayoutType (TensorLayout): The layout type of the destination storage.
  • ​other_origin (MutOrigin): The origin of the destination storage.
  • ​other_address_space (AddressSpace): The address space of the destination storage.
  • ​src_dtype (DType): The element data type of the source storage.
  • ​dst_dtype (DType): The element data type of the destination storage.
  • ​OtherEngine (TensorEngine): The engine of the destination.

Args:

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community