Tensor Access Patterns (taplib)

taplib describes how DMAs walk tensors. A DMA buffer descriptor executes a strided walk: an element offset plus parallel sizes and strides, outermost dimension first. A TensorAccessPattern is exactly that walk over a tensor of known shape. The tilings designs use can be built from TensorAccessPattern.full(dims) (the row-major walk over the whole tensor) with a few operations on those integers.

Building patterns

To	Use
Walk a whole tensor row by row	TensorAccessPattern.full(dims)
Restrict the walk with NumPy indexing	tap[2:6, ::2], tap[i]
Reorder or reverse the dimensions	tap.permute((1, 0)), tap.T
Split or merge a dimension	tap.split(dim, inner), tap.merge(dim)
Use the fewest dimensions for the same walk	tap.coalesce()
Walk the same data again	tap.repeat(n) (a stride-0 outermost dimension)
Have a memtile pad the stream	tap.pad([(before, after), ...])
Cut a tensor into equal tiles	tap.tile(tile_dims)
Cut one dimension into k equal chunks	tap.partition(k)
Read a tile-blocked buffer back in row-major order	tap.inverse()

tile() turns a rank-r pattern into a rank-2r one: the grid dimensions first, then the tile dimensions, so the whole result walks every tile in row-major order. Everything else is ordinary indexing and reordering of those dimensions:

To	Use
One tile	tiles[i, j]
A row of tiles	tiles[i]
A block of tiles, or every S-th tile	tiles[a:b, j::S]
Visit tiles column by column	tiles.permute((1, 0, 2, 3))
Walk each tile column-major	tiles.permute((0, 1, 3, 2))
Walk a group of tiles again	tiles[i].repeat(n)

Patterns are immutable: every operation returns a new one, and a list of patterns is just a Python list.

from aie.helpers.taplib import TensorAccessPattern

tiles = TensorAccessPattern.full((16, 16)).tile((4, 4))
print(tiles)        # TensorAccessPattern([16, 16], offset=0, sizes=[4, 4, 4, 4], strides=[64, 4, 16, 1])
print(tiles[1, 2])  # TensorAccessPattern([16, 16], offset=72, sizes=[4, 4], strides=[16, 1])
print(tiles[:2, :2])
                    # TensorAccessPattern([16, 16], offset=0, sizes=[2, 2, 4, 4], strides=[64, 4, 16, 1])
print(TensorAccessPattern.full((1024,)).partition(4)[2])
                    # TensorAccessPattern([1024], offset=512, sizes=[256], strides=[1])

Tiling a tiled pattern tiles hierarchically. tiles.tile((1, 1, 4, 4)) keeps the grid and cuts each tile into 4 x 4 sub-tiles, so the walk goes tile by tile and, inside each tile, sub-tile by sub-tile. tiles[i, j].tile((4, 4)) does the same for one tile. Such a walk can need more dimensions than one buffer descriptor holds, even after coalesce(). A shim fill() or drain() with a static walk is then split into several transfers. An ObjectFifo pattern has to fit in one descriptor (four dimensions on a memtile, three on a core tile).

What a DMA can walk

A DMA moves whole 4-byte words. taplib itself accepts any walk, and the compiler rejects one the hardware cannot do when it lowers the design:

• Every size, and every stride other than an innermost 1, must span whole words. A bf16 run is an even number of elements, and an int8 run a multiple of four.
• The innermost dimension steps a word at a time. For elements that are not 32 bits wide, its stride must be 1.

So for bf16 or int8 data, .T and permute() of a tile are rejected, and so is a ::2 slice of the innermost dimension. For tiles[0, 0].T on a bf16 tensor, a shim fill() fails with:

error: 'aie.dma_bd' op Stride 1 is 1 elements * 2 bytes = 2 bytes, which is not divisible by 4.

The same walk as a memtile's to_stream fails with:

error: 'aie.dma_bd' op For <32b width datatypes, inner-most dim stride must be 1

To transpose sub-word data, let the DMA move s x s blocks whose rows are whole words, and transpose each block in the kernel. The transposes example below does this.

Using patterns

A pattern goes wherever IRON takes a DMA walk:

• fill() and drain() in a runtime sequence take it as tap=.
• An ObjectFifo takes it as to_stream (how the producer's DMA reads its object onto the stream) and as from_stream / from_stream_per_cons (how a consumer's DMA writes the stream into its object). The pattern walks a tensor the size of what each transfer moves, from offset 0: one object, or one segment of it on a join's output or a distribute's input. A padded pattern as to_stream also pads the stream on a MemTile.
• tap.transformation_dims gives the ((size, stride), ...) pairs, for code that still wants them.

Checking a data path on the host

On its way from the host to a core, a tensor passes through up to four DMAs, each walking it with its own pattern: the shim reads the host tensor onto the stream, the memtile writes it into an object, the memtile reads that object back out, and the core writes it into its own object. On hardware a mistake in any of them just looks like scrambled data in the kernel.

gather(tensor) returns the stream a DMA emits when it walks tensor with a pattern, and scatter(stream) the object a DMA stores when it writes stream with a pattern. Chaining them with the patterns a design uses shows exactly what each core receives, with no hardware. This is the transposes design's --strategy=combined path: the memtile shuffles each tile into s x s blocks so that the kernel only has to transpose each block in place.

import numpy as np
from aie.helpers.taplib import TensorAccessPattern

M, K, m, n, s = 64, 64, 16, 16, 8
host = np.arange(M * K).reshape(M, K)
shim = TensorAccessPattern.full((M, K)).tile((m, n))
memtile_in = TensorAccessPattern.full((n, m)).tile((s, s)).permute((1, 2, 0, 3))

