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

Mapping Expressions

A mapping expression is a Rust type that encodes a mapping. This page defines its constructors and equivalence rules.

Axis Sizes

The axes! macro declares axis identifiers and their sizes. The following declaration applies throughout this section:

#![allow(unused)]
fn main() {
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 8, B = 512];
}

Mapping Interface

A mapping expression like m![H, W] is a Rust type that assigns each tensor index to a buffer position. Every mapping expression implements the M trait, which provides the buffer size and a function from buffer positions to tensor indices:

#![allow(unused)]
fn main() {
// Inside `furiosa_opt_std::prelude`...
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
use std::fmt::Debug;
pub trait M: Debug + Clone {
    /// The computed size for the given shape.
    const SIZE: usize;

    /// Converts the mapping expression type into a value.
    fn to_value() -> Mapping;

    /// Converts a buffer index to a tensor index, returning `None` if out-of-bounds.
    fn map(i: usize) -> Option<Index>;
}

/// Tensor index: a map from axis identifiers to coordinate values.
pub struct Index { /* ... */ }

/// Constructs tensor indices.
/// `i![A: 2, B: 3]` creates an `Index` with A = 2 and B = 3.
macro_rules! i {
    () => {};
    /* ... */
}
}

Usage Example: Host Tensor

The simplest concrete type built on the M trait is HostTensor<D, E>: a host memory buffer of element type D whose layout is fully determined by mapping E. E determines both the buffer size (E::SIZE) and the correspondence from buffer positions to tensor indices (E::map). HostTensor<bf16, m![A, B]> contains 4,096 elements of bf16 data. A HostTensor<D, E> holds a tensor \(T\) when:

  • for every buffer index i and tensor index ti where E::map(i) = Some(ti),
  • the i-th element stores the value of tensor \(T\) at index ti.

Device tensors such as HbmTensor and DmTensor have more complex layouts spanning multiple mapping expressions; see Spatial and Temporal Dimensions for details.

Constructors

