[QST] Question about different number of N in warpgroup_wait<N> in existing code

Hello,

I am fixing code for my own functionality and have some questions about wargroup_wait called in mma function of sm90_mma_tma_gmma_x_warpspecialized_x.hpp files. Why do below files use different amount of N for warpgroup_wait? Just to be clear, I understand warpgroup_wait to be waiting until there is N outstanding warpgroup_commit_batch remaining. While warpgroup_arrive is needed to make all warps ready before wgmma.

• sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp : warpgroup_wait<K_BLOCK_MAX - 1>();
• sm90_mma_tma_gmma_ss_warpspecialized.hpp : warpgroup_wait<K_PIPE_MMAS>();
• sm90_mma_tma_gmma_rs_warpspecialized.hpp : warpgroup_wait<2>();

Another question is why is a single thread in a warp elected to process load functions? Are all other 31 threads idle? For example, below is the code for load function in "sm90_mma_tma_gmma_rs_warpspecialized.hpp".

load(...) {
int lane_predicate = cute::elect_one_sync();
if (lane_predicate) { ... }
}

Thank you!

Start by comparing the cited implementations in sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp, sm90_mma_tma_gmma_ss_warpspecialized.hpp, and sm90_mma_tma_gmma_rs_warpspecialized.hpp. Trace the mma and load entry points and the warpgroup_wait, warpgroup_arrive, and elect_one_sync calls. Done means documenting why the wait counts differ and how the elected-thread load behavior works.