for t, tile in enumerate(shim.gather(host).reshape(-1, m * n)):
    i, j = divmod(t, K // n)
    obj = memtile_in.scatter(tile)  # the memtile's (n, m) object
    blocks = obj.reshape(n // s, s, m // s, s)
    kernel_out = blocks.transpose(0, 3, 2, 1).reshape(n, m)
    assert (kernel_out == host[i * m : (i + 1) * m, j * n : (j + 1) * n].T).all()

For a padded pattern, gather(tensor, pad_value=) fills the padded positions, and padded_sizes is the shape the receiving object must have. To inspect a single walk, use accesses(), access_order(), access_count(), compare_access_orders() and visualize().

Staged patterns in a dispatch-time sequence

Every operation also accepts staged values: an aie.ir.Value (a DispatchTime[T] scalar a runtime sequence receives, or any arithmetic on one) can stand in for a size, stride, offset, index, slice bound or repeat count. The arithmetic is emitted as arith ops where it is used, and each check that would have raised ValueError becomes a cf.assert guard. A fully static specialization folds the guard away, or fails at generation time if it is false. The dispatch-time builder instead returns no stream, and the host refuses the call with the guard's message.

def seq(a_h, b_h, start, n, in_prod, out_cons):
    # The buffer as max_tiles equal chunks; the chunk index is staged
    # arithmetic (start + loop iv), so the tap's offset is too.
    chunks = TensorAccessPattern.full((max_tiles * tile_size,)).partition(max_tiles)
    for tile in range_(n):
        tap = chunks[start + tile]
        tg = TaskGroup()
        out_cons.drain(b_h, tap=tap, wait=True, group=tg)
        in_prod.fill(a_h, tap=tap, group=tg)
        tg.finish()

require(cond, message) from aie.iron adds a guard of your own, such as a shape constraint the design depends on. cond can depend on dispatch-time values. whole_array.py (below) has require(M % (m * n_aie_rows) == 0, "M must be a multiple of m * n_aie_rows") with a dispatch-time M, and a call whose M breaks it is refused with that message. The message is a plain string fixed when the design is generated, so it cannot include the value the call passed.

Some rules of thumb:

• Sizes, strides, offsets and loop bounds can all be dispatch-time values, but the number of transfers a loop step issues cannot. A range_ over a dispatch-time bound is emitted as a loop, so its body is generated once. Every step issues the same transfers, waited on the same way, in the same order, and the TaskGroups carried from one step to the next are a fixed set of values (see Runtime tasks). This comes from how the sequence is generated, not from the DMA. So when the data does not divide evenly into blocks, the loop covers the whole blocks, and the ragged last block is one more transfer after it, under with if_(rem > 0):, with its size computed from rem.
• The structure of a pattern stays a Python value: its rank, slice steps, padding amounts, and the sizes and strides of the dimensions merge() combines.
• Staged values of any integer type are accepted. The compiler hoists their arithmetic out of the buffer-descriptor block, and a value too wide for its descriptor field is refused at dispatch instead of being truncated.

programming_examples/basic/matrix_multiplication/whole_array/whole_array.py is the whole-array GEMM written this way, with dispatch-time M, K and N.

Checking a dispatch-time design

instructions(**scalars) on an @iron.jit design returns the instruction words a call with those values would run, built exactly as a call builds them, so no NPU is needed. A dispatch-time stream is not byte-identical to a static specialization's: it draws buffer descriptors from a pool, polls for room in a channel's task queue and assembles descriptor words at build time. aie.utils.txn_trace reduces either stream to the events the hardware acts on (queue pushes resolved to the chain of transfers they start, token waits, and other register writes such as runtime parameters) and compares those:

from aie.utils.txn_trace import compare, explain

words = tiled_copy.specialize().instructions(n_tiles=3, start_tile=1)
static = tiled_copy.specialize(n_tiles=3, start_tile=1).instructions()
assert compare(words, static) == []
print(explain(words))  # one line per event

python -m aie.utils.txn_trace insts.bin [other.bin] does the same from the command line. test/python/dispatch_taplib_copy.py and the whole-array GEMM's tests/dispatch_txn.py are complete examples.

API reference

TensorAccessPattern

TensorAccessPattern(
    tensor_dims: Sequence[IntLike],
    offset: IntLike,
    sizes: Sequence[IntLike],
    strides: Sequence[IntLike],
)

A strided walk over a tensor: offset plus outermost-first sizes and strides.

Build one with full() and refine it, or give the numbers directly. Instances are immutable; every method returns a new pattern.

Two patterns are equal when they walk the same tensor the same way, i.e. when their coalesce() forms match. Staged patterns are equal only to themselves.

Create an access pattern.

Values may be Python ints or staged runtime values; a check on a staged value becomes a dispatch-time guard.

Parameters:

Name	Type	Description	Default
tensor_dims	Sequence[IntLike]	Shape of the tensor the pattern walks.	required
offset	IntLike	Element offset of the first element visited.	required
sizes	Sequence[IntLike]	Extent of each dimension, outermost first.	required
strides	Sequence[IntLike]	Element step of each dimension, outermost first.	required

Raises:

Type	Description
TypeError	If a value is not an integer.
ValueError	If sizes and strides differ in length or a concrete value is out of range.

Source code in python/helpers/taplib/tap.py

def __init__(
    self,
    tensor_dims: Sequence[IntLike],
    offset: IntLike,
    sizes: Sequence[IntLike],
    strides: Sequence[IntLike],
):
    """Create an access pattern.

    Values may be Python ints or staged runtime values; a check on a staged
    value becomes a dispatch-time guard.

    Args:
        tensor_dims (Sequence[IntLike]): Shape of the tensor the pattern walks.
        offset (IntLike): Element offset of the first element visited.
        sizes (Sequence[IntLike]): Extent of each dimension, outermost first.
        strides (Sequence[IntLike]): Element step of each dimension, outermost first.

    Raises:
        TypeError: If a value is not an integer.
        ValueError: If `sizes` and `strides` differ in length or a
            concrete value is out of range.
    """
    tensor_dims = [sint(d) for d in tensor_dims]
    offset = sint(offset)
    sizes = [sint(s) for s in sizes]
    strides = [sint(s) for s in strides]
    self._tensor_dims = tuple(validate_tensor_dims(tensor_dims))
    self._offset = validate_offset(offset, tensor_dims)
    sizes, strides = validate_and_clean_sizes_strides(sizes, strides)
    self._sizes: tuple[IntLike, ...] = tuple(sizes)
    self._strides: tuple[IntLike, ...] = tuple(strides)
    self._padding: tuple[tuple[int, int], ...] | None = None

tensor_dims property

tensor_dims: tuple[IntLike, ...]

Shape of the tensor the pattern walks.

offset property

offset: IntLike

Element offset of the first element visited.

sizes property

sizes: tuple[IntLike, ...]

Extent of each dimension, outermost first.

strides property

strides: tuple[IntLike, ...]

Element step of each dimension, outermost first.

transformation_dims property

transformation_dims: tuple[tuple[IntLike, IntLike], ...]

The pattern as ((size, stride), ...), outermost first.

padding property

padding: tuple[tuple[int, int], ...] | None

(before, after) padding per dimension, or None if the pattern is unpadded.

padded_sizes property

padded_sizes: tuple[IntLike, ...]

Extent of each dimension on the (possibly padded) stream.

rank property

rank: int

Number of dimensions of the walk.

numel property

numel: IntLike

Number of tensor elements the walk visits (repeats counted, padding not).

is_symbolic property

is_symbolic: bool

Whether any value of the pattern is staged.

T property

T: TensorAccessPattern

The walk with its dimensions reversed, like numpy.ndarray.T.

TensorAccessPattern.full((M, N)).T walks an (M, N) tensor column by column.

full classmethod

full(tensor_dims: Sequence[IntLike]) -> TensorAccessPattern

Return the row-major walk over a whole tensor of shape tensor_dims.

Parameters:

Name	Type	Description	Default
tensor_dims	Sequence[IntLike]	Shape of the tensor.	required

Returns:

Name	Type	Description
TensorAccessPattern	TensorAccessPattern	A pattern that visits every element once, in order.

Source code in python/helpers/taplib/tap.py

@classmethod
def full(cls, tensor_dims: Sequence[IntLike]) -> TensorAccessPattern:
    """Return the row-major walk over a whole tensor of shape `tensor_dims`.

    Args:
        tensor_dims (Sequence[IntLike]): Shape of the tensor.

    Returns:
        TensorAccessPattern: A pattern that visits every element once, in order.
    """
    dims = [sint(d) for d in tensor_dims]
    return cls(dims, 0, dims, row_major_strides(dims))

permute

permute(axes: Sequence[int]) -> TensorAccessPattern

Reorder dimensions: result dimension i is this pattern's dimension axes[i].

Whether a DMA can execute the result depends on its 4-byte word granule: every non-unit stride must span whole words, and for elements other than 32 bits the innermost stride must be 1. The DMA verifier and the dynamic lowering enforce this, so .T of a bf16 or int8 tile is rejected when the design is lowered.

Parameters:

Name	Type	Description	Default
axes	Sequence[int]	A permutation of range(rank).	required

Returns:

Name	Type	Description
TensorAccessPattern	TensorAccessPattern	The reordered walk.

Raises:

Type	Description
ValueError	If axes is not a permutation of range(rank).

Source code in python/helpers/taplib/tap.py

def permute(self, axes: Sequence[int]) -> TensorAccessPattern:
    """Reorder dimensions: result dimension `i` is this pattern's dimension `axes[i]`.

    Whether a DMA can execute the result depends on its 4-byte word
    granule: every non-unit stride must span whole words, and for elements
    other than 32 bits the innermost stride must be 1. The DMA verifier and
    the dynamic lowering enforce this, so `.T` of a bf16 or int8 tile is
    rejected when the design is lowered.

    Args:
        axes (Sequence[int]): A permutation of `range(rank)`.

    Returns:
        TensorAccessPattern: The reordered walk.

    Raises:
        ValueError: If `axes` is not a permutation of `range(rank)`.
    """
    axes = validate_permutation(axes, self.rank, "axes")
    return self._with(
        sizes=[self._sizes[a] for a in axes],
        strides=[self._strides[a] for a in axes],
    )

split

split(dim: int, inner: IntLike) -> TensorAccessPattern

Split dimension dim of size n into (n // inner, inner).

The outer part strides by inner * stride; the inner part keeps the original stride.

Parameters:

Name	Type	Description	Default
dim	int	The dimension to split.	required
inner	IntLike	Size of the new inner dimension; must divide n.	required

Returns:

Name	Type	Description
TensorAccessPattern	TensorAccessPattern	The walk with one more dimension.

Raises:

Type	Description
ValueError	If inner does not divide the dimension.

Source code in python/helpers/taplib/tap.py

def split(self, dim: int, inner: IntLike) -> TensorAccessPattern:
    """Split dimension `dim` of size `n` into `(n // inner, inner)`.

    The outer part strides by `inner * stride`; the inner part keeps the
    original stride.

    Args:
        dim (int): The dimension to split.
        inner (IntLike): Size of the new inner dimension; must divide `n`.

    Returns:
        TensorAccessPattern: The walk with one more dimension.

    Raises:
        ValueError: If `inner` does not divide the dimension.
    """
    dim = self._axis(dim)
    inner = sint(inner)
    n, s = self._sizes[dim], self._strides[dim]
    require(
        n % inner == 0,
        f"dimension {dim} of size {show(n)} is not divisible by {show(inner)}",
    )
    sizes = (*self._sizes[:dim], n // inner, inner, *self._sizes[dim + 1 :])
    strides = (*self._strides[:dim], s * inner, s, *self._strides[dim + 1 :])
    return self._with(sizes=sizes, strides=strides)

merge

merge(dim: int) -> TensorAccessPattern

Merge dimensions dim and dim + 1 into one.

Only legal when they are contiguous, i.e. strides[dim] == sizes[dim + 1] * strides[dim + 1], or when either has size 1. Whether to merge is a structural decision, so the values must be concrete.

Parameters:

Name	Type	Description	Default
dim	int	The outer of the two dimensions.	required

Returns:

Name	Type	Description
TensorAccessPattern	TensorAccessPattern	The walk with one fewer dimension.

Raises:

Type	Description
ValueError	If the two dimensions are not contiguous.
TypeError	If either dimension is staged.

Source code in python/helpers/taplib/tap.py

def merge(self, dim: int) -> TensorAccessPattern:
    """Merge dimensions `dim` and `dim + 1` into one.

    Only legal when they are contiguous, i.e. ``strides[dim] ==
    sizes[dim + 1] * strides[dim + 1]``, or when either has size 1.
    Whether to merge is a structural
    decision, so the values must be concrete.

    Args:
        dim (int): The outer of the two dimensions.

    Returns:
        TensorAccessPattern: The walk with one fewer dimension.

    Raises:
        ValueError: If the two dimensions are not contiguous.
        TypeError: If either dimension is staged.
    """
    dim = self._axis(dim)
    if dim + 1 >= self.rank:
        raise ValueError(f"dimension {dim} has no successor to merge with")
    n0, s0 = self._sizes[dim], self._strides[dim]
    n1, s1 = self._sizes[dim + 1], self._strides[dim + 1]
    if sym_any([n0, s0, n1, s1]):
        raise TypeError(
            "merge() is a structural decision; it needs concrete sizes/strides"
        )
    if n0 != 1 and n1 != 1 and s0 != n1 * s1:
        raise ValueError(
            f"dimensions {dim} and {dim + 1} are not contiguous "
            f"(stride {s0} != {n1} * {s1}); cannot merge"
        )
    sizes = (*self._sizes[:dim], n0 * n1, *self._sizes[dim + 2 :])
    stride = s0 if n1 == 1 else s1
    strides = (*self._strides[:dim], stride, *self._strides[dim + 2 :])
    return self._with(sizes=sizes, strides=strides)

coalesce

coalesce() -> TensorAccessPattern

Return the same walk in the fewest dimensions.

Size-1 dimensions are dropped and every contiguous adjacent pair is merged; the result visits the same elements in the same order. A pair with a staged size or stride is kept. The DMA builders apply this themselves to a walk deeper than the hardware takes.

Returns:

Name	Type	Description
TensorAccessPattern	TensorAccessPattern	The same walk with no size-1 or mergeable dimensions.

Source code in python/helpers/taplib/tap.py

def coalesce(self) -> TensorAccessPattern:
    """Return the same walk in the fewest dimensions.

    Size-1 dimensions are dropped and every contiguous adjacent pair is
    merged; the result visits the same elements in the same order. A pair
    with a staged size or stride is kept. The DMA builders apply this
    themselves to a walk deeper than the hardware takes.

    Returns:
        TensorAccessPattern: The same walk with no size-1 or mergeable dimensions.
    """
    out = self._drop_unit_dims()
    i = 0
    while i + 1 < out.rank:
        pair = out._sizes[i : i + 2] + out._strides[i : i + 2]
        if not sym_any(pair) and (
            out._strides[i] == out._sizes[i + 1] * out._strides[i + 1]
        ):
            out = out.merge(i)
        else:
            i += 1
    return out

repeat

repeat(count: IntLike) -> TensorAccessPattern

Walk the whole pattern count times: a new outermost dimension with stride 0.

On a shim DMA the outermost dimension becomes the queue repeat; on a memtile it is a plain zero-stride dimension.

Parameters:

Name	Type	Description	Default
count	IntLike	Number of walks; must be >= 1.	required

Returns:

Name	Type	Description
TensorAccessPattern	TensorAccessPattern	The walk with a new outermost dimension of count.

Raises:

Type	Description
ValueError	If count is below 1.

Source code in python/helpers/taplib/tap.py

def repeat(self, count: IntLike) -> TensorAccessPattern:
    """Walk the whole pattern `count` times: a new outermost dimension with stride 0.

    On a shim DMA the outermost dimension becomes the queue repeat; on a
    memtile it is a plain zero-stride dimension.

    Args:
        count (IntLike): Number of walks; must be >= 1.

    Returns:
        TensorAccessPattern: The walk with a new outermost dimension of `count`.

    Raises:
        ValueError: If `count` is below 1.
    """
    count = sint(count)
    require(count >= 1, f"repeat count must be >= 1, got {show(count)}")
    return self._with(sizes=(count, *self._sizes), strides=(0, *self._strides))

tile

tile(tile_dims: Sequence[IntLike]) -> TensorAccessPattern

Divide every dimension into tiles of tile_dims.

Dimension i of size n_i becomes a grid dimension of n_i // t_i tiles and a tile dimension of t_i elements. The grid dimensions come first, so a rank-r pattern becomes a rank-2r one that walks the tiles in row-major order: tiles[i, j] is the tile at grid position (i, j) and tiles[i] the i-th row of tiles. Reorder the grid with permute() ((1, 0, 2, 3) walks a column of tiles at a time) and take strided or partial groups of tiles by slicing.

Parameters:

Name	Type	Description	Default
tile_dims	Sequence[IntLike]	One tile extent per dimension; each must divide the dimension.	required

Returns:

Name	Type	Description
TensorAccessPattern	TensorAccessPattern	The grid dimensions followed by the tile dimensions.

Raises:

Type	Description
ValueError	If tile_dims does not have one entry per dimension, or an entry is below 1 or does not divide its dimension.

Source code in python/helpers/taplib/tap.py

def tile(self, tile_dims: Sequence[IntLike]) -> TensorAccessPattern:
    """Divide every dimension into tiles of `tile_dims`.

    Dimension `i` of size `n_i` becomes a grid dimension of `n_i // t_i`
    tiles and a tile dimension of `t_i` elements. The grid dimensions
    come first, so a rank-`r` pattern becomes a rank-`2r` one that walks
    the tiles in row-major order: `tiles[i, j]` is the tile at grid
    position `(i, j)` and `tiles[i]` the `i`-th row of tiles. Reorder
    the grid with `permute()` (`(1, 0, 2, 3)` walks a column of tiles at
    a time) and take strided or partial groups of tiles by slicing.

    Args:
        tile_dims (Sequence[IntLike]): One tile extent per dimension; each
            must divide the dimension.

    Returns:
        TensorAccessPattern: The grid dimensions followed by the tile dimensions.

    Raises:
        ValueError: If `tile_dims` does not have one entry per dimension,
            or an entry is below 1 or does not divide its dimension.
    """
    tile_dims = [sint(t) for t in tile_dims]
    if len(tile_dims) != self.rank:
        raise ValueError(
            f"tile_dims has {len(tile_dims)} entries for a pattern of rank {self.rank}"
        )
    grid_sizes: list[IntLike] = []
    grid_strides: list[IntLike] = []
    for dim, t in enumerate(tile_dims):
        n, s = self._sizes[dim], self._strides[dim]
        require(t >= 1, f"tile_dims[{dim}] must be >= 1")
        require(
            n % t == 0,
            f"dimension {dim} of size {show(n)} is not divisible by tile size {show(t)}",
        )
        grid_sizes.append(n // t)
        grid_strides.append(s * t)
    return self._with(
        sizes=(*grid_sizes, *tile_dims), strides=(*grid_strides, *self._strides)
    )

partition

partition(
    parts: IntLike, dim: int = 0
) -> TensorAccessPattern

Split dimension dim into parts equal contiguous pieces.

Like np.array_split on an evenly divisible axis, it splits the outermost dimension unless told otherwise. The result has a new outermost dimension of parts, so TensorAccessPattern.full((N,)).partition(k)[i] is the i-th of k equal chunks of a flat range. Other dimensions are kept whole.

Parameters:

Name	Type	Description	Default
parts	IntLike	Number of pieces; must divide the dimension.	required
dim	int	The dimension to split. Defaults to 0.	0

Returns:

Name	Type	Description
TensorAccessPattern	TensorAccessPattern	The pieces, one per index of the new outermost dimension.

Raises:

Type	Description
ValueError	If parts is below 1 or does not divide the dimension.

Source code in python/helpers/taplib/tap.py

def partition(self, parts: IntLike, dim: int = 0) -> TensorAccessPattern:
    """Split dimension `dim` into `parts` equal contiguous pieces.

    Like `np.array_split` on an evenly divisible axis, it splits the
    outermost dimension unless told otherwise. The result has a new
    outermost dimension of `parts`, so
    `TensorAccessPattern.full((N,)).partition(k)[i]` is the `i`-th of
    `k` equal chunks of a flat range. Other dimensions are kept whole.

    Args:
        parts (IntLike): Number of pieces; must divide the dimension.
        dim (int, optional): The dimension to split. Defaults to 0.

    Returns:
        TensorAccessPattern: The pieces, one per index of the new outermost dimension.

    Raises:
        ValueError: If `parts` is below 1 or does not divide the dimension.
    """
    dim = self._axis(dim)
    parts = sint(parts)
    n, s = self._sizes[dim], self._strides[dim]
    require(parts >= 1, "parts must be >= 1")
    require(
        n % parts == 0,
        f"dimension {dim} of size {show(n)} is not divisible into {show(parts)} parts",
    )
    chunk = n // parts
    sizes = list(self._sizes)
    sizes[dim] = chunk
    return self._with(sizes=(parts, *sizes), strides=(s * chunk, *self._strides))

inverse

inverse() -> TensorAccessPattern

Return the walk that puts this walk's stream back in row-major order.

A DMA that streams a tensor with a pattern visiting every element exactly once stores a buffer B with B[k] = tensor.flat[walk[k]]. The inverse walks B and emits the tensor in row-major order, so p.inverse().gather(p.gather(x)) is x.reshape(-1). For the tiling full((m, n)).tile((r, t)) it is the "un-blocking" walk a memtile applies to a core's blocked output: sizes [m//r, r, n//t, t] with strides [r*n, t, r*t, 1].

Returns:

Name	Type	Description
TensorAccessPattern	TensorAccessPattern	The inverse walk, over a buffer of the same shape.

Raises:

Type	Description
TypeError	If the pattern is staged.
ValueError	If the pattern is padded or does not visit every element of its tensor exactly once.

Source code in python/helpers/taplib/tap.py

def inverse(self) -> TensorAccessPattern:
    """Return the walk that puts this walk's stream back in row-major order.

    A DMA that streams a tensor with a pattern visiting every element
    exactly once stores a buffer `B` with `B[k] = tensor.flat[walk[k]]`.
    The inverse walks `B` and emits the tensor in row-major order, so
    `p.inverse().gather(p.gather(x))` is `x.reshape(-1)`. For the
    tiling `full((m, n)).tile((r, t))` it is the "un-blocking" walk a
    memtile applies to a core's blocked output: sizes `[m//r, r, n//t, t]`
    with strides `[r*n, t, r*t, 1]`.

    Returns:
        TensorAccessPattern: The inverse walk, over a buffer of the same shape.

    Raises:
        TypeError: If the pattern is staged.
        ValueError: If the pattern is padded or does not visit every
            element of its tensor exactly once.
    """
    self._require_concrete("inverse()")
    if self._padding is not None:
        raise ValueError("a padded pattern has no inverse")
    steps = [i for i, n in enumerate(self._sizes) if n != 1]
    steps.sort(key=lambda i: self._strides[i], reverse=True)
    covered = 1
    for i in reversed(steps):
        if self._strides[i] != covered:
            covered = 0
            break
        covered *= self._sizes[i]
    if self._offset != 0 or covered != sprod(self._tensor_dims):
        raise ValueError(
            f"{self} does not visit every element exactly once, so it has no inverse"
        )
    if not steps:
        return TensorAccessPattern.full(self._tensor_dims)
    stream_strides = row_major_strides(self._sizes)
    return TensorAccessPattern._raw(
        self._tensor_dims,
        0,
        [self._sizes[i] for i in steps],
        [stream_strides[i] for i in steps],
    )

pad

pad(
    padding: Sequence[Sequence[int]],
) -> TensorAccessPattern

Surround every walk of each dimension with constant elements.

A memtile MM2S channel can pad the stream it emits: for dimension i it inserts before constant elements ahead of each pass over the dimension and after behind it, so the stream is prod(padded_sizes) elements long. Give a padded pattern to an ObjectFifo as to_stream and it sets the padding too (the pad value is set on the fifo). Padding is applied last: a padded pattern cannot be reshaped further.

Parameters:

Name	Type	Description	Default
padding	Sequence[Sequence[int]]	One (before, after) pair per dimension, outermost first, as compile-time ints.	required

Returns:

Name	Type	Description
TensorAccessPattern	TensorAccessPattern	This walk with its padding.

Raises:

Type	Description
TypeError	If a padding count is staged.
ValueError	If the pattern is already padded, an entry is not a pair of counts >= 0, or there is not one entry per dimension.

Source code in python/helpers/taplib/tap.py

def pad(self, padding: Sequence[Sequence[int]]) -> TensorAccessPattern:
    """Surround every walk of each dimension with constant elements.

    A memtile MM2S channel can pad the stream it emits: for dimension
    `i` it inserts `before` constant elements ahead of each pass over
    the dimension and `after` behind it, so the stream is
    `prod(padded_sizes)` elements long. Give a padded pattern to an
    ObjectFifo as `to_stream` and it sets the padding too (the pad value
    is set on the fifo). Padding is applied last: a padded pattern cannot
    be reshaped further.

    Args:
        padding (Sequence[Sequence[int]]): One `(before, after)` pair per
            dimension, outermost first, as compile-time ints.

    Returns:
        TensorAccessPattern: This walk with its padding.

    Raises:
        TypeError: If a padding count is staged.
        ValueError: If the pattern is already padded, an entry is not a
            pair of counts >= 0, or there is not one entry per dimension.
    """
    if self._padding is not None:
        raise ValueError("this pattern is already padded")
    pads = []
    for entry in padding:
        if len(entry) != 2:
            raise ValueError("each padding entry is a (before, after) pair")
        before, after = (sint(v) for v in entry)
        if is_sym(before) or is_sym(after):
            raise TypeError("padding counts must be compile-time ints")
        if before < 0 or after < 0:
            raise ValueError(f"padding counts must be >= 0, got {tuple(entry)}")
        pads.append((before, after))
    if len(pads) != self.rank:
        raise ValueError(
            f"padding has {len(pads)} entries for a pattern of rank {self.rank}"
        )
    if not any(b or a for b, a in pads):
        return self
    return TensorAccessPattern._raw(
        self._tensor_dims, self._offset, self._sizes, self._strides, tuple(pads)
    )

gather

gather(tensor: Any, pad_value: Any = 0) -> ndarray

Return the stream a DMA walking tensor with this pattern emits.

This is what an ObjectFifo given this pattern as to_stream sends, padding included.

Parameters:

Name	Type	Description	Default
tensor	array_like	A tensor of shape tensor_dims (any shape with the same number of elements is accepted).	required
pad_value	optional	The value padded positions hold. Defaults to 0.	0

Returns:

Type	Description
ndarray	np.ndarray: A 1-D array of prod(padded_sizes) elements.

Raises:

Type	Description
TypeError	If the pattern is staged.
ValueError	If tensor does not have prod(tensor_dims) elements or the walk leaves the tensor.

Source code in python/helpers/taplib/tap.py

def gather(self, tensor: Any, pad_value: Any = 0) -> np.ndarray:
    """Return the stream a DMA walking `tensor` with this pattern emits.

    This is what an ObjectFifo given this pattern as `to_stream` sends,
    padding included.

    Args:
        tensor (array_like): A tensor of shape `tensor_dims` (any
            shape with the same number of elements is accepted).
        pad_value (optional): The value padded positions hold. Defaults to 0.

    Returns:
        np.ndarray: A 1-D array of `prod(padded_sizes)` elements.

    Raises:
        TypeError: If the pattern is staged.
        ValueError: If `tensor` does not have `prod(tensor_dims)` elements
            or the walk leaves the tensor.
    """
    arr = np.asarray(tensor)
    idx = self._walk().reshape(-1)
    self._check_numel(arr.size, "tensor")
    out = arr.reshape(-1)[np.maximum(idx, 0)]
    if self._padding is not None:
        out = np.where(idx >= 0, out, np.asarray(pad_value, dtype=arr.dtype))
    return out

scatter

scatter(stream: Any, out: ndarray | None = None) -> ndarray

Write stream into a tensor the way a DMA walking this pattern does.

This is what an ObjectFifo given this pattern as from_stream stores: stream element k lands at the k-th position of the walk.

Parameters:

Name	Type	Description	Default
stream	array_like	The stream, one element per step of the walk.	required
out	ndarray | None	The tensor to write into; elements the walk does not visit keep their values. Defaults to a zero tensor of shape tensor_dims.	None

Returns:

Type	Description
ndarray	np.ndarray: The written tensor.

Raises:

Type	Description
ValueError	If the pattern is padded (only an emitting DMA pads) or the stream length does not match the walk.

Source code in python/helpers/taplib/tap.py

def scatter(self, stream: Any, out: np.ndarray | None = None) -> np.ndarray:
    """Write `stream` into a tensor the way a DMA walking this pattern does.

    This is what an ObjectFifo given this pattern as `from_stream` stores:
    stream element `k` lands at the `k`-th position of the walk.

    Args:
        stream (array_like): The stream, one element per step of the walk.
        out (np.ndarray | None, optional): The tensor to write into; elements the
            walk does not visit keep their values. Defaults to a zero tensor
            of shape `tensor_dims`.

    Returns:
        np.ndarray: The written tensor.

    Raises:
        ValueError: If the pattern is padded (only an emitting DMA pads) or
            the stream length does not match the walk.
    """
    if self._padding is not None:
        raise ValueError("only a to_stream pattern can pad")
    stream = np.asarray(stream).reshape(-1)
    idx = self._walk().reshape(-1)
    if stream.size != idx.size:
        raise ValueError(
            f"stream has {stream.size} elements for a walk of {idx.size}"
        )
    if out is None:
        out = np.zeros(self._tensor_dims, dtype=stream.dtype)
    self._check_numel(out.size, "out")
    flat = out.reshape(-1)
    flat[idx] = stream
    if not np.shares_memory(flat, out):
        out[...] = flat.reshape(out.shape)
    return out

accesses

accesses() -> tuple[ndarray, ndarray]

Return the access_order and access_count arrays.

The access_order array numbers the accesses to each element of the tensor in walk order, -1 where the walk never goes; an element accessed more than once holds its last number. The access_count array holds the number of times the walk accesses each element.

Returns:

Type	Description
tuple[ndarray, ndarray]	tuple[np.ndarray, np.ndarray]: access_order and access_count, each of shape tensor_dims.

Raises:

Type	Description
TypeError	If the pattern is staged.

Source code in python/helpers/taplib/tap.py

def accesses(self) -> tuple[np.ndarray, np.ndarray]:
    """Return the access_order and access_count arrays.

    The access_order array numbers the accesses to each element of the
    tensor in walk order, -1 where the walk never goes; an element accessed
    more than once holds its last number. The access_count array holds the
    number of times the walk accesses each element.

    Returns:
        tuple[np.ndarray, np.ndarray]: access_order and access_count, each
            of shape `tensor_dims`.

    Raises:
        TypeError: If the pattern is staged.
    """
    idx = self._walk().reshape(-1)
    idx = idx[idx >= 0]
    numel = int(sprod(self._tensor_dims))
    order = np.full(numel, -1, dtype=self._DTYPE)
    np.maximum.at(order, idx, np.arange(idx.size, dtype=order.dtype))
    count = np.bincount(idx, minlength=numel).astype(self._DTYPE)
    return order.reshape(self._tensor_dims), count.reshape(self._tensor_dims)

access_order

access_order() -> ndarray

Return the access_order array of accesses().

Returns:

Type	Description
ndarray	np.ndarray: The walk number of each element's last access, -1 if unvisited.

Source code in python/helpers/taplib/tap.py

def access_order(self) -> np.ndarray:
    """Return the access_order array of `accesses()`.

    Returns:
        np.ndarray: The walk number of each element's last access, -1 if unvisited.
    """
    return self.accesses()[0]

access_count

access_count() -> ndarray

Return the access_count array of accesses().

Returns:

Type	Description
ndarray	np.ndarray: The number of times the walk accesses each element.

Source code in python/helpers/taplib/tap.py

def access_count(self) -> np.ndarray:
    """Return the access_count array of `accesses()`.

    Returns:
        np.ndarray: The number of times the walk accesses each element.
    """
    return self.accesses()[1]

compare_access_orders

compare_access_orders(other: TensorAccessPattern) -> bool

Return whether two patterns walk the same elements in the same order.

Patterns with different sizes and strides can still be functionally equivalent; this compares the walks themselves, padding included.

Parameters:

Name	Type	Description	Default
other	TensorAccessPattern	The pattern to compare to.	required

Returns:

Name	Type	Description
bool	bool	Whether the two walks visit the same elements in the same order.

Raises:

Type	Description
TypeError	If other is not a TensorAccessPattern.

Source code in python/helpers/taplib/tap.py

def compare_access_orders(self, other: TensorAccessPattern) -> bool:
    """Return whether two patterns walk the same elements in the same order.

    Patterns with different sizes and strides can still be functionally
    equivalent; this compares the walks themselves, padding included.

    Args:
        other (TensorAccessPattern): The pattern to compare to.

    Returns:
        bool: Whether the two walks visit the same elements in the same order.

    Raises:
        TypeError: If `other` is not a TensorAccessPattern.
    """
    if not isinstance(other, TensorAccessPattern):
        raise TypeError(
            "Can only compare access order against another TensorAccessPattern"
        )
    return np.array_equal(self._walk().reshape(-1), other._walk().reshape(-1))

visualize

visualize(
    show_arrows: bool | None = None,
    title: str | None = None,
    file_path: str | None = None,
    show_plot: bool = True,
    plot_access_count: bool = False,
) -> None

Plot the access order (and optionally count) of the walk over its tensor.

Parameters:

Name	Type	Description	Default
show_arrows	bool | None	Draw arrows between consecutively accessed elements. Defaults to None (only for small tensors).	None
title	str | None	Title of the plot. Defaults to the pattern's repr.	None
file_path	str | None	Path to save the plot to. Defaults to None.	None
show_plot	bool	Show the plot, e.g. in a Jupyter notebook. Defaults to True.	True
plot_access_count	bool	Plot the access count as well as the access order. Defaults to False.	False

Raises:

Type	Description
NotImplementedError	If the tensor has more than 2 dimensions.

Source code in python/helpers/taplib/tap.py

def visualize(
    self,
    show_arrows: bool | None = None,
    title: str | None = None,
    file_path: str | None = None,
    show_plot: bool = True,
    plot_access_count: bool = False,
) -> None:
    """Plot the access order (and optionally count) of the walk over its tensor.

    Args:
        show_arrows (bool | None, optional): Draw arrows between consecutively
            accessed elements. Defaults to None (only for small tensors).
        title (str | None, optional): Title of the plot. Defaults to the pattern's repr.
        file_path (str | None, optional): Path to save the plot to. Defaults to None.
        show_plot (bool, optional): Show the plot, e.g. in a Jupyter notebook.
            Defaults to True.
        plot_access_count (bool, optional): Plot the access count as well as
            the access order. Defaults to False.

    Raises:
        NotImplementedError: If the tensor has more than 2 dimensions.
    """
    from .visualization2d import visualize_from_accesses

    if len(self._tensor_dims) > 2:
        raise NotImplementedError(
            "Visualization is only currently supported for 1- or 2-dimensional tensors"
        )
    access_order, access_count = self.accesses()
    if title is None:
        title = str(self)
    visualize_from_accesses(
        np.atleast_2d(access_order),
        np.atleast_2d(access_count) if plot_access_count else None,
        title=title,
        show_arrows=show_arrows,
        file_path=file_path,
        show_plot=show_plot,
    )

animate

animate(
    frame_dims: int = 1,
    title: str | None = None,
    animate_access_count: bool = False,
) -> FuncAnimation

Animate the walk, one frame per index of its frame_dims outermost dimensions.

For a tiling full((M, N)).tile((m, n)), frame_dims=2 shows one tile per frame and frame_dims=1 one row of tiles.

Parameters:

Name	Type	Description	Default
frame_dims	int	Number of outermost dimensions the frames step through. Defaults to 1.	1
title	str | None	The title of the animation. Defaults to None.	None
animate_access_count	bool	Animate the access count as well as the access order. Defaults to False.	False

Returns:

Name	Type	Description
FuncAnimation	FuncAnimation	A handle to the animation.

Raises:

Type	Description
NotImplementedError	If the tensor has more than 2 dimensions.
ValueError	If frame_dims is not in [0, rank].

Source code in python/helpers/taplib/tap.py

def animate(
    self,
    frame_dims: int = 1,
    title: str | None = None,
    animate_access_count: bool = False,
) -> FuncAnimation:
    """Animate the walk, one frame per index of its `frame_dims` outermost dimensions.

    For a tiling `full((M, N)).tile((m, n))`, `frame_dims=2` shows one
    tile per frame and `frame_dims=1` one row of tiles.

    Args:
        frame_dims (int, optional): Number of outermost dimensions the
            frames step through. Defaults to 1.
        title (str | None, optional): The title of the animation. Defaults to None.
        animate_access_count (bool, optional): Animate the access count
            as well as the access order. Defaults to False.

    Returns:
        FuncAnimation: A handle to the animation.

    Raises:
        NotImplementedError: If the tensor has more than 2 dimensions.
        ValueError: If `frame_dims` is not in `[0, rank]`.
    """
    from .visualization2d import animate_from_accesses

    if not 0 <= frame_dims <= self.rank:
        raise ValueError(
            f"frame_dims must be in [0, {self.rank}], got {frame_dims}"
        )
    self._require_concrete("animate()")
    if len(self._tensor_dims) > 2:
        raise NotImplementedError(
            "Visualization is only currently supported for 1- or 2-dimensional tensors"
        )
    shape = (
        self._tensor_dims
        if len(self._tensor_dims) == 2
        else (1, *self._tensor_dims)
    )
    order_frames = [np.full(shape, -1, self._DTYPE)]
    count_frames = [np.zeros(shape, self._DTYPE)] if animate_access_count else None
    for index in np.ndindex(*self._sizes[:frame_dims]):
        order, count = self[index].accesses()
        order_frames.append(order.reshape(shape))
        if count_frames is not None:
            count_frames.append(count.reshape(shape))
    return animate_from_accesses(
        order_frames,
        count_frames,
        title="TensorAccessPattern Animation" if title is None else title,
    )

Utilities

Validation and stride helpers shared by the access-pattern algebra.

validate_and_clean_sizes_strides

validate_and_clean_sizes_strides(
    sizes: Sequence[IntLike], strides: Sequence[IntLike]
) -> tuple[list[IntLike], list[IntLike]]

Validate sizes and strides, and zero the strides of leading unit dimensions.

A check on a staged value becomes a dispatch-time guard.

Parameters:

Name	Type	Description	Default
sizes	Sequence[IntLike]	Extent of each dimension, outermost first.	required
strides	Sequence[IntLike]	Element step of each dimension, outermost first.	required

Returns:

Type	Description
tuple[list[IntLike], list[IntLike]]	tuple[list[IntLike], list[IntLike]]: The sizes and the cleaned strides.

Raises:

Type	Description
ValueError	If the lists are empty or differ in length, a size is below 1 or a stride below 0.

Source code in python/helpers/taplib/utils.py

def validate_and_clean_sizes_strides(
    sizes: Sequence[IntLike], strides: Sequence[IntLike]
) -> tuple[list[IntLike], list[IntLike]]:
    """Validate sizes and strides, and zero the strides of leading unit dimensions.

    A check on a staged value becomes a dispatch-time guard.

    Args:
        sizes (Sequence[IntLike]): Extent of each dimension, outermost first.
        strides (Sequence[IntLike]): Element step of each dimension, outermost first.

    Returns:
        tuple[list[IntLike], list[IntLike]]: The sizes and the cleaned strides.

    Raises:
        ValueError: If the lists are empty or differ in length, a size is
            below 1 or a stride below 0.
    """
    sizes, strides = list(sizes), list(strides)
    if not sizes:
        raise ValueError("len(sizes) must be >0")
    if len(strides) != len(sizes):
        raise ValueError(f"len(sizes) ({len(sizes)}) != len(strides) ({len(strides)})")
    for s in sizes:
        require(s >= 1, f"All sizes must be >= 1, but got {show(sizes)}")
    for s in strides:
        require(s >= 0, f"All strides must be >= 0, but got {show(strides)}")
    return sizes, zero_leading_unit_strides(sizes, strides)

zero_leading_unit_strides

zero_leading_unit_strides(
    sizes: Sequence, strides: Sequence, start: int = 0
) -> list

Zero the stride of each unit dimension from start up to the first that steps.

A unit dimension never steps, so this makes equal walks compare equal. The innermost stride is left as is. Rank and unit-ness are structural, so a staged size ends the scan.

Parameters:

Name	Type	Description	Default
sizes	Sequence	Extent of each dimension, outermost first.	required
strides	Sequence	Element step of each dimension, outermost first.	required
start	int	The first dimension to consider. Defaults to 0.	0

Returns:

Name	Type	Description
list	list	The strides, with those of the leading unit dimensions zeroed.

Source code in python/helpers/taplib/utils.py

def zero_leading_unit_strides(
    sizes: Sequence, strides: Sequence, start: int = 0
) -> list:
    """Zero the stride of each unit dimension from `start` up to the first that steps.

    A unit dimension never steps, so this makes equal walks compare equal.
    The innermost stride is left as is. Rank and unit-ness are structural,
    so a staged size ends the scan.

    Args:
        sizes (Sequence): Extent of each dimension, outermost first.
        strides (Sequence): Element step of each dimension, outermost first.
        start (int, optional): The first dimension to consider. Defaults to 0.

    Returns:
        list: The strides, with those of the leading unit dimensions zeroed.
    """
    strides = list(strides)
    for i in range(start, len(sizes) - 1):
        if is_sym(sizes[i]) or sizes[i] != 1:
            break
        strides[i] = 0
    return strides

validate_tensor_dims

validate_tensor_dims(
    tensor_dims: Sequence[IntLike],
) -> list[IntLike]

Check that a tensor has at least one dimension and every dimension is >= 1.

Parameters:

Name	Type	Description	Default
tensor_dims	Sequence[IntLike]	Tensor dimensions to check.	required

Returns:

Type	Description
list[IntLike]	list[IntLike]: The tensor dimensions.

Raises:

Type	Description
ValueError	If there are no dimensions or a dimension is below 1.

Source code in python/helpers/taplib/utils.py

def validate_tensor_dims(tensor_dims: Sequence[IntLike]) -> list[IntLike]:
    """Check that a tensor has at least one dimension and every dimension is >= 1.

    Args:
        tensor_dims (Sequence[IntLike]): Tensor dimensions to check.

    Returns:
        list[IntLike]: The tensor dimensions.

    Raises:
        ValueError: If there are no dimensions or a dimension is below 1.
    """
    tensor_dims = list(tensor_dims)
    if not tensor_dims:
        raise ValueError(
            f"Number of tensor dimensions must be >= 1 (dimensions={show(tensor_dims)})"
        )
    for d in tensor_dims:
        require(
            d >= 1,
            f"Each tensor dimension must be >= 1 (dimensions={show(tensor_dims)})",
        )
    return tensor_dims

validate_offset

validate_offset(
    offset: IntLike, tensor_dims: Sequence[IntLike]
) -> IntLike

Check that offset is an element index into a tensor of shape tensor_dims.

Parameters:

Name	Type	Description	Default
offset	IntLike	The offset to check.	required
tensor_dims	Sequence[IntLike]	Shape of the tensor.	required

Returns:

Name	Type	Description
IntLike	IntLike	The offset.

Raises:

Type	Description
ValueError	If the offset is negative or past the last element.

Source code in python/helpers/taplib/utils.py

def validate_offset(offset: IntLike, tensor_dims: Sequence[IntLike]) -> IntLike:
    """Check that `offset` is an element index into a tensor of shape `tensor_dims`.

    Args:
        offset (IntLike): The offset to check.
        tensor_dims (Sequence[IntLike]): Shape of the tensor.

    Returns:
        IntLike: The offset.

    Raises:
        ValueError: If the offset is negative or past the last element.
    """
    require(offset >= 0, f"Offset must be >= 0 (offset={show(offset)})")
    numel = sprod(tensor_dims)
    require(
        offset < numel,
        f"Offset too large: {show(offset)}. Max value allowed for tensor: {show(numel)}",
    )
    return offset

row_major_strides

row_major_strides(dims: Sequence) -> list

Row-major (C-order) element strides for a tensor of shape dims.

Parameters:

Name	Type	Description	Default
dims	Sequence	Tensor dimensions; entries may be staged values.	required

Returns:

Name	Type	Description
list	list	One stride per dimension, outermost first.

Source code in python/helpers/taplib/utils.py

def row_major_strides(dims: Sequence) -> list:
    """Row-major (C-order) element strides for a tensor of shape `dims`.

    Args:
        dims (Sequence): Tensor dimensions; entries may be staged values.

    Returns:
        list: One stride per dimension, outermost first.
    """
    strides: list = [1] * len(dims)
    for axis in range(len(dims) - 2, -1, -1):
        strides[axis] = strides[axis + 1] * dims[axis + 1]
    return strides

validate_permutation

validate_permutation(
    axes: Sequence[int], rank: int, what: str
) -> tuple[int, ...]

Check that axes is a permutation of range(rank).

Parameters:

Name	Type	Description	Default
axes	Sequence[int]	The permutation to check.	required
rank	int	Number of dimensions permuted.	required
what	str	Name of the argument, for the error message.	required

Returns:

Type	Description
tuple[int, ...]	tuple[int, ...]: The permutation as a tuple of ints.

Raises:

Type	Description
ValueError	If axes is not a permutation of range(rank).

Source code in python/helpers/taplib/utils.py

def validate_permutation(axes: Sequence[int], rank: int, what: str) -> tuple[int, ...]:
    """Check that `axes` is a permutation of `range(rank)`.

    Args:
        axes (Sequence[int]): The permutation to check.
        rank (int): Number of dimensions permuted.
        what (str): Name of the argument, for the error message.

    Returns:
        tuple[int, ...]: The permutation as a tuple of ints.

    Raises:
        ValueError: If `axes` is not a permutation of `range(rank)`.
    """
    axes = tuple(int(a) for a in axes)
    if sorted(axes) != list(range(rank)):
        raise ValueError(f"{what} must be a permutation of range({rank}), got {axes}")
    return axes

Instruction-stream tracing

Decode and semantically compare NPU TXN instruction streams.

A runtime sequence compiled once with DispatchTime scalars and a fully static specialization of the same sequence program the same DMA transfers, but their word streams are not byte-identical: the dynamic path draws buffer descriptors from a free-list pool (different BD ids), polls for room in a channel's task queue, assembles BD words at build time, and leaves the buffer-address word to the address patch. This module replays a stream into the register state it programs and reduces it to a list of events, the things the hardware actually acts on:

• push: a queue write on a channel, resolved to the chain of BDs it runs (for each, transfer length, address as (host argument, byte offset) or an absolute address, the addressing dimensions in a normalized form, and iteration), the repeat count, whether a task-completion token is issued, and the channel control bits.
• wait: a task-completion-token wait.
• write: a register write outside the DMA BD and channel registers (runtime parameters, locks, anything else the sequence programs).
• loadpdi / preempt / scratchpad / update_reg: kept verbatim.

Polls only wait for state another op produces, so they are not events. A push of a BD the stream never wrote is an error rather than a transfer of zeros: its contents come from somewhere the stream does not show.

Two streams are equivalent when their headers name the same device and their event lists are equal. compare reports the first divergence; explain prints the events for a human.

The decoding covers the AIE2 / AIE2p shim, memtile and core DMA register layouts used by npu1 and npu2 (register fields from aie-rt's xaiemlgbl_params.h).

Opcode

Bases: IntEnum

TXN opcodes, from include/aie/Runtime/TxnEncoding.h.

Header dataclass

Header(
    major: int,
    minor: int,
    dev_gen: int,
    num_rows: int,
    num_cols: int,
    num_mem_tile_rows: int,
    num_ops: int,
    size: int,
)

The 4-word TXN header.

Op dataclass

Op(
    kind: str,
    pos: int,
    words: tuple[int, ...],
    addr: int = 0,
    value: int = 0,
    mask: int = 0,
    data: tuple[int, ...] = (),
)

One decoded TXN instruction.

Transfer dataclass

Transfer(
    length: int,
    address: tuple,
    dims: tuple[tuple[int, int], ...],
    outer_stride: int,
    iteration: tuple[int, int],
    packet: tuple[int, ...],
    flags: tuple[int, ...],
    burst_axcache: tuple[int, int],
    padding: tuple[tuple[int, int], ...] = (),
)

A buffer descriptor as the DMA sees it.

The form is independent of how the stream encoded it.

decode

decode(words: Sequence[int] | ndarray) -> list[Op]

Split a TXN stream (header included) into instructions.

The header's size and op count must match the stream.

Source code in python/utils/txn_trace.py

def decode(words: Sequence[int] | np.ndarray) -> list[Op]:
    """Split a TXN stream (header included) into instructions.

    The header's size and op count must match the stream.
    """
    w = [int(x) & 0xFFFFFFFF for x in words]
    if len(w) < 4:
        raise ValueError("stream shorter than the 4-word TXN header")
    header = Header.of(w)
    if header.size != 4 * len(w):
        raise ValueError(
            f"TXN header gives {header.size} bytes, the stream has {4 * len(w)}"
        )
    ops: list[Op] = [Op("header", 0, tuple(w[:4]))]
    pos = 4
    n = len(w)
    while pos < n:
        opc = w[pos]
        base = opc & 0xFF
        if base == Opcode.WRITE:
            ops.append(
                Op(
                    "write32",
                    pos,
                    tuple(w[pos : pos + 6]),
                    addr=w[pos + 2],
                    value=w[pos + 4],
                )
            )
            pos += 6
        elif base in (Opcode.MASKWRITE, Opcode.MASKPOLL):
            ops.append(
                Op(
                    "maskwrite32" if base == Opcode.MASKWRITE else "maskpoll32",
                    pos,
                    tuple(w[pos : pos + 7]),
                    addr=w[pos + 2],
                    value=w[pos + 4],
                    mask=w[pos + 5],
                )
            )
            pos += 7
        elif base == Opcode.BLOCKWRITE:
            total = w[pos + 3] // 4
            if total < 4:
                raise ValueError(f"blockwrite at word {pos} shorter than its header")
            ops.append(
                Op(
                    "blockwrite",
                    pos,
                    tuple(w[pos : pos + total]),
                    addr=w[pos + 2],
                    data=tuple(w[pos + 4 : pos + total]),
                )
            )
            pos += total
        elif base == Opcode.TCT:
            total = w[pos + 1] // 4
            ops.append(Op("tct", pos, tuple(w[pos : pos + total])))
            pos += total
        elif base == Opcode.DDR_PATCH:
            total = w[pos + 1] // 4
            arg_plus = w[pos + 10] | (w[pos + 11] << 32)
            ops.append(
                Op(
                    "patch",
                    pos,
                    tuple(w[pos : pos + total]),
                    addr=w[pos + 6],
                    data=(w[pos + 8], arg_plus),
                )
            )
            pos += total
        elif base == Opcode.LOADPDI:
            ops.append(Op("loadpdi", pos, tuple(w[pos : pos + 4])))
            pos += 4
        elif base == Opcode.PREEMPT:
            ops.append(Op("preempt", pos, (opc,), value=opc >> 8))
            pos += 1
        elif base == Opcode.CREATE_SCRATCHPAD:
            # Words 2-3 are the DDR address the runtime patches in.
            ops.append(
                Op(
                    "scratchpad",
                    pos,
                    tuple(w[pos : pos + 4]),
                    value=w[pos + 1],
                    data=((opc >> 8) & 0xFF,),
                )
            )
            pos += 4
        elif base == Opcode.UPDATE_REG:
            ops.append(
                Op(
                    "update_reg",
                    pos,
                    tuple(w[pos : pos + 3]),
                    addr=w[pos + 2],
                    value=w[pos + 1],
                    data=((opc >> 8) & 0xFF, (opc >> 16) & 0xFF),
                )
            )
            pos += 3
        else:
            raise ValueError(f"unknown TXN opcode {opc:#x} at word {pos}")
    if pos != n:
        raise ValueError(f"the op at word {ops[-1].pos} runs past the stream")
    if header.num_ops != len(ops) - 1:
        raise ValueError(
            f"TXN header counts {header.num_ops} ops, the stream has {len(ops) - 1}"
        )
    return ops

trace

trace(words: Sequence[int] | ndarray) -> list[Event]

Replay a stream and return the events it triggers, in order.

Source code in python/utils/txn_trace.py

def trace(words: Sequence[int] | np.ndarray) -> list[Event]:
    """Replay a stream and return the events it triggers, in order."""
    ops = decode(words)
    header = Header.of(ops[0].words)
    regs: dict[int, int] = {}
    patches: dict[int, tuple] = {}  # BD address-word register -> ("arg", idx, plus)
    events: list[Event] = []

    def locate(addr: int) -> tuple[int, int, str, int]:
        row = (addr >> _ROW_SHIFT) & 0x1F
        return (addr >> _COL_SHIFT) & 0x7F, row, header.tile_kind(row), addr & _REG_MASK

    def queue(kind: str, reg: int) -> tuple[str, int] | None:
        """Return the (direction, channel) whose queue register `reg` is, if any."""
        layout = _LAYOUT[kind]
        for direction, delta in (("S2MM", 0), ("MM2S", layout.mm2s_delta)):
            for ch in range(layout.channels):
                if reg == layout.ctrl_base + delta + ch * 8 + 4:
                    return direction, ch
        return None

    def bd_chain(
        op: Op, col: int, row: int, kind: str, bd_id: int
    ) -> tuple[list[Transfer], int | None]:
        layout = _LAYOUT[kind]
        tile = (col << _COL_SHIFT) | (row << _ROW_SHIFT)
        transfers: list[Transfer] = []
        seen: dict[int, int] = {}
        while bd_id not in seen:
            seen[bd_id] = len(transfers)
            bd_base = tile | (layout.bd_base + bd_id * _BD_STRIDE)
            addrs = [bd_base + 4 * i for i in range(ctypes.sizeof(layout.bd) // 4)]
            # Only shim BDs address host memory, through the word-1 patch.
            patched = patches.get(bd_base + 4) if kind == "shim" else None
            missing = [
                i
                for i, a in enumerate(addrs)
                if a not in regs and not (patched and i == 1)
            ]
            if missing:
                raise ValueError(
                    f"the push at word {op.pos} runs BD {bd_id} of tile "
                    f"({col},{row}), but the stream never writes its words {missing}"
                )
            w = [regs.get(a, 0) for a in addrs]
            transfers.append(Transfer.from_words(kind, w, patched))
            bd = layout.bd.of(w)
            if not bd.use_next_bd:
                return transfers, None
            bd_id = bd.next_bd
        return transfers, seen[bd_id]

    def write(addr: int, value: int, mask: int) -> None:
        col, row, kind, reg = locate(addr)
        if not _LAYOUT[kind].owns(reg):
            events.append(Event("write", col=col, row=row, raw=(reg, value, mask)))

    for op in ops:
        if op.kind == "write32":
            regs[op.addr] = op.value
        elif op.kind == "maskwrite32":
            regs[op.addr] = (regs.get(op.addr, 0) & ~op.mask) | (op.value & op.mask)
            write(op.addr, op.value & op.mask, op.mask)
        elif op.kind == "blockwrite":
            # A fresh BD image supersedes an earlier patch of its address word.
            for i, d in enumerate(op.data):
                regs[op.addr + 4 * i] = d
                patches.pop(op.addr + 4 * i, None)
                write(op.addr + 4 * i, d, 0xFFFFFFFF)
        elif op.kind == "patch":
            patches[op.addr] = ("arg", op.data[0], op.data[1])
        elif op.kind == "tct":
            w2, w3 = op.words[2], op.words[3]
            direction = "S2MM" if (w2 & 0xFF) == 0 else "MM2S"
            # word 3: nrow [15:8], ncol [23:16], channel [31:24].
            events.append(
                Event(
                    "wait",
                    col=(w2 >> 16) & 0xFF,
                    row=(w2 >> 8) & 0xFF,
                    direction=direction,
                    channel=(w3 >> 24) & 0xFF,
                    raw=((w3 >> 16) & 0xFF, (w3 >> 8) & 0xFF),
                )
            )
        elif op.kind == "scratchpad":
            events.append(Event(op.kind, raw=(op.data[0], op.value)))
        elif op.kind == "update_reg":
            events.append(Event(op.kind, raw=(*op.data, op.value, op.addr)))
        elif op.kind in ("loadpdi", "preempt"):
            events.append(Event(op.kind, raw=op.words))
        if op.kind != "write32":
            continue
        col, row, kind, reg = locate(op.addr)
        hit = queue(kind, reg)
        if hit is None:
            write(op.addr, op.value, 0xFFFFFFFF)
            continue
        direction, ch = hit
        push = _LAYOUT[kind].queue_push.of([op.value])
        repeat = push.repeat_count
        (bd, *chain), loop = bd_chain(op, col, row, kind, push.start_bd_id)
        # A linear BD re-run repeat+1 times with an iteration dimension that
        # advances by its own length is one linear transfer that long; the
        # static emitter folds a contiguous repeat dimension into the length.
        iter_wrap, iter_stride = bd.iteration[0] + 1, bd.iteration[1]
        if (
            kind == "shim"
            and not chain
            and loop is None
            and not bd.dims
            and bd.outer_stride == 1
            and repeat
            and iter_wrap == repeat + 1
            and iter_stride == bd.length
        ):
            bd = dataclasses.replace(bd, length=bd.length * iter_wrap, iteration=(0, 1))
            repeat = 0
        events.append(
            Event(
                "push",
                col=col,
                row=row,
                direction=direction,
                channel=ch,
                bd=bd,
                chain=tuple(chain),
                chain_loop=loop,
                repeat=repeat,
                issue_token=bool(push.enable_token_issue),
                ctrl=regs.get(op.addr - 4, 0),
            )
        )
    return events

compare

compare(
    a: Sequence[int] | ndarray,
    b: Sequence[int] | ndarray,
    *,
    names=("a", "b")
) -> list[str]

Return the differences between two streams' devices and events.

Empty when the streams are equivalent.

Source code in python/utils/txn_trace.py

def compare(
    a: Sequence[int] | np.ndarray, b: Sequence[int] | np.ndarray, *, names=("a", "b")
) -> list[str]:
    """Return the differences between two streams' devices and events.

    Empty when the streams are equivalent.
    """
    ha, hb = Header.of(decode(a)[0].words), Header.of(decode(b)[0].words)
    if ha.device != hb.device:
        return [f"headers differ:\n  {names[0]}: {ha}\n  {names[1]}: {hb}"]
    ea, eb = trace(a), trace(b)
    out: list[str] = []
    for i, (x, y) in enumerate(zip(ea, eb)):
        if x != y:
            out.append(f"event {i} differs:\n  {names[0]}: {x}\n  {names[1]}: {y}")
            break
    if len(ea) != len(eb):
        out.append(f"{names[0]} has {len(ea)} events, {names[1]} has {len(eb)}")
    return out

explain

explain(
    words: Sequence[int] | ndarray, *, raw: bool = False
) -> str

Human-readable listing of a stream's events (or raw ops).

Source code in python/utils/txn_trace.py

def explain(words: Sequence[int] | np.ndarray, *, raw: bool = False) -> str:
    """Human-readable listing of a stream's events (or raw ops)."""
    if raw:
        return "\n".join(str(op) for op in decode(words))
    return "\n".join(f"{i:4d} {e}" for i, e in enumerate(trace(words)))

load

load(path) -> ndarray

Read an insts.bin / .txn file as a uint32 array.

Source code in python/utils/txn_trace.py

def load(path) -> np.ndarray:
    """Read an insts.bin / .txn file as a uint32 array."""
    return np.fromfile(path, dtype=np.uint32)