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")