fix L2 shared-memory indexing under the ir_mul layout - #212
Open
asglover wants to merge 1 commit into
Open
Conversation
Under layout='ir_mul', load_ir_segments stages the second input into shared
memory in ir_mul order ([ir][mul]), but every L2 read in loop_unroll_tp.cuh
was hardcoded to mul_ir indexing:
l2_vec[j] = L2_smem[j + start + k * ir.dim]
L1 and L3 were unaffected because all of their shared-memory accesses go
through layout_load / layout_store, which branch on problem.layout. L2 had no
such branch, so the kernel read the block transposed.
The two orderings coincide when L2 has mul == 1 or a scalar irrep, which is why
this went unnoticed: spherical harmonics always have multiplicity 1, so no
MACE- or NequIP-shaped model hits it, and every existing ir_mul test used
L2 mul == 1. Anything with L2 mul > 1 and l >= 1 silently produced wrong
results in the forward pass, both input gradients, and the weight gradients.
Input validation accepted these problems; it only rejects uvw under ir_mul.
Adds l2_smem_index, mirroring layout_load/layout_store, and routes all five L2
accesses (forward uvu and uvw, backward, double-backward, and the L2 gradient
accumulation) through it.
Adds L2 multiplicities of 2, 3, 8 and 40 to the ir_mul suites in batch_test and
conv_test; 40 also crosses the 32-wide chunking threshold.
Introduced in df24066 (ir_mul layout support, #192).
Verified symbolically only -- no GPU available here. For each layout the
generated CUDA was parsed back into a shared-memory index map and checked
against the staging performed by load_ir_segments: 72/72 (mul, dim, layout)
configurations correct after the fix, and all 48 kernels for the new test
problems build. The tests themselves have not been executed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
vbharadwaj-bk
requested changes
Sep 4, 2026
vbharadwaj-bk
left a comment
vbharadwaj-bk
left a comment
Member
There was a problem hiding this comment.
Nice work - I like the macro - did you run the full test suite under this change?
| ), | ||
| oeq.TPProblem( | ||
| "32x1e", | ||
| "3x1e", |
Member
There was a problem hiding this comment.
Hmm - ok now there are enough copies of oeq.TPProblem that we might like to write this list compactly using a list comprehension.
| internal_weights=False, | ||
| label="ir_mul_repr_13x1x13_l535", | ||
| ), | ||
| oeq.TPProblem( |
Member
There was a problem hiding this comment.
Ditto, list comprehension here now that the list is long.
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.
Under layout='ir_mul', load_ir_segments stages the second input into shared memory in ir_mul order ([ir][mul]), but every L2 read in loop_unroll_tp.cuh was hardcoded to mul_ir indexing. L1 and L3 are unaffected because all of their shared-memory accesses go through layout_load / layout_store. L2 had no such branch, so the kernel read the block transposed.
Alternately, I could mark L2 > 1 as unsupported, instead of the fix.
Let me know your thoughts