From 72573bf9f27d56f4dc159000e8a58a145265efb8 Mon Sep 17 00:00:00 2001 From: asglover <140220574+asglover@users.noreply.github.com> Date: Sun, 9 Aug 2026 22:10:06 -0400 Subject: [PATCH] WIP: fix L2 shared-memory indexing under the ir_mul layout 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) --- .../templates/loop_unroll_tp.cuh | 13 +++---- .../openequivariance/templates/macros.jinja | 8 +++++ tests/batch_test.py | 36 +++++++++++++++++++ tests/conv_test.py | 36 +++++++++++++++++++ 4 files changed, 87 insertions(+), 6 deletions(-) diff --git a/openequivariance/openequivariance/templates/loop_unroll_tp.cuh b/openequivariance/openequivariance/templates/loop_unroll_tp.cuh index 52b9219..ea94762 100644 --- a/openequivariance/openequivariance/templates/loop_unroll_tp.cuh +++ b/openequivariance/openequivariance/templates/loop_unroll_tp.cuh @@ -1,4 +1,5 @@ -{%- from 'macros.jinja' import layout_load, layout_store, reg_store with context %} +{%- from 'macros.jinja' import layout_load, layout_store, reg_store, + l2_smem_index with context %} {%- from 'wmm.cuh' import generate_matmul %} {%- macro generate_segment_kernel_forward(id, segment, warp_size) %} @@ -51,7 +52,7 @@ __device__ __forceinline__ void forward_loop_unroll_{{id}}(IRREP_T* __restrict__ #pragma unroll for(int j = 0; j < {{L2[v].ir.dim}}; j++) - l2_vec[j] = L2_smem[j + {{L2.slices()[v].start}} + k * {{L2[v].ir.dim}}] * weight; + l2_vec[j] = L2_smem[{{ l2_smem_index(problem.layout, L2[v].mul, L2[v].ir.dim, L2.slices()[v].start, 'k', 'j') }}] * weight; {%- elif problem.instructions[k].connection_mode == "uvw" %} {# Stream weights here #} {%- set slice_size = L3[w].mul * L1[u].mul %} @@ -61,7 +62,7 @@ __device__ __forceinline__ void forward_loop_unroll_{{id}}(IRREP_T* __restrict__ } #pragma unroll for(int j = 0; j < {{L2[v].ir.dim}}; j++) - l2_vec[j] = L2_smem[j + {{L2.slices()[v].start}} + k * {{L2[v].ir.dim}}]; + l2_vec[j] = L2_smem[{{ l2_smem_index(problem.layout, L2[v].mul, L2[v].ir.dim, L2.slices()[v].start, 'k', 'j') }}]; {%- endif %} // ----------------- CORE CALCULATION ----------------- @@ -184,11 +185,11 @@ __device__ __forceinline__ void forward_loop_unroll_{{id}}(IRREP_T* __restrict__ {%- if k == 0 or interactions[k][1] != interactions[k-1][1] or L2[v].mul > 1 or L1[u].mul != L1[interactions[k-1][0]].mul %} #pragma unroll for(int j = 0; j < {{L2[v].ir.dim}}; j++) { - l2_vec[j] = L2_smem[j + {{L2.slices()[v].start}} + k * {{L2[v].ir.dim}}]; + l2_vec[j] = L2_smem[{{ l2_smem_index(problem.layout, L2[v].mul, L2[v].ir.dim, L2.slices()[v].start, 'k', 'j') }}]; l2_grad[j] = 0.0; {%- if double_bwd %} - l2_original[j] = L2_original[j + {{L2.slices()[v].start}} + k * {{L2[v].ir.dim}}]; + l2_original[j] = L2_original[{{ l2_smem_index(problem.layout, L2[v].mul, L2[v].ir.dim, L2.slices()[v].start, 'k', 'j') }}]; {%- endif %} } {%- endif %} @@ -287,7 +288,7 @@ __device__ __forceinline__ void forward_loop_unroll_{{id}}(IRREP_T* __restrict__ if(lane_id == 0) { #pragma unroll for(int j = 0; j < {{L2[v].ir.dim}}; j++) - L2_grad_smem[j + {{L2.slices()[v].start}} + k * {{L2[v].ir.dim}}] += l2_grad[j]; + L2_grad_smem[{{ l2_smem_index(problem.layout, L2[v].mul, L2[v].ir.dim, L2.slices()[v].start, 'k', 'j') }}] += l2_grad[j]; } {%- endif %} diff --git a/openequivariance/openequivariance/templates/macros.jinja b/openequivariance/openequivariance/templates/macros.jinja index 59727e8..ce2a840 100644 --- a/openequivariance/openequivariance/templates/macros.jinja +++ b/openequivariance/openequivariance/templates/macros.jinja @@ -66,6 +66,14 @@ Keys map to lists of tuples with (name, dtype, num_elements) of each subarray. {%- endif %} {%- endmacro %} +{%- macro l2_smem_index(layout, mul, dim, start, mul_var, dim_var) -%} + {%- if layout == "ir_mul" -%} + {{mul_var}} + {{start}} + {{dim_var}} * {{mul}} + {%- else -%} + {{dim_var}} + {{start}} + {{mul_var}} * {{dim}} + {%- endif -%} +{%- endmacro %} + {%- macro declare_smem_variables(segment, smem_base) %} {%- for name in segment.smem %} {%- if name != "total" %} diff --git a/tests/batch_test.py b/tests/batch_test.py index 7ec6333..6b631b8 100644 --- a/tests/batch_test.py +++ b/tests/batch_test.py @@ -330,6 +330,42 @@ class TestIrMul(TPCorrectness): internal_weights=False, label="ir_mul_repr_13x1x13_l535", ), + oeq.TPProblem( + "32x1e", + "3x1e", + "32x1e", + [(0, 0, 0, "uvu", True)], + shared_weights=False, + internal_weights=False, + label="ir_mul_L2mul3_l111", + ), + oeq.TPProblem( + "32x1e", + "2x2e", + "32x1e", + [(0, 0, 0, "uvu", True)], + shared_weights=False, + internal_weights=False, + label="ir_mul_L2mul2_l121", + ), + oeq.TPProblem( + "16x2e", + "8x2e", + "16x2e", + [(0, 0, 0, "uvu", True)], + shared_weights=False, + internal_weights=False, + label="ir_mul_L2mul8_l222", + ), + oeq.TPProblem( + "16x1e", + "40x1e", + "16x1e", + [(0, 0, 0, "uvu", True)], + shared_weights=False, + internal_weights=False, + label="ir_mul_L2mul40_l111", + ), ] @pytest.fixture(params=tpps, ids=lambda x: x.label, scope="class") diff --git a/tests/conv_test.py b/tests/conv_test.py index 446d0f3..f8dec95 100644 --- a/tests/conv_test.py +++ b/tests/conv_test.py @@ -345,6 +345,42 @@ class TestIrMulLayout(ConvCorrectness): internal_weights=False, label="ir_mul_repr_13x1x13_l535", ), + oeq.TPProblem( + "32x1e", + "3x1e", + "32x1e", + [(0, 0, 0, "uvu", True)], + shared_weights=False, + internal_weights=False, + label="ir_mul_L2mul3_l111", + ), + oeq.TPProblem( + "32x1e", + "2x2e", + "32x1e", + [(0, 0, 0, "uvu", True)], + shared_weights=False, + internal_weights=False, + label="ir_mul_L2mul2_l121", + ), + oeq.TPProblem( + "16x2e", + "8x2e", + "16x2e", + [(0, 0, 0, "uvu", True)], + shared_weights=False, + internal_weights=False, + label="ir_mul_L2mul8_l222", + ), + oeq.TPProblem( + "16x1e", + "40x1e", + "16x1e", + [(0, 0, 0, "uvu", True)], + shared_weights=False, + internal_weights=False, + label="ir_mul_L2mul40_l111", + ), ] @pytest.fixture(params=production_model_tpps, ids=lambda x: x.label, scope="class")