Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Computing Tensors

Computing Tensors explains how a valid Tensor Unit stream reaches the Vector or Contraction Engine. A stationary operand stays in a register file while the other operand streams through the engine. The Tensor Register File (TRF) stores contraction operands, and the Vector Register File (VRF) stores vector operands. This chapter explains the engine behavior that places operands and preserves stream dimensions.

Selecting a Compute Route

Kernel needSelectBehavior to read next
Elementwise operation or reductionVector EngineVector stream and reduction contracts
Contraction such as matmul or convolutionTensor Register File (TRF) plus Contraction EngineTRF layout and contraction mapping
Cross-slice redistribution or reductionSwitch Engine or Inter-Slice ReducerSlice movement and reducer contracts
Layout or precision changeFetch/Commit Adapter, Cast Engine, or Transpose EngineAdapter stage and output mapping contracts

Tensor Unit

The Tensor Unit is the on-chip compute pipeline. It reads tensor data from DM, transforms it through ten engines, and writes results back to DM.

Each tensor flows through the pipeline as a stream of packets, one packet per cycle. The engines consume and produce these streams, reshaping the per-cycle layout and the iteration order along the way. The Collect Engine normalizes incoming packets to 32-byte flits. Every downstream engine operates on these flits. The pipeline includes Contraction, Vector, Cast, Transpose, Commit Adapter, and Commit. See the linked engine pages for Contraction, Vector, Cast, Transpose, Commit Adapter, and Commit.

flowchart TB
    subgraph SRAM
        DM[(DM)] & TRF[(TRF)] & VRF[(VRF)]
    end

    subgraph TU[Tensor Unit]
        direction LR
        FE[Fetch] --> FA[Fetch Adapter] --> SW[Switching] --> CO[Collect] --> CE[Contraction] --> VE[Vector] --> CA[Cast] --> TR[Transpose] --> CMA[Commit Adapter] --> CM[Commit]
    end

    DM --> FE
    CM --> DM
    CO --> TRF --> CE
    CO --> VRF --> VE

    click FE "../moving-tensors/fetch-engine.html" "Fetch Engine"
    click FA "./fetch-adapter.html" "Fetch Adapter"
    click SW "./switch-engine.html" "Switch Engine"
    click CO "./collect-engine.html" "Collect Engine"
    click CE "./contraction-engine/index.html" "Contraction Engine"
    click VE "./vector-engine/index.html" "Vector Engine"
    click CA "./cast-engine.html" "Cast Engine"
    click TR "./transpose-engine.html" "Transpose Engine"
    click CMA "./commit-adapter.html" "Commit Adapter"
    click CM "../moving-tensors/commit-engine.html" "Commit Engine"
EngineFunctionKey Constraint
FetchLoad data from DM into the pipelinePacket must be 8-byte aligned. Slice is preserved unless axis lifting replaces a broadcast with an axis from Time
Fetch AdapterPer-element transforms after fetch (table lookup, cast)Optional. Identity if skipped
SwitchingMove data across slicesRing network, Slice can change
CollectNormalize packets to 32-byte flitsOutput = exactly one flit
ContractionEinsum: matmul, convolution, attentionOne operand resident in TRF. The other streams
VectorElementwise, binary, reduce operationsOnly i32/f32 input
CastPrecision lowering with batchingOutput = exactly one flit
TransposeReorder elements within a flitWithin-flit only
Commit AdapterPer-element transforms before commit (cast, ReLU, trim)Optional. Chained before .commit()
CommitWrite results back to DMFlit-aligned writes

Each tensor stream inside the Tensor Unit carries five dimensions, [Chip, Cluster, Slice, Time, Packet], that split into two groups. Chip, Cluster, and Slice are spatial dimensions: each slice runs its own pipeline instance, with slices grouped by cluster and clusters grouped by chip. Time and Packet describe the per-slice stream (see Spatial and Temporal Dimensions for the definitions). The engines above reshape Time / Packet along the pipeline. Most engines preserve the spatial dimensions. Switch moves data across slices. Fetch’s axis lifting assigns different DM read offsets along Chip, Cluster, or Slice, moving an axis out of Time. The Vector inter-slice reducer combines the 256 slices in a cluster.

The Contraction and Vector Engines each take one operand from the pipeline stream and the other operand from a dedicated per-slice register file. TRF (Tensor Register File) feeds the Contraction Engine, and VRF (Vector Register File) feeds the Vector Engine. The Collect Engine writes into TRF via .to_trf() and into VRF via .to_vrf(). For an end-to-end example using both files, see Quick Start.

Fetch reads from DM and Commit writes back to DM. Their detailed sequencer behavior is documented in Moving Tensors rather than here.

Execution Context

The scheduler treats each execution context as an independent stream of operations. The hardware exposes three:

  • Main drives the Tensor Unit pipeline for the kernel’s primary computation.
  • Sub drives a subset of the same pipeline, typically prefetching operands into TRF / VRF while main computes.
  • DMA drives the DMA Engine alone, external to the Tensor Unit.

The main context can drive every Tensor Unit engine. The sub context drops the Contraction Engine and a handful of other features. Everything else carries over from main.

Context ordering, overlap, resource conflicts, and memory-rule scheduling are defined in Schedule.