Tuning

Tuning is a controlled experiment: preserve correctness, record a baseline, change one supported lever, and compare the same schedule metric before deciding whether to keep the change.

Control the experiment

Follow this order for every candidate:

Run the type-checking path to validate mapping and shape constraints.
Run a CPU test with a host oracle, including the relevant boundary cases.
Compile the baseline and record its schedule makespan.
Change one lever while keeping release, shapes, data types, and mappings fixed.
Run the same checks and dump the candidate schedule.
Keep the change only if correctness still passes and the measured metric improves; otherwise restore the baseline.

Schedule makespan is the primary static metric in this chapter. Do not report a throughput or cycle improvement without a reproducible schedule comparison or separately documented device evidence.

Choose one lever

Choose a lever only after diagnosis names the dependency or resource it targets:

change the execution-engine path when the current resource is the bottleneck;
change a tile or split shape when the schedule exposes avoidable serial work or an ill-fitting reduction;
change a mapping, padding, or transfer boundary when the limiting interval is movement or an address dependency.

The Moving Tensors, Computing Tensors, and Mapping Tensors chapters define these choices.

Unroll a loop

Use #[unroll] to expose every iteration of a small static loop to the scheduler.

#![feature(proc_macro_hygiene, register_tool, stmt_expr_attributes)]
#![register_tool(furiosa_opt)]
extern crate furiosa_opt_std;
extern crate tokio;
use furiosa_opt_std::prelude::*;
axes![Group = 4];
fn consume(_: usize) {}
#[device]
fn example(_: &mut Device) {
#[unroll]
for group in 0..Group::SIZE {
    consume(group);
}
}
#[tokio::main]
async fn main() -> Result<(), Error> {
    let mut device = Device::new(example.topology())?;
    launch(example, &mut device).await?;

    Ok(())
}

The kernel crate requires #![feature(proc_macro_hygiene, stmt_expr_attributes)]. The loop must start at zero, run at least once, and have a compile-time constant end. #[unroll] fully expands that trip count; there is no source-level partial-unroll factor or hard maximum trip count. In practice, generated IR and code grow with the trip count, so reserve it for small loops.

A runtime if inside the body is supported and is copied into every expanded iteration. A runtime loop end such as 0..valid_len is rejected before unrolling because the compiler has no fixed number of body copies to create. The loop may be in a device entry point or a helper called by one. Nested loops are not unrolled automatically. To unroll every level, add #[unroll] to each loop. A nested loop with its own constant end is supported. An end that depends on an outer loop counter, such as for j in 0..i, is a runtime loop end and is rejected even when the outer loop has #[unroll]; loop construction currently happens before unrolling.

The compiler checks for remaining loops after applying #[unroll]. Register-file reuse is enabled only when no loop remains anywhere in the kernel. Unrolling only some loops can still expose those iterations to scheduling, but a remaining outer or nested loop does not unlock loop-free optimizations.

Double buffering

Double buffering can overlap staging the next weight group on SubContext with contracting the current group on MainContext. The scheduler can consider this overlap only when both operations are visible in the same scheduling unit. Use a manual two-stage loop or #[unroll] to create that opportunity.

The comparison uses 20 weight groups and the same schedule window for every case.

Loop form	Groups visible per scheduling unit	Cross-iteration overlap	Makespan
Ordinary rolled loop	1	No	19,106 cycles
Manual two-stage software pipeline	2	Yes	15,477 cycles
Full unrolling	20	Yes	15,086 cycles

StoTrf stages a weight group in the tensor register file (TRF). The TRF can place one group in each half, but both halves share banks. See Register Files: Double Buffering for capacity and address-mode details. Use the Schedule Viewer to verify the resulting overlap.

Blue boxes mark StoTrf, orange boxes mark contraction, and green bands mark measured overlap.

Ordinary rolled loop: 19,106 cycles

The scheduler builds one schedule for the rolled loop body and repeats that schedule for every iteration. Therefore, Contraction(i) and StoTrf(i + 1) cannot overlap.

#![allow(incomplete_features)]
#![feature(adt_const_params, proc_macro_hygiene, register_tool)]
#![register_tool(furiosa_opt)]
extern crate furiosa_opt_std;
extern crate tokio;
use furiosa_opt_std::prelude::*;
axes![Tok = 16, Red = 64, Out = 8, Group = 20];
use furiosa_opt_std::prelude::*;

