diff --git a/Project.toml b/Project.toml index 4ffe88e..131ed21 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "FixedEffects" uuid = "c8885935-8500-56a7-9867-7708b20db0eb" -version = "3.3.1" +version = "3.4.0" [deps] GroupedArrays = "6407cd72-fade-4a84-8a1e-56e431fc1533" diff --git a/benchmarks/akm_worker_firm.jl b/benchmarks/akm_worker_firm.jl new file mode 100644 index 0000000..836a27c --- /dev/null +++ b/benchmarks/akm_worker_firm.jl @@ -0,0 +1,209 @@ +using FixedEffects +using LinearAlgebra +using Printf +using Random +using StatsBase + +# Simulated counterpart to the xhdfe AKM-style performance example: +# log wage on seniority controls, absorbing worker, firm, and year fixed effects. +# This package is the residualization backend, so worker-clustered SEs are out of scope here. +# Usage: julia --project -t auto benchmarks/akm_worker_firm.jl [--seed=1234] [--method=cpu] +# [--double-precision=true] [--tol=1e-8] [--maxiter=Inf] +const N_WORKERS = 50_000 +const N_FIRMS = 7_000 +const N_YEARS = 36 +const OBS_PER_WORKER = 8 +const FIRM_WINDOW = 10 +const MOVE_PROBABILITY = 0.35 +const CONTROL_NAMES = ("tenure", "tenure_sq", "experience", "experience_sq", "mover") + +struct AKMPanel + worker::Vector{Int32} + firm::Vector{Int32} + year::Vector{Int32} + y::Vector{Float64} + x::Matrix{Float64} +end + +Base.length(panel::AKMPanel) = length(panel.y) + +function parse_options(args) + options = Dict{String,String}() + + for arg in args + startswith(arg, "--") || error("Unexpected positional argument '$arg'. Use --key=value options.") + key_value = split(arg[3:end], "=", limit = 2) + length(key_value) == 2 || error("Expected --key=value, got $arg") + options[key_value[1]] = key_value[2] + end + + return options +end + +option(options, key, default::Int) = parse(Int, get(options, key, string(default))) +option(options, key, default::Float64) = parse(Float64, get(options, key, string(default))) + +function option(options, key, default::Bool) + value = lowercase(get(options, key, string(default))) + value in ("true", "yes", "1") && return true + value in ("false", "no", "0") && return false + error("Expected boolean for --$key, got '$value'") +end + +function method_option(options) + value = lowercase(get(options, "method", "cpu")) + value == "cpu" && return :cpu + value == "cuda" && return :CUDA + value == "metal" && return :Metal + error("Expected --method=cpu, --method=CUDA, or --method=Metal, got '$value'") +end + +function maxiter_option(options) + value = lowercase(get(options, "maxiter", "inf")) + value in ("inf", "infinity") && return typemax(Int) + return parse(Int, value) +end + +function load_backend!(method::Symbol) + if method == :CUDA + @eval using CUDA + CUDA.functional() || error("CUDA was requested but CUDA.functional() is false") + elseif method == :Metal + @eval using Metal + end + return nothing +end + +function local_firm(rng::AbstractRNG, anchor::Int) + lo = max(1, anchor - FIRM_WINDOW) + hi = min(N_FIRMS, anchor + FIRM_WINDOW) + return rand(rng, lo:hi) +end + +function simulate_akm_panel(; seed::Int) + rng = MersenneTwister(seed) + n = N_WORKERS * OBS_PER_WORKER + worker = Vector{Int32}(undef, n) + firm = Vector{Int32}(undef, n) + year = Vector{Int32}(undef, n) + y = Vector{Float64}(undef, n) + x = Matrix{Float64}(undef, n, length(CONTROL_NAMES)) + + worker_effect = randn(rng, N_WORKERS) + firm_effect = 0.6 .* randn(rng, N_FIRMS) + year_effect = [0.02 * (t - 1) + 0.05 * sin(2 * pi * (t - 1) / N_YEARS) for t in 1:N_YEARS] + + row = 1 + max_start_year = max(1, N_YEARS - min(OBS_PER_WORKER, N_YEARS) + 1) + for w in 1:N_WORKERS + anchor = clamp(1 + fld((w - 1) * N_FIRMS, N_WORKERS), 1, N_FIRMS) + current_firm = local_firm(rng, anchor) + start_year = rand(rng, 1:max_start_year) + base_experience = rand(rng, 1:20) + tenure = 0 + + for spell_t in 1:OBS_PER_WORKER + moved = spell_t > 1 && rand(rng) < MOVE_PROBABILITY + if moved + current_firm = local_firm(rng, anchor) + tenure = 0 + elseif spell_t > 1 + tenure += 1 + end + + calendar_year = 1 + mod(start_year + spell_t - 2, N_YEARS) + experience = base_experience + spell_t - 1 + tenure_sq = tenure^2 / 100 + experience_sq = experience^2 / 100 + mover = moved ? 1.0 : 0.0 + + worker[row] = Int32(w) + firm[row] = Int32(current_firm) + year[row] = Int32(calendar_year) + x[row, 1] = tenure + x[row, 2] = tenure_sq + x[row, 3] = experience + x[row, 4] = experience_sq + x[row, 5] = mover + y[row] = 0.04 * tenure - 0.03 * tenure_sq + + 0.015 * experience - 0.02 * experience_sq + + 0.05 * mover + + worker_effect[w] + firm_effect[current_firm] + + year_effect[calendar_year] + 0.2 * randn(rng) + row += 1 + end + end + + order = randperm(rng, n) + return AKMPanel(worker[order], firm[order], year[order], y[order], x[order, :]) +end + +function akm_estimator_call(panel::AKMPanel; + method::Symbol, + double_precision::Bool, + tol::Real, + maxiter::Integer) + y = copy(panel.y) + x = copy(panel.x) + fes = [FixedEffect(panel.worker), FixedEffect(panel.firm), FixedEffect(panel.year)] + T = double_precision ? Float64 : Float32 + solver = AbstractFixedEffectSolver{T}(fes, uweights(T, length(y)), Val{method}) + variables = Vector{AbstractVector{Float64}}(undef, 1 + size(x, 2)) + variables[1] = y + for j in axes(x, 2) + variables[j + 1] = view(x, :, j) + end + _, iterations, converged = solve_residuals!(variables, solver; + tol = tol, + maxiter = maxiter, + progress_bar = false) + beta = x \ y + return (beta = beta, iterations = iterations, converged = converged) +end + +function print_result(result) + @printf(" beta:") + for (name, value) in zip(CONTROL_NAMES, result.beta) + @printf(" %s=% .4f", name, value) + end + println() + println(" iterations: ", join(result.iterations, ", ")) + println(" converged: ", join(result.converged, ", ")) +end + +options = parse_options(ARGS) +seed = option(options, "seed", 1234) +method = method_option(options) +double_precision = option(options, "double-precision", method == :cpu) +tol = option(options, "tol", double_precision ? 1e-8 : 1e-6) +maxiter = maxiter_option(options) + +load_backend!(method) + +n = N_WORKERS * OBS_PER_WORKER +println("Simulated AKM benchmark") +println(" observations: ", n) +println(" worker FE: ", N_WORKERS) +println(" firm FE: ", N_FIRMS) +println(" year FE: ", N_YEARS) +println(" controls: ", join(CONTROL_NAMES, ", ")) +println(" method: ", method) +println(" double precision:", double_precision) +println(" tol: ", tol) +println(" maxiter: ", maxiter == typemax(Int) ? "Inf" : maxiter) +println(" note: this times FE construction, residualization of y and controls, and dense OLS; clustered SEs are not included.") + +panel = simulate_akm_panel(; seed = seed) +@printf(" panel memory: %.1f MiB\n", Base.summarysize(panel) / 2.0^20) + +println("\nWarmup") +warmup = akm_estimator_call(panel; method, double_precision, tol, maxiter) +print_result(warmup) + +println("\nTimed run") +GC.gc() +timed = @timed akm_estimator_call(panel; method, double_precision, tol, maxiter) +@printf(" time: %.3f s\n", timed.time) +@printf(" allocated: %.1f MiB\n", timed.bytes / 2.0^20) +@printf(" gc time: %.3f s\n", timed.gctime) +print_result(timed.value) diff --git a/benchmarks/bench_gather.jl b/benchmarks/bench_gather.jl deleted file mode 100644 index 6c3f83d..0000000 --- a/benchmarks/bench_gather.jl +++ /dev/null @@ -1,169 +0,0 @@ -using Random, BenchmarkTools, Base.Threads -println("Julia ", VERSION, " — ", nthreads(), " threads") -Random.seed!(1234) - -############################################################################## -# Current serial gather (baseline) -############################################################################## -function gather_serial!(fecoef, refs, α, y, cache) - @fastmath @inbounds @simd for i in eachindex(y) - fecoef[refs[i]] += α * y[i] * cache[i] - end -end - -############################################################################## -# Approach 1: CSC-style transposed gather -# Precompute a CSC structure: for each group k, store obs indices. -# Each group k can be processed independently → trivially parallel, zero conflicts. -############################################################################## -struct CSCIndex - offsets::Vector{Int} - indices::Vector{Int} -end - -function build_csc(refs::AbstractVector{<:Integer}, n::Int) - N = length(refs) - counts = zeros(Int, n) - @inbounds for i in 1:N - counts[refs[i]] += 1 - end - offsets = Vector{Int}(undef, n + 1) - offsets[1] = 1 - @inbounds for k in 1:n - offsets[k+1] = offsets[k] + counts[k] - end - indices = Vector{Int}(undef, N) - fill!(counts, 0) - @inbounds for i in 1:N - k = refs[i] - counts[k] += 1 - indices[offsets[k] + counts[k] - 1] = i - end - return CSCIndex(offsets, indices) -end - -function gather_csc_parallel!(fecoef::AbstractVector{T}, csc::CSCIndex, α, y, cache) where T - offsets, indices = csc.offsets, csc.indices - n = length(fecoef) - Threads.@threads for k in 1:n - s = zero(T) - @fastmath @inbounds for j in offsets[k]:(offsets[k+1]-1) - i = indices[j] - s += y[i] * cache[i] - end - @inbounds fecoef[k] += α * s - end -end - -function gather_csc_serial!(fecoef::AbstractVector{T}, csc::CSCIndex, α, y, cache) where T - offsets, indices = csc.offsets, csc.indices - n = length(fecoef) - for k in 1:n - s = zero(T) - @fastmath @inbounds for j in offsets[k]:(offsets[k+1]-1) - i = indices[j] - s += y[i] * cache[i] - end - @inbounds fecoef[k] += α * s - end -end - -############################################################################## -# Approach 2: Per-thread accumulators with manual chunking (@spawn) -############################################################################## -struct PerThreadBuffers{T} - buffers::Vector{Vector{T}} -end -PerThreadBuffers{T}(n::Int, nt::Int) where T = PerThreadBuffers([zeros(T, n) for _ in 1:nt]) - -function gather_perthread!(fecoef::AbstractVector{T}, refs, α, y, cache, ptb::PerThreadBuffers{T}) where T - nt = length(ptb.buffers) - N = length(y) - for buf in ptb.buffers - fill!(buf, zero(T)) - end - chunk = cld(N, nt) - @sync for t in 1:nt - Threads.@spawn begin - buf = ptb.buffers[t] - lo = (t-1)*chunk + 1 - hi = min(t*chunk, N) - @fastmath @inbounds for i in lo:hi - buf[refs[i]] += y[i] * cache[i] - end - end - end - @inbounds for buf in ptb.buffers - @simd for k in eachindex(fecoef) - fecoef[k] += α * buf[k] - end - end -end - -############################################################################## -# Benchmarks -############################################################################## -function run_bench(label, N, n_groups) - println("\n", "="^60) - println("$label: N=$N, n_groups=$n_groups (avg group size=$(N÷n_groups))") - println("="^60) - - refs = rand(1:n_groups, N) - y = rand(N) - cache = rand(N) - α = 1.0 - nt = nthreads() - - csc = build_csc(refs, n_groups) - ptb = PerThreadBuffers{Float64}(n_groups, nt) - - # Verify correctness - out_ref = zeros(n_groups) - gather_serial!(out_ref, refs, α, y, cache) - - for (name, fn!) in [ - ("CSC parallel", (out) -> gather_csc_parallel!(out, csc, α, y, cache)), - ("CSC serial", (out) -> gather_csc_serial!(out, csc, α, y, cache)), - ("Per-thread chunked", (out) -> gather_perthread!(out, refs, α, y, cache, ptb)), - ] - out_test = zeros(n_groups) - fn!(out_test) - if !isapprox(out_test, out_ref, rtol=1e-10) - println(" WARNING $name: INCORRECT (max diff = $(maximum(abs.(out_test .- out_ref))))") - end - end - - print(" Serial (baseline): ") - b0 = @benchmark gather_serial!(out, $refs, $α, $y, $cache) setup=(out=zeros($n_groups)) evals=1 samples=30 - show(stdout, MIME("text/plain"), b0); println() - - print(" CSC serial: ") - b1 = @benchmark gather_csc_serial!(out, $csc, $α, $y, $cache) setup=(out=zeros($n_groups)) evals=1 samples=30 - show(stdout, MIME("text/plain"), b1); println() - - print(" CSC parallel: ") - b2 = @benchmark gather_csc_parallel!(out, $csc, $α, $y, $cache) setup=(out=zeros($n_groups)) evals=1 samples=30 - show(stdout, MIME("text/plain"), b2); println() - - print(" Per-thread chunked: ") - b3 = @benchmark gather_perthread!(out, $refs, $α, $y, $cache, $ptb) setup=(out=zeros($n_groups)) evals=1 samples=30 - show(stdout, MIME("text/plain"), b3); println() - - t0 = median(b0).time - println("\n Speedups vs serial baseline:") - println(" CSC serial: $(round(t0/median(b1).time, digits=2))x") - println(" CSC parallel: $(round(t0/median(b2).time, digits=2))x") - println(" Per-thread chunked: $(round(t0/median(b3).time, digits=2))x") -end - -# Scenario 1: Few large groups (like year FE) -run_bench("Few large groups", 10_000_000, 100) - -# Scenario 2: Many medium groups -run_bench("Many medium groups", 10_000_000, 100_000) - -# Scenario 3: Many small groups (worker FE) -run_bench("Many small groups (worker FE)", 800_000, 400_000) - -# Scenario 4: Moderate groups (firm FE) -run_bench("Moderate groups (firm FE)", 800_000, 50_000) diff --git a/benchmark/bench_gather.jl b/benchmarks/gather_strategies.jl similarity index 100% rename from benchmark/bench_gather.jl rename to benchmarks/gather_strategies.jl diff --git a/benchmark/run.jl b/benchmarks/solve_backends.jl similarity index 100% rename from benchmark/run.jl rename to benchmarks/solve_backends.jl diff --git a/benchmarks/benchmark_CPU.jl b/benchmarks/solve_cpu.jl similarity index 100% rename from benchmarks/benchmark_CPU.jl rename to benchmarks/solve_cpu.jl diff --git a/benchmarks/worker_firm_spectrum.jl b/benchmarks/worker_firm_spectrum.jl new file mode 100644 index 0000000..7ebb955 --- /dev/null +++ b/benchmarks/worker_firm_spectrum.jl @@ -0,0 +1,85 @@ +# Measure the leading canonical correlations sigma_k between the two FE subspaces for the +# benchmark worker-firm data. After Jacobi scaling A'A = [[I,C],[C',I]] with sigma_k = svd(C); +# slow LSMR convergence is driven by sigma_k -> 1. This quantifies how many slow modes exist +# (=> how many deflation vectors k would be needed and the achievable iteration reduction). +using Random, LinearAlgebra +Random.seed!(1234) + +# --- benchmark hard scenario --- +N = 800_000; M = 400_000; O = 50_000 +refs1 = rand(1:M, N) # worker +refs2 = [rand(max(1, div(x, 8)-10):min(O, div(x, 8)+10)) for x in refs1] # firm +n1 = maximum(refs1); n2 = maximum(refs2) + +# Jacobi scales (unweighted, no interaction): scale[g] = 1/sqrt(group count) +function scales(refs, n) + c = zeros(Int, n); @inbounds for r in refs; c[r] += 1; end + s = zeros(n); @inbounds for g in 1:n; s[g] = c[g] > 0 ? 1 / sqrt(c[g]) : 0.0; end + s +end +s1 = scales(refs1, n1); s2 = scales(refs2, n2) + +# C v = gather_FE1(scatter_FE2(v)), maps firm-space (n2) -> worker-space (n1) +function Cmul!(out, v, refs1, refs2, s1, s2) + fill!(out, 0.0) + @inbounds for i in eachindex(refs1) + out[refs1[i]] += s2[refs2[i]] * v[refs2[i]] + end + @inbounds for g in eachindex(out); out[g] *= s1[g]; end + out +end + +# C' u, maps worker-space (n1) -> firm-space (n2) +function Ctmul!(out, u, refs1, refs2, s1, s2) + fill!(out, 0.0) + @inbounds for i in eachindex(refs1) + out[refs2[i]] += s1[refs1[i]] * u[refs1[i]] + end + @inbounds for h in eachindex(out); out[h] *= s2[h]; end + out +end + +# Subspace (block) iteration on C'C (acts on firm space n2=50k) for the top-k sigma^2. +function top_sigmas(k, iters) + Y = randn(n2, k) + tmp1 = zeros(n1); ritz = zeros(k) + for it in 1:iters + # Z = (C'C) Y + Z = similar(Y) + for j in 1:k + Cmul!(tmp1, view(Y, :, j), refs1, refs2, s1, s2) + Ctmul!(view(Z, :, j), tmp1, refs1, refs2, s1, s2) + end + F = qr(Z); Q = Matrix(F.Q) + # Rayleigh-Ritz on Q + AQ = similar(Q) + for j in 1:k + Cmul!(tmp1, view(Q, :, j), refs1, refs2, s1, s2) + Ctmul!(view(AQ, :, j), tmp1, refs1, refs2, s1, s2) + end + H = Symmetric(Q' * AQ) + E = eigen(H); ev = E.values; perm = sortperm(ev, rev = true) + ritz = ev[perm] + Y = Q * E.vectors[:, perm] + end + sqrt.(clamp.(ritz, 0, Inf)) +end + +println("worker-firm: N=$N, n1(worker)=$n1, n2(firm)=$n2") +k = 30 +sig = top_sigmas(k, 60) +println("\nTop $k canonical correlations sigma_k (descending):") +for (j, s) in enumerate(sig) + gap = 1 - s + println(" k=$(lpad(j, 2)) sigma=", round(s, digits = 6), " 1-sigma=", round(gap, sigdigits = 3)) +end + +# condition-number proxy and crude iteration estimate (iters ~ sqrt((1+s)/(1-s)) for sigma_max) +s2nd = sig[2] +println("\nsigma_1=", round(sig[1], digits = 6), " (constant mode, deflated by rank-deficiency)") +println("sigma_2=", round(s2nd, digits = 6), " => kappa~", round((1 + s2nd) / (1 - s2nd), digits = 1), + " sqrt(kappa)~", round(sqrt((1 + s2nd) / (1 - s2nd)), digits = 1)) +nbig = count(>(0.99), sig) +println("count(sigma>0.99) in top $k: ", nbig) +nbig999 = count(>(0.999), sig) +println("count(sigma>0.999) in top $k: ", nbig999) diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl index cfac1d9..1a8c3e4 100644 --- a/ext/CUDAExt.jl +++ b/ext/CUDAExt.jl @@ -1,119 +1,128 @@ module CUDAExt using FixedEffects, CUDA -using FixedEffects: FixedEffectCoefficients, AbstractWeights, UnitWeights, LinearAlgebra, Adjoint, mul!, rmul!, lsmr!, AbstractFixedEffectLinearMap, copy_internal!, AtomicGather +using FixedEffects: FixedEffectCoefficients, AbstractWeights, UnitWeights, LinearAlgebra, Adjoint, mul!, rmul!, AbstractFixedEffectLinearMap, copy_internal!, AbsorptionPlan, AbsorbedBlock, block_width CUDA.allowscalar(false) ############################################################################## ## -## Conversion FixedEffect between CPU and GPU +## CUDA backend — same layout as src/CPU.jl and ext/MetalExt.jl: +## 1. FixedEffectLinearMapCUDA: plan transfer, mul!, kernels; +## 2. FixedEffectSolverCUDA: solver storage and interface. +## The AbsorptionPlan (block transforms and whitened row values) is built on +## the CPU; refs and qrows are moved to the device and consumed by fused +## block kernels. ## ############################################################################## -# https://github.com/JuliaGPU/CUDA.jl/issues/142 -function _cu(T::Type, fe::FixedEffect) - refs = CuArray(fe.refs) - interaction = _cu(T, fe.interaction) - FixedEffect{typeof(refs), typeof(interaction)}(refs, interaction, fe.n) -end -_cu(T::Type, w::UnitWeights) = fill!(CuVector{T}(undef, length(w)), w[1]) -_cu(T::Type, w::AbstractVector) = CuVector{T}(convert(Vector{T}, w)) - ############################################################################## ## -## FixedEffectLinearMap on the GPU (code by Paul Schrimpf) -## -## Model matrix of categorical variables -## mutiplied by diag(1/sqrt(∑w * interaction^2, ..., ∑w * interaction^2) (Jacobi preconditoner) -## -## We define these methods used in lsmr! (duck typing): -## eltype -## size -## mul! +## 1. FixedEffectLinearMapCUDA ## ############################################################################## -mutable struct FixedEffectLinearMapCUDA{T} <: AbstractFixedEffectLinearMap{T} +## 1a) FixedEffectLinearMapCUDA Constructor + +_cu(T::Type, w::UnitWeights) = fill!(CuVector{T}(undef, length(w)), w[1]) +_cu(T::Type, w::AbstractVector) = CuVector{T}(convert(Vector{T}, w)) + +mutable struct FixedEffectLinearMapCUDA{T,P<:AbsorptionPlan} <: AbstractFixedEffectLinearMap{T} fes::Vector{<:FixedEffect} - scales::Vector{<:AbstractVector} - caches::Vector{<:AbstractVector} - gathers::Vector{AtomicGather} + plan::P end -function FixedEffectLinearMapCUDA{T}(fes::Vector{<:FixedEffect}) where {T} - fes = [_cu(T, fe) for fe in fes] - scales = [CUDA.zeros(T, fe.n) for fe in fes] - caches = [CUDA.zeros(T, length(fes[1].interaction)) for fe in fes] - gathers = [AtomicGather() for fe in fes] - return FixedEffectLinearMapCUDA{T}(fes, scales, caches, gathers) +function FixedEffectLinearMapCUDA{T}(fes::Vector{<:FixedEffect}, weights::AbstractWeights) where {T} + plan = _cu_plan(T, fes, weights) + return FixedEffectLinearMapCUDA{T,typeof(plan)}(fes, plan) end -function FixedEffects.gather!(fecoef::CuVector, refs::CuVector, α::Number, y::CuVector, cache::CuVector, ::AtomicGather) - nthreads = 256 - nblocks = cld(length(y), nthreads) - @cuda threads=nthreads blocks=nblocks gather_kernel!(fecoef, refs, α, y, cache) +function _cu_plan(::Type{T}, fes::Vector{<:FixedEffect}, weights::AbstractWeights) where {T} + cpu_plan = AbsorptionPlan(T, fes, weights) + blocks = [AbsorbedBlock(CuArray(block.refs), block.interactions, block.n, block.input_terms) + for block in cpu_plan.blocks] + qrows = [CuArray(q) for q in cpu_plan.qrows] + return AbsorptionPlan(blocks, cpu_plan.transforms, cpu_plan.ranks, qrows) end -function gather_kernel!(fecoef, refs, α, y, cache) - index = (blockIdx().x - Int32(1)) * blockDim().x + threadIdx().x - stride = blockDim().x * gridDim().x - i = index - @inbounds while i <= length(y) - CUDA.@atomic fecoef[refs[i]] += α * y[i] * cache[i] - i += stride - end -end +## 1b) FixedEffectLinearMapCUDA mul! -function FixedEffects.scatter!(y::CuVector, α::Number, fecoef::CuVector, refs::CuVector, cache::CuVector) - nthreads = 256 - nblocks = cld(length(y), nthreads) - @cuda threads=nthreads blocks=nblocks scatter_kernel!(y, α, fecoef, refs, cache) +## Implement right multiplication +function LinearAlgebra.mul!(y::CuVector, fem::FixedEffectLinearMapCUDA{T}, + fecoefs::FixedEffectCoefficients, α::Number, β::Number) where {T} + if iszero(β) + fill!(y, zero(T)) + β = one(β) + end + for (coef_block, block, qrows) in zip(fecoefs.x, fem.plan.blocks, fem.plan.qrows) + _scatter_block!(y, block.refs, qrows, coef_block, α, β) + β = one(β) + end + return y end -function FixedEffects.scatter!(y::CuVector, α::Number, fecoef::CuVector, refs::CuVector, cache::CuVector, β::Number) +function _scatter_block!(y::CuVector, refs::CuVector, qrows::CuMatrix, + coef_block::CuMatrix, α::Number, β::Number) nthreads = 256 nblocks = cld(length(y), nthreads) - if iszero(β) - @cuda threads=nthreads blocks=nblocks scatter_kernel_zero!(y, α, fecoef, refs, cache) - elseif isone(β) - @cuda threads=nthreads blocks=nblocks scatter_kernel!(y, α, fecoef, refs, cache) - else - @cuda threads=nthreads blocks=nblocks scatter_kernel_scaled!(y, α, fecoef, refs, cache, β) - end + @cuda threads=nthreads blocks=nblocks scatter_block_kernel!(y, refs, qrows, coef_block, α, β, size(coef_block, 1)) + return y end -function scatter_kernel!(y, α, fecoef, refs, cache) +function scatter_block_kernel!(y, refs, qrows, coef_block, α, β, k) index = (blockIdx().x - Int32(1)) * blockDim().x + threadIdx().x stride = blockDim().x * gridDim().x i = index @inbounds while i <= length(y) - y[i] += α * fecoef[refs[i]] * cache[i] + g = refs[i] + fit = zero(eltype(y)) + for c in 1:k + fit += coef_block[c, g] * qrows[c, i] + end + y[i] = β * y[i] + α * fit i += stride end + return nothing end -function scatter_kernel_zero!(y, α, fecoef, refs, cache) - index = (blockIdx().x - Int32(1)) * blockDim().x + threadIdx().x - stride = blockDim().x * gridDim().x - i = index - @inbounds while i <= length(y) - y[i] = α * fecoef[refs[i]] * cache[i] - i += stride +## 1c) FixedEffectLinearMapCUDA mul!, Adjoint + +## Implement left multiplication +function LinearAlgebra.mul!(fecoefs::FixedEffectCoefficients, + Cfem::Adjoint{T, <:FixedEffectLinearMapCUDA{T}}, + y::CuVector, α::Number, β::Number) where {T} + fem = adjoint(Cfem) + rmul!(fecoefs, β) + for (coef_block, block, qrows) in zip(fecoefs.x, fem.plan.blocks, fem.plan.qrows) + _gather_block!(coef_block, block.refs, qrows, y, α) end + return fecoefs +end + +function _gather_block!(coef_block::CuMatrix, refs::CuVector, qrows::CuMatrix, + y::CuVector, α::Number) + nthreads = 256 + nblocks = cld(length(y), nthreads) + @cuda threads=nthreads blocks=nblocks gather_block_kernel!(coef_block, refs, qrows, y, α, size(coef_block, 1)) + return coef_block end -function scatter_kernel_scaled!(y, α, fecoef, refs, cache, β) +function gather_block_kernel!(coef_block, refs, qrows, y, α, k) index = (blockIdx().x - Int32(1)) * blockDim().x + threadIdx().x stride = blockDim().x * gridDim().x i = index @inbounds while i <= length(y) - y[i] = β * y[i] + α * fecoef[refs[i]] * cache[i] + g = refs[i] + yi = α * y[i] + for c in 1:k + CUDA.@atomic coef_block[c, g] += yi * qrows[c, i] + end i += stride end + return nothing end ############################################################################## ## -## Implement AbstractFixedEffectSolver interface +## 2. FixedEffectSolverCUDA ## ############################################################################## @@ -122,72 +131,31 @@ mutable struct FixedEffectSolverCUDA{T} <: FixedEffects.AbstractFixedEffectSolve weights::CuVector{T} b::CuVector{T} r::CuVector{T} - x::FixedEffectCoefficients{<: AbstractVector{T}} - v::FixedEffectCoefficients{<: AbstractVector{T}} - h::FixedEffectCoefficients{<: AbstractVector{T}} - hbar::FixedEffectCoefficients{<: AbstractVector{T}} + x::FixedEffectCoefficients + v::FixedEffectCoefficients + h::FixedEffectCoefficients + hbar::FixedEffectCoefficients tmp::Vector{T} # used to convert AbstractVector to Vector{T} - fes::Vector{<:FixedEffect} end function FixedEffects.AbstractFixedEffectSolver{T}(fes::Vector{<:FixedEffect}, weights::AbstractWeights, ::Type{Val{:CUDA}}) where {T} - m = FixedEffectLinearMapCUDA{T}(fes) + m = FixedEffectLinearMapCUDA{T}(fes, weights) b = CUDA.zeros(T, length(weights)) r = CUDA.zeros(T, length(weights)) - x = FixedEffectCoefficients([CUDA.zeros(T, fe.n) for fe in fes]) - v = FixedEffectCoefficients([CUDA.zeros(T, fe.n) for fe in fes]) - h = FixedEffectCoefficients([CUDA.zeros(T, fe.n) for fe in fes]) - hbar = FixedEffectCoefficients([CUDA.zeros(T, fe.n) for fe in fes]) + x = FixedEffectCoefficients([CUDA.zeros(T, block_width(block), block.n) for block in m.plan.blocks]) + v = FixedEffectCoefficients([CUDA.zeros(T, block_width(block), block.n) for block in m.plan.blocks]) + h = FixedEffectCoefficients([CUDA.zeros(T, block_width(block), block.n) for block in m.plan.blocks]) + hbar = FixedEffectCoefficients([CUDA.zeros(T, block_width(block), block.n) for block in m.plan.blocks]) tmp = zeros(T, length(weights)) - feM = FixedEffectSolverCUDA{T}(m, CUDA.zeros(T, length(weights)), b, r, x, v, h, hbar, tmp, fes) - FixedEffects.update_weights!(feM, weights) + return FixedEffectSolverCUDA{T}(m, _cu(T, weights), b, r, x, v, h, hbar, tmp) end function FixedEffects.update_weights!(feM::FixedEffectSolverCUDA{T}, weights::AbstractWeights) where {T} copyto!(feM.weights, _cu(T, weights)) - for (scale, fe) in zip(feM.m.scales, feM.m.fes) - scale!(scale, fe.refs, fe.interaction, feM.weights) - end - for (cache, scale, fe) in zip(feM.m.caches, feM.m.scales, feM.m.fes) - cache!(cache, fe.refs, fe.interaction, feM.weights, scale) - end + feM.m.plan = _cu_plan(T, feM.m.fes, weights) return feM end -function scale!(scale::CuVector, refs::CuVector, interaction::CuVector, weights::CuVector) - nthreads = 256 - nblocks = cld(length(refs), nthreads) - fill!(scale, 0) - @cuda threads=nthreads blocks=nblocks scale_kernel!(scale, refs, interaction, weights) - map!(x -> x > 0 ? 1 / sqrt(x) : zero(eltype(scale)), scale, scale) -end - -function scale_kernel!(scale, refs, interaction, weights) - index = (blockIdx().x - Int32(1)) * blockDim().x + threadIdx().x - stride = blockDim().x * gridDim().x - i = index - @inbounds while i <= length(interaction) - CUDA.@atomic scale[refs[i]] += abs2(interaction[i]) * weights[i] - i += stride - end -end - -function cache!(cache::CuVector, refs::CuVector, interaction::CuVector, weights::CuVector, scale::CuVector) - nthreads = 256 - nblocks = cld(length(cache), nthreads) - @cuda threads=nthreads blocks=nblocks cache!_kernel!(cache, refs, interaction, weights, scale) -end - -function cache!_kernel!(cache, refs, interaction, weights, scale) - index = (blockIdx().x - Int32(1)) * blockDim().x + threadIdx().x - stride = blockDim().x * gridDim().x - i = index - @inbounds while i <= length(cache) - cache[i] = interaction[i] * sqrt(weights[i]) * scale[refs[i]] - i += stride - end -end - function FixedEffects.copy_internal!(feM::FixedEffectSolverCUDA, field::Symbol, r::AbstractVector) copyto!(feM.tmp, r) copyto!(getfield(feM, field), feM.tmp) diff --git a/ext/MetalExt.jl b/ext/MetalExt.jl index c4a7225..3f99ff7 100644 --- a/ext/MetalExt.jl +++ b/ext/MetalExt.jl @@ -1,48 +1,75 @@ module MetalExt using FixedEffects, Metal -using FixedEffects: FixedEffectCoefficients, AbstractWeights, UnitWeights, LinearAlgebra, Adjoint, mul!, rmul!, lsmr!, AbstractFixedEffectLinearMap, copy_internal!, AtomicGather, BucketGather +using FixedEffects: FixedEffectCoefficients, AbstractWeights, UnitWeights, LinearAlgebra, Adjoint, mul!, rmul!, AbstractFixedEffectLinearMap, copy_internal!, AbsorptionPlan, AbsorbedBlock, block_width Metal.allowscalar(false) ############################################################################## ## -## Conversion FixedEffect between CPU and Metal +## Metal backend — same layout as src/CPU.jl and ext/CUDAExt.jl: +## 1. FixedEffectLinearMapMetal: gather strategies, plan transfer, mul!, kernels; +## 2. FixedEffectSolverMetal: solver storage and interface. +## The AbsorptionPlan (block transforms and whitened row values) is built on +## the CPU; refs and qrows are moved to the device and consumed by fused +## block kernels. ## ############################################################################## -function _mtl(T::Type, fe::FixedEffect) - refs = MtlArray(fe.refs) - interaction = _mtl(T, fe.interaction) - FixedEffect{typeof(refs), typeof(interaction)}(refs, interaction, fe.n) -end -_mtl(T::Type, w::UnitWeights) = Metal.ones(T, length(w)) -_mtl(T::Type, w::AbstractVector) = MtlVector{T}(convert(Vector{T}, w)) - ############################################################################## ## -## FixedEffectLinearMap on Metal -## -## Model matrix of categorical variables -## mutiplied by diag(1/sqrt(∑w * interaction^2, ..., ∑w * interaction^2) (Jacobi preconditoner) -## -## We define these methods used in lsmr! (duck typing): -## eltype -## size -## mul! +## 1. FixedEffectLinearMapMetal ## ############################################################################## -mutable struct FixedEffectLinearMapMetal{T} <: AbstractFixedEffectLinearMap{T} - fes::Vector{<:FixedEffect} - scales::Vector{MtlVector{T}} - caches::Vector{MtlVector{T}} - gathers::Vector{Union{AtomicGather, BucketGather}} -end +## 1a) FixedEffectLinearMapMetal Constructor + +_mtl(T::Type, w::UnitWeights) = Metal.ones(T, length(w)) +_mtl(T::Type, w::AbstractVector) = MtlVector{T}(convert(Vector{T}, w)) function _metal_threadgroup_width() width = Int(device().maxThreadsPerThreadgroup.width) return prevpow(2, width) end +# Per-block plan for the adjoint gather (A'u), chosen once at construction: +# bucketize (one threadgroup per group) for low cardinality, else atomic adds. +struct AtomicGather end +struct BucketGather{V<:AbstractVector} + perm::V # observation indices sorted by group + offsets::V # CSR offsets into perm (length ngroups + 1) +end + +mutable struct FixedEffectLinearMapMetal{T,P<:AbsorptionPlan} <: AbstractFixedEffectLinearMap{T} + fes::Vector{<:FixedEffect} + plan::P + gathers::Vector{Union{AtomicGather, BucketGather}} +end + +function FixedEffectLinearMapMetal{T}(fes::Vector{<:FixedEffect}, weights::AbstractWeights) where {T} + plan = _mtl_plan(T, fes, weights) + G = Union{AtomicGather, BucketGather} + gathers = Vector{G}(undef, length(plan.blocks)) + Threads.@threads for i in 1:length(plan.blocks) + refs = fes[plan.blocks[i].input_terms[1]].refs + n = plan.blocks[i].n + # bucketize (one threadgroup per group) for low cardinality; else atomic adds + if n < min(100_000, div(length(refs), 16)) + perm, offsets = bucketize_refs(refs, n) + gathers[i] = BucketGather(perm, offsets) + else + gathers[i] = AtomicGather() + end + end + return FixedEffectLinearMapMetal{T,typeof(plan)}(fes, plan, gathers) +end + +function _mtl_plan(::Type{T}, fes::Vector{<:FixedEffect}, weights::AbstractWeights) where {T} + cpu_plan = AbsorptionPlan(T, fes, weights) + blocks = [AbsorbedBlock(MtlArray(block.refs), block.interactions, block.n, block.input_terms) + for block in cpu_plan.blocks] + qrows = [MtlArray(q) for q in cpu_plan.qrows] + return AbsorptionPlan(blocks, cpu_plan.transforms, cpu_plan.ranks, qrows) +end + function bucketize_refs(refs::AbstractVector{<:Integer}, n::Int) # count the number of obs per group counts = zeros(Int, n) @@ -67,136 +94,132 @@ function bucketize_refs(refs::AbstractVector{<:Integer}, n::Int) return MtlVector{Int}(perm), MtlVector{Int}(offsets) end -function FixedEffectLinearMapMetal{T}(fes::Vector{<:FixedEffect}) where {T} - fes2 = [_mtl(T, fe) for fe in fes] - scales = [Metal.zeros(T, fe.n) for fe in fes] - caches = [Metal.zeros(T, length(fe.refs)) for fe in fes] - G = Union{AtomicGather, BucketGather} - gathers = Vector{G}(undef, length(fes)) - Threads.@threads for i in 1:length(fes) - refs = fes[i].refs - n = fes[i].n - # bucketize (one threadgroup per group) for low cardinality; else atomic adds - if n < min(100_000, div(length(refs), 16)) - perm, offsets = bucketize_refs(refs, n) - gathers[i] = BucketGather(perm, offsets) - else - gathers[i] = AtomicGather() - end - end - return FixedEffectLinearMapMetal{T}(fes2, scales, caches, gathers) -end +## 1b) FixedEffectLinearMapMetal mul! -function FixedEffects.gather!(fecoef::MtlVector, refs::MtlVector, α::Number, y::MtlVector, cache::MtlVector, g::BucketGather) - n = length(fecoef) - nthreads = _metal_threadgroup_width() - Metal.@sync @metal threads=nthreads groups=n gather_kernel_bin!(fecoef, refs, α, y, cache, g.perm, g.offsets, Val(nthreads)) +## Implement right multiplication +function LinearAlgebra.mul!(y::MtlVector, fem::FixedEffectLinearMapMetal{T}, + fecoefs::FixedEffectCoefficients, α::Number, β::Number) where {T} + if iszero(β) + fill!(y, zero(T)) + β = one(β) + end + for (coef_block, block, qrows) in zip(fecoefs.x, fem.plan.blocks, fem.plan.qrows) + _scatter_block!(y, block.refs, qrows, coef_block, α, β) + β = one(β) + end + return y end -function FixedEffects.gather!(fecoef::MtlVector, refs::MtlVector, α::Number, y::MtlVector, cache::MtlVector, ::AtomicGather) +function _scatter_block!(y::MtlVector, refs::MtlVector, qrows::MtlMatrix, + coef_block::MtlMatrix, α::Number, β::Number) nthreads = _metal_threadgroup_width() nblocks = cld(length(y), nthreads) - Metal.@sync @metal threads=nthreads groups=nblocks gather_kernel!(fecoef, refs, α, y, cache) + Metal.@sync @metal threads=nthreads groups=nblocks scatter_block_kernel!(y, refs, qrows, coef_block, α, β, size(coef_block, 1)) + return y end -function gather_kernel_bin!(fecoef, refs, α, y, cache, perm, offsets, ::Val{NT}) where {NT} - k = Int(threadgroup_position_in_grid().x) - tid = Int(thread_position_in_threadgroup().x) - nt = Int(threads_per_threadgroup().x) - - # threadgroup scratch - T = eltype(fecoef) - shared = Metal.MtlThreadGroupArray(T, NT) - - start = @inbounds offsets[k] - stop = @inbounds offsets[k+1] - 1 - - acc = zero(T) - - # each thread walks its portion of the bucket - j = start + tid - 1 - while j <= stop - i = @inbounds perm[j] - @inbounds acc += (α * y[i] * cache[i]) - j += nt - end - - @inbounds shared[tid] = acc - Metal.threadgroup_barrier(Metal.MemoryFlagThreadGroup) - - # tree reduction in shared memory - offset = nt ÷ 2 - while offset > 0 - if tid <= offset - @inbounds shared[tid] += shared[tid + offset] - end - Metal.threadgroup_barrier(Metal.MemoryFlagThreadGroup) - offset ÷= 2 - end - - # one write per coefficient (no atomics needed if groups == n and 1 group per k) - if tid == 1 - @inbounds fecoef[k] += shared[1] - end - - return nothing -end - -function gather_kernel!(fecoef, refs, α, y, cache) +function scatter_block_kernel!(y, refs, qrows, coef_block, α, β, k) i = thread_position_in_grid_1d() - if i <= length(refs) - @inbounds Metal.atomic_fetch_add_explicit(pointer(fecoef, refs[i]), α * y[i] * cache[i]) + if i <= length(y) + @inbounds begin + g = refs[i] + fit = zero(eltype(y)) + for c in 1:k + fit += coef_block[c, g] * qrows[c, i] + end + y[i] = β * y[i] + α * fit + end end return nothing end -function FixedEffects.scatter!(y::MtlVector, α::Number, fecoef::MtlVector, refs::MtlVector, cache::MtlVector) +## 1c) FixedEffectLinearMapMetal mul!, Adjoint + +## Implement left multiplication +function LinearAlgebra.mul!(fecoefs::FixedEffectCoefficients, + Cfem::Adjoint{T, <:FixedEffectLinearMapMetal{T}}, + y::MtlVector, α::Number, β::Number) where {T} + fem = adjoint(Cfem) + rmul!(fecoefs, β) + for (coef_block, block, qrows, gather) in zip(fecoefs.x, fem.plan.blocks, fem.plan.qrows, fem.gathers) + _gather_block!(coef_block, block.refs, qrows, y, α, gather) + end + return fecoefs +end + +function _gather_block!(coef_block::MtlMatrix, refs::MtlVector, qrows::MtlMatrix, + y::MtlVector, α::Number, gather::BucketGather) + n = size(coef_block, 2) nthreads = _metal_threadgroup_width() - nblocks = cld(length(y), nthreads) - Metal.@sync @metal threads=nthreads groups=nblocks scatter_kernel!(y, α, fecoef, refs, cache) + Metal.@sync @metal threads=nthreads groups=n gather_block_kernel_bin!(coef_block, α, y, qrows, gather.perm, gather.offsets, Val(nthreads), size(coef_block, 1)) + return coef_block end -function FixedEffects.scatter!(y::MtlVector, α::Number, fecoef::MtlVector, refs::MtlVector, cache::MtlVector, β::Number) +function _gather_block!(coef_block::MtlMatrix, refs::MtlVector, qrows::MtlMatrix, + y::MtlVector, α::Number, ::AtomicGather) nthreads = _metal_threadgroup_width() nblocks = cld(length(y), nthreads) - if iszero(β) - Metal.@sync @metal threads=nthreads groups=nblocks scatter_kernel_zero!(y, α, fecoef, refs, cache) - elseif isone(β) - Metal.@sync @metal threads=nthreads groups=nblocks scatter_kernel!(y, α, fecoef, refs, cache) - else - Metal.@sync @metal threads=nthreads groups=nblocks scatter_kernel_scaled!(y, α, fecoef, refs, cache, β) - end + Metal.@sync @metal threads=nthreads groups=nblocks gather_block_kernel!(coef_block, refs, α, y, qrows, size(coef_block, 1)) + return coef_block end -function scatter_kernel!(y, α, fecoef, refs, cache) - i = thread_position_in_grid_1d() - if i <= length(y) - @inbounds y[i] += α * fecoef[refs[i]] * cache[i] - end - return nothing -end +function gather_block_kernel_bin!(coef_block, α, y, qrows, perm, offsets, ::Val{NT}, k) where {NT} + g = Int(threadgroup_position_in_grid().x) + tid = Int(thread_position_in_threadgroup().x) + nt = Int(threads_per_threadgroup().x) + T = eltype(coef_block) + shared = Metal.MtlThreadGroupArray(T, NT) + + start = @inbounds offsets[g] + stop = @inbounds offsets[g + 1] - 1 + + for c in 1:k + acc = zero(T) + j = start + tid - 1 + while j <= stop + i = @inbounds perm[j] + @inbounds acc += α * y[i] * qrows[c, i] + j += nt + end -function scatter_kernel_zero!(y, α, fecoef, refs, cache) - i = thread_position_in_grid_1d() - if i <= length(y) - @inbounds y[i] = α * fecoef[refs[i]] * cache[i] + @inbounds shared[tid] = acc + Metal.threadgroup_barrier(Metal.MemoryFlagThreadGroup) + + offset = nt ÷ 2 + while offset > 0 + if tid <= offset + @inbounds shared[tid] += shared[tid + offset] + end + Metal.threadgroup_barrier(Metal.MemoryFlagThreadGroup) + offset ÷= 2 + end + + if tid == 1 + @inbounds coef_block[c, g] += shared[1] + end end + return nothing end -function scatter_kernel_scaled!(y, α, fecoef, refs, cache, β) +function gather_block_kernel!(coef_block, refs, α, y, qrows, k) i = thread_position_in_grid_1d() if i <= length(y) - @inbounds y[i] = β * y[i] + α * fecoef[refs[i]] * cache[i] + @inbounds begin + g = refs[i] + yi = α * y[i] + for c in 1:k + idx = c + (g - 1) * k + Metal.atomic_fetch_add_explicit(pointer(coef_block, idx), yi * qrows[c, i]) + end + end end return nothing end - - ############################################################################## ## -## Implement AbstractFixedEffectSolver interface +## 2. FixedEffectSolverMetal ## ############################################################################## @@ -205,77 +228,37 @@ mutable struct FixedEffectSolverMetal{T} <: FixedEffects.AbstractFixedEffectSolv weights::MtlVector{T} b::MtlVector{T} r::MtlVector{T} - x::FixedEffectCoefficients{<: AbstractVector{T}} - v::FixedEffectCoefficients{<: AbstractVector{T}} - h::FixedEffectCoefficients{<: AbstractVector{T}} - hbar::FixedEffectCoefficients{<: AbstractVector{T}} + x::FixedEffectCoefficients + v::FixedEffectCoefficients + h::FixedEffectCoefficients + hbar::FixedEffectCoefficients tmp::Vector{T} - fes::Vector{<:FixedEffect} end - + function FixedEffects.AbstractFixedEffectSolver{T}(fes::Vector{<:FixedEffect}, weights::AbstractWeights, ::Type{Val{:Metal}}) where {T} T === Float32 || throw(ArgumentError("The Metal backend supports Float32 solves only; pass double_precision=false or use method=:cpu for Float64.")) - m = FixedEffectLinearMapMetal{T}(fes) + m = FixedEffectLinearMapMetal{T}(fes, weights) b = Metal.zeros(T, length(weights)) r = Metal.zeros(T, length(weights)) - x = FixedEffectCoefficients([Metal.zeros(T, fe.n) for fe in fes]) - v = FixedEffectCoefficients([Metal.zeros(T, fe.n) for fe in fes]) - h = FixedEffectCoefficients([Metal.zeros(T, fe.n) for fe in fes]) - hbar = FixedEffectCoefficients([Metal.zeros(T, fe.n) for fe in fes]) + x = FixedEffectCoefficients([Metal.zeros(T, block_width(block), block.n) for block in m.plan.blocks]) + v = FixedEffectCoefficients([Metal.zeros(T, block_width(block), block.n) for block in m.plan.blocks]) + h = FixedEffectCoefficients([Metal.zeros(T, block_width(block), block.n) for block in m.plan.blocks]) + hbar = FixedEffectCoefficients([Metal.zeros(T, block_width(block), block.n) for block in m.plan.blocks]) tmp = zeros(T, length(weights)) - feM = FixedEffectSolverMetal{T}(m, Metal.zeros(T, length(weights)), b, r, x, v, h, hbar, tmp, fes) - FixedEffects.update_weights!(feM, weights) + return FixedEffectSolverMetal{T}(m, _mtl(T, weights), b, r, x, v, h, hbar, tmp) end function FixedEffects.update_weights!(feM::FixedEffectSolverMetal{T}, weights::AbstractWeights) where {T} copyto!(feM.weights, _mtl(T, weights)) - for (scale, fe) in zip(feM.m.scales, feM.m.fes) - scale!(scale, fe.refs, fe.interaction, feM.weights) - end - for (cache, scale, fe) in zip(feM.m.caches, feM.m.scales, feM.m.fes) - cache!(cache, fe.refs, fe.interaction, feM.weights, scale) - end + feM.m.plan = _mtl_plan(T, feM.m.fes, weights) return feM end -function scale!(scale::MtlVector, refs::MtlVector, interaction::MtlVector, weights::MtlVector) - nthreads = _metal_threadgroup_width() - nblocks = cld(length(refs), nthreads) - fill!(scale, 0) - Metal.@sync @metal threads=nthreads groups=nblocks scale_kernel!(scale, refs, interaction, weights) - Metal.@sync @metal threads=nthreads groups=nblocks inv_kernel!(scale, eltype(scale)) -end - -function scale_kernel!(scale, refs, interaction, weights) - i = thread_position_in_grid_1d() - if i <= length(refs) - @inbounds Metal.atomic_fetch_add_explicit(pointer(scale, refs[i]), interaction[i]^2 * weights[i]) - end - return nothing -end - -function inv_kernel!(scale, T) - i = thread_position_in_grid_1d() - if i <= length(scale) - @inbounds scale[i] = (scale[i] > 0) ? (1 / sqrt(scale[i])) : zero(T) - end - return nothing -end - -function cache!(cache::MtlVector, refs::MtlVector, interaction::MtlVector, weights::MtlVector, scale::MtlVector) - nthreads = _metal_threadgroup_width() - nblocks = cld(length(cache), nthreads) - Metal.@sync @metal threads=nthreads groups=nblocks cache!_kernel!(cache, refs, interaction, weights, scale) -end - -function cache!_kernel!(cache, refs, interaction, weights, scale) - i = thread_position_in_grid_1d() - if i <= length(cache) - @inbounds cache[i] = interaction[i] * sqrt(weights[i]) * scale[refs[i]] - end - return nothing +function FixedEffects.recover_coefficients(feM::FixedEffectSolverMetal{T}, ::Type{Tout}) where {T,Tout} + synchronize() + return FixedEffects.recover_coefficients(T, feM.m.fes, feM.m.plan, Matrix{T}[Array(x) for x in feM.x.x], Tout) end function FixedEffects.copy_internal!(feM::FixedEffectSolverMetal{T}, field::Symbol, r::AbstractVector) where {T} diff --git a/src/AbsorptionPlan.jl b/src/AbsorptionPlan.jl new file mode 100644 index 0000000..7cfc01d --- /dev/null +++ b/src/AbsorptionPlan.jl @@ -0,0 +1,253 @@ +############################################################################## +## +## AbsorptionPlan +## +############################################################################## + +## 1a) Types + +# One physical operator unit: all fixed-effect terms that share the same refs, +# e.g. fe(id) + fe(id)&x is one block of width 2 with interactions (1, x). +struct AbsorbedBlock{R<:AbstractVector{<:Integer},I<:Tuple} + refs::R + interactions::I # one column per term: 1, x, x^2, ... + n::Int # number of groups (= fe.n of the input terms) + input_terms::Vector{Int} # indices into the original fes vector, one per column +end + +""" +Whitened representation of a set of fixed-effect terms, grouped into +`AbsorbedBlock`s of terms sharing the same refs. + +For block `j` with columns `Z` (the interactions) and group `g`, let +`G_g = Z_g' W Z_g` be the weighted Gram matrix. The plan stores a rank-revealing +transform `R_g = transforms[j][:, 1:ranks[j][g], g]` with `R_g' G_g R_g = I`, +and the whitened rows `qrows[j][c, i] = sqrt(w_i) z_i' R_g[:, c]`, whose +columns are orthonormal within each group. + +The solvers apply the whitened operator `A = W^(1/2) Z R`: column scaling +(`R` and `sqrt(w)`) is baked into `qrows` here, while `solve_residuals!` and +`solve_coefficients!` scale the RHS by `sqrt(w)` at solve time — both scalings +are needed, and neither may be applied twice. Whitened coefficients transform +back to original coordinates via `β_g = R_g θ_g` (see `recover_coefficients`). +For a single scalar fixed effect this reduces to the usual Jacobi (diagonal) +preconditioner. +""" +struct AbsorptionPlan{B<:AbstractVector{<:AbsorbedBlock},TR<:AbstractVector,RA<:AbstractVector,RV<:AbstractVector} + blocks::B + transforms::TR # per block: k × k × n, R_g in columns 1:ranks[g] + ranks::RA # per block: rank of each group's Gram matrix + qrows::RV # per block: k × nobs whitened rows sqrt(w) Z R +end + +block_width(block::AbsorbedBlock) = length(block.interactions) +_ncoef(plan::AbsorptionPlan) = sum(block_width(block) * block.n for block in plan.blocks) + +## 1b) Constructors + +function AbsorptionPlan(::Type{T}, fes::Vector{<:FixedEffect}, weights::AbstractVector; + ranktol::Union{Nothing,Real} = nothing) where {T} + blocks = _build_absorbed_blocks(fes) + transforms, ranks, qrows = _build_transforms(T, blocks, weights, ranktol) + return AbsorptionPlan(blocks, transforms, ranks, qrows) +end + +# Rebuild for new weights, reusing the block structure (which does not depend on weights). +function AbsorptionPlan(::Type{T}, plan::AbsorptionPlan, weights::AbstractVector; + ranktol::Union{Nothing,Real} = nothing) where {T} + transforms, ranks, qrows = _build_transforms(T, plan.blocks, weights, ranktol) + return AbsorptionPlan(plan.blocks, transforms, ranks, qrows) +end + +function _build_absorbed_blocks(fes::Vector{<:FixedEffect}) + blocks = AbsorbedBlock[] + for (j, fe) in enumerate(fes) + block_id = findfirst(block -> block.n == fe.n && block.refs == fe.refs, blocks) + if block_id === nothing + push!(blocks, AbsorbedBlock(fe.refs, (fe.interaction,), fe.n, [j])) + else + block = blocks[block_id] + interactions = (block.interactions..., fe.interaction) + input_terms = copy(block.input_terms) + push!(input_terms, j) + blocks[block_id] = AbsorbedBlock(block.refs, interactions, block.n, input_terms) + end + end + return blocks +end + +function _build_transforms(::Type{T}, blocks::AbstractVector{<:AbsorbedBlock}, + weights::AbstractVector, ranktol::Union{Nothing,Real}) where {T} + transforms = Vector{Array{T,3}}(undef, length(blocks)) + ranks = Vector{Vector{Int}}(undef, length(blocks)) + qrows = Vector{Matrix{T}}(undef, length(blocks)) + for j in eachindex(blocks) + transforms[j], ranks[j], qrows[j] = _build_block_transform(T, blocks[j], weights, ranktol) + end + return transforms, ranks, qrows +end + +## 1c) Per-block transform build (rank-revealing Gram-Schmidt) + +function _build_block_transform(::Type{T}, block::AbsorbedBlock, weights::AbstractVector, + ranktol::Union{Nothing,Real}) where {T} + k = block_width(block) + nlevels = block.n + nobs = length(block.refs) + transforms = zeros(T, k, k, nlevels) + ranks = zeros(Int, nlevels) + qrows = zeros(T, k, nobs) + if k == 1 + interaction = block.interactions[1] + gram = zeros(T, nlevels) + @inbounds for i in eachindex(block.refs) + g = block.refs[i] + z = T(interaction[i]) + gram[g] += T(weights[i]) * abs2(z) + end + @inbounds for g in 1:nlevels + if gram[g] > zero(T) + transforms[1, 1, g] = inv(sqrt(gram[g])) + ranks[g] = 1 + end + end + @spawn_for_chunks 100_000 for i in eachindex(block.refs) + @inbounds begin + g = block.refs[i] + qrows[1, i] = sqrt(T(weights[i])) * T(interaction[i]) * transforms[1, 1, g] + end + end + return transforms, ranks, qrows + end + + counts, offsets, perm = _group_permutation(block.refs, nlevels) + maxrows = isempty(counts) ? 0 : maximum(counts) + if nthreads() > 1 && nobs >= 100_000 + nchunks = max(1, min(nthreads(), nlevels)) + else + nchunks = 1 + end + chunks = _row_chunks(nlevels, nchunks) + if nchunks == 1 + _build_block_transform_chunk!(transforms, ranks, qrows, block, weights, + ranktol, counts, offsets, perm, chunks[1], maxrows) + else + # Groups are disjoint row segments of perm, so chunks can be processed in parallel. + @sync for chunk in chunks + let chunk = chunk + Base.Threads.@spawn _build_block_transform_chunk!(transforms, ranks, qrows, + block, weights, ranktol, counts, offsets, perm, chunk, maxrows) + end + end + end + return transforms, ranks, qrows +end + +function _build_block_transform_chunk!(transforms::AbstractArray{T,3}, ranks::AbstractVector{Int}, + qrows::AbstractMatrix{T}, block::AbsorbedBlock, weights::AbstractVector, + ranktol::Union{Nothing,Real}, counts::AbstractVector{Int}, offsets::AbstractVector{Int}, + perm::AbstractVector{Int}, chunk, maxrows::Int) where {T} + k = block_width(block) + tmp = Vector{T}(undef, maxrows) + colnorms = Vector{T}(undef, k) + coef = zeros(T, k) + for g in chunk + firstrow = offsets[g] + lastrow = offsets[g + 1] - 1 + nrows = counts[g] + nrows > 0 || continue + tol = ranktol === nothing ? T(k) * sqrt(eps(T)) : T(ranktol) + @inbounds for a in 1:k + s = zero(T) + zvals = block.interactions[a] + for p in firstrow:lastrow + i = perm[p] + z = T(zvals[i]) + s += T(weights[i]) * abs2(z) + end + colnorms[a] = sqrt(s) + end + rank = 0 + for a in 1:k + colnorm = colnorms[a] + colnorm > zero(T) || continue + invcolnorm = inv(colnorm) + fill!(coef, zero(T)) + coef[a] = invcolnorm + zvals = block.interactions[a] + @inbounds for (localrow, p) in enumerate(firstrow:lastrow) + i = perm[p] + tmp[localrow] = sqrt(T(weights[i])) * T(zvals[i]) * invcolnorm + end + for _ in 1:2 + for c in 1:rank + h = zero(T) + @inbounds for (localrow, p) in enumerate(firstrow:lastrow) + i = perm[p] + h += qrows[c, i] * tmp[localrow] + end + iszero(h) && continue + @inbounds for (localrow, p) in enumerate(firstrow:lastrow) + i = perm[p] + tmp[localrow] -= h * qrows[c, i] + end + @inbounds for b in 1:k + coef[b] -= h * transforms[b, c, g] + end + end + end + nrm2 = zero(T) + @inbounds for localrow in 1:nrows + nrm2 += abs2(tmp[localrow]) + end + nrm = sqrt(nrm2) + if nrm > tol + rank += 1 + invnrm = inv(nrm) + @inbounds for (localrow, p) in enumerate(firstrow:lastrow) + i = perm[p] + qrows[rank, i] = tmp[localrow] * invnrm + end + @inbounds for b in 1:k + transforms[b, rank, g] = coef[b] * invnrm + end + end + end + ranks[g] = rank + end + return nothing +end + +function _group_permutation(refs::AbstractVector, nlevels::Integer) + counts = zeros(Int, nlevels) + @inbounds for i in eachindex(refs) + counts[refs[i]] += 1 + end + offsets = Vector{Int}(undef, nlevels + 1) + offsets[1] = 1 + @inbounds for g in 1:nlevels + offsets[g + 1] = offsets[g] + counts[g] + end + cursor = copy(offsets) + perm = Vector{Int}(undef, length(refs)) + @inbounds for i in eachindex(refs) + g = refs[i] + p = cursor[g] + perm[p] = i + cursor[g] = p + 1 + end + return counts, offsets, perm +end + +function _row_chunks(n::Int, k::Int) + base, rem = divrem(n, k) + out = Vector{UnitRange{Int}}(undef, k) + s = 1 + for t in 1:k + len = base + (t <= rem ? 1 : 0) + out[t] = s:(s + len - 1) + s += len + end + return out +end + diff --git a/src/AbstractFixedEffectLinearMap.jl b/src/AbstractFixedEffectLinearMap.jl index 9f19875..216069e 100644 --- a/src/AbstractFixedEffectLinearMap.jl +++ b/src/AbstractFixedEffectLinearMap.jl @@ -1,71 +1,21 @@ ############################################################################## -## -## -## Implement AbstractFixedEffectLinearMap ## -## Model matrix of categorical variables -## mutiplied by diag(1/sqrt(∑w * interaction^2, ..., ∑w * interaction^2) (Jacobi preconditoner) +## AbstractFixedEffectLinearMap ## -## We define these methods used in lsmr! (duck typing): -## eltyp -## size -## mul! +## The whitened operator A = W^(1/2) * (fixed-effect design) * R used by lsmr!, +## where R is the per-group block transform stored in the AbsorptionPlan. +## +## Each backend stores an AbsorptionPlan in a `plan` field and defines mul! +## for itself and its adjoint. lsmr! needs (duck typing): eltype, size, mul!. ## ############################################################################## abstract type AbstractFixedEffectLinearMap{T} end -# Per-fixed-effect plan for the adjoint gather (A'u). Chosen once at construction and -# dispatched on by each backend's `gather!`. The forward map (A = scatter) needs only the -# preconditioner `caches[i]` and is shared; the gather plan lives in `gathers[i]`. -# CPU uses Serial/Threaded; the GPU backends use Atomic/Bucket. -struct SerialGather end -struct ThreadedGather{V<:AbstractVector} - buffers::Vector{V} # one length-fe.n accumulator per thread - ranges::Vector{UnitRange{Int}} # contiguous row chunks -end -struct AtomicGather end -struct BucketGather{V<:AbstractVector} - perm::V # observation indices sorted by group - offsets::V # CSR offsets into perm (length ngroups + 1) -end - Base.adjoint(fem::AbstractFixedEffectLinearMap) = Adjoint(fem) function Base.size(fem::AbstractFixedEffectLinearMap, dim::Integer) - (dim == 1) ? length(fem.fes[1].refs) : (dim == 2) ? sum(fe.n for fe in fem.fes) : 1 + (dim == 1) ? length(fem.fes[1].refs) : (dim == 2) ? _ncoef(fem.plan) : 1 end Base.eltype(x::AbstractFixedEffectLinearMap{T}) where {T} = T - -function scatter!(y, α, fecoef, refs, cache, β) - isone(β) && return scatter!(y, α, fecoef, refs, cache) - iszero(β) ? fill!(y, zero(eltype(y))) : rmul!(y, β) - scatter!(y, α, fecoef, refs, cache) -end - -function LinearAlgebra.mul!(fecoefs::FixedEffectCoefficients, - Cfem::Adjoint{T, <:AbstractFixedEffectLinearMap{T}}, - y::AbstractVector, α::Number, β::Number) where {T} - fem = adjoint(Cfem) - rmul!(fecoefs, β) - for (fecoef, fe, cache, gather) in zip(fecoefs.x, fem.fes, fem.caches, fem.gathers) - gather!(fecoef, fe.refs, α, y, cache, gather) - end - return fecoefs -end - -function LinearAlgebra.mul!(y::AbstractVector, fem::AbstractFixedEffectLinearMap, - fecoefs::FixedEffectCoefficients, α::Number, β::Number) - βj = β - any_fe = false - for (fecoef, fe, cache) in zip(fecoefs.x, fem.fes, fem.caches) - scatter!(y, α, fecoef, fe.refs, cache, βj) - βj = one(βj) - any_fe = true - end - if !any_fe - rmul!(y, β) - end - return y -end diff --git a/src/AbstractFixedEffectSolver.jl b/src/AbstractFixedEffectSolver.jl index 8e09180..430a5c2 100644 --- a/src/AbstractFixedEffectSolver.jl +++ b/src/AbstractFixedEffectSolver.jl @@ -1,8 +1,14 @@ ############################################################################## -## -## AbstractFixedEffectSolver ## -## this type must defined solve_residuals!, solve_coefficients! +## AbstractFixedEffectSolver and the public API +## +## solve_residuals! and solve_coefficients! (defined here) are generic over any +## AbstractFixedEffectSolver. A backend provides: +## AbstractFixedEffectSolver{T}(fes, weights, ::Type{Val{method}}) +## update_weights!(feM, weights) +## copy_internal! (both directions, host <-> solver storage) +## mul! for its linear map and adjoint (used by lsmr!) +## and may override recover_coefficients. ## ############################################################################## abstract type AbstractFixedEffectSolver{T} end @@ -44,8 +50,6 @@ function solve_residuals!(y::AbstractVector{<: Real}, fes::AbstractVector{<: Fix solve_residuals!(y, feM; maxiter = maxiter, tol = tol) end - - function solve_residuals!(r::AbstractVector{<:Real}, feM::AbstractFixedEffectSolver{T}; tol::Real = sqrt(eps(T)), maxiter::Integer = 100_000) where {T} maxiter >= 0 || throw(ArgumentError("maxiter must be non-negative")) # One cannot copy view of Vector (r) on GPU, so first collect the vector @@ -56,7 +60,7 @@ function solve_residuals!(r::AbstractVector{<:Real}, feM::AbstractFixedEffectSol copyto!(feM.b, feM.r) fill!(feM.x, zero(T)) iter, converged = 0, true - if length(feM.x.x) == 1 + if length(feM.m.plan.blocks) == 1 mul!(feM.x, feM.m', feM.b, 1, 0) else _, ch = lsmr!(feM.x, feM.m, feM.b, feM.v, feM.h, feM.hbar; atol = tol, btol = tol, maxiter = maxiter) @@ -100,7 +104,10 @@ function solve_residuals!(xs, feM::AbstractFixedEffectSolver; progress_bar = tru return xs, iterations, convergeds end - +# Guard: without this, a matrix would fall into the collection method above, +# be iterated element-wise, and recurse until a StackOverflowError. +solve_residuals!(::AbstractMatrix, ::AbstractFixedEffectSolver; kwargs...) = + throw(ArgumentError("pass the columns, e.g. eachcol(X), rather than a matrix")) """ @@ -158,9 +165,69 @@ function solve_coefficients!(r::AbstractVector, feM::AbstractFixedEffectSolver{T fill!(feM.x, zero(T)) _, ch = lsmr!(feM.x, feM.m, feM.b, feM.v, feM.h, feM.hbar; atol = tol, btol = tol, maxiter = maxiter) ch.isconverged || @warn "solve_coefficients! did not converge within maxiter LSMR iterations; returned values may be inaccurate." iterations=ch.mvps maxiter tol - for (x, scale) in zip(feM.x.x, feM.m.scales) - x .*= scale + recover_coefficients(feM, eltype(r)), ch.mvps, ch.isconverged +end + + +function recover_coefficients(feM::AbstractFixedEffectSolver{T}, ::Type{Tout}) where {T, Tout} + return recover_coefficients(T, feM.m.fes, feM.m.plan, Matrix{T}[Array(x) for x in feM.x.x], Tout) +end + +# Transform whitened block coefficients back to one vector per input FixedEffect, +# expanded to observation level. `fes` and `coef_blocks` must live on the CPU. +function recover_coefficients(::Type{T}, fes::Vector{<:FixedEffect}, plan::AbsorptionPlan, + coef_blocks::Vector{<:Matrix}, ::Type{Tout}) where {T, Tout} + group_coefs = [zeros(T, fe.n) for fe in fes] + for (coef_block, block, transform) in zip(coef_blocks, plan.blocks, plan.transforms) + k = block_width(block) + β = zeros(T, k) + @inbounds for g in 1:block.n + for a in 1:k + s = zero(T) + for c in 1:k + s += transform[a, c, g] * coef_block[c, g] + end + β[a] = s + end + for (column, term_id) in enumerate(block.input_terms) + group_coefs[term_id][g] = β[column] + end + end + end + normalize!(group_coefs, fes) + return Vector{Tout}[Tout.(coef[fe.refs]) for (coef, fe) in zip(group_coefs, fes)] +end + +# Fixed-effect coefficients are generally not unique: within each connected +# component, a constant can be shifted between the scalar (non-interacted) +# fixed effects. Pin down a solution by demeaning every scalar fixed effect but +# the first within each component (uses `components` from FixedEffect.jl). +function normalize!(fecoefs::AbstractVector{<: Vector{<: Real}}, fes::AbstractVector{<:FixedEffect}) + idx = findall(fe -> isa(fe.interaction, UnitWeights), fes) + length(idx) >= 2 && rescale!(view(fecoefs, idx), view(fes, idx)) + return fecoefs +end + +function rescale!(fecoefs::AbstractVector{<: Vector{<: Real}}, fes::AbstractVector{<:FixedEffect}) + for component_vec in components(fes) + m = 0.0 + # demean all fixed effects except the first + for j in length(fecoefs):(-1):2 + fecoef, component = fecoefs[j], component_vec[j] + mj = 0.0 + for k in component + mj += fecoef[k] + end + mj = mj / length(component) + for k in component + fecoef[k] -= mj + end + m += mj + end + # rescale the first fixed effects + fecoef, component = fecoefs[1], component_vec[1] + for k in component + fecoef[k] += m + end end - x = Vector{eltype(r)}[collect(x) for x in feM.x.x] - full(normalize!(x, feM.m.fes), feM.m.fes), ch.mvps, ch.isconverged end diff --git a/src/CPU.jl b/src/CPU.jl new file mode 100644 index 0000000..ffda428 --- /dev/null +++ b/src/CPU.jl @@ -0,0 +1,237 @@ +############################################################################## +## +## CPU backend +## +## Same layout as ext/CUDAExt.jl and ext/MetalExt.jl: +## 1. FixedEffectLinearMapCPU — the whitened operator over the +## AbsorptionPlan: gather strategies, gather/scatter kernels, and mul! +## for the map and its adjoint (all that lsmr! needs); +## 2. FixedEffectSolverCPU — solver storage and interface. +## +############################################################################## + +############################################################################## +## +## 1. FixedEffectLinearMapCPU +## +############################################################################## + + + +## 1a) FixedEffectLinearMapCPU Constructor + + + +# Per-block plan for the adjoint gather (A'u), chosen once at construction and +# dispatched on by gather_block!. +struct SerialGather end +struct ThreadedGather{M<:AbstractMatrix} + buffers::Vector{M} # one k × n accumulator per thread + ranges::Vector{UnitRange{Int}} # contiguous row chunks +end + +mutable struct FixedEffectLinearMapCPU{T,F<:Vector{<:FixedEffect},P<:AbsorptionPlan,G<:AbstractVector} <: AbstractFixedEffectLinearMap{T} + fes::F + plan::P + gathers::G +end + +# The struct definitions must precede this constructor (a constructor method +# signature is evaluated at definition time, unlike ordinary function calls). +function FixedEffectLinearMapCPU{T}(fes::Vector{<:FixedEffect}, + weights::AbstractVector = uweights(T, length(fes[1].refs))) where {T} + plan = AbsorptionPlan(T, fes, weights) + N = length(fes[1].refs) + nt = nthreads() + ranges = _row_chunks(N, nt) + G = Union{SerialGather, ThreadedGather{Matrix{T}}} + gathers = G[_gather_strategy(T, block, N, nt, ranges) for block in plan.blocks] + return FixedEffectLinearMapCPU{T,typeof(fes),typeof(plan),typeof(gathers)}(fes, plan, gathers) +end + +# Toggle to force the serial baseline (e.g. for benchmarking); threading is on by default. +const _USE_THREADED_GATHER = Ref(true) +# Threading the gather pays off only when the nt per-thread accumulators of size +# k × n fit in cache; beyond that the fill/merge memory traffic dominates +# and serial is faster. +const _GATHER_BUFFER_BUDGET = 8 * 1024 * 1024 # bytes +const _GATHER_MIN_ROWS = 100_000 # below this, threading overhead isn't worth it + +# Per block, thread the gather only if the accumulators fit in cache and N is large. +function _gather_strategy(::Type{T}, block::AbsorbedBlock, N::Int, nt::Int, + ranges::Vector{UnitRange{Int}}) where {T} + k = block_width(block) + if _USE_THREADED_GATHER[] && nt > 1 && N >= _GATHER_MIN_ROWS && + nt * k * block.n * sizeof(T) <= _GATHER_BUFFER_BUDGET + return ThreadedGather([zeros(T, k, block.n) for _ in 1:nt], ranges) + else + return SerialGather() + end +end + + +## 1b) FixedEffectLinearMapCPU mul! + +## Implement right multiplication +function LinearAlgebra.mul!(y::AbstractVector, fem::FixedEffectLinearMapCPU{T}, + fecoefs::FixedEffectCoefficients, α::Number, β::Number) where {T} + # β applies once, fused into the first scatter; later blocks accumulate + for (coef_block, block, qrows) in zip(fecoefs.x, fem.plan.blocks, fem.plan.qrows) + scatter_block!(y, block, coef_block, qrows, α, β) + β = one(β) + end + return y +end + +# y[i] += α * sum over c of coef_block[c, refs[i]] * qrows[c, i], the forward map A x. +function scatter_block!(y::AbstractVector, block::AbsorbedBlock, coef_block::AbstractMatrix, + qrows::AbstractMatrix{T}, + α::Number = one(T)) where {T} + @spawn_for_chunks 100_000 for i in eachindex(y) + @inbounds begin + g = block.refs[i] + fit = zero(T) + for c in 1:block_width(block) + fit += coef_block[c, g] * qrows[c, i] + end + y[i] += α * fit + end + end + return y +end + +# Fused y = β * y + α * (A x) so mul! avoids a separate scaling pass over y. +function scatter_block!(y::AbstractVector, block::AbsorbedBlock, coef_block::AbstractMatrix, + qrows::AbstractMatrix{T}, α::Number, β::Number) where {T} + if isone(β) + return scatter_block!(y, block, coef_block, qrows, α) + end + if iszero(β) + @spawn_for_chunks 100_000 for i in eachindex(y) + @inbounds begin + g = block.refs[i] + fit = zero(T) + for c in 1:block_width(block) + fit += coef_block[c, g] * qrows[c, i] + end + y[i] = α * fit + end + end + else + @spawn_for_chunks 100_000 for i in eachindex(y) + @inbounds begin + g = block.refs[i] + fit = zero(T) + for c in 1:block_width(block) + fit += coef_block[c, g] * qrows[c, i] + end + y[i] = β * y[i] + α * fit + end + end + end + return y +end + +## 1c) FixedEffectLinearMapCPU mul!, Adjoint + + +## Implement left multiplication +function LinearAlgebra.mul!(fecoefs::FixedEffectCoefficients, + Cfem::Adjoint{T, <:FixedEffectLinearMapCPU{T}}, + y::AbstractVector, α::Number, β::Number) where {T} + fem = adjoint(Cfem) + rmul!(fecoefs, β) + for (coef_block, block, qrows, gather) in zip(fecoefs.x, fem.plan.blocks, fem.plan.qrows, fem.gathers) + gather_block!(coef_block, block, qrows, y, α, gather) + end + return fecoefs +end + + +# Serial: one pass over all rows straight into coef_block (which already holds β * old). +gather_block!(coef_block::AbstractMatrix, block::AbsorbedBlock, qrows::AbstractMatrix, + y::AbstractVector, α::Number, ::SerialGather) = + _gather_block!(coef_block, block, qrows, y, α, eachindex(y)) + +# Threaded: each thread reduces its row chunk into a private (cache-resident) buffer, +# then the buffers are summed into coef_block. +function gather_block!(coef_block::AbstractMatrix, block::AbsorbedBlock, qrows::AbstractMatrix, + y::AbstractVector, α::Number, g::ThreadedGather) + @threads for t in eachindex(g.buffers) + buf = g.buffers[t] + fill!(buf, zero(eltype(buf))) + _gather_block!(buf, block, qrows, y, α, g.ranges[t]) + end + @inbounds for buf in g.buffers + @simd for j in eachindex(coef_block) + coef_block[j] += buf[j] + end + end + return coef_block +end + + +# Kernels. block_width is the length of the interactions tuple, so it is a +# compile-time constant: the inner loops over columns unroll separately for +# each block width, and no hand-written k = 1 or k = 2 specializations are +# needed. Inside the @spawn_for_chunks closures the loop bound must be written +# block_width(block) (not a captured integer) for the constant to survive. + +# coef_block[c, refs[i]] += α * y[i] * qrows[c, i] over one row range. +# No @simd: distinct i may write the same coef_block column. +function _gather_block!(coef_block::AbstractMatrix, block::AbsorbedBlock, + qrows::AbstractMatrix{T}, y::AbstractVector, α::Number, range) where {T} + k = block_width(block) + @fastmath @inbounds for i in range + g = block.refs[i] + yi = α * y[i] + for c in 1:k + coef_block[c, g] += yi * qrows[c, i] + end + end + return coef_block +end + +############################################################################## +## +## 2. FixedEffectSolverCPU +## +############################################################################## + +mutable struct FixedEffectSolverCPU{T,M<:FixedEffectLinearMapCPU{T},C<:FixedEffectCoefficients{Matrix{T}}} <: AbstractFixedEffectSolver{T} + m::M + weights::AbstractVector + b::Vector{T} + r::Vector{T} + x::C + v::C + h::C + hbar::C +end + +function AbstractFixedEffectSolver{T}(fes::Vector{<:FixedEffect}, weights::AbstractWeights, ::Type{Val{:cpu}}) where {T} + m = FixedEffectLinearMapCPU{T}(fes, weights) + b = zeros(T, length(weights)) + r = zeros(T, length(weights)) + blocks = m.plan.blocks + x = FixedEffectCoefficients([zeros(T, block_width(block), block.n) for block in blocks]) + v = FixedEffectCoefficients([zeros(T, block_width(block), block.n) for block in blocks]) + h = FixedEffectCoefficients([zeros(T, block_width(block), block.n) for block in blocks]) + hbar = FixedEffectCoefficients([zeros(T, block_width(block), block.n) for block in blocks]) + return FixedEffectSolverCPU(m, weights, b, r, x, v, h, hbar) +end + +function update_weights!(feM::FixedEffectSolverCPU{T}, weights::AbstractWeights) where {T} + feM.m.plan = AbsorptionPlan(T, feM.m.plan, weights) + feM.weights = weights + return feM +end + +function copy_internal!(feM::FixedEffectSolverCPU, field::Symbol, r::AbstractVector) + copyto!(getfield(feM, field), r) +end + +function copy_internal!(r::AbstractVector, feM::FixedEffectSolverCPU, field::Symbol) + copyto!(r, getfield(feM, field)) +end + diff --git a/src/FixedEffect.jl b/src/FixedEffect.jl index 48e205b..70f5e54 100644 --- a/src/FixedEffect.jl +++ b/src/FixedEffect.jl @@ -5,7 +5,7 @@ ############################################################################## struct FixedEffect{R <: AbstractVector{<:Integer}, I <: AbstractVector{<:Real}} - refs::R # refs must be between 0 and n + refs::R # group of each observation, in 1:n (0 marks a missing group; such rows must be dropped before solving) interaction::I # the continuous interaction n::Int # Number of potential values (= maximum(refs)) function FixedEffect{R, I}(refs, interaction, n) where {R <: AbstractVector{<:Integer}, I <: AbstractVector{<: Real}} @@ -103,44 +103,3 @@ function components(fes::AbstractVector{<:FixedEffect}) return out end -############################################################################## -## -## normalize! a vector of fixedeffect coefficients using connected components -## -############################################################################## - -function normalize!(fecoefs::AbstractVector{<: Vector{<: Real}}, fes::AbstractVector{<:FixedEffect}) - # The solution is generally not unique. Find connected components and scale accordingly - idx = findall(fe -> isa(fe.interaction, UnitWeights), fes) - length(idx) >= 2 && rescale!(view(fecoefs, idx), view(fes, idx)) - return fecoefs -end - -function rescale!(fecoefs::AbstractVector{<: Vector{<: Real}}, fes::AbstractVector{<:FixedEffect}) - for component_vec in components(fes) - m = 0.0 - # demean all fixed effects except the first - for j in length(fecoefs):(-1):2 - fecoef, component = fecoefs[j], component_vec[j] - mj = 0.0 - for k in component - mj += fecoef[k] - end - mj = mj / length(component) - for k in component - fecoef[k] -= mj - end - m += mj - end - # rescale the first fixed effects - fecoef, component = fecoefs[1], component_vec[1] - for k in component - fecoef[k] += m - end - end -end - -function full(fecoefs::AbstractVector{<: Vector{<: Real}}, fes::AbstractVector{<:FixedEffect}) - # add collect in case Metal, since then fecoef is Vector while fe.refs is GPU - [fecoef[collect(fe.refs)] for (fecoef, fe) in zip(fecoefs, fes)] -end diff --git a/src/FixedEffectCoefficients.jl b/src/FixedEffectCoefficients.jl index e7bd0a6..42c7ff0 100644 --- a/src/FixedEffectCoefficients.jl +++ b/src/FixedEffectCoefficients.jl @@ -1,17 +1,17 @@ -# Define methods used in LSMR - ############################################################################## -## -## FixedEffectCoefficients : vector x in A'Ax = A'b ## -## We define these methods used in lsmr! (duck typing): -## copyto!, fill!, rmul!, axpy!, norm +## FixedEffectCoefficients: the whitened coefficient vector x seen by lsmr!, +## stored as one k × n matrix per AbsorbedBlock (column g holds the +## whitened coordinates of group g). +## +## We define the methods lsmr! needs (duck typing): +## copyto!, fill!, rmul!, axpy!, norm, similar ## -## Do not define iteration on each fixedeffect since it would conflict with eltype +## Do not define iteration on each block since it would conflict with eltype ## ############################################################################## -struct FixedEffectCoefficients{U <: AbstractVector} +struct FixedEffectCoefficients{U <: AbstractArray} x::Vector{U} end @@ -53,3 +53,7 @@ function LinearAlgebra.axpy!(α::Number, fecoefs1::FixedEffectCoefficients, feco end return fecoefs2 end + +function Base.similar(fecoefs::FixedEffectCoefficients, ::Type{T} = eltype(fecoefs)) where {T} + return FixedEffectCoefficients([similar(x, T) for x in fecoefs.x]) +end diff --git a/src/FixedEffects.jl b/src/FixedEffects.jl index 7852041..86ea0e4 100644 --- a/src/FixedEffects.jl +++ b/src/FixedEffects.jl @@ -23,12 +23,23 @@ include("utils/lsmr.jl") include("utils/progressbar.jl") include("FixedEffect.jl") +include("AbsorptionPlan.jl") include("AbstractFixedEffectSolver.jl") include("FixedEffectCoefficients.jl") include("AbstractFixedEffectLinearMap.jl") -include("SolverCPU.jl") +include("CPU.jl") +############################################################################## +## +## Compatibility shims +## +############################################################################## +function AbstractFixedEffectSolver{T}(fes::Vector{<:FixedEffect}, weights::AbstractWeights, + method::Type{Val{M}}, nthreads::Integer) where {T,M} + nthreads >= 1 || throw(ArgumentError("nthreads must be positive")) + return AbstractFixedEffectSolver{T}(fes, weights, method) +end include("precompile.jl") diff --git a/src/SolverCPU.jl b/src/SolverCPU.jl deleted file mode 100644 index 0b98383..0000000 --- a/src/SolverCPU.jl +++ /dev/null @@ -1,178 +0,0 @@ -############################################################################## -## -## Implement AbstractFixedEffectLinearMap -## -############################################################################## - -mutable struct FixedEffectLinearMapCPU{T,F<:Vector{<:FixedEffect},S<:AbstractVector,C<:AbstractVector,G<:AbstractVector} <: AbstractFixedEffectLinearMap{T} - fes::F - scales::S - caches::C - gathers::G -end - -# Toggle to force the serial baseline (e.g. for benchmarking); threading is on by default. -const _USE_THREADED_GATHER = Ref(true) -# Threading the gather pays off only when the nt per-thread accumulators of length fe.n fit -# in cache; beyond that the fill/merge memory traffic dominates and serial is faster. -const _GATHER_BUFFER_BUDGET = 8 * 1024 * 1024 # bytes -const _GATHER_MIN_ROWS = 100_000 # below this, threading overhead isn't worth it - -function _row_chunks(n::Int, k::Int) - base, rem = divrem(n, k) - out = Vector{UnitRange{Int}}(undef, k) - s = 1 - for t in 1:k - len = base + (t <= rem ? 1 : 0) - out[t] = s:(s + len - 1) - s += len - end - return out -end - -# Per fixed effect, thread the gather only if the accumulators fit in cache and N is large. -function _gather_strategy(::Type{T}, fe::FixedEffect, N::Int, nt::Int, - ranges::Vector{UnitRange{Int}}) where {T} - if _USE_THREADED_GATHER[] && nt > 1 && N >= _GATHER_MIN_ROWS && - nt * fe.n * sizeof(T) <= _GATHER_BUFFER_BUDGET - return ThreadedGather([zeros(T, fe.n) for _ in 1:nt], ranges) - else - return SerialGather() - end -end - -function FixedEffectLinearMapCPU{T}(fes::Vector{<:FixedEffect}) where {T} - scales = [zeros(T, fe.n) for fe in fes] - caches = [zeros(T, length(fes[1].interaction)) for fe in fes] - N = length(fes[1].refs) - nt = nthreads() - ranges = _row_chunks(N, nt) - G = Union{SerialGather, ThreadedGather{Vector{T}}} - gathers = G[_gather_strategy(T, fe, N, nt, ranges) for fe in fes] - return FixedEffectLinearMapCPU{T,typeof(fes),typeof(scales),typeof(caches),typeof(gathers)}(fes, scales, caches, gathers) -end - - -# The one place the gather arithmetic lives: scatter-add a row range into `out`, -# out[refs[i]] += α * y[i] * cache[i]. No @simd: distinct i may write the same out[refs[i]]. -function _gather!(out::AbstractVector, refs::AbstractVector, α::Number, - y::AbstractVector, cache::AbstractVector, range) - @fastmath @inbounds for i in range - out[refs[i]] += α * y[i] * cache[i] - end - return out -end - -# Serial: one pass over all rows straight into fecoef (which already holds β * old). -gather!(fecoef::AbstractVector, refs::AbstractVector, α::Number, - y::AbstractVector, cache::AbstractVector, ::SerialGather) = - _gather!(fecoef, refs, α, y, cache, eachindex(y)) - -# Threaded: each thread reduces its row chunk into a private (cache-resident) buffer, -# then the buffers are summed into fecoef. -function gather!(fecoef::AbstractVector, refs::AbstractVector, α::Number, - y::AbstractVector, cache::AbstractVector, g::ThreadedGather) - @threads for t in eachindex(g.buffers) - buf = g.buffers[t] - fill!(buf, zero(eltype(buf))) - _gather!(buf, refs, α, y, cache, g.ranges[t]) - end - @inbounds for buf in g.buffers - @simd for k in eachindex(fecoef) - fecoef[k] += buf[k] - end - end - return fecoef -end - -function scatter!(y::AbstractVector, α::Number, fecoef::AbstractVector, - refs::AbstractVector, cache::AbstractVector) - @spawn_for_chunks 100_000 for i in eachindex(y) - @inbounds y[i] += α * fecoef[refs[i]] * cache[i] - end -end - -function scatter!(y::Vector, α::Number, fecoef::Vector, - refs::Vector, cache::Vector, β::Number) - if iszero(β) - @spawn_for_chunks 100_000 for i in eachindex(y) - @inbounds y[i] = α * fecoef[refs[i]] * cache[i] - end - elseif isone(β) - scatter!(y, α, fecoef, refs, cache) - else - @spawn_for_chunks 100_000 for i in eachindex(y) - # Fuse y[i] = β * y[i] + α * fecoef[refs[i]] * cache[i] to avoid a separate scaling pass. - @inbounds y[i] = β * y[i] + α * fecoef[refs[i]] * cache[i] - end - end -end - - -############################################################################## -## -## Implement AbstractFixedEffectSolver interface -## -############################################################################## - -mutable struct FixedEffectSolverCPU{T,M<:FixedEffectLinearMapCPU{T},C<:FixedEffectCoefficients{Vector{T}}} <: AbstractFixedEffectSolver{T} - m::M - weights::AbstractVector - b::Vector{T} - r::Vector{T} - x::C - v::C - h::C - hbar::C -end - - -function AbstractFixedEffectSolver{T}(fes::Vector{<:FixedEffect}, weights::AbstractWeights, ::Type{Val{:cpu}}) where {T} - m = FixedEffectLinearMapCPU{T}(fes) - b = zeros(T, length(weights)) - r = zeros(T, length(weights)) - x = FixedEffectCoefficients([zeros(T, fe.n) for fe in fes]) - v = FixedEffectCoefficients([zeros(T, fe.n) for fe in fes]) - h = FixedEffectCoefficients([zeros(T, fe.n) for fe in fes]) - hbar = FixedEffectCoefficients([zeros(T, fe.n) for fe in fes]) - return update_weights!(FixedEffectSolverCPU(m, weights, b, r, x, v, h, hbar), weights) -end - - -function update_weights!(feM::FixedEffectSolverCPU, weights::AbstractWeights) - for (scale, fe) in zip(feM.m.scales, feM.m.fes) - scale!(scale, fe.refs, fe.interaction, weights) - end - for (cache, scale, fe) in zip(feM.m.caches, feM.m.scales, feM.m.fes) - cache!(cache, fe.refs, fe.interaction, weights, scale) - end - feM.weights = weights - return feM -end - -function scale!(scale::AbstractVector, refs::AbstractVector, interaction::AbstractVector, weights::AbstractVector) - fill!(scale, 0) - # no @simd: multiple i may write to the same scale[refs[i]] - @fastmath @inbounds for i in eachindex(refs) - scale[refs[i]] += abs2(interaction[i]) * weights[i] - end - # Case of interaction variable equal to zero in the category (issue #97) - T = eltype(scale) - @fastmath @inbounds @simd for i in eachindex(scale) - scale[i] = scale[i] > 0 ? (1 / sqrt(scale[i])) : zero(T) - end -end - -function cache!(cache::AbstractVector, refs::AbstractVector, interaction::AbstractVector, weights::AbstractVector, scale::AbstractVector) - @spawn_for_chunks 100_000 for i in eachindex(cache) - @inbounds @fastmath cache[i] = interaction[i] * sqrt(weights[i]) * scale[refs[i]] - end -end - -function copy_internal!(feM::FixedEffectSolverCPU, field::Symbol, r::AbstractVector) - copyto!(getfield(feM, field), r) -end - -function copy_internal!(r::AbstractVector, feM::FixedEffectSolverCPU, field::Symbol) - copyto!(r, getfield(feM, field)) -end diff --git a/test/runtests.jl b/test/runtests.jl index 00fc320..12a3be1 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -12,4 +12,4 @@ for test in tests showerror(stdout, e, backtrace()) rethrow(e) end -end \ No newline at end of file +end diff --git a/test/solve.jl b/test/solve.jl index 19c2d07..9c3e505 100644 --- a/test/solve.jl +++ b/test/solve.jl @@ -1,4 +1,4 @@ - +using LinearAlgebra p1 = repeat(1:5, inner = 2) p2 = repeat(1:5, outer = 2) @@ -57,6 +57,18 @@ function _residual_from_coefs(y, coefs) return out end +function _residual_from_coefs(y, fes, coefs) + out = copy(y) + for (fe, coef) in zip(fes, coefs) + if fe.interaction isa UnitWeights + out .-= coef + else + out .-= fe.interaction .* coef + end + end + return out +end + @testset "GPU parity" begin n_gpu = 2048 p1_gpu = mod1.(1:n_gpu, 32) @@ -66,6 +78,12 @@ end interaction_gpu = 0.5 .+ mod.(1:n_gpu, 11) ./ 13 fes_gpu = [FixedEffect(p1_gpu), FixedEffect(p2_gpu)] fes_interact_gpu = [FixedEffect(p1_gpu, interaction = interaction_gpu), FixedEffect(p2_gpu)] + fes_block_gpu = [ + FixedEffect(p1_gpu), + FixedEffect(p1_gpu, interaction = interaction_gpu), + FixedEffect(p2_gpu), + FixedEffect(p2_gpu, interaction = 2 .* interaction_gpu), + ] fes_bin_gpu = [FixedEffect(p1_gpu)] atol_gpu = 1e-3 rtol_gpu = 1e-3 @@ -83,10 +101,19 @@ end gpu_interact_r = solve_residuals!(deepcopy(x_gpu), fes_interact_gpu, weights_gpu; method = method, double_precision = false)[1] @test gpu_interact_r ≈ cpu_interact_r atol=atol_gpu rtol=rtol_gpu + cpu_block_r = solve_residuals!(deepcopy(x_gpu), fes_block_gpu, weights_gpu; double_precision = false)[1] + gpu_block_r = solve_residuals!(deepcopy(x_gpu), fes_block_gpu, weights_gpu; method = method, double_precision = false)[1] + @test gpu_block_r ≈ cpu_block_r atol=atol_gpu rtol=rtol_gpu + cpu_coefs = solve_coefficients!(deepcopy(x_gpu), fes_gpu, weights_gpu; double_precision = false)[1] gpu_coefs = solve_coefficients!(deepcopy(x_gpu), fes_gpu, weights_gpu; method = method, double_precision = false)[1] @test _residual_from_coefs(x_gpu, gpu_coefs) ≈ _residual_from_coefs(x_gpu, cpu_coefs) atol=atol_gpu rtol=rtol_gpu + cpu_block_coefs = solve_coefficients!(deepcopy(x_gpu), fes_block_gpu, weights_gpu; double_precision = false)[1] + gpu_block_coefs = solve_coefficients!(deepcopy(x_gpu), fes_block_gpu, weights_gpu; method = method, double_precision = false)[1] + @test _residual_from_coefs(x_gpu, fes_block_gpu, gpu_block_coefs) ≈ + _residual_from_coefs(x_gpu, fes_block_gpu, cpu_block_coefs) atol=atol_gpu rtol=rtol_gpu + cpu_bin_r = solve_residuals!(deepcopy(x_gpu), fes_bin_gpu, weights_gpu; double_precision = false)[1] gpu_bin_r = solve_residuals!(deepcopy(x_gpu), fes_bin_gpu, weights_gpu; method = method, double_precision = false)[1] @test gpu_bin_r ≈ cpu_bin_r atol=atol_gpu rtol=rtol_gpu @@ -114,6 +141,11 @@ solve_residuals!(deepcopy(x), feM)[1] ≈ solve_residuals!(deepcopy(x), fes, wei weights = Weights(reverse([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])) FixedEffects.update_weights!(feM, weights) solve_residuals!(deepcopy(x), feM)[1] ≈ solve_residuals!(deepcopy(x), fes, weights)[1] +feM_compat = FixedEffects.AbstractFixedEffectSolver{Float64}(fes, weights, Val{:cpu}, Base.Threads.nthreads()) +@test solve_residuals!(deepcopy(x), feM_compat)[1] ≈ solve_residuals!(deepcopy(x), fes, weights)[1] + +# a matrix must be passed as columns (e.g. eachcol), not raw +@test_throws ArgumentError solve_residuals!(rand(10, 2), feM) # test interacted fixed effects interaction = [0.2, 0.8, 0.3, 0.7, 0.5, 0.5, 0.4, 0.6, 0.1, 0.9] @@ -125,3 +157,146 @@ fes_interact = [FixedEffect(p1, interaction = interaction)] fes_both = [FixedEffect(p1), FixedEffect(p1, interaction = interaction)] (r_both, iter, conv) = solve_residuals!(deepcopy(x), fes_both) @test conv + +@testset "block absorption algebra" begin + id = repeat(1:4, inner = 3) + slope = [-1.0, 0.0, 1.0, 0.5, 1.0, 1.5, -2.0, -1.0, 0.0, 2.0, 2.0, 2.0] + y = [2.0, 1.5, 1.0, -1.0, 0.5, 1.5, 3.0, 2.0, 1.0, -2.0, -1.5, -1.0] + weights = Weights(1 .+ (1:length(y)) ./ 20) + fes_block = [ + FixedEffect(id), + FixedEffect(id, interaction = slope), + FixedEffect(id, interaction = slope .^ 2), + ] + plan = FixedEffects.AbsorptionPlan(Float64, fes_block, weights) + @test length(plan.blocks) == 1 + @test FixedEffects.block_width(plan.blocks[1]) == 3 + @test maximum(plan.ranks[1]) == 3 + @test minimum(plan.ranks[1]) < 3 + + X = zeros(length(y), 12) + for i in eachindex(y) + X[i, id[i]] = 1 + X[i, 4 + id[i]] = slope[i] + X[i, 8 + id[i]] = slope[i]^2 + end + sqrtw = sqrt.(weights) + Xw = X .* reshape(sqrtw, :, 1) + yw = y .* sqrtw + r_dense = (yw - Xw * (pinv(Xw) * yw)) ./ sqrtw + r_block, _, conv = solve_residuals!(copy(y), fes_block, weights) + @test conv + @test r_block ≈ r_dense atol = 1e-10 + + coefs = solve_coefficients!(copy(y), fes_block, weights)[1] + @test y .- coefs[1] .- slope .* coefs[2] .- slope.^2 .* coefs[3] ≈ r_block atol = 1e-10 + + # several variables through the collection fallback, sharing one solver + feM_block = FixedEffects.AbstractFixedEffectSolver{Float64}(fes_block, weights, Val{:cpu}) + cols = [copy(y), 2 .* y .+ 1] + solve_residuals!(cols, feM_block; progress_bar = false) + @test cols[1] ≈ r_block atol = 1e-10 + @test cols[2] ≈ solve_residuals!(2 .* y .+ 1, fes_block, weights)[1] atol = 1e-10 + + id_big = repeat(1:2, inner = 3) + slope_big = [100_000.0, 100_001.0, 100_002.0, 200_000.0, 200_001.0, 200_002.0] + y_big = [1.0, 2.0, 4.0, -1.0, 0.5, 3.0] + fes_big = [FixedEffect(id_big), FixedEffect(id_big, interaction = slope_big)] + plan_big = FixedEffects.AbsorptionPlan(Float64, fes_big, uweights(length(y_big))) + @test plan_big.ranks[1] == [2, 2] + X_big = zeros(length(y_big), 4) + for i in eachindex(y_big) + X_big[i, id_big[i]] = 1 + X_big[i, 2 + id_big[i]] = slope_big[i] + end + @test solve_residuals!(copy(y_big), fes_big)[1] ≈ y_big - X_big * (pinv(X_big) * y_big) atol = 1e-8 +end + +# Independent implementation of the exact one-block projection residual, +# y ← y - Z G⁺ Z' W y with G⁺ = R R' from the plan transforms: the oracle the +# operator tests below are checked against. +function project_block!(y::AbstractVector, block, transform, weights) + k = FixedEffects.block_width(block) + coef = zeros(k, block.n) + for i in eachindex(y) + g = block.refs[i] + for c in 1:k + coef[c, g] += weights[i] * block.interactions[c][i] * y[i] + end + end + for g in 1:block.n + R = transform[:, :, g] + coef[:, g] = R * (R' * coef[:, g]) + end + for i in eachindex(y) + g = block.refs[i] + for c in 1:k + y[i] -= block.interactions[c][i] * coef[c, g] + end + end + return y +end + +@testset "block operator identities" begin + id = [1, 1, 1, 2, 2, 3, 3, 3] + slope = [0.0, 1.0, 2.0, 1.0, 1.0, -1.0, 0.0, 1.0] + y = [1.0, 0.5, 2.0, -1.0, 3.0, 2.5, -0.5, 1.5] + z = [-2.0, 1.0, 0.0, 4.0, 3.0, -1.0, 2.0, 0.5] + weights = Weights([1.0, 1.5, 2.0, 0.75, 1.25, 1.1, 0.9, 1.3]) + fes_block = [FixedEffect(id), FixedEffect(id, interaction = slope), FixedEffect(id, interaction = 2 .* slope)] + feM = FixedEffects.AbstractFixedEffectSolver{Float64}(fes_block, weights, Val{:cpu}) + + coef = FixedEffects.FixedEffectCoefficients([randn(size(blockcoef)) for blockcoef in feM.x.x]) + Acoef = zeros(length(y)) + mul!(Acoef, feM.m, coef, 1.0, 0.0) + adj = similar(coef) + fill!(adj, 0.0) + mul!(adj, feM.m', z, 1.0, 0.0) + @test dot(Acoef, z) ≈ sum(dot(a, b) for (a, b) in zip(coef.x, adj.x)) atol = 1e-10 + + r1 = copy(y) + project_block!(r1, feM.m.plan.blocks[1], feM.m.plan.transforms[1], weights) + r2 = copy(r1) + project_block!(r2, feM.m.plan.blocks[1], feM.m.plan.transforms[1], weights) + @test r2 ≈ r1 atol = 1e-10 + rz = copy(z) + project_block!(rz, feM.m.plan.blocks[1], feM.m.plan.transforms[1], weights) + @test dot(weights .* r1, z) ≈ dot(weights .* y, rz) atol = 1e-10 + for g in 1:3 + for s in (ones(length(y)), slope, 2 .* slope) + @test sum(weights[i] * s[i] * r1[i] for i in eachindex(y) if id[i] == g) ≈ 0 atol = 1e-10 + end + end +end + +@testset "rank tolerance with large collinear groups" begin + n = 200_000 + id = mod1.(1:n, 4) + slope = randn(n) + y = randn(n) + fes_dup = [FixedEffect(id), FixedEffect(id, interaction = slope), FixedEffect(id, interaction = 2 .* slope)] + plan = FixedEffects.AbsorptionPlan(Float64, fes_dup, uweights(n)) + # slope and 2 * slope are exactly collinear: the rank must be 2 even though the + # orthogonalization residual of the duplicated column in a 50_000-row group is + # rounding noise far above eps + @test plan.ranks[1] == fill(2, 4) + r_dup = solve_residuals!(copy(y), fes_dup)[1] + r_ref = solve_residuals!(copy(y), [FixedEffect(id), FixedEffect(id, interaction = slope)])[1] + @test r_dup ≈ r_ref atol = 1e-8 + coefs = solve_coefficients!(copy(y), fes_dup)[1] + @test all(maximum(abs, c) < 1e6 for c in coefs) +end + +@testset "threaded gather parity" begin + n = 150_000 + id1 = mod1.(1:n, 100) + id2 = mod1.(7 .* (1:n) .+ 3, 31) + xslope = randn(n) + fes = [FixedEffect(id1), FixedEffect(id1, interaction = xslope), FixedEffect(id2)] + y = randn(n) + r_default = solve_residuals!(copy(y), fes)[1] + FixedEffects._USE_THREADED_GATHER[] = false + r_serial = solve_residuals!(copy(y), fes)[1] + FixedEffects._USE_THREADED_GATHER[] = true + @test r_default ≈ r_serial atol = 1e-8 +end