Legalise over-wide immediates so geometries at MLEN 512 compile - #75
Open
Shreyas8612 wants to merge 1 commit into
Open
Legalise over-wide immediates so geometries at MLEN 512 compile#75Shreyas8612 wants to merge 1 commit into
Shreyas8612 wants to merge 1 commit into
Conversation
S_ADDI_INT encodes an 18-bit immediate, and mlen * mlen is exactly 2^18 at MLEN 512. Around 30 emitter sites across ffn, projection, gemv and batched matmul templates pass that value as a raw immediate rather than going through the _imm helpers, so the assembler correctly refuses the program and no geometry at or above MLEN 512 can be compiled. Adds a pass over emitted assembly that rewrites over-wide immediates through S_LUI_INT. A gp0 source becomes a wide load, a non-aliasing destination becomes a wide load into the destination plus one add, and only an aliasing destination falls back to a chunked relative add. None of the three needs a scratch register, so the pass is safe to run after register allocation. The chunked fallback is capped so a pathological immediate fails with a clear error rather than emitting an unbounded instruction sequence.
There was a problem hiding this comment.
Pull request overview
This PR addresses a compiler/assembler incompatibility at MLEN=512 where S_ADDI_INT’s 18-bit immediate field overflows (e.g. mlen * mlen == 2^18), preventing compilation of larger geometries. The fix adds a post-emission legalization pass that rewrites over-wide S_ADDI_INT immediates into sequences using S_LUI_INT (and bounded fallbacks) without requiring scratch registers, making it safe after register allocation.
Changes:
- Wrap
PlenaCompiler.compile()output with a text-levellegalize_immediates()rewrite pass. - Extend
_imm.pyto support legalization via regex matching plus a bounded chunked fallback for aliasing cases. - Update/add unit tests around large-immediate behavior and chunk-limit failure modes.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
aten/plena/compiler.py |
Runs emitted assembly through legalize_immediates() in compile() so over-wide S_ADDI_INT no longer blocks compilation. |
asm_templates/_imm.py |
Adds the textual legalization pass and refines large-immediate add behavior (dest-as-temp when non-aliasing; capped chunk fallback when aliasing). |
asm_templates/tests/test_large_immediate.py |
Updates large-immediate tests and adds coverage for new behaviors (aliasing vs non-aliasing, chunk limits, template register expectations). |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+111
to
123
| def test_large_add_aliasing_chunk_limit(self): | ||
| """A pathological aliasing immediate fails loudly instead of flooding.""" | ||
| from asm_templates._imm import CHUNK_LIMIT, IMM2_BOUND | ||
|
|
||
| over_limit = (IMM2_BOUND - 1) * CHUNK_LIMIT + 1 | ||
| with self.assertRaises(ValueError): | ||
| _add_large_int(5, 5, over_limit, temp_reg=None) | ||
| # exactly at the limit still succeeds | ||
| at_limit = (IMM2_BOUND - 1) * CHUNK_LIMIT | ||
| result = _add_large_int(5, 5, at_limit, temp_reg=None) | ||
| self.assertEqual(len(result), CHUNK_LIMIT) | ||
|
|
||
|
|
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
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.
Bug: Over-wide immediates block compilation at MLEN 512
S_ADDI_INTencodes an 18-bit immediate, andmlen * mlenis exactly ( 2^{18} ) atMLEN512. Around 30 emitter sites across the ffn, projection, gemv, and batched matmul templates pass that value as a raw immediate rather than going through the_immhelpers, so the assembler correctly refuses the program, and no geometry at or aboveMLEN512 can be compiled.Fix
Adds a pass over emitted assembly that rewrites over-wide immediates through
S_LUI_INT:gp0source becomes a wide load.None of the three cases needs a scratch register, so the pass is safe to run after register allocation.
The chunked fallback is capped so a pathological immediate fails with a clear error rather than emitting an unbounded instruction sequence.