Skip to content

Reduce allocation in mass matrix adaptor #427

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Conversation

ErikQQY
Copy link
Collaborator

@ErikQQY ErikQQY commented Apr 13, 2025

Instead of allocating using zeros and ones when resizing the adaptor, we can directly use in-place operations like resize! and fill! to reduce allocations.

@yebai yebai requested a review from devmotion April 17, 2025 21:09
@@ -102,20 +102,50 @@ function WelfordVar(sz::Union{Tuple{Int},Tuple{Int,Int}}; kwargs...)
return WelfordVar{Float64}(sz; kwargs...)
end

function Base.resize!(wv::WelfordVar, θ::AbstractVecOrMat{T}) where {T<:AbstractFloat}
if size(θ) != size(wv.var)
function Base.resize!(wv::WelfordVar, θ::AbstractMatrix{T}) where {T<:AbstractFloat}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IMO this function must be renamed, it's a misuse of Base.resize! as the second argument of resize! is supposed to be an integer. It should just be an internal function.

n = length(θ)
if n != size(wv.var, 1)
@assert wv.n == 0 "Cannot resize a var estimator when it contains samples."
__resize_fill!(wv.μ, n, T(0))
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't really see the need for a custom function, it seems this is just

Suggested change
__resize_fill!(wv.μ, n, T(0))
fill!(resize!(wv.μ, n), T(0))

Comment on lines +136 to +140
if length_diff > 0
append!(x, Vector{T}(undef, length_diff))
else
resize!(x, new_length)
end
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both branches seem to be just

Suggested change
if length_diff > 0
append!(x, Vector{T}(undef, length_diff))
else
resize!(x, new_length)
end
resize!(x, new_length)

? As mentioned above, in that case fill!(resize!(..), ..) would be simpler than a separate function IMO.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants