Struct TensorMap 

pub struct TensorMap<E: CubePrimitive, K: TensorMapKind> { /* private fields */ }

A CUDA CUtensorMap object. Represents a tensor encoded with a lot of metadata, and is an opaque packed object at runtime. Does not support retrieving any shapes or strides, nor does it give access to the pointer. So these need to be passed separately in an aliased Tensor if needed.

Also see [cubecl_common::tma].

Implementations

impl<T: CubePrimitive, K: TensorMapKind> TensorMap<T, K>

pub fn buffer(&self) -> Tensor<Line<T>>

Get a reference to the underlying buffer for the tensor map.

pub fn stride<C: Index>(&self, _dim: C) -> u32

Obtain the stride of input at dimension dim

pub fn shape<C: Index>(&self, _dim: C) -> u32

Obtain the shape of input at dimension dim

pub fn coordinate<I: Index, D: Index>(&self, _index: I, _dim: D) -> u32

Obtain the coordinate corresponding to the given index of the tensor at dimension dim.

A coordinate is a list of indices corresponding to the multi-dimensional position of an element in the tensor. The dim element in a coordinate is the position along the dim dimension of the tensor.

pub fn len(&self) -> u32

The number of vectorized elements in the tensor.

Warning

The length will be affected by the vectorization factor. To obtain the number of elements, you should multiply the length by the vectorization factor.

pub fn buffer_len(&self) -> u32

The length of the buffer representing the tensor in terms of vectorized elements.

Warning

The buffer length will be affected by the vectorization factor. To obtain the number of elements, you should multiply the length by the vectorization factor.

pub fn rank(&self) -> u32

Returns the rank of the tensor.

pub fn try_cast_unchecked<E: CubePrimitive>(&self) -> TensorMap<E, K>

Try to cast the tensormap to the given type and panic if the type isn’t the same.

This function should only be used to satisfy the Rust type system, when two generic types are supposed to be the same.

pub fn __expand_buffer( scope: &mut Scope, expand: ExpandElementTyped<TensorMap<T, K>>, ) -> ExpandElementTyped<Tensor<Line<T>>>

pub fn __expand_stride<C: Index>( scope: &mut Scope, expand: ExpandElementTyped<TensorMap<T, K>>, dim: ExpandElementTyped<u32>, ) -> ExpandElementTyped<u32>

pub fn __expand_shape<C: Index>( scope: &mut Scope, expand: ExpandElementTyped<TensorMap<T, K>>, dim: ExpandElementTyped<u32>, ) -> ExpandElementTyped<u32>

pub fn __expand_coordinate<I: Index, D: Index>( scope: &mut Scope, expand: ExpandElementTyped<TensorMap<T, K>>, index: ExpandElementTyped<u32>, dim: ExpandElementTyped<u32>, ) -> ExpandElementTyped<u32>

pub fn __expand_len<C: Index>( scope: &mut Scope, expand: ExpandElementTyped<TensorMap<T, K>>, ) -> ExpandElementTyped<u32>

pub fn __expand_buffer_len<C: Index>( scope: &mut Scope, expand: ExpandElementTyped<TensorMap<T, K>>, ) -> ExpandElementTyped<u32>

pub fn __expand_rank<C: Index>( scope: &mut Scope, expand: ExpandElementTyped<TensorMap<T, K>>, ) -> ExpandElementTyped<u32>

Trait Implementations

impl<E: Clone + CubePrimitive, K: Clone + TensorMapKind> Clone for TensorMap<E, K>

fn clone(&self) -> TensorMap<E, K>

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<E: CubePrimitive, K: TensorMapKind> CubeType for TensorMap<E, K>

type ExpandType = ExpandElementTyped<TensorMap<E, K>>

fn into_mut(scope: &mut Scope, expand: Self::ExpandType) -> Self::ExpandType

Wrapper around the init method, necessary to type inference.

impl<E: CubePrimitive, K: TensorMapKind> ExpandElementIntoMut for TensorMap<E, K>

fn elem_into_mut(_scope: &mut Scope, elem: ExpandElement) -> ExpandElement

impl<E: CubePrimitive, K: TensorMapKind> LaunchArg for TensorMap<E, K>

type RuntimeArg<'a, R: Runtime> = TensorMapArg<'a, R, K>

The runtime argument for the kernel.

type CompilationArg = TensorMapCompilationArg

Compilation argument.

fn compilation_arg<R: Runtime>( _runtime_arg: &Self::RuntimeArg<'_, R>, ) -> Self::CompilationArg

fn expand( _arg: &Self::CompilationArg, builder: &mut KernelBuilder, ) -> ExpandElementTyped<TensorMap<E, K>>

Register an input variable during compilation that fill the KernelBuilder.

fn expand_output( _arg: &Self::CompilationArg, builder: &mut KernelBuilder, ) -> ExpandElementTyped<TensorMap<E, K>>

Register an output variable during compilation that fill the KernelBuilder.

impl<E: CubePrimitive, K: TensorMapKind> Lined for TensorMap<E, K>

fn line_size(&self) -> u32

fn __expand_line_size(_scope: &mut Scope, this: Self::ExpandType) -> u32

impl<E: CubePrimitive, K: TensorMapKind> Copy for TensorMap<E, K>