[muh] fix scale_reg: 补上 CCCL scale_reg_bound 的 threads SMEM cap
从 CCCL util_arch.cuh 读入 scale_reg_bound 精确实现:
items = max(1, nominal * 4 / max(4, type_size))
threads = min(nominal, round_up(48KB / (type_size * items), 32))
之前 muh 的 scale_reg 漏了第二行 (threads cap):
return {nominal_threads, items} // 没有 cap!
对 type_size=8/16 (double/int128) 可能导致 threads 过多, 寄存器溢出到 SMEM
超过 48KB 限制。当前 radix sort 参数 (type_size=4, nominal=256-384) 下不触发
但必须修正以保证 type_size=8 (float64 key) 的正确性。
同时加了 threads 下限 32 (一个 warp) 防止 SMEM 极端紧张时 threads=0。
This commit is contained in:
@@ -114,11 +114,22 @@ struct reg_scaled {
|
||||
int items_per_thread;
|
||||
};
|
||||
|
||||
// Matches CCCL util_arch.cuh scale_reg_bound exactly:
|
||||
// items = max(1, nominal * 4 / max(4, type_size)) [no expand beyond nominal]
|
||||
// threads = min(nominal, round_up(48KB / (type_size * items), 32)) [SMEM spill cap]
|
||||
constexpr reg_scaled scale_reg(int nominal_threads, int nominal_4b_items, int type_size) {
|
||||
int items = nominal_4b_items * 4 / (type_size > 4 ? type_size : 4);
|
||||
int ts = type_size > 4 ? type_size : 4;
|
||||
int items = nominal_4b_items * 4 / ts;
|
||||
if (items < 1) items = 1;
|
||||
if (items > nominal_4b_items) items = nominal_4b_items;
|
||||
return {nominal_threads, items};
|
||||
// CCCL caps threads by SMEM spill prevention (48KB / (type_size * items))
|
||||
int smem_per_thread = type_size * items;
|
||||
int max_threads = smem_per_thread > 0
|
||||
? ((49152 / smem_per_thread + 31) / 32) * 32 // round_up to warp multiple
|
||||
: nominal_threads;
|
||||
int threads = nominal_threads < max_threads ? nominal_threads : max_threads;
|
||||
if (threads < 32) threads = 32;
|
||||
return {threads, items};
|
||||
}
|
||||
|
||||
// Scale histogram private_partitions: more partitions for small types
|
||||
|
||||
Reference in New Issue
Block a user