From 177b1b3c7857a0d285140f5ed4663015ab379fea Mon Sep 17 00:00:00 2001 From: Qian Gong Date: Wed, 22 Jul 2026 04:04:58 -0400 Subject: [PATCH] Fix MGARD-X HIP Huffman lossless (codec correctness + code-length OOB) Four independent bugs in the MGARD-X HIP ParallelHuffman lossless path: three in the Huffman codec (canonical-codeword generation, header serialization, and codeword-length overflow handling) made GPU-compressed streams undecodable/oversized or abort, and one out-of-bounds device read in the parallel code-length generation faulted on the GPU. All verified on gfx90a (Frontier MI250X/MI210): HIP-compressed streams now round-trip and are byte-compatible with the SERIAL decoder. HuffmanWorkspace.hpp: allocate `status` as non-managed device memory. It was managed (hipMallocManaged); atomicMin on fine-grained memory is unreliable on ROCm/gfx90a, so GenerateCW::Operation4's atomic min never updated newCDPI. The Huffman length-group boundary was never found, all symbols collapsed into one group, and the canonical codewords came out with a constant offset -- codebook and decodebook mutually inconsistent, so the stream was undecodable by any backend. Huffman.hpp: default-initialize `outlier_count = 0`. ComposedLosslessCompressor::Compress -> CompressPrimary never writes outlier_count, but Serialize reads it to size compressed_data.resize(). The uninitialized value produced a garbage-sized allocation (hipMalloc OOM, or an invalid Copy1D). GetCodebook.hpp: throw instead of exit(1) when the longest codeword exceeds the H-type budget (sizeof(H)*8 - 8 bits), so callers can catch it and fall back to another lossless backend or a smaller dict_size. GenerateCL.hpp: bounds-check histogram[_lNodesCur + _curLeavesNum] in the parallel code-length generation; when every remaining leaf joins the batch the index reached dict_size (one past end of the nz_dict_size histogram), faulting on GPU. Co-Authored-By: Claude Opus 4.8 --- .../Lossless/ParallelHuffman/GenerateCL.hpp | 29 ++++++++++++++----- .../Lossless/ParallelHuffman/GetCodebook.hpp | 14 ++++++++- .../Lossless/ParallelHuffman/Huffman.hpp | 2 +- .../ParallelHuffman/HuffmanWorkspace.hpp | 4 ++- 4 files changed, 39 insertions(+), 10 deletions(-) diff --git a/include/mgard-x/Lossless/ParallelHuffman/GenerateCL.hpp b/include/mgard-x/Lossless/ParallelHuffman/GenerateCL.hpp index bf3bd5497e..617a3fa1c5 100644 --- a/include/mgard-x/Lossless/ParallelHuffman/GenerateCL.hpp +++ b/include/mgard-x/Lossless/ParallelHuffman/GenerateCL.hpp @@ -349,15 +349,30 @@ class GenerateCLFunctor : public HuffmanCLCustomizedFunctor { (*status((IDX)_iNodesFront)) = (*status((IDX)_iNodesRear)); } /* Odd number of nodes to merge - leave out one*/ - // If number of participating leaf node is zero OR - // The highest frequency leaf node is less than the - // highest frequency internal node + // Remove one internal node (keep it for the next outer iteration) if: + // (a) there are no leaf nodes in the batch (_curLeavesNum == 0), OR + // (b) the first leaf node NOT in the batch has a frequency <= the + // highest-frequency internal node in the batch, meaning it is + // cheaper to defer that internal node. + // + // BUG FIX: The original code read histogram[_lNodesCur + _curLeavesNum] + // unconditionally. When _curLeavesNum == dict_size - _lNodesCur (i.e., + // all remaining leaf nodes qualify for the batch), the index equals + // dict_size, which is one past the end of the nz_dict_size-sized + // histogram array, causing a GPU memory access fault. + // FIX: Guard the histogram access with the bounds check + // (_lNodesCur + _curLeavesNum < dict_size). + // When there is no next leaf outside the batch, the condition is treated + // as false and execution falls through to the else branch (remove one + // leaf from the batch instead). else if (((*status((IDX)_iNodesSize)) != 0) and ((*status((IDX)_curLeavesNum)) == 0 or - (*histogram((IDX)(*status((IDX)_lNodesCur)) + - (*status((IDX)_curLeavesNum))) <= - *iNodesFreq( - (IDX)MOD((*status((IDX)_iNodesRear)) - 1, dict_size))))) { + (((*status((IDX)_lNodesCur)) + (*status((IDX)_curLeavesNum)) < + dict_size) and + (*histogram((IDX)(*status((IDX)_lNodesCur)) + + (*status((IDX)_curLeavesNum))) <= + *iNodesFreq((IDX)MOD((*status((IDX)_iNodesRear)) - 1, + dict_size)))))) { (*status((IDX)_mergeRear)) = MOD((*status((IDX)_mergeRear)) - 1, dict_size); (*status((IDX)_iNodesFront)) = diff --git a/include/mgard-x/Lossless/ParallelHuffman/GetCodebook.hpp b/include/mgard-x/Lossless/ParallelHuffman/GetCodebook.hpp index 0acc4c170f..fc3343d377 100644 --- a/include/mgard-x/Lossless/ParallelHuffman/GetCodebook.hpp +++ b/include/mgard-x/Lossless/ParallelHuffman/GetCodebook.hpp @@ -14,6 +14,9 @@ #include "ReorderByIndex.hpp" #include "ReverseArray.hpp" +#include +#include + #ifndef MGARD_X_GET_CODEBOOK_TEMPLATE_HPP #define MGARD_X_GET_CODEBOOK_TEMPLATE_HPP @@ -119,7 +122,16 @@ void GetCodebook(int dict_size, << "Huffman codeword representation requires at least " << max_CL + 8 << " bits (longest codeword: " << max_CL << " bits)" << std::endl; - exit(1); + // Throw (instead of exit) so callers can catch and fall back to another + // lossless backend (e.g. raw Zstd) or retry with a smaller huff_dict_size. + // A longer dictionary makes the tree deeper, so a degenerate/low-entropy + // input can produce codewords exceeding the H-type budget (sizeof(H)*8 - + // 8). + throw std::runtime_error( + "MGARD-X Huffman: longest codeword (" + std::to_string(max_CL) + + " bits) exceeds the " + std::to_string(max_CW_bits) + + "-bit budget of the H code type; retry with a smaller huff_dict_size " + "or a different lossless backend"); } DeviceLauncher::Execute( diff --git a/include/mgard-x/Lossless/ParallelHuffman/Huffman.hpp b/include/mgard-x/Lossless/ParallelHuffman/Huffman.hpp index 1c607b2752..4f2105f649 100644 --- a/include/mgard-x/Lossless/ParallelHuffman/Huffman.hpp +++ b/include/mgard-x/Lossless/ParallelHuffman/Huffman.hpp @@ -446,7 +446,7 @@ class Huffman : public LosslessCompressorInterface { bool initialized; SIZE max_size; size_t primary_count; - ATOMIC_IDX outlier_count; + ATOMIC_IDX outlier_count = 0; int dict_size; int chunk_size; size_t huffmeta_size; diff --git a/include/mgard-x/Lossless/ParallelHuffman/HuffmanWorkspace.hpp b/include/mgard-x/Lossless/ParallelHuffman/HuffmanWorkspace.hpp index d64e5f76f8..dcefc0fd53 100644 --- a/include/mgard-x/Lossless/ParallelHuffman/HuffmanWorkspace.hpp +++ b/include/mgard-x/Lossless/ParallelHuffman/HuffmanWorkspace.hpp @@ -135,7 +135,9 @@ class HuffmanWorkspace { copyIsLeaf_array = Array<1, int, DeviceType>({dict_size}); copyIndex_array = Array<1, int, DeviceType>({dict_size}); _d_codebook_array_org = Array<1, H, DeviceType>({dict_size}); - status_array = Array<1, int, DeviceType>({(SIZE)16}, false, true); + status_array = Array<1, int, DeviceType>( + {(SIZE)16}, false, + false); // non-managed: atomicMin on managed mem unreliable on ROCm SIZE mblocks = (DeviceRuntime::GetMaxNumThreadsPerTB() / DeviceRuntime::GetWarpSize()) * DeviceRuntime::GetNumSMs();