ezmsg.sigproc.util.blockdiag#

Find block-diagonal structure in a weight matrix, then decide whether to use it.

For y = x @ W with W of shape (n_in, n_out): when W is block-diagonal the product decomposes into independent per-block matmuls that touch only the weights inside the blocks. That is fewer FLOPs, but it is not automatically faster — each block costs a separate kernel launch, and below a few hundred channels a dense matmul against an L2-resident weight matrix wins outright (see ezmsg-org/ezmsg-sigproc#210).

So this module separates two questions:

  1. What structure does W have? contiguous_block_partition() answers it with numpy alone, in one pass over the nonzero mask. Contiguous blocks are the case worth optimizing: they slice into views, so there is no gather, no scatter, and the output can be filled in place.

  2. Is that structure worth exploiting? plan_block_matmul() answers it with the cost model below, and returns None for “just do a dense matmul”.

Structure is always read off W itself, never taken on a caller’s word — a hint that disagrees with the weights used to silently compute the wrong answer (ezmsg-org/ezmsg-sigproc#198).

Cost model#

Runtime of one matmul formulation is modelled as:

(n_samples + WEIGHT_LOAD_SAMPLES) * (weight elements touched)
    + (number of kernel calls) * CALL_COST_MACS
    + (elements gathered) * GATHER_COST_MACS

in units of multiply-accumulates. The first term charges both the arithmetic (n_samples MACs per weight) and the one-off cost of streaming the weights into cache (worth about WEIGHT_LOAD_SAMPLES samples of arithmetic at typical FLOPs-per-byte ratios). The constants are ratios, so only their relative size matters; they were fit to benchmarks/benchmark_affine_kernels.py on an Apple M-series CPU with float32 data, and the decisions they drive are insensitive to moderate error — the model only has to get the ordering right, and the formulations are within ~2x of each other near every crossover.

Known limitation: the model counts weights touched, not access patterns, so it does not see that a many-block loop re-walks a strided view of a chunk too large to cache. Above ~2000 channels with several thousand samples per chunk it picks a finer blocking than optimal, costing up to ~40% against the best merge (still several times faster than dense). Pass kernel="dense" if that combination is your hot path.

Module Attributes

WEIGHT_LOAD_SAMPLES

Streaming the weight matrix once costs about this many samples of arithmetic.

CALL_COST_MACS

Cost of one extra matmul call, in MACs, for eager CPU backends (numpy).

CALL_COST_MACS_DISPATCHED

Ditto for backends with heavier per-op dispatch (MLX, torch, cupy).

GATHER_COST_MACS

Cost of moving one element through a fancy-index gather, in MACs.

PERMUTED_SEARCH_MIN_WEIGHTS

Skip the connected-components search (and its scipy import) for matrices smaller than this.

Functions

contiguous_block_partition(weights)[source]#

Split weights into the finest tiling of contiguous diagonal blocks.

Every nonzero of weights lies inside one returned (rows, cols) block, and the blocks tile the full row and column ranges. A single block spanning everything means “no exploitable contiguous structure”.

All-zero rows and columns are absorbed into a neighbouring block rather than dropped, which keeps the tiling gap-free: an omitted output column would otherwise have to be zero-filled separately.

Parameters:

weights (ndarray)

Return type:

list[tuple[slice, slice]]

plan_block_matmul(weights, n_samples, *, force=False, dispatched=False)[source]#

Choose between a dense matmul and a block-diagonal one.

Parameters:

  • weights (ndarray) – 2-D weight matrix in (n_in, n_out) orientation.

  • n_samples (int) – Representative number of samples per message (everything on the message except the channel axis). Feeds the cost model; short chunks favour the dense kernel because they cannot amortize the per-block call overhead.

  • force (bool) – Return a block plan whenever any structure exists, ignoring the cost model. For tests and benchmarks.

  • dispatched (bool) – Set for backends with heavier per-op overhead than numpy.

Returns:

A BlockPlan, or None meaning “use a dense matmul”.

Return type:

BlockPlan | None

Classes

class BlockPlan(blocks, in_perm=None, out_perm=None)[source]#

Bases: object

How to evaluate x @ W as a sequence of per-block matmuls.

Each entry of blocks pairs a contiguous input slice with the contiguous output slice it writes. The slices tile 0..n_in and 0..n_out exactly, so the output buffer needs no zero-fill.

When the blocks are only contiguous after reordering channels, the permutations say how: weights[in_perm][:, out_perm] is what the blocks tile. At runtime that means gathering the input by in_perm and undoing out_perm on the result.

Parameters:
blocks: tuple[tuple[slice, slice], ...]#
in_perm: ndarray | None = None#
out_perm: ndarray | None = None#
property n_blocks: int#
__init__(blocks, in_perm=None, out_perm=None)#
Parameters:
Return type:

None