WIP: move make_zero to generated (and hopefully type stable/runtime h…

[wsmoses]

…appy) fn

[wsmoses]

@gbaraldi @oscardssmith

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.

diff --git a/src/typeutils/make_zero.jl b/src/typeutils/make_zero.jl
index c72a05f6..0c8a61d6 100644
--- a/src/typeutils/make_zero.jl
+++ b/src/typeutils/make_zero.jl
@@ -216,18 +216,20 @@ end
         for i in 1:nf
             ST = fieldtype(RT, i)
             is_const = Base.isconst(RT, i)
-            push!(exprs, quote
-                if isdefined(prev, $i)
-                    xi = getfield(prev, $i)
-                    T = $(isconcretetype(ST) ? ST : :(Core.Typeof(xi)))
-                    xi_zero = EnzymeCore.make_zero(T, seen, xi, Val(copy_if_inactive))
-                    if $is_const
-                        ccall(:jl_set_nth_field, Cvoid, (Any, Csize_t, Any), y, $(i - 1), xi_zero)
-                    else
-                        setfield!(y, $i, xi_zero)
-                    end
+            push!(
+                exprs, quote
+                    if isdefined(prev, $i)
+                        xi = getfield(prev, $i)
+                        T = $(isconcretetype(ST) ? ST : :(Core.Typeof(xi)))
+                        xi_zero = EnzymeCore.make_zero(T, seen, xi, Val(copy_if_inactive))
+                        if $is_const
+                            ccall(:jl_set_nth_field, Cvoid, (Any, Csize_t, Any), y, $(i - 1), xi_zero)
+                        else
+                            setfield!(y, $i, xi_zero)
+                        end
+                end
                 end
-            end)
+            )
         end
         return quote
             if haskey(seen, prev)
@@ -250,27 +252,31 @@ end
             for i in 1:nf
                 ST = fieldtype(RT, i)
                 sym = Symbol("f_", i)
-                push!(evals, quote
-                    if continue_flag && isdefined(prev, $i)
-                        xi = getfield(prev, $i)
-                        T = $(isconcretetype(ST) ? ST : :(Core.Typeof(xi)))
-                        $sym = EnzymeCore.make_zero(T, seen, xi, Val(copy_if_inactive))
-                        nf_defined = $i
-                    else
-                        continue_flag = false
+                push!(
+                    evals, quote
+                        if continue_flag && isdefined(prev, $i)
+                            xi = getfield(prev, $i)
+                            T = $(isconcretetype(ST) ? ST : :(Core.Typeof(xi)))
+                            $sym = EnzymeCore.make_zero(T, seen, xi, Val(copy_if_inactive))
+                            nf_defined = $i
+                        else
+                            continue_flag = false
+                        end
                     end
-                end)
+                )
             end
             branches = []
             for k in 0:nf
                 args = [Symbol("f_", i) for i in 1:k]
-                push!(branches, quote
-                    if nf_defined == $k
-                        y = $(Expr(:new, RT, args...))
-                        seen[prev] = y
-                        return y
+                push!(
+                    branches, quote
+                        if nf_defined == $k
+                            y = $(Expr(:new, RT, args...))
+                            seen[prev] = y
+                            return y
+                        end
                     end
-                end)
+                )
             end
             return quote
                 if haskey(seen, prev)
@@ -287,16 +293,18 @@ end
             exprs = []
             for i in 1:nf
                 ST = fieldtype(RT, i)
-                push!(exprs, quote
-                    if continue_flag && isdefined(prev, $i)
-                        xi = getfield(prev, $i)
-                        T = $(isconcretetype(ST) ? ST : :(Core.Typeof(xi)))
-                        flds[$i] = EnzymeCore.make_zero(T, seen, xi, Val(copy_if_inactive))
-                        nf_defined = $i
-                    else
-                        continue_flag = false
+                push!(
+                    exprs, quote
+                        if continue_flag && isdefined(prev, $i)
+                            xi = getfield(prev, $i)
+                            T = $(isconcretetype(ST) ? ST : :(Core.Typeof(xi)))
+                            flds[$i] = EnzymeCore.make_zero(T, seen, xi, Val(copy_if_inactive))
+                            nf_defined = $i
+                        else
+                            continue_flag = false
+                        end
                     end
-                end)
+                )
             end
             return quote
                 if haskey(seen, prev)

[github-actions]

Benchmark Results

	main	c5a69eb...	main / c5a69eb...
basics/make_zero/namedtuple	0.0584 ± 0.0039 μs	0.0581 ± 0.0059 μs	1 ± 0.12
basics/make_zero/struct	0.262 ± 0.0075 μs	0.147 ± 0.0059 μs	1.78 ± 0.088
basics/overhead	4.54 ± 0.01 ns	4.54 ± 0.01 ns	1 ± 0.0031
basics/remake_zero!/namedtuple	0.232 ± 0.012 μs	0.232 ± 0.015 μs	0.999 ± 0.083
basics/remake_zero!/struct	0.231 ± 0.01 μs	0.23 ± 0.012 μs	1 ± 0.07
fold_broadcast/multidim_sum_bcast/1D	10.3 ± 0.42 μs	10.4 ± 0.71 μs	0.99 ± 0.079
fold_broadcast/multidim_sum_bcast/2D	12.6 ± 0.36 μs	12.7 ± 0.39 μs	0.991 ± 0.042
time_to_load	1.11 ± 0.017 s	1.12 ± 0.0051 s	0.994 ± 0.016

Benchmark Plots

A plot of the benchmark results has been uploaded as an artifact at https://github.com/EnzymeAD/Enzyme.jl/actions/runs/26251301384/artifacts/7146655409.