From ec1c85cd9a80e1e8ed9d80870f0dbc1b9487e03d Mon Sep 17 00:00:00 2001 From: dylanyunlon Date: Sat, 1 Aug 2026 00:31:22 +0800 Subject: [PATCH] =?UTF-8?q?fix(common.cuh):=20scale=5Fmem=5Fbound=20?= =?UTF-8?q?=E2=80=94=203=20bugs=20vs=20CCCL=20original?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. Return order: {threads, items} → {items, threads} matching CCCL scaling_result 2. Upper clamp: nominal*1 → nominal*2 (CCCL allows small types to double items) 3. Add threads SMEM cap: min(nominal, round_up(48KB/(ts*items), 32)) Verified against all 8 test vectors from CCCL catch2_test_util_arch.cu. The old code was only safe because current bi100_* structs don't hit the edge cases — but any future CCCL code copy would silently produce wrong values. --- muh/include/muh/tuning/common.cuh | 56 +++++++++++++++++++++++++------ 1 file changed, 45 insertions(+), 11 deletions(-) diff --git a/muh/include/muh/tuning/common.cuh b/muh/include/muh/tuning/common.cuh index b3ad429e..e319b575 100644 --- a/muh/include/muh/tuning/common.cuh +++ b/muh/include/muh/tuning/common.cuh @@ -119,21 +119,55 @@ enum class op_kind_t { plus, min, max, other }; namespace muh::tuning { -/// Memory-bound scaling: given nominal params for 4-byte types, -/// scale items_per_thread inversely with actual type size -/// to keep shared memory footprint constant. -/// Directly mirrors cub::detail::MemBoundScaling. -struct scaled_params { - int threads_per_block; +/// Scaling result — matches CCCL's cub::detail::scaling_result field order: +/// { items_per_thread, threads_per_block } +/// Callers destructure as: auto [items, threads] = scale_mem_bound(...); +struct scaling_result { int items_per_thread; + int threads_per_block; }; -constexpr scaled_params scale_mem_bound( - int nominal_threads, int nominal_4b_items, int type_size) { - int items = (nominal_4b_items * 4) / type_size; +/// Memory-bound scaling for non-4-byte types. +/// +/// Mirrors cub::detail::scale_mem_bound() from cub/util_arch.cuh lines 153-161. +/// Returns {items_per_thread, threads_per_block} — items-first, matching CCCL. +/// +/// Three operations: +/// 1. Scale items inversely with type size (4B nominal → 8B halves, 1B doubles) +/// 2. Clamp items to [1, nominal*2] (allows small types to increase items) +/// 3. Cap threads by SMEM: min(nominal, round_up(48KB / (type_size * items), 32)) +/// +/// CCCL source reference: +/// items = clamp(nominal_items * 4 / target_size, 1, nominal_items * 2) +/// threads = min(nominal_threads, round_up(max_smem / (target_size * items), 32)) +/// +/// The previous muh version had three bugs: +/// a) Return order was {threads, items} — should be {items, threads} +/// b) Upper clamp was nominal*1 — should be nominal*2 +/// c) No SMEM cap on threads — CCCL caps threads to prevent SMEM overflow +constexpr scaling_result scale_mem_bound( + int nominal_4B_threads, int nominal_4B_items, int target_type_size) { + constexpr int max_smem = 48 * 1024; // 49152 bytes + + // Step 1+2: scale items, clamp to [1, nominal*2] + int items = nominal_4B_items * 4 / target_type_size; if (items < 1) items = 1; - if (items > nominal_4b_items) items = nominal_4b_items; - return {nominal_threads, items}; + if (items > nominal_4B_items * 2) items = nominal_4B_items * 2; + + // Step 3: cap threads by SMEM + // round_up(x, 32) = ((x + 31) / 32) * 32 + int smem_per_thread = target_type_size * items; + int max_threads_by_smem; + if (smem_per_thread > 0) { + int raw = max_smem / smem_per_thread; + max_threads_by_smem = ((raw + 31) / 32) * 32; + } else { + max_threads_by_smem = nominal_4B_threads; + } + int threads = nominal_4B_threads < max_threads_by_smem + ? nominal_4B_threads : max_threads_by_smem; + + return {items, threads}; // items-first, matching CCCL scaling_result } } // namespace muh::tuning