Register Files

The Collect Engine streams into the Contraction Engine and the Vector Engine. Both engines also take input from a register file (one per slice): the Tensor Register File (TRF) feeds the Contraction Engine, and the Vector Register File (VRF) feeds the Vector Engine. These register files must be populated before their consumer engine runs.

Tensor Register File

Interface

A TrfTensor is a tensor stored in the TRF:

/// Tensor stored in the tensor register file.
#[primitive(TrfTensor)]
#[derive(Debug)]
pub struct TrfTensor<D: Scalar, Chip: M, Cluster: M, Slice: M, Lane: M, Element: M, B: Backend = CurrentBackend> {
    pub(crate) inner: Tensor<D, Pair<Chip, Pair<Cluster, Pair<Slice, Pair<Lane, Element>>>>, B>,
    _marker: PhantomData<(D, Chip, Cluster, Slice, Lane, Element)>,
}

Chip / Cluster / Slice pass through from the source. Lane indexes the spatial parallelism (1, 2, 4, or 8 active lanes). Element holds the per-lane layout.

From Collect Engine

.to_trf::<Lane, Element>() on CollectTensor produces a TrfTensor. Where in the register file it lands is the compiler’s to decide, so the call names no region:

impl<'l, const T: Tu, P: CanApplyToTrf, D: Scalar, Chip: M, Cluster: M, Slice: M, Time: M, Packet: M, B: Backend>
    TuTensor<'l, T, P, D, Chip, Cluster, Slice, Time, Packet, B>
{
    /// Stores to the tensor register file.
    #[primitive(TuTensor::to_trf)]
    pub fn to_trf<Lane: M, Element: M>(self) -> TrfTensor<D, Chip, Cluster, Slice, Lane, Element, B> {
        verify_to_trf::<D, Lane, Time, Packet, Element>();
        TrfTensor::new(self.inner.transpose(false))
    }
}

.to_trf reshapes the streaming Time / Packet into Lane / Element:

Lane    = Time / FlitsPerLane
Element = [Time % FlitsPerLane, Packet]

for some FlitsPerLane that the compiler derives from Lane and Time, so each lane is filled by FlitsPerLane consecutive flits.

The following simple case occupies the full TRF:

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![B = 32];

fn load_trf<'l, const T: Tu>(
    input: CollectTensor<'l, T, i8, m![1], m![1 # 2], m![1 # 256], m![1], m![B]>,
) -> TrfTensor<i8, m![1], m![1 # 2], m![1 # 256], m![1], m![B]> {
    input.to_trf()
}

let mut ctx = Context::acquire();

let c: CollectTensor<'_, _, i8, m![1], m![1 # 2], m![1 # 256], m![1], m![B]> = CollectTensor::new(&mut ctx.main, Tensor::zero());
let _o = load_trf(c);
}

For example, in a matmul kernel Lane holds output channels and Element holds the contracted axis.

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![V = 32, M = 32, N = 8, K = 32];

type Chip    = m![1];
type Cluster = m![V / 16];
type Slice   = m![V % 16 # 256];
type Lane    = m![N];

/// Stores matmul weights into TRF for consumption by `bmatmul` in
/// [Contraction Engine: Example: Batched MatMul](./contraction-engine/index.md#example-batched-matmul).
fn store_bmatmul_trf<'l, const T: Tu>(
    input: CollectTensor<'l, T, bf16, Chip, Cluster, Slice, m![N, K / 16], m![K % 16]>,
) -> TrfTensor<bf16, Chip, Cluster, Slice, Lane, m![K]> {
    input.to_trf()
}

let mut ctx = Context::acquire();

let c: CollectTensor<'_, _, bf16, Chip, Cluster, Slice, m![N, K / 16], m![K % 16]> = CollectTensor::new(&mut ctx.main, Tensor::zero());
let _o = store_bmatmul_trf(c);
}

From the Vector Engine

A store does not have to happen off the collect. A main-context stream can run the Vector Engine first and store what the pass computed:

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![B = 64];

fn store_pass(ctx: &mut Context) -> VrfTensor<f32, m![1], m![1 # 2], m![1 # 256], m![B]> {
    let dm: DmTensor<f32, m![1], m![1 # 2], m![1 # 256], m![B]> = DmTensor::new();
    ctx.main
        .begin(dm.view())
        .fetch::<m![1], m![B]>()
        .collect::<m![B / 8], m![B % 8]>()
        .vector_init()
        .vector_intra_slice_tag(TagMode::Zero)
        .vector_narrow_split::<m![B / 8, B % 8 / 4 % 2], m![B % 4]>()
        .vector_fp_unary(FpUnaryOp::Sqrt)
        .vector_widen_concat::<m![B / 8], m![B % 8]>()
        .vector_final()
        .to_vrf(&mut ctx.sub)
}

let mut ctx = Context::acquire();
let _o = store_pass(&mut ctx);
}