type Chip = m![1];
type Cluster = m![Tok / 8 % 2];
type Slice = m![Tok % 8 # 256];

/// Stages and contracts one weight group per rolled iteration.
#[device(chip = 1)]
pub fn rolled(
    device: &mut Device,
    activation: &HbmTensor<bf16, Chip, m![Tok, Red]>,
    weight: &HbmTensor<bf16, Chip, m![Group, Out, Red]>,
) -> HbmTensor<bf16, Chip, m![Tok, Group, Out]> {
    let activation: DmTensor<bf16, Chip, Cluster, Slice, m![Red]> = activation.to_dm(&mut device.tdma);
    let weight: DmTensor<bf16, Chip, Cluster, Slice, m![Group, Out, Red]> = weight.to_dm(&mut device.tdma);
    let mut output: DmTensor<bf16, Chip, Cluster, Slice, m![Group, Out]> = DmTensor::new();

    for g in 0..Group::SIZE {
        let weight_group = weight.view().tile::<m![Group], 1, m![1 # 20, Out, Red]>(g);
        let weight_trf: TrfTensor<bf16, Chip, Cluster, Slice, m![Out], m![Red]> = device
            .sub
            .begin(weight_group)
            .fetch::<m![Out, Red / 16], m![Red % 16]>()
            .collect::<m![Out, Red / 16], m![Red % 16]>()
            .to_trf();

        device
            .main
            .begin(activation.view())
            .fetch::<m![Red / 16], m![Red % 16]>()
            .collect::<m![Red / 16], m![Red % 16]>()
            .contract_outer::<m![Red / 32], m![Red % 32], _, _, _>(&weight_trf)
            .contract_packet::<m![1]>()
            .contract_time::<m![1]>()
            .contract_lane::<m![1], m![Out]>(LaneMode::Interleaved)
            .cast::<bf16, m![Out # 16]>()
            .commit_trim::<m![Out]>()
            .commit_view(output.view_mut().tile::<m![Group], 1, m![1 #{!} 20, Out]>(g));
    }

    let mut result = HbmTensor::<bf16, Chip, m![Tok, Group, Out]>::new();
    output.view().to_hbm_view(&mut device.tdma, result.view_mut());
    result
}
#[tokio::main]
async fn main() -> Result<(), Error> {
let mut device = Device::new(rolled.topology())?;
let activation = HbmTensor::<bf16, m![1], m![Tok, Red]>::new();
let weight = HbmTensor::<bf16, m![1], m![Group, Out, Red]>::new();
let _output = launch(rolled, (&mut device, &activation, &weight)).await?;

    Ok(())
}

The rolled loop schedules each StoTrf before its matching contraction without cross-iteration overlap.

Manual two-stage software pipeline: 15,477 cycles

The manual pipeline places two groups in one loop body. This makes Contraction(first) and StoTrf(second) available to the scheduler together.

#![allow(incomplete_features)]
#![feature(adt_const_params, proc_macro_hygiene, register_tool)]
#![register_tool(furiosa_opt)]
extern crate furiosa_opt_std;
extern crate tokio;
use furiosa_opt_std::prelude::*;
axes![Tok = 16, Red = 64, Out = 8, Group = 20, Pairs = 10];
use furiosa_opt_std::prelude::*;

type Chip = m![1];
type Cluster = m![Tok / 8 % 2];
type Slice = m![Tok % 8 # 256];

/// Exposes two weight groups per iteration for software pipelining.
#[device(chip = 1)]
pub fn software_pipelined(
    device: &mut Device,
    activation: &HbmTensor<bf16, Chip, m![Tok, Red]>,
    weight: &HbmTensor<bf16, Chip, m![Group, Out, Red]>,
) -> HbmTensor<bf16, Chip, m![Tok, Group, Out]> {
    let activation: DmTensor<bf16, Chip, Cluster, Slice, m![Red]> = activation.to_dm(&mut device.tdma);
    let weight: DmTensor<bf16, Chip, Cluster, Slice, m![Group, Out, Red]> = weight.to_dm(&mut device.tdma);
    let mut output: DmTensor<bf16, Chip, Cluster, Slice, m![Group, Out]> = DmTensor::new();

    for pair in 0..Pairs::SIZE {
        let first = pair * 2;
        let second = first + 1;
        let first_weight = weight.view().tile::<m![Group], 1, m![1 # 20, Out, Red]>(first);
        let first_trf: TrfTensor<bf16, Chip, Cluster, Slice, m![Out], m![Red]> = device
            .sub
            .begin(first_weight)
            .fetch::<m![Out, Red / 16], m![Red % 16]>()
            .collect::<m![Out, Red / 16], m![Red % 16]>()
            .to_trf();
        let second_weight = weight.view().tile::<m![Group], 1, m![1 # 20, Out, Red]>(second);
        let second_trf: TrfTensor<bf16, Chip, Cluster, Slice, m![Out], m![Red]> = device
            .sub
            .begin(second_weight)
            .fetch::<m![Out, Red / 16], m![Red % 16]>()
            .collect::<m![Out, Red / 16], m![Red % 16]>()
            .to_trf();

        device
            .main
            .begin(activation.view())
            .fetch::<m![Red / 16], m![Red % 16]>()
            .collect::<m![Red / 16], m![Red % 16]>()
            .contract_outer::<m![Red / 32], m![Red % 32], _, _, _>(&first_trf)
            .contract_packet::<m![1]>()
            .contract_time::<m![1]>()
            .contract_lane::<m![1], m![Out]>(LaneMode::Interleaved)
            .cast::<bf16, m![Out # 16]>()
            .commit_trim::<m![Out]>()
            .commit_view(output.view_mut().tile::<m![Group], 1, m![1 #{!} 20, Out]>(first));
        device
            .main
            .begin(activation.view())
            .fetch::<m![Red / 16], m![Red % 16]>()
            .collect::<m![Red / 16], m![Red % 16]>()
            .contract_outer::<m![Red / 32], m![Red % 32], _, _, _>(&second_trf)
            .contract_packet::<m![1]>()
            .contract_time::<m![1]>()
            .contract_lane::<m![1], m![Out]>(LaneMode::Interleaved)
            .cast::<bf16, m![Out # 16]>()
            .commit_trim::<m![Out]>()
            .commit_view(output.view_mut().tile::<m![Group], 1, m![1 #{!} 20, Out]>(second));
    }

    let mut result = HbmTensor::<bf16, Chip, m![Tok, Group, Out]>::new();
    output.view().to_hbm_view(&mut device.tdma, result.view_mut());
    result
}
#[tokio::main]
async fn main() -> Result<(), Error> {
let mut device = Device::new(software_pipelined.topology())?;
let activation = HbmTensor::<bf16, m![1], m![Tok, Red]>::new();
let weight = HbmTensor::<bf16, m![1], m![Group, Out, Red]>::new();
let _output = launch(software_pipelined, (&mut device, &activation, &weight)).await?;

    Ok(())
}

The software-pipelined loop overlaps each current contraction with the next StoTrf.

Full unrolling: 15,086 cycles

#[unroll] exposes all 20 groups without manually pairing them.

#![allow(incomplete_features)]
#![feature(adt_const_params, proc_macro_hygiene, register_tool)]
#![register_tool(furiosa_opt)]
extern crate furiosa_opt_std;
extern crate tokio;
use furiosa_opt_std::prelude::*;
axes![Tok = 16, Red = 64, Out = 8, Group = 20];
use furiosa_opt_std::prelude::*;

type Chip = m![1];
type Cluster = m![Tok / 8 % 2];
type Slice = m![Tok % 8 # 256];

/// Exposes every weight group by fully unrolling the group loop.
#[device(chip = 1)]
pub fn unrolled(
    device: &mut Device,
    activation: &HbmTensor<bf16, Chip, m![Tok, Red]>,
    weight: &HbmTensor<bf16, Chip, m![Group, Out, Red]>,
) -> HbmTensor<bf16, Chip, m![Tok, Group, Out]> {
    let activation: DmTensor<bf16, Chip, Cluster, Slice, m![Red]> = activation.to_dm(&mut device.tdma);
    let weight: DmTensor<bf16, Chip, Cluster, Slice, m![Group, Out, Red]> = weight.to_dm(&mut device.tdma);
    let mut output: DmTensor<bf16, Chip, Cluster, Slice, m![Group, Out]> = DmTensor::new();

    #[unroll]
    for g in 0..Group::SIZE {
        let weight_group = weight.view().tile::<m![Group], 1, m![1 # 20, Out, Red]>(g);
        let weight_trf: TrfTensor<bf16, Chip, Cluster, Slice, m![Out], m![Red]> = device
            .sub
            .begin(weight_group)
            .fetch::<m![Out, Red / 16], m![Red % 16]>()
            .collect::<m![Out, Red / 16], m![Red % 16]>()
            .to_trf();

        device
            .main
            .begin(activation.view())
            .fetch::<m![Red / 16], m![Red % 16]>()
            .collect::<m![Red / 16], m![Red % 16]>()
            .contract_outer::<m![Red / 32], m![Red % 32], _, _, _>(&weight_trf)
            .contract_packet::<m![1]>()
            .contract_time::<m![1]>()
            .contract_lane::<m![1], m![Out]>(LaneMode::Interleaved)
            .cast::<bf16, m![Out # 16]>()
            .commit_trim::<m![Out]>()
            .commit_view(output.view_mut().tile::<m![Group], 1, m![1 #{!} 20, Out]>(g));
    }

    let mut result = HbmTensor::<bf16, Chip, m![Tok, Group, Out]>::new();
    output.view().to_hbm_view(&mut device.tdma, result.view_mut());
    result
}
#[tokio::main]
async fn main() -> Result<(), Error> {
let mut device = Device::new(unrolled.topology())?;
let activation = HbmTensor::<bf16, m![1], m![Tok, Red]>::new();
let weight = HbmTensor::<bf16, m![1], m![Group, Out, Red]>::new();
let _output = launch(unrolled, (&mut device, &activation, &weight)).await?;

    Ok(())
}

The unrolled loop overlaps the first contraction with the next iteration’s StoTrf.

All captures show cycles 4,000 through 8,000.

Keyboard shortcuts

Programming Tensor Contraction Processors