From 9629507587fe0e01d433098e1f7f0c53cf499dfe Mon Sep 17 00:00:00 2001 From: Trent Nelson Date: Tue, 29 Sep 2026 19:08:37 -0700 Subject: [PATCH] Honor nvCOMP requirements in batched Deflate decoding Align raw payloads before decoding and retain their device buffers until stream work completes. Supply status and actual-size scratch for nvCOMP compatibility, including pooled and failure paths. Signed-off-by: Trent Nelson --- src/cuphoton/xdr/nvcomp_batch.py | 177 ++++++++++++++------ src/cuphoton/xdr/src/nvcomp_batch_ext.cpp | 129 +++++++++------ tests/xdr/test_deflate_gpu.py | 188 ++++++++++++++++++++++ tests/xdr/test_nvcomp_batch.py | 2 +- 4 files changed, 396 insertions(+), 100 deletions(-) create mode 100644 tests/xdr/test_deflate_gpu.py diff --git a/src/cuphoton/xdr/nvcomp_batch.py b/src/cuphoton/xdr/nvcomp_batch.py index 790e7a93..5704dbd4 100644 --- a/src/cuphoton/xdr/nvcomp_batch.py +++ b/src/cuphoton/xdr/nvcomp_batch.py @@ -27,6 +27,7 @@ import numpy as np +_REPACK_DEFLATE_KERNEL = None _DEFLATE_CODEC = None _CPP_EXT = None _CPP_EXT_PROBED = False @@ -746,6 +747,65 @@ def _checked_output_offsets( return offsets, total +def _align_deflate_inputs(d_concat, offsets, lengths): + """Pack raw DEFLATE payloads at nvCOMP's required four-byte alignment.""" + if not np.any((int(d_concat.data.ptr) + offsets) % 4): + return d_concat, offsets, () + + import cupy as cp + + global _REPACK_DEFLATE_KERNEL + if _REPACK_DEFLATE_KERNEL is None: + _REPACK_DEFLATE_KERNEL = cp.RawKernel( + r""" + extern "C" __global__ void align_deflate( + const unsigned char* src, const long long* offsets, + const long long* lengths, const long long* out_offsets, + unsigned char* dst) { + const int tile = blockIdx.x; + for (long long i = threadIdx.x; i < lengths[tile]; + i += blockDim.x) { + dst[out_offsets[tile] + i] = src[offsets[tile] + i]; + } + } + """, + "align_deflate", + ) + if np.any(lengths > np.iinfo(np.int64).max - 3): + raise ValueError("aligned DEFLATE size exceeds the int64 range") + aligned_offsets, total = _checked_output_offsets((lengths + 3) & ~3) + packed = cp.empty(total, dtype=cp.uint8) + tables = [] + try: + for table in (offsets, lengths, aligned_offsets): + tables.append(cp.asarray(table)) + _REPACK_DEFLATE_KERNEL( + (len(offsets),), (256,), (d_concat, *tables, packed) + ) + except Exception: + # Keep the input and metadata alive through any queued copies. + cp.cuda.get_current_stream().synchronize() + raise + return packed, aligned_offsets, (d_concat, packed, *tables) + + +def _retain_decode_inputs(d_out, owners): + """Keep repacked input alive until the decoded output is released.""" + import cupy as cp + + memory = cp.cuda.UnownedMemory( + int(d_out.data.ptr), + int(d_out.nbytes), + (d_out, *owners), + device_id=int(d_out.device.id), + ) + return cp.ndarray( + d_out.shape, + dtype=d_out.dtype, + memptr=cp.cuda.MemoryPointer(memory, 0), + ) + + def gpu_gzip_decompress_batch( d_concat, rel_offsets: Sequence[int], @@ -861,66 +921,81 @@ def gpu_gzip_decompress_batch( if n == 0: return cp.empty(0, dtype=cp.uint8), out_offsets - # Allocate one concatenated output buffer and slice it per tile. - d_out = ( - _native_device_empty_uint8(total, ext=ext) - if use_native_pool and ext is not None - else cp.empty(total, dtype=cp.uint8) + d_concat, deflate_offsets, aligned_owners = _align_deflate_inputs( + d_concat, deflate_offsets, deflate_lengths ) if keepalive is not None: - keepalive.append(d_out) - - if ext is not None: - # C++ path: device pointers + length arrays, no per-tile Python loop. - stream_ptr = int(cp.cuda.get_current_stream().ptr) - d_concat_ptr = int(d_concat.data.ptr) - d_out_ptr = int(d_out.data.ptr) - pooled_fn = getattr(ext, "batch_deflate_decompress_pooled", None) - if ( - use_native_pool - and pooled_fn is not None - and keepalive is not None - ): - scratch_owner = pooled_fn( + keepalive.extend(aligned_owners) + + try: + # Allocate one concatenated output buffer and slice it per tile. + d_out = ( + _native_device_empty_uint8(total, ext=ext) + if use_native_pool and ext is not None + else cp.empty(total, dtype=cp.uint8) + ) + if aligned_owners and keepalive is None: + d_out = _retain_decode_inputs(d_out, aligned_owners) + if keepalive is not None: + keepalive.append(d_out) + + if ext is not None: + # C++ path: device pointers and tables, without a per-tile loop. + stream_ptr = int(cp.cuda.get_current_stream().ptr) + d_concat_ptr = int(d_concat.data.ptr) + d_out_ptr = int(d_out.data.ptr) + pooled_fn = getattr(ext, "batch_deflate_decompress_pooled", None) + if ( + use_native_pool + and pooled_fn is not None + and keepalive is not None + ): + scratch_owner = pooled_fn( + d_concat_ptr, + deflate_offsets, + deflate_lengths, + d_out_ptr, + out_offsets, + uncompressed_sizes, + stream_ptr, + ) + keepalive.append(scratch_owner) + return d_out, out_offsets + + ext.batch_deflate_decompress( d_concat_ptr, deflate_offsets, deflate_lengths, d_out_ptr, out_offsets, - uncompressed_sizes, + size_array, stream_ptr, ) - keepalive.append(scratch_owner) return d_out, out_offsets - ext.batch_deflate_decompress( - d_concat_ptr, - deflate_offsets, - deflate_lengths, - d_out_ptr, - out_offsets, - size_array, - stream_ptr, - ) - return d_out, out_offsets - - # Python fallback: original path, retained for compile-less checkouts. - import nvidia.nvcomp as nvcomp + # Python fallback: original path, retained for compile-less checkouts. + import nvidia.nvcomp as nvcomp - src_arrays = [ - nvcomp.as_array( - d_concat[ - deflate_offsets[i] : deflate_offsets[i] + deflate_lengths[i] - ] - ) - for i in range(n) - ] - out_arrays = [ - nvcomp.as_array( - d_out[out_offsets[i] : out_offsets[i] + size_array[i]] - ) - for i in range(n) - ] - codec = _get_codec() - codec.decode(src_arrays, out=out_arrays) - return d_out, out_offsets + src_arrays = [ + nvcomp.as_array( + d_concat[ + deflate_offsets[i] : deflate_offsets[i] + + deflate_lengths[i] + ] + ) + for i in range(n) + ] + out_arrays = [ + nvcomp.as_array( + d_out[out_offsets[i] : out_offsets[i] + size_array[i]] + ) + for i in range(n) + ] + codec = _get_codec() + codec.decode(src_arrays, out=out_arrays) + return d_out, out_offsets + except Exception: + # A failed decode cannot return an output to retain its input owners. + if keepalive is None and (aligned_owners or use_native_pool): + cp.cuda.get_current_stream().synchronize() + raise diff --git a/src/cuphoton/xdr/src/nvcomp_batch_ext.cpp b/src/cuphoton/xdr/src/nvcomp_batch_ext.cpp index 450249f6..218a6526 100644 --- a/src/cuphoton/xdr/src/nvcomp_batch_ext.cpp +++ b/src/cuphoton/xdr/src/nvcomp_batch_ext.cpp @@ -130,6 +130,9 @@ py::object batch_deflate_decompress_impl( } h_comp_ptrs[i] = reinterpret_cast(d_concat_ptr + static_cast(ro(i))); + if (reinterpret_cast(h_comp_ptrs[i]) % 4 != 0) { + throw std::invalid_argument("raw DEFLATE inputs must be 4-byte aligned"); + } h_out_ptrs[i] = reinterpret_cast(d_out_ptr + static_cast(oo(i))); h_comp_sizes[i] = static_cast(ln(i)); h_out_sizes[i] = static_cast(os(i)); @@ -143,11 +146,12 @@ py::object batch_deflate_decompress_impl( void* d_comp_sizes = nullptr; void* d_out_sizes = nullptr; void* d_temp = nullptr; + void* d_statuses = nullptr; + void* d_actual_sizes = nullptr; auto opts = nvcompBatchedDeflateDecompressDefaultOpts; - // The hardware backend requires non-stream-ordered scratch plus explicit - // actual-size and status buffers. Keep this path on CUDA until those - // ownership and reporting semantics are implemented together. + // The hardware backend requires non-stream-ordered scratch allocations. + // Keep this path on CUDA until that ownership contract is supported. opts.backend = NVCOMP_DECOMPRESS_BACKEND_CUDA; std::size_t temp_bytes = 0; check_nvcomp( @@ -156,39 +160,67 @@ py::object batch_deflate_decompress_impl( "get temp size"); std::vector> pooled_keepalive; - if (use_native_pool) { - int device_id = -1; - check_cuda(cudaGetDevice(&device_id), "cudaGetDevice native nvcomp scratch"); - auto comp_ptrs = acquire_native_device_allocation(device_id, n * sizeof(void*)); - auto out_ptrs = acquire_native_device_allocation(device_id, n * sizeof(void*)); - auto comp_sizes = acquire_native_device_allocation(device_id, n * sizeof(std::size_t)); - auto out_sizes_buf = acquire_native_device_allocation(device_id, n * sizeof(std::size_t)); - d_comp_ptrs = comp_ptrs.data; - d_out_ptrs = out_ptrs.data; - d_comp_sizes = comp_sizes.data; - d_out_sizes = out_sizes_buf.data; - pooled_keepalive.push_back(std::move(comp_ptrs.owner)); - pooled_keepalive.push_back(std::move(out_ptrs.owner)); - pooled_keepalive.push_back(std::move(comp_sizes.owner)); - pooled_keepalive.push_back(std::move(out_sizes_buf.owner)); - if (temp_bytes > 0) { - auto temp = acquire_native_device_allocation(device_id, temp_bytes); - d_temp = temp.data; - pooled_keepalive.push_back(std::move(temp.owner)); - } - } else { - check_cuda(cudaMallocAsync(&d_comp_ptrs, n * sizeof(void*), stream), "alloc comp_ptrs"); - check_cuda(cudaMallocAsync(&d_out_ptrs, n * sizeof(void*), stream), "alloc out_ptrs"); - check_cuda( - cudaMallocAsync(&d_comp_sizes, n * sizeof(std::size_t), stream), "alloc comp_sizes"); - check_cuda( - cudaMallocAsync(&d_out_sizes, n * sizeof(std::size_t), stream), "alloc out_sizes"); - if (temp_bytes > 0) { - check_cuda(cudaMallocAsync(&d_temp, temp_bytes, stream), "alloc temp"); + auto release_async_scratch = [&]() { + for (void* ptr : {d_temp, + d_comp_ptrs, + d_out_ptrs, + d_comp_sizes, + d_out_sizes, + d_statuses, + d_actual_sizes}) { + if (ptr != nullptr) { + cudaFreeAsync(ptr, stream); + } } - } + }; try { + if (use_native_pool) { + int device_id = -1; + check_cuda(cudaGetDevice(&device_id), "cudaGetDevice native nvcomp scratch"); + auto comp_ptrs = acquire_native_device_allocation(device_id, n * sizeof(void*)); + auto out_ptrs = acquire_native_device_allocation(device_id, n * sizeof(void*)); + auto comp_sizes = acquire_native_device_allocation(device_id, n * sizeof(std::size_t)); + auto out_sizes_buf = + acquire_native_device_allocation(device_id, n * sizeof(std::size_t)); + auto statuses = acquire_native_device_allocation(device_id, n * sizeof(nvcompStatus_t)); + auto actual_sizes = + acquire_native_device_allocation(device_id, n * sizeof(std::size_t)); + d_comp_ptrs = comp_ptrs.data; + d_out_ptrs = out_ptrs.data; + d_comp_sizes = comp_sizes.data; + d_out_sizes = out_sizes_buf.data; + d_statuses = statuses.data; + d_actual_sizes = actual_sizes.data; + pooled_keepalive.push_back(std::move(comp_ptrs.owner)); + pooled_keepalive.push_back(std::move(out_ptrs.owner)); + pooled_keepalive.push_back(std::move(comp_sizes.owner)); + pooled_keepalive.push_back(std::move(out_sizes_buf.owner)); + pooled_keepalive.push_back(std::move(statuses.owner)); + pooled_keepalive.push_back(std::move(actual_sizes.owner)); + if (temp_bytes > 0) { + auto temp = acquire_native_device_allocation(device_id, temp_bytes); + d_temp = temp.data; + pooled_keepalive.push_back(std::move(temp.owner)); + } + } else { + check_cuda(cudaMallocAsync(&d_comp_ptrs, n * sizeof(void*), stream), "alloc comp_ptrs"); + check_cuda(cudaMallocAsync(&d_out_ptrs, n * sizeof(void*), stream), "alloc out_ptrs"); + check_cuda( + cudaMallocAsync(&d_comp_sizes, n * sizeof(std::size_t), stream), + "alloc comp_sizes"); + check_cuda( + cudaMallocAsync(&d_out_sizes, n * sizeof(std::size_t), stream), "alloc out_sizes"); + check_cuda( + cudaMallocAsync(&d_statuses, n * sizeof(nvcompStatus_t), stream), "alloc statuses"); + check_cuda( + cudaMallocAsync(&d_actual_sizes, n * sizeof(std::size_t), stream), + "alloc actual sizes"); + if (temp_bytes > 0) { + check_cuda(cudaMallocAsync(&d_temp, temp_bytes, stream), "alloc temp"); + } + } + check_cuda( cudaMemcpyAsync( d_comp_ptrs, h_comp_ptrs.data(), n * sizeof(void*), cudaMemcpyHostToDevice, stream), @@ -218,43 +250,44 @@ py::object batch_deflate_decompress_impl( // can progress while the async kernel launches + returns. { py::gil_scoped_release release; + // Supply both output buffers for nvCOMP 5.2 compatibility. check_nvcomp( nvcompBatchedDeflateDecompressAsync( static_cast(d_comp_ptrs), static_cast(d_comp_sizes), static_cast(d_out_sizes), - nullptr, // device_uncompressed_chunk_bytes — optional for CUDA backend + static_cast(d_actual_sizes), n, d_temp, temp_bytes, static_cast(d_out_ptrs), opts, - nullptr, // device_statuses — optional + static_cast(d_statuses), stream), "nvcompBatchedDeflateDecompressAsync"); } + + if (use_native_pool) { + auto holder = std::make_unique>>(); + py::capsule owner(holder.get(), [](void* p) { + delete reinterpret_cast>*>(p); + }); + holder->swap(pooled_keepalive); + holder.release(); + return owner; + } } catch (...) { if (use_native_pool) { cudaStreamSynchronize(stream); + } else { + release_async_scratch(); } throw; } - if (use_native_pool) { - auto holder = new std::vector>(std::move(pooled_keepalive)); - return py::capsule(holder, [](void* p) { - delete reinterpret_cast>*>(p); - }); - } - // Stream-ordered frees — nvcomp consumes its inputs asynchronously and // these frees will be sequenced after the kernel completes. - if (d_temp) - cudaFreeAsync(d_temp, stream); - cudaFreeAsync(d_comp_ptrs, stream); - cudaFreeAsync(d_out_ptrs, stream); - cudaFreeAsync(d_comp_sizes, stream); - cudaFreeAsync(d_out_sizes, stream); + release_async_scratch(); return py::none(); } diff --git a/tests/xdr/test_deflate_gpu.py b/tests/xdr/test_deflate_gpu.py new file mode 100644 index 00000000..17afc60d --- /dev/null +++ b/tests/xdr/test_deflate_gpu.py @@ -0,0 +1,188 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Raw DEFLATE alignment and asynchronous decode ownership.""" + +from __future__ import annotations + +import gzip +import struct +import zlib +from types import SimpleNamespace + +import numpy as np +import pytest + +from cuphoton.xdr import nvcomp_batch + + +@pytest.fixture +def cp(): + module = pytest.importorskip("cupy") + try: + if module.cuda.runtime.getDeviceCount() == 0: + pytest.skip("CUDA device unavailable") + except module.cuda.runtime.CUDARuntimeError: + pytest.skip("CUDA runtime unavailable") + ext = nvcomp_batch._try_get_cpp_ext() + if ext is None or not hasattr(ext, "batch_deflate_decompress"): + pytest.skip("native DEFLATE extension unavailable") + return module + + +@pytest.mark.parametrize("pooled", [False, True]) +@pytest.mark.parametrize("preparsed", [False, True]) +def test_native_gzip_decodes_optional_headers(cp, pooled, preparsed): + payloads, originals, header_sizes = [], [], [] + for flags in range(32): + original = bytes(range(256)) * 9 + bytes([flags]) * 11 + payload = _gzip_with_flags(original, flags) + assert gzip.decompress(payload) == original + payloads.append(payload) + originals.append(original) + header_sizes.append(nvcomp_batch._parse_gzip_header_len(payload)) + lengths = np.array([len(payload) for payload in payloads], dtype="i8") + offsets = np.cumsum(lengths) - lengths + sizes = np.array([len(value) for value in originals], dtype="i8") + # Dense tile starts deliberately do not obey raw DEFLATE's alignment. + assert np.any(offsets % 4 != 0) + stream = cp.cuda.Stream(non_blocking=True) + owners = [] + with stream: + data = cp.asarray(np.frombuffer(b"".join(payloads), dtype="u1")) + result, out_offsets = nvcomp_batch.gpu_gzip_decompress_batch( + data, + offsets, + lengths, + sizes, + use_cpp_helper=True, + use_native_pool=pooled, + keepalive=owners, + header_sizes=header_sizes if preparsed else None, + ) + stream.synchronize() + np.testing.assert_array_equal(out_offsets, np.cumsum(sizes) - sizes) + np.testing.assert_array_equal( + cp.asnumpy(result), np.frombuffer(b"".join(originals), dtype="u1") + ) + + +def _gzip_with_flags(data, flags): + header = b"\x1f\x8b\x08" + bytes([flags]) + bytes(6) + if flags & 4: + extra = b"extra\x00payload" + header += struct.pack("