The two positions are different data paths, not two spellings of one. Off the collect, the Fetch Network hands the flits to the register file and no engine runs in between. Off a vector_final, the pass ends at the Vector Engine’s own write port, so the value that lands in the register is the one the pass computed: Element2 flattens the post-pass Time / Packet, which a reduce or a widen may have reshaped along the way.

That is the only way a computed value reaches the VRF, since the DM-direct store and the collect-position store both read Data Memory. Nothing after the Vector Engine can run in the same command, so a store cannot be combined with a cast, a transpose, or a commit.

Which context a store occupies

A store issued from the main context runs its write through the sub context, whichever of the two positions above it takes, so it occupies both and takes &mut ctx.sub as well. ctx.main and ctx.sub are separate fields, so a kernel (which holds &mut Context) borrows both at once, as the example above does.

The borrow lasts the statement, which is the extent of the occupancy: the sub context is free again once the command retires. How long the register stays occupied is a different question with a different answer (until the operand’s last read), and it is the compiler’s allocation to make, which is why the returned VrfTensor carries no borrow.

A store issued from the sub context needs no argument, since it already holds that context.

From Data Memory

For completely contiguous input access (no gaps or reordering), TRF supports a short command (StoTRF), a compact hardware instruction that loads data from Data Memory directly into the TRF, bypassing the full Fetch → Switch → Collect → to_trf() pipeline. The shortcut trades arbitrary-layout support for lower setup overhead.

To Contraction Engine

Each read covers 8 lanes × (1 or 2) banks × 1 row × 320 bits per bank, producing 320 bits per lane per cycle for a narrow read (one bank) or 640 bits per lane per cycle for a wide read (both banks): all active lanes access one or both banks of the same row in parallel (one bank for narrow reads, both for wide reads). Per slice that totals 320 bytes/cycle (narrow) or 640 bytes/cycle (wide) across all 8 lanes. See TRF Sequencer for how the sequencer iterates these reads across rows and broadcasts.

Architecture

The TRF is a banked SRAM with the structure 8 lanes × 2 banks × 128 rows × 320 bits = 80 KB per slice. The 8 lanes operate in parallel, with 1, 2, 4, or 8 active per access.

How many elements pack into a single 320-bit row depends on the data type:

Only integers promote on store: i4 → i5 (5 bits) and i8 → i9 (9 bits) leave room for the fetch adapter’s optional zero-point subtraction, which widens an integer weight by one bit. i8 / f8 and bf16 stay at their native widths (8 and 16 bits respectively); the 320-bit row holds extra slack relative to a flat 8- or 16-bit packing so the same physical row width serves all types.

With fewer than 8 active lanes, each active lane sees more rows (as if the row count grew). Halving the active count doubles the rows per active lane (e.g., 4 active → 256 rows per bank, 1 active → 1024).

Double Buffering

The TRF enables double-buffering by splitting each bank into two halves: the TRF Sequencer loads from one half while a store fills the other, and the two can be flipped between iterations. Three address modes name the region, and the scheduler picks one at store time: Full uses all 128 rows per bank, FirstHalf uses rows 0–63, and SecondHalf uses rows 64–127. The half modes cap per-slice capacity at 40 KB.

See Scheduling and Tuning: Double Buffering for the tuning pattern that uses these halves across the main and sub contexts.

Both halves share the same banks, so reads and writes contend at the bank level even though they target different rows. When both target the same bank in a cycle, the read takes priority because the contraction pipeline needs the data this cycle while the store can wait.

The TRF mitigates this contention with a read cache and bank alternation.

TRF reads have heavy reuse: the same data is typically broadcast across many cycles, so a direct-mapped read cache (8 lanes × 2 banks × 4 rows × 320 bits = 2.5 KB) sits in front of the banks and absorbs repeated reads. The cache also relieves contention with the concurrent store. On a hit the read skips the bank, so the store can use it that cycle. On a miss the cache refills from the bank, occupying it for that cycle.

For narrow reads (≤ 32 bytes), bank alternation adds a second mitigation. Reads use only one bank, so they can alternate at 32-byte granularity across the two banks. Reads and writes then end up on different banks on successive cycles, avoiding contention even on cache misses. Wide reads (64 bytes) occupy both banks every cycle, so each cache miss blocks the concurrent store; narrow reads preserve half-bandwidth alternation even when the cache misses.

