Bound CUDA PUC chunk size to available GPU memory

[aarmey]

Problem

bayesian_blocks (the default discretizer) has no upper bound on the number of bins it selects. On datasets with many samples — e.g. single-cell data with hundreds of thousands of cells, exactly the use case NetworkInference/PIDC targets and that GPU acceleration is meant to make tractable — it can pick thousands of bins per gene.

compute_puc_full_cuda's intermediate buffers scale as:

• counts_chunk_gpu: k_bins^2 * num_nodes * chunk_size (Int32)
• si_chunk_gpu: k_bins * num_nodes * chunk_size (Float64)

so the previous fixed chunk_size = 256 could request far more GPU memory than exists, even for small gene panels. Confirmed empirically: infer_network with the CUDA backend on a 43-gene x 265,188-sample matrix (continuous, mostly-unique log-normalized single-cell values) triggered:

ERROR: LoadError: Out of GPU memory trying to allocate 117.123 GiB
Effective GPU memory usage: 70.94% (16.683 GiB/23.516 GiB)
...
 [11] compute_puc_full_cuda(nodes::Vector{Node}, config::PIDCConfig, base::Int64)
    @ FastPIDCCUDAExt ~/.julia/packages/FastPIDC/kNsZf/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl:193

(bayesian_blocks had picked k_bins ≈ 1690 for this data.) Repository issues are disabled, so opening this directly as a PR rather than filing a report first.

Fix

Size the chunk to fit the GPU memory actually free at call time (CUDA.memory_info()), clamped to the previous default of at most 256, with an 0.8 safety factor for the fixed-size buffers (data_gpu, marginals_gpu, puc_scores_gpu, mi_matrix_gpu) and allocator overhead. If even a single-gene chunk (chunk_size = 1) wouldn't fit, raise a clear ErrorException naming the cause (large k_bins from an adaptive discretizer) and suggesting remedies (uniform_width with a fixed bin count, or config.backend = :cpu), instead of surfacing a raw CUDA OOM stacktrace.

No public API changes — compute_puc_full_cuda's signature and behavior for well-behaved (small k_bins) inputs are unchanged; verified below.

Testing

• test/cuda_numeric_tests.jl's CPU/GPU equivalence, symmetry, and determinism checks pass unchanged on toy_small_200.txt (small k_bins, so chunk_size still resolves to the old default of 200/256 there — confirmed the sizing logic is a no-op in the well-behaved case).
• The previously-failing 43-gene/265,188-sample case now completes: k_bins=1690 → auto-selected chunk_size=11 (4 chunk iterations) instead of the 117 GiB allocation attempt.

Suggested follow-ups (not included here, happy to open separately if useful)

• A hard cap on bayesian_blocks' bin count (or a max_bins keyword) so CPU-backend runs and overall runtime aren't also blown out by very large k_bins.
• Documenting the k_bins^2 * num_nodes * chunk_size GPU memory scaling and recommending uniform_width explicitly for large-sample-count single-cell inputs in the README/CLI help.

🤖 Generated with Claude Code

https://claude.ai/code/session_0151Ww8ZmYpes58vMJNbXHpA