From 7d731ca91e66d3c42c8194506d1242691912fa90 Mon Sep 17 00:00:00 2001 From: jjschirle Date: Fri, 11 Sep 2026 11:17:29 -0700 Subject: [PATCH 1/6] C kernel memory allocation bugfix. Updated test suite for it. CLI bugfix. Updated dependencies --- CLI_fastpidc.jl | 85 +++++++++++++--- Project.toml | 8 +- ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl | 33 +++++- python/pyproject.toml | 6 +- python/src/fastpidc/cuda.py | 33 ++++++ python/src/fastpidc/kernels/pidc_kernels.cu | 92 +++++++++++++---- python/tests/test_cuda.py | 103 +++++++++++++++++++ test/cli_argument_tests.jl | 78 +++++++++++++++ test/cuda_largemem_tests.jl | 105 ++++++++++++++++++++ test/runtests.jl | 2 + 10 files changed, 510 insertions(+), 35 deletions(-) create mode 100644 test/cli_argument_tests.jl create mode 100644 test/cuda_largemem_tests.jl diff --git a/CLI_fastpidc.jl b/CLI_fastpidc.jl index 1734f2f..9cf350d 100644 --- a/CLI_fastpidc.jl +++ b/CLI_fastpidc.jl @@ -34,6 +34,55 @@ function parse_args() return args end +# Every flag main() actually reads (see the `get`/`haskey` calls below). Kept +# in sync with those by hand, since parse_args() has no way to know which keys +# are meaningful - it happily stores (and silently drops) anything. +const VALID_ARG_KEYS = Set([ + "help", + "infile", + "outfile", + "delim", + "discretizer", + "estimator", + "n-bins", + "base", + "backend", + "bb-backend", + "output-format", + "dump-mi-path", + "dump-puc-path", + "verbose", +]) + +""" + validate_args(args) + +`parse_args` accepts any `--key value` pair, so a misspelled or +wrongly-punctuated flag was +previously stored under a key `main()` never reads, silently keeping its +default rather than erroring - a caller could ask for `--discretizer +uniform_width --n-bins 6` and get 10 bins with no indication anything was +wrong. This checks every parsed key against [`VALID_ARG_KEYS`](@ref) and +errors out, naming the likely intended flag when the mismatch is only a +`-`/`_` swap. +""" +function validate_args(args::Dict{String,String}) + unknown = sort(collect(setdiff(keys(args), VALID_ARG_KEYS))) + isempty(unknown) && return nothing + + lines = String[] + for key in unknown + swapped = replace(key, "-" => "_", "_" => "-") + hint = swapped in VALID_ARG_KEYS ? " (did you mean --$swapped?)" : "" + push!(lines, " --$key$hint") + end + error( + "Unrecognized command-line argument(s):\n" * + join(lines, "\n") * + "\nRun with --help to see the full list of supported arguments.", + ) +end + function parse_delim(s::AbstractString) s_l = lowercase(strip(s)) if s_l == "space" || s_l == " " @@ -76,7 +125,7 @@ Basic options: Default: 'bayesian_blocks' --estimator STR e.g. 'maximum_likelihood' Default: 'maximum_likelihood' - --n_bins INT Number of bins (ignored by bayesian_blocks). Default: 10 + --n-bins INT Number of bins (ignored by bayesian_blocks). Default: 10 --base INT Log base for MI (2, e, 10). Default: 2 Execution / Environment: @@ -109,6 +158,8 @@ function main() return end + validate_args(args) + # ----------------- Required arguments ----------------- infile = get(args, "infile", nothing) outfile = get(args, "outfile", nothing) @@ -126,7 +177,7 @@ function main() delim = parse_delim(delim_str) discretizer = get(args, "discretizer", "bayesian_blocks") estimator = get(args, "estimator", "maximum_likelihood") - n_bins = parse(Int, get(args, "n_bins", "10")) + n_bins = parse(Int, get(args, "n-bins", "10")) base = parse(Int, get(args, "base", "2")) verbose_flag = parse_bool(get(args, "verbose", "false")) @@ -185,7 +236,7 @@ function main() println(" delim = $delim_str") println(" discretizer = $discretizer") println(" estimator = $estimator") - println(" n_bins = $n_bins") + println(" n-bins = $n_bins") println(" base = $base") println(" backend = $(cfg.backend)") println(" bb_backend = $(cfg.bb_backend)") @@ -222,15 +273,21 @@ function main() @say @sprintf("All done. Total runtime: %.1f s", t_total) end -# Ensure we get a traceback for errors -try - main() -catch e - bt = catch_backtrace() - @say "ERROR: $(sprint(showerror, e))" - println("\nStacktrace:") - Base.show_backtrace(stdout, bt) - println() - @say "Tip: If the error mentions discretization, try --discretizer uniform_width and --n_bins 10-20." - exit(1) +# Only run when invoked as a script (`julia CLI_fastpidc.jl ...`), not when +# `include()`d - e.g. by a test file exercising `parse_args`/`validate_args` +# in isolation - which would otherwise execute `main()` against the including +# process's `ARGS` and `exit(1)` out from under it on any error. +if abspath(PROGRAM_FILE) == @__FILE__ + # Ensure we get a traceback for errors + try + main() + catch e + bt = catch_backtrace() + @say "ERROR: $(sprint(showerror, e))" + println("\nStacktrace:") + Base.show_backtrace(stdout, bt) + println() + @say "Tip: If the error mentions discretization, try --discretizer uniform_width and --n-bins 10-20." + exit(1) + end end diff --git a/Project.toml b/Project.toml index ff75e52..00ce13e 100644 --- a/Project.toml +++ b/Project.toml @@ -25,9 +25,15 @@ Statistics = "1.11.1" julia = "≥ 1.0.0" [extras] +Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["DelimitedFiles", "LinearAlgebra", "Test"] +# Dates and Printf are needed because test/cli_argument_tests.jl includes +# CLI_fastpidc.jl (to exercise its argument parsing/validation in isolation), +# which `using`s both; they are otherwise resolvable stdlibs, not real deps of +# the package itself. +test = ["Dates", "DelimitedFiles", "LinearAlgebra", "Printf", "Test"] diff --git a/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl b/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl index 7cd04e0..32e1add 100644 --- a/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl +++ b/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl @@ -92,6 +92,26 @@ end # --- Host implementation --- +""" + _MAX_K_BINS + +Largest bins-per-gene the shared kernels can index. They compute every flat +buffer offset in 64-bit, with one exception: `joint_counts_kernel` forms the +bin-pair index `u * k_bins + v` in Int32, which stays exact only while +`k_bins^2` fits in Int32 (isqrt(typemax(Int32)) == 46340). +""" +const _MAX_K_BINS = 46340 + +function _check_kernel_index_limits(k_bins::Integer) + k_bins <= _MAX_K_BINS || error( + "compute_puc_full_cuda: the discretizer selected k_bins=$k_bins bins " * + "per gene, but the CUDA kernels index bin pairs in 32-bit and support " * + "at most $(_MAX_K_BINS). Use discretizer=\"uniform_width\" with a " * + "fixed, small number_of_bins, or config.backend = :cpu.", + ) + return nothing +end + function _smallest_unsigned_type(max_value::Integer) max_value >= 0 || throw(ArgumentError("max_value must be nonnegative")) @@ -122,7 +142,12 @@ 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, matching the kernel source). Raises an `ErrorException` with a suggested remedy if even -a single-gene chunk would not fit in the currently-free GPU memory. +a single-gene chunk would not fit in the currently-free GPU memory, or if +the discretizer selected more than `_MAX_K_BINS` bins per gene. + +The chunked intermediates legitimately exceed 2^31 elements on large gene +sets - `counts` alone is `k_bins^2 * num_nodes * chunk_size` - so the shared +kernels index them in 64-bit (see the indexing contract in the kernel source). Device buffers use Julia's column-major layout with dimensions reversed relative to the kernel source's documented (row-major) shapes - e.g. a @@ -140,6 +165,7 @@ 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) + _check_kernel_index_limits(k_bins) # Prepare static data on CPU and move to GPU. The shared CUDA C kernels use # 0-indexed Int32 bin ids; FastPIDC.jl's bin ids are 1-indexed, so shift @@ -163,6 +189,11 @@ function FastPIDC.compute_puc_full_cuda(nodes, config, base) # joint counts and k_bins * num_nodes * chunk_size for specific information. # Size the target-gene chunk from currently-free memory rather than always # allocating a fixed 256-gene chunk. + # + # This bounds MEMORY AVAILABILITY only - it is not, and must not be turned + # back into, a bound on the flat element index. The kernels index these + # buffers in 64-bit precisely so the chunk can be sized from free memory + # without an int32 element-count cap. bytes_per_chunk_col = k_bins^2 * num_nodes * sizeof(Int32) + # counts_chunk_gpu k_bins * num_nodes * sizeof(Float64) # si_chunk_gpu diff --git a/python/pyproject.toml b/python/pyproject.toml index 31e87fc..be47ba9 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -27,7 +27,11 @@ artifacts = ["src/fastpidc/kernels/*.cu"] [tool.pytest.ini_options] testpaths = ["tests"] -markers = ["julia: cross-validates against a freshly run FastPIDC.jl (requires julia on PATH)"] +markers = [ + "julia: cross-validates against a freshly run FastPIDC.jl (requires julia on PATH)", + "largemem: allocates >8 GiB of GPU memory to exercise flat indices past 2^31 (deselected by default; run with -m largemem on an idle GPU)", +] +addopts = "-m 'not largemem'" [tool.ruff] line-length = 120 diff --git a/python/src/fastpidc/cuda.py b/python/src/fastpidc/cuda.py index 84fe0e5..4f8a37b 100644 --- a/python/src/fastpidc/cuda.py +++ b/python/src/fastpidc/cuda.py @@ -9,6 +9,10 @@ Genes are processed along the target ("z") axis in chunks sized to fit the currently free device memory, mirroring ``FastPIDCCUDAExt.compute_puc_full_cuda`` in FastPIDC.jl. + +Those chunked buffers legitimately exceed 2^31 elements on large gene sets +(``counts`` alone is ``k_bins**2 * n * chunk_size``), so the shared kernels +index them in 64-bit; see the indexing contract in the kernel source. """ from __future__ import annotations @@ -26,6 +30,10 @@ _KERNEL_SOURCE_PATH = Path(__file__).with_name("kernels") / "pidc_kernels.cu" _MAX_CHUNK_SIZE = 256 +# Largest bins-per-gene the shared kernels can index: joint_counts_kernel forms +# the bin-pair index u * k_bins + v in int32, so k_bins**2 must fit there +# (isqrt(2**31 - 1) == 46340). Every other flat offset is 64-bit. +_MAX_K_BINS = 46340 # Headroom for the fixed buffers, allocator overhead and fragmentation, matching # the safety factor used by FastPIDC.jl's CUDA extension. _CHUNK_MEMORY_SAFETY_FACTOR = 0.8 @@ -74,6 +82,24 @@ def _load_module(): return cp.RawModule(code=source, options=("--std=c++11",)) +def _check_kernel_index_limits(k_bins: int) -> None: + """Validate the one kernel limit that is not lifted by 64-bit indexing. + + The shared kernels compute every flat buffer offset in 64-bit, with a single + exception: ``joint_counts_kernel`` forms the bin-pair index + ``u * k_bins + v`` in ``int``, which is exact only while ``k_bins**2`` fits + in int32. FastPIDC.jl's ``_check_kernel_index_limits`` enforces the same + bound. + """ + if k_bins > _MAX_K_BINS: + raise RuntimeError( + f"compute_puc_full_cuda: the discretizer selected k_bins={k_bins} bins per " + f"gene, but the CUDA kernels index bin pairs in 32-bit and support at most " + f'{_MAX_K_BINS}. Use discretizer="uniform_width" with a fixed, small ' + f"number_of_bins, or config.backend = 'cpu'." + ) + + def _chunk_size_for_free_memory(n: int, k_bins: int, free_bytes: int) -> int: """Largest target-gene chunk whose intermediate buffers fit in ``free_bytes`` of device memory, capped at :data:`_MAX_CHUNK_SIZE`. @@ -82,6 +108,10 @@ def _chunk_size_for_free_memory(n: int, k_bins: int, free_bytes: int) -> int: ``k_bins * n`` (specific information) per target gene, so an adaptive discretizer that picks many bins can make even one gene per chunk too large; that case raises instead of failing inside the allocator. + + This bounds *memory availability* only - it is deliberately not a bound on + the flat element index. The kernels index these buffers in 64-bit precisely + so the chunk can be sized from free memory without an int32 element cap. """ bytes_per_chunk_column = k_bins**2 * n * np.dtype(np.int32).itemsize + k_bins * n * np.dtype(np.float64).itemsize usable_bytes = free_bytes * _CHUNK_MEMORY_SAFETY_FACTOR @@ -131,6 +161,9 @@ def compute_puc_full_cuda( n = len(nodes) m = nodes[0].binned_values.size k_bins = max(node.number_of_bins for node in nodes) + # Checked here rather than in _chunk_size_for_free_memory, which is skipped + # entirely when the caller passes an explicit chunk_size. + _check_kernel_index_limits(k_bins) if chunk_size is None: chunk_size = _chunk_size_for_free_memory(n, k_bins, int(cp.cuda.runtime.memGetInfo()[0])) diff --git a/python/src/fastpidc/kernels/pidc_kernels.cu b/python/src/fastpidc/kernels/pidc_kernels.cu index d640104..a34ed03 100644 --- a/python/src/fastpidc/kernels/pidc_kernels.cu +++ b/python/src/fastpidc/kernels/pidc_kernels.cu @@ -29,12 +29,25 @@ // // All arrays are indexed 0-based, row-major (C order), and passed as flat // buffers with the shapes documented per kernel. +// +// Indexing contract: the scalar parameters (n, m, k_bins, z_start, z_chunk_size) +// are int32 on both hosts, but the BUFFERS ARE NOT BOUNDED BY INT32. `counts` +// alone holds k_bins^2 * n * z_chunk_size elements, which exceeds 2^31 on +// ordinary single-cell datasets - 12k genes x 33 bins x a 256-gene chunk is +// 3.4e9 elements, and computing that index in `int` wrapped it negative and +// faulted with CUDA_ERROR_ILLEGAL_ADDRESS. Every composed flat offset below is +// therefore `long long`, hoisted out of the hot loops as a base offset plus a +// loop-invariant stride so the inner bodies cost 64-bit adds rather than 32-bit +// multiply-add chains. The single exception is the bin-pair index +// `u * k_bins + v` in joint_counts_kernel, which stays 32-bit and so requires +// k_bins <= 46340; both hosts check that before launching. extern "C" { // data: (m, n) int32 -- data[s * n + x] = bin id of gene x, sample s // counts: (k_bins, k_bins, n, chunk) int32, zero-initialized by the caller // counts[((u * k_bins + v) * n + x) * chunk + z_local] +// Both of these routinely exceed 2^31 elements; see the indexing contract above. // One thread handles one (x, z_local) pair, looping over all m samples. __global__ void joint_counts_kernel( const int* __restrict__ data, @@ -49,11 +62,24 @@ __global__ void joint_counts_kernel( int z_global = z_start + z_local; if (z_global >= n || x == z_global) return; - for (int s = 0; s < m; ++s) { - int u = data[s * n + x]; - int v = data[s * n + z_global]; + // 64-bit offsets: counts holds k_bins^2 * n * z_chunk_size elements, which + // exceeds 2^31 on ordinary datasets. The strides are loop-invariant, so they + // are hoisted here and the sample loop costs one widening multiply-add. + const long long plane_stride = (long long)n * z_chunk_size; // one (u, v) plane + const long long cell = (long long)x * z_chunk_size + z_local; // this thread's (x, z_local) + + // Walking the sample row also keeps data's m * n index 64-bit. One + // accumulator serves both reads, since they differ only by a fixed column. + long long data_row = 0; + + for (int s = 0; s < m; ++s, data_row += n) { + int u = data[data_row + x]; + int v = data[data_row + z_global]; if (u >= 0 && u < k_bins && v >= 0 && v < k_bins) { - int idx = ((u * k_bins + v) * n + x) * z_chunk_size + z_local; + // u, v < k_bins was just checked, so u * k_bins + v < k_bins^2, which + // both hosts keep inside int (k_bins <= 46340); plane_stride carries + // the 64-bit range. + long long idx = (long long)(u * k_bins + v) * plane_stride + cell; atomicAdd(&counts[idx], 1); } } @@ -65,6 +91,7 @@ __global__ void joint_counts_kernel( // si_matrix: (k_bins, n, chunk) float64 -- si_matrix[(v * n + x) * chunk + z_local] // specific information of source x with respect to target z_global, // at target bin v. +// counts may exceed 2^31 elements; see the indexing contract above. // One thread handles one (x, z_local) pair. __global__ void mi_si_kernel( const int* __restrict__ counts, @@ -81,19 +108,35 @@ __global__ void mi_si_kernel( int z_global = z_start + z_local; if (z_global >= n || x == z_global) return; + // 64-bit offsets, hoisted once per thread. counts (k_bins, k_bins, n, chunk) + // advances by plane_stride per v and u_stride per u; si_matrix (k_bins, n, + // chunk) advances by plane_stride per v. Both loops then cost 64-bit adds + // only. The advances live in the for-increment clause so the `continue`s + // below cannot skip them. + const long long plane_stride = (long long)n * z_chunk_size; + const long long u_stride = plane_stride * k_bins; + const long long cell = (long long)x * z_chunk_size + z_local; + double inv_m = 1.0 / (double)m; double mi_val = 0.0; - for (int v = 0; v < k_bins; ++v) { - double p_z_v = marginals[v * n + z_global]; + long long counts_v = cell; // counts[((0 * k_bins + v) * n + x) * chunk + z_local] + long long si_idx = cell; // si_matrix[(v * n + x) * chunk + z_local] + long long marg_z = z_global; // marginals[v * n + z_global] + + for (int v = 0; v < k_bins; ++v, + counts_v += plane_stride, si_idx += plane_stride, marg_z += n) { + double p_z_v = marginals[marg_z]; if (p_z_v <= 0.0) continue; double si_v = 0.0; - for (int u = 0; u < k_bins; ++u) { - double p_x_u = marginals[u * n + x]; + long long counts_uv = counts_v; // u = 0 + long long marg_x = x; // marginals[u * n + x] + for (int u = 0; u < k_bins; ++u, counts_uv += u_stride, marg_x += n) { + double p_x_u = marginals[marg_x]; if (p_x_u <= 0.0) continue; - int c_uv = counts[((u * k_bins + v) * n + x) * z_chunk_size + z_local]; + int c_uv = counts[counts_uv]; double p_uv = (double)c_uv * inv_m; if (p_uv > 0.0) { @@ -102,10 +145,10 @@ __global__ void mi_si_kernel( si_v += p_u_cond_v * log2(p_u_cond_v / p_x_u); } } - si_matrix[(v * n + x) * z_chunk_size + z_local] = si_v; + si_matrix[si_idx] = si_v; } - mi_matrix[x * n + z_global] = mi_val; + mi_matrix[(long long)x * n + z_global] = mi_val; } // si_matrix: (k_bins, n, chunk) float64, as produced above @@ -129,20 +172,33 @@ __global__ void puc_accumulation_kernel( int z_global = z_start + z_local; if (z_global >= n || x == z_global) return; - double mi_xz = mi_matrix[x * n + z_global]; + double mi_xz = mi_matrix[(long long)x * n + z_global]; if (mi_xz <= 1e-12) return; + // si_matrix is (k_bins, n, chunk): one bin plane is plane_stride apart, one + // source gene is z_chunk_size apart. Both walks are 64-bit, and the advances + // live in the for-increment clauses so the `continue`s cannot skip them. + const long long plane_stride = (long long)n * z_chunk_size; + const long long si_x_cell = (long long)x * z_chunk_size + z_local; + double local_puc = 0.0; - for (int y = 0; y < n; ++y) { + long long si_y_cell = z_local; // y = 0: (0 * n + y) * chunk + z_local + + for (int y = 0; y < n; ++y, si_y_cell += z_chunk_size) { if (y == x || y == z_global) continue; double redundancy = 0.0; - for (int k = 0; k < k_bins; ++k) { - double p_z_k = marginals[k * n + z_global]; + long long si_x_idx = si_x_cell; + long long si_y_idx = si_y_cell; + long long marg_idx = z_global; // marginals[k * n + z_global] + + for (int k = 0; k < k_bins; ++k, + si_x_idx += plane_stride, si_y_idx += plane_stride, marg_idx += n) { + double p_z_k = marginals[marg_idx]; if (p_z_k <= 0.0) continue; - double si_x = si_matrix[(k * n + x) * z_chunk_size + z_local]; - double si_y = si_matrix[(k * n + y) * z_chunk_size + z_local]; + double si_x = si_matrix[si_x_idx]; + double si_y = si_matrix[si_y_idx]; redundancy += p_z_k * fmin(si_x, si_y); } @@ -152,7 +208,7 @@ __global__ void puc_accumulation_kernel( } } - puc_scores[x * n + z_global] = local_puc; + puc_scores[(long long)x * n + z_global] = local_puc; } } // extern "C" diff --git a/python/tests/test_cuda.py b/python/tests/test_cuda.py index 59a94d8..85bfc2f 100644 --- a/python/tests/test_cuda.py +++ b/python/tests/test_cuda.py @@ -15,11 +15,13 @@ from fastpidc.cuda import ( _MAX_CHUNK_SIZE, + _MAX_K_BINS, _bb_kernel_name, _bb_memory_batches, _bb_problem_bytes, _bb_quantile_buckets, _bb_threads_for_max_u, + _check_kernel_index_limits, _chunk_size_for_free_memory, _smallest_unsigned_dtype, cuda_available, @@ -297,3 +299,104 @@ def test_bayesian_blocks_falls_back_to_cpu_without_a_gpu(monkeypatch, julia_test cpu_nodes = get_nodes(path, bb_backend="cpu") for from_fallback, from_cpu in zip(fallback_nodes, cpu_nodes): np.testing.assert_array_equal(from_fallback.binned_values, from_cpu.binned_values) + + +# --- Flat index range: the >2^31 regression --------------------------------- +# +# A 12,071-gene x 38,176-cell production run crashed with +# CUDA_ERROR_ILLEGAL_ADDRESS because the shared kernels computed the `counts` +# index in 32-bit: k_bins^2 * n * chunk_size = 33^2 * 12071 * 256 is +# 3,365,201,664 elements, whose top index wraps past INT32_MAX. All flat offsets +# are 64-bit now; these tests pin that. + +INT32_MAX = 2**31 - 1 + +# n=64 genes at 725 bins with a 64-gene chunk is the cheapest configuration that +# clears 2^31 counts elements: 8.02 GiB, versus ~22 MiB of everything else. +_OVERFLOW_N_NODES = 64 +_OVERFLOW_N_SAMPLES = 2000 +_OVERFLOW_N_BINS = 725 +_OVERFLOW_CHUNK = 64 +# A chunk small enough that the same problem stays provably inside int32, +# giving a control the overflow cannot have touched. +_CONTROL_CHUNK = 32 +# The overflowing chunk needs ~8 GiB plus room for cupy's pool and the other +# buffers; require real headroom so the test never competes for a full device. +_OVERFLOW_FREE_MEMORY_FLOOR = 12 * GIB + + +def _counts_elements(k_bins: int, n: int, chunk: int) -> int: + """Element count of the per-chunk joint-count buffer the kernels index.""" + return k_bins**2 * n * chunk + + +def test_production_configuration_exceeds_int32_indexing(): + # Documents why the kernels must index in 64-bit, and pins the arithmetic + # that the large-memory test below relies on. No GPU needed. + assert _counts_elements(k_bins=33, n=12071, chunk=256) - 1 > INT32_MAX + assert _counts_elements(_OVERFLOW_N_BINS, _OVERFLOW_N_NODES, _OVERFLOW_CHUNK) - 1 > INT32_MAX + assert _counts_elements(_OVERFLOW_N_BINS, _OVERFLOW_N_NODES, _CONTROL_CHUNK) - 1 <= INT32_MAX + + +def test_check_kernel_index_limits_accepts_the_largest_supported_k_bins(): + _check_kernel_index_limits(_MAX_K_BINS) + assert _MAX_K_BINS**2 <= INT32_MAX + assert (_MAX_K_BINS + 1) ** 2 > INT32_MAX + + +def test_check_kernel_index_limits_rejects_a_larger_k_bins(): + # joint_counts_kernel forms u * k_bins + v in int32; everything else is 64-bit. + with pytest.raises(RuntimeError, match="bin pairs in 32-bit"): + _check_kernel_index_limits(_MAX_K_BINS + 1) + + +def test_chunk_sizing_is_not_capped_by_int32_element_count(): + # The memory guard bounds bytes, not the element index. Regression guard + # against anyone "fixing" the overflow by reimposing an int32 element cap, + # which would silently shrink chunks on large GPUs. + chunk = _chunk_size_for_free_memory(n=12071, k_bins=33, free_bytes=int(19.22 * GIB)) + assert chunk == 256 + assert _counts_elements(33, 12071, chunk) - 1 > INT32_MAX + + +@pytest.mark.largemem +def test_puc_indexing_past_int32_matches_a_smaller_chunk(): + """Drive the PUC kernels past 2^31 counts elements and require the result to + match a chunking that stays inside int32. + + Same kernels, same nodes, two buffer layouts: only the flat offsets differ, + so any 32-bit truncation shows up as a mismatch (before the fix it faulted + outright with CUDA_ERROR_ILLEGAL_ADDRESS). + """ + if not cuda_available(): + pytest.skip("no functional GPU / cupy backend available") + + import cupy as cp + + from fastpidc.cuda import compute_puc_full_cuda + + free_bytes = int(cp.cuda.runtime.memGetInfo()[0]) + if free_bytes < _OVERFLOW_FREE_MEMORY_FLOOR: + pytest.skip( + f"needs {_OVERFLOW_FREE_MEMORY_FLOOR / GIB:.0f} GiB free device memory, have {free_bytes / GIB:.1f} GiB" + ) + + rng = np.random.default_rng(0) + values = rng.random((_OVERFLOW_N_SAMPLES, _OVERFLOW_N_NODES)) + nodes = [ + Node.from_raw_values(f"N{i}", values[:, i], "uniform_width", "maximum_likelihood", _OVERFLOW_N_BINS) + for i in range(_OVERFLOW_N_NODES) + ] + k_bins = max(node.number_of_bins for node in nodes) + + # Fail loudly rather than pass vacuously if the configuration drifts below + # the threshold this test exists to cross. + assert _counts_elements(k_bins, len(nodes), _OVERFLOW_CHUNK) - 1 > INT32_MAX + + overflow_mi, overflow_puc = compute_puc_full_cuda(nodes, chunk_size=_OVERFLOW_CHUNK) + control_mi, control_puc = compute_puc_full_cuda(nodes, chunk_size=_CONTROL_CHUNK) + + np.testing.assert_array_equal(overflow_mi, control_mi) + np.testing.assert_array_equal(overflow_puc, control_puc) + assert np.all(np.isfinite(overflow_mi)) + assert np.all(overflow_puc >= 0) diff --git a/test/cli_argument_tests.jl b/test/cli_argument_tests.jl new file mode 100644 index 0000000..5ae4aac --- /dev/null +++ b/test/cli_argument_tests.jl @@ -0,0 +1,78 @@ +using Test + +# CLI_fastpidc.jl runs `main()` against the process's own `ARGS` and calls +# `exit(1)` on error when invoked as a script; it is guarded by +# `abspath(PROGRAM_FILE) == @__FILE__` so that `include()`ing it here (to reach +# `parse_args`/`validate_args` in isolation) is safe. +include(joinpath(dirname(@__FILE__), "..", "CLI_fastpidc.jl")) + +# Regression test: parse_args() stores any "--key value" pair with no +# validation, so a flag typed with the wrong punctuation (e.g. "--n-bins" +# instead of the documented "--n_bins") used to be silently ignored - main() +# would fall back to n_bins's default with no warning, so a request like +# `--discretizer uniform_width --n-bins 6` silently produced 10 bins instead +# of 6. validate_args() must catch every such mismatch loudly. +@testset "CLI argument validation" begin + @testset "accepts every flag main() reads" begin + args = Dict( + "infile" => "x", + "outfile" => "y", + "delim" => "space", + "discretizer" => "bayesian_blocks", + "estimator" => "maximum_likelihood", + "n_bins" => "10", + "base" => "2", + "backend" => "cpu", + "bb-backend" => "cpu", + "output-format" => "tsv", + "dump-mi-path" => "mi.npy", + "dump-puc-path" => "puc.npy", + "verbose" => "true", + "help" => "false", + ) + @test validate_args(args) === nothing + end + + @testset "rejects the exact hyphen/underscore mismatch that was silently dropped" begin + # This is the concrete failure this test exists to prevent: --n-bins + # used to be accepted and stored, then never read by main(), leaving + # n_bins at its default with no indication anything was wrong. + @test_throws ErrorException validate_args(Dict("n-bins" => "6")) + @test_throws ErrorException validate_args(Dict("bb_backend" => "cpu")) + end + + @testset "error message names the likely intended flag" begin + err = try + validate_args(Dict("n-bins" => "6")) + nothing + catch e + e + end + @test err isa ErrorException + @test occursin("--n-bins", err.msg) + @test occursin("--n_bins", err.msg) + end + + @testset "rejects an unrelated typo without a hint" begin + err = try + validate_args(Dict("verbosee" => "true")) + nothing + catch e + e + end + @test err isa ErrorException + @test occursin("--verbosee", err.msg) + @test !occursin("did you mean", err.msg) + end + + @testset "parse_args followed by validate_args catches a real command line" begin + empty!(ARGS) + append!(ARGS, ["--infile", "x", "--outfile", "y", "--n-bins", "6"]) + try + parsed = parse_args() + @test_throws ErrorException validate_args(parsed) + finally + empty!(ARGS) + end + end +end diff --git a/test/cuda_largemem_tests.jl b/test/cuda_largemem_tests.jl new file mode 100644 index 0000000..63bb8d2 --- /dev/null +++ b/test/cuda_largemem_tests.jl @@ -0,0 +1,105 @@ +using FastPIDC +using Test +using CUDA + +# Regression test for the 32-bit flat-index overflow that crashed a 12,071-gene +# x 38,176-cell run with CUDA_ERROR_ILLEGAL_ADDRESS. The shared kernels compute +# the joint-count index as k_bins^2 * n * chunk_size; at 33 bins, 12,071 genes +# and a 256-gene chunk that is 3,365,201,664 elements, whose top index wraps +# past typemax(Int32). All flat offsets are 64-bit now. +# +# The overflow cannot be reproduced at small sizes - the buffer must genuinely +# hold more than 2^31 Int32 elements, i.e. at least ~8 GiB of device memory - so +# this file is opt-in and skipped unless explicitly requested: +# +# FASTPIDC_LARGEMEM_TESTS=1 julia --project=. -e 'using Pkg; Pkg.test()' +# +# Do not enable it on a shared GPU: it allocates ~8 GiB. + +const _LARGEMEM_ENABLED = get(ENV, "FASTPIDC_LARGEMEM_TESTS", "0") == "1" + +# 64 genes at 725 bins with a 64-gene chunk is the cheapest configuration that +# clears 2^31 counts elements: 8.02 GiB, versus ~23 MiB of everything else. +const _LARGEMEM_NODES = 64 +const _LARGEMEM_SAMPLES = 2000 +const _LARGEMEM_BINS = 725 +# The host derives chunk_size from free memory; it only reaches 64 (rather than +# 63, which stays inside Int32) with roughly 11 GiB free, so require headroom. +const _LARGEMEM_FREE_FLOOR = 12 * 2^30 + +_counts_elements(k_bins, n, chunk) = big(k_bins)^2 * n * chunk + +if _LARGEMEM_ENABLED && CUDA.functional() + @testset "PUC flat indexing past 2^31" begin + free_bytes = Int(CUDA.free_memory()) + + if free_bytes < _LARGEMEM_FREE_FLOOR + @warn "Skipping large-memory PUC index test" required_gib = + _LARGEMEM_FREE_FLOOR / 2^30 free_gib = free_bytes / 2^30 + else + values = [ + (sin(sample * 0.37 + gene) + 1) / 2 for sample = 1:_LARGEMEM_SAMPLES, + gene = 1:_LARGEMEM_NODES + ] + nodes = [ + Node( + "N$gene", + values[:, gene], + "uniform_width", + "maximum_likelihood", + _LARGEMEM_BINS, + ) for gene = 1:_LARGEMEM_NODES + ] + k_bins = maximum(node -> node.number_of_bins, nodes) + + # The host picks the chunk itself, so fail loudly rather than pass + # vacuously if this configuration drifts below the threshold the + # test exists to cross. + cuda_ext = Base.get_extension(FastPIDC, :FastPIDCCUDAExt) + @test cuda_ext !== nothing + bytes_per_chunk_col = + k_bins^2 * length(nodes) * sizeof(Int32) + + k_bins * length(nodes) * sizeof(Float64) + chunk_size = clamp( + floor(Int, free_bytes * 0.8 / bytes_per_chunk_col), + 1, + min(256, length(nodes)), + ) + @test _counts_elements(k_bins, length(nodes), chunk_size) - 1 > + big(typemax(Int32)) + + gpu_mi, gpu_puc = FastPIDC.compute_puc_full( + nodes; + estimator = "maximum_likelihood", + base = 2, + config = PIDCConfig(backend = :cuda), + ) + cpu_mi, cpu_puc = FastPIDC.compute_puc_full( + nodes; + estimator = "maximum_likelihood", + base = 2, + config = PIDCConfig(backend = :cpu), + ) + + @test all(isfinite, gpu_mi) + @test all(>=(0), gpu_puc) + @test isapprox(gpu_mi, cpu_mi; atol = 1e-9, rtol = 1e-12) + @test isapprox(gpu_puc, cpu_puc; atol = 1e-6, rtol = 1e-9) + end + end +end + +# Cheap, always-run companion: the one kernel limit 64-bit indexing does not +# lift. joint_counts_kernel forms the bin-pair index u * k_bins + v in Int32. +if CUDA.functional() + @testset "Kernel bin-pair index limit" begin + cuda_ext = Base.get_extension(FastPIDC, :FastPIDCCUDAExt) + @test cuda_ext !== nothing + @test cuda_ext._MAX_K_BINS^2 <= typemax(Int32) + @test big(cuda_ext._MAX_K_BINS + 1)^2 > big(typemax(Int32)) + @test cuda_ext._check_kernel_index_limits(cuda_ext._MAX_K_BINS) === nothing + @test_throws ErrorException cuda_ext._check_kernel_index_limits( + cuda_ext._MAX_K_BINS + 1, + ) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 3efd609..01d3859 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -13,6 +13,7 @@ using Test using DelimitedFiles include("cuda_smoke_tests.jl") +include("cli_argument_tests.jl") include("baseline_helpers.jl") include("baseline_smoke_tests.jl") include("bayesian_blocks_tests.jl") @@ -20,6 +21,7 @@ include("cuda_bayesian_blocks_tests.jl") include("diagnostic_dump_tests.jl") include("cuda_numeric_tests.jl") include("cuda_numeric_tests_bb.jl") +include("cuda_largemem_tests.jl") include("benchmark_puc.jl") include("benchmark_puc_bb.jl") From bc48caaff588606764e41481d5d4564fabb8dbe7 Mon Sep 17 00:00:00 2001 From: jjschirle Date: Fri, 11 Sep 2026 11:29:38 -0700 Subject: [PATCH 2/6] cli argument test bugfix --- CLI_fastpidc.jl | 4 +-- test/cli_argument_tests.jl | 61 +++++++++++++++++++++++++++++--------- 2 files changed, 49 insertions(+), 16 deletions(-) diff --git a/CLI_fastpidc.jl b/CLI_fastpidc.jl index 9cf350d..1ec768c 100644 --- a/CLI_fastpidc.jl +++ b/CLI_fastpidc.jl @@ -61,7 +61,7 @@ const VALID_ARG_KEYS = Set([ wrongly-punctuated flag was previously stored under a key `main()` never reads, silently keeping its default rather than erroring - a caller could ask for `--discretizer -uniform_width --n-bins 6` and get 10 bins with no indication anything was +uniform_width --n_bins 6` and get 10 bins with no indication anything was wrong. This checks every parsed key against [`VALID_ARG_KEYS`](@ref) and errors out, naming the likely intended flag when the mismatch is only a `-`/`_` swap. @@ -140,7 +140,7 @@ Diagnostics Dumps: --dump-puc-path PATH If set, dump pre-context PUC scores here (TSV). Other: - --verbose Print detailed progress information + --verbose BOOL Print detailed progress information. Default: false --help, -h Show this help and exit. Example: diff --git a/test/cli_argument_tests.jl b/test/cli_argument_tests.jl index 5ae4aac..10546f2 100644 --- a/test/cli_argument_tests.jl +++ b/test/cli_argument_tests.jl @@ -7,10 +7,10 @@ using Test include(joinpath(dirname(@__FILE__), "..", "CLI_fastpidc.jl")) # Regression test: parse_args() stores any "--key value" pair with no -# validation, so a flag typed with the wrong punctuation (e.g. "--n-bins" -# instead of the documented "--n_bins") used to be silently ignored - main() +# validation, so a flag typed with the wrong punctuation (e.g. "--n_bins" +# instead of the documented "--n-bins") used to be silently ignored - main() # would fall back to n_bins's default with no warning, so a request like -# `--discretizer uniform_width --n-bins 6` silently produced 10 bins instead +# `--discretizer uniform_width --n_bins 6` silently produced 10 bins instead # of 6. validate_args() must catch every such mismatch loudly. @testset "CLI argument validation" begin @testset "accepts every flag main() reads" begin @@ -20,7 +20,7 @@ include(joinpath(dirname(@__FILE__), "..", "CLI_fastpidc.jl")) "delim" => "space", "discretizer" => "bayesian_blocks", "estimator" => "maximum_likelihood", - "n_bins" => "10", + "n-bins" => "10", "base" => "2", "backend" => "cpu", "bb-backend" => "cpu", @@ -33,24 +33,30 @@ include(joinpath(dirname(@__FILE__), "..", "CLI_fastpidc.jl")) @test validate_args(args) === nothing end - @testset "rejects the exact hyphen/underscore mismatch that was silently dropped" begin - # This is the concrete failure this test exists to prevent: --n-bins - # used to be accepted and stored, then never read by main(), leaving - # n_bins at its default with no indication anything was wrong. - @test_throws ErrorException validate_args(Dict("n-bins" => "6")) - @test_throws ErrorException validate_args(Dict("bb_backend" => "cpu")) + @testset "rejects underscore spellings of hyphenated flags" begin + # Public CLI flags use kebab-case. Snake-case spellings must fail rather + # than being stored under an unused key and silently falling back to a default. + for (wrong_key, value) in ( + ("n_bins", "6"), + ("bb_backend", "cpu"), + ("output_format", "npy"), + ("dump_mi_path", "mi.npy"), + ("dump_puc_path", "puc.npy"), + ) + @test_throws ErrorException validate_args(Dict(wrong_key => value)) + end end @testset "error message names the likely intended flag" begin err = try - validate_args(Dict("n-bins" => "6")) + validate_args(Dict("n_bins" => "6")) nothing catch e e end @test err isa ErrorException - @test occursin("--n-bins", err.msg) @test occursin("--n_bins", err.msg) + @test occursin("--n-bins", err.msg) end @testset "rejects an unrelated typo without a hint" begin @@ -65,9 +71,36 @@ include(joinpath(dirname(@__FILE__), "..", "CLI_fastpidc.jl")) @test !occursin("did you mean", err.msg) end - @testset "parse_args followed by validate_args catches a real command line" begin + @testset "parse_args accepts a representative documented command line" begin + empty!(ARGS) + append!(ARGS, [ + "--infile", "x", + "--outfile", "y", + "--output-format", "npy", + "--dump-mi-path", "mi.npy", + "--dump-puc-path", "puc.npy", + "--discretizer", "bayesian_blocks", + "--estimator", "maximum_likelihood", + "--n-bins", "6", + "--base", "2", + "--backend", "cuda", + "--bb-backend", "cuda", + "--verbose", "true", + ]) + try + parsed = parse_args() + @test parsed["n-bins"] == "6" + @test parsed["bb-backend"] == "cuda" + @test parsed["output-format"] == "npy" + @test validate_args(parsed) === nothing + finally + empty!(ARGS) + end + end + + @testset "parse_args followed by validate_args catches underscore spelling" begin empty!(ARGS) - append!(ARGS, ["--infile", "x", "--outfile", "y", "--n-bins", "6"]) + append!(ARGS, ["--infile", "x", "--outfile", "y", "--n_bins", "6"]) try parsed = parse_args() @test_throws ErrorException validate_args(parsed) From 783d7d9641fdf99136563120f0ca09e770c0f1aa Mon Sep 17 00:00:00 2001 From: jjschirle Date: Fri, 11 Sep 2026 11:30:17 -0700 Subject: [PATCH 3/6] python cli and cli test --- python/src/fastpidc/cli.py | 4 ++-- python/tests/test_cli.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/python/src/fastpidc/cli.py b/python/src/fastpidc/cli.py index 752bfe7..e859344 100644 --- a/python/src/fastpidc/cli.py +++ b/python/src/fastpidc/cli.py @@ -60,7 +60,7 @@ def _build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--estimator", default="maximum_likelihood", help="Default: maximum_likelihood") parser.add_argument( - "--n_bins", type=int, default=10, help="Number of bins (ignored by bayesian_blocks). Default: 10" + "--n-bins", type=int, default=10, help="Number of bins (ignored by bayesian_blocks). Default: 10" ) parser.add_argument("--base", type=int, default=2, help="Log base for MI (2, e-like int, 10). Default: 2") parser.add_argument("--backend", default="cuda", choices=("cuda", "cpu"), help="PUC/PIDC backend. Default: cuda") @@ -149,7 +149,7 @@ def main(argv: list[str] | None = None) -> int: _say(f"ERROR: {e}") print("\nTraceback:") traceback.print_exc() - _say("Tip: If the error mentions discretization, try --discretizer uniform_width and --n_bins 10-20.") + _say("Tip: If the error mentions discretization, try --discretizer uniform_width and --n-bins 10-20.") return 1 _say(f"Wrote edges to {args.outfile}") diff --git a/python/tests/test_cli.py b/python/tests/test_cli.py index 4d141ee..94e3cad 100644 --- a/python/tests/test_cli.py +++ b/python/tests/test_cli.py @@ -43,7 +43,7 @@ def test_cli_runs_end_to_end(tmp_path, text_data_file, capsys): "cpu", "--discretizer", "uniform_width", - "--n_bins", + "--n-bins", "2", ] ) @@ -109,7 +109,7 @@ def test_cli_npy_output(tmp_path, text_data_file): "cpu", "--discretizer", "uniform_width", - "--n_bins", + "--n-bins", "2", "--output-format", "npy", From 7a75df2b9f2466f414b2aad5b9d98c91b9293106 Mon Sep 17 00:00:00 2001 From: jjschirle Date: Fri, 11 Sep 2026 13:01:33 -0700 Subject: [PATCH 4/6] Boundary memory handling. Memory boundary test suite expanded --- ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl | 452 +++++++++++++------- python/src/fastpidc/cuda.py | 446 +++++++++++-------- python/src/fastpidc/kernels/pidc_kernels.cu | 26 +- python/tests/test_cuda.py | 117 +++-- test/cuda_bayesian_blocks_tests.jl | 12 +- test/cuda_largemem_tests.jl | 57 ++- 6 files changed, 723 insertions(+), 387 deletions(-) diff --git a/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl b/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl index 32e1add..d7864e2 100644 --- a/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl +++ b/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl @@ -92,26 +92,124 @@ end # --- Host implementation --- +const _KERNEL_INT_MAX = typemax(Int32) +const _GPU_MEMORY_BUDGET_NUMERATOR = 65 +const _GPU_MEMORY_BUDGET_DENOMINATOR = 100 +const _MAX_CHUNK_SIZE = 256 + +function _gpu_memory_budget_bytes(free_bytes::Integer) + free_bytes > 0 || throw(ArgumentError("free_bytes must be positive")) + return Int( + div( + big(free_bytes) * _GPU_MEMORY_BUDGET_NUMERATOR, + _GPU_MEMORY_BUDGET_DENOMINATOR, + ), + ) +end + +function _reusable_gpu_memory_bytes( + driver_free_bytes::Integer, + pool_cached_bytes::Integer = 0, + pool_used_bytes::Integer = 0, +) + driver_free_bytes >= 0 || throw(ArgumentError("driver_free_bytes must be nonnegative")) + pool_cached_bytes >= 0 || throw(ArgumentError("pool_cached_bytes must be nonnegative")) + pool_used_bytes >= 0 || throw(ArgumentError("pool_used_bytes must be nonnegative")) + pool_used_bytes <= pool_cached_bytes || throw( + ArgumentError("pool_used_bytes cannot exceed pool_cached_bytes"), + ) + + # CUDA.free_memory() reports memory still free at the driver level. CUDA.jl's + # stream-ordered pool may additionally hold reserved-but-unused bytes that are + # immediately reusable by this process, so include those in the planning view. + return Int(big(driver_free_bytes) + big(pool_cached_bytes - pool_used_bytes)) +end + +function _available_gpu_memory_bytes() + driver_free_bytes = Int(CUDA.free_memory()) + + # CUDA.jl 5.4+ exposes pool accounting helpers. On allocators/devices that do + # not use the stream-ordered pool they return `missing`, in which case driver + # free memory is already the only reusable pool we can account for. + cached = CUDA.cached_memory() + used = CUDA.used_memory() + if cached isa Integer && used isa Integer + return _reusable_gpu_memory_bytes(driver_free_bytes, Int(cached), Int(used)) + end + return driver_free_bytes +end + """ - _MAX_K_BINS + _check_kernel_scalar_limits(num_nodes, num_samples, k_bins) -Largest bins-per-gene the shared kernels can index. They compute every flat -buffer offset in 64-bit, with one exception: `joint_counts_kernel` forms the -bin-pair index `u * k_bins + v` in Int32, which stays exact only while -`k_bins^2` fits in Int32 (isqrt(typemax(Int32)) == 46340). +The shared CUDA kernels receive launch dimensions as signed 32-bit integers. +All composed flat-buffer offsets are 64-bit, but these scalar dimensions must +still fit exactly in `Int32` before crossing the host/device boundary. """ -const _MAX_K_BINS = 46340 - -function _check_kernel_index_limits(k_bins::Integer) - k_bins <= _MAX_K_BINS || error( - "compute_puc_full_cuda: the discretizer selected k_bins=$k_bins bins " * - "per gene, but the CUDA kernels index bin pairs in 32-bit and support " * - "at most $(_MAX_K_BINS). Use discretizer=\"uniform_width\" with a " * - "fixed, small number_of_bins, or config.backend = :cpu.", +function _check_kernel_scalar_limits( + num_nodes::Integer, + num_samples::Integer, + k_bins::Integer, +) + for (name, value) in ( + ("num_nodes", num_nodes), + ("num_samples", num_samples), + ("k_bins", k_bins), ) + value >= 1 || throw(ArgumentError("$name must be positive; got $value")) + value <= _KERNEL_INT_MAX || throw( + ArgumentError( + "compute_puc_full_cuda: $name=$value exceeds the CUDA kernel's " * + "signed 32-bit scalar limit of $(_KERNEL_INT_MAX).", + ), + ) + end return nothing end +function _puc_memory_plan( + num_nodes::Integer, + num_samples::Integer, + k_bins::Integer, + free_bytes::Integer, +) + # Use BigInt for the planning arithmetic so an impossible input cannot + # overflow on the host while we are trying to decide whether it fits. + n = big(num_nodes) + m = big(num_samples) + k = big(k_bins) + budget_bytes = big(_gpu_memory_budget_bytes(free_bytes)) + + fixed_bytes = + n * m * sizeof(Int32) + # discretized data + n * k * sizeof(Float64) + # marginals + 2 * n * n * sizeof(Float64) # MI + PUC output matrices + bytes_per_chunk_col = + k * k * n * sizeof(Int32) + # joint counts + k * n * sizeof(Float64) # specific information + + minimum_bytes = fixed_bytes + bytes_per_chunk_col + minimum_bytes <= budget_bytes || error( + "compute_puc_full_cuda: the fixed GPU buffers plus a one-gene chunk " * + "would require $(round(Float64(minimum_bytes) / 2^30, digits = 2)) GiB, " * + "which exceeds the configured 65% memory budget " * + "($(round(Float64(budget_bytes) / 2^30, digits = 2)) GiB of " * + "$(round(free_bytes / 2^30, digits = 2)) GiB currently reusable). " * + "Reduce the number of genes/samples/bins, use a fixed small-bin " * + "discretizer, or use config.backend = :cpu.", + ) + + max_chunk_size = div(budget_bytes - fixed_bytes, bytes_per_chunk_col) + chunk_size = min(max_chunk_size, big(_MAX_CHUNK_SIZE), n) + + return ( + chunk_size = Int(chunk_size), + fixed_bytes = Int(fixed_bytes), + bytes_per_chunk_col = Int(bytes_per_chunk_col), + budget_bytes = Int(budget_bytes), + ) +end + function _smallest_unsigned_type(max_value::Integer) max_value >= 0 || throw(ArgumentError("max_value must be nonnegative")) @@ -130,24 +228,22 @@ end FastPIDC.compute_puc_full_cuda(nodes, config, base) -> (mi_scores, puc_scores) 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 (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 the -shared `joint_counts_kernel`, `mi_si_kernel` and `puc_accumulation_kernel` -(see the `FastPIDCCUDAExt` module docstring) 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, matching the -kernel source). Raises an `ErrorException` with a suggested remedy if even -a single-gene chunk would not fit in the currently-free GPU memory, or if -the discretizer selected more than `_MAX_K_BINS` bins per gene. - -The chunked intermediates legitimately exceed 2^31 elements on large gene -sets - `counts` alone is `k_bins^2 * num_nodes * chunk_size` - so the shared -kernels index them in 64-bit (see the indexing contract in the kernel source). +pairwise MI matrix and pre-context PUC matrix for `nodes` on the GPU. +Before allocating device buffers it plans the complete footprint - fixed data, +marginals and output matrices plus chunked intermediates - against 65% of the +memory currently reusable by the process (driver-free plus unused CUDA.jl pool +blocks). The target (`z`) chunk is capped at 256 genes and shrunk +as needed, leaving 35% headroom for allocator fragmentation, CUDA/runtime +workspaces and concurrent users. If the fixed buffers plus a one-gene chunk do +not fit that budget, the function raises a descriptive error before allocating +the large device arrays. + +All composed flat-buffer offsets in the shared CUDA kernels are 64-bit, +including the bin-pair term. The scalar launch dimensions remain signed +32-bit for ABI compatibility and are validated on the host before conversion. +`config.verbose` prints the memory plan; `base` is currently unused (mutual +information is always computed in base 2 on the GPU, matching the kernel +source). Device buffers use Julia's column-major layout with dimensions reversed relative to the kernel source's documented (row-major) shapes - e.g. a @@ -157,6 +253,10 @@ with manual pointer arithmetic is identical in both languages, with no transposition needed at the call boundary. """ function FastPIDC.compute_puc_full_cuda(nodes, config, base) + isempty(nodes) && throw(ArgumentError("compute_puc_full_cuda requires at least one node")) + + # Compile/load the module before measuring free memory so its device-side + # footprint is already reflected in the runtime memory budget. md = _get_module() joint_counts_kernel = CuFunction(md, "joint_counts_kernel") mi_si_kernel = CuFunction(md, "mi_si_kernel") @@ -164,138 +264,160 @@ function FastPIDC.compute_puc_full_cuda(nodes, config, base) num_nodes = length(nodes) num_samples = length(nodes[1].binned_values) + all(n -> length(n.binned_values) == num_samples, nodes) || throw( + ArgumentError("all nodes must contain the same number of discretized samples"), + ) k_bins = maximum(n -> n.number_of_bins, nodes) - _check_kernel_index_limits(k_bins) - - # Prepare static data on CPU and move to GPU. The shared CUDA C kernels use - # 0-indexed Int32 bin ids; FastPIDC.jl's bin ids are 1-indexed, so shift - # them down at this boundary. + _check_kernel_scalar_limits(num_nodes, num_samples, k_bins) + + # Plan the *entire* device footprint before allocating anything. The budget + # is 65% of currently reusable VRAM (driver-free plus unused CUDA.jl pool + # blocks), leaving 35% for allocator fragmentation, + # runtime/library workspaces, and other users/processes on a shared GPU. + free_bytes = _available_gpu_memory_bytes() + memory_plan = _puc_memory_plan(num_nodes, num_samples, k_bins, free_bytes) + chunk_size = memory_plan.chunk_size + + # Prepare static data on CPU. The shared CUDA C kernels use 0-indexed Int32 + # bin ids; FastPIDC.jl's bin ids are 1-indexed, so shift them down at this + # boundary. data_cpu = zeros(Int32, num_nodes, num_samples) # kernel shape (m, n), reversed marginals_cpu = zeros(Float64, num_nodes, k_bins) # kernel shape (k_bins, n), reversed for i = 1:num_nodes - data_cpu[i, :] .= Int32.(nodes[i].binned_values) .- Int32(1) - p = nodes[i].probabilities - marginals_cpu[i, 1:length(p)] .= Float64.(p) - end + node = nodes[i] + node.number_of_bins >= 1 || throw( + ArgumentError("node $(node.label) has non-positive number_of_bins"), + ) + maximum(node.binned_values) <= node.number_of_bins || throw( + ArgumentError("node $(node.label) contains a bin id above number_of_bins"), + ) + minimum(node.binned_values) >= 1 || throw( + ArgumentError("node $(node.label) contains a bin id below 1"), + ) - data_gpu = CuArray(data_cpu) - marginals_gpu = CuArray(marginals_cpu) - - # Global output matrices (kernel shape (n, n); square, so no reversal needed). - puc_scores_gpu = CUDA.zeros(Float64, num_nodes, num_nodes) - mi_matrix_gpu = CUDA.zeros(Float64, num_nodes, num_nodes) - - # Chunked intermediates scale with k_bins^2 * num_nodes * chunk_size for - # joint counts and k_bins * num_nodes * chunk_size for specific information. - # Size the target-gene chunk from currently-free memory rather than always - # allocating a fixed 256-gene chunk. - # - # This bounds MEMORY AVAILABILITY only - it is not, and must not be turned - # back into, a bound on the flat element index. The kernels index these - # buffers in 64-bit precisely so the chunk can be sized from free memory - # without an int32 element-count cap. - 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 = Int(CUDA.free_memory()) - safety_factor = 0.8 # headroom for fixed buffers + 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.", + data_cpu[i, :] .= Int32.(node.binned_values) .- Int32(1) + p = node.probabilities + length(p) <= k_bins || throw( + ArgumentError("node $(node.label) has more probabilities than k_bins"), ) + marginals_cpu[i, 1:length(p)] .= Float64.(p) end - # Chunked intermediate buffers (pre-allocated once), with dimensions - # reversed to preserve the row-major flat layout expected by the CUDA C kernels. - counts_chunk_gpu = CUDA.zeros(Int32, chunk_size, num_nodes, k_bins, k_bins) - si_chunk_gpu = CUDA.zeros(Float64, chunk_size, num_nodes, k_bins) + data_gpu = nothing + marginals_gpu = nothing + puc_scores_gpu = nothing + mi_matrix_gpu = nothing + counts_chunk_gpu = nothing + si_chunk_gpu = nothing - if config.verbose - 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 + try + data_gpu = CuArray(data_cpu) + marginals_gpu = CuArray(marginals_cpu) + + # Global output matrices (kernel shape (n, n); square, so no reversal needed). + puc_scores_gpu = CUDA.zeros(Float64, num_nodes, num_nodes) + mi_matrix_gpu = CUDA.zeros(Float64, num_nodes, num_nodes) + + # Chunked intermediate buffers, pre-allocated once. Dimensions are + # reversed to preserve the row-major flat layout expected by CUDA C. + counts_chunk_gpu = CUDA.zeros(Int32, chunk_size, num_nodes, k_bins, k_bins) + si_chunk_gpu = CUDA.zeros(Float64, chunk_size, num_nodes, k_bins) + + if config.verbose + println( + "[FastPIDC] GPU Chunked PUC: Processing $num_nodes x $num_nodes pairs " * + "(k_bins=$k_bins)...", + ) + println( + "[FastPIDC] GPU memory: $(round(free_bytes / 2^30, digits = 2)) GiB reusable; " * + "65% budget=$(round(memory_plan.budget_bytes / 2^30, digits = 2)) GiB; " * + "fixed=$(round(memory_plan.fixed_bytes / 2^30, digits = 2)) GiB", + ) + println( + "[FastPIDC] Using chunk size of $chunk_size " * + "(approx. $(ceil(Int, num_nodes / chunk_size)) iterations)", + ) + end - threads = (16, 16) + threads = (16, 16) - # Iterate over the Z-axis in chunks. The shared kernels use 0-based target - # indices, so convert z_start at the call boundary. - for z_start_1 in 1:chunk_size:num_nodes - z_start = z_start_1 - 1 - z_end = min(z_start_1 + chunk_size - 1, num_nodes) - z_curr_chunk_size = z_end - z_start_1 + 1 + # Iterate over the Z-axis in chunks. The shared kernels use 0-based target + # indices, so convert z_start at the call boundary. + for z_start_1 in 1:chunk_size:num_nodes + z_start = z_start_1 - 1 + z_end = min(z_start_1 + chunk_size - 1, num_nodes) + z_curr_chunk_size = z_end - z_start_1 + 1 - CUDA.fill!(counts_chunk_gpu, Int32(0)) - CUDA.fill!(si_chunk_gpu, Float64(0)) + CUDA.fill!(counts_chunk_gpu, Int32(0)) + CUDA.fill!(si_chunk_gpu, Float64(0)) - blocks = (cld(num_nodes, 16), cld(z_curr_chunk_size, 16)) + blocks = (cld(num_nodes, 16), cld(z_curr_chunk_size, 16)) - cudacall( - joint_counts_kernel, - (CuPtr{Cint}, CuPtr{Cint}, Cint, Cint, Cint, Cint, Cint), - data_gpu, counts_chunk_gpu, - Cint(num_nodes), Cint(num_samples), Cint(k_bins), - Cint(z_start), Cint(z_curr_chunk_size); - blocks=blocks, threads=threads, - ) + cudacall( + joint_counts_kernel, + (CuPtr{Cint}, CuPtr{Cint}, Cint, Cint, Cint, Cint, Cint), + data_gpu, counts_chunk_gpu, + Cint(num_nodes), Cint(num_samples), Cint(k_bins), + Cint(z_start), Cint(z_curr_chunk_size); + blocks=blocks, threads=threads, + ) - cudacall( - mi_si_kernel, - ( - CuPtr{Cint}, CuPtr{Cdouble}, CuPtr{Cdouble}, CuPtr{Cdouble}, - Cint, Cint, Cint, Cint, Cint, - ), - counts_chunk_gpu, marginals_gpu, mi_matrix_gpu, si_chunk_gpu, - Cint(num_nodes), Cint(num_samples), Cint(k_bins), - Cint(z_start), Cint(z_curr_chunk_size); - blocks=blocks, threads=threads, - ) + cudacall( + mi_si_kernel, + ( + CuPtr{Cint}, CuPtr{Cdouble}, CuPtr{Cdouble}, CuPtr{Cdouble}, + Cint, Cint, Cint, Cint, Cint, + ), + counts_chunk_gpu, marginals_gpu, mi_matrix_gpu, si_chunk_gpu, + Cint(num_nodes), Cint(num_samples), Cint(k_bins), + Cint(z_start), Cint(z_curr_chunk_size); + blocks=blocks, threads=threads, + ) - cudacall( - puc_accumulation_kernel, - ( - CuPtr{Cdouble}, CuPtr{Cdouble}, CuPtr{Cdouble}, CuPtr{Cdouble}, - Cint, Cint, Cint, Cint, - ), - si_chunk_gpu, mi_matrix_gpu, puc_scores_gpu, marginals_gpu, - Cint(num_nodes), Cint(k_bins), - Cint(z_start), Cint(z_curr_chunk_size); - blocks=blocks, threads=threads, - ) - end + cudacall( + puc_accumulation_kernel, + ( + CuPtr{Cdouble}, CuPtr{Cdouble}, CuPtr{Cdouble}, CuPtr{Cdouble}, + Cint, Cint, Cint, Cint, + ), + si_chunk_gpu, mi_matrix_gpu, puc_scores_gpu, marginals_gpu, + Cint(num_nodes), Cint(k_bins), + Cint(z_start), Cint(z_curr_chunk_size); + blocks=blocks, threads=threads, + ) + end - # Kernels write row-major (x, z) into a Julia array whose dimensions were - # reversed above, so transpose the square outputs back to Julia's convention. - mi_matrix_cpu = permutedims(Array(mi_matrix_gpu)) - puc_scores_cpu = permutedims(Array(puc_scores_gpu)) + # Kernels write row-major (x, z) into Julia arrays whose dimensions are + # reversed above, so transpose the square outputs back to Julia's convention. + mi_matrix_cpu = permutedims(Array(mi_matrix_gpu)) + puc_scores_cpu = permutedims(Array(puc_scores_gpu)) + + # Symmetrize PUC scores: each ordered pair contains one directional + # contribution from the shared kernel. + for i = 1:num_nodes + for j = (i+1):num_nodes + val = puc_scores_cpu[i, j] + puc_scores_cpu[j, i] + puc_scores_cpu[i, j] = val + puc_scores_cpu[j, i] = val + end + end - # Symmetrize PUC scores: each ordered pair contains one directional - # contribution from the shared kernel. - for i = 1:num_nodes - for j = (i+1):num_nodes - val = puc_scores_cpu[i, j] + puc_scores_cpu[j, i] - puc_scores_cpu[i, j] = val - puc_scores_cpu[j, i] = val + return mi_matrix_cpu, puc_scores_cpu + finally + # Return allocations to CUDA.jl's pool immediately, including on a kernel + # error. This prevents a failed or repeated run from retaining pressure + # until Julia's GC notices the arrays. + for array in ( + counts_chunk_gpu, + si_chunk_gpu, + puc_scores_gpu, + mi_matrix_gpu, + marginals_gpu, + data_gpu, + ) + array === nothing || CUDA.unsafe_free!(array) end end - - return mi_matrix_cpu, puc_scores_cpu end # --- Bayesian-block CUDA backend ------------------------------------------- @@ -392,7 +514,11 @@ function _bb_problem_bytes( sizeof(Float64) * (u + 1) + # block lengths sizeof(CountT) * u + # prefix counts sizeof(Float64) * u + # best scores - sizeof(IndexT) * u # back-pointers + sizeof(IndexT) * u + # back-pointers + sizeof(Int64) + # state offset + sizeof(Int64) + # block offset + sizeof(Int32) + # unique count + sizeof(Float64) # final score ) end @@ -518,6 +644,12 @@ function _solve_bb_cuda_batch_with_priors( "from 32, 64, 128, or 256; got $threads", ), ) + length(problem_indices) <= _KERNEL_INT_MAX || throw( + ArgumentError( + "CUDA Bayesian blocks batch has $(length(problem_indices)) genes, " * + "which exceeds the signed 32-bit block-index limit.", + ), + ) prefix_counts, block_lengths, state_offsets, block_offsets, unique_counts = _flatten_bb_batch(problems, problem_indices, CountT) @@ -613,8 +745,18 @@ function FastPIDC.solve_bayesian_blocks_cuda( CUDA.functional() || return nothing isempty(problems) && return FastPIDC.BayesianBlocksSolution[] + # Load the module before measuring reusable memory so kernel/module residency + # is already reflected in the driver's and CUDA.jl pool's accounting. + _get_module() + sample_count = maximum(p -> Int(round(p.prefix_counts[end])), problems) max_u = maximum(p -> length(p.prefix_counts), problems) + max_u <= _KERNEL_INT_MAX || throw( + ArgumentError( + "Bayesian blocks CUDA backend supports at most $(_KERNEL_INT_MAX) " * + "unique values per gene; got $max_u", + ), + ) # A cumulative prefix count can reach the number of cells, so select the # smallest exact unsigned type that guards against overflow for this input. @@ -622,14 +764,13 @@ function FastPIDC.solve_bayesian_blocks_cuda( # Back-pointers only need to represent candidate indices up to U_g. IndexT = _smallest_unsigned_type(max_u) - free_bytes = Int(CUDA.free_memory()) - # Keep headroom for the CUDA context, allocator bookkeeping, and other - # active package allocations while still using most of the currently free - # device memory. - budget_bytes = max(1, floor(Int, 0.65 * Float64(free_bytes))) - buckets = _bb_quantile_buckets(problems) solutions = Vector{FastPIDC.BayesianBlocksSolution}(undef, length(problems)) + + # The priors remain live across all batches. Allocate them first, then size + # each bucket from the *remaining* reusable memory at runtime. Every batch is + # kept within 65% of what is free at that moment, and _bb_problem_bytes + # includes all per-gene device metadata, not just the large state arrays. priors_gpu = CuArray(_bb_prior_values(max_u)) if verbose @@ -640,16 +781,15 @@ function FastPIDC.solve_bayesian_blocks_cuda( "U_g median=$median_u, max=$(unique_counts[end]), " * "prefix counts=$(CountT), back-pointers=$(IndexT)", ) - println( - "[FastPIDC] CUDA Bayesian blocks memory budget: " * - "$(round(budget_bytes / 2.0^30; digits = 2)) GiB", - ) end try for (bucket_number, bucket) in enumerate(buckets) bucket_max_u = maximum(i -> length(problems[i].prefix_counts), bucket) threads = _bb_threads_for_max_u(bucket_max_u) + + free_bytes = _available_gpu_memory_bytes() + budget_bytes = _gpu_memory_budget_bytes(free_bytes) batches = _bb_memory_batches( bucket, problems, @@ -663,7 +803,9 @@ function FastPIDC.solve_bayesian_blocks_cuda( println( "[FastPIDC] CUDA BB bucket $bucket_number/$(length(buckets)): " * "$(length(bucket)) genes, U_g=$bucket_min_u:$bucket_max_u, " * - "threads=$threads, batches=$(length(batches))", + "threads=$threads, batches=$(length(batches)), " * + "reusable=$(round(free_bytes / 2.0^30; digits = 2)) GiB, " * + "65% budget=$(round(budget_bytes / 2.0^30; digits = 2)) GiB", ) end diff --git a/python/src/fastpidc/cuda.py b/python/src/fastpidc/cuda.py index 4f8a37b..48ca91f 100644 --- a/python/src/fastpidc/cuda.py +++ b/python/src/fastpidc/cuda.py @@ -30,13 +30,9 @@ _KERNEL_SOURCE_PATH = Path(__file__).with_name("kernels") / "pidc_kernels.cu" _MAX_CHUNK_SIZE = 256 -# Largest bins-per-gene the shared kernels can index: joint_counts_kernel forms -# the bin-pair index u * k_bins + v in int32, so k_bins**2 must fit there -# (isqrt(2**31 - 1) == 46340). Every other flat offset is 64-bit. -_MAX_K_BINS = 46340 -# Headroom for the fixed buffers, allocator overhead and fragmentation, matching -# the safety factor used by FastPIDC.jl's CUDA extension. -_CHUNK_MEMORY_SAFETY_FACTOR = 0.8 +_KERNEL_INT_MAX = np.iinfo(np.int32).max +_GPU_MEMORY_BUDGET_NUMERATOR = 65 +_GPU_MEMORY_BUDGET_DENOMINATOR = 100 _FALLBACK_CUDA_HEADER_DIRS = ("/usr", "/usr/local/cuda") @@ -82,51 +78,121 @@ def _load_module(): return cp.RawModule(code=source, options=("--std=c++11",)) -def _check_kernel_index_limits(k_bins: int) -> None: - """Validate the one kernel limit that is not lifted by 64-bit indexing. - - The shared kernels compute every flat buffer offset in 64-bit, with a single - exception: ``joint_counts_kernel`` forms the bin-pair index - ``u * k_bins + v`` in ``int``, which is exact only while ``k_bins**2`` fits - in int32. FastPIDC.jl's ``_check_kernel_index_limits`` enforces the same - bound. +def _gpu_memory_budget_bytes(free_bytes: int) -> int: + if free_bytes <= 0: + raise ValueError("free_bytes must be positive") + return free_bytes * _GPU_MEMORY_BUDGET_NUMERATOR // _GPU_MEMORY_BUDGET_DENOMINATOR + + +def _reusable_gpu_memory_bytes( + driver_free_bytes: int, + pool_free_bytes: int = 0, + *, + pool_used_bytes: int = 0, + pool_limit_bytes: int = 0, +) -> int: + """Memory this process can reuse without exceeding the device/pool limits.""" + for name, value in ( + ("driver_free_bytes", driver_free_bytes), + ("pool_free_bytes", pool_free_bytes), + ("pool_used_bytes", pool_used_bytes), + ("pool_limit_bytes", pool_limit_bytes), + ): + if value < 0: + raise ValueError(f"{name} must be nonnegative") + + available = driver_free_bytes + pool_free_bytes + if pool_limit_bytes: + available = min(available, max(pool_limit_bytes - pool_used_bytes, 0)) + return available + + +def _available_gpu_memory_bytes(cp) -> int: + """Return VRAM reusable by CuPy, including cached free pool blocks. + + ``cudaMemGetInfo`` does not include blocks CuPy has already reserved but is + no longer using. Counting those blocks prevents repeated FastPIDC calls from + becoming artificially more conservative. A configured CuPy pool limit is + also honored so the planner cannot approve a batch the allocator will reject. """ - if k_bins > _MAX_K_BINS: - raise RuntimeError( - f"compute_puc_full_cuda: the discretizer selected k_bins={k_bins} bins per " - f"gene, but the CUDA kernels index bin pairs in 32-bit and support at most " - f'{_MAX_K_BINS}. Use discretizer="uniform_width" with a fixed, small ' - f"number_of_bins, or config.backend = 'cpu'." - ) + driver_free_bytes = int(cp.cuda.runtime.memGetInfo()[0]) + pool = cp.get_default_memory_pool() + return _reusable_gpu_memory_bytes( + driver_free_bytes, + int(pool.free_bytes()), + pool_used_bytes=int(pool.used_bytes()), + pool_limit_bytes=int(pool.get_limit()), + ) -def _chunk_size_for_free_memory(n: int, k_bins: int, free_bytes: int) -> int: - """Largest target-gene chunk whose intermediate buffers fit in - ``free_bytes`` of device memory, capped at :data:`_MAX_CHUNK_SIZE`. +def _check_kernel_scalar_limits(n: int, m: int, k_bins: int) -> None: + """Validate the signed-int32 launch scalars used by the shared kernels. + + Every composed flat-buffer offset is 64-bit. Only the ABI-level dimensions + remain int32, so reject impossible values before numpy/cupy can narrow them. + """ + for name, value in (("n", n), ("m", m), ("k_bins", k_bins)): + if value < 1: + raise ValueError(f"{name} must be positive; got {value}") + if value > _KERNEL_INT_MAX: + raise ValueError( + f"compute_puc_full_cuda: {name}={value} exceeds the CUDA kernel's " + f"signed 32-bit scalar limit of {_KERNEL_INT_MAX}." + ) - The chunked intermediates scale with ``k_bins**2 * n`` (joint counts) and - ``k_bins * n`` (specific information) per target gene, so an adaptive - discretizer that picks many bins can make even one gene per chunk too - large; that case raises instead of failing inside the allocator. - This bounds *memory availability* only - it is deliberately not a bound on - the flat element index. The kernels index these buffers in 64-bit precisely - so the chunk can be sized from free memory without an int32 element cap. +def _puc_memory_plan( + n: int, + m: int, + k_bins: int, + free_bytes: int, + *, + requested_chunk_size: int | None = None, +) -> tuple[int, int, int, int]: + """Plan the entire PUC device footprint against 65% of currently reusable VRAM. + + Returns ``(chunk_size, fixed_bytes, bytes_per_chunk_column, budget_bytes)``. + Python integers are unbounded, so the planning arithmetic itself cannot wrap. """ - bytes_per_chunk_column = k_bins**2 * n * np.dtype(np.int32).itemsize + k_bins * n * np.dtype(np.float64).itemsize - usable_bytes = free_bytes * _CHUNK_MEMORY_SAFETY_FACTOR - max_chunk_size = int(usable_bytes // bytes_per_chunk_column) - if max_chunk_size < 1: + budget_bytes = _gpu_memory_budget_bytes(free_bytes) + + fixed_bytes = ( + n * m * np.dtype(np.int32).itemsize + + n * k_bins * np.dtype(np.float64).itemsize + + 2 * n * n * np.dtype(np.float64).itemsize + ) + bytes_per_chunk_column = ( + k_bins * k_bins * n * np.dtype(np.int32).itemsize + + k_bins * n * np.dtype(np.float64).itemsize + ) + + if fixed_bytes + bytes_per_chunk_column > budget_bytes: raise RuntimeError( - f"compute_puc_full_cuda: even a single-gene chunk would require " - f"{bytes_per_chunk_column / 2**30:.2f} GiB of GPU memory (only " - f"{usable_bytes / 2**30:.2f} GiB usable), because the discretizer selected " - f"k_bins={k_bins} bins per gene. This is usually caused by an adaptive " - f'discretizer (e.g. "bayesian_blocks", the default) picking an unbounded ' - f'number of bins on a dataset with many samples. Try discretizer="uniform_width" ' - f"with a fixed, small number_of_bins (e.g. 10-20), or config.backend = 'cpu'." + "compute_puc_full_cuda: the fixed GPU buffers plus a one-gene chunk " + f"would require {(fixed_bytes + bytes_per_chunk_column) / 2**30:.2f} GiB, " + "which exceeds the configured 65% memory budget " + f"({budget_bytes / 2**30:.2f} GiB of {free_bytes / 2**30:.2f} GiB currently reusable). " + "Reduce the number of genes/samples/bins, use a fixed small-bin discretizer, " + "or use config.backend = 'cpu'." ) - return min(max_chunk_size, _MAX_CHUNK_SIZE, n) + + safe_max_chunk = min((budget_bytes - fixed_bytes) // bytes_per_chunk_column, _MAX_CHUNK_SIZE, n) + if requested_chunk_size is None: + chunk_size = int(safe_max_chunk) + else: + if requested_chunk_size < 1: + raise ValueError("chunk_size must be positive") + if requested_chunk_size > n: + raise ValueError(f"chunk_size={requested_chunk_size} exceeds n={n}") + if requested_chunk_size > safe_max_chunk: + raise RuntimeError( + f"compute_puc_full_cuda: requested chunk_size={requested_chunk_size} would exceed " + f"the configured 65% GPU-memory budget; the largest safe chunk is {safe_max_chunk} " + f"with {free_bytes / 2**30:.2f} GiB currently reusable." + ) + chunk_size = requested_chunk_size + + return int(chunk_size), int(fixed_bytes), int(bytes_per_chunk_column), int(budget_bytes) def compute_puc_full_cuda( @@ -134,14 +200,10 @@ def compute_puc_full_cuda( ) -> tuple[np.ndarray, np.ndarray]: """GPU implementation of :func:`fastpidc.puc.compute_puc_full`. - Parameters - ---------- - base : accepted for signature parity with the CPU path but unused: mutual - information is always computed in base 2 on the GPU, matching - FastPIDC.jl's CUDA extension. - chunk_size : number of target genes processed per pass. ``None`` (the - default) sizes it from the currently free device memory, as - FastPIDC.jl does. + The complete device footprint is planned before large allocations against + 65% of VRAM currently reusable by CuPy (driver-free memory plus unused pool + blocks, bounded by any configured pool limit). ``chunk_size=None`` chooses + the largest safe chunk up to 256; an explicit chunk must also fit the same budget. """ del base if not cuda_available(): @@ -150,6 +212,8 @@ def compute_puc_full_cuda( "GPU was detected. Install the 'cuda' extra and ensure a GPU is " "available, or use config.backend = 'cpu'." ) + if not nodes: + raise ValueError("compute_puc_full_cuda requires at least one node") import cupy as cp @@ -159,96 +223,121 @@ def compute_puc_full_cuda( puc_accumulation_kernel = module.get_function("puc_accumulation_kernel") n = len(nodes) - m = nodes[0].binned_values.size + m = int(nodes[0].binned_values.size) + if any(node.binned_values.size != m for node in nodes): + raise ValueError("all nodes must contain the same number of discretized samples") k_bins = max(node.number_of_bins for node in nodes) - # Checked here rather than in _chunk_size_for_free_memory, which is skipped - # entirely when the caller passes an explicit chunk_size. - _check_kernel_index_limits(k_bins) - - if chunk_size is None: - chunk_size = _chunk_size_for_free_memory(n, k_bins, int(cp.cuda.runtime.memGetInfo()[0])) + _check_kernel_scalar_limits(n, m, k_bins) + + free_bytes = _available_gpu_memory_bytes(cp) + chunk_size, fixed_bytes, _, budget_bytes = _puc_memory_plan( + n, + m, + k_bins, + free_bytes, + requested_chunk_size=chunk_size, + ) data_host = np.zeros((m, n), dtype=np.int32) marginals_host = np.zeros((k_bins, n), dtype=np.float64) for i, node in enumerate(nodes): + if node.number_of_bins < 1: + raise ValueError(f"node {node.label!r} has non-positive number_of_bins") + if node.binned_values.size and ( + node.binned_values.min() < 0 or node.binned_values.max() >= node.number_of_bins + ): + raise ValueError(f"node {node.label!r} contains a bin id outside [0, number_of_bins)") + if node.probabilities.size > k_bins: + raise ValueError(f"node {node.label!r} has more probabilities than k_bins") + data_host[:, i] = node.binned_values marginals_host[: node.number_of_bins, i] = node.probabilities - data_gpu = cp.asarray(data_host) - marginals_gpu = cp.asarray(marginals_host) - - puc_scores_gpu = cp.zeros((n, n), dtype=cp.float64) - mi_matrix_gpu = cp.zeros((n, n), dtype=cp.float64) + data_gpu = marginals_gpu = None + puc_scores_gpu = mi_matrix_gpu = None + counts_chunk_gpu = si_chunk_gpu = None - counts_chunk_gpu = cp.zeros((k_bins, k_bins, n, chunk_size), dtype=cp.int32) - si_chunk_gpu = cp.zeros((k_bins, n, chunk_size), dtype=cp.float64) - - threads = (16, 16) - - if verbose: - n_chunks = -(-n // chunk_size) - print(f"[fastpidc] GPU chunked PUC: processing {n} x {n} pairs (k_bins={k_bins})...") - print(f"[fastpidc] Using chunk size {chunk_size} ({n_chunks} iterations)") - - for z_start in range(0, n, chunk_size): - z_chunk_size = min(chunk_size, n - z_start) - - counts_chunk_gpu.fill(0) - si_chunk_gpu.fill(0.0) - - blocks = (-(-n // threads[0]), -(-z_chunk_size // threads[1])) - - joint_counts_kernel( - blocks, - threads, - ( - data_gpu, - counts_chunk_gpu, - np.int32(n), - np.int32(m), - np.int32(k_bins), - np.int32(z_start), - np.int32(z_chunk_size), - ), - ) - mi_si_kernel( - blocks, - threads, - ( - counts_chunk_gpu, - marginals_gpu, - mi_matrix_gpu, - si_chunk_gpu, - np.int32(n), - np.int32(m), - np.int32(k_bins), - np.int32(z_start), - np.int32(z_chunk_size), - ), - ) - puc_accumulation_kernel( - blocks, - threads, - ( - si_chunk_gpu, - mi_matrix_gpu, - puc_scores_gpu, - marginals_gpu, - np.int32(n), - np.int32(k_bins), - np.int32(z_start), - np.int32(z_chunk_size), - ), - ) + try: + data_gpu = cp.asarray(data_host) + marginals_gpu = cp.asarray(marginals_host) + puc_scores_gpu = cp.zeros((n, n), dtype=cp.float64) + mi_matrix_gpu = cp.zeros((n, n), dtype=cp.float64) + counts_chunk_gpu = cp.zeros((k_bins, k_bins, n, chunk_size), dtype=cp.int32) + si_chunk_gpu = cp.zeros((k_bins, n, chunk_size), dtype=cp.float64) - puc_scores = cp.asnumpy(puc_scores_gpu) - mi_matrix = cp.asnumpy(mi_matrix_gpu) + threads = (16, 16) - # Symmetrize: each ordered pair (x, z) only holds one of the two - # directional contributions (see the module docstring in pidc_kernels.cu). - puc_scores = puc_scores + puc_scores.T + if verbose: + n_chunks = -(-n // chunk_size) + print(f"[fastpidc] GPU chunked PUC: processing {n} x {n} pairs (k_bins={k_bins})...") + print( + f"[fastpidc] GPU memory: {free_bytes / 2**30:.2f} GiB reusable; " + f"65% budget={budget_bytes / 2**30:.2f} GiB; " + f"fixed={fixed_bytes / 2**30:.2f} GiB" + ) + print(f"[fastpidc] Using chunk size {chunk_size} ({n_chunks} iterations)") + + for z_start in range(0, n, chunk_size): + z_chunk_size = min(chunk_size, n - z_start) + + counts_chunk_gpu.fill(0) + si_chunk_gpu.fill(0.0) + + blocks = (-(-n // threads[0]), -(-z_chunk_size // threads[1])) + + joint_counts_kernel( + blocks, + threads, + ( + data_gpu, + counts_chunk_gpu, + np.int32(n), + np.int32(m), + np.int32(k_bins), + np.int32(z_start), + np.int32(z_chunk_size), + ), + ) + mi_si_kernel( + blocks, + threads, + ( + counts_chunk_gpu, + marginals_gpu, + mi_matrix_gpu, + si_chunk_gpu, + np.int32(n), + np.int32(m), + np.int32(k_bins), + np.int32(z_start), + np.int32(z_chunk_size), + ), + ) + puc_accumulation_kernel( + blocks, + threads, + ( + si_chunk_gpu, + mi_matrix_gpu, + puc_scores_gpu, + marginals_gpu, + np.int32(n), + np.int32(k_bins), + np.int32(z_start), + np.int32(z_chunk_size), + ), + ) - return mi_matrix, puc_scores + puc_scores = cp.asnumpy(puc_scores_gpu) + mi_matrix = cp.asnumpy(mi_matrix_gpu) + puc_scores = puc_scores + puc_scores.T + return mi_matrix, puc_scores + finally: + # Drop references immediately even on a failed launch/allocation. CuPy's + # pool may cache the released blocks, but they are no longer live and can + # be reused or returned under memory pressure. + del counts_chunk_gpu, si_chunk_gpu + del puc_scores_gpu, mi_matrix_gpu, marginals_gpu, data_gpu # --- Bayesian blocks -------------------------------------------------------- @@ -259,9 +348,6 @@ def compute_puc_full_cuda( # flattened into the packed buffers the shared kernel indexes with per-gene # offsets. Back-pointers come back 0-based and are backtracked on the host. -# Keep headroom for the CUDA context, allocator bookkeeping and other live -# allocations while still using most of the free device memory. -_BB_MEMORY_BUDGET_FRACTION = 0.65 _BB_VALID_THREAD_COUNTS = (32, 64, 128, 256) _BB_TYPE_SUFFIX = { np.dtype(np.uint8): "u8", @@ -351,6 +437,10 @@ def _bb_problem_bytes(problem: BayesianBlocksProblem, count_dtype: np.dtype, ind + count_dtype.itemsize * u # prefix counts + 8 * u # best scores + index_dtype.itemsize * u # back-pointers + + 8 # state offset + + 8 # block offset + + 4 # unique count + + 8 # final score ) @@ -391,9 +481,12 @@ def _flatten_bb_batch( ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Pack a batch of problems into the flat buffers the shared kernel reads, with 0-based per-gene offsets.""" - unique_counts = np.array([problems[i].prefix_counts.size for i in problem_indices], dtype=np.int32) - if unique_counts.size and unique_counts.max() > np.iinfo(np.int32).max: - raise ValueError("Bayesian-blocks CUDA backend supports at most 2**31-1 unique values per gene") + unique_count_values = [int(problems[i].prefix_counts.size) for i in problem_indices] + if any(u > _KERNEL_INT_MAX for u in unique_count_values): + raise ValueError( + f"Bayesian-blocks CUDA backend supports at most {_KERNEL_INT_MAX} unique values per gene" + ) + unique_counts = np.asarray(unique_count_values, dtype=np.int32) state_offsets = np.zeros(len(problem_indices), dtype=np.int64) block_offsets = np.zeros(len(problem_indices), dtype=np.int64) @@ -452,6 +545,11 @@ def _solve_bb_batch( if threads not in _BB_VALID_THREAD_COUNTS: raise ValueError(f"CUDA Bayesian blocks requires a thread count from {_BB_VALID_THREAD_COUNTS}; got {threads}") + if len(problem_indices) > _KERNEL_INT_MAX: + raise ValueError( + f"CUDA Bayesian blocks batch has {len(problem_indices)} genes, " + "which exceeds the signed 32-bit block-index limit." + ) prefix_counts, block_lengths, state_offsets, block_offsets, unique_counts = _flatten_bb_batch( problems, problem_indices, count_dtype @@ -510,11 +608,9 @@ def solve_bayesian_blocks_cuda( :func:`fastpidc.discretizers.solve_bayesian_blocks_cpu`, so the selected change points agree exactly with the CPU reference. - Raises - ------ - RuntimeError - If no GPU backend is available, or if one gene's problem alone exceeds - the device memory budget. + Each workload bucket is sized from 65% of the device memory reusable by + CuPy at that point in the run. The per-problem byte estimate includes every + packed device buffer and metadata array. """ if not cuda_available(): raise RuntimeError( @@ -526,17 +622,22 @@ def solve_bayesian_blocks_cuda( import cupy as cp - # A cumulative prefix count reaches the observation count, and back-pointers - # only need to index candidates up to U_g, so pick the narrowest exact type - # for each rather than paying for 64-bit buffers on every dataset. + # Compile/load before measuring reusable memory so module residency is + # already reflected in the driver/pool accounting. + _load_module() + sample_count = max(int(round(float(p.prefix_counts[-1]))) for p in problems) - max_u = max(p.prefix_counts.size for p in problems) + max_u = max(int(p.prefix_counts.size) for p in problems) + if max_u > _KERNEL_INT_MAX: + raise ValueError( + f"Bayesian-blocks CUDA backend supports at most {_KERNEL_INT_MAX} unique values per gene" + ) + + # A cumulative prefix count reaches the observation count, and back-pointers + # only need to index candidates up to U_g, so pick the narrowest exact type. count_dtype = _smallest_unsigned_dtype(sample_count) index_dtype = _smallest_unsigned_dtype(max_u) - free_bytes = int(cp.cuda.runtime.memGetInfo()[0]) - budget_bytes = max(1, int(_BB_MEMORY_BUDGET_FRACTION * free_bytes)) - if verbose: unique_counts = sorted(p.prefix_counts.size for p in problems) median_u = unique_counts[(len(unique_counts) - 1) // 2] @@ -545,27 +646,36 @@ def solve_bayesian_blocks_cuda( f"U_g median={median_u}, max={unique_counts[-1]}, " f"prefix counts={count_dtype}, back-pointers={index_dtype}" ) - print(f"[fastpidc] CUDA Bayesian blocks memory budget: {budget_bytes / 2**30:.2f} GiB") priors_gpu = cp.asarray(_bb_prior_values(max_u)) solutions: list[BayesianBlocksSolution | None] = [None] * len(problems) - for bucket_number, bucket in enumerate(_bb_quantile_buckets(problems), start=1): - bucket_max_u = max(problems[i].prefix_counts.size for i in bucket) - threads = _bb_threads_for_max_u(bucket_max_u) - batches = _bb_memory_batches(bucket, problems, budget_bytes, count_dtype, index_dtype) - - if verbose: - bucket_min_u = min(problems[i].prefix_counts.size for i in bucket) - print( - f"[fastpidc] CUDA BB bucket {bucket_number}: {len(bucket)} genes, " - f"U_g={bucket_min_u}:{bucket_max_u}, threads={threads}, batches={len(batches)}" - ) - - for batch in batches: - for problem_index, solution in zip( - batch, _solve_bb_batch(problems, batch, threads, count_dtype, index_dtype, priors_gpu) - ): - solutions[problem_index] = solution - - return solutions # type: ignore[return-value] + try: + for bucket_number, bucket in enumerate(_bb_quantile_buckets(problems), start=1): + bucket_max_u = max(problems[i].prefix_counts.size for i in bucket) + threads = _bb_threads_for_max_u(bucket_max_u) + + # priors_gpu is already live. Re-evaluate memory available to CuPy + # for every bucket, including reusable free blocks in its memory + # pool and respecting any configured pool limit. + free_bytes = _available_gpu_memory_bytes(cp) + budget_bytes = _gpu_memory_budget_bytes(free_bytes) + batches = _bb_memory_batches(bucket, problems, budget_bytes, count_dtype, index_dtype) + + if verbose: + bucket_min_u = min(problems[i].prefix_counts.size for i in bucket) + print( + f"[fastpidc] CUDA BB bucket {bucket_number}: {len(bucket)} genes, " + f"U_g={bucket_min_u}:{bucket_max_u}, threads={threads}, batches={len(batches)}, " + f"reusable={free_bytes / 2**30:.2f} GiB, 65% budget={budget_bytes / 2**30:.2f} GiB" + ) + + for batch in batches: + for problem_index, solution in zip( + batch, _solve_bb_batch(problems, batch, threads, count_dtype, index_dtype, priors_gpu) + ): + solutions[problem_index] = solution + + return solutions # type: ignore[return-value] + finally: + del priors_gpu diff --git a/python/src/fastpidc/kernels/pidc_kernels.cu b/python/src/fastpidc/kernels/pidc_kernels.cu index a34ed03..e13e071 100644 --- a/python/src/fastpidc/kernels/pidc_kernels.cu +++ b/python/src/fastpidc/kernels/pidc_kernels.cu @@ -38,9 +38,9 @@ // faulted with CUDA_ERROR_ILLEGAL_ADDRESS. Every composed flat offset below is // therefore `long long`, hoisted out of the hot loops as a base offset plus a // loop-invariant stride so the inner bodies cost 64-bit adds rather than 32-bit -// multiply-add chains. The single exception is the bin-pair index -// `u * k_bins + v` in joint_counts_kernel, which stays 32-bit and so requires -// k_bins <= 46340; both hosts check that before launching. +// multiply-add chains. The bin-pair term `u * k_bins + v` is widened before +// multiplication as well, so no composed flat-buffer offset relies on 32-bit +// arithmetic. extern "C" { @@ -76,10 +76,10 @@ __global__ void joint_counts_kernel( int u = data[data_row + x]; int v = data[data_row + z_global]; if (u >= 0 && u < k_bins && v >= 0 && v < k_bins) { - // u, v < k_bins was just checked, so u * k_bins + v < k_bins^2, which - // both hosts keep inside int (k_bins <= 46340); plane_stride carries - // the 64-bit range. - long long idx = (long long)(u * k_bins + v) * plane_stride + cell; + // Widen before forming the bin-pair index: k_bins is an int32 + // launch scalar, but k_bins^2 need not fit in int32. + const long long bin_pair = (long long)u * k_bins + v; + long long idx = bin_pair * plane_stride + cell; atomicAdd(&counts[idx], 1); } } @@ -321,16 +321,20 @@ __device__ void fastpidc_bayesian_blocks_dp( double local_best = fastpidc_negative_infinity(); int local_i = FASTPIDC_BB_NO_CANDIDATE; - for (int i = tid; i <= k; i += nthreads) { + // Use a 64-bit loop cursor so the final `i += nthreads` cannot wrap if + // n_unique approaches the signed-int32 ABI ceiling. Candidate values + // themselves are still <= INT32_MAX and are narrowed only after checking. + for (long long i64 = tid; i64 <= (long long)k; i64 += nthreads) { + const int i = (int)i64; const double prefix_before = - (i == 0) ? 0.0 : (double)prefix_counts[state_start + i - 1]; + (i == 0) ? 0.0 : (double)prefix_counts[state_start + i64 - 1]; const double count = prefix_k - prefix_before; - const double width = block_lengths[block_start + i] - block_length_end; + const double width = block_lengths[block_start + i64] - block_length_end; // Fitness function (eq. 19) and prior (eq. 21) from Scargle 2012. double fit = count * log(count / width) - prior; if (i > 0) { - fit += best[state_start + i - 1]; + fit += best[state_start + i64 - 1]; } if (fastpidc_bb_take_other(fit, i, local_best, local_i)) { diff --git a/python/tests/test_cuda.py b/python/tests/test_cuda.py index 85bfc2f..ec189d8 100644 --- a/python/tests/test_cuda.py +++ b/python/tests/test_cuda.py @@ -14,15 +14,17 @@ import pytest from fastpidc.cuda import ( + _KERNEL_INT_MAX, _MAX_CHUNK_SIZE, - _MAX_K_BINS, _bb_kernel_name, _bb_memory_batches, _bb_problem_bytes, _bb_quantile_buckets, _bb_threads_for_max_u, - _check_kernel_index_limits, - _chunk_size_for_free_memory, + _check_kernel_scalar_limits, + _gpu_memory_budget_bytes, + _puc_memory_plan, + _reusable_gpu_memory_bytes, _smallest_unsigned_dtype, cuda_available, ) @@ -58,27 +60,64 @@ def test_cuda_available_returns_a_bool(): assert isinstance(cuda_available(), bool) +def test_gpu_memory_budget_is_65_percent_of_currently_free_memory(): + assert _gpu_memory_budget_bytes(1000) == 650 + + +def test_reusable_gpu_memory_includes_cached_pool_blocks(): + assert _reusable_gpu_memory_bytes(1000, 400, pool_used_bytes=250) == 1400 + + +def test_reusable_gpu_memory_honors_a_cupy_pool_limit(): + # The pool can reuse its free blocks and grow only until the configured + # limit. Here the physical device would allow 1400 bytes, but the pool has + # only 650 bytes of headroom from its current live usage. + assert ( + _reusable_gpu_memory_bytes(1000, 400, pool_used_bytes=250, pool_limit_bytes=900) + == 650 + ) + + +def test_puc_memory_plan_counts_fixed_and_chunk_buffers_exactly(): + n, m, k_bins = 3, 5, 7 + _, fixed_bytes, bytes_per_chunk_column, _ = _puc_memory_plan( + n=n, m=m, k_bins=k_bins, free_bytes=64 * GIB + ) + assert fixed_bytes == n * m * 4 + n * k_bins * 8 + 2 * n * n * 8 + assert bytes_per_chunk_column == k_bins * k_bins * n * 4 + k_bins * n * 8 + + def test_chunk_size_is_capped_by_the_maximum(): - # Plenty of memory for a small problem: the fixed cap applies. - assert _chunk_size_for_free_memory(n=1000, k_bins=4, free_bytes=64 * GIB) == _MAX_CHUNK_SIZE + chunk, _, _, _ = _puc_memory_plan(n=1000, m=1000, k_bins=4, free_bytes=64 * GIB) + assert chunk == _MAX_CHUNK_SIZE def test_chunk_size_never_exceeds_the_number_of_genes(): - assert _chunk_size_for_free_memory(n=10, k_bins=4, free_bytes=64 * GIB) == 10 + chunk, _, _, _ = _puc_memory_plan(n=10, m=1000, k_bins=4, free_bytes=64 * GIB) + assert chunk == 10 def test_chunk_size_shrinks_when_memory_is_tight(): - tight = _chunk_size_for_free_memory(n=5000, k_bins=64, free_bytes=8 * GIB) - roomy = _chunk_size_for_free_memory(n=5000, k_bins=64, free_bytes=64 * GIB) + tight, _, _, _ = _puc_memory_plan(n=5000, m=1000, k_bins=64, free_bytes=8 * GIB) + roomy, _, _, _ = _puc_memory_plan(n=5000, m=1000, k_bins=64, free_bytes=64 * GIB) assert 1 <= tight < roomy <= _MAX_CHUNK_SIZE -def test_chunk_size_raises_when_a_single_gene_chunk_does_not_fit(): - # An adaptive discretizer choosing thousands of bins makes the per-target - # intermediates exceed device memory; that must be an explicit error rather - # than an allocator failure deep inside the kernel launch loop. - with pytest.raises(RuntimeError, match="single-gene chunk"): - _chunk_size_for_free_memory(n=20000, k_bins=4000, free_bytes=8 * GIB) +def test_memory_plan_raises_before_allocating_when_one_gene_chunk_does_not_fit(): + with pytest.raises(RuntimeError, match="one-gene chunk"): + _puc_memory_plan(n=20000, m=1000, k_bins=4000, free_bytes=8 * GIB) + + +def test_explicit_chunk_size_must_fit_the_same_memory_budget(): + safe, _, _, _ = _puc_memory_plan(n=5000, m=1000, k_bins=64, free_bytes=8 * GIB) + with pytest.raises(RuntimeError, match="requested chunk_size"): + _puc_memory_plan( + n=5000, + m=1000, + k_bins=64, + free_bytes=8 * GIB, + requested_chunk_size=safe + 1, + ) @pytest.mark.skipif(not cuda_available(), reason="no functional GPU / cupy backend available") @@ -170,7 +209,9 @@ def test_bb_quantile_buckets_of_empty_input(): def test_bb_problem_bytes_counts_every_device_buffer(): problem = _bb_problems(sizes=(40,))[0] u = problem.prefix_counts.size - expected = 8 * (u + 1) + 1 * u + 8 * u + 1 * u # lengths + prefix + best + back-pointers + expected = ( + 8 * (u + 1) + 1 * u + 8 * u + 1 * u + 8 + 8 + 4 + 8 + ) # state arrays + per-gene metadata assert _bb_problem_bytes(problem, np.dtype(np.uint8), np.dtype(np.uint8)) == expected @@ -322,7 +363,7 @@ def test_bayesian_blocks_falls_back_to_cpu_without_a_gpu(monkeypatch, julia_test _CONTROL_CHUNK = 32 # The overflowing chunk needs ~8 GiB plus room for cupy's pool and the other # buffers; require real headroom so the test never competes for a full device. -_OVERFLOW_FREE_MEMORY_FLOOR = 12 * GIB +_OVERFLOW_FREE_MEMORY_FLOOR = 14 * GIB def _counts_elements(k_bins: int, n: int, chunk: int) -> int: @@ -338,25 +379,39 @@ def test_production_configuration_exceeds_int32_indexing(): assert _counts_elements(_OVERFLOW_N_BINS, _OVERFLOW_N_NODES, _CONTROL_CHUNK) - 1 <= INT32_MAX -def test_check_kernel_index_limits_accepts_the_largest_supported_k_bins(): - _check_kernel_index_limits(_MAX_K_BINS) - assert _MAX_K_BINS**2 <= INT32_MAX - assert (_MAX_K_BINS + 1) ** 2 > INT32_MAX +def test_kernel_scalar_limits_allow_bin_pair_products_past_int32(): + # k_bins itself fits int32; k_bins**2 deliberately does not. The shared + # kernel widens u before multiplying by k_bins, so this is now valid + # indexing arithmetic (memory availability is a separate concern). + _check_kernel_scalar_limits(12_071, 38_176, 50_000) + assert 50_000**2 > INT32_MAX -def test_check_kernel_index_limits_rejects_a_larger_k_bins(): - # joint_counts_kernel forms u * k_bins + v in int32; everything else is 64-bit. - with pytest.raises(RuntimeError, match="bin pairs in 32-bit"): - _check_kernel_index_limits(_MAX_K_BINS + 1) +@pytest.mark.parametrize( + ("n", "m", "k_bins"), + [ + (_KERNEL_INT_MAX + 1, 1, 1), + (1, _KERNEL_INT_MAX + 1, 1), + (1, 1, _KERNEL_INT_MAX + 1), + ], +) +def test_kernel_scalar_limits_reject_values_that_would_narrow(n, m, k_bins): + with pytest.raises(ValueError, match="signed 32-bit scalar limit"): + _check_kernel_scalar_limits(n, m, k_bins) def test_chunk_sizing_is_not_capped_by_int32_element_count(): - # The memory guard bounds bytes, not the element index. Regression guard - # against anyone "fixing" the overflow by reimposing an int32 element cap, - # which would silently shrink chunks on large GPUs. - chunk = _chunk_size_for_free_memory(n=12071, k_bins=33, free_bytes=int(19.22 * GIB)) + # Memory, not int32 indexing, controls the chunk. With enough free VRAM the + # production-sized shape still reaches the 256-gene cap even though the + # joint-count buffer contains >2^31 elements. + chunk, _, _, _ = _puc_memory_plan( + n=12_071, + m=38_176, + k_bins=33, + free_bytes=32 * GIB, + ) assert chunk == 256 - assert _counts_elements(33, 12071, chunk) - 1 > INT32_MAX + assert _counts_elements(33, 12_071, chunk) - 1 > INT32_MAX @pytest.mark.largemem @@ -373,9 +428,9 @@ def test_puc_indexing_past_int32_matches_a_smaller_chunk(): import cupy as cp - from fastpidc.cuda import compute_puc_full_cuda + from fastpidc.cuda import _available_gpu_memory_bytes, compute_puc_full_cuda - free_bytes = int(cp.cuda.runtime.memGetInfo()[0]) + free_bytes = _available_gpu_memory_bytes(cp) if free_bytes < _OVERFLOW_FREE_MEMORY_FLOOR: pytest.skip( f"needs {_OVERFLOW_FREE_MEMORY_FLOOR / GIB:.0f} GiB free device memory, have {free_bytes / GIB:.1f} GiB" diff --git a/test/cuda_bayesian_blocks_tests.jl b/test/cuda_bayesian_blocks_tests.jl index 163e471..faa11d9 100644 --- a/test/cuda_bayesian_blocks_tests.jl +++ b/test/cuda_bayesian_blocks_tests.jl @@ -196,6 +196,12 @@ if CUDA.functional() end end + @testset "GPU memory accounting" begin + @test cuda_ext._gpu_memory_budget_bytes(1000) == 650 + @test cuda_ext._reusable_gpu_memory_bytes(1000, 400, 250) == 1150 + @test_throws ArgumentError cuda_ext._reusable_gpu_memory_bytes(100, 50, 51) + end + @testset "Lightweight U_g bucketing and batching" begin buckets = cuda_ext._bb_quantile_buckets(problems) @test sort(vcat(buckets...)) == collect(eachindex(problems)) @@ -212,7 +218,11 @@ if CUDA.functional() sizeof(Float64) * (u + 1) + sizeof(UInt8) * u + sizeof(Float64) * u + - sizeof(UInt8) * u + sizeof(UInt8) * u + + sizeof(Int64) + + sizeof(Int64) + + sizeof(Int32) + + sizeof(Float64) @test cuda_ext._bb_problem_bytes(problems[1], UInt8, UInt8) == expected_bytes diff --git a/test/cuda_largemem_tests.jl b/test/cuda_largemem_tests.jl index 63bb8d2..3028d3d 100644 --- a/test/cuda_largemem_tests.jl +++ b/test/cuda_largemem_tests.jl @@ -23,15 +23,18 @@ const _LARGEMEM_ENABLED = get(ENV, "FASTPIDC_LARGEMEM_TESTS", "0") == "1" const _LARGEMEM_NODES = 64 const _LARGEMEM_SAMPLES = 2000 const _LARGEMEM_BINS = 725 -# The host derives chunk_size from free memory; it only reaches 64 (rather than -# 63, which stays inside Int32) with roughly 11 GiB free, so require headroom. -const _LARGEMEM_FREE_FLOOR = 12 * 2^30 +# The host plans the entire PUC device footprint against 65% of currently +# reusable memory. Require enough headroom for the >8 GiB counts buffer plus +# fixed arrays. +const _LARGEMEM_FREE_FLOOR = 14 * 2^30 _counts_elements(k_bins, n, chunk) = big(k_bins)^2 * n * chunk if _LARGEMEM_ENABLED && CUDA.functional() @testset "PUC flat indexing past 2^31" begin - free_bytes = Int(CUDA.free_memory()) + cuda_ext = Base.get_extension(FastPIDC, :FastPIDCCUDAExt) + @test cuda_ext !== nothing + free_bytes = cuda_ext._available_gpu_memory_bytes() if free_bytes < _LARGEMEM_FREE_FLOOR @warn "Skipping large-memory PUC index test" required_gib = @@ -55,16 +58,13 @@ if _LARGEMEM_ENABLED && CUDA.functional() # The host picks the chunk itself, so fail loudly rather than pass # vacuously if this configuration drifts below the threshold the # test exists to cross. - cuda_ext = Base.get_extension(FastPIDC, :FastPIDCCUDAExt) - @test cuda_ext !== nothing - bytes_per_chunk_col = - k_bins^2 * length(nodes) * sizeof(Int32) + - k_bins * length(nodes) * sizeof(Float64) - chunk_size = clamp( - floor(Int, free_bytes * 0.8 / bytes_per_chunk_col), - 1, - min(256, length(nodes)), + memory_plan = cuda_ext._puc_memory_plan( + length(nodes), + length(nodes[1].binned_values), + k_bins, + free_bytes, ) + chunk_size = memory_plan.chunk_size @test _counts_elements(k_bins, length(nodes), chunk_size) - 1 > big(typemax(Int32)) @@ -89,17 +89,32 @@ if _LARGEMEM_ENABLED && CUDA.functional() end end -# Cheap, always-run companion: the one kernel limit 64-bit indexing does not -# lift. joint_counts_kernel forms the bin-pair index u * k_bins + v in Int32. +# Cheap, always-run companion: all composed flat offsets are 64-bit now. The +# remaining ABI limit is that launch scalars themselves must fit signed Int32. if CUDA.functional() - @testset "Kernel bin-pair index limit" begin + @testset "Kernel scalar index limits" begin cuda_ext = Base.get_extension(FastPIDC, :FastPIDCCUDAExt) @test cuda_ext !== nothing - @test cuda_ext._MAX_K_BINS^2 <= typemax(Int32) - @test big(cuda_ext._MAX_K_BINS + 1)^2 > big(typemax(Int32)) - @test cuda_ext._check_kernel_index_limits(cuda_ext._MAX_K_BINS) === nothing - @test_throws ErrorException cuda_ext._check_kernel_index_limits( - cuda_ext._MAX_K_BINS + 1, + + # k_bins^2 may exceed Int32 because the bin-pair term is widened before + # multiplication in joint_counts_kernel. + @test cuda_ext._check_kernel_scalar_limits(12_071, 38_176, 50_000) === nothing + @test big(50_000)^2 > big(typemax(Int32)) + + @test_throws ArgumentError cuda_ext._check_kernel_scalar_limits( + Int(typemax(Int32)) + 1, + 1, + 1, + ) + @test_throws ArgumentError cuda_ext._check_kernel_scalar_limits( + 1, + Int(typemax(Int32)) + 1, + 1, + ) + @test_throws ArgumentError cuda_ext._check_kernel_scalar_limits( + 1, + 1, + Int(typemax(Int32)) + 1, ) end end From 8c7bd3bf25e50b42e70b2b1ff3143f0e44319a3e Mon Sep 17 00:00:00 2001 From: jjschirle Date: Fri, 11 Sep 2026 13:50:18 -0700 Subject: [PATCH 5/6] Removed hardcoded references to using 65% of memory. Better maintainable if we change the allocation in the future --- ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl | 55 +++++++--- python/src/fastpidc/cuda.py | 48 +++++++-- python/tests/test_cuda.py | 136 +++++++++++++++++++++++++ test/cuda_bayesian_blocks_tests.jl | 78 ++++++++++++++ 4 files changed, 295 insertions(+), 22 deletions(-) diff --git a/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl b/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl index d7864e2..5415481 100644 --- a/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl +++ b/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl @@ -97,6 +97,16 @@ const _GPU_MEMORY_BUDGET_NUMERATOR = 65 const _GPU_MEMORY_BUDGET_DENOMINATOR = 100 const _MAX_CHUNK_SIZE = 256 +function _gpu_memory_budget_percent_label( + numerator_value::Integer = _GPU_MEMORY_BUDGET_NUMERATOR, + denominator_value::Integer = _GPU_MEMORY_BUDGET_DENOMINATOR, +) + denominator_value > 0 || throw(ArgumentError("denominator_value must be positive")) + percent = 100 * numerator_value // denominator_value + return denominator(percent) == 1 ? + "$(numerator(percent))%" : "$(round(Float64(percent), digits = 2))%" +end + function _gpu_memory_budget_bytes(free_bytes::Integer) free_bytes > 0 || throw(ArgumentError("free_bytes must be positive")) return Int( @@ -192,7 +202,7 @@ function _puc_memory_plan( minimum_bytes <= budget_bytes || error( "compute_puc_full_cuda: the fixed GPU buffers plus a one-gene chunk " * "would require $(round(Float64(minimum_bytes) / 2^30, digits = 2)) GiB, " * - "which exceeds the configured 65% memory budget " * + "which exceeds the configured $(_gpu_memory_budget_percent_label()) memory budget " * "($(round(Float64(budget_bytes) / 2^30, digits = 2)) GiB of " * "$(round(free_bytes / 2^30, digits = 2)) GiB currently reusable). " * "Reduce the number of genes/samples/bins, use a fixed small-bin " * @@ -230,10 +240,10 @@ 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. Before allocating device buffers it plans the complete footprint - fixed data, -marginals and output matrices plus chunked intermediates - against 65% of the -memory currently reusable by the process (driver-free plus unused CUDA.jl pool -blocks). The target (`z`) chunk is capped at 256 genes and shrunk -as needed, leaving 35% headroom for allocator fragmentation, CUDA/runtime +marginals and output matrices plus chunked intermediates - against the configured +fraction of memory currently reusable by the process (driver-free plus unused +CUDA.jl pool blocks). The target (`z`) chunk is capped at 256 genes and shrunk +as needed, leaving the remaining headroom for allocator fragmentation, CUDA/runtime workspaces and concurrent users. If the fixed buffers plus a one-gene chunk do not fit that budget, the function raises a descriptive error before allocating the large device arrays. @@ -271,8 +281,8 @@ function FastPIDC.compute_puc_full_cuda(nodes, config, base) _check_kernel_scalar_limits(num_nodes, num_samples, k_bins) # Plan the *entire* device footprint before allocating anything. The budget - # is 65% of currently reusable VRAM (driver-free plus unused CUDA.jl pool - # blocks), leaving 35% for allocator fragmentation, + # is the configured fraction of currently reusable VRAM (driver-free plus unused + # CUDA.jl pool blocks), leaving the remainder for allocator fragmentation, # runtime/library workspaces, and other users/processes on a shared GPU. free_bytes = _available_gpu_memory_bytes() memory_plan = _puc_memory_plan(num_nodes, num_samples, k_bins, free_bytes) @@ -330,7 +340,7 @@ function FastPIDC.compute_puc_full_cuda(nodes, config, base) ) println( "[FastPIDC] GPU memory: $(round(free_bytes / 2^30, digits = 2)) GiB reusable; " * - "65% budget=$(round(memory_plan.budget_bytes / 2^30, digits = 2)) GiB; " * + "$(_gpu_memory_budget_percent_label()) budget=$(round(memory_plan.budget_bytes / 2^30, digits = 2)) GiB; " * "fixed=$(round(memory_plan.fixed_bytes / 2^30, digits = 2)) GiB", ) println( @@ -556,6 +566,24 @@ function _bb_memory_batches( return batches end +function _bb_memory_plan( + bucket::Vector{Int}, + problems::Vector{FastPIDC.BayesianBlocksProblem}, + free_bytes::Integer, + ::Type{CountT}, + ::Type{IndexT}, +) where {CountT<:Integer,IndexT<:Integer} + budget_bytes = _gpu_memory_budget_bytes(free_bytes) + batches = _bb_memory_batches( + bucket, + problems, + budget_bytes, + CountT, + IndexT, + ) + return (batches = batches, budget_bytes = budget_bytes) +end + function _flatten_bb_batch( problems::Vector{FastPIDC.BayesianBlocksProblem}, problem_indices::Vector{Int}, @@ -769,7 +797,7 @@ function FastPIDC.solve_bayesian_blocks_cuda( # The priors remain live across all batches. Allocate them first, then size # each bucket from the *remaining* reusable memory at runtime. Every batch is - # kept within 65% of what is free at that moment, and _bb_problem_bytes + # kept within the configured fraction of what is free at that moment, and _bb_problem_bytes # includes all per-gene device metadata, not just the large state arrays. priors_gpu = CuArray(_bb_prior_values(max_u)) @@ -789,14 +817,15 @@ function FastPIDC.solve_bayesian_blocks_cuda( threads = _bb_threads_for_max_u(bucket_max_u) free_bytes = _available_gpu_memory_bytes() - budget_bytes = _gpu_memory_budget_bytes(free_bytes) - batches = _bb_memory_batches( + memory_plan = _bb_memory_plan( bucket, problems, - budget_bytes, + free_bytes, CountT, IndexT, ) + budget_bytes = memory_plan.budget_bytes + batches = memory_plan.batches if verbose bucket_min_u = minimum(i -> length(problems[i].prefix_counts), bucket) @@ -805,7 +834,7 @@ function FastPIDC.solve_bayesian_blocks_cuda( "$(length(bucket)) genes, U_g=$bucket_min_u:$bucket_max_u, " * "threads=$threads, batches=$(length(batches)), " * "reusable=$(round(free_bytes / 2.0^30; digits = 2)) GiB, " * - "65% budget=$(round(budget_bytes / 2.0^30; digits = 2)) GiB", + "$(_gpu_memory_budget_percent_label()) budget=$(round(budget_bytes / 2.0^30; digits = 2)) GiB", ) end diff --git a/python/src/fastpidc/cuda.py b/python/src/fastpidc/cuda.py index 48ca91f..92acde7 100644 --- a/python/src/fastpidc/cuda.py +++ b/python/src/fastpidc/cuda.py @@ -78,6 +78,16 @@ def _load_module(): return cp.RawModule(code=source, options=("--std=c++11",)) +def _gpu_memory_budget_percent_label( + numerator: int = _GPU_MEMORY_BUDGET_NUMERATOR, + denominator: int = _GPU_MEMORY_BUDGET_DENOMINATOR, +) -> str: + if denominator <= 0: + raise ValueError("denominator must be positive") + percent = 100 * numerator / denominator + return f"{percent:g}%" + + def _gpu_memory_budget_bytes(free_bytes: int) -> int: if free_bytes <= 0: raise ValueError("free_bytes must be positive") @@ -149,7 +159,7 @@ def _puc_memory_plan( *, requested_chunk_size: int | None = None, ) -> tuple[int, int, int, int]: - """Plan the entire PUC device footprint against 65% of currently reusable VRAM. + """Plan the entire PUC device footprint against the configured reusable-VRAM fraction. Returns ``(chunk_size, fixed_bytes, bytes_per_chunk_column, budget_bytes)``. Python integers are unbounded, so the planning arithmetic itself cannot wrap. @@ -170,7 +180,7 @@ def _puc_memory_plan( raise RuntimeError( "compute_puc_full_cuda: the fixed GPU buffers plus a one-gene chunk " f"would require {(fixed_bytes + bytes_per_chunk_column) / 2**30:.2f} GiB, " - "which exceeds the configured 65% memory budget " + f"which exceeds the configured {_gpu_memory_budget_percent_label()} memory budget " f"({budget_bytes / 2**30:.2f} GiB of {free_bytes / 2**30:.2f} GiB currently reusable). " "Reduce the number of genes/samples/bins, use a fixed small-bin discretizer, " "or use config.backend = 'cpu'." @@ -187,7 +197,8 @@ def _puc_memory_plan( if requested_chunk_size > safe_max_chunk: raise RuntimeError( f"compute_puc_full_cuda: requested chunk_size={requested_chunk_size} would exceed " - f"the configured 65% GPU-memory budget; the largest safe chunk is {safe_max_chunk} " + f"the configured {_gpu_memory_budget_percent_label()} GPU-memory budget; " + f"the largest safe chunk is {safe_max_chunk} " f"with {free_bytes / 2**30:.2f} GiB currently reusable." ) chunk_size = requested_chunk_size @@ -201,7 +212,7 @@ def compute_puc_full_cuda( """GPU implementation of :func:`fastpidc.puc.compute_puc_full`. The complete device footprint is planned before large allocations against - 65% of VRAM currently reusable by CuPy (driver-free memory plus unused pool + the configured fraction of VRAM currently reusable by CuPy (driver-free memory plus unused pool blocks, bounded by any configured pool limit). ``chunk_size=None`` chooses the largest safe chunk up to 256; an explicit chunk must also fit the same budget. """ @@ -272,7 +283,7 @@ def compute_puc_full_cuda( print(f"[fastpidc] GPU chunked PUC: processing {n} x {n} pairs (k_bins={k_bins})...") print( f"[fastpidc] GPU memory: {free_bytes / 2**30:.2f} GiB reusable; " - f"65% budget={budget_bytes / 2**30:.2f} GiB; " + f"{_gpu_memory_budget_percent_label()} budget={budget_bytes / 2**30:.2f} GiB; " f"fixed={fixed_bytes / 2**30:.2f} GiB" ) print(f"[fastpidc] Using chunk size {chunk_size} ({n_chunks} iterations)") @@ -476,6 +487,19 @@ def _bb_memory_batches( return batches +def _bb_memory_plan( + bucket: list[int], + problems: list[BayesianBlocksProblem], + free_bytes: int, + count_dtype: np.dtype, + index_dtype: np.dtype, +) -> tuple[list[list[int]], int]: + """Plan Bayesian-block batches against the configured reusable-memory fraction.""" + budget_bytes = _gpu_memory_budget_bytes(free_bytes) + batches = _bb_memory_batches(bucket, problems, budget_bytes, count_dtype, index_dtype) + return batches, budget_bytes + + def _flatten_bb_batch( problems: list[BayesianBlocksProblem], problem_indices: list[int], count_dtype: np.dtype ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: @@ -608,7 +632,7 @@ def solve_bayesian_blocks_cuda( :func:`fastpidc.discretizers.solve_bayesian_blocks_cpu`, so the selected change points agree exactly with the CPU reference. - Each workload bucket is sized from 65% of the device memory reusable by + Each workload bucket is sized from the configured fraction of device memory reusable by CuPy at that point in the run. The per-problem byte estimate includes every packed device buffer and metadata array. """ @@ -659,15 +683,21 @@ def solve_bayesian_blocks_cuda( # for every bucket, including reusable free blocks in its memory # pool and respecting any configured pool limit. free_bytes = _available_gpu_memory_bytes(cp) - budget_bytes = _gpu_memory_budget_bytes(free_bytes) - batches = _bb_memory_batches(bucket, problems, budget_bytes, count_dtype, index_dtype) + batches, budget_bytes = _bb_memory_plan( + bucket, + problems, + free_bytes, + count_dtype, + index_dtype, + ) if verbose: bucket_min_u = min(problems[i].prefix_counts.size for i in bucket) print( f"[fastpidc] CUDA BB bucket {bucket_number}: {len(bucket)} genes, " f"U_g={bucket_min_u}:{bucket_max_u}, threads={threads}, batches={len(batches)}, " - f"reusable={free_bytes / 2**30:.2f} GiB, 65% budget={budget_bytes / 2**30:.2f} GiB" + f"reusable={free_bytes / 2**30:.2f} GiB, " + f"{_gpu_memory_budget_percent_label()} budget={budget_bytes / 2**30:.2f} GiB" ) for batch in batches: diff --git a/python/tests/test_cuda.py b/python/tests/test_cuda.py index ec189d8..d65f5ae 100644 --- a/python/tests/test_cuda.py +++ b/python/tests/test_cuda.py @@ -16,13 +16,16 @@ from fastpidc.cuda import ( _KERNEL_INT_MAX, _MAX_CHUNK_SIZE, + _available_gpu_memory_bytes, _bb_kernel_name, _bb_memory_batches, + _bb_memory_plan, _bb_problem_bytes, _bb_quantile_buckets, _bb_threads_for_max_u, _check_kernel_scalar_limits, _gpu_memory_budget_bytes, + _gpu_memory_budget_percent_label, _puc_memory_plan, _reusable_gpu_memory_bytes, _smallest_unsigned_dtype, @@ -60,8 +63,28 @@ def test_cuda_available_returns_a_bool(): assert isinstance(cuda_available(), bool) +def test_gpu_memory_budget_label_is_derived_from_the_configured_fraction(): + import fastpidc.cuda as cuda_module + + expected = ( + 100 + * cuda_module._GPU_MEMORY_BUDGET_NUMERATOR + / cuda_module._GPU_MEMORY_BUDGET_DENOMINATOR + ) + assert _gpu_memory_budget_percent_label() == f"{expected:g}%" + assert _gpu_memory_budget_percent_label(80, 100) == "80%" + + def test_gpu_memory_budget_is_65_percent_of_currently_free_memory(): assert _gpu_memory_budget_bytes(1000) == 650 + assert _gpu_memory_budget_bytes(1001) == 650 # integer floor, never rounds upward + + +def test_gpu_memory_budget_rejects_nonpositive_inputs(): + with pytest.raises(ValueError, match="positive"): + _gpu_memory_budget_bytes(0) + with pytest.raises(ValueError, match="positive"): + _gpu_memory_budget_bytes(-1) def test_reusable_gpu_memory_includes_cached_pool_blocks(): @@ -78,6 +101,40 @@ def test_reusable_gpu_memory_honors_a_cupy_pool_limit(): ) +def test_available_gpu_memory_uses_driver_free_pool_cache_and_pool_limit(): + class FakeRuntime: + @staticmethod + def memGetInfo(): + return (1000, 2000) + + class FakePool: + @staticmethod + def free_bytes(): + return 400 + + @staticmethod + def used_bytes(): + return 250 + + @staticmethod + def get_limit(): + return 900 + + class FakeCuda: + runtime = FakeRuntime() + + class FakeCupy: + cuda = FakeCuda() + + @staticmethod + def get_default_memory_pool(): + return FakePool() + + # Physical reusable memory is 1000 + 400 = 1400 bytes, but the CuPy + # pool has only 900 - 250 = 650 bytes of allocatable headroom. + assert _available_gpu_memory_bytes(FakeCupy()) == 650 + + def test_puc_memory_plan_counts_fixed_and_chunk_buffers_exactly(): n, m, k_bins = 3, 5, 7 _, fixed_bytes, bytes_per_chunk_column, _ = _puc_memory_plan( @@ -103,6 +160,40 @@ def test_chunk_size_shrinks_when_memory_is_tight(): assert 1 <= tight < roomy <= _MAX_CHUNK_SIZE +@pytest.mark.parametrize( + ("free_gib", "expected_chunk"), + [(12, 75), (16, 125), (24, 225), (32, 256)], +) +def test_production_sized_puc_plan_is_dynamic_and_never_exceeds_budget(free_gib, expected_chunk): + # Shape from the production run that exposed the old >2^31 flat-index bug. + # Exact expected chunks make this a regression test for both the 65% policy + # and the absence of a hidden/fixed ~8 GiB allocation ceiling. + n, m, k_bins = 12_071, 38_176, 33 + chunk, fixed_bytes, bytes_per_chunk_column, budget_bytes = _puc_memory_plan( + n=n, m=m, k_bins=k_bins, free_bytes=free_gib * GIB + ) + + assert chunk == expected_chunk + assert budget_bytes == free_gib * GIB * 65 // 100 + assert fixed_bytes + chunk * bytes_per_chunk_column <= budget_bytes + + # Unless the independent 256-gene cap is what stopped growth, one more + # target gene must be the first chunk size that would exceed the budget. + if chunk < min(_MAX_CHUNK_SIZE, n): + assert fixed_bytes + (chunk + 1) * bytes_per_chunk_column > budget_bytes + + +def test_puc_plan_can_safely_approve_more_than_eight_gib_of_device_buffers(): + chunk, fixed_bytes, bytes_per_chunk_column, budget_bytes = _puc_memory_plan( + n=12_071, m=38_176, k_bins=33, free_bytes=32 * GIB + ) + required_bytes = fixed_bytes + chunk * bytes_per_chunk_column + + assert chunk == _MAX_CHUNK_SIZE + assert required_bytes > 8 * GIB + assert required_bytes <= budget_bytes + + def test_memory_plan_raises_before_allocating_when_one_gene_chunk_does_not_fit(): with pytest.raises(RuntimeError, match="one-gene chunk"): _puc_memory_plan(n=20000, m=1000, k_bins=4000, free_bytes=8 * GIB) @@ -245,6 +336,39 @@ def test_bb_memory_batches_raises_when_one_problem_cannot_fit(): _bb_memory_batches([0], problems, 1024, np.dtype(np.uint16), np.dtype(np.uint16)) +def test_bb_memory_plan_applies_65_percent_before_packing_batches(): + # Make every problem the same size and choose free memory so exactly two + # problems fit in the 65% budget. This pins the composition of the budget + # policy and the batch packer, not just each helper in isolation. + problems = _bb_problems(sizes=(50,) * 6) + bucket = list(range(len(problems))) + count_dtype = index_dtype = np.dtype(np.uint16) + per_problem = _bb_problem_bytes(problems[0], count_dtype, index_dtype) + target_budget = 2 * per_problem + free_bytes = (target_budget * 100 + 64) // 65 + + batches, budget_bytes = _bb_memory_plan( + bucket, problems, free_bytes, count_dtype, index_dtype + ) + + assert budget_bytes == target_budget + assert batches == [[0, 1], [2, 3], [4, 5]] + for batch in batches: + batch_bytes = sum(_bb_problem_bytes(problems[i], count_dtype, index_dtype) for i in batch) + assert batch_bytes <= budget_bytes + + +def test_bb_memory_plan_rejects_a_problem_that_cannot_fit_inside_65_percent(): + problems = _bb_problems(sizes=(900,)) + count_dtype = index_dtype = np.dtype(np.uint16) + per_problem = _bb_problem_bytes(problems[0], count_dtype, index_dtype) + + # Giving the planner only `per_problem` bytes of reusable memory means its + # actual allocation budget is 65% of that, so this must fail pre-allocation. + with pytest.raises(RuntimeError, match="exceeds the CUDA batch budget"): + _bb_memory_plan([0], problems, per_problem, count_dtype, index_dtype) + + # --- Bayesian blocks: GPU behavior ----------------------------------------- @@ -387,6 +511,18 @@ def test_kernel_scalar_limits_allow_bin_pair_products_past_int32(): assert 50_000**2 > INT32_MAX +def test_shared_kernel_keeps_composed_flat_offsets_64_bit(): + # The true >8-GiB execution regression is opt-in, so normal Python CI also + # pins the critical widening expressions in the canonical shared source. + import fastpidc.cuda as cuda_module + + source = cuda_module._KERNEL_SOURCE_PATH.read_text() + assert "const long long bin_pair = (long long)u * k_bins + v;" in source + assert "for (long long i64 = tid;" in source + assert "mi_matrix[(long long)x * n + z_global]" in source + assert "puc_scores[(long long)x * n + z_global]" in source + + @pytest.mark.parametrize( ("n", "m", "k_bins"), [ diff --git a/test/cuda_bayesian_blocks_tests.jl b/test/cuda_bayesian_blocks_tests.jl index faa11d9..3caf81e 100644 --- a/test/cuda_bayesian_blocks_tests.jl +++ b/test/cuda_bayesian_blocks_tests.jl @@ -36,6 +36,84 @@ function _bb_test_values() return values end +const _CUDA_EXT = Base.get_extension(FastPIDC, :FastPIDCCUDAExt) + +@testset "CUDA memory planning (device-independent)" begin + @test _CUDA_EXT !== nothing + cuda_ext = _CUDA_EXT + + @testset "65% budget and reusable-memory arithmetic" begin + @test cuda_ext._gpu_memory_budget_percent_label() == + "$(100 * cuda_ext._GPU_MEMORY_BUDGET_NUMERATOR ÷ cuda_ext._GPU_MEMORY_BUDGET_DENOMINATOR)%" + @test cuda_ext._gpu_memory_budget_percent_label(80, 100) == "80%" + @test cuda_ext._gpu_memory_budget_bytes(1000) == 650 + @test cuda_ext._gpu_memory_budget_bytes(1001) == 650 + @test_throws ArgumentError cuda_ext._gpu_memory_budget_bytes(0) + @test_throws ArgumentError cuda_ext._gpu_memory_budget_bytes(-1) + + @test cuda_ext._reusable_gpu_memory_bytes(1000, 400, 250) == 1150 + @test_throws ArgumentError cuda_ext._reusable_gpu_memory_bytes(100, 50, 51) + end + + @testset "Bayesian-block 65% batch planner" begin + values = Float64.(1:50) + problem = FastPIDC.prepare_bayesian_blocks(values) + problems = [problem for _ = 1:6] + bucket = collect(eachindex(problems)) + per_problem = cuda_ext._bb_problem_bytes(problem, UInt16, UInt16) + u = length(problem.prefix_counts) + @test per_problem == + sizeof(Float64) * (u + 1) + + sizeof(UInt16) * u + + sizeof(Float64) * u + + sizeof(UInt16) * u + + sizeof(Int64) + + sizeof(Int64) + + sizeof(Int32) + + sizeof(Float64) + + # Prefix counts and back-pointers must widen before their values overflow. + @test cuda_ext._smallest_unsigned_type(255) == UInt8 + @test cuda_ext._smallest_unsigned_type(256) == UInt16 + @test cuda_ext._smallest_unsigned_type(65_535) == UInt16 + @test cuda_ext._smallest_unsigned_type(65_536) == UInt32 + @test cuda_ext._smallest_unsigned_type(big(typemax(UInt32)) + 1) == UInt64 + + # Pick reusable memory whose 65% budget is exactly two problems. The + # resulting batches must therefore contain exactly two genes each. + target_budget = 2 * per_problem + free_bytes = cld(target_budget * 100, 65) + plan = cuda_ext._bb_memory_plan( + bucket, + problems, + free_bytes, + UInt16, + UInt16, + ) + + @test plan.budget_bytes == target_budget + @test plan.batches == [[1, 2], [3, 4], [5, 6]] + for batch in plan.batches + batch_bytes = sum( + i -> cuda_ext._bb_problem_bytes(problems[i], UInt16, UInt16), + batch, + ) + @test batch_bytes <= plan.budget_bytes + end + + # One problem requiring more than the 65% budget is rejected before any + # device allocation is attempted. + @test_throws ArgumentError cuda_ext._bb_memory_plan( + [1], + [problem], + per_problem, + UInt16, + UInt16, + ) + end + +end + if CUDA.functional() @testset "CUDA Bayesian blocks equivalence and determinism" begin cuda_ext = Base.get_extension(FastPIDC, :FastPIDCCUDAExt) From c53062be1549410d09c0790c1d9140fdb0551282 Mon Sep 17 00:00:00 2001 From: jjschirle Date: Fri, 11 Sep 2026 14:16:47 -0700 Subject: [PATCH 6/6] docs update --- docs/src/api.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/src/api.md b/docs/src/api.md index a20aa0c..f44631c 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -146,6 +146,7 @@ FastPIDCCUDAExt FastPIDCCUDAExt._kernel_source_path FastPIDCCUDAExt._compile_ptx FastPIDCCUDAExt._get_module +FastPIDCCUDAExt._check_kernel_scalar_limits FastPIDCCUDAExt._bb_kernel_name FastPIDC.bayesian_blocks_cuda_available FastPIDC.solve_bayesian_blocks_cuda