gemm: check tile sizes against the kernel's divisibility rule, not a bound - #189
Merged
Merged
Conversation
andrej
requested changes
Sep 9, 2026
Comment on lines
+64
to
+70
| # r, s, t are the aie::mmul tile dims the bf16 kernel is built from | ||
| # (aie_kernels/aie2p/mm.cc, matmul_vectorized_2x2_mmul): it expands A | ||
| # and B 2x in m and n, so the static_asserts there require | ||
| # m % (2*r) == 0, k % s == 0, n % (2*t) == 0 -- a divisibility rule, | ||
| # not a lower bound. A tile_m/tile_n that merely meets the old ">=" | ||
| # check (e.g. 8 or 12) passes here and then fails that static_assert | ||
| # at kernel compile time, in a file this class never names. |
Collaborator
There was a problem hiding this comment.
Suggested change
| # r, s, t are the aie::mmul tile dims the bf16 kernel is built from | |
| # (aie_kernels/aie2p/mm.cc, matmul_vectorized_2x2_mmul): it expands A | |
| # and B 2x in m and n, so the static_asserts there require | |
| # m % (2*r) == 0, k % s == 0, n % (2*t) == 0 -- a divisibility rule, | |
| # not a lower bound. A tile_m/tile_n that merely meets the old ">=" | |
| # check (e.g. 8 or 12) passes here and then fails that static_assert | |
| # at kernel compile time, in a file this class never names. | |
| # r, s, t are the aie::mmul tile dims the bf16 kernel is built from | |
| # (aie_kernels/aie2p/mm.cc, matmul_vectorized_2x2_mmul) |
| tile_k is genuinely undoubled (k % s == 0, s=8 in both regimes), so its old | ||
| `>= 8` bound was already correct; it is included here for the same divisor, | ||
| expressed the same way as tile_m/tile_n. | ||
| """ |
Collaborator
There was a problem hiding this comment.
Can we please stop documenting code changes in comments?
|
|
||
| def test_tile_k_not_a_multiple_of_8_is_rejected(): | ||
| # tile_k=4 still divides K=512 evenly (the outer K % tile_k check), so | ||
| # this exercises the new s-divisibility check rather than that one. |
…ones __post_init__ checked tile_m/tile_k/tile_n >= a minimum, but aie_kernels/aie2p/mm.cc static_asserts m % (2*r) == 0 and n % (2*t) == 0 for the (r,t) mmul shape emulate_bf16_mmul_with_bfp16 selects -- a divisibility rule, not a lower bound. tile_m=8 under the default emulation passed the old check and failed that static_assert at kernel compile time. Check the real rule and name the offending value.
Add the tile_m=4 (no-emulation) and tile_n-under-emulation cases so both tile_m and tile_n are exercised in both emulate_bf16_mmul_with_bfp16 settings, not just the default.
atassis
force-pushed
the
fix/gemm-tile-divisibility-check
branch
from
September 11, 2026 17:14
e733bf7 to
a3d5181
Compare
andrej
approved these changes
Sep 11, 2026
andrej
enabled auto-merge
September 11, 2026 21:33
Contributor
CI Test Results8fd5fe0 (2026_09_11_22_17_26) IRON - CI SummaryExamplesiron/applications/llama_3.2_1b
Smalliron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/leaky_relu
iron/operators/mem_copy
iron/operators/mha
iron/operators/relu
iron/operators/repeat
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/strided_copy
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Krackan - SmallIRONTested on iron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/leaky_relu
iron/operators/mem_copy
iron/operators/mha
iron/operators/relu
iron/operators/repeat
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/strided_copy
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Krackan - ExamplesIRONTested on iron/applications/llama_3.2_1b
Phoenix - SmallIRONTested on iron/operators/axpy
iron/operators/dequant
iron/operators/elementwise_add
iron/operators/elementwise_mul
iron/operators/gelu
iron/operators/gemm
iron/operators/gemv
iron/operators/layer_norm
iron/operators/leaky_relu
iron/operators/mem_copy
iron/operators/relu
iron/operators/repeat
iron/operators/rms_norm
iron/operators/rope
iron/operators/sigmoid
iron/operators/silu
iron/operators/softmax
iron/operators/strided_copy
iron/operators/swiglu_decode
iron/operators/swiglu_prefill
iron/operators/tanh
iron/operators/transpose
Phoenix - ExamplesIRONTested on Trend tables omitted, the comment hit GitHub's size limit. Full report in the workflow run. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
GEMM.__post_init__checkstile_m,tile_k,tile_nagainst a minimum with>=, butaie_kernels/aie2p/mm.cc'smatmul_vectorized_*_bf16_*kernelsstatic_assert(m % (2*r) == 0)andstatic_assert(n % (2*t) == 0)for the(r, t)pairemulate_bf16_mmul_with_bfp16selects (r=t=8when enabled, the default;r=4, t=8otherwise) -- a divisibility rule, not a lower bound. Bothtile_mandtile_nare affected, in both regimes:tile_m=8passes the old check under the default emulation but needs a multiple of 16;tile_m=4passes it without emulation but needs a multiple of 8.GEMM(tile_m=8, emulate_bf16_mmul_with_bfp16=True, ...)constructs without error and fails the static_assert at kernel compile time, in a file this class never names.tile_kis the one dimensionmatmul_vectorized_2x2_mmuldoes not expand, so its rule isk % s == 0withs=8in both regimes. That is still a divisibility rule rather than a bound, so this does tightentile_kas well:tile_k=12passed the old>= 8check and is now rejected, correctly --mm.ccwould have refused it.The existing test suite doesn't cover this:
test.pyhardcodesemulate_bf16_mmul_with_bfp16=False, and its onetile_m=8case (extensive_params,M=2048,K=2048,N=2048,m=8,k=16,n=32) only survives because2*r == 8in that branch -- the default-emulation boundary was untested.Added
iron/tests/operators/gemm_tile_divisibility.py: construction-only tests (no device, no MLIR generation) covering bothtile_mandtile_nunder bothemulate_bf16_mmul_with_bfp16settings, plustile_k.Changed
iron/operators/gemm/op.py:__post_init__now checkstile_m % (2*r) == 0,tile_k % s == 0,tile_n % (2*t) == 0instead of>=, and the error message names the required multiple and the source file.Removed
Known gap
design.py's own guard (if not use_scalar: assert m % r == 0; assert n % t == 0) is a second, independently-too-weak check -- it's missing the same2*factor, anduse_scalardoesn't gate it out: theextern "C"block always emits both the scalar and vectorized symbols (combos(matmul_vectorized_c_func) combos(matmul_scalar_c_func) ...), so the vectorized static_assert fires at compile time even for a scalar-only build. Not fixed here sinceGEMM(this PR's scope) always callsdesign.pywith already-validated tiles; a directdesign.pyCLI/API caller is still exposed.Evidence
Ran with
--noconftest: this repo's rootconftest.pyopens the NPU unconditionally at collection time (see the companion PR fixing that), which this change has nothing to do with, so it's bypassed here rather than exercised.Confirmed the defect against unmodified
devel(deb6e1e7) before writing the fix:GEMM(M=512, K=512, N=512, tile_m=8, tile_k=64, tile_n=64, emulate_bf16_mmul_with_bfp16=True, context=None)constructs with no error.PR Merge Checklist
develcommit and pointing todevel.