Vector Register File

The VRF is written either from the Collect Engine or directly from Data Memory, and read by the Vector Engine.

Interface

A VrfTensor is a tensor stored in the VRF:

/// Tensor stored in the vector register file (VRF).
#[primitive(VrfTensor)]
#[derive(Debug, Clone)]
pub struct VrfTensor<D: VeScalar, Chip: M, Cluster: M, Slice: M, Element: M, B: Backend = CurrentBackend> {
    pub(crate) inner: Tensor<D, Pair<Chip, Pair<Cluster, Pair<Slice, Element>>>, B>,
    _marker: PhantomData<(D, Chip, Cluster, Slice, Element)>,
}

Chip / Cluster / Slice pass through from the source. Element holds the per-(slice) layout.

Those three are what a Vector Engine op matches an operand against, so an operand carried under a different partition than the stream it feeds is a compile error rather than a read of another slice’s data (see Operands). .reshape::<Chip2, Cluster2, Slice2, Element2>() restates one under another partition of the same slices, which is how an operand replicated under an anonymous broadcast (m![256]) feeds a stream partitioned by a named axis. It relabels the mapping and moves no data, so it is unsafe for the same reason DmTensor::reshape is: every physical position must already hold what the new mapping claims of it. Regrouping axes and naming a broadcast distribution axis are the two forms that satisfy that; the # Safety section on VrfTensor::reshape states both.

From Collect Engine

.to_vrf::<Element2>() on CollectTensor stores the flits into the VRF and produces a VrfTensor. Where in the register file they land is the compiler’s to decide, so the call names no address:

impl<'l, P: CanApplyToVrf, D: VeScalar, Chip: M, Cluster: M, Slice: M, Time: M, Packet: M, B: Backend>
    TuTensor<'l, { Tu::Sub }, P, D, Chip, Cluster, Slice, Time, Packet, B>
{
    /// Stores to the vector register file.
    #[primitive(TuTensor::to_vrf)]
    pub fn to_vrf<Element: M>(self) -> VrfTensor<D, Chip, Cluster, Slice, Element, B> {
        verify_to_vrf::<D, Time, Packet, Element>();
        VrfTensor::new(self.inner.transpose(false))
    }
}

impl<'l, P: CanApplyToVrf, D: VeScalar, Chip: M, Cluster: M, Slice: M, Time: M, Packet: M, B: Backend>
    TuTensor<'l, { Tu::Main }, P, D, Chip, Cluster, Slice, Time, Packet, B>
{
    /// Stores to the vector register file, occupying the sub context the write runs through.
    ///
    /// The borrow lasts the statement. How long the *register* stays occupied is the compiler's
    /// allocation to make, so the returned [`VrfTensor`] keeps no borrow.
    #[primitive(TuTensor::to_vrf_from_main)]
    pub fn to_vrf<Element: M>(
        self,
        _sub: &mut TuContext<{ Tu::Sub }>,
    ) -> VrfTensor<D, Chip, Cluster, Slice, Element, B> {
        verify_to_vrf::<D, Time, Packet, Element>();
        VrfTensor::new(self.inner.transpose(false))
    }
}

.to_vrf flattens the streaming Time / Packet into Element2:

Element2 = [Time, Packet]

The user picks Element2.

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![B = 64];

fn store_vrf<'l>(
    input: CollectTensor<'l, { Tu::Sub }, i32, m![1], m![1 # 2], m![1 # 256], m![B / 8], m![B % 8]>,
) -> VrfTensor<i32, m![1], m![1 # 2], m![1 # 256], m![B]> {
    input.to_vrf()
}

let mut ctx = Context::acquire();

let c: CollectTensor<'_, { Tu::Sub }, i32, m![1], m![1 # 2], m![1 # 256], m![B / 8], m![B % 8]> = CollectTensor::new(&mut ctx.sub, Tensor::zero());
let _o = store_vrf(c);
}

The store has one signature per context, and the example above takes the sub context’s. A store issued from the main context takes ctx.sub as an argument, because the write registers live in the sub context’s register map: see Which context a store occupies.

From Data Memory

For completely contiguous input access (no gaps or reordering), VRF supports a short command (StoVRF), a compact hardware instruction that loads data from Data Memory directly into the VRF, bypassing the full Fetch → Switch → Collect → to_vrf() pipeline. The shortcut trades arbitrary-layout support for lower setup overhead.

Architecture