From cff135538f58019595ad6ac84e890e54c56fc0df Mon Sep 17 00:00:00 2001 From: Matthieu Gomez Date: Tue, 25 Aug 2026 09:58:24 -0400 Subject: [PATCH] Speed up fixed-effect scatter multiplies Add a six-argument scatter path so forward multiplies fuse the beta scaling with the first fixed-effect scatter on CPU, CUDA, and Metal. This avoids a separate pass over the observation vector and handles beta = 0 and beta = 1 directly. Make CPU solver buffers concrete to reduce hot-path dispatch and allocations. Bump version to 3.3.1. --- Project.toml | 2 +- ext/CUDAExt.jl | 32 ++++++++++++++++++++ ext/MetalExt.jl | 28 ++++++++++++++++++ src/AbstractFixedEffectLinearMap.jl | 17 +++++++++-- src/SolverCPU.jl | 46 +++++++++++++++++++---------- 5 files changed, 105 insertions(+), 20 deletions(-) diff --git a/Project.toml b/Project.toml index 2ef19b4..4ffe88e 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "FixedEffects" uuid = "c8885935-8500-56a7-9867-7708b20db0eb" -version = "3.3.0" +version = "3.3.1" [deps] GroupedArrays = "6407cd72-fade-4a84-8a1e-56e431fc1533" diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl index c398d49..cfac1d9 100644 --- a/ext/CUDAExt.jl +++ b/ext/CUDAExt.jl @@ -69,6 +69,18 @@ function FixedEffects.scatter!(y::CuVector, α::Number, fecoef::CuVector, refs:: @cuda threads=nthreads blocks=nblocks scatter_kernel!(y, α, fecoef, refs, cache) end +function FixedEffects.scatter!(y::CuVector, α::Number, fecoef::CuVector, refs::CuVector, cache::CuVector, β::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 +end + function scatter_kernel!(y, α, fecoef, refs, cache) index = (blockIdx().x - Int32(1)) * blockDim().x + threadIdx().x stride = blockDim().x * gridDim().x @@ -79,6 +91,26 @@ function scatter_kernel!(y, α, fecoef, refs, cache) end 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 + end +end + +function scatter_kernel_scaled!(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] = β * y[i] + α * fecoef[refs[i]] * cache[i] + i += stride + end +end + ############################################################################## ## ## Implement AbstractFixedEffectSolver interface diff --git a/ext/MetalExt.jl b/ext/MetalExt.jl index 070d0bc..c4a7225 100644 --- a/ext/MetalExt.jl +++ b/ext/MetalExt.jl @@ -156,6 +156,18 @@ function FixedEffects.scatter!(y::MtlVector, α::Number, fecoef::MtlVector, refs Metal.@sync @metal threads=nthreads groups=nblocks scatter_kernel!(y, α, fecoef, refs, cache) end +function FixedEffects.scatter!(y::MtlVector, α::Number, fecoef::MtlVector, refs::MtlVector, cache::MtlVector, β::Number) + 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 +end + function scatter_kernel!(y, α, fecoef, refs, cache) i = thread_position_in_grid_1d() if i <= length(y) @@ -164,6 +176,22 @@ function scatter_kernel!(y, α, fecoef, refs, cache) return nothing 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] + end + return nothing +end + +function scatter_kernel_scaled!(y, α, fecoef, refs, cache, β) + i = thread_position_in_grid_1d() + if i <= length(y) + @inbounds y[i] = β * y[i] + α * fecoef[refs[i]] * cache[i] + end + return nothing +end + ############################################################################## diff --git a/src/AbstractFixedEffectLinearMap.jl b/src/AbstractFixedEffectLinearMap.jl index e3ac757..9f19875 100644 --- a/src/AbstractFixedEffectLinearMap.jl +++ b/src/AbstractFixedEffectLinearMap.jl @@ -38,6 +38,12 @@ 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} @@ -51,10 +57,15 @@ end function LinearAlgebra.mul!(y::AbstractVector, fem::AbstractFixedEffectLinearMap, fecoefs::FixedEffectCoefficients, α::Number, β::Number) - rmul!(y, β) + βj = β + any_fe = false for (fecoef, fe, cache) in zip(fecoefs.x, fem.fes, fem.caches) - scatter!(y, α, fecoef, fe.refs, cache) + 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/SolverCPU.jl b/src/SolverCPU.jl index bb0d63a..0b98383 100644 --- a/src/SolverCPU.jl +++ b/src/SolverCPU.jl @@ -4,11 +4,11 @@ ## ############################################################################## -mutable struct FixedEffectLinearMapCPU{T} <: AbstractFixedEffectLinearMap{T} - fes::Vector{<:FixedEffect} - scales::Vector{<:AbstractVector} - caches::Vector{<:AbstractVector} - gathers::Vector{Union{SerialGather, ThreadedGather{Vector{T}}}} +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. @@ -49,7 +49,7 @@ function FixedEffectLinearMapCPU{T}(fes::Vector{<:FixedEffect}) where {T} 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}(fes, scales, caches, gathers) + return FixedEffectLinearMapCPU{T,typeof(fes),typeof(scales),typeof(caches),typeof(gathers)}(fes, scales, caches, gathers) end @@ -92,6 +92,22 @@ function scatter!(y::AbstractVector, α::Number, fecoef::AbstractVector, 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 + ############################################################################## ## @@ -99,15 +115,15 @@ end ## ############################################################################## -mutable struct FixedEffectSolverCPU{T} <: AbstractFixedEffectSolver{T} - m::FixedEffectLinearMapCPU{T} +mutable struct FixedEffectSolverCPU{T,M<:FixedEffectLinearMapCPU{T},C<:FixedEffectCoefficients{Vector{T}}} <: AbstractFixedEffectSolver{T} + m::M weights::AbstractVector - b::AbstractVector{T} - r::AbstractVector{T} - x::FixedEffectCoefficients{<: AbstractVector{T}} - v::FixedEffectCoefficients{<: AbstractVector{T}} - h::FixedEffectCoefficients{<: AbstractVector{T}} - hbar::FixedEffectCoefficients{<: AbstractVector{T}} + b::Vector{T} + r::Vector{T} + x::C + v::C + h::C + hbar::C end @@ -160,5 +176,3 @@ end function copy_internal!(r::AbstractVector, feM::FixedEffectSolverCPU, field::Symbol) copyto!(r, getfield(feM, field)) end - -