From 1730e2a0036dba75eff557bf570f583d0539ca1a Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Tue, 1 Sep 2026 20:58:14 -0700 Subject: [PATCH] Bound CUDA PUC chunk size to available GPU memory bayesian_blocks (the default discretizer) has no upper bound on the number of bins it selects, and on datasets with many samples (e.g. single-cell data with hundreds of thousands of cells) it can pick thousands of bins per gene. compute_puc_full_cuda's intermediate counts_chunk_gpu/si_chunk_gpu buffers scale with k_bins^2 and k_bins respectively, so the previous fixed chunk_size=256 could request far more GPU memory than is available even for small gene panels (confirmed empirically: a 43-gene x 265,188-sample dataset triggered a 117 GiB allocation request). Size the chunk to fit within the GPU memory actually free at call time instead, clamped to the previous default of at most 256, and raise an actionable error (suggesting uniform_width with a fixed bin count, or the CPU backend) if even a single-gene chunk would not fit, rather than surfacing a raw CUDA out-of-memory error. Verified against test/cuda_numeric_tests.jl's CPU/GPU equivalence, symmetry, and determinism checks (unchanged), and against the previously-failing 43-gene/265,188-sample bayesian_blocks case (now completes, auto-selecting chunk_size=11 for k_bins=1690). Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_0151Ww8ZmYpes58vMJNbXHpA --- ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl | 54 ++++++++++++++++++++------ 1 file changed, 43 insertions(+), 11 deletions(-) diff --git a/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl b/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl index 3df5595..7fb6e1a 100644 --- a/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl +++ b/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl @@ -156,23 +156,24 @@ end GPU implementation of [`FastPIDC.compute_puc_full`](@ref): computes the full pairwise MI matrix and pre-context PUC matrix for `nodes` on the GPU, -processing genes along the target (`z`) axis in chunks of 256 to bound -device memory use. Moves discretized data and marginal probabilities to -the GPU once, then for each chunk launches +processing genes along the target (`z`) axis in chunks (up to 256 genes at +a time) sized to fit the GPU memory currently free, to bound device memory +use even when the discretizer has picked a large number of bins (see the +chunk-sizing comment in the implementation). Moves discretized data and +marginal probabilities to the GPU once, then for each chunk launches [`joint_counts_kernel_chunked!`](@ref), [`mi_si_kernel_chunked!`](@ref) and [`puc_accumulation_kernel_chunked!`](@ref) in sequence, symmetrizing the resulting PUC matrix before returning both matrices to the CPU. `config.verbose` enables progress printouts; `base` is currently unused -(mutual information is always computed in base 2 on the GPU path). +(mutual information is always computed in base 2 on the GPU path). Raises +an `ErrorException` with a suggested remedy if even a single-gene chunk +would not fit in the currently-free GPU memory. """ function FastPIDC.compute_puc_full_cuda(nodes, config, base) num_nodes = length(nodes) num_samples = length(nodes[1].binned_values) k_bins = maximum(n -> n.number_of_bins, nodes) - - # Chunk configuration (512 genes at a time) - chunk_size = 256 - + # Prepare static data on CPU and move to GPU data_cpu = zeros(Int32, num_samples, num_nodes) marginals_cpu = zeros(Float64, k_bins, num_nodes) @@ -188,14 +189,45 @@ function FastPIDC.compute_puc_full_cuda(nodes, config, base) # Global output matrices puc_scores_gpu = CUDA.zeros(Float64, num_nodes, num_nodes) mi_matrix_gpu = CUDA.zeros(Float64, num_nodes, num_nodes) - + + # Chunk configuration: `counts_chunk_gpu`/`si_chunk_gpu` scale as + # k_bins^2 * num_nodes * chunk_size / k_bins * num_nodes * chunk_size + # respectively, so an adaptive discretizer that picks a large k_bins + # (e.g. "bayesian_blocks", the package default, on a dataset with many + # samples -- it has no upper bound on the number of bins it selects) + # can make even a modest fixed chunk_size request far larger than the + # GPU's total memory. Size the chunk to fit within what's actually free + # right now instead of always requesting a fixed chunk_size=256, and + # fail with an actionable message rather than a raw CUDA OOM error if + # even a single-gene chunk doesn't fit. + bytes_per_chunk_col = + k_bins^2 * num_nodes * sizeof(Int32) + # counts_chunk_gpu + k_bins * num_nodes * sizeof(Float64) # si_chunk_gpu + free_bytes, _ = CUDA.memory_info() + safety_factor = 0.8 # headroom for the fixed buffers above + allocator overhead/fragmentation + max_chunk_size = floor(Int, free_bytes * safety_factor / bytes_per_chunk_col) + chunk_size = clamp(max_chunk_size, 1, min(256, num_nodes)) + + if max_chunk_size < 1 + error( + "compute_puc_full_cuda: even a single-gene chunk would require " * + "$(round(bytes_per_chunk_col / 2^30, digits = 2)) GiB of GPU memory " * + "(only $(round(free_bytes * safety_factor / 2^30, digits = 2)) GiB " * + "usable), because the discretizer selected k_bins=$k_bins bins per " * + "gene. This is usually caused by an adaptive discretizer (e.g. " * + "\"bayesian_blocks\", the default) picking an unbounded number of " * + "bins on a dataset with many samples. Try discretizer=\"uniform_width\" " * + "with a fixed, small number_of_bins (e.g. 10-20), or config.backend = :cpu.", + ) + end + # Chunked intermediate buffers (Pre-allocated once!) counts_chunk_gpu = CUDA.zeros(Int32, k_bins, k_bins, num_nodes, chunk_size) si_chunk_gpu = CUDA.zeros(Float64, k_bins, num_nodes, chunk_size) if config.verbose - println("[FastPIDC] GPU Chunked PUC: Processing $num_nodes x $num_nodes pairs...") - println("[FastPIDC] Using chunk size of $chunk_size (approx. $(ceil(Int, num_nodes/chunk_size)) iterations)") + println("[FastPIDC] GPU Chunked PUC: Processing $num_nodes x $num_nodes pairs (k_bins=$k_bins)...") + println("[FastPIDC] Using chunk size of $chunk_size (approx. $(ceil(Int, num_nodes/chunk_size)) iterations), sized to fit $(round(free_bytes / 2^30, digits = 2)) GiB free GPU memory") end # Iterate over the Z-axis in chunks