Mapping expressions, including the layout E in HostTensor<D, E>, are built by composing small constructors, each of which transforms or combines simpler mappings. These expressions use arithmetic-like operators (/, %, and # for padding) to concisely define the mapping between tensor and linear buffer indices.

Symbol

A symbol is a single uppercase letter whose size comes from the shape declaration. The mapping m![A] maps 8 buffer indices linearly to tensor indices along the axis:

#![allow(unused)]
fn main() {
extern crate furiosa_opt_std;
extern crate furiosa_mapping;
use furiosa_opt_std::prelude::*;

axes![A = 8];

type E = m![A]; // Symbol<Ident::A, 8>

fn test_symbol() {
    assert_eq!(E::map(0), Some(i![A: 0]));
    assert_eq!(E::map(1), Some(i![A: 1]));
    assert_eq!(E::map(2), Some(i![A: 2]));
    for i in 0..E::SIZE {
        assert_eq!(E::map(i), Some(i![A: i]));
    }
    assert_eq!(E::map(E::SIZE), None);
}

test_symbol();
}
impl<S: AxisName> M for Symbol<S> {
    const SIZE: usize = S::SIZE;

    fn to_value() -> Mapping {
        Mapping::Symbol {
            symbol: S::NAME,
            size: S::SIZE,
        }
    }

    fn map(i: usize) -> Option<Index> {
        (i < S::SIZE).then(|| {
            let mut index = Index::new();
            index.add_term(
                Term {
                    inner: Atom::Symbol {
                        symbol: S::NAME,
                        size: S::SIZE,
                    },
                    stride: 1,
                    modulo: S::SIZE,
                },
                i,
            );
            index
        })
    }
}

Note

For every symbol A, the zeroth index i![A: 0] is equivalent to the empty tensor index i![].

Pair

The pair mapping m![A, B] stores a 2D tensor with shape \(\{A=8, B=512\}\) as a buffer of 4,096 elements. The mapping Pair<L, R> maps the Cartesian product of two spaces into a linear buffer where L is the major dimension and R is the minor dimension. The size is L::SIZE * R::SIZE, and the mapping uses floor division and modulo to decompose indices. m![A, B, C, D] expands to Pair<A, Pair<B, Pair<C, D>>> and is right-associative.

#![allow(unused)]
fn main() {
extern crate furiosa_opt_std;
extern crate furiosa_mapping;
use furiosa_opt_std::prelude::*;

axes![A = 8, B = 512];

type E = m![A, B]; // Pair<m![A], m![B]>

fn test_pair() {
    // First 512 elements hold A=0, next 512 hold A=1
    assert_eq!(E::map(0),   Some(i![A: 0, B: 0]));
    assert_eq!(E::map(511), Some(i![A: 0, B: 511]));
    assert_eq!(E::map(512), Some(i![A: 1, B: 0]));
    assert_eq!(E::map(519), Some(i![A: 1, B: 7])); // 519 == 512 * 1 + 7
    for i in 0..E::SIZE {
        assert_eq!(E::map(i), Some(i![A: i / <m![B]>::SIZE, B: i % <m![B]>::SIZE]));
    }
    assert_eq!(E::map(E::SIZE), None);
}

test_pair();
}
impl<L, R> M for Pair<L, R>
where
    L: M,
    R: M,
{
    const SIZE: usize = L::SIZE * R::SIZE;

    fn to_value() -> Mapping {
        Mapping::Pair {
            left: RBox::new(L::to_value()),
            right: RBox::new(R::to_value()),
        }
    }

    fn map(i: usize) -> Option<Index> {
        let mut l = L::map(i / R::SIZE)?;
        let r = R::map(i % R::SIZE)?;
        l.add(r);
        Some(l)
    }
}

Identity

The identity mapping m![1] creates a single-element buffer that maps buffer index 0 to the empty tensor index i![]. It serves as the identity element for Pair: m![1, A] and m![A, 1] are both equivalent to m![A].

#![allow(unused)]
fn main() {
extern crate furiosa_opt_std;
extern crate furiosa_mapping;
use furiosa_opt_std::prelude::*;

type E = m![1]; // Identity

fn test_identity() {
    assert_eq!(E::map(0), Some(i![]));
    assert_eq!(E::map(1), None);
}

test_identity();
}
/// The identity mapping (size-1 broadcast), the unit written `m![1]`.
pub type Identity = Broadcast<1>;

Padding

Padding aligns data to hardware requirements by adding unused buffer space. For example, the DMA engine requires rows to start on 64-byte boundaries. With axes![C = 13, D = 61], m![C, D] creates misaligned rows since 61 is not divisible by 64. m![C, D # 64] fixes this by aligning each row to 64-byte boundaries, using 3 extra elements per row.

#![allow(unused)]
fn main() {
extern crate furiosa_opt_std;
extern crate furiosa_mapping;
use furiosa_opt_std::prelude::*;

axes![C = 13, D = 61];

type E = m![C, D # 64]; // Pair<m![C], Padding<m![D], 64>>

fn test_padding() {
    assert_eq!(E::map(0),  Some(i![C: 0, D: 0]));
    assert_eq!(E::map(60), Some(i![C: 0, D: 60]));
    assert_eq!(E::map(61), None); // padding
    assert_eq!(E::map(62), None); // padding
    assert_eq!(E::map(63), None); // padding
    assert_eq!(E::map(64), Some(i![C: 1, D: 0]));
}

test_padding();
}
impl<L, const SIZE: usize, const KIND: PaddingKind> M for Padding<L, SIZE, KIND>
where
    L: M,
{
    const SIZE: usize = SIZE;

    fn to_value() -> Mapping {
        Mapping::Padding {
            inner: RBox::new(L::to_value()),
            padding: SIZE,
            kind: KIND,
        }
    }

    fn map(i: usize) -> Option<Index> {
        L::map(i)
    }
}

The padded slots’ content is part of the type, not just their count. Three kinds are tracked.

  • m![A # m] (or m![A #{*} m]) is top padding to size m. Slots are accessible but hold arbitrary values. Raw DM tensors carry this. # is the shorthand; #{*} spells the kind out explicitly.
  • m![A #{0} m] is zero-filled padding to size m. Slots are accessible and known to hold zero. The Fetch Adapter’s masking stage produces this from # m.
  • m![A #{!} m] is bottom padding to size m. Slots are inaccessible and reads/writes are undefined behavior. This models addresses the compiler must avoid.

# defaults to top kind. The Rust type level mirrors this via a const generic of PaddingKind on Padding<L, SIZE, KIND>. Padding<L, N> is KIND = PaddingKind::Top, Padding<L, N, { PaddingKind::Zero }> is the zero-filled variant, and Padding<L, N, { PaddingKind::Bottom }> is inaccessible.

Resize

Resize constrains a mapping to a smaller logical size by truncating indices beyond the new size, discarding elements outside that range. Unlike padding, which expands the buffer, Resize shrinks the logical view. The mapping m![D = 2] takes only the first 2 elements of axis D, producing indices D = 0 and D = 1.

#![allow(unused)]
fn main() {
extern crate furiosa_opt_std;
extern crate furiosa_mapping;
use furiosa_opt_std::prelude::*;

axes![C = 2, D = 3];
type E = m![C, D = 2]; // Pair<m![C], Resize<m![D], 2>>

fn test_resize() {
    assert_eq!(E::map(0), Some(i![C: 0, D: 0]));
    assert_eq!(E::map(1), Some(i![C: 0, D: 1]));
    assert_eq!(E::map(2), Some(i![C: 1, D: 0]));
    assert_eq!(E::map(3), Some(i![C: 1, D: 1]));
    assert_eq!(E::map(4), None);
}

test_resize();
}
impl<L, const SIZE: usize> M for Resize<L, SIZE>
where
    L: M,
{
    const SIZE: usize = SIZE;

    fn to_value() -> Mapping {
        Mapping::Resize {
            inner: RBox::new(L::to_value()),
            resize: SIZE,
        }
    }

    fn map(i: usize) -> Option<Index> {
        if i < SIZE { L::map(i) } else { None }
    }
}

Tiling

Tiling is implemented through indexed views, pure metadata transformations without data copies. The .tile() method extracts a tile by resizing one dimension to the tile size and offsetting into the buffer.

#![allow(unused)]
fn main() {
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;

axes![A = 8, B = 512];

let tensor = unsafe { HbmTensor::<bf16, m![1], m![A, B]>::from_addr(0) };
let view = tensor.view(); // HbmTensorView::<'_, bf16, m![1], m![A, B]>
let tile01 = view.tile::<m![B], 2, m![A, B = 2 # 512]>(0); // HbmTensorView::<'_, bf16, m![1], m![A, B = 2 # 512]>
let tile23 = view.tile::<m![B], 2, m![A, B = 2 # 512]>(2); // HbmTensorView::<'_, bf16, m![1], m![A, B = 2 # 512]>
}

The .tile() method takes three type parameters and one value parameter.

  • The tile dimension m![B] specifies which dimension to divide along.
  • The tile size 2 specifies the number of elements per tile.
  • The tile mapping m![A, B = 2 # 512] defines the resulting view’s mapping. The mapping B = 2 # 512 signifies that dimension B has a logical size of 2 within the view but exists within a physical footprint of 512. Without # 512, the stride between tiles would be 2 instead of 512, causing the view to read from wrong buffer positions.
  • The starting index specifies which tile to extract. Passing 0 captures the range 0..2 for tile01, while passing 2 captures the range 2..4 for tile23.

Stride and Modulo

Stride (/) and modulo (%) decompose a single dimension into two: the outer (block index) and the inner (position within block). Consider the 512-element axis B divided into 8 blocks of 64 elements each. The mapping m![B / 64, B % 64] creates an 8 × 64 grid where the first dimension selects which block and the second dimension selects the position within that block:

#![allow(unused)]
fn main() {
extern crate furiosa_opt_std;
extern crate furiosa_mapping;
use furiosa_opt_std::prelude::*;
axes![A = 8, B = 512];
type D1 = m![B / 64]; // stride with size 8
type D2 = m![B % 64]; // modulo with size 64

type E = m![B / 64, B % 64]; // equivalent to `m![B]`

fn test_stride_modulo() {
    assert_eq!(E::map(130), Some(i![B / 64: 2, B % 64: 2])); // block 2, position 2: B = 64*2 + 2 = 130
    assert_eq!(E::map(130), <m![B]>::map(130));               // same result as flat m![B]

    for i in 0..8 {
        assert_eq!(D1::map(i), Some(i![B / 64: i]));
    }
    assert_eq!(D1::map(8), None);

    for j in 0..64 {
        assert_eq!(D2::map(j), Some(i![B % 64: j]));
    }
    assert_eq!(D2::map(64), None);

    for i in 0..8 {
        for j in 0..64 {
            assert_eq!(
                E::map(64 * i + j),
                <m![B]>::map(64 * i + j),
            );
        }
    }
    assert_eq!(E::map(512), None);
}

test_stride_modulo();
}
impl<L, const SIZE: usize> M for Stride<L, SIZE>
where
    L: M,
{
    const SIZE: usize = {
        assert!(L::SIZE % SIZE == 0, "Stride size must divide the original size");
        L::SIZE / SIZE
    };

    fn to_value() -> Mapping {
        Mapping::Stride {
            inner: RBox::new(L::to_value()),
            stride: SIZE,
        }
    }

    fn map(i: usize) -> Option<Index> {
        if i < Self::SIZE { L::map(i * SIZE) } else { None }
    }
}

impl<L, const SIZE: usize> M for Modulo<L, SIZE>
where
    L: M,
{
    const SIZE: usize = {
        assert!(L::SIZE % SIZE == 0, "Modulo size must divide the original size");
        SIZE
    };

    fn to_value() -> Mapping {
        Mapping::Modulo {
            inner: RBox::new(L::to_value()),
            modulo: SIZE,
        }
    }

    fn map(i: usize) -> Option<Index> {
        if i < Self::SIZE { L::map(i % L::SIZE) } else { None }
    }
}

Stride and modulo mappings can be visualized in tabular form. Consider the mapping m![B / 4, B % 4] with B::SIZE = 16. The following table shows how buffer indices are arranged: each row corresponds to a specific index of B / 4 (the stride axis), and each column corresponds to an index of B % 4 (the modulo axis):

i![B % 4: 0]i![B % 4: 1]i![B % 4: 2]i![B % 4: 3]
i![B / 4: 0]i![B: 0]i![B: 1]i![B: 2]i![B: 3]
i![B / 4: 1]i![B: 4]i![B: 5]i![B: 6]i![B: 7]
i![B / 4: 2]i![B: 8]i![B: 9]i![B: 10]i![B: 11]
i![B / 4: 3]i![B: 12]i![B: 13]i![B: 14]i![B: 15]

Modulo differs from resize in how it handles buffer size:

  • Resize shrinks the buffer by truncating indices beyond the new size.
  • Modulo preserves the original buffer size while partitioning it into equal-sized blocks.

These operations can be nested for complex decompositions. The following example splits B into three dimensions where the buffer’s bit layout differs from that of the tensor index.

#![allow(unused)]
fn main() {
extern crate furiosa_opt_std;
extern crate furiosa_mapping;
use furiosa_opt_std::prelude::*;
axes![A = 8, B = 512];
// B's bits: 6 - 8,  0 - 4,          5
// Values:   0 - 7, 0 - 31,      0 - 1
type E = m![B / 64, B % 32, B / 32 % 2];

fn test_nested_stride() {
    assert_eq!(E::map(67), Some(i![B: 97])); // 67 = 64*1 + 2*1 + 1 (i=1,j=1,k=1) → B = 64*1 + 1 + 32*1 = 97
    // Verify B=97 round-trips: 97/64=1, 97%32=1, (97/32)%2=1
    assert_eq!(97 / 64, 1);
    assert_eq!(97 % 32, 1);
    assert_eq!((97 / 32) % 2, 1);

    // buffer index: 64 * i + 2 * j + k (i = block, j = position within block, k = sub-block)
    // tensor index B: 64 * i + j + 32 * k (rearranges bit positions)
    for i in 0..8 {
        for j in 0..32 {
            for k in 0..2 {
                assert_eq!(
                    E::map(64 * i + 2 * j + k),
                    Some(i![B: 64 * i + j + 32 * k]),
                );
            }
        }
    }
    assert_eq!(E::map(512), None);
}

test_nested_stride();
}

This kind of bit rearrangement maps naturally to hardware memory layouts where address bits are reordered for bank interleaving or cache efficiency. In binary, this rearranges bit positions: buffer 001_00001_1 becomes B = 001_1_00001. The buffer groups bits as [8:6]_[5:1]_[0] while B groups them as [8:6]_[5]_[4:0].

Tiling can operate on blocks rather than individual elements. The following example tiles by block using m![B / 32] and creates overlapping tiles:

#![allow(unused)]
fn main() {
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 8, B = 512];
let tensor = unsafe { HbmTensor::<bf16, m![1], m![A, B]>::from_addr(0) };
for i in 0..15 {
    let tile = tensor.view().tile::<m![B / 32], 2, m![A, B / 32 = 2 # 16, B % 32]>(i);
}
}

With B = 512, the dimension B / 32 has 16 blocks numbered 0-15. Each tile takes 2 consecutive blocks starting at index i. Tile 0 covers blocks {0, 1}, tile 1 covers blocks {1, 2}, and so on through tile 14 covering blocks {14, 15}. These tiles overlap because consecutive tiles share one block.

The tile mapping B / 32 = 2 resizes the block dimension to 2 since each tile contains exactly 2 blocks. When tiling with a single block, B / 32 = 1 simplifies to the identity m![1] since the dimension has only one value.

Escape

For complex mappings, define type aliases and reference them using { ... }. With separate mappings L = m![A] and R = m![B], combining them as m![{ L }, { R }] produces the same result as m![A, B]:

#![allow(unused)]
fn main() {
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![A = 8, B = 512];
type L = m![A];
type R = m![B];
type E = m![{ L }, { R }]; // equivalent to `m![A, B]`

fn test_escape() {
    for i in 0..E::SIZE {
        assert_eq!(E::map(i), <m![A, B]>::map(i));
    }
}

test_escape();
}

This escape syntax breaks down complex mappings into named, reusable components.

Advanced Constructors

Skewed axis

A skewed axis creates a diagonal access pattern across two dimensions. Skewed axes introduce derived axis labels defined by arithmetic differences between existing axes; for example, B' = B - A defines a new axis B' whose coordinate at any point equals B minus A. Algorithms that process data along diagonals use this pattern, such as certain wavefront computations.

The expression m![A, B' = 4] with B' = B - A creates a mapping where each row is shifted relative to the previous one. The = operator specifies the logical size after skewing. The result wraps around using modular arithmetic.

For example, with axes![A = 4, B = 4] and B' = B - A:

(A, B’)(A, B)
(0, 0)(0, 0)
(0, 1)(0, 1)
(0, 2)(0, 2)
(0, 3)(0, 3)
(1, 0)(1, 1)
(1, 1)(1, 2)
(1, 2)(1, 3)
(1, 3)(1, 0)

When A = 1 and B' = 3, the original B coordinate wraps to 0 via modular arithmetic since B = (B' + A) % 4 = (3 + 1) % 4 = 0.

Indirect sequencing

Sliding (linear combination)

Note

Linear combination expressions $(e1:n1, ..., ed:nd) combine multiple dimensions with specified strides. Formal definition: size_S($(e1:n1, ..., ed:nd)) = 1 + sum_k((size_S(ek) - 1) * nk). The mapping S, $(e1:n1, ..., ed:nd) |- si ~ ti holds if there exist si1...sid, ti1...tid such that for all k: S, ek |- sik ~ tik, si = sum_k(sik * nk), and ti = sum_k(tik * nk).

Linear combinations can encode outer sum: e1 * e2 is equivalent to $(e1 : size_S(e2), e2 : 1). However, outer sum is preferred because it’s more resilient to axis reordering. Changing e1 * e2 to e2 * e1 doesn’t require manual stride updates.

Sliding operations access overlapping data blocks, essential for convolutional neural networks. Consider a buffer of 9 elements representing a tensor with shape \(\{N=5, F=3\}\) where each row is a 3-element slice that slides one element at a time. The tensor element at \((N, F)\) maps to buffer index \(N + 2F\):

$$ \begin{array}{c|ccc} & F=0 & F=1 & F=2 \\ \hline N=0 & 0 & 2 & 4 \\ N=1 & 1 & 3 & 5 \\ N=2 & 2 & 4 & 6 \\ N=3 & 3 & 5 & 7 \\ N=4 & 4 & 6 & 8 \\ \end{array} $$

Note

In this sliding pattern, a single space index can map to multiple tensor indices. For example, space index 4 maps to {4_N}, {2_N, 1_F}, and {2_F} simultaneously. This illustrates the non-one-to-one nature of (S, e).maps(si, ti).

This can be expressed using a linear combination expression where the N axis has stride 1 and the F axis has stride 2, yielding a total size of 1 + (5-1)*1 + (3-1)*2 = 9.

Equivalent Mapping

Different constructor combinations can produce the same mapping. Specifically, mappings E1 and E2 are equivalent when:

  • E1::SIZE == E2::SIZE, and
  • For every i, E1::map(i) == E2::map(i).

The equivalence relation is reflexive, symmetric, and transitive. The following identities capture common equivalences:

  • Identity of pairs: for every E, E is equivalent both to m![{ E }, 1] and m![1, { E }].
  • Stride-modulo decomposition: for every E whose size E::SIZE is divisible by n, E and m![{ E } / n, { E } % n] are equivalent.
  • Pair projection: for every A and B, m![[{ A }, { B }] / B::SIZE] is equivalent to m![A] and m![[{ A }, { B }] % B::SIZE] is equivalent to m![B].
  • Associativity of pairs: for every E1, E2, E3, m![{ E1 }, { E2 }, { E3 }], m![[{ E1 }, { E2 }], { E3 }], and m![{ E1 }, [{ E2 }, { E3 }]] are equivalent.
  • Idempotent operations: for every E, E is equivalent to m![{ E } / 1], to m![{ E } # E::SIZE], and to m![{ E } = E::SIZE].
  • Modulo by 1: for every E, m![E % 1] is equivalent to the identity mapping m![1].