Skip to content

Commit 8ac2022

Browse files
spcypptmeta-codesync[bot]
spcyppt
authored and
meta-codesync[bot]
committed
Add PROCESS_BLOCK macro for grad accumulation loop unrolling (#5835)
Summary: Pull Request resolved: #5835 X-link: https://github.com/facebookresearch/FBGEMM/pull/2758 On ROCm, replace the single-element gradient accumulation loop in `compute_grad_sum` with a cascading block-size unrolling scheme (`PROCESS_BLOCK` with block sizes 8/4/2/1) that processes multiple segment elements per iteration. Each block manually unrolls the per-element `SHFL_SYNC` metadata broadcasts and `Vec4TAcc` gradient loads, so several independent `grad_output` loads are issued before accumulation. This increases instruction-level parallelism, hides global-memory latency across the block, and reduces loop/branch overhead. The 8-wide block is gated on `sizeof(grad_t) <= 2` (fp16) to bound register pressure. `compute_grad_sum` lives in the shared `embedding_backward_split_device_kernel_template.cuh`, which is `#include`d by both the CTA and warp backward kernels, so both benefit. The CTA kernel processes long, high-collision segments (many loop iterations to amortize the unrolling over) and gains the most; the warp kernel processes short segments and sees a smaller but consistent gain. The NVIDIA code path is unchanged. The helper macros are gated on `not gen_once` (they are only used by `compute_grad_sum`) and `#undef`'d after use, so they never leak into downstream translation units. Benchmarks on MI350 (fp32 weights/output, `rowwise_adagrad`, no flag), backward kernel time in us, positive = faster. ## TBE sweep over T, D, B Backward CTA kernel: ``` T D B baseline optimized speedup 1 8 2048 103 60 +42.4% 1 8 4096 201 113 +43.8% 1 8 131072 527 299 +43.3% 1 128 2048 844 67 +92.1% 1 128 4096 1195 124 +89.6% 1 128 131072 697 402 +42.4% 1 256 2048 n/a 67 (*) 1 256 4096 259 126 +51.3% 1 256 131072 834 725 +13.1% 1 512 2048 339 76 +77.5% 1 512 4096 375 163 +56.6% 1 512 131072 1435 1427 +0.6% 1 1024 2048 398 172 +56.9% 1 1024 4096 828 341 +58.9% 1 1024 131072 2942 3092 -5.1% 10 8 2048 367 71 +80.6% 10 8 4096 1466 145 +90.1% 10 8 131072 4835 2404 +50.3% 10 128 2048 605 89 +85.2% 10 128 4096 5(*) 184 (*) 10 128 131072 6323 3694 +41.6% 10 256 2048 207 106 +48.9% 10 256 4096 431 228 +47.0% 10 256 131072 8241 7047 +14.5% 10 512 2048 252 178 +29.2% 10 512 4096 540 383 +29.0% 10 512 131072 17449 14523 +16.8% ``` Backward warp kernel: ``` T D B baseline optimized speedup 1 8 2048 48 43 +10.4% 1 8 4096 75 70 +6.0% 1 8 131072 1353 1196 +11.6% 1 128 2048 544 47 +91.4% 1 128 4096 284 78 +72.6% 1 128 131072 1457 1398 +4.0% 1 256 2048 n/a 49 (*) 1 256 4096 95 86 +9.2% 1 256 131072 1772 1688 +4.7% 1 512 2048 76 65 +14.3% 1 512 4096 134 119 +11.1% 1 512 131072 2462 2353 +4.4% 1 1024 2048 146 119 +18.5% 1 1024 4096 254 219 +14.0% 1 1024 131072 4412 4242 +3.9% 10 8 2048 899 349 +61.2% 10 8 4096 696 666 +4.3% 10 8 131072 12261 11890 +3.0% 10 128 2048 4461 423 +90.5% 10 128 4096 2069 773 +62.7% 10 128 131072 15498 15427 +0.5% 10 256 2048 513 492 +4.1% 10 256 4096 921 892 +3.1% 10 256 131072 23360 17836 +23.6% 10 512 2048 686 648 +5.5% 10 512 4096 1236 1183 +4.3% 10 512 131072 25658 29404 -14.6% ``` ## bench_list (representative production mixed-D configs) Backward CTA kernel: ``` bench T Ds baseline optimized speedup bench_103 2 8 18 15 +18.1% bench_0 10 128 488 140 +71.3% bench_2 14 12/24/128 256 138 +46.3% bench_1 18 20/128 400 142 +64.6% bench_27 22 20/24/128 311 154 +50.4% TOTAL 1474 588 +60.1% ``` Backward warp kernel: ``` bench T Ds baseline optimized speedup bench_103 2 8 57 51 +9.8% bench_0 10 128 364 252 +30.6% bench_2 14 12/24/128 236 237 -0.5% bench_1 18 20/128 173 167 +3.5% bench_27 22 20/24/128 276 270 +2.2% TOTAL 1106 978 +11.6% ``` ## VBE (variable batch size, by alpha skew) ``` kernel a baseline optimized speedup CTA 1 4 4 -0.1% CTA 1.15 299 130 +56.6% warp 1 126 129 -2.4% warp 1.15 65 58 +11.5% ``` (*) Two CTA sweep baseline cells are measurement artifacts on this run and have no meaningful speedup: `T=1 D=256 B=2048` has no baseline sample, and `T=10 D=128 B=4096` recorded a spurious 5us baseline. The few negative outliers (e.g. CTA `T=1 D=1024 B=131072`, warp `T=10 D=512 B=131072`, VBE `a=1`) are within the cross-run measurement noise observed on this setup; `a=1` VBE is ~4us where noise dominates. CTA gains are largest for high-collision configs (long segments); warp gains are smaller but broadly positive. Reviewed By: q10 Differential Revision: D107551685 fbshipit-source-id: d46d46a3c1357ca662d1baf91d648ba024d4da50
1 parent a48a374 commit 8ac2022

1 file changed

Lines changed: 136 additions & 1 deletion

File tree

‎fbgemm_gpu/codegen/training/backward/embedding_backward_split_device_kernel_template.cuh‎

Lines changed: 136 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,103 @@
1414

1515
using namespace fbgemm_gpu;
1616

17+
{#- /* These helper macros are only used by `compute_grad_sum` below, which is
18+
rendered multiple times based on TBE config variants.
19+
`#undef` block is set later after the calls. */ #}
20+
{%- if is_rocm and not gen_once %}
21+
// Helper macro: Generate block_size grad_offset_j_i variables (i from 1 to block_size-1)
22+
#define GRAD_OFFSET(i, j) const auto grad_offset_j_##i = SHFL_SYNC(grad_offset, j + i);
23+
#define L(i, j) int32_t l_j_##i = SHFL_SYNC(l, j + i);
24+
#define B(i, j) int32_t b_j_##i = SHFL_SYNC(b, j + i);
25+
#define D_START(i, j) int32_t D_start_j_##i = SHFL_SYNC(D_start, j + i);
26+
#define IDX_WEIGHT(i, j) at::acc_type<cache_t, true> idx_weight_j_##i = SHFL_SYNC(idx_weight, j + i);
27+
28+
#define REPEAT_8(X, j) X(1, j); X(2, j); X(3, j); X(4, j); X(5, j); X(6, j); X(7, j);
29+
#define REPEAT_4(X, j) X(1, j); X(2, j); X(3, j);
30+
#define REPEAT_2(X, j) X(1, j);
31+
#define REPEAT_1(X, j) // No additional variables needed for block size 1
32+
33+
#define REPEAT_I_S_8(X, j, m, n) X(1, j, m, n); X(2, j, m, n); X(3, j, m, n); X(4, j, m, n); X(5, j, m, n); X(6, j, m, n); X(7, j, m, n);
34+
#define REPEAT_I_S_4(X, j, m, n) X(1, j, m, n); X(2, j, m, n); X(3, j, m, n);
35+
#define REPEAT_I_S_2(X, j, m, n) X(1, j, m, n);
36+
#define REPEAT_I_S_1(X, j, m, n) // No additional variables needed for block size 1
37+
38+
// Helper macro: Generate block_size Vec4TAcc objects (i from 1 to block_size-1)
39+
// if nobag and is_index_select
40+
#define GRAD_VEC_N_I(i, grad_offset, grad_stride, d) Vec4TAcc<grad_t> grad_out_vec_##i(&grad_output[grad_offset + l_j_##i * grad_stride + d]);
41+
// elif nobag
42+
#define GRAD_VEC_N(i, d) Vec4TAcc<grad_t> grad_out_vec_##i(&grad_output[l_j_##i][d]);
43+
// elif vbe
44+
#define GRAD_VEC_V(i, d) Vec4TAcc<grad_t> grad_out_vec_##i(&grad_output[0][grad_offset_j_##i + d]);
45+
// else
46+
#define GRAD_VEC(i, d) Vec4TAcc<grad_t> grad_out_vec_##i(&grad_output[b_j_##i][0] + D_start_j_##i + d);
47+
48+
// Helper macro: Generate block_size fma_ calls (i from 1 to block_size-1)
49+
#define FMA_GRAD(i, vec) grad_sum[vec].fma_(grad_out_vec_##i, idx_weight_j_##i);
50+
// Helper macro: Generate block_size add_ calls (i from 1 to block_size-1)
51+
#define ADD_GRAD(i, vec) grad_sum[vec].add_(grad_out_vec_##i);
52+
53+
// Core macro: Process blocks of specified size (block_size = 8/4/2/1)
54+
// Parameters:
55+
// - block_size: Size of each block to process
56+
// - unroll_count: Number of unroll iterations for the inner loop
57+
#define PROCESS_BLOCK(block_size, unroll_count, grad_sum, grad_output, grad_offset, vec_start, kThreadGroupSize, threadIdx_x, VEC_WIDTH, D, j, sl, sl_end) \
58+
for (; j + (block_size - 1) < kThreadGroupSize && sl + j + (block_size - 1) < sl_end; j += block_size) { \
59+
{%- if nobag %}
60+
int32_t l_j_0 = SHFL_SYNC(l, j); \
61+
REPEAT_##block_size(L, j) \
62+
{%- elif vbe %}
63+
/* Generate block_size grad_offset_j_0 ~ grad_offset_j_(block_size-1) */ \
64+
const auto grad_offset_j_0 = SHFL_SYNC(grad_offset, j); \
65+
/* Generate subsequent grad_offset_j_1 ~ grad_offset_j_(block_size-1) based on block size */ \
66+
REPEAT_##block_size(GRAD_OFFSET, j) \
67+
{%- else %}
68+
int32_t b_j_0 = SHFL_SYNC(b, j); \
69+
REPEAT_##block_size(B, j) \
70+
int32_t D_start_j_0 = SHFL_SYNC(D_start, j); \
71+
REPEAT_##block_size(D_START, j) \
72+
{%- endif %}
73+
{%- if weighted %}
74+
at::acc_type<cache_t, true> idx_weight_j_0 = SHFL_SYNC(idx_weight, j); \
75+
REPEAT_##block_size(IDX_WEIGHT, j) \
76+
{%- endif %}
77+
{%- set d = "(((vec + vec_start) * kThreadGroupSize + threadIdx.x) * VEC_WIDTH)" %}
78+
\
79+
for (int32_t vec = 0; vec < unroll_count && (((vec + vec_start) * kThreadGroupSize + threadIdx_x) * VEC_WIDTH) < D; ++vec) { \
80+
const int32_t d = (((vec + vec_start) * kThreadGroupSize + threadIdx_x) * VEC_WIDTH); \
81+
/* Generate block_size Vec4TAcc objects and accumulate them */ \
82+
Vec4TAcc<grad_t> grad_out_vec_0( \
83+
{%- if nobag and is_index_select %}
84+
&grad_output[grad_offset + l_j_0 * grad_stride + d] \
85+
{%- elif nobag %}
86+
&grad_output[l_j_0][d] \
87+
{%- elif vbe %}
88+
&grad_output[0][grad_offset_j_0 + d] \
89+
{%- else %}
90+
&grad_output[b_j_0][0] + D_start_j_0 + d \
91+
{%- endif %}
92+
); \
93+
{%- if nobag and is_index_select %}
94+
REPEAT_I_S_##block_size(GRAD_VEC_N_I, grad_offset, grad_stride, d) \
95+
{%- elif nobag %}
96+
REPEAT_##block_size(GRAD_VEC_N, d) \
97+
{%- elif vbe %}
98+
REPEAT_##block_size(GRAD_VEC_V, d) \
99+
{%- else %}
100+
REPEAT_##block_size(GRAD_VEC, d) \
101+
{%- endif %}
102+
\
103+
{%- if weighted %}
104+
grad_sum[vec].fma_(grad_out_vec_0, idx_weight_j_0); \
105+
REPEAT_##block_size(FMA_GRAD, vec) \
106+
{%- else %}
107+
grad_sum[vec].add_(grad_out_vec_0); \
108+
REPEAT_##block_size(ADD_GRAD, vec) \
109+
{%- endif %}
110+
} \
111+
}
112+
{%- endif %}
113+
17114
{%- if gen_once %}
18115
{#- /*
19116
The kernels in this section will be generated only once for all TBE configs
@@ -141,7 +238,44 @@ DEVICE_INLINE void compute_grad_sum_{{ kdesc }}(
141238
? sorted_indice_weights[segment_start + sl_j]
142239
: 0.0;
143240
{%- endif %}
144-
for (int32_t j = 0; j < kThreadGroupSize && sl + j < sl_end; ++j) {
241+
int32_t j = 0;
242+
243+
{%- if is_rocm %}
244+
// Process blocks of different sizes with loop unrolling
245+
if constexpr (sizeof(grad_t) <= 2) {
246+
PROCESS_BLOCK(8, kFixedMaxVecsPerThread, grad_sum, grad_output, grad_offset, \
247+
vec_start, kThreadGroupSize, threadIdx.x, VEC_WIDTH, D, j, sl, sl_end)
248+
}
249+
PROCESS_BLOCK(4, kFixedMaxVecsPerThread, grad_sum, grad_output, grad_offset, \
250+
vec_start, kThreadGroupSize, threadIdx.x, VEC_WIDTH, D, j, sl, sl_end)
251+
PROCESS_BLOCK(2, kFixedMaxVecsPerThread, grad_sum, grad_output, grad_offset, \
252+
vec_start, kThreadGroupSize, threadIdx.x, VEC_WIDTH, D, j, sl, sl_end)
253+
PROCESS_BLOCK(1, kFixedMaxVecsPerThread, grad_sum, grad_output, grad_offset, \
254+
vec_start, kThreadGroupSize, threadIdx.x, VEC_WIDTH, D, j, sl, sl_end)
255+
256+
#undef PROCESS_BLOCK
257+
#undef GRAD_OFFSET
258+
#undef L
259+
#undef B
260+
#undef D_START
261+
#undef IDX_WEIGHT
262+
#undef REPEAT_8
263+
#undef REPEAT_4
264+
#undef REPEAT_2
265+
#undef REPEAT_1
266+
#undef REPEAT_I_S_8
267+
#undef REPEAT_I_S_4
268+
#undef REPEAT_I_S_2
269+
#undef REPEAT_I_S_1
270+
#undef GRAD_VEC_N_I
271+
#undef GRAD_VEC_N
272+
#undef GRAD_VEC_V
273+
#undef GRAD_VEC
274+
#undef FMA_GRAD
275+
#undef ADD_GRAD
276+
277+
{%- else %}
278+
for (; j < kThreadGroupSize && sl + j < sl_end; ++j) {
145279
{%- if nobag %}
146280
int32_t l_j = SHFL_SYNC(l, j);
147281
{%- elif vbe %}
@@ -180,6 +314,7 @@ DEVICE_INLINE void compute_grad_sum_{{ kdesc }}(
180314
{%- endif %}
181315
}
182316
}
317+
{%- endif %}
183318
}
184319
{%- set d_vec = "((vec + vec_start) * kThreadGroupSize + threadIdx.x)" %}
185320

0 commit comments

Comments
 (0)