GPU: kernels read host buffers in place on unified-memory systems

gpu_direct_host_read() (CUDAWrapper) is true where the device accesses pageable
memory through the host's page tables (cudaDevAttrPageableMemoryAccess and
...UsesHostPageTables: GB10, GH200, GB200), asked once per device;
JFJOCH_FORCE_DIRECT_HOST_READ forces it on for testing with page-locked inputs.
Where it holds:
- BSLZ4DecoderGPU: the LZ4/fused kernels and the tail copy read the reader's
  compressed chunk directly, no pinned staging memcpy and no H2D copy;
- ImagePreprocessorGPU: Analyze and the UINT16_WIDENED path hand the host
  buffer to preprocess_kernel (if aligned to the pixel type), and nothing is
  page-locked;
- RotationScaleMergeGPU::IngestUpload: IngestExtractKernel reads the pinned
  stage in place, without device staging buffers; the stage is reused only
  after the kernel has run.
Discrete cards take the unchanged path. Every caller synchronises its stream
before the host buffer goes away. DirectHostReadGPUTest checks the three
sites against the CPU on page-locked inputs, in either mode.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
2026-10-14 21:49:37 +02:00
co-authored by Claude Opus 5.5
parent fe4db26716
commit e3b33591ca
10 changed files with 263 additions and 29 deletions
@@ -7,6 +7,7 @@
#include "ImagePreprocessorGPU.h"
#include "BSLZ4DecodeWarp.h"
#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"
namespace {
@@ -14,6 +15,12 @@ namespace {
if (val != cudaSuccess)
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
}
// A kernel reads pixels of type T, so a host buffer it reads in place must be aligned to T. A
// decompressed buffer always is; an image inside a received message need not be.
bool aligned_to(const void *ptr, size_t bytes) {
return reinterpret_cast<uintptr_t>(ptr) % bytes == 0;
}
}
template<class T>
@@ -338,6 +345,7 @@ ImagePreprocessorGPU::ImagePreprocessorGPU(const DiffractionExperiment &experime
threads = 128;
blocks = 4 * prop.multiProcessorCount;
direct_host_read = gpu_direct_host_read();
}
float ImagePreprocessorGPU::GetLastDecompressionTime_s() const {
@@ -359,7 +367,7 @@ ImagePreprocessorGPU::~ImagePreprocessorGPU() {
}
void ImagePreprocessorGPU::PinInputRegion(const void *ptr, size_t bytes) {
if (ptr == pinned_input && bytes == pinned_input_bytes)
if (direct_host_read || (ptr == pinned_input && bytes == pinned_input_bytes))
return;
if (pinned_input)
cuda_err(cudaHostUnregister(const_cast<void *>(pinned_input)));
@@ -385,7 +393,8 @@ void ImagePreprocessorGPU::PinInputBuffer(std::vector<uint8_t> &buffer, size_t s
// Unregister before the resize, which can move the buffer.
input_reg.unregister();
buffer.resize(size);
input_reg.rebind(buffer);
if (!direct_host_read)
input_reg.rebind(buffer);
}
ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *image_ptr,
@@ -527,6 +536,11 @@ ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed
const uint8_t *input,
T err_value,
T sat_value) {
// Read in place where the card reads host memory directly; the input outlives this call, which
// waits for the stream.
if (direct_host_read && aligned_to(input, sizeof(T)))
return AnalyzeOnDevice<T>(processed_image, input, err_value, sat_value);
// Allocated here rather than in the constructor: only the host-upload path needs it, and the
// device-decode path - which is what BSHUF_LZ4 images take - never touches it. Overshoot to
// 4 bytes per pixel so a 1- or 2-byte image fits the same buffer.
@@ -537,14 +551,11 @@ ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed
// every blocking stream in the process, which serialised all workers behind whichever one was
// uploading. The stream is synchronised at the end of this function, so the ordering is unchanged.
cuda_err(cudaMemcpyAsync(gpu_decompressed_image, input, npixels * sizeof(T), cudaMemcpyHostToDevice, *stream));
return AnalyzeOnDevice<T>(processed_image, err_value, sat_value);
return AnalyzeOnDevice<T>(processed_image, gpu_decompressed_image.get(), err_value, sat_value);
}
ImageStatistics ImagePreprocessorGPU::AnalyzeStoredInt32(ImagePreprocessorBuffer &processed_image,
const CompressedImage &image) {
if (!gpu_decompressed_image.get())
gpu_decompressed_image = CudaDevicePtr<uint8_t>(npixels * sizeof(uint32_t));
if (image.GetCompressionAlgorithm() == CompressionAlgorithm::UINT16_WIDENED) {
// Uploaded as stored, two bytes a pixel, and widened by the preprocessing kernel itself: with
// int32's error marker and int32's saturation, each pixel is decided exactly as the int32 image
@@ -552,31 +563,40 @@ ImageStatistics ImagePreprocessorGPU::AnalyzeStoredInt32(ImagePreprocessorBuffer
last_decode = LastDecode::None;
if (image.GetCompressedSize() != npixels * sizeof(uint16_t))
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mismatch in size");
// Read in place where the card reads host memory directly: nothing is page-locked, so engines
// sharing the frame do not interfere.
if (direct_host_read && aligned_to(image.GetCompressed(), sizeof(uint16_t)))
return AnalyzeOnDevice<uint16_t>(processed_image, image.GetCompressed(), INT32_MIN, INT32_MAX);
// Staged through this engine's own page-locked buffer, as the CBF and BSLZ4 decoders stage
// theirs, and not page-locked in place: the pre-scan hands one kept frame to several engines
// at once, and an engine that unregisters a frame another one is still copying from takes the
// whole context down. AnalyzeOnDevice waits for the stream, so the buffer is free again here.
if (!gpu_decompressed_image.get())
gpu_decompressed_image = CudaDevicePtr<uint8_t>(npixels * sizeof(uint32_t));
if (!host_stored.get())
host_stored = CudaHostPtr<uint8_t>(npixels * sizeof(uint16_t));
std::memcpy(host_stored.get(), image.GetCompressed(), image.GetCompressedSize());
cuda_err(cudaMemcpyAsync(gpu_decompressed_image, host_stored.get(), image.GetCompressedSize(),
cudaMemcpyHostToDevice, *stream));
return AnalyzeOnDevice<uint16_t>(processed_image, INT32_MIN, INT32_MAX);
return AnalyzeOnDevice<uint16_t>(processed_image, gpu_decompressed_image.get(), INT32_MIN, INT32_MAX);
}
if (!gpu_decompressed_image.get())
gpu_decompressed_image = CudaDevicePtr<uint8_t>(npixels * sizeof(uint32_t));
last_decode = LastDecode::CBF;
if (!cbf_decoder)
cbf_decoder = std::make_unique<CBFDecoderGPU>(stream);
cbf_decoder->Decode(image, reinterpret_cast<int32_t *>(gpu_decompressed_image.get()));
const ImageStatistics ret = AnalyzeOnDevice<int32_t>(processed_image, INT32_MIN, INT32_MAX);
const ImageStatistics ret = AnalyzeOnDevice<int32_t>(processed_image, gpu_decompressed_image.get(),
INT32_MIN, INT32_MAX);
cbf_decoder->ThrowIfDecodeFailed();
return ret;
}
// Everything after the image is on the device, shared by the host-upload and the device-decode
// entry points so the two cannot drift apart.
// Everything after the image is on the device - or in host memory the card reads directly - shared by
// the host-upload and the device-decode entry points so the two cannot drift apart.
template<class T>
ImageStatistics ImagePreprocessorGPU::AnalyzeOnDevice(ImagePreprocessorBuffer &processed_image,
ImageStatistics ImagePreprocessorGPU::AnalyzeOnDevice(ImagePreprocessorBuffer &processed_image, const void *input,
int64_t err_value, int64_t sat_value) {
if (sat_value > saturation_limit)
sat_value = saturation_limit;
@@ -584,7 +604,7 @@ ImageStatistics ImagePreprocessorGPU::AnalyzeOnDevice(ImagePreprocessorBuffer &p
cpu_stats[0] = ImageStatistics{.max_value = INT64_MIN, .min_value = INT64_MAX};
cuda_err(cudaMemcpyAsync(gpu_stats, cpu_stats.data(), sizeof(ImageStatistics), cudaMemcpyHostToDevice, *stream));
preprocess_kernel<T> <<< blocks, threads, 0, *stream >>>(
reinterpret_cast<const T *>(gpu_decompressed_image.get()),
reinterpret_cast<const T *>(input),
gpu_mask->get(),
processed_image.getGPUBuffer(),
gpu_stats,