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

Lane Folder

The Lane Folder is the Contraction Engine’s final stage. It eliminates the Lane dimension by relocating its 8 values into either OutPacket (Interleaved) or OutTime (Sequential). No values are summed: the stage folds Lane into another axis rather than reducing it.

Interface

.contract_lane(mode) invokes the Lane Folder. The stage drains the upstream Time Reducer’s buffer through an 8-element-wide output bus one cycle at a time, and LaneMode selects what each cycle’s flit carries.

impl<
    'l,
    const T: Tu,
    D: ContractionCast + MaterializableScalar,
    Chip: M,
    Cluster: M,
    Slice: M,
    Lane: M,
    Time: M,
    Packet: M,
    B: Backend,
> ContractTimeTensor<'l, T, D, Chip, Cluster, Slice, Lane, Time, Packet, B>
{
    /// Folds the `Lane` dimension into the output stream.
    /// `LaneMode::Interleaved` relocates `Lane` into `OutPacket`;
    /// `LaneMode::Sequential` relocates `Lane` into `OutTime`.
    #[primitive(ContractTimeTensor::contract_lane)]
    pub fn contract_lane<OutTime: M, OutPacket: M>(
        self,
        mode: LaneMode,
    ) -> ContractTensor<'l, T, D, Chip, Cluster, Slice, OutTime, OutPacket, B> {
        verify_contract_lane(
            Lane::to_value(),
            Time::to_value(),
            Packet::to_value(),
            OutTime::to_value(),
            OutPacket::to_value(),
            self.pre_reduce_time,
            mode,
        );
        // Finalize the carried operands with ONE fused contraction onto this stage's input mapping
        // `[Chip, Cluster, Slice, Lane, Time, Packet]` (Packet/Time were never actually reduced by the
        // earlier stages, only relabeled to their post-stage extents). The Lane fold relayout
        // (`transpose(false)`) then runs on the result, exactly as it always has.
        //
        // `out` is rebuilt here from this stage's own type params, independently of the `pre_reduce`
        // stashed by `contract_outer`; `Backend::contraction` reduces `pre_reduce` onto `out` via
        // `pre_reduce.carve(out)`, the same mapping-algebra carve `reduce` uses elsewhere, which is the
        // authority on whether `out` is a valid restriction of `pre_reduce` -- NOT a manual re-check here.
        // A naive per-symbol `.axes()` comparison is unsound for that: a contracted symbol can split
        // across a spatial slot this fold never touches (e.g. `Cluster`, carrying a `K`-fragment that
        // survives to `out` unreduced) and the slots that actually get contracted (`Time`/`Packet`,
        // carrying the rest of `K`); `pre_reduce`'s canonical `K` term then legitimately has a different
        // shape (wider modulo) than `out`'s, even though `out` is a correct restriction of `pre_reduce`.
        let contraction = self.inner;
        let out = <m![{ Chip }, { Cluster }, { Slice }, { Lane }, { Time }, { Packet }]>::to_value();
        let reduced: Tensor<D, m![{ Chip }, { Cluster }, { Slice }, { Lane }, { Time }, { Packet }], B> =
            Tensor::from_inner(B::contraction(
                &contraction.lhs,
                &contraction.rhs,
                &contraction.lhs_map,
                &contraction.rhs_map,
                &contraction.pre_reduce,
                &out,
            ));
        ContractTensor::new(self.ctx, reduced.transpose(false))
    }
}

The minimal examples below take a ContractTimeTensor (the output of the upstream Time Reducer) and call only .contract_lane(...), so each example shows the Lane Folder in isolation. The input Packet carries the size that survived the Packet Reducer, one of {1, 2, 4, 8, 16, 32} elements per lane.

Interleaved

The Lane dimension folds into OutPacket: each cycle reads one column position across all 8 lanes (one value per lane, 8 values per flit), with Lane materialized as the innermost OutPacket.

