This is a part of a series of patches that rework OpenMP cross-team
reductions.
This patch changes the cross-team reduction runtime to no longer work
through larger number of teams in chunks. Instead, we allocate a
suitable-sized global buffer for the team values and let all teams run
at once. The last team that finishes uses a strided loop to reduce the
team values from the global buffer.
We also use `mapping::getNumberOfThreadsInBlock()` instead of
`omp_get_num_threads()` because the reduction of the team values runs
outside of the parallel region device code, which would make
`omp_get_num_threads()` always return 1. For Generic-SPMD mode, we also
want to use all available threads, which means that we need to copy the
reduction data from LDS (where it lives in that mode by default) to
scratch in codegen before calling the cross-team reduction.
Performance:
All performance tests can be reproduced with
https://github.com/ro-i/xteam-test @ commit
2f09ae1f566f13189d20b414d8b83bf2986564d2. To reproduce, simply create a
`local.mk` file in the cloned directory with a suitable `OFFLOAD_ARCH`
for your machine and `CXX_trunk` + `CXX_trunk_dev` set to the paths of
the clang++ binaries for llvm/main and this patch. (llvm/main should
best be at the commit that is currently the base for this PR. At the
moment, this is fe87c971bf). Then, run
`make trunk trunk_dev` to build the benchmark binaries for 208 and 10400
teams. Run them with `./run_bench.sh -rq -n10 red_trunk_208
red_trunk_dev_208 red_trunk_10400 red_trunk_dev_10400` to get the avg
performance numbers over 10 rounds.
This tests multiple reduction workloads, including reductions that run
in the Generic-SPMD mode, with 208 teams and with 10400 teams, both à
512 threads, and with a reduction array size of 177,777,777. I tested on
a gfx942 and didn't find any regressions - here are some numbers showing
the performance of this patch relative to the baseline:
```
red_comb_sep_arr_32 double dev_208: +84.37% dev_10400: +54.06%
red_sum_arr_32 double dev_208: +9.72% dev_10400: +51.29%
red_comb double dev_208: +2.73% dev_10400: +8.73%
red_comb_sep double dev_208: +1.89% dev_10400: +87.13%
red_dot double dev_208: +1.46% dev_10400: +25.66%
red_indirect double dev_208: +1.51% dev_10400: +23.54%
red_kernel_part double dev_208: +1.27% dev_10400: +88.28%
red_max double dev_208: +1.29% dev_10400: +22.69%
red_mult double dev_208: +1.96% dev_10400: +22.02%
red_sum double dev_208: +2.03% dev_10400: +22.00%
red_pi double dev_208: +1.08% dev_10400: +1.22%
red_comb_sep_arr_32 uint dev_208: +139.10% dev_10400: +88.02%
red_sum_arr_32 uint dev_208: +49.70% dev_10400: +143.24%
red_dot uint dev_208: +0.77% dev_10400: +24.91%
red_max uint dev_208: +3.43% dev_10400: +36.61%
red_sum uint dev_208: +4.56% dev_10400: +36.95%
red_comb_sep_arr_32 ulong dev_208: +85.22% dev_10400: +58.17%
red_sum_arr_32 ulong dev_208: +9.01% dev_10400: +55.14%
red_dot ulong dev_208: +2.94% dev_10400: +26.89%
red_max ulong dev_208: +2.83% dev_10400: +23.10%
red_sum ulong dev_208: +1.59% dev_10400: +22.82%
red_comb_sep_arr_32 Value dev_208: +29.88% dev_10400: +24.17%
red_sum_arr_32 Value dev_208: +7.26% dev_10400: +38.14%
red_dot Value dev_208: +2.65% dev_10400: +8.60%
red_max Value dev_208: +2.86% dev_10400: +4.77%
red_sum Value dev_208: +0.47% dev_10400: +10.22%
```
This patch uses ideas from Johannes Doerfert, ideas from the AOMP
cross-team reduction implementation, and was assisted by Claude.
e9acb01904