[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:
muh-bot
2026-08-05 03:26:36 +00:00
parent faedfab7e9
commit 8d26e23e8e

View File

@@ -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