Speed up fixed-effect scatter multiplies

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"

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

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
+
 
 
 ##############################################################################

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
-

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,22 +92,38 @@ 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
+
 
 ##############################################################################
 ##
 ## Implement AbstractFixedEffectSolver interface
 ##
 ##############################################################################
 
-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
-
-