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_runtime::tma.

Implementations

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

pub fn stride(&self, _dim: usize) -> usize

Obtain the stride of input at dimension dim

pub fn shape(&self, _dim: usize) -> usize

Obtain the shape of input at dimension dim

pub fn coordinate(&self, _index: usize, _dim: usize) -> usize

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) -> usize

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) -> usize

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) -> usize

Returns the rank of the tensor.

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

Downcast 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_stride( scope: &Scope, expand: NativeExpand<TensorMap<T, K>>, dim: NativeExpand<usize>, ) -> NativeExpand<usize>

pub fn __expand_shape( scope: &Scope, expand: NativeExpand<TensorMap<T, K>>, dim: NativeExpand<usize>, ) -> NativeExpand<usize>

pub fn __expand_coordinate( scope: &Scope, expand: NativeExpand<TensorMap<T, K>>, index: NativeExpand<usize>, dim: NativeExpand<usize>, ) -> NativeExpand<usize>

pub fn __expand_len( scope: &Scope, expand: NativeExpand<TensorMap<T, K>>, ) -> NativeExpand<usize>

pub fn __expand_buffer_len( scope: &Scope, expand: NativeExpand<TensorMap<T, K>>, ) -> NativeExpand<usize>

pub fn __expand_rank( scope: &Scope, expand: NativeExpand<TensorMap<T, K>>, ) -> NativeExpand<usize>

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> Copy for TensorMap<E, K>

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

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

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

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

The runtime argument for the kernel.

type CompilationArg = ()

Compilation argument.

fn register<R: Runtime>( arg: Self::RuntimeArg<R>, launcher: &mut KernelLauncher<R>, ) -> Self::CompilationArg

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

Register a variable during compilation that fill the KernelBuilder.

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

fn vector_size(&self) -> VectorSize

fn __expand_vector_size(_scope: &Scope, this: &Self::ExpandType) -> VectorSize