From 5e95e8dc9847317d53d6ce51d79531aede38c7e8 Mon Sep 17 00:00:00 2001 From: Matthieu Gomez Date: Thu, 27 Aug 2026 16:58:32 -0400 Subject: [PATCH 1/2] Speed up fixed-effect absorption --- ext/CUDAExt.jl | 78 ++++++++++++++++++++++++-- ext/MetalExt.jl | 50 +++++------------ src/AbsorptionPlan.jl | 87 ++++++++++++++++++++++------- src/AbstractFixedEffectLinearMap.jl | 15 ++++- test/solve.jl | 36 ++++++++++++ 5 files changed, 202 insertions(+), 64 deletions(-) diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl index 1a8c3e4..3702688 100644 --- a/ext/CUDAExt.jl +++ b/ext/CUDAExt.jl @@ -1,6 +1,6 @@ module CUDAExt using FixedEffects, CUDA -using FixedEffects: FixedEffectCoefficients, AbstractWeights, UnitWeights, LinearAlgebra, Adjoint, mul!, rmul!, AbstractFixedEffectLinearMap, copy_internal!, AbsorptionPlan, AbsorbedBlock, block_width +using FixedEffects: FixedEffectCoefficients, AbstractWeights, UnitWeights, LinearAlgebra, Adjoint, mul!, rmul!, AbstractFixedEffectLinearMap, copy_internal!, AbsorptionPlan, AbsorbedBlock, _group_permutation, block_width CUDA.allowscalar(false) ############################################################################## @@ -25,14 +25,38 @@ CUDA.allowscalar(false) _cu(T::Type, w::UnitWeights) = fill!(CuVector{T}(undef, length(w)), w[1]) _cu(T::Type, w::AbstractVector) = CuVector{T}(convert(Vector{T}, w)) +# Per-block plan for the adjoint gather (A'u), chosen once at construction: +# bucketize (one thread block per group) for low cardinality, else atomic adds. +struct AtomicGather end +struct BucketGather{V<:AbstractVector} + perm::V + offsets::V +end + mutable struct FixedEffectLinearMapCUDA{T,P<:AbsorptionPlan} <: AbstractFixedEffectLinearMap{T} fes::Vector{<:FixedEffect} plan::P + gathers::Vector{Union{AtomicGather, BucketGather}} end function FixedEffectLinearMapCUDA{T}(fes::Vector{<:FixedEffect}, weights::AbstractWeights) where {T} plan = _cu_plan(T, fes, weights) - return FixedEffectLinearMapCUDA{T,typeof(plan)}(fes, plan) + G = Union{AtomicGather, BucketGather} + gathers = Vector{G}(undef, length(plan.blocks)) + for i in eachindex(plan.blocks) + refs = fes[plan.blocks[i].input_terms[1]].refs + gathers[i] = _gather_strategy(refs, plan.blocks[i].n) + end + return FixedEffectLinearMapCUDA{T,typeof(plan)}(fes, plan, gathers) +end + +function _gather_strategy(refs::AbstractVector{<:Integer}, nlevels::Int) + if nlevels < min(100_000, div(length(refs), 16)) + _, offsets, perm = _group_permutation(refs, nlevels) + return BucketGather(CuVector{Int}(perm), CuVector{Int}(offsets)) + else + return AtomicGather() + end end function _cu_plan(::Type{T}, fes::Vector{<:FixedEffect}, weights::AbstractWeights) where {T} @@ -91,20 +115,64 @@ function LinearAlgebra.mul!(fecoefs::FixedEffectCoefficients, 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, α) + 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::CuMatrix, refs::CuVector, qrows::CuMatrix, - y::CuVector, α::Number) + y::CuVector, α::Number, gather::BucketGather) + nthreads = 256 + nblocks = size(coef_block, 2) + @cuda threads=nthreads blocks=nblocks gather_block_kernel_bin!(coef_block, α, y, qrows, + gather.perm, gather.offsets, Val(nthreads), size(coef_block, 1)) + return coef_block +end + +function _gather_block!(coef_block::CuMatrix, refs::CuVector, qrows::CuMatrix, + y::CuVector, α::Number, ::AtomicGather) 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 gather_block_kernel_bin!(coef_block, α, y, qrows, perm, offsets, + ::Val{NT}, k) where {NT} + g = Int(blockIdx().x) + tid = Int(threadIdx().x) + T = eltype(coef_block) + shared = CUDA.CuStaticSharedArray(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 + + @inbounds shared[tid] = acc + CUDA.sync_threads() + offset = NT ÷ 2 + while offset > 0 + if tid <= offset + @inbounds shared[tid] += shared[tid + offset] + end + CUDA.sync_threads() + offset ÷= 2 + end + if tid == 1 + @inbounds coef_block[c, g] += shared[1] + end + end + return nothing +end + function gather_block_kernel!(coef_block, refs, qrows, y, α, k) index = (blockIdx().x - Int32(1)) * blockDim().x + threadIdx().x stride = blockDim().x * gridDim().x diff --git a/ext/MetalExt.jl b/ext/MetalExt.jl index 3f99ff7..5e31d65 100644 --- a/ext/MetalExt.jl +++ b/ext/MetalExt.jl @@ -1,6 +1,6 @@ module MetalExt using FixedEffects, Metal -using FixedEffects: FixedEffectCoefficients, AbstractWeights, UnitWeights, LinearAlgebra, Adjoint, mul!, rmul!, AbstractFixedEffectLinearMap, copy_internal!, AbsorptionPlan, AbsorbedBlock, block_width +using FixedEffects: FixedEffectCoefficients, AbstractWeights, UnitWeights, LinearAlgebra, Adjoint, mul!, rmul!, AbstractFixedEffectLinearMap, copy_internal!, AbsorptionPlan, AbsorbedBlock, _group_permutation, block_width Metal.allowscalar(false) ############################################################################## @@ -34,8 +34,8 @@ end # 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) + perm::V + offsets::V end mutable struct FixedEffectLinearMapMetal{T,P<:AbsorptionPlan} <: AbstractFixedEffectLinearMap{T} @@ -48,20 +48,22 @@ function FixedEffectLinearMapMetal{T}(fes::Vector{<:FixedEffect}, weights::Abstr 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) + for i in eachindex(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 + gathers[i] = _gather_strategy(refs, plan.blocks[i].n) end return FixedEffectLinearMapMetal{T,typeof(plan)}(fes, plan, gathers) end +function _gather_strategy(refs::AbstractVector{<:Integer}, nlevels::Int) + if nlevels < min(100_000, div(length(refs), 16)) + _, offsets, perm = _group_permutation(refs, nlevels) + return BucketGather(MtlVector{Int}(perm), MtlVector{Int}(offsets)) + else + return AtomicGather() + end +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) @@ -70,30 +72,6 @@ function _mtl_plan(::Type{T}, fes::Vector{<:FixedEffect}, weights::AbstractWeigh 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) - @inbounds for r in refs - counts[r] += 1 - end - # offsets is vcat(1, cumsum(counts)) - offsets = Vector{Int}(undef, n + 1) - offsets[1] = 1 - @inbounds for k in 1:n - offsets[k+1] = offsets[k] + counts[k] - end - - perm = Vector{Int}(undef, length(refs)) - next = offsets[1:n] - @inbounds for i in eachindex(refs) - r = refs[i] - p = next[r] - perm[p] = i - next[r] = p + 1 - end - return MtlVector{Int}(perm), MtlVector{Int}(offsets) -end - ## 1b) FixedEffectLinearMapMetal mul! ## Implement right multiplication diff --git a/src/AbsorptionPlan.jl b/src/AbsorptionPlan.jl index 7cfc01d..87d986a 100644 --- a/src/AbsorptionPlan.jl +++ b/src/AbsorptionPlan.jl @@ -41,7 +41,6 @@ struct AbsorptionPlan{B<:AbstractVector{<:AbsorbedBlock},TR<:AbstractVector,RA<: end block_width(block::AbsorbedBlock) = length(block.interactions) -_ncoef(plan::AbsorptionPlan) = sum(block_width(block) * block.n for block in plan.blocks) ## 1b) Constructors @@ -118,29 +117,78 @@ function _build_block_transform(::Type{T}, block::AbsorbedBlock, weights::Abstra end end return transforms, ranks, qrows - end + elseif k == 2 && count(interaction -> interaction isa UnitWeights, block.interactions) == 1 + # Stable weighted moments avoid cancellation in within-group slope variation. + intercept = block.interactions[1] isa UnitWeights ? 1 : 2 + slope = 3 - intercept + z = block.interactions[slope] + sumw = zeros(T, nlevels) + anchor = zeros(T, nlevels) + mean = zeros(T, nlevels) + m2 = zeros(T, nlevels) + @inbounds for i in eachindex(block.refs) + g = block.refs[i] + w = T(weights[i]) + zi = T(z[i]) + new_sumw = sumw[g] + w + if new_sumw > zero(T) + if iszero(sumw[g]) + anchor[g] = zi + end + centered_z = zi - anchor[g] + delta = centered_z - mean[g] + new_mean = mean[g] + w * delta / new_sumw + m2[g] += w * delta * (centered_z - new_mean) + mean[g] = new_mean + sumw[g] = new_sumw + end + 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) + tol = ranktol === nothing ? T(2) * sqrt(eps(T)) : T(ranktol) + @inbounds for g in 1:nlevels + if sumw[g] > zero(T) && one(T) > tol + transforms[intercept, 1, g] = inv(sqrt(sumw[g])) + ranks[g] = 1 + end + centered_sumsq = max(zero(T), m2[g]) + group_mean = anchor[g] + mean[g] + slope_sumsq = centered_sumsq + sumw[g] * abs2(group_mean) + if ranks[g] == 1 && slope_sumsq > zero(T) && + sqrt(centered_sumsq / slope_sumsq) > tol + transforms[slope, 2, g] = inv(sqrt(centered_sumsq)) + transforms[intercept, 2, g] = -group_mean * transforms[slope, 2, g] + ranks[g] = 2 + end + end + @inbounds for i in eachindex(block.refs) + g = block.refs[i] + sqrtw = sqrt(T(weights[i])) + qrows[1, i] = sqrtw * transforms[intercept, 1, g] + qrows[2, i] = sqrtw * (T(z[i]) - anchor[g] - mean[g]) * transforms[slope, 2, g] + end + return transforms, ranks, qrows 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) + 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 + if nchunks == 1 + _build_block_transform_chunk!(transforms, ranks, qrows, block, weights, + ranktol, counts, offsets, perm, 1:nlevels, maxrows) + else + # Groups are disjoint row segments of perm, so chunks can be processed in parallel. + @sync for chunk in _row_chunks(nlevels, nchunks) + 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 - return transforms, ranks, qrows end function _build_block_transform_chunk!(transforms::AbstractArray{T,3}, ranks::AbstractVector{Int}, @@ -250,4 +298,3 @@ function _row_chunks(n::Int, k::Int) end return out end - diff --git a/src/AbstractFixedEffectLinearMap.jl b/src/AbstractFixedEffectLinearMap.jl index 216069e..cb09442 100644 --- a/src/AbstractFixedEffectLinearMap.jl +++ b/src/AbstractFixedEffectLinearMap.jl @@ -10,12 +10,21 @@ ## ############################################################################## +# Concrete subtypes must be mutable and expose two fields used by shared code: +# fes — the original fixed effects, used for the observation count and coefficient recovery; +# plan — the AbsorptionPlan used by the operator, replaceable when weights change. abstract type AbstractFixedEffectLinearMap{T} end +Base.eltype(x::AbstractFixedEffectLinearMap{T}) where {T} = T + Base.adjoint(fem::AbstractFixedEffectLinearMap) = Adjoint(fem) function Base.size(fem::AbstractFixedEffectLinearMap, dim::Integer) - (dim == 1) ? length(fem.fes[1].refs) : (dim == 2) ? _ncoef(fem.plan) : 1 + if dim == 1 + return length(fem.fes[1].refs) + elseif dim == 2 + return sum(block_width(block) * block.n for block in fem.plan.blocks) + else + 1 + end end - -Base.eltype(x::AbstractFixedEffectLinearMap{T}) where {T} = T diff --git a/test/solve.jl b/test/solve.jl index 9c3e505..1b78592 100644 --- a/test/solve.jl +++ b/test/solve.jl @@ -210,6 +210,42 @@ fes_both = [FixedEffect(p1), FixedEffect(p1, interaction = interaction)] 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 + + # The stable moment path also handles weights, a large slope offset, and the + # less common input order in which the slope precedes the intercept. + id_stable = repeat(1:2, inner = 3) + slope_stable = 1.0e7 .+ [0.0, 1.0, 2.0, 3.0, 5.0, 8.0] + weights_stable = Weights([1.0, 2.0, 4.0, 1.5, 2.5, 3.5]) + y_stable = [1.0, -2.0, 4.0, 0.5, 3.0, -1.0] + X_stable = zeros(length(id_stable), 4) + for i in eachindex(id_stable) + g = id_stable[i] + X_stable[i, g] = 1 + X_stable[i, 2 + g] = slope_stable[i] - slope_stable[firstindex(slope_stable)] + end + sqrtw_stable = sqrt.(weights_stable) + r_stable = (sqrtw_stable .* y_stable - + (X_stable .* sqrtw_stable) * (pinv(X_stable .* sqrtw_stable) * + (sqrtw_stable .* y_stable))) ./ sqrtw_stable + for fes_stable in ([FixedEffect(id_stable), FixedEffect(id_stable, interaction = slope_stable)], + [FixedEffect(id_stable, interaction = slope_stable), FixedEffect(id_stable)]) + plan_stable = FixedEffects.AbsorptionPlan(Float64, fes_stable, weights_stable) + @test plan_stable.ranks[1] == [2, 2] + for g in 1:2 + rows = findall(==(g), id_stable) + Q = permutedims(plan_stable.qrows[1][:, rows]) + @test Q' * Q ≈ I atol = 1e-12 + end + @test solve_residuals!(copy(y_stable), fes_stable, weights_stable)[1] ≈ + r_stable atol = 1e-10 + end + + # Two slopes without an intercept still use the general block path. + slope2 = [1.0, 2.0, 4.0, 2.0, 3.0, 5.0] + fes_slopes = [FixedEffect(id_stable, interaction = slope2), + FixedEffect(id_stable, interaction = slope2 .^ 2)] + plan_slopes = FixedEffects.AbsorptionPlan(Float64, fes_slopes, weights_stable) + @test plan_slopes.ranks[1] == [2, 2] end # Independent implementation of the exact one-block projection residual, From 274f1dac6fd1e048f6d085e77d82b35a4f8e778f Mon Sep 17 00:00:00 2001 From: Matthieu Gomez Date: Thu, 27 Aug 2026 17:05:25 -0400 Subject: [PATCH 2/2] Bump version to 3.4.1 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 131ed21..3d75c66 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "FixedEffects" uuid = "c8885935-8500-56a7-9867-7708b20db0eb" -version = "3.4.0" +version = "3.4.1" [deps] GroupedArrays = "6407cd72-fade-4a84-8a1e-56e431fc1533"