Skip to content

repeat: split cols against every constraint the BD imposes, and add test coverage - #161

Open
atassis wants to merge 3 commits into
amd:develfrom
atassis:pr/repeat-tests
Open

repeat: split cols against every constraint the BD imposes, and add test coverage#161
atassis wants to merge 3 commits into
amd:develfrom
atassis:pr/repeat-tests

Conversation

@atassis

@atassis atassis commented Aug 26, 2026

Copy link
Copy Markdown
Contributor

Stacked on #157: the new tests gate at exact equality, which #157 is what makes
expressible. Its two commits are in this diff and drop out when it merges.

repeat has no tests, and could not have had: rows=8, cols=512, repeat=4 does not
build.

Both taps split cols as

sizes=[repeat, rows, cols // cols_split, cols_split]   strides=[..., cols_split, 1]

so the innermost dim counts chunks and the contiguous run sits one dim out. cols_split
is 1 whenever no split is needed, which makes the innermost dim one element. At bf16
that is two bytes, not a whole 32-bit word, and verifyStridesWraps rejects it:

'aie.dma_bd' op Transfer sizes must be multiples of 4 bytes. 1 elements at 2 bytes
each equal 2 bytes, which is not divisible by 4.
aie.dma_bd(%arg0 : memref<8x512xbf16> ... sizes = [4, 8, 512, 1] strides = [0, 512, 1, 1])

The 512-element stride-1 run the transfer is actually made of is right there in dim 2.

The reason this survived: llama_npu.py passes cols = prompt_len * head_dim, far above
1023, so it lands at cols_split = 256 and its innermost dim is a healthy 256 elements.
The only caller in the tree never trips it; every small-cols shape does, which is what a
test would use.

Added

  • iron/operators/repeat/test.py. Five device arms at rel_tol=abs_tol=0 (repeat moves
    data and computes nothing) spanning cols_split 1, 2, 4 and 256 -- the last being the
    shape llama_npu.py actually dispatches -- plus three arms pinning the shapes that have
    no legal split. 8/8.

Changed

  • iron/operators/repeat/design.py: swap the two inner dims so the innermost carries
    cols // cols_split at stride 1. The address sequence is byte-identical either way --
    the nested (outer, inner) walk of a linear range is the same walk regardless of which
    factor is outer -- so the change is only in which dim the hardware sees as innermost.

    The swap moves cols_split onto the next dim out, which carries the same 10-bit wrap
    field, so bounding the chunk length alone is not sufficient. The selection loop now
    searches for a divisor satisfying all three constraints at once: chunk <= 1023, chunk a
    whole number of 32-bit words, and count <= 1023. Where no divisor satisfies them -- an
    odd cols, a prime above 1023, or twice such a prime -- it raises rather than emitting a
    tap the BD verifier rejects less legibly. Every cols that worked before selects the
    same cols_split, including llama's 131072.

Removed

Evidence

8/8 on device, run without the not extensive filter so both extensive arms -- including
the llama shape -- actually executed. The emitted BD for the llama arm is
sizes = [4, 8, 256, 512] strides = [0, 131072, 512, 1].

The CSV reporting hook requires a [...] suffix on every nodeid and raises
otherwise. A test with no parameters has none, and the iteration parametrize
that would otherwise supply one is only added when --iterations > 1. On devel
today, `pytest iron/tests/stream/names.py --iterations 1` aborts the whole
session with an INTERNALERROR on
test_module_parameters_cover_the_golden_weights.
verify_buffer thresholds on `diff >= max(abs_tol, rel_tol * norm)`. At
rel_tol=abs_tol=0 the threshold is 0, so `diff >= 0` holds for every element and
a bit-identical buffer scores 100% errors. That makes exact equality
inexpressible, which is the gate an operator that does no arithmetic wants and
none can currently ask for. Compare strictly, matching what np.isclose means by
its tolerances: the new mask is a subset of the old one, so no result that
passes today can start failing.
The taps split cols as sizes [..., cols // cols_split, cols_split] with strides
[..., cols_split, 1], so the innermost dim counts chunks and the contiguous run sits
one dim out. In the unsplit case that is a 1-element innermost dim: at bf16 two bytes,
not a whole 32-bit word, and the BD verifier rejects it -- rows=8 cols=512 repeat=4
does not build at all.

Swapping the two inner dims gives a byte-identical address sequence (the nested
(outer, inner) walk of a linear range is the same walk either way) while making the
innermost dim the contiguous run.

The swap moves cols_split onto the next dim out, which carries the SAME 10-bit wrap
field, so bounding the chunk length alone is not enough -- and the old assert, which
bounded cols_split, was the only thing checking it. Rather than move the assert, the
selection loop now searches for a divisor satisfying all three constraints at once:
chunk <= 1023, chunk a whole number of 32-bit words, and count <= 1023. Where no
divisor satisfies them -- an odd cols, a prime above 1023, or 2 x such a prime -- it
raises instead of emitting a tap the verifier will reject less legibly. Every cols
that worked before selects the same cols_split, llama's 131072 included.

Adds the operator's first test coverage: five device arms at exact equality spanning
cols_split 1, 2, 4 and 256 (the llama shape), plus three rejection arms. 8/8.
@atassis atassis changed the title repeat: put the chunk length innermost, and add test coverage repeat: split cols against every constraint the BD imposes, and add test coverage Aug 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant