[mpc] fix out-of-bounds output write in MPCcompress (issue #317)

[zjin-lcf]

Summary

Fixes #317 — a heap/global buffer overflow in the MPCcompress kernel reported with compute-sanitizer on a crafted small input.

The compressor sized the output buffer's residual region by the input word count (insize):

outsize = insize + 1 + (insize + 63) / 64;

but MPCcompress can emit up to one residual word per thread across all launched threads. Small inputs still transpose full 64-word bit-planes from the padded tail, so a tiny input (e.g. one 8-byte word → 24-byte buffer) makes the kernel write past d_out at compressed[start + tid], which the sanitizer flags as an invalid __global__ write.

Fix

Size the residual region to the worst case the kernel can emit — ceil(insize/TPB) chunks × TPB threads — plus the header and per-64 bitmap words:

const int chunks = (insize + TPB - 1) / TPB;
outsize = 1 + (insize + 63) / 64 + chunks * TPB;

This is the "size d_out to the worst case the kernel can emit" option from the issue's suggested fix. The kernel is untouched, so there is no ABI change and no truncation of output.

Applied to the mpc-cuda, mpc-hip, and mpc-sycl variants.

Notes

• No performance impact: the allocation is outside the timed region, and the device→host copy still uses the actual compressed length from the header (output[0] >> 32), so program output is unchanged.
• Only the compression path (argc == 3) is affected; decompression is unchanged.

Test plan

• Built the SYCL variant with the fix.
• Ran a pathological 8-byte input (0xAA…, maximizing bit-plane residuals) — now completes cleanly.
• Rebuild mpc-cuda with -lineinfo and re-run compute-sanitizer --tool=memcheck ./main ./mpc-memcheck.txt 1 to confirm no invalid writes (no CUDA toolchain on my machine).

[shufengc]

Rebuild mpc-cuda with -lineinfo and re-run compute-sanitizer --tool=memcheck ./main ./mpc-memcheck.txt 1 to confirm no invalid writes (no CUDA toolchain on my machine).

Fix has been confirmed:
Ran this on an A30, CUDA 12.9, Compute Sanitizer 2025.2.0.0, with the PoC from #317.

make clean && make ARCH=sm_80 EXTRA_CFLAGS="-lineinfo"
compute-sanitizer --tool=memcheck ./main ./mpc-memcheck.txt 1

This PR: 64e1d71 has 0 errors.

[zjin-lcf]

Thank you for the tests.

[Copilot]

Pull request overview

Fixes MPC compression buffer overflows by allocating for worst-case per-thread residual output.

Changes:

• Computes chunk-rounded residual capacity.
• Applies the fix consistently across CUDA, HIP, and SYCL.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

File	Description
src/mpc-cuda/main.cu	Corrects CUDA output sizing.
src/mpc-hip/main.cu	Corrects HIP output sizing.
src/mpc-sycl/main.cpp	Corrects SYCL output sizing.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[keitaTN]

This fix looks reasonable to me when handling out-of-bound access for arrays whose size is not multiple of 64. BTW, it assume warp/wavefront/vector size is 64 or less (31,16,8,4,2, or 1)? Do you have any back-up plan for future architecture especially RISC-V that can accommodate bigger vector size? (At ORNL, we are thinking abou tobtaining Tenstorrent).

[zjin-lcf]

Thanks for the review, Keita. Short answer: the sizing fix itself carries no vector-width assumption — the 64s in it are the data word width, not the lane count. The width assumption you're sensing is real, but it is pre-existing and lives in the kernel rather than in the allocation.

1. The 64s in outsize = 1 + (insize + 63) / 64 + chunks * TPB are sizeof(long) * 8.
MPC's BIT stage transposes a 64-value × 64-bit block, so the bitmap holds exactly one 64-bit word per 64 input values. That is part of the MPC container format (the decompressor reads compressed[1 + idx / 64]), so it stays 64 on any hardware. The transpose goes through shared memory plus __syncthreads() — lanex = tid & 63 and warpx = tid & 0x3c0 merely index a 64-thread tile inside the 1024-thread block — so it does not require that tile to be a warp.

2. The residual term is a per-thread bound, so it is width-agnostic.
Within a chunk each thread contributes at most one residual word (if (v2 != 0) sbuf2[loc - 1] = v2; followed by if (tid < top) compressed[start + tid] = sbuf2[tid];), so the worst case over chunks = ceil(insize / TPB) chunks of TPB threads is chunks * TPB words. That bound holds for any subgroup width, which is why this line would not need revisiting if the vector length changed.

3. Where the width assumption actually is: the ballot that assembles the bitmap.
CUDA hard-codes 32 and stitches two 32-bit halves into the 64-bit word:

unsigned int bitmap = __ballot_sync(0xffffffff, loc);
if (lanex == 32) sbuf2[tid] = bitmap;
__syncthreads();
if (lanex == 0) {
  if (idx < n) compressed[1 + idx / 64] = (sbuf2[tid + 32] << 32) + bitmap;
}

HIP and SYCL are templated on WarpSize with separate 32 and 64 paths, and HIP bails out on anything else ("Only a warp size of 32 or 64 is supported"). So the current requirement is a subgroup of exactly 32 or 64 lanes, i.e. the 64-value bit-plane tile must coincide with either one 64-lane wavefront or two 32-lane warps. More than 64 lanes per subgroup would break that assembly — a single ballot would span more than one bitmap word — and would also break the shuffle-based prefixsum* scans. It would not, however, invalidate the buffer size computed here.

4. Back-up plan for wider vectors. Two width-agnostic replacements, both local to the kernel:

• Bitmap: drop the ballot and build each 64-bit word in shared memory, e.g. every thread with v2 != 0 does atomicOr(&smem_bitmap[tid >> 6], 1ull << lanex) and one thread per tile stores compressed[1 + idx / 64]. That decouples the format's 64 from the hardware lane count entirely, at the cost of one shared-memory atomic per non-zero residual. Where available, a sycl::group_ballot over a fixed 64-lane tile achieves the same.
• Scans: replace the shuffle scans with work-group collectives — sycl::inclusive_scan_over_group / sycl::joint_inclusive_scan, or CUB BlockScan / rocPRIM block_scan — none of which encode a lane width in the source.