diff --git a/CLI_fastpidc.jl b/CLI_fastpidc.jl index e854641..1734f2f 100644 --- a/CLI_fastpidc.jl +++ b/CLI_fastpidc.jl @@ -80,7 +80,8 @@ Basic options: --base INT Log base for MI (2, e, 10). Default: 2 Execution / Environment: - --backend STR 'cuda' (default) or 'cpu'. + --backend STR PUC backend: 'cuda' (default) or 'cpu'. + --bb-backend STR Bayesian-block backend: 'cuda' (default) or 'cpu'. --output-format STR 'tsv' (default) or NumPy binary 'npy' Note: To run on multiple CPU threads, use the Julia flag: `julia -t auto command_line_fastpidc.jl ...` @@ -96,7 +97,7 @@ Other: Example: julia --project=. command_line_fastpidc.jl \\ --infile X.txt --outfile edges.tsv \\ - --backend cuda + --backend cuda --bb-backend cuda """ function main() @@ -132,16 +133,22 @@ function main() # ----------------- Execution Environment ---------------- n_threads_act = Threads.nthreads() backend = Symbol(lowercase(get(args, "backend", "cuda"))) + bb_backend = Symbol(lowercase(get(args, "bb-backend", "cuda"))) + + backend in (:cpu, :cuda) || + error("Unsupported --backend=$backend. Use 'cpu' or 'cuda'.") + bb_backend in (:cpu, :cuda) || + error("Unsupported --bb-backend=$bb_backend. Use 'cpu' or 'cuda'.") # --- FAST FAIL GPU CHECK --- - if backend == :cuda + if backend == :cuda || bb_backend == :cuda @say "Checking for CUDA availability..." try Core.eval(Main, :(import CUDA)) is_functional = Core.eval(Main, :(CUDA.functional())) if !is_functional error( - "CUDA.jl is installed, but no functional GPU was detected. Try running with --backend cpu", + "CUDA.jl is installed, but no functional GPU was detected. Try running with --backend cpu --bb-backend cpu", ) end @say "CUDA GPU detected successfully." @@ -150,7 +157,7 @@ function main() rethrow(e) else error( - "Failed to load CUDA.jl. Please ensure CUDA is installed in your Julia environment, or run with --backend cpu.", + "Failed to load CUDA.jl. Please ensure CUDA is installed in your Julia environment, or run with --backend cpu --bb-backend cpu.", ) end end @@ -163,6 +170,7 @@ function main() # ----------------- Build PIDCConfig ---------------- cfg = PIDCConfig( backend = backend, + bb_backend = bb_backend, discretizer = discretizer, estimator = estimator, dump_mi_path = dump_mi_path, @@ -180,6 +188,7 @@ function main() println(" n_bins = $n_bins") println(" base = $base") println(" backend = $(cfg.backend)") + println(" bb_backend = $(cfg.bb_backend)") println(" JULIA_NUM_THREADS= $n_threads_act") println( " dump_mi_path = $(cfg.dump_mi_path === nothing ? "none" : cfg.dump_mi_path)", diff --git a/docs/src/api.md b/docs/src/api.md index f601dec..f1fcaba 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -98,6 +98,16 @@ DiscretizeBayesianBlocks binedges ``` +### Bayesian Blocks internals + +```@docs +FastPIDC.BayesianBlocksProblem +FastPIDC.BayesianBlocksSolution +FastPIDC.prepare_bayesian_blocks +FastPIDC.solve_bayesian_blocks_cpu +FastPIDC._build_nodes +``` + ### Information measures ```@docs @@ -136,4 +146,5 @@ FastPIDCCUDAExt FastPIDCCUDAExt.joint_counts_kernel_chunked! FastPIDCCUDAExt.mi_si_kernel_chunked! FastPIDCCUDAExt.puc_accumulation_kernel_chunked! +FastPIDCCUDAExt.bayesian_blocks_dp_kernel! ``` diff --git a/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl b/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl index 7fb6e1a..32a7485 100644 --- a/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl +++ b/ext/FastPIDCCUDAExt/FastPIDCCUDAExt.jl @@ -33,13 +33,16 @@ function joint_counts_kernel_chunked!(data, counts, n, m, k_bins, z_start, z_chu z_global = z_start + z_local - 1 if z_global > n || x == z_global; return nothing; end - # Each thread computes joint counts for pair (x, z_global) + # Each thread computes joint counts for pair (x, z_global). Derive the + # increment from the count buffer so UInt16/UInt32 storage remains generic. + count_increment = one(eltype(counts)) for s in 1:m - u = data[s, x] - v = data[s, z_global] + # Bin IDs may use a compact unsigned storage type on the GPU; widen + # them for bounds checks and native N-D indexing. + u = Int32(data[s, x]) + v = Int32(data[s, z_global]) if u >= 1 && u <= k_bins && v >= 1 && v <= k_bins - # Using native N-D indexing - counts[u, v, x, z_local] += Int32(1) + counts[u, v, x, z_local] += count_increment end end @@ -151,60 +154,58 @@ end # --- Host Implementation --- -""" - FastPIDC.compute_puc_full_cuda(nodes, config, base) -> (mi_scores, puc_scores) +function _smallest_unsigned_type(max_value::Integer) + max_value >= 0 || throw(ArgumentError("max_value must be nonnegative")) -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 -[`joint_counts_kernel_chunked!`](@ref), [`mi_si_kernel_chunked!`](@ref) and -[`puc_accumulation_kernel_chunked!`](@ref) in sequence, symmetrizing the -resulting PUC matrix before returning both matrices to the CPU. -`config.verbose` enables progress printouts; `base` is currently unused -(mutual information is always computed in base 2 on the GPU path). Raises -an `ErrorException` with a suggested remedy if even a single-gene chunk -would not fit in the currently-free GPU memory. -""" -function FastPIDC.compute_puc_full_cuda(nodes, config, base) + if max_value <= 255 + return UInt8 + elseif max_value <= 65_535 + return UInt16 + elseif max_value <= 4_294_967_295 + return UInt32 + else + return UInt64 + end +end + +function _compute_puc_full_cuda_typed( + nodes, + config, + base, + ::Type{BinT}, + ::Type{CountT}, +) where {BinT<:Integer,CountT<:Integer} num_nodes = length(nodes) num_samples = length(nodes[1].binned_values) k_bins = maximum(n -> n.number_of_bins, nodes) - # Prepare static data on CPU and move to GPU - data_cpu = zeros(Int32, num_samples, num_nodes) + # Prepare static data on CPU and move to GPU. + data_cpu = Matrix{BinT}(undef, num_samples, num_nodes) marginals_cpu = zeros(Float64, k_bins, num_nodes) for i = 1:num_nodes - data_cpu[:, i] .= Int32.(nodes[i].binned_values) + data_cpu[:, i] .= nodes[i].binned_values p = nodes[i].probabilities marginals_cpu[1:length(p), i] .= Float64.(p) end data_gpu = CuArray(data_cpu) marginals_gpu = CuArray(marginals_cpu) - - # Global output matrices + + # Global output matrices. puc_scores_gpu = CUDA.zeros(Float64, num_nodes, num_nodes) mi_matrix_gpu = CUDA.zeros(Float64, num_nodes, num_nodes) - # Chunk configuration: `counts_chunk_gpu`/`si_chunk_gpu` scale as - # k_bins^2 * num_nodes * chunk_size / k_bins * num_nodes * chunk_size - # respectively, so an adaptive discretizer that picks a large k_bins - # (e.g. "bayesian_blocks", the package default, on a dataset with many - # samples -- it has no upper bound on the number of bins it selects) - # can make even a modest fixed chunk_size request far larger than the - # GPU's total memory. Size the chunk to fit within what's actually free - # right now instead of always requesting a fixed chunk_size=256, and - # fail with an actionable message rather than a raw CUDA OOM error if - # even a single-gene chunk doesn't fit. + # Chunk configuration: `counts_chunk_gpu`/`si_chunk_gpu` scale with + # k_bins^2 * num_nodes * chunk_size and k_bins * num_nodes * chunk_size, + # respectively. Size the chunk to fit the GPU memory currently free rather + # than always requesting a fixed chunk size of 256. Because BB_prefix uses + # the smallest exact joint-count type, account for CountT rather than + # assuming Int32 storage. bytes_per_chunk_col = - k_bins^2 * num_nodes * sizeof(Int32) + # counts_chunk_gpu + k_bins^2 * num_nodes * sizeof(CountT) + # counts_chunk_gpu k_bins * num_nodes * sizeof(Float64) # si_chunk_gpu free_bytes, _ = CUDA.memory_info() - safety_factor = 0.8 # headroom for the fixed buffers above + allocator overhead/fragmentation + 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)) @@ -221,45 +222,72 @@ function FastPIDC.compute_puc_full_cuda(nodes, config, base) ) end - # Chunked intermediate buffers (Pre-allocated once!) - counts_chunk_gpu = CUDA.zeros(Int32, k_bins, k_bins, num_nodes, chunk_size) + # Chunked intermediate buffers, pre-allocated once. + counts_chunk_gpu = CUDA.zeros(CountT, k_bins, k_bins, num_nodes, chunk_size) si_chunk_gpu = CUDA.zeros(Float64, k_bins, num_nodes, chunk_size) - + if config.verbose - println("[FastPIDC] GPU Chunked PUC: Processing $num_nodes x $num_nodes pairs (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") + 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", + ) + println( + "[FastPIDC] GPU storage types: bin IDs=$(BinT) (max bins=$k_bins), " * + "joint counts=$(CountT) (cells=$num_samples)", + ) end - + # Iterate over the Z-axis in chunks for z_start in 1:chunk_size:num_nodes z_end = min(z_start + chunk_size - 1, num_nodes) z_curr_chunk_size = z_end - z_start + 1 - - # Wipe the intermediate buffers clean before the next chunk! - CUDA.fill!(counts_chunk_gpu, Int32(0)) - CUDA.fill!(si_chunk_gpu, Float64(0)) - + + # Wipe the intermediate buffers clean before the next chunk. + CUDA.fill!(counts_chunk_gpu, zero(CountT)) + CUDA.fill!(si_chunk_gpu, 0.0) + threads = (16, 16) blocks = (cld(num_nodes, 16), cld(z_curr_chunk_size, 16)) - + @cuda threads=threads blocks=blocks joint_counts_kernel_chunked!( - data_gpu, counts_chunk_gpu, - Int32(num_nodes), Int32(num_samples), Int32(k_bins), - Int32(z_start), Int32(z_curr_chunk_size) + data_gpu, + counts_chunk_gpu, + Int32(num_nodes), + Int32(num_samples), + Int32(k_bins), + Int32(z_start), + Int32(z_curr_chunk_size), ) - + @cuda threads=threads blocks=blocks mi_si_kernel_chunked!( - counts_chunk_gpu, marginals_gpu, mi_matrix_gpu, si_chunk_gpu, - Int32(num_nodes), Int32(num_samples), Int32(k_bins), - Int32(z_start), Int32(z_curr_chunk_size) + counts_chunk_gpu, + marginals_gpu, + mi_matrix_gpu, + si_chunk_gpu, + Int32(num_nodes), + Int32(num_samples), + Int32(k_bins), + Int32(z_start), + Int32(z_curr_chunk_size), ) - + @cuda threads=threads blocks=blocks puc_accumulation_kernel_chunked!( - si_chunk_gpu, mi_matrix_gpu, puc_scores_gpu, marginals_gpu, - Int32(num_nodes), Int32(k_bins), - Int32(z_start), Int32(z_curr_chunk_size) + si_chunk_gpu, + mi_matrix_gpu, + puc_scores_gpu, + marginals_gpu, + Int32(num_nodes), + Int32(k_bins), + Int32(z_start), + Int32(z_curr_chunk_size), ) end + # Copy results back puc_scores_cpu = Array(puc_scores_gpu) mi_matrix_cpu = Array(mi_matrix_gpu) @@ -276,4 +304,531 @@ function FastPIDC.compute_puc_full_cuda(nodes, config, base) return mi_matrix_cpu, puc_scores_cpu 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 +[`joint_counts_kernel_chunked!`](@ref), [`mi_si_kernel_chunked!`](@ref) and +[`puc_accumulation_kernel_chunked!`](@ref) in sequence, symmetrizing the +resulting PUC matrix before returning both matrices to the CPU. +`config.verbose` enables progress printouts; `base` is currently unused +(mutual information is always computed in base 2 on the GPU path). Raises +an `ErrorException` with a suggested remedy if even a single-gene chunk +would not fit in the currently-free GPU memory. +""" +function FastPIDC.compute_puc_full_cuda(nodes, config, base) + num_samples = length(nodes[1].binned_values) + k_bins = maximum(n -> n.number_of_bins, nodes) + + # Bin IDs only need to represent the largest per-gene bin index. + BinT = _smallest_unsigned_type(k_bins) + + # A joint count can reach the total number of cells. Choose the smallest + # exact unsigned type that can hold num_samples, guarding against overflow + # as datasets with larger cell counts are processed. + CountT = _smallest_unsigned_type(num_samples) + + return _compute_puc_full_cuda_typed(nodes, config, base, BinT, CountT) +end + + +# --- Bayesian-block CUDA backend ------------------------------------------- + +FastPIDC.bayesian_blocks_cuda_available() = CUDA.functional() + +# Deterministic comparison used by the Bayesian Blocks reduction: higher +# score wins, with exact ties resolved in favor of the smaller candidate index. +@inline function _bb_take_other( + other_score::Float64, + other_i::Int32, + current_score::Float64, + current_i::Int32, +)::Bool + return other_score > current_score || + (other_score == current_score && other_i < current_i) +end + +""" + bayesian_blocks_dp_kernel!(...) + +Assign one CUDA block to each gene. Endpoints `K` remain sequential because +`best[K]` depends on earlier endpoints, while threads within the block evaluate +candidate starts `i <= K` in parallel. The reduction uses a deterministic +first-maximum rule: higher score wins, and exact ties choose the smaller `i`. +""" +function bayesian_blocks_dp_kernel!( + prefix_counts::CuDeviceArray{CountT,1}, + block_lengths::CuDeviceArray{Float64,1}, + state_offsets::CuDeviceArray{Int64,1}, + block_offsets::CuDeviceArray{Int64,1}, + unique_counts::CuDeviceArray{Int32,1}, + best::CuDeviceArray{Float64,1}, + last::CuDeviceArray{IndexT,1}, + final_scores::CuDeviceArray{Float64,1}, + priors::CuDeviceArray{Float64,1}, +) where {CountT<:Unsigned,IndexT<:Unsigned} + gene = Int32(blockIdx().x) + tid = Int32(threadIdx().x) + nthreads = Int32(blockDim().x) + lane = ((tid - 1) & Int32(31)) + 1 + warp = ((tid - 1) >>> 5) + 1 + nwarps = nthreads >>> 5 + + # Keep one candidate per thread and one reduced candidate per warp. The + # reduction uses shared memory rather than generic shuffle helpers, keeping + # device dispatch fully concrete while requiring only two block barriers. + thread_scores = CuStaticSharedArray(Float64, 256) + thread_indices = CuStaticSharedArray(Int32, 256) + warp_scores = CuStaticSharedArray(Float64, 8) + warp_indices = CuStaticSharedArray(Int32, 8) + + @inbounds begin + state_start = state_offsets[gene] + block_start = block_offsets[gene] + n_unique = unique_counts[gene] + + # Match the CPU reference's singleton behavior explicitly. The generic + # event-fitness expression has zero block width when U_g == 1, whereas a + # constant gene should deterministically return its two outer edges and + # an objective score of zero. + if n_unique == 1 + if tid == 1 + best[state_start] = 0.0 + last[state_start] = one(IndexT) + final_scores[gene] = 0.0 + end + return nothing + end + + K = Int32(1) + while K <= n_unique + block_length_K1 = block_lengths[block_start + Int64(K)] + prefix_K = Float64(prefix_counts[state_start + Int64(K) - 1]) + prior = priors[K] + + local_best = -Inf + local_i = typemax(Int32) + i = tid + while i <= K + prefix_before = i == 1 ? 0.0 : + Float64(prefix_counts[state_start + Int64(i) - 2]) + count = prefix_K - prefix_before + width = + block_lengths[block_start + Int64(i) - 1] - block_length_K1 + + fit = count * log(count / width) - prior + if i > 1 + fit += best[state_start + Int64(i) - 2] + end + + if _bb_take_other(fit, i, local_best, local_i) + local_best = fit + local_i = i + end + i += nthreads + end + + thread_scores[tid] = local_best + thread_indices[tid] = local_i + sync_threads() + + # Each warp leader scans its 32 thread-local candidates in a fixed + # order. Exact ties still choose the smaller i, matching the CPU + # strict-`>` scan independently of launch size. + if lane == 1 + warp_start = (warp - 1) * Int32(32) + 1 + warp_end = warp_start + Int32(31) + warp_best = thread_scores[warp_start] + warp_i = thread_indices[warp_start] + slot = warp_start + 1 + while slot <= warp_end + other_score = thread_scores[slot] + other_i = thread_indices[slot] + if _bb_take_other(other_score, other_i, warp_best, warp_i) + warp_best = other_score + warp_i = other_i + end + slot += 1 + end + warp_scores[warp] = warp_best + warp_indices[warp] = warp_i + end + sync_threads() + + if tid == 1 + block_best = warp_scores[1] + block_i = warp_indices[1] + warp_slot = Int32(2) + while warp_slot <= nwarps + other_score = warp_scores[warp_slot] + other_i = warp_indices[warp_slot] + if _bb_take_other(other_score, other_i, block_best, block_i) + block_best = other_score + block_i = other_i + end + warp_slot += 1 + end + + state_index = state_start + Int64(K) - 1 + best[state_index] = block_best + # IndexT was selected from max(U_g), so this modular conversion + # is exact and avoids a checked integer constructor in device code. + last[state_index] = block_i % IndexT + if K == n_unique + final_scores[gene] = block_best + end + end + + # `best[K]` is stored in global memory and is required by every + # thread during the next endpoint. Block synchronization makes that + # write visible before advancing K. + sync_threads() + K += 1 + end + end + + return nothing +end + +function _bb_threads_for_max_u(max_u::Integer) + if max_u <= 32 + return 32 + elseif max_u <= 512 + return 64 + elseif max_u <= 4_096 + return 128 + else + return 256 + end +end + +function _bb_quantile_buckets(problems::Vector{FastPIDC.BayesianBlocksProblem}) + n = length(problems) + n == 0 && return Vector{Vector{Int}}() + + # U_g is already available from required preprocessing. Sorting only these + # gene indices is a lightweight O(G log G) operation and avoids a second + # scan of the expression matrix merely to choose GPU workload buckets. + order = sortperm(eachindex(problems); by = i -> length(problems[i].prefix_counts)) + n_buckets = min(4, n) + buckets = Vector{Vector{Int}}() + for bucket = 1:n_buckets + lo = fld((bucket - 1) * n, n_buckets) + 1 + hi = fld(bucket * n, n_buckets) + lo <= hi && push!(buckets, collect(order[lo:hi])) + end + return buckets +end + +function _bb_prior_values(max_u::Integer) + # The prior depends only on endpoint K, not on the gene. Compute it once on + # the CPU with the reference expression, avoiding one pow/log pair per gene + # per endpoint and eliminating that source of CPU/CUDA numeric variation. + return [4 - log(73.53 * 0.05 * ((K)^-0.478)) for K = 1:max_u] +end + +function _bb_problem_bytes( + problem::FastPIDC.BayesianBlocksProblem, + ::Type{CountT}, + ::Type{IndexT}, +) where {CountT<:Integer,IndexT<:Integer} + u = length(problem.prefix_counts) + return ( + sizeof(Float64) * (u + 1) + # block lengths + sizeof(CountT) * u + # prefix counts + sizeof(Float64) * u + # best scores + sizeof(IndexT) * u # back-pointers + ) +end + +function _bb_memory_batches( + bucket::Vector{Int}, + problems::Vector{FastPIDC.BayesianBlocksProblem}, + budget_bytes::Integer, + ::Type{CountT}, + ::Type{IndexT}, +) where {CountT<:Integer,IndexT<:Integer} + batches = Vector{Vector{Int}}() + current = Int[] + current_bytes = 0 + + for problem_index in bucket + problem_bytes = _bb_problem_bytes(problems[problem_index], CountT, IndexT) + problem_bytes <= budget_bytes || throw( + ArgumentError( + "One Bayesian-block problem requires $(problem_bytes) bytes, " * + "which exceeds the CUDA batch budget of $(budget_bytes) bytes. " * + "Reduce the number of unique input values for that gene.", + ), + ) + + if !isempty(current) && current_bytes + problem_bytes > budget_bytes + push!(batches, current) + current = Int[] + current_bytes = 0 + end + push!(current, problem_index) + current_bytes += problem_bytes + end + + !isempty(current) && push!(batches, current) + return batches +end + +function _flatten_bb_batch( + problems::Vector{FastPIDC.BayesianBlocksProblem}, + problem_indices::Vector{Int}, + ::Type{CountT}, +) where {CountT<:Integer} + n_genes = length(problem_indices) + total_states = sum(i -> length(problems[i].prefix_counts), problem_indices) + total_blocks = total_states + n_genes + + prefix_counts = Vector{CountT}(undef, total_states) + block_lengths = Vector{Float64}(undef, total_blocks) + state_offsets = Vector{Int64}(undef, n_genes) + block_offsets = Vector{Int64}(undef, n_genes) + unique_counts = Vector{Int32}(undef, n_genes) + + state_cursor = 1 + block_cursor = 1 + for (local_gene, problem_index) in enumerate(problem_indices) + problem = problems[problem_index] + u = length(problem.prefix_counts) + u <= typemax(Int32) || throw( + ArgumentError("Bayesian blocks CUDA backend supports at most $(typemax(Int32)) unique values per gene"), + ) + + state_offsets[local_gene] = state_cursor + block_offsets[local_gene] = block_cursor + unique_counts[local_gene] = Int32(u) + + @inbounds for j = 1:u + prefix_counts[state_cursor+j-1] = CountT(problem.prefix_counts[j]) + end + edge_end = problem.edges[end] + @inbounds for j = 1:(u+1) + block_lengths[block_cursor+j-1] = edge_end - problem.edges[j] + end + + state_cursor += u + block_cursor += u + 1 + end + + return prefix_counts, block_lengths, state_offsets, block_offsets, unique_counts +end + +function _change_points_from_last(last_values, offset::Int, n::Int) + n >= 1 || throw(ArgumentError("Bayesian-block backtracking requires n >= 1")) + + # A valid partition may place every unique value in its own block. In that + # case the returned edge-index path contains U_g + 1 entries, so allocating + # only U_g slots can underflow to Julia index zero during backtracking. + change_points = Vector{Int64}(undef, n + 1) + i_cp = n + 2 + ind = n + 1 + while true + i_cp -= 1 + change_points[i_cp] = ind + ind == 1 && break + + state = ind - 1 + 1 <= state <= n || throw( + ArgumentError("invalid Bayesian-block state $state while backtracking"), + ) + next_ind = Int(last_values[offset + state - 1]) + 1 <= next_ind <= state || throw( + ArgumentError( + "invalid Bayesian-block back-pointer $next_ind for state $state", + ), + ) + ind = next_ind + end + return change_points[i_cp:end] +end + +function _solve_bb_cuda_batch_with_priors( + problems::Vector{FastPIDC.BayesianBlocksProblem}, + problem_indices::Vector{Int}, + threads::Int, + ::Type{CountT}, + ::Type{IndexT}, + priors_gpu, +) where {CountT<:Integer,IndexT<:Integer} + threads in (32, 64, 128, 256) || throw( + ArgumentError( + "CUDA Bayesian blocks requires a power-of-two thread count " * + "from 32, 64, 128, or 256; got $threads", + ), + ) + + prefix_counts, block_lengths, state_offsets, block_offsets, unique_counts = + _flatten_bb_batch(problems, problem_indices, CountT) + + prefix_gpu = CuArray(prefix_counts) + block_gpu = CuArray(block_lengths) + state_offsets_gpu = CuArray(state_offsets) + block_offsets_gpu = CuArray(block_offsets) + unique_counts_gpu = CuArray(unique_counts) + best_gpu = CUDA.zeros(Float64, length(prefix_counts)) + last_gpu = CUDA.zeros(IndexT, length(prefix_counts)) + final_scores_gpu = CUDA.zeros(Float64, length(problem_indices)) + + try + @cuda threads=threads blocks=length(problem_indices) bayesian_blocks_dp_kernel!( + prefix_gpu, + block_gpu, + state_offsets_gpu, + block_offsets_gpu, + unique_counts_gpu, + best_gpu, + last_gpu, + final_scores_gpu, + priors_gpu, + ) + + last_values = Array(last_gpu) + final_scores = Array(final_scores_gpu) + + solutions = + Vector{FastPIDC.BayesianBlocksSolution}(undef, length(problem_indices)) + for local_gene = eachindex(problem_indices) + offset = state_offsets[local_gene] + n = Int(unique_counts[local_gene]) + change_points = _change_points_from_last(last_values, offset, n) + solutions[local_gene] = FastPIDC.BayesianBlocksSolution( + change_points, + final_scores[local_gene], + ) + end + return solutions + finally + # Explicitly return batch allocations to CUDA's pool. The CUDA backend + # may process many U_g buckets, so relying on a later GC cycle can retain + # unnecessary pressure between batches or after an exception. + for array in ( + prefix_gpu, + block_gpu, + state_offsets_gpu, + block_offsets_gpu, + unique_counts_gpu, + best_gpu, + last_gpu, + final_scores_gpu, + ) + CUDA.unsafe_free!(array) + end + end +end + +function _solve_bb_cuda_batch( + problems::Vector{FastPIDC.BayesianBlocksProblem}, + problem_indices::Vector{Int}, + threads::Int, + ::Type{CountT}, + ::Type{IndexT}, +) where {CountT<:Integer,IndexT<:Integer} + max_u = maximum(i -> length(problems[i].prefix_counts), problem_indices) + priors_gpu = CuArray(_bb_prior_values(max_u)) + try + return _solve_bb_cuda_batch_with_priors( + problems, + problem_indices, + threads, + CountT, + IndexT, + priors_gpu, + ) + finally + CUDA.unsafe_free!(priors_gpu) + end +end + +function FastPIDC.solve_bayesian_blocks_cuda( + problems::Vector{FastPIDC.BayesianBlocksProblem}, + verbose::Bool, +) + CUDA.functional() || return nothing + isempty(problems) && return FastPIDC.BayesianBlocksSolution[] + + sample_count = maximum(p -> Int(round(p.prefix_counts[end])), problems) + max_u = maximum(p -> length(p.prefix_counts), problems) + + # A cumulative prefix count can reach the number of cells, so select the + # smallest exact unsigned type that guards against overflow for this input. + CountT = _smallest_unsigned_type(sample_count) + # 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)) + priors_gpu = CuArray(_bb_prior_values(max_u)) + + if verbose + unique_counts = sort!(collect(length(p.prefix_counts) for p in problems)) + median_u = unique_counts[cld(length(unique_counts), 2)] + println( + "[FastPIDC] CUDA Bayesian blocks: $(length(problems)) genes, " * + "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) + batches = _bb_memory_batches( + bucket, + problems, + budget_bytes, + CountT, + IndexT, + ) + + if verbose + bucket_min_u = minimum(i -> length(problems[i].prefix_counts), bucket) + 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))", + ) + end + + for batch in batches + batch_solutions = _solve_bb_cuda_batch_with_priors( + problems, + batch, + threads, + CountT, + IndexT, + priors_gpu, + ) + for (problem_index, solution) in zip(batch, batch_solutions) + solutions[problem_index] = solution + end + end + end + return solutions + finally + CUDA.unsafe_free!(priors_gpu) + end +end + end # module diff --git a/src/common.jl b/src/common.jl index e82a194..92c130a 100644 --- a/src/common.jl +++ b/src/common.jl @@ -39,7 +39,8 @@ Runtime configuration for PUC/PIDC network inference. Throws an `ArgumentError` if `backend` is not `:cpu` or `:cuda`. """ Base.@kwdef struct PIDCConfig - backend::Symbol = :cuda # :cuda (default) or :cpu + backend::Symbol = :cuda # PUC backend: :cuda or :cpu + bb_backend::Symbol = :cuda # Bayesian blocks backend: :cuda or :cpu discretizer::String = "bayesian_blocks" # mirrors existing default estimator::String = "maximum_likelihood" # mirrors existing default dump_mi_path::Union{Nothing,String} = nothing # Output stem/path; writes *_mi.npy @@ -48,6 +49,7 @@ Base.@kwdef struct PIDCConfig # Inner constructor for automatic validation function PIDCConfig( backend, + bb_backend, discretizer, estimator, dump_mi_path, @@ -58,8 +60,12 @@ Base.@kwdef struct PIDCConfig if !(backend in (:cpu, :cuda)) throw(ArgumentError("backend must be :cpu or :cuda, got :$backend")) end + if !(bb_backend in (:cpu, :cuda)) + throw(ArgumentError("bb_backend must be :cpu or :cuda, got :$bb_backend")) + end new( backend, + bb_backend, discretizer, estimator, dump_mi_path, @@ -69,7 +75,7 @@ Base.@kwdef struct PIDCConfig end end -# NumPy output helpers +# --- NumPy output helpers ----------------------------------------- """ _npy_output_path(file_path) -> String @@ -140,35 +146,57 @@ end """ - Node(line::AbstractArray, discretizer, estimator, number_of_bins) -> Node + Node(label::AbstractString, raw_values::AbstractVector{<:Real}, + discretizer, estimator, number_of_bins) -Construct a [`Node`](@ref) from one row of a data file: `line` is an array -whose first element is the node's label and whose remaining elements are -its raw (continuous) data values. The raw values are discretized using -`discretizer` (overwriting `number_of_bins` if it is `"bayesian_blocks"`, -which chooses its own bin count) and the per-bin probabilities are then -estimated using `estimator`. +Construct a `Node` directly from a label and typed numeric values. This path +avoids materializing the legacy mixed-type `Matrix{Any}` row when loading HDF5 +columns, reducing allocation and type instability without changing the +underlying discretization or probability calculations. """ -function Node(line::AbstractArray, discretizer, estimator, number_of_bins) - - label = string(line[1]) - raw_values = Array{Float64}(line[2:end]) - - # Raw values are mapped to their bin IDs - binned_values = zeros(Int, length(raw_values)) +function Node( + label::AbstractString, + raw_values::AbstractVector{<:Real}, + discretizer, + estimator, + number_of_bins, +) + values = collect(Float64, raw_values) + + # Raw values are mapped to their bin IDs. + binned_values = zeros(Int, length(values)) # If the discretizer is Bayesian blocks, number_of_bins will be # overwritten by the ideal number of bins. Otherwise, it will remain # the same as the value passed in. - number_of_bins = get_bin_ids!(raw_values, discretizer, number_of_bins, binned_values) + number_of_bins = get_bin_ids!(values, discretizer, number_of_bins, binned_values) probabilities = get_probabilities( estimator, get_frequencies_from_bin_ids(binned_values, number_of_bins), ) - return Node(label, binned_values, number_of_bins, probabilities) + return Node(String(label), binned_values, number_of_bins, probabilities) +end + +""" + Node(line::AbstractArray, discretizer, estimator, number_of_bins) -> Node +Construct a [`Node`](@ref) from one row of a data file: `line` is an array +whose first element is the node's label and whose remaining elements are +its raw (continuous) data values. The raw values are discretized using +`discretizer` (overwriting `number_of_bins` if it is `"bayesian_blocks"`, +which chooses its own bin count) and the per-bin probabilities are then +estimated using `estimator`. +""" +function Node(line::AbstractArray, discretizer, estimator, number_of_bins) + return Node( + string(line[1]), + collect(Float64, line[2:end]), + discretizer, + estimator, + number_of_bins, + ) end diff --git a/src/discretizers.jl b/src/discretizers.jl index 15753b3..b41b3ec 100644 --- a/src/discretizers.jl +++ b/src/discretizers.jl @@ -247,20 +247,45 @@ advance; it is determined by [`binedges`](@ref). struct DiscretizeBayesianBlocks <: DiscretizationAlgorithm end """ - binedges(alg::DiscretizeBayesianBlocks, data) -> Vector + BayesianBlocksProblem + +CPU-prepared inputs for the Bayesian-block dynamic program. Preparation is +shared by the CPU and CUDA backends so backend comparisons isolate the dynamic +program itself rather than sorting or unique-value compression. +""" +struct BayesianBlocksProblem + edges::Vector{Float64} + prefix_counts::Vector{Float64} +end + +""" + BayesianBlocksSolution + +Selected change-point indices and the final dynamic-program objective value. +The change points index `BayesianBlocksProblem.edges`. +""" +struct BayesianBlocksSolution + change_points::Vector{Int64} + score::Float64 +end + +# Placeholders extended by FastPIDCCUDAExt when CUDA.jl is loaded. +function bayesian_blocks_cuda_available end +function solve_bayesian_blocks_cuda end -Compute Bayesian-blocks bin edges for `data`, following the histogram -variant of the algorithm in Scargle (2012) (event data, sorted then binned -by maximizing a fitness function via dynamic programming). The number of -edges returned, and hence the number of bins, is chosen adaptively. """ -function binedges(alg::DiscretizeBayesianBlocks, data::AbstractArray{N}) where {N<:AbstractFloat} + prepare_bayesian_blocks(data) +Sort one gene's observations, collapse repeated values, and prepare the shared +prefix-count representation used by both Bayesian-block backends. +""" +function prepare_bayesian_blocks(data::AbstractArray{N}) where {N<:AbstractFloat} # Single sorted pass to get unique values together with their multiplicities, # rather than sorting/uniquing separately and then re-scanning the full # data array once per unique value (which is O(n_unique * length(data))). sorted_data = sort(vec(data)) m = length(sorted_data) + m > 0 || throw(ArgumentError("Bayesian blocks requires at least one observation")) unique_data = Vector{Float64}(undef, m) nn_vec = Vector{Float64}(undef, m) @@ -280,47 +305,72 @@ function binedges(alg::DiscretizeBayesianBlocks, data::AbstractArray{N}) where { resize!(unique_data, n) resize!(nn_vec, n) - edges = zeros(n + 1) + edges = zeros(Float64, n + 1) edges[1] = unique_data[1] for i = 1:(n-1) edges[i+1] = 0.5 * (unique_data[i] + unique_data[i+1]) end edges[end] = unique_data[end] - block_length = unique_data[end] .- edges - count_vec = zeros(n) - best = zeros(n) - last = zeros(Int64, n) + prefix_counts = cumsum(nn_vec) + return BayesianBlocksProblem(edges, prefix_counts) +end - # Reused across iterations so the O(n^2) DP does not also pay for O(n^2) - # bytes of temporary array allocation (one fresh slice/broadcast per K). - widths = Vector{Float64}(undef, n) - fit_vec = Vector{Float64}(undef, n) +function prepare_bayesian_blocks(data::AbstractArray{N}) where {N<:Integer} + return prepare_bayesian_blocks(convert(Array{Float64}, data)) +end + +""" + solve_bayesian_blocks_cpu(problem) +Solve one prepared Bayesian-block problem with the exact prefix-count CPU +dynamic program. This remains the reference implementation for CUDA +conformance testing. +""" +function solve_bayesian_blocks_cpu(problem::BayesianBlocksProblem) + prefix_counts = problem.prefix_counts + # Block lengths are derived from retained edges only when the CPU solver + # needs them, avoiding a second U_g-sized Float64 vector per gene while + # CUDA batches are being prepared. + block_length = problem.edges[end] .- problem.edges + n = length(prefix_counts) + + if n == 1 + return BayesianBlocksSolution(Int64[1, 2], 0.0) + end + + best = zeros(Float64, n) + last = zeros(Int64, n) + + # Prefix counts let each candidate block count be recovered in O(1), + # avoiding the repeated count-vector update and extra full-prefix passes. + # Multiplicities and their cumulative sums are integer-valued Float64s, so + # the resulting counts are bit-exact while the legacy fitness expression + # and first-maximum tie-breaking order remain unchanged. @inbounds for K = 1:n block_length_K1 = block_length[K+1] - for i = 1:K - widths[i] = block_length[i] - block_length_K1 - end - for i = 1:K - count_vec[i] += nn_vec[K] - end # Prior (eq. 21 from Scargle 2012) prior = 4 - log(73.53 * 0.05 * ((K)^-0.478)) - # Fitness function (eq. 19 from Scargle 2012) - for i = 1:K - fit_vec[i] = count_vec[i] * log(count_vec[i] / widths[i]) - prior - end - for i = 2:K - fit_vec[i] += best[i-1] - end + # Initialize from the first candidate exactly as the legacy loop did, + # then scan the remaining candidates in the same order and retain the + # first maximum on ties. + count = prefix_counts[K] + width = block_length[1] - block_length_K1 + best_val = count * log(count / width) - prior i_max = 1 - best_val = fit_vec[1] + for i = 2:K - if fit_vec[i] > best_val - best_val = fit_vec[i] + count = prefix_counts[K] - prefix_counts[i-1] + width = block_length[i] - block_length_K1 + + # Fitness function (eq. 19 from Scargle 2012) + fit = count * log(count / width) - prior + fit += best[i-1] + + if fit > best_val + best_val = fit i_max = i end end @@ -328,8 +378,11 @@ function binedges(alg::DiscretizeBayesianBlocks, data::AbstractArray{N}) where { best[K] = best_val end - change_points = zeros(Int64, n) - i_cp = n + 1 + # The maximal partition has one block per unique value and therefore + # contains n + 1 edge indices. Reserve that full path length so valid + # all-singleton partitions cannot underflow the backtracking buffer. + change_points = zeros(Int64, n + 1) + i_cp = n + 2 ind = n + 1 while true i_cp -= 1 @@ -340,10 +393,26 @@ function binedges(alg::DiscretizeBayesianBlocks, data::AbstractArray{N}) where { ind = last[ind-1] end change_points = change_points[i_cp:end] - edges[change_points] + return BayesianBlocksSolution(change_points, best[end]) +end +""" + binedges(alg::DiscretizeBayesianBlocks, data) -> Vector + +Compute Bayesian-blocks bin edges for `data`, following the histogram +variant of the algorithm in Scargle (2012) (event data, sorted then binned +by maximizing a fitness function via dynamic programming). The number of +edges returned, and hence the number of bins, is chosen adaptively. +""" +function binedges( + alg::DiscretizeBayesianBlocks, + data::AbstractArray{N}, +) where {N<:AbstractFloat} + problem = prepare_bayesian_blocks(data) + solution = solve_bayesian_blocks_cpu(problem) + return problem.edges[solution.change_points] end + function binedges(alg::DiscretizeBayesianBlocks, data::AbstractArray{N}) where {N<:Integer} - data = convert(Array{Float64}, data) - return binedges(alg, data) + return binedges(alg, convert(Array{Float64}, data)) end diff --git a/src/infer_network.jl b/src/infer_network.jl index bc52d86..af7f8a5 100644 --- a/src/infer_network.jl +++ b/src/infer_network.jl @@ -1,5 +1,221 @@ # Helper functions for inferring a network from a data file +function _validate_bb_backend(bb_backend::Symbol) + bb_backend in (:cpu, :cuda) || + throw(ArgumentError("bb_backend must be :cpu or :cuda, got :$bb_backend")) + return bb_backend +end + + +function _bayesian_blocks_cuda_available() + return hasmethod(bayesian_blocks_cuda_available, Tuple{}) && + bayesian_blocks_cuda_available() +end + +function _solve_bayesian_blocks_batch( + problems::Vector{BayesianBlocksProblem}; + bb_backend::Symbol, + verbose::Bool, +) + _validate_bb_backend(bb_backend) + isempty(problems) && return BayesianBlocksSolution[] + + if bb_backend == :cuda + if _bayesian_blocks_cuda_available() && + hasmethod(solve_bayesian_blocks_cuda, (typeof(problems), Bool)) + solutions = solve_bayesian_blocks_cuda(problems, verbose) + solutions !== nothing && return solutions + end + + @warn "CUDA Bayesian blocks requested, but no functional CUDA GPU was found. Falling back to the CPU reference implementation." + end + + solutions = Vector{BayesianBlocksSolution}(undef, length(problems)) + Threads.@threads for i = eachindex(problems) + solutions[i] = solve_bayesian_blocks_cpu(problems[i]) + end + return solutions +end + +function _node_from_bayesian_solution( + label::AbstractString, + values::Vector{Float64}, + problem::BayesianBlocksProblem, + solution::BayesianBlocksSolution, + estimator, +) + edges = problem.edges[solution.change_points] + binned_values = encode(LinearDiscretizer(edges), values) + number_of_bins = length(edges) - 1 + probabilities = get_probabilities( + estimator, + get_frequencies_from_bin_ids(binned_values, number_of_bins), + ) + return Node(String(label), binned_values, number_of_bins, probabilities) +end + +""" + _build_nodes(labels, value_at; ...) + +Construct nodes from a callable returning one gene's observations. Bayesian +blocks are prepared once on the CPU, then solved as a batch by the selected +backend. Quantization or other input transformations remain upstream of +FastPIDC; this function operates on the values exactly as supplied. +""" +function _build_nodes( + labels::AbstractVector, + value_at; + discretizer, + estimator, + number_of_bins, + bb_backend::Symbol, + verbose::Bool, +) + number_of_nodes = length(labels) + _validate_bb_backend(bb_backend) + + if discretizer == "bayesian_blocks" && bb_backend == :cuda && + !_bayesian_blocks_cuda_available() + @warn "CUDA Bayesian blocks requested, but no functional CUDA GPU was found. Falling back to the CPU reference implementation." + bb_backend = :cpu + end + + if discretizer != "bayesian_blocks" || bb_backend == :cpu + # Keep the CPU reference on the original per-gene construction path. + # This avoids retaining every gene's prepared unique-value arrays in + # memory when batching is unnecessary. + nodes = Array{Node}(undef, number_of_nodes) + Threads.@threads for i = 1:number_of_nodes + nodes[i] = Node( + string(labels[i]), + collect(Float64, vec(value_at(i))), + discretizer, + estimator, + number_of_bins, + ) + end + return nodes + end + + preparation_start_ns = time_ns() + + # Sorting and unique-value compression are shared by the CPU and CUDA + # solvers. The resulting U_g values are also all the CUDA backend needs to + # form lightweight workload buckets; no extra scan of the expression data + # is required for bucket selection. + problems = Vector{Union{Nothing,BayesianBlocksProblem}}(undef, number_of_nodes) + fallback = zeros(UInt8, number_of_nodes) + Threads.@threads for i = 1:number_of_nodes + values = collect(Float64, vec(value_at(i))) + try + problems[i] = prepare_bayesian_blocks(values) + catch + problems[i] = nothing + fallback[i] = 1 + end + end + + preparation_seconds = (time_ns() - preparation_start_ns) / 1.0e9 + if verbose + unique_counts = sort!([ + problem === nothing ? 0 : length(problem.prefix_counts) for problem in problems + ]) + nonzero_counts = [u for u in unique_counts if u != 0] + if !isempty(nonzero_counts) + percentile_index(p) = clamp(ceil(Int, p * length(nonzero_counts)), 1, length(nonzero_counts)) + candidate_work = UInt128(0) + for u in nonzero_counts + candidate_work += UInt128(u) * UInt128(u + 1) ÷ UInt128(2) + end + println( + "[FastPIDC] Bayesian-block preparation: " * + "$(round(preparation_seconds; digits = 2)) s", + ) + println( + "[FastPIDC] Bayesian-block U_g: min=$(first(nonzero_counts)), " * + "median=$(nonzero_counts[percentile_index(0.50)]), " * + "p90=$(nonzero_counts[percentile_index(0.90)]), " * + "p99=$(nonzero_counts[percentile_index(0.99)]), " * + "max=$(last(nonzero_counts)); candidate evaluations=$candidate_work", + ) + end + end + + active_indices = Int[] + active_problems = BayesianBlocksProblem[] + for i = 1:number_of_nodes + problem = problems[i] + if fallback[i] == 0 && problem !== nothing && length(problem.prefix_counts) > 1 + push!(active_indices, i) + push!(active_problems, problem) + end + end + + solve_start_ns = time_ns() + active_solutions = _solve_bayesian_blocks_batch( + active_problems; + bb_backend = bb_backend, + verbose = verbose, + ) + solve_seconds = (time_ns() - solve_start_ns) / 1.0e9 + verbose && println( + "[FastPIDC] Bayesian-block dynamic program ($bb_backend): " * + "$(round(solve_seconds; digits = 2)) s", + ) + + solutions = Vector{Union{Nothing,BayesianBlocksSolution}}(undef, number_of_nodes) + fill!(solutions, nothing) + for (i, solution) in zip(active_indices, active_solutions) + solutions[i] = solution + end + empty!(active_problems) + + encoding_start_ns = time_ns() + nodes = Array{Node}(undef, number_of_nodes) + Threads.@threads for i = 1:number_of_nodes + values = collect(Float64, vec(value_at(i))) + problem = problems[i] + solution = solutions[i] + + if fallback[i] != 0 || problem === nothing + nodes[i] = Node( + string(labels[i]), + values, + "uniform_width", + estimator, + number_of_bins, + ) + println("Bayesian blocks failed for $(labels[i]), fell back to uniform width") + elseif length(problem.prefix_counts) == 1 + binned_values = ones(Int, length(values)) + probabilities = get_probabilities( + estimator, + get_frequencies_from_bin_ids(binned_values, 1), + ) + nodes[i] = Node(String(labels[i]), binned_values, 1, probabilities) + else + nodes[i] = _node_from_bayesian_solution( + string(labels[i]), + values, + problem, + solution::BayesianBlocksSolution, + estimator, + ) + end + problems[i] = nothing + end + + encoding_seconds = (time_ns() - encoding_start_ns) / 1.0e9 + verbose && println( + "[FastPIDC] Bayesian-block bin encoding: " * + "$(round(encoding_seconds; digits = 2)) s", + ) + + return nodes +end + +# Entry point: checks extension and routes to the right loader + """ get_nodes(data_file_path::String; ) -> Vector{Node} @@ -14,11 +230,28 @@ function get_nodes( discretizer = "bayesian_blocks", estimator = "maximum_likelihood", number_of_bins = 10, + bb_backend::Symbol = :cuda, + verbose::Bool = false, ) if endswith(data_file_path, ".h5") - return get_nodes_h5(data_file_path; discretizer, estimator, number_of_bins) + return get_nodes_h5( + data_file_path; + discretizer, + estimator, + number_of_bins, + bb_backend, + verbose, + ) else - return get_nodes_text(data_file_path; delim, discretizer, estimator, number_of_bins) + return get_nodes_text( + data_file_path; + delim, + discretizer, + estimator, + number_of_bins, + bb_backend, + verbose, + ) end end @@ -51,6 +284,8 @@ function get_nodes_h5( discretizer = "bayesian_blocks", estimator = "maximum_likelihood", number_of_bins = 10, + bb_backend::Symbol = :cuda, + verbose::Bool = false, ) nodes = Node[] @@ -98,31 +333,26 @@ function get_nodes_h5( read(data_obj["indices"]) .+ 1, read(data_obj["data"]) ) - number_of_cells = size(X_raw, 1) elseif isa(data_obj, HDF5.Dataset) # Python saved (Cells, Genes) C-Order. Julia reads (Genes, Cells). # We permute dimensions to make it (Cells, Genes) so our slicing logic is uniform. X_raw = permutedims(read(data_obj)) - number_of_cells = size(X_raw, 1) else throw(ArgumentError("Object at '$matrix_key' is neither an HDF5 Group nor a Dataset.")) end - # Build the Nodes - nodes = Array{Node}(undef, number_of_nodes) - Threads.@threads for i = 1:number_of_nodes - label = gene_names[i] - - # Since X_raw is (Cells, Genes), Column `i` is Gene `i` - data_gene = @view X_raw[:, i] - - legacy_format = Matrix{Any}(undef, 1, number_of_cells + 1) - legacy_format[1, 1] = label - legacy_format[1, 2:end] .= data_gene - - nodes[i] = Node(legacy_format, discretizer, estimator, number_of_bins) - end + # Since X_raw is (Cells, Genes), column `i` is gene `i`. + value_at = i -> (@view X_raw[:, i]) + nodes = _build_nodes( + gene_names, + value_at; + discretizer, + estimator, + number_of_bins, + bb_backend, + verbose, + ) end return nodes @@ -147,6 +377,8 @@ Arguments: * `discretizer="bayesian_blocks"`: algorithm for discretizing the data * `estimator="maximum_likelihood"`: algorithm for estimating probabilities * `number_of_bins=10`: will be overwritten if using "bayesian_blocks" +* `bb_backend=:cuda`: Bayesian-block dynamic-program backend (`:cuda` or `:cpu`) +* `verbose=false`: print Bayesian-block phase, workload, and bucket diagnostics The "maximum_likelihood" estimator is recommended for PUC and PIDC. """ @@ -156,6 +388,8 @@ function get_nodes_text( discretizer = "bayesian_blocks", estimator = "maximum_likelihood", number_of_bins = 10, + bb_backend::Symbol = :cuda, + verbose::Bool = false, ) lines = open(data_file_path) do io if delim == false @@ -165,15 +399,17 @@ function get_nodes_text( end end - number_of_nodes = size(lines, 1) - nodes = Array{Node}(undef, number_of_nodes) - - Threads.@threads for i = 1:number_of_nodes - # Note: lines[i:i, 1:end] is a memory trap; it allocates a new matrix for every gene. - nodes[i] = Node(lines[i:i, 1:end], discretizer, estimator, number_of_bins) - end - - return nodes + labels = string.(lines[:, 1]) + value_at = i -> (@view lines[i, 2:end]) + return _build_nodes( + labels, + value_at; + discretizer, + estimator, + number_of_bins, + bb_backend, + verbose, + ) end @@ -356,6 +592,7 @@ Arguments: * `discretizer="bayesian_blocks"`: algorithm for discretizing the data * `estimator="maximum_likelihood"`: algorithm for estimating probabilities * `number_of_bins=10`: will be overwritten if using "bayesian_blocks" +* `config.bb_backend`: Bayesian-block backend used while constructing nodes * `base=2`: base for the information measures * `out_file_path=""`: path to output file. If empty, will not write a file @@ -381,6 +618,8 @@ function infer_network( discretizer = discretizer, estimator = estimator, number_of_bins = number_of_bins, + bb_backend = config.bb_backend, + verbose = config.verbose, ) println("Inferring network...") diff --git a/test/bayesian_blocks_tests.jl b/test/bayesian_blocks_tests.jl new file mode 100644 index 0000000..0b40d0c --- /dev/null +++ b/test/bayesian_blocks_tests.jl @@ -0,0 +1,155 @@ +using FastPIDC +using Test + +# Frozen copy of the Bayesian-block dynamic program immediately before the +# prefix-count refactor. This remains test-only so optimized implementations +# can be required to reproduce the previous edges bit for bit. +function bayesian_blocks_reference(data::AbstractArray{<:AbstractFloat}) + sorted_data = sort(vec(data)) + m = length(sorted_data) + + unique_data = Vector{Float64}(undef, m) + nn_vec = Vector{Float64}(undef, m) + n = 0 + i = 1 + @inbounds while i <= m + v = sorted_data[i] + j = i + 1 + while j <= m && sorted_data[j] == v + j += 1 + end + n += 1 + unique_data[n] = v + nn_vec[n] = j - i + i = j + end + resize!(unique_data, n) + resize!(nn_vec, n) + + edges = zeros(n + 1) + edges[1] = unique_data[1] + for i = 1:(n-1) + edges[i+1] = 0.5 * (unique_data[i] + unique_data[i+1]) + end + edges[end] = unique_data[end] + block_length = unique_data[end] .- edges + + count_vec = zeros(n) + best = zeros(n) + last = zeros(Int64, n) + widths = Vector{Float64}(undef, n) + fit_vec = Vector{Float64}(undef, n) + + @inbounds for K = 1:n + block_length_K1 = block_length[K+1] + for i = 1:K + widths[i] = block_length[i] - block_length_K1 + end + for i = 1:K + count_vec[i] += nn_vec[K] + end + + prior = 4 - log(73.53 * 0.05 * ((K)^-0.478)) + for i = 1:K + fit_vec[i] = count_vec[i] * log(count_vec[i] / widths[i]) - prior + end + for i = 2:K + fit_vec[i] += best[i-1] + end + + i_max = 1 + best_val = fit_vec[1] + for i = 2:K + if fit_vec[i] > best_val + best_val = fit_vec[i] + i_max = i + end + end + last[K] = i_max + best[K] = best_val + end + + change_points = zeros(Int64, n) + i_cp = n + 1 + ind = n + 1 + while true + i_cp -= 1 + change_points[i_cp] = ind + if ind == 1 + break + end + ind = last[ind-1] + end + change_points = change_points[i_cp:end] + return edges[change_points] +end + +function float_bits(values::AbstractVector{Float64}) + return collect(reinterpret(UInt64, values)) +end + +@testset "Bayesian blocks prefix-count equivalence" begin + cases = Vector{Float64}[ + [0.0, 0.0, 0.0, 1.0, 1.0, 2.0, 3.0, 3.0, 8.0, 13.0], + [-4.0, -2.0, -2.0, -1.0, 0.0, 0.0, 0.25, 0.5, 2.0, 9.0], + [0.0, 0.0, 1.0e-12, 2.0e-12, 0.1, 0.1, 1.0, 10.0, 100.0], + collect(range(-3.0, 5.0; length = 24)), + ] + + # Deterministic zero-inflated pseudo-random cases without adding a test + # dependency on Random.jl. + for seed = 1:20 + state = UInt64(seed) + values = Vector{Float64}(undef, 64) + for i = eachindex(values) + state = state * 6364136223846793005 + 1442695040888963407 + value = Float64(Int(state % UInt64(2001)) - 1000) / 1000 + values[i] = state % UInt64(5) == 0 ? 0.0 : value + end + push!(cases, values) + end + + algorithm = FastPIDC.DiscretizeBayesianBlocks() + for values in cases + reference_edges = bayesian_blocks_reference(values) + optimized_edges = FastPIDC.binedges(algorithm, values) + + @test optimized_edges == reference_edges + @test float_bits(optimized_edges) == float_bits(reference_edges) + + reference_ids = FastPIDC.encode( + FastPIDC.LinearDiscretizer(reference_edges), + values, + ) + optimized_ids = FastPIDC.encode( + FastPIDC.LinearDiscretizer(optimized_edges), + values, + ) + @test optimized_ids == reference_ids + end +end + +@testset "Typed Node constructor preserves legacy results" begin + values = [0.0, 0.0, 0.5, 1.0, 1.0, 2.0, 3.5, 8.0, 8.0, 13.0] + + # Reproduce the mixed-type row accepted by the legacy constructor. + legacy_line = Matrix{Any}(undef, 1, length(values) + 1) + legacy_line[1, 1] = "G1" + legacy_line[1, 2:end] .= values + + typed_node = Node("G1", values, "bayesian_blocks", "maximum_likelihood", 10) + legacy_node = Node(legacy_line, "bayesian_blocks", "maximum_likelihood", 10) + + @test typed_node.label == legacy_node.label + @test typed_node.number_of_bins == legacy_node.number_of_bins + @test typed_node.binned_values == legacy_node.binned_values + @test typed_node.probabilities == legacy_node.probabilities + @test float_bits(typed_node.probabilities) == float_bits(legacy_node.probabilities) +end + + +@testset "Bayesian blocks backend configuration" begin + @test PIDCConfig().bb_backend == :cuda + @test PIDCConfig(bb_backend = :cpu).bb_backend == :cpu + @test_throws ArgumentError PIDCConfig(bb_backend = :invalid) +end diff --git a/test/benchmark_puc_bb.jl b/test/benchmark_puc_bb.jl new file mode 100644 index 0000000..6fddce8 --- /dev/null +++ b/test/benchmark_puc_bb.jl @@ -0,0 +1,180 @@ +using FastPIDC +using CUDA +using Distributed +using Statistics + +# Manual benchmark: run this file directly. Do not include it from test/runtests.jl. +const TARGET_CPU_WORKERS = 10 +const CURRENT_CPU_WORKERS = max(nprocs() - 1, 0) + +if CURRENT_CPU_WORKERS < TARGET_CPU_WORKERS + needed = TARGET_CPU_WORKERS - CURRENT_CPU_WORKERS + active_project = Base.active_project() + + if active_project === nothing + addprocs(needed) + else + addprocs(needed; exeflags = "--project=$(dirname(active_project))") + end +end + +@everywhere using FastPIDC + +CUDA.functional() || error("CUDA is not functional on this system.") + +const DATA_DIR = joinpath(dirname(@__FILE__), "data") +const DATASET = joinpath(DATA_DIR, "toy_large_1k.txt") +const WARMUP_DATASET = joinpath(DATA_DIR, "toy_small_200.txt") + +function load_nodes(path::String, config::PIDCConfig) + return get_nodes( + path; + discretizer = config.discretizer, + estimator = config.estimator, + bb_backend = config.bb_backend, + verbose = config.verbose, + ) +end + +function infer_puc(nodes, config::PIDCConfig) + return InferredNetwork( + PUCNetworkInference(), + nodes; + estimator = config.estimator, + config = config, + ) +end + +function timed_cuda(f::F) where {F} + CUDA.synchronize() + result = nothing + elapsed = @elapsed begin + result = f() + CUDA.synchronize() + end + return result, elapsed +end + +function edge_key(edge) + label_1 = edge.nodes[1].label + label_2 = edge.nodes[2].label + return label_1 <= label_2 ? (label_1, label_2) : (label_2, label_1) +end + +function numeric_diagnostics(cpu_network, gpu_network; top_k::Int = 250) + cpu_weights = Dict(edge_key(edge) => edge.weight for edge in cpu_network.edges) + gpu_weights = Dict(edge_key(edge) => edge.weight for edge in gpu_network.edges) + + Set(keys(cpu_weights)) == Set(keys(gpu_weights)) || + error("CPU and GPU networks do not contain the same edge set.") + + edge_keys = collect(keys(cpu_weights)) + cpu_values = [cpu_weights[key] for key in edge_keys] + gpu_values = [gpu_weights[key] for key in edge_keys] + + absolute_errors = abs.(gpu_values .- cpu_values) + nonzero_reference = abs.(cpu_values) .> 1.0e-12 + relative_errors = absolute_errors[nonzero_reference] ./ abs.(cpu_values[nonzero_reference]) + ratios = gpu_values[nonzero_reference] ./ cpu_values[nonzero_reference] + + effective_top_k = min(top_k, length(cpu_network.edges), length(gpu_network.edges)) + cpu_top = Set(edge_key(edge) for edge in cpu_network.edges[1:effective_top_k]) + gpu_top = Set(edge_key(edge) for edge in gpu_network.edges[1:effective_top_k]) + + println("\n--- Correctness Check ---") + println( + "Top $effective_top_k Edge Overlap: ", + length(intersect(cpu_top, gpu_top)), + " / ", + effective_top_k, + ) + println("Max Absolute Error: ", maximum(absolute_errors)) + println( + "Max Relative Error: ", + isempty(relative_errors) ? NaN : maximum(relative_errors), + ) + println("Mean Ratio (GPU/CPU): ", isempty(ratios) ? NaN : mean(ratios)) +end + +function main() + isfile(DATASET) || error("Benchmark dataset not found: $DATASET") + + config_cpu = PIDCConfig( + backend = :cpu, + bb_backend = :cpu, + discretizer = "bayesian_blocks", + estimator = "maximum_likelihood", + verbose = false, + ) + + config_cuda = PIDCConfig( + backend = :cuda, + bb_backend = :cuda, + discretizer = "bayesian_blocks", + estimator = "maximum_likelihood", + verbose = false, + ) + + # Compile both execution paths before timing the full dataset. + if isfile(WARMUP_DATASET) + println("Warming up CPU and CUDA paths with $WARMUP_DATASET ...") + + warm_cpu_nodes = load_nodes(WARMUP_DATASET, config_cpu) + infer_puc(warm_cpu_nodes[1:min(100, length(warm_cpu_nodes))], config_cpu) + + warm_cuda_nodes = load_nodes(WARMUP_DATASET, config_cuda) + infer_puc(warm_cuda_nodes[1:min(100, length(warm_cuda_nodes))], config_cuda) + + warm_cpu_nodes = nothing + warm_cuda_nodes = nothing + GC.gc() + CUDA.reclaim() + else + println( + "Warmup dataset not found; reported timings include first-call compilation.", + ) + end + + println("\nDataset: $DATASET") + println("CPU workers for PUC: $(nprocs() - 1)") + println("Julia threads for BB preparation/CPU BB: $(Threads.nthreads())") + + println("\n--- CPU BB + CPU PUC ---") + GC.gc() + cpu_nodes = nothing + t_cpu_bb = @elapsed cpu_nodes = load_nodes(DATASET, config_cpu) + println("CPU Bayesian-block time: $t_cpu_bb seconds") + println("Loaded $(length(cpu_nodes)) CPU-binned nodes.") + + cpu_network = nothing + t_cpu_puc = @elapsed cpu_network = infer_puc(cpu_nodes, config_cpu) + println("CPU PUC time: $t_cpu_puc seconds") + + t_cpu_total = t_cpu_bb + t_cpu_puc + println("CPU total time: $t_cpu_total seconds") + + println("\n--- CUDA BB + CUDA PUC ---") + GC.gc() + CUDA.reclaim() + + gpu_nodes, t_cuda_bb = timed_cuda(() -> load_nodes(DATASET, config_cuda)) + println("CUDA Bayesian-block time: $t_cuda_bb seconds") + println("Loaded $(length(gpu_nodes)) CUDA-binned nodes.") + + gpu_network, t_cuda_puc = timed_cuda(() -> infer_puc(gpu_nodes, config_cuda)) + println("CUDA PUC time: $t_cuda_puc seconds") + + t_cuda_total = t_cuda_bb + t_cuda_puc + println("CUDA total time: $t_cuda_total seconds") + + numeric_diagnostics(cpu_network, gpu_network) + + println("\n--- Speedup ---") + println("Bayesian blocks: $(round(t_cpu_bb / t_cuda_bb; digits = 2))x") + println("PUC scoring: $(round(t_cpu_puc / t_cuda_puc; digits = 2))x") + println("End to end: $(round(t_cpu_total / t_cuda_total; digits = 2))x") + + return nothing +end + +main() diff --git a/test/cuda_bayesian_blocks_tests.jl b/test/cuda_bayesian_blocks_tests.jl new file mode 100644 index 0000000..dc89af7 --- /dev/null +++ b/test/cuda_bayesian_blocks_tests.jl @@ -0,0 +1,265 @@ +using FastPIDC +using Test +using CUDA + +function _bb_test_values() + values = Vector{Float64}[ + [ + 0.0, 0.0, 0.0, 0.0, + 0.1, 0.1, 0.2, 0.2, + 3.0, 3.0, 3.1, 3.1, + 8.0, 8.1, 8.1, 8.2, + ], + [ + -5.0, -5.0, -4.9, -4.8, + -0.2, -0.1, 0.0, 0.1, + 4.8, 4.9, 5.0, 5.0, + 12.0, 12.0, 12.1, 12.2, + ], + vcat(fill(0.0, 12), collect(0.5:0.5:8.0), fill(15.0, 12)), + [Float64(i^2) / 17 for i = 0:39], + fill(2.5, 32), + ] + + # Add deterministic zero-inflated and clustered cases with varying U_g so + # conformance exercises more than a few hand-selected partitions. + for case = 1:16 + n = 48 + 3 * case + x = [ + sin((sample + case) / 5) + + 0.2 * cos((2 * sample + case) / 7) + + (sample % (5 + case % 4) == 0 ? 3.0 : 0.0) for sample = 1:n + ] + x[case:(4 + case % 5):end] .= 0.0 + push!(values, x) + end + return values +end + +if CUDA.functional() + @testset "CUDA Bayesian blocks equivalence and determinism" begin + cuda_ext = Base.get_extension(FastPIDC, :FastPIDCCUDAExt) + @test cuda_ext !== nothing + + values_by_gene = _bb_test_values() + problems = FastPIDC.prepare_bayesian_blocks.(values_by_gene) + cpu_solutions = FastPIDC.solve_bayesian_blocks_cpu.(problems) + gpu_solutions_1 = FastPIDC.solve_bayesian_blocks_cuda(problems, false) + gpu_solutions_2 = FastPIDC.solve_bayesian_blocks_cuda(problems, false) + + @test gpu_solutions_1 !== nothing + @test gpu_solutions_2 !== nothing + + @testset "CPU vs GPU Bayesian-block result" begin + for i = eachindex(problems) + cpu_solution = cpu_solutions[i] + gpu_solution = gpu_solutions_1[i] + + # CUDA and CPU libdevice/libm logarithms may differ by a few + # bits, but robust test cases should retain the same optimum. + @test isapprox( + gpu_solution.score, + cpu_solution.score; + atol = 1e-9, + rtol = 1e-12, + ) + @test gpu_solution.change_points == cpu_solution.change_points + + cpu_edges = problems[i].edges[cpu_solution.change_points] + gpu_edges = problems[i].edges[gpu_solution.change_points] + @test gpu_edges == cpu_edges + + # A constant gene has one unique value, so its two outer + # Bayesian-block edges are equal. The CPU and CUDA solvers can + # still be compared above, but a zero-width interval is not a + # valid LinearDiscretizer and therefore has no encoding step. + if length(problems[i].prefix_counts) == 1 + @test all(==(only(unique(values_by_gene[i]))), values_by_gene[i]) + continue + end + + cpu_ids = FastPIDC.encode( + FastPIDC.LinearDiscretizer(cpu_edges), + values_by_gene[i], + ) + gpu_ids = FastPIDC.encode( + FastPIDC.LinearDiscretizer(gpu_edges), + values_by_gene[i], + ) + @test gpu_ids == cpu_ids + end + end + + @testset "CUDA Bayesian blocks determinism" begin + for i = eachindex(gpu_solutions_1) + @test gpu_solutions_1[i].change_points == + gpu_solutions_2[i].change_points + @test reinterpret(UInt64, [gpu_solutions_1[i].score]) == + reinterpret(UInt64, [gpu_solutions_2[i].score]) + end + + # Candidate assignment changes with block size, but the ordered + # first-maximum reduction must produce the same deterministic DP. + all_indices = collect(eachindex(problems)) + solutions_32 = cuda_ext._solve_bb_cuda_batch( + problems, + all_indices, + 32, + UInt8, + UInt8, + ) + solutions_64 = cuda_ext._solve_bb_cuda_batch( + problems, + all_indices, + 64, + UInt8, + UInt8, + ) + for i = eachindex(problems) + @test solutions_32[i].change_points == solutions_64[i].change_points + @test reinterpret(UInt64, [solutions_32[i].score]) == + reinterpret(UInt64, [solutions_64[i].score]) + end + + # Compile and exercise the UInt16 prefix-count/back-pointer kernel + # variant used when either cells or U_g exceed UInt8 capacity. + wide_values = [Float64(i^2) / 101 for i = 0:299] + wide_problem = FastPIDC.prepare_bayesian_blocks(wide_values) + wide_cpu = FastPIDC.solve_bayesian_blocks_cpu(wide_problem) + wide_gpu = only( + cuda_ext._solve_bb_cuda_batch( + [wide_problem], + [1], + 64, + UInt16, + UInt16, + ), + ) + @test wide_gpu.change_points == wide_cpu.change_points + @test isapprox(wide_gpu.score, wide_cpu.score; atol = 1e-9, rtol = 1e-12) + + @test_throws ArgumentError cuda_ext._solve_bb_cuda_batch( + problems, + all_indices, + 48, + UInt8, + UInt8, + ) + + reversed_solutions = FastPIDC.solve_bayesian_blocks_cuda( + reverse(problems), + false, + ) + @test reverse([s.change_points for s in reversed_solutions]) == + [s.change_points for s in gpu_solutions_1] + end + + @testset "Bayesian-block backtracking edge cases" begin + @test cuda_ext._change_points_from_last(UInt8[1], 1, 1) == + Int64[1, 2] + @test cuda_ext._change_points_from_last(UInt8[1, 2], 1, 2) == + Int64[1, 2, 3] + @test cuda_ext._change_points_from_last(UInt8[1, 1], 1, 2) == + Int64[1, 3] + @test_throws ArgumentError cuda_ext._change_points_from_last( + UInt8[1, 0], + 1, + 2, + ) + end + + @testset "Lightweight U_g bucketing and batching" begin + buckets = cuda_ext._bb_quantile_buckets(problems) + @test sort(vcat(buckets...)) == collect(eachindex(problems)) + @test cuda_ext._bb_threads_for_max_u(32) == 32 + @test cuda_ext._bb_threads_for_max_u(33) == 64 + @test cuda_ext._bb_threads_for_max_u(513) == 128 + @test cuda_ext._bb_threads_for_max_u(4_097) == 256 + + # The batching helper requires a concrete byte count. Keep this + # focused check so an accidental bare `return` cannot silently + # turn the estimate into `nothing` again. + u = length(problems[1].prefix_counts) + expected_bytes = + sizeof(Float64) * (u + 1) + + sizeof(UInt8) * u + + sizeof(Float64) * u + + sizeof(UInt8) * u + @test cuda_ext._bb_problem_bytes(problems[1], UInt8, UInt8) == + expected_bytes + + batches = cuda_ext._bb_memory_batches( + buckets[1], + problems, + typemax(Int), + UInt8, + UInt8, + ) + @test batches == [buckets[1]] + end + + @testset "End-to-end nodes and PUC integrity" begin + # Use equal-length columns so both backends receive identical + # observations through the same batched node-construction path. + matrix = hcat( + repeat(values_by_gene[1], 3), + repeat(values_by_gene[2], 3), + vcat(fill(0.0, 16), collect(0.5:0.5:8.0), fill(15.0, 16)), + vcat(values_by_gene[4], values_by_gene[4][1:8]), + ) + labels = ["G1", "G2", "G3", "G4"] + value_at = i -> (@view matrix[:, i]) + + nodes_cpu = FastPIDC._build_nodes( + labels, + value_at; + discretizer = "bayesian_blocks", + estimator = "maximum_likelihood", + number_of_bins = 10, + bb_backend = :cpu, + verbose = false, + ) + nodes_gpu = FastPIDC._build_nodes( + labels, + value_at; + discretizer = "bayesian_blocks", + estimator = "maximum_likelihood", + number_of_bins = 10, + bb_backend = :cuda, + verbose = false, + ) + + for i = eachindex(nodes_cpu) + @test nodes_gpu[i].number_of_bins == nodes_cpu[i].number_of_bins + @test nodes_gpu[i].binned_values == nodes_cpu[i].binned_values + @test nodes_gpu[i].probabilities == nodes_cpu[i].probabilities + end + + # Exercise the public loader and its independent bb_backend option. + mktemp() do path, io + println(io, join(vcat("gene", ["S$i" for i = 1:size(matrix, 1)]), " ")) + for gene = 1:size(matrix, 2) + println(io, join(vcat(labels[gene], string.(matrix[:, gene])), " ")) + end + close(io) + + public_cpu = get_nodes(path; bb_backend = :cpu) + public_gpu = get_nodes(path; bb_backend = :cuda) + for i = eachindex(public_cpu) + @test public_gpu[i].number_of_bins == public_cpu[i].number_of_bins + @test public_gpu[i].binned_values == public_cpu[i].binned_values + end + end + + config_cuda = PIDCConfig(backend = :cuda, bb_backend = :cuda) + mi_cpu_bins, puc_cpu_bins = + FastPIDC.compute_puc_full_cuda(nodes_cpu, config_cuda, 2) + mi_gpu_bins, puc_gpu_bins = + FastPIDC.compute_puc_full_cuda(nodes_gpu, config_cuda, 2) + @test mi_gpu_bins == mi_cpu_bins + @test puc_gpu_bins == puc_cpu_bins + end + end +else + @warn "CUDA unavailable. Skipping CUDA Bayesian-block tests." +end diff --git a/test/cuda_numeric_tests.jl b/test/cuda_numeric_tests.jl index f3913da..2a450ef 100644 --- a/test/cuda_numeric_tests.jl +++ b/test/cuda_numeric_tests.jl @@ -13,6 +13,33 @@ if CUDA.functional() config_cpu = PIDCConfig(backend = :cpu, verbose = false) config_cuda = PIDCConfig(backend = :cuda, verbose = false) + cuda_ext = Base.get_extension(FastPIDC, :FastPIDCCUDAExt) + @test cuda_ext !== nothing + + @testset "Compact GPU integer type selection" begin + @test cuda_ext._smallest_unsigned_type(34) == UInt8 + @test cuda_ext._smallest_unsigned_type(255) == UInt8 + @test cuda_ext._smallest_unsigned_type(256) == UInt16 + @test cuda_ext._smallest_unsigned_type(38_176) == UInt16 + @test cuda_ext._smallest_unsigned_type(65_535) == UInt16 + @test cuda_ext._smallest_unsigned_type(65_536) == UInt32 + @test_throws ArgumentError cuda_ext._smallest_unsigned_type(-1) + end + + @testset "Compact storage matches Int32 reference" begin + mi_compact, puc_compact = + FastPIDC.compute_puc_full_cuda(nodes, config_cuda, 2) + mi_int32, puc_int32 = cuda_ext._compute_puc_full_cuda_typed( + nodes, + config_cuda, + 2, + Int32, + Int32, + ) + + @test mi_compact == mi_int32 + @test puc_compact == puc_int32 + end # We need to call the internal matrix generators directly, not just the Network wrapper mi_cpu, puc_cpu = FastPIDC.compute_puc_full(nodes, config = config_cpu, base = 2) @@ -28,7 +55,7 @@ if CUDA.functional() # --- TEST 2: Determinism (Race Condition Check) --- @testset "GPU Determinism" begin - # Run the GPU calculation twice on the exact same data + # Compare the two loader paths on the same data mi_gpu1, puc_gpu1 = FastPIDC.compute_puc_full_cuda(nodes, config_cuda, 2) mi_gpu2, puc_gpu2 = FastPIDC.compute_puc_full_cuda(nodes, config_cuda, 2) @@ -69,10 +96,10 @@ if CUDA.functional() # --- TEST 4: txt vs. H5 Determinism --- @testset "txt vs. H5 Determinism" begin node_txt = get_nodes(dataset) - node_h5 = get_nodes(joinpath(DATA_DIR, "toy_small_200.h5")) - # Run the GPU calculation twice on the exact same data - mi_txt, puc_txt = FastPIDC.compute_puc_full_cuda(nodes, config_cuda, 2) - mi_h5, puc_h5 = FastPIDC.compute_puc_full_cuda(nodes, config_cuda, 2) + node_h5 = get_nodes(joinpath(DATA_DIR, "toy_small_200.h5")) + # Compare the two loader paths on the same data + mi_txt, puc_txt = FastPIDC.compute_puc_full_cuda(node_txt, config_cuda, 2) + mi_h5, puc_h5 = FastPIDC.compute_puc_full_cuda(node_h5, config_cuda, 2) # These should be bit-for-bit identical. If they aren't, # there is an atomic race condition or uninitialized memory. diff --git a/test/cuda_numeric_tests_bb.jl b/test/cuda_numeric_tests_bb.jl new file mode 100644 index 0000000..3ed17eb --- /dev/null +++ b/test/cuda_numeric_tests_bb.jl @@ -0,0 +1,139 @@ +using FastPIDC +using Test +using CUDA +using Statistics + +const BB_NUMERIC_DATA_DIR = joinpath(dirname(@__FILE__), "data") + +""" +Return the unique undirected edge weights and their `(i, j)` node-index pairs. +Only the strict upper triangle is used so the diagonal and symmetric duplicate +entries do not distort the numerical diagnostics. +""" +function _bb_undirected_scores(scores::AbstractMatrix{<:Real}) + n, m = size(scores) + n == m || throw(ArgumentError("score matrix must be square")) + + number_of_edges = n * (n - 1) ÷ 2 + weights = Vector{Float64}(undef, number_of_edges) + pairs = Vector{Tuple{Int,Int}}(undef, number_of_edges) + + edge = 1 + @inbounds for j = 2:n + for i = 1:(j-1) + weights[edge] = Float64(scores[i, j]) + pairs[edge] = (i, j) + edge += 1 + end + end + + return weights, pairs +end + +function _bb_top_edge_set( + weights::AbstractVector{<:Real}, + pairs::AbstractVector{Tuple{Int,Int}}, + k::Int, +) + top_k = min(k, length(weights)) + order = sortperm(eachindex(weights); by = i -> (-weights[i], i)) + return Set(pairs[order[1:top_k]]), top_k +end + +if CUDA.functional() + @testset "Bayesian-block CPU/CPU vs GPU/GPU numeric diagnostics" begin + dataset = joinpath(BB_NUMERIC_DATA_DIR, "toy_small_200.txt") + + config_cpu = PIDCConfig( + backend = :cpu, + bb_backend = :cpu, + verbose = false, + ) + config_gpu = PIDCConfig( + backend = :cuda, + bb_backend = :cuda, + verbose = false, + ) + + # CPU/CPU means CPU Bayesian blocks followed by CPU PUC scoring. + nodes_cpu = get_nodes( + dataset; + discretizer = "bayesian_blocks", + estimator = "maximum_likelihood", + number_of_bins = 10, + bb_backend = config_cpu.bb_backend, + verbose = false, + ) + mi_cpu, puc_cpu = FastPIDC.compute_puc_full( + nodes_cpu; + config = config_cpu, + base = 2, + ) + + # GPU/GPU means CUDA Bayesian blocks followed by CUDA PUC scoring. + nodes_gpu = get_nodes( + dataset; + discretizer = "bayesian_blocks", + estimator = "maximum_likelihood", + number_of_bins = 10, + bb_backend = config_gpu.bb_backend, + verbose = false, + ) + mi_gpu, puc_gpu = FastPIDC.compute_puc_full( + nodes_gpu; + config = config_gpu, + base = 2, + ) + + @test getfield.(nodes_gpu, :label) == getfield.(nodes_cpu, :label) + @test size(mi_gpu) == size(mi_cpu) + @test size(puc_gpu) == size(puc_cpu) + @test mi_cpu == transpose(mi_cpu) + @test puc_cpu == transpose(puc_cpu) + @test isapprox(mi_gpu, transpose(mi_gpu); atol = 1e-8, rtol = 1e-10) + @test isapprox(puc_gpu, transpose(puc_gpu); atol = 1e-8, rtol = 1e-10) + + cpu_weights, edge_pairs = _bb_undirected_scores(puc_cpu) + gpu_weights, gpu_edge_pairs = _bb_undirected_scores(puc_gpu) + @test gpu_edge_pairs == edge_pairs + + absolute_errors = abs.(gpu_weights .- cpu_weights) + max_absolute_error = maximum(absolute_errors) + + # Relative error and ratios are undefined or uninformative when the CPU + # reference is effectively zero, so exclude only those values. + reference_floor = 1e-12 + nonzero_reference = abs.(cpu_weights) .> reference_floor + relative_errors = absolute_errors[nonzero_reference] ./ + abs.(cpu_weights[nonzero_reference]) + ratios = gpu_weights[nonzero_reference] ./ cpu_weights[nonzero_reference] + + max_relative_error = + isempty(relative_errors) ? 0.0 : maximum(relative_errors) + mean_ratio = isempty(ratios) ? 1.0 : mean(ratios) + + cpu_top, top_k = _bb_top_edge_set(cpu_weights, edge_pairs, 250) + gpu_top, _ = _bb_top_edge_set(gpu_weights, edge_pairs, 250) + top_overlap = length(intersect(cpu_top, gpu_top)) + + println("\n--- Bayesian Blocks CPU/CPU vs GPU/GPU Diagnostic ---") + println("[Diagnostics] Top $top_k Edge Overlap: $top_overlap / $top_k") + println("[Diagnostics] Max Absolute Error: $max_absolute_error") + println("[Diagnostics] Max Relative Error: $max_relative_error") + println("[Diagnostics] Mean Ratio (GPU/CPU): $mean_ratio") + + # Retain broad integrity assertions while leaving the printed metrics + # available for judging smaller CUDA/libm rank and score differences. + @test top_overlap == top_k + @test all( + isapprox.( + cpu_weights, + gpu_weights; + atol = 1e-8, + rtol = 1e-10, + ), + ) + end +else + @warn "CUDA unavailable. Skipping Bayesian-block CPU/CPU vs GPU/GPU numeric diagnostics." +end diff --git a/test/runtests.jl b/test/runtests.jl index 0312ea3..3efd609 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -15,9 +15,13 @@ using DelimitedFiles include("cuda_smoke_tests.jl") include("baseline_helpers.jl") include("baseline_smoke_tests.jl") +include("bayesian_blocks_tests.jl") +include("cuda_bayesian_blocks_tests.jl") include("diagnostic_dump_tests.jl") include("cuda_numeric_tests.jl") +include("cuda_numeric_tests_bb.jl") include("benchmark_puc.jl") +include("benchmark_puc_bb.jl") # These tests use a dataset generated from the 10-node Yeast1 network from http://gnw.sourceforge.net/