OutTime   = [Time, Packet]
OutPacket = [Lane # 8]

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

/// Lane folds into OutPacket.
fn lane_interleaved<'l, const T: Tu>(
    // Input from upstream Time Reducer: Lane = m![N], Time = m![M], Packet = m![P].
    input: ContractTimeTensor<'l, T, f32, m![1], m![1 # 2], m![1 # 256], m![N], m![M], m![P]>,
    // Output: OutTime = m![M, P] = [Time, Packet], OutPacket = m![N] = [Lane].
) -> ContractTensor<'l, T, f32, m![1], m![1 # 2], m![1 # 256], m![M, P], m![N]> {
    input.contract_lane::<m![M, P], m![N]>(LaneMode::Interleaved)
}

let mut ctx = Context::acquire();

let a: CollectTensor<'_, _, bf16, m![1], m![1 # 2], m![1 # 256], m![M], m![P]> = CollectTensor::new(&mut ctx.main, Tensor::zero());
let b: TrfTensor<bf16, m![1], m![1 # 2], m![1 # 256], m![N], m![P]> = unsafe { TrfTensor::from_addr(TrfAddress::Full) };

let i: ContractTimeTensor<'_, _, f32, m![1], m![1 # 2], m![1 # 256], m![N], m![M], m![P]> = a 
    .contract_outer::<m![M], m![P], m![N], m![P], _>(&b)
    .contract_packet::<m![P]>()
    .contract_time::<m![M]>();

let _o = lane_interleaved(i);
}

Sequential

The Lane dimension folds into OutTime: each cycle reads 8 column positions from one lane’s Packet (8 values per flit), with Lane iterating across successive cycles. Since each cycle is 8 elements wide, Packet is first padded up to a multiple of 8 (the bus width), then split into PadPacket / 8 cycles per lane and PadPacket % 8 elements per cycle.

PadPacket = Packet # align_up(Packet::SIZE, 8)   (pad Packet up to the next multiple of 8)
OutTime   = [Time, Lane, PadPacket / 8]
OutPacket = [PadPacket % 8]

For Packet::SIZE < 32, [PadPacket / 8]::SIZE = ceil(Packet::SIZE / 8) is the number of cycles per packet (e.g., 1 cycle for Packet::SIZE = 4, 2 cycles for Packet::SIZE = 16).

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

/// Lane folds into OutTime.
fn lane_sequential<'l, const T: Tu>(
    // Input from upstream Time Reducer: Lane = m![N], Time = m![M], Packet = m![P].
    input: ContractTimeTensor<'l, T, f32, m![1], m![1 # 2], m![1 # 256], m![N], m![M], m![P]>,
    // Output: OutTime = m![M, N, P / 8] = [Time, Lane, Packet / 8], OutPacket = m![P % 8] = [Packet % 8].
) -> ContractTensor<'l, T, f32, m![1], m![1 # 2], m![1 # 256], m![M, N, P / 8], m![P % 8]> {
    input.contract_lane::<m![M, N, P / 8], m![P % 8]>(LaneMode::Sequential)
}

let mut ctx = Context::acquire();

let a: CollectTensor<'_, _, bf16, m![1], m![1 # 2], m![1 # 256], m![M], m![P]> = CollectTensor::new(&mut ctx.main, Tensor::zero());
let b: TrfTensor<bf16, m![1], m![1 # 2], m![1 # 256], m![N], m![P]> = unsafe { TrfTensor::from_addr(TrfAddress::Full) };

let i: ContractTimeTensor<'_, _, f32, m![1], m![1 # 2], m![1 # 256], m![N], m![M], m![P]> = a 
    .contract_outer::<m![M], m![P], m![N], m![P], _>(&b)
    .contract_packet::<m![P]>()
    .contract_time::<m![M]>();

let _o = lane_sequential(i);
}

Constraints

The Lane Folder has no constraints of its own. The LaneMode selected here determines the slot capacity bound that the upstream Time Reducer enforces (see Time Reducer Constraints).

Performance

In Interleaved mode, throughput drops by Lane::SIZE / 8 when Lane < 8 (inactive lanes leave bus positions empty).

In Sequential mode, when Packet::SIZE < 8 (e.g., Packet::SIZE = 4 after the Packet Reducer collapses half of an bf16 packet), each cycle carries exactly Packet::SIZE elements rather than the full 8-element bus width: the output is narrower per cycle, but there is no padding and no wasted bus slots.

Latency is negligible: the Lane Folder reshapes per-lane outputs and does not add cycles beyond the buffer’s drain time.