Skip to content

rewrite make_zero! and remake_zero! with ntuple over fields #2607

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.

Sign up for GitHub

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

Merged
vchuravy merged 1 commit into from
Sep 26, 2025
Merged

rewrite make_zero! and remake_zero! with ntuple over fields #2607

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions src/typeutils/make_zero.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -572,13 +572,13 @@ end
return nothing
end
push!(seen, prev)
for i in 1:nf
ntuple(Val(nf)) do i
Copy link
Member

@wsmoses wsmoses Sep 22, 2025

Choose a reason for hiding this comment

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

honestly this feels like it should become a generated function (and we can also therefore conditionally create the seen dict if there is > 1 potentially aliasing piece of data)

Copy link
Member Author

@vchuravy vchuravy Sep 22, 2025

Choose a reason for hiding this comment

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

ntuple is a generated function, just one that is limited in scope. I still prefer this over the headache that generated function bring with them otherwise.

I also don't think it would be unreasonable for a user to provide their own make_zero function / remake_zero! function.

Copy link
Member

@wsmoses wsmoses Sep 22, 2025

Choose a reason for hiding this comment

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

yeah we definitely want user definable overrides. I'm just dreaming of a world where we don't need the explicit specialization for make_zero(Array) in addition to make_zero(Array, seen).

in any case we can figure this out later, this is already clearly a good change

Copy link
Member Author

@vchuravy vchuravy Sep 22, 2025

Choose a reason for hiding this comment

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

Also sadly while we might be able to use a generated function, that would only allow us to guard against recursion in the type-domain and not recursion in the value domain :/

Copy link
Member

@wsmoses wsmoses Sep 22, 2025

Choose a reason for hiding this comment

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

Perhaps but things like make_zero of a tuple of floats is sufficiently common that we really want it to use the right thing without the is check

@inline
if isdefined(prev, i)
xi = getfield(prev, i)
SBT = Core.Typeof(xi)
activitystate = active_reg_nothrow(SBT)
if activitystate == AnyState # guaranteed_const
continue
elseif ismutabletype(T) && !ismutabletype(SBT)
yi = make_zero_immutable!(xi, seen)
if Base.isconst(T, i)
Expand All @@ -593,6 +593,7 @@ end
throw(ArgumentError(msg))
end
end
return nothing
end
return nothing
end
Expand All @@ -611,13 +612,13 @@ end
return nothing
end
push!(seen, prev)
for i in 1:nf
ntuple(Val(nf)) do i
@inline
if isdefined(prev, i)
xi = getfield(prev, i)
SBT = Core.Typeof(xi)
activitystate = active_reg_nothrow(SBT)
if activitystate == AnyState # guaranteed_const
continue
elseif ismutabletype(T) && !ismutabletype(SBT)
yi = make_zero_immutable!(xi, seen)
if Base.isconst(T, i)
Expand All @@ -631,6 +632,7 @@ end
EnzymeCore.remake_zero!(xi, seen)
end
end
return nothing
end
return nothing
end
Loading