From c33f80878fb2d4280fcd63a6afa2752437685d46 Mon Sep 17 00:00:00 2001 From: Pavan Balaji Date: Sun, 8 Mar 2026 01:00:55 -0800 Subject: [PATCH] Fix UB in allgather with zero-size output Summary: When the output buffer is empty (outBytes == 0), the input buffer pointer may be null. The memcpy at line 50 was called before the early return check, passing a null pointer to memcpy which is undefined behavior even with size 0. UBSAN catches this in the Count_0 test cases. Move the outBytes == 0 early return before the memcpy to avoid the null pointer. Split the context->size == 1 check to remain after the memcpy, since single-process still needs the input-to-output copy. Differential Revision: D95712700 --- gloo/allgather.cc | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/gloo/allgather.cc b/gloo/allgather.cc index 8ab4e6f4f..e79ea7c7c 100644 --- a/gloo/allgather.cc +++ b/gloo/allgather.cc @@ -44,6 +44,11 @@ void allgather(AllgatherOptions& opts) { const size_t inBytes = out->size / context->size; const size_t outBytes = out->size; + // Short circuit if the output is empty. + if (outBytes == 0) { + return; + } + // If the input buffer is specified, this is NOT an in place operation, // and the output buffer needs to be primed with the input. if (in != nullptr) { @@ -53,8 +58,8 @@ void allgather(AllgatherOptions& opts) { in->size); } - // Short circuit if there is only a single process or the output is empty. - if (context->size == 1 || outBytes == 0) { + // Short circuit if there is only a single process. + if (context->size == 1) { return; }