Fix buffer safety, MPI type size correctness, and empty-data edge cases

ZCCL/include/ZCCL.h

@@ -41,6 +41,16 @@ extern "C" {
 #include "ZCCLd_float.h"
 #include "ZCCL_TypeManager.h"
 
+/* Expansion factor for compressed-data buffers — varies by data size.
+ * Smaller data needs more headroom because compression overhead dominates. */
+static inline int get_buf_exp(size_t nbytes) {
+    if (nbytes <= 4096)    return 24;   /* ≤ 4 KB    — very small, max overhead */
+    if (nbytes <= 65536)   return 12;   /* 4–64 KB   — small */
+    if (nbytes <= 1048576) return 6;    /* 64 KB–1 MB — medium */
+    return 4;                            /* > 1 MB    — large, compresses well */
+}
+// #define VERBOSE 1
+
 typedef union lint16
 {
 	unsigned short usvalue;

ZCCL/include/ZCCL_utils.h

@@ -7,6 +7,26 @@
 
 #ifndef _ZCCL_UTILS_H
 #define _ZCCL_UTILS_H
+#include <mpi.h>  // because used MPI_Datatype
+#include <stddef.h>  // because using sizeof
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * correctly get sizes for MPI_Datatype
+ */
+static inline size_t mpi_sizeof(MPI_Datatype dttype) {
+    int size;
+    int ret = MPI_Type_size(dttype, &size);
+    return (size_t)size;
+}
+
+#ifdef __cplusplus
+}
+#endif
+
 
 void get_4_digits(int input, int *output);
 

ZCCL/src/ZCCL_broadcast.c

@@ -34,17 +34,19 @@ int MPI_Bcast_ZCCL(void *buffer,
     MPI_Status stas[2];
     int is_contig;
     MPI_Aint type_size;
+    int buf_exp = 0;
     void *tmp_buf = NULL;
 
     MPI_Comm_rank(comm, &rank);
     MPI_Comm_size(comm, &comm_size);
 
     type_size = sizeof(data_type);
+    buf_exp = get_buf_exp(tmp_len * type_size);
 
     size_t outSize;
     size_t byteLength;
     int *compressed_sizes = (int *) malloc(comm_size * sizeof(int));
-    unsigned char *outputBytes = (unsigned char *) malloc(tmp_len * type_size);
+    unsigned char *outputBytes = (unsigned char *) malloc(buf_exp * tmp_len * type_size);
     if (rank == root) {
         ZCCL_float_single_thread_arg(outputBytes, buffer, &outSize, absErrBound, count, blockSize);
     }
@@ -123,17 +125,19 @@ int MPI_Bcast_ZCCL_mt(void *buffer,
     MPI_Status stas[2];
     int is_contig;
     MPI_Aint type_size;
+    int buf_exp = 0;
     void *tmp_buf = NULL;
 
     MPI_Comm_rank(comm, &rank);
     MPI_Comm_size(comm, &comm_size);
 
     type_size = sizeof(data_type);
+    buf_exp = get_buf_exp(tmp_len * type_size);
 
     size_t outSize;
     size_t byteLength;
     int *compressed_sizes = (int *) malloc(comm_size * sizeof(int));
-    unsigned char *outputBytes = (unsigned char *) malloc(tmp_len * type_size);
+    unsigned char *outputBytes = (unsigned char *) malloc(buf_exp * tmp_len * type_size);
     if (rank == root) {
         ZCCL_float_openmp_threadblock_arg(
             outputBytes, buffer, &outSize, absErrBound, count, blockSize);

ZCCL/src/ZCCL_float.c

@@ -26,6 +26,12 @@ unsigned char *ZCCL_float_openmp_direct_predict_quantization(float *oriData,
     size_t nbEle,
     int blockSize)
 {
+
+if (nbEle == 0) {
+    *outSize = 0;
+    return NULL;  
+}
+
 #ifdef _OPENMP
 
     float *op = oriData;
@@ -78,6 +84,10 @@ unsigned char *ZCCL_float_openmp_threadblock_predict_quantization(float *oriData
     size_t nbEle,
     int blockSize)
 {
+if (nbEle == 0) {
+    *outSize = 0;
+    return NULL;  
+}
 #ifdef _OPENMP
 
     float *op = oriData;
@@ -150,6 +160,10 @@ unsigned char *ZCCL_fast_compress_args(int fastMode,
 {
     unsigned char *bytes = NULL;
     size_t length = computeDataLength(r5, r4, r3, r2, r1);
+    if (length == 0) {
+        *outSize = 0;
+        return NULL;  
+    }
     size_t i = 0;
     int blockSize = 128;
     if (dataType == SZ_FLOAT) {
@@ -195,6 +209,10 @@ unsigned char *ZCCL_float_openmp_threadblock(float *oriData,
     size_t nbEle,
     int blockSize)
 {
+if (nbEle == 0) {
+    *outSize = 0;
+    return NULL;  
+}
 #ifdef _OPENMP
 
     float *op = oriData;
@@ -505,8 +523,23 @@ void ZCCL_float_openmp_threadblock_arg(unsigned char *outputBytes,
     size_t nbEle,
     int blockSize)
 {
+if (nbEle == 0) {
+    *outSize = 0;
+    return NULL;  
+}
+
 #ifdef _OPENMP
 
+    int world_rank;
+    MPI_Comm_rank(MPI_COMM_WORLD, &world_rank);  // 获取进程号
+
+    #ifdef VERBOSE
+    printf("===== Rank %d: 开始压缩 =====\n", world_rank);
+    printf("原始数据: nbEle=%zu, 原始大小=%zu 字节\n", 
+           nbEle, nbEle * sizeof(float));
+    fflush(stdout);
+    #endif
+
     float *op = oriData;
 
     size_t maxPreservedBufferSize = sizeof(float) * nbEle;
@@ -536,7 +569,7 @@ void ZCCL_float_openmp_threadblock_arg(unsigned char *outputBytes,
             outSize_perthread_arr = (size_t *) malloc(nbThreads * sizeof(size_t));
             offsets_perthread_arr = (size_t *) malloc(nbThreads * sizeof(size_t));
 
-            maxPreservedBufferSize_perthread = (sizeof(float) * nbEle + nbThreads - 1) / nbThreads;
+            maxPreservedBufferSize_perthread = 2 * (sizeof(float) * nbEle + nbThreads - 1) / nbThreads;
             inver_bound = 1 / absErrBound;
             threadblocksize = nbEle / nbThreads;
             remainder = nbEle % nbThreads;
@@ -772,6 +805,11 @@ void ZCCL_float_openmp_threadblock_arg(unsigned char *outputBytes,
         }
 
         outSize_perthread_arr[tid] = outSize_perthread;
+        #ifdef VERBOSE
+        printf("Rank %d, Thread %d: outSize_perthread=%zu\n", 
+            world_rank, tid, outSize_perthread);
+        fflush(stdout);
+        #endif
 #pragma omp barrier
 
 #pragma omp single
@@ -784,6 +822,14 @@ void ZCCL_float_openmp_threadblock_arg(unsigned char *outputBytes,
             (*outSize) +=
                 offsets_perthread_arr[nbThreads - 1] + outSize_perthread_arr[nbThreads - 1];
             memcpy(outputBytes, offsets_perthread_arr, nbThreads * sizeof(size_t));
+            #ifdef VERBOSE
+            printf("Rank %d: 压缩后大小=%zu 字节\n", 
+                world_rank, *outSize);
+            printf("Rank %d: 压缩比=%.2f%%\n", 
+                world_rank, (double)(*outSize) / (nbEle * sizeof(float)) * 100);
+            printf("=============================\n");
+            #endif
+            fflush(stdout);
         }
 #pragma omp barrier
         memcpy(real_outputBytes + offsets_perthread_arr[tid],
@@ -813,6 +859,10 @@ void ZCCL_float_single_thread_arg(unsigned char *outputBytes,
     size_t nbEle,
     int blockSize)
 {
+if (nbEle == 0) {
+    *outSize = 0;
+    return NULL;  
+}
     float *op = oriData;
 
     size_t maxPreservedBufferSize = sizeof(float) * nbEle;
@@ -838,6 +888,7 @@ void ZCCL_float_single_thread_arg(unsigned char *outputBytes,
     offsets_perthread_arr = (size_t *) malloc(nbThreads * sizeof(size_t));
 
     maxPreservedBufferSize_perthread = (sizeof(float) * nbEle + nbThreads - 1) / nbThreads;
+    assert(maxPreservedBufferSize_perthread);
     inver_bound = 1 / absErrBound;
     threadblocksize = nbEle / nbThreads;
     remainder = nbEle % nbThreads;
@@ -994,6 +1045,10 @@ size_t ZCCL_float_single_thread_arg_record(unsigned char *outputBytes,
     size_t nbEle,
     int blockSize)
 {
+if (nbEle == 0) {
+    *outSize = 0;
+    return NULL;  
+}
     size_t total_memaccess = 0;
 
     float *op = oriData;
@@ -1215,6 +1270,10 @@ unsigned char *ZCCL_float_openmp_threadblock_randomaccess(float *oriData,
     size_t nbEle,
     int blockSize)
 {
+if (nbEle == 0) {
+    *outSize = 0;
+    return NULL;  
+}
 #ifdef _OPENMP
 
     float *op = oriData;
@@ -1536,6 +1595,10 @@ void ZCCL_float_single_thread_arg_split_record(unsigned char *outputBytes,
     unsigned char *chunk_arr,
     size_t chunk_iter)
 {
+if (nbEle == 0) {
+    *outSize = 0;
+    return NULL;  
+}
     float *op = oriData;
 
     size_t *arr = (size_t *) chunk_arr;