Skip to content

Fix xDR Deflate alignment and nvCOMP output buffers #81

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tpn merged 1 commit into from
Sep 30, 2026
Merged

Fix xDR Deflate alignment and nvCOMP output buffers #81

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
177 changes: 126 additions & 51 deletions src/cuphoton/xdr/nvcomp_batch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@

import numpy as np

_REPACK_DEFLATE_KERNEL = None
_DEFLATE_CODEC = None
_CPP_EXT = None
_CPP_EXT_PROBED = False
Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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
129 changes: 81 additions & 48 deletions src/cuphoton/xdr/src/nvcomp_batch_ext.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,9 @@ py::object batch_deflate_decompress_impl(
}
h_comp_ptrs[i] =
reinterpret_cast<const void*>(d_concat_ptr + static_cast<std::size_t>(ro(i)));
if (reinterpret_cast<std::uintptr_t>(h_comp_ptrs[i]) % 4 != 0) {
throw std::invalid_argument("raw DEFLATE inputs must be 4-byte aligned");
}
h_out_ptrs[i] = reinterpret_cast<void*>(d_out_ptr + static_cast<std::size_t>(oo(i)));
h_comp_sizes[i] = static_cast<std::size_t>(ln(i));
h_out_sizes[i] = static_cast<std::size_t>(os(i));
Expand All @@ -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(
Expand All @@ -156,39 +160,67 @@ py::object batch_deflate_decompress_impl(
"get temp size");

std::vector<std::shared_ptr<void>> 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),
Expand Down Expand Up @@ -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<const void* const*>(d_comp_ptrs),
static_cast<const std::size_t*>(d_comp_sizes),
static_cast<const std::size_t*>(d_out_sizes),
nullptr, // device_uncompressed_chunk_bytes — optional for CUDA backend
static_cast<std::size_t*>(d_actual_sizes),
n,
d_temp,
temp_bytes,
static_cast<void* const*>(d_out_ptrs),
opts,
nullptr, // device_statuses — optional
static_cast<nvcompStatus_t*>(d_statuses),
stream),
"nvcompBatchedDeflateDecompressAsync");
}

if (use_native_pool) {
auto holder = std::make_unique<std::vector<std::shared_ptr<void>>>();
py::capsule owner(holder.get(), [](void* p) {
delete reinterpret_cast<std::vector<std::shared_ptr<void>>*>(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::shared_ptr<void>>(std::move(pooled_keepalive));
return py::capsule(holder, [](void* p) {
delete reinterpret_cast<std::vector<std::shared_ptr<void>>*>(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();
}

Expand Down
Loading
Loading