From 64e1d71b4e89c99fa6cc13a0ca82f7068311385d Mon Sep 17 00:00:00 2001 From: Zheming Jin Date: Mon, 10 Aug 2026 16:39:27 -0700 Subject: [PATCH] [mpc] fix out-of-bounds output write in MPCcompress (issue #317) The compressor sized the output buffer's residual region by the input word count (insize), 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 word) makes the kernel write past d_out, which compute-sanitizer flags as an invalid __global__ write. 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. The device->host copy still uses the actual compressed length from the header, so output is unchanged and there is no perf impact. Applied to the cuda, hip, and sycl variants. Fixes #317 Co-authored-by: Cursor --- src/mpc-cuda/main.cu | 9 ++++++++- src/mpc-hip/main.cu | 9 ++++++++- src/mpc-sycl/main.cpp | 9 ++++++++- 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/src/mpc-cuda/main.cu b/src/mpc-cuda/main.cu index 4a63556fc..de31fa666 100644 --- a/src/mpc-cuda/main.cu +++ b/src/mpc-cuda/main.cu @@ -383,7 +383,14 @@ int main(int argc, char *argv[]) if (argc == 3) { dim = atoi(argv[2]); - outsize = insize + 1 + (insize + 63) / 64; + // Worst-case output size the compressor can emit: a header word, one bitmap + // word per 64 inputs, and up to one residual word per thread. Each of the + // ceil(insize/TPB) chunks launches TPB threads, so the residual region can + // grow to chunks*TPB words. Small inputs still transpose full 64-word + // bit-planes from the padded tail, so sizing residuals by insize alone + // under-allocates and lets MPCcompress write past d_out. + const int chunks = (insize + TPB - 1) / TPB; + outsize = 1 + (insize + 63) / 64 + chunks * TPB; } else { assert(((input[0] >> 8) & 0xffffff) == 0x43504d); dim = (input[0] & 31) + 1; diff --git a/src/mpc-hip/main.cu b/src/mpc-hip/main.cu index 667d7c874..21134bf1c 100644 --- a/src/mpc-hip/main.cu +++ b/src/mpc-hip/main.cu @@ -389,7 +389,14 @@ int main(int argc, char *argv[]) if (argc == 3) { dim = atoi(argv[2]); - outsize = insize + 1 + (insize + 63) / 64; + // Worst-case output size the compressor can emit: a header word, one bitmap + // word per 64 inputs, and up to one residual word per thread. Each of the + // ceil(insize/TPB) chunks launches TPB threads, so the residual region can + // grow to chunks*TPB words. Small inputs still transpose full 64-word + // bit-planes from the padded tail, so sizing residuals by insize alone + // under-allocates and lets MPCcompress write past d_out. + const int chunks = (insize + TPB - 1) / TPB; + outsize = 1 + (insize + 63) / 64 + chunks * TPB; } else { assert(((input[0] >> 8) & 0xffffff) == 0x43504d); dim = (input[0] & 31) + 1; diff --git a/src/mpc-sycl/main.cpp b/src/mpc-sycl/main.cpp index 4f82ebe39..a04c9a605 100644 --- a/src/mpc-sycl/main.cpp +++ b/src/mpc-sycl/main.cpp @@ -423,7 +423,14 @@ int main(int argc, char *argv[]) if (argc == 3) { dim = atoi(argv[2]); - outsize = insize + 1 + (insize + 63) / 64; + // Worst-case output size the compressor can emit: a header word, one bitmap + // word per 64 inputs, and up to one residual word per thread. Each of the + // ceil(insize/TPB) chunks launches TPB threads, so the residual region can + // grow to chunks*TPB words. Small inputs still transpose full 64-word + // bit-planes from the padded tail, so sizing residuals by insize alone + // under-allocates and lets MPCcompress write past d_out. + const int chunks = (insize + TPB - 1) / TPB; + outsize = 1 + (insize + 63) / 64 + chunks * TPB; } else { assert(((input[0] >> 8) & 0xffffff) == 0x43504d); dim = (input[0] & 31) + 1;