Speed up fixed-effect absorption

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"

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

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

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
-

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

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,