Skip to main content

Mojo struct

NullableTileTensor

struct NullableTileTensor[mut: Bool, //, dtype: DType, LayoutType: TensorLayout, origin: Origin[mut=mut], *, address_space: AddressSpace = AddressSpace.GENERIC, linear_idx_type: DType = _get_index_type[LayoutType](address_space), element_size: Int = 1]

A TileTensor variant whose pointer may be absent (null).

NullableTileTensor carries the same layout metadata as TileTensor but explicitly represents its pointer as nullable.

Only layout-query methods are provided. To perform loads, stores, or other data operations, first check self.ptr and then call value() to obtain a regular TileTensor.

Parameters

  • mut (Bool): The inferred mutability of the underlying pointer.
  • dtype (DType): The data type of tensor elements.
  • LayoutType (TensorLayout): A type implementing TensorLayout that defines the tensor's shape and stride structure.
  • origin (Origin): The origin of the underlying pointer for lifetime tracking.
  • address_space (AddressSpace): Memory address space. Defaults to GENERIC.
  • linear_idx_type (DType): Integer type for memory indexing.
  • element_size (Int): The number of scalar elements per logical element.

Fields

  • ptr (Optional[NullableTileTensor[dtype, LayoutType, origin, address_space=address_space, linear_idx_type=linear_idx_type, element_size=element_size].PtrType]): Optional pointer to the tensor's underlying data storage. When None, represents a tensor with layout metadata but no backing memory (e.g. an output buffer that the callee should allocate).
  • layout (LayoutType): The layout instance defining shape and stride mappings.

Implemented traits

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDestructible, Movable, RegisterPassable

comptime members

all_dims_known

comptime all_dims_known = LayoutType.all_dims_known

True if both shape and stride are fully known at compile time.

ElementType

comptime ElementType = SIMD[dtype, element_size]

The SIMD type used for element access.

flat_rank

comptime flat_rank = TypeList[#kgen.variadic.reduce(#kgen.variadic.reduce(#kgen.variadic.reduce(#kgen.variadic.reduce(LayoutType._shape_types, base=, reducer=[PrevV: KGENParamList[CoordLike], VA: KGENParamList[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, VA[idx].ParamListType if VA[idx].is_tuple else VA[idx])), base=, reducer=[PrevV: KGENParamList[CoordLike], VA: KGENParamList[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, VA[idx].ParamListType if VA[idx].is_tuple else VA[idx])), base=, reducer=[PrevV: KGENParamList[CoordLike], VA: KGENParamList[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, VA[idx].ParamListType if VA[idx].is_tuple else VA[idx])), base=, reducer=[PrevV: KGENParamList[CoordLike], VA: KGENParamList[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, VA[idx].ParamListType if VA[idx].is_tuple else VA[idx]))].size

The flattened rank.

GenericType

comptime GenericType = NullableTileTensor[dtype, LayoutType, origin, linear_idx_type=linear_idx_type]

Type alias for this tensor with GENERIC address space.

Used by constructors that create tensors from Span, DeviceBuffer, or HostBuffer, which all produce GENERIC address space tensors.

is_compatible_with

comptime is_compatible_with[C: KGENParamList[CoordLike]] = False if (TypeList[LayoutType._shape_types].size != TypeList[C].size) else #kgen.variadic.reduce(C, base=True, reducer=[PrevV: KGENParamList[Bool], VA: KGENParamList[CoordLike], idx: __mlir_type.index] PrevV[0] and True if not C[idx].is_tuple.__bool__() else False if not LayoutType._shape_types[idx].is_tuple.__bool__() else False if (TypeList[LayoutType._shape_types[idx].ParamListType].size != TypeList[C[idx].ParamListType].size) else #kgen.variadic.reduce(C[idx].ParamListType, base=True, reducer=[PrevV: KGENParamList[Bool], VA: KGENParamList[CoordLike], idx: __mlir_type.index] PrevV[0] and True if not C[idx].ParamListType[idx].is_tuple.__bool__() else False if not LayoutType._shape_types[idx].ParamListType[idx].is_tuple.__bool__() else False if (TypeList[LayoutType._shape_types[idx].ParamListType[idx].ParamListType].size != TypeList[C[idx].ParamListType[idx].ParamListType].size) else #kgen.variadic.reduce(C[idx].ParamListType[idx].ParamListType, base=True, reducer=[PrevV: KGENParamList[Bool], VA: KGENParamList[CoordLike], idx: __mlir_type.index] PrevV[0] and True if not C[idx].ParamListType[idx].ParamListType[idx].is_tuple.__bool__() else False if not LayoutType._shape_types[idx].ParamListType[idx].ParamListType[idx].is_tuple.__bool__() else (TypeList[LayoutType._shape_types[idx].ParamListType[idx].ParamListType[idx].ParamListType].size == TypeList[C[idx].ParamListType[idx].ParamListType[idx].ParamListType].size))[0])[0])[0]

True if coordinate types C are structurally compatible with this tensor's layout shape.

A scalar coordinate element is always compatible. A tuple coordinate element requires the corresponding layout shape element to also be a tuple of the same length, checked recursively up to 4 levels of nesting.

Parameters

  • C (KGENParamList): The coordinate element types to check against.

is_row_major

comptime is_row_major = (Coord[#kgen.variadic.reduce(LayoutType._stride_types, base=, reducer=[PrevV: KGENParamList[CoordLike], VA: KGENParamList[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, ComptimeInt[1 if (VA[idx].static_value == #kgen.variadic.reduce(#kgen.variadic.reduce(LayoutType._shape_types[0].ParamListType if LayoutType._shape_types[0].is_tuple if (TypeList[LayoutType._shape_types].size == 1) else (TypeList[LayoutType._shape_types].size == 1) else LayoutType._shape_types, base=, reducer=[PrevV: KGENParamList[CoordLike], VA: KGENParamList[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(PrevV, VA[(add (mul idx, -1), len(VA), -1)])), base=, reducer=[PrevV: KGENParamList[CoordLike], VA: KGENParamList[CoordLike], idx: __mlir_type.index] #kgen.variadic.concat(ComptimeInt[1] if (idx == 0) else RuntimeInt[VA[(add idx, -1)].DTYPE if not VA[(add idx, -1)].is_static_value.__bool__() else PrevV[0].DTYPE] if not VA[(add idx, -1)].is_static_value.__bool__() if not VA[(add idx, -1)].is_static_value.__bool__() else not PrevV[0].is_static_value.__bool__() else ComptimeInt[(VA[(add idx, -1)].static_value * PrevV[0].static_value)], PrevV))[idx].static_value) else 0]))].static_product == 1 if (TypeList[LayoutType._shape_types].size == 0) else Coord[#kgen.variadic.tabulate(TypeList[LayoutType._shape_types], [idx: __mlir_type.index] ComptimeInt[1])].static_product)

True if the tensor has row-major (contiguous) strides.

PtrType

comptime PtrType = NonNullUnsafePointer[Scalar[dtype], origin, address_space=address_space]

The non-null pointer type for the underlying data storage.

rank

comptime rank = LayoutType.rank

The number of dimensions in the tensor's layout.

shape_known

comptime shape_known = LayoutType.shape_known

True if all shape dimensions are compile-time constants.

static_shape

comptime static_shape[i: Int] = LayoutType.static_shape[i]

Get the compile-time shape value for dimension i, or -1 if dynamic.

Parameters

  • i (Int): The dimension index.

static_stride

comptime static_stride[i: Int] = LayoutType.static_stride[i]

Get the compile-time stride value for dimension i, or -1 if dynamic.

Parameters

  • i (Int): The dimension index.

stride_known

comptime stride_known = LayoutType.stride_known

True if all stride dimensions are compile-time constants.

Methods

__init__

@implicit __init__(other: NullableTileTensor[other.dtype, other.LayoutType, other.origin, address_space=other.address_space, linear_idx_type=other.linear_idx_type, element_size=other.element_size]) -> NullableTileTensor[other.dtype, other.LayoutType, origin_of(other.origin), address_space=other.address_space, linear_idx_type=other.linear_idx_type, element_size=other.element_size]

Implicitly cast a mutable NullableTileTensor to immutable.

Args:

Returns:

NullableTileTensor

@implicit __init__(other: TileTensor[other.dtype, other.LayoutType, other.origin, address_space=other.address_space, linear_idx_type=other.linear_idx_type, element_size=other.element_size]) -> NullableTileTensor[other.dtype, other.LayoutType, other.origin, address_space=other.address_space, linear_idx_type=other.linear_idx_type, element_size=other.element_size]

Implicitly cast a TileTensor to a NullableTileTensor.

Args:

  • other (TileTensor): The TileTensor to cast from.

Returns:

NullableTileTensor

@implicit __init__(other: TileTensor[other.dtype, other.LayoutType, other.origin, address_space=other.address_space, linear_idx_type=other.linear_idx_type, element_size=other.element_size]) -> NullableTileTensor[other.dtype, other.LayoutType, origin_of(other.origin), address_space=other.address_space, linear_idx_type=other.linear_idx_type, element_size=other.element_size]

Implicitly cast a mutable TileTensor to an immutable NullableTileTensor.

Args:

  • other (TileTensor): The mutable TileTensor to cast from.

Returns:

NullableTileTensor

value

value(self) -> TileTensor[dtype, LayoutType, origin, address_space=address_space, linear_idx_type=linear_idx_type, element_size=element_size]

Returns a regular TileTensor with the underlying pointer.

The caller must ensure the underlying pointer is non-null before calling this method.

Returns:

TileTensor: A TileTensor backed by the stored pointer and layout.

dim

dim[i: Int](self) -> Scalar[linear_idx_type]

Returns the size of dimension i.

Parameters:

  • i (Int): The dimension index (compile-time constant).

Returns:

Scalar: The size of dimension i as a scalar.

dim[IndexType: Indexer](self, index: IndexType) -> Scalar[linear_idx_type]

Returns the size of the specified dimension.

Parameters:

  • IndexType (Indexer): The type of the index argument.

Args:

  • index (IndexType): The dimension index (runtime value).

Returns:

Scalar: The size of the specified dimension as a scalar.

to_layout_tensor

to_layout_tensor(self) -> LayoutTensor[dtype, Layout(coord_to_int_tuple[LayoutType._shape_types](), coord_to_int_tuple[LayoutType._stride_types]()), origin, address_space=address_space]

Return a LayoutTensor with the same shape, stride, and address space of this tensor. Currently it expects flat layouts.

This is a utility to help with porting LayoutTensor methods to this type.

Returns:

LayoutTensor: A LayoutTensor with the same shape, stride, and address space of this tensor.

Was this page helpful?