Struct TensorCore 

pub struct TensorCore {
    pub dims: (usize, usize, usize),
    pub threads: usize,
    pub elements_per_thread: (usize, usize, usize),
    pub dtype_in: DType,
    pub dtype_out: DType,
    pub opts: SmallVec<[TcOpt; 8]>,
    pub swizzle: ((SmallVec<[SwizzleAxis; 8]>, SmallVec<[SwizzleAxis; 8]>, SmallVec<[SwizzleAxis; 8]>), (SmallVec<[SwizzleAxis; 8]>, SmallVec<[SwizzleAxis; 8]>, SmallVec<[SwizzleAxis; 8]>)),
    pub pack_a: bool,
    pub tile_grid: (usize, usize),
}

Tensor core configuration for hardware-accelerated matrix multiplication.

Describes a specific matrix multiplication unit with fixed dimensions and data types. Based on NVIDIA’s WMMA (Warp Matrix Multiply-Accumulate) API and similar accelerators.

Matrix Dimensions

Tensor cores perform: C[M,N] += A[M,K] × B[K,N]

• dims.0 (N): Number of output columns
• dims.1 (M): Number of output rows
• dims.2 (K): Reduction dimension size

Example

NVIDIA Tensor Core 16x16x16:

• Processes 16×16 output tile
• Accumulates across 16 K elements
• Uses 32 threads (warp size)
• Each thread handles multiple elements via opts

Fields

dims: (usize, usize, usize)

Matrix dimensions (N, M, K).

threads: usize

Number of threads required (typically warp size: 32 for CUDA, 64 for AMD).

elements_per_thread: (usize, usize, usize)

Elements per thread in each dimension (N, M, K).

Describes how the matrix is distributed across threads. Example: (2, 2, 4) means each thread handles 2×2 output elements and processes 4 K elements.

dtype_in: DType

Input matrix data type (A and B matrices).

dtype_out: DType

Output/accumulator data type (C matrix).

opts: SmallVec<[TcOpt; 8]>

Optimization sequence for tensor core application.

A sequence of operations to transform ranges. Each operation splits a dimension (N, M, or K) and assigns it to a new axis type.

Example: [Upcast(0), Local(0), Local(0), Local(1), Local(1), Local(1), Upcast(1)]

• Upcast N once
• Local split N twice
• Local split M three times
• Upcast M once

Uses SmallVec to avoid heap allocation for typical tensor cores (≤8 ops).

swizzle: ((SmallVec<[SwizzleAxis; 8]>, SmallVec<[SwizzleAxis; 8]>, SmallVec<[SwizzleAxis; 8]>), (SmallVec<[SwizzleAxis; 8]>, SmallVec<[SwizzleAxis; 8]>, SmallVec<[SwizzleAxis; 8]>))

Swizzle patterns for input permutation.

Describes how to permute input matrices to match hardware layout. Format: ((A_local, A_upcast, A_reduce), (B_local, B_upcast, B_reduce))

Each tuple contains axis references that describe the permutation pattern for optimal memory access. The first tuple is for matrix A, second for B.

Uses SmallVec to avoid heap allocation for typical swizzles (≤8 axes per vec).

pack_a: bool

Pre-pack operand A into contiguous scratch buffer before the reduction loop. Beneficial when the A operand has non-unit stride access (e.g., AMX row-major matmul).

tile_grid: (usize, usize)

Tile grid for multi-FMA batching (tile_y_count, tile_x_count).

When > (1, 1), the codegen emits load-pair instructions and multiple FMAs per K iteration to compute a grid of output tiles simultaneously. Default is (1, 1) for single-tile operation.

Implementations

impl TensorCore

pub fn get_reduce_axes(&self) -> Vec<(usize, usize)>

Get the axes for reduction unrolling.

Returns pairs of (dimension_index, unroll_amount) for the K dimension. Used during TC application to unroll the reduction dimension.

pub fn upcast_axes(&self) -> (Vec<usize>, Vec<usize>, Vec<usize>)

Get the upcast axes configuration for WMMA construction.

Returns axes configuration for CONTRACT operations. Format: (A_axes, B_axes, output_axes)

pub fn sm75_tensor_cores() -> Vec<TensorCore>

Get all tensor cores for NVIDIA SM75 architecture (Turing).

pub fn sm80_tensor_cores(allow_tf32: bool) -> Vec<TensorCore>

Get all tensor cores for NVIDIA SM80 architecture (Ampere).

pub fn sm89_tensor_cores(allow_tf32: bool) -> Vec<TensorCore>

Get all tensor cores for NVIDIA SM89 architecture (Hopper).

pub fn rdna3_tensor_cores() -> Vec<TensorCore>

Get all tensor cores for AMD RDNA3 architecture (RX 7000 series).

pub fn rdna4_tensor_cores() -> Vec<TensorCore>

Get all tensor cores for AMD RDNA4 architecture.

pub fn cdna3_tensor_cores() -> Vec<TensorCore>

Get all tensor cores for AMD CDNA3 architecture (MI300).

pub fn cdna4_tensor_cores() -> Vec<TensorCore>

Get all tensor cores for AMD CDNA4 architecture.

pub fn metal_tensor_cores() -> Vec<TensorCore>

Get all tensor cores for Apple Metal (M1/M2/M3).

pub fn amx_tensor_cores() -> Vec<TensorCore>

Get all tensor cores for Apple AMX (M1/M2/M3 matrix accelerators).

pub fn intel_tensor_cores() -> Vec<TensorCore>

Get all tensor cores for Intel Xe architecture.

Trait Implementations

impl Clone for TensorCore

fn clone(&self) -> TensorCore

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for TensorCore

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.