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 struct

StableAddrEngine

struct StableAddrEngine[*, element_width: Int = Int(1)]

Implements TensorEngine over a device graph stable address slot.

StableAddrEngine is the slot-backed analogue of DevicePointerEngine. Its StorageType handle is a StableStorage: the address of a slot plus an element offset. The slot is graph-lifetime device memory holding one device data pointer; a recorded graph bakes in the slot's address and DeviceGraphBuilder.place_stable() rewrites its contents before each replay. Because the handle stores the slot rather than the pointer, the indirection is preserved across every TileTensor view derived from it.

Operations that dereference storage (load, store, copy_from, unsafe_ptr) read the slot and are device only: they are compile errors on the host. Operations that only manipulate the handle (offset, distance, unsafe_cast) work on both host and device, so a tile over a slot can be constructed, tiled, and sliced on the host and handed to a kernel.

The engine does not implement TensorOps; the elementwise arithmetic on TileTensor is unavailable for slot-backed tiles until it does.

Parameters​

  • ​element_width (Int): Number of scalar elements per logical element. A value of 1 (the default) is a non-vectorized tensor; larger values describe a vectorized view whose logical elements are SIMD vectors. The handle's offset is always in scalar-element units, so element_width only sets element_size (and thus the tile's vectorized ElementType), exactly as for DefaultEngine.

Implemented traits​

AnyType, Deinitable, Movable, TensorEngine

comptime members​

element_size​

comptime element_size = element_width

Number of scalar elements per logical element (alias of element_width).

OffsetResultType​

comptime OffsetResultType[offset_types: TypeList[values]] = StableAddrEngine[element_width=element_width]

The storage type produced by offsetting with a given coordinate.

Offsetting never changes the engine, so this is Self.

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] = StableStorage[dtype, origin, address_space]

A StableStorage handle: the slot plus a scalar-element offset.

Parameters​

  • ​mut (Bool): The mutability of the viewed storage, inferred from origin.
  • ​dtype (DType): The element data type of the viewed storage.
  • ​origin (Origin[mut=mut]): The origin the view is attributed to.
  • ​address_space (AddressSpace): The address space of the viewed storage.

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.

unsafe_cast​

static def unsafe_cast[to_mut: Bool, //, to_dtype: DType, to_origin: Origin[mut=to_mut], to_address_space: AddressSpace](storage: StableStorage) -> StableStorage[to_dtype, to_origin, to_address_space]

Reinterprets a handle with new type parameters.

The slot and offset are untyped, so this rebuilds the handle under the new parameters without touching either. Note the offset stays in scalar elements of the old dtype; a dtype cast on a handle with a non-zero offset is the caller's responsibility, as for the other engines.

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:

Returns:

StableStorage[to_dtype, to_origin, to_address_space]: A handle over the same slot and offset with the new parameters.

load​

static def load[dtype: DType, //, width: SIMDLength, alignment: Int, invariant: Bool = False, non_temporal: Bool = False](storage: StableStorage[dtype]) -> SIMD[dtype, width]

Loads a SIMD value from the storage.

Device only: reads the slot to find the data pointer.

Parameters:

  • ​dtype (DType): The element data type of the storage.
  • ​width (SIMDLength): The number of elements to load.
  • ​alignment (Int): The alignment guarantee for the load.
  • ​invariant (Bool): If True, the compiler may assume the memory won't be modified during the kernel, enabling load hoisting and caching.
  • ​non_temporal (Bool): If True, indicates the data will not be reused soon, allowing the hardware to bypass caches (e.g., streaming loads).

Args:

Returns:

SIMD[dtype, width]: The loaded SIMD value.

static def load[dtype: DType, //, width: SIMDLength, alignment: Int, invariant: Bool = False, non_temporal: Bool = False](storage: StableStorage[dtype], offset: T) -> SIMD[dtype, width]

Loads a SIMD value at a scalar-element offset from the storage.

Device only: reads the slot to find the data pointer.

Parameters:

  • ​dtype (DType): The element data type of the storage.
  • ​width (SIMDLength): The number of elements to load.
  • ​alignment (Int): The alignment guarantee for the load.
  • ​invariant (Bool): If True, the compiler may assume the memory won't be modified during the kernel, enabling load hoisting and caching.
  • ​non_temporal (Bool): If True, indicates the data will not be reused soon, allowing the hardware to bypass caches (e.g., streaming loads).

Args:

  • ​storage (StableStorage[dtype]): The storage to load from.
  • ​offset (T): The scalar-element offset to load at.

Returns:

SIMD[dtype, width]: The loaded SIMD value.

store​

static def store[dtype: DType, alignment: Int, *, non_temporal: Bool = False](storage: StableStorage[dtype], value: SIMD[dtype])

Stores a SIMD value into the storage.

Device only: reads the slot to find the data pointer.

Parameters:

  • ​dtype (DType): The element data type of the storage.
  • ​alignment (Int): The alignment guarantee for the store.
  • ​non_temporal (Bool): If True, indicates the data will not be reused soon, allowing the hardware to bypass caches (e.g., streaming stores).

Args:

static def store[dtype: DType, alignment: Int, *, non_temporal: Bool = False](storage: StableStorage[dtype], offset: T, value: SIMD[dtype])

Stores a SIMD value at a scalar-element offset in the storage.

Device only: reads the slot to find the data pointer.

Parameters:

  • ​dtype (DType): The element data type of the storage.
  • ​alignment (Int): The alignment guarantee for the store.
  • ​non_temporal (Bool): If True, indicates the data will not be reused soon, allowing the hardware to bypass caches (e.g., streaming stores).

Args:

  • ​storage (StableStorage[dtype]): The storage to store into.
  • ​offset (T): The scalar-element offset to store at.
  • ​value (SIMD[dtype]): The SIMD value to store.

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: StableStorage[offset_dtype, offset_origin, offset_address_space], var offset_coord: Coord[offset_types]) -> StableStorage[offset_dtype, offset_origin, offset_address_space]

Returns a handle advanced by a number of scalar elements.

Only the handle's offset changes; the slot is untouched, so the result still resolves through the slot when it is eventually dereferenced. Works on host and device.

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:

Returns:

StableStorage[offset_dtype, offset_origin, offset_address_space]: A handle over the same slot, that many scalar elements further in.

distance​

static def distance[dtype: DType, address_space: AddressSpace, //](storage: StableStorage[dtype, address_space], other: StableStorage[dtype, address_space]) -> Int

Returns the scalar-element distance from other to storage.

Both handles must refer to the same slot, since the distance between two slots' contents is unknowable without reading them (and is not stable across replays anyway). Works on host and device.

Parameters:

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

Args:

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.

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[StableStorage[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.

Delegates to the shared _copy_from loop; destination accesses go through this engine's load/store, so they are device only.

  • 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:

View source

View source

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community