Update iterate unwrap

EnzymeAD · Dec 30, 2024 · 960bc0f · 960bc0f
1 parent fb36b65
commit 960bc0f
Showing 1 changed file with 19 additions and 19 deletions.
diff --git a/src/rules/jitrules.jl b/src/rules/jitrules.jl
@@ -817,31 +817,31 @@ function push_if_not_ref(
     return darg
 end
 
+@inline function push_inner(@nospecialize(vals), @nospecialize(arg), @nospecialize(darg))
+    ty = Core.Typeof(arg)
+    actreg = active_reg_nothrow(ty, Val(nothing))
+    if actreg == AnyState
+        Const(arg)
+    elseif actreg == ActiveState
+        Active(arg)
+    elseif actreg == MixedState
+        darg = Base.inferencebarrier(darg)
+        MixedDuplicated(
+            arg,
+            push_if_not_ref(Val(reverse), vals, darg, ty)::Base.RefValue{ty},
+        )
+    else
+        Duplicated(arg, darg)
+    end
+end
+
 @inline function iterate_unwrap_augfwd_dup(
     ::Val{reverse},
     vals,
     args,
     dargs,
 ) where {reverse}
-    ntuple(Val(length(args))) do i
-        Base.@_inline_meta
-        arg = args[i]
-        ty = Core.Typeof(arg)
-        actreg = active_reg_nothrow(ty, Val(nothing))
-        if actreg == AnyState
-            Const(arg)
-        elseif actreg == ActiveState
-            Active(arg)
-        elseif actreg == MixedState
-            darg = Base.inferencebarrier(dargs[i])
-            MixedDuplicated(
-                arg,
-                push_if_not_ref(Val(reverse), vals, darg, ty)::Base.RefValue{ty},
-            )
-        else
-            Duplicated(arg, dargs[i])
-        end
-    end
+    map(Base.Fix1(push_inner, vals), args, dargs)
 end
 
 @inline function iterate_unwrap_augfwd_batchdup(