From 960bc0f9eaa8c51cbad938596d7bc5f0d089159c Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Mon, 30 Dec 2024 18:49:35 -0500 Subject: [PATCH 1/5] Update iterate unwrap --- src/rules/jitrules.jl | 38 +++++++++++++++++++------------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/src/rules/jitrules.jl b/src/rules/jitrules.jl index 3ca0d0a3cd..2387647dad 100644 --- 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( From 1995456532b0a515c48a522570c57aa3f3b2e841 Mon Sep 17 00:00:00 2001 From: "William S. Moses" Date: Mon, 30 Dec 2024 19:10:21 -0500 Subject: [PATCH 2/5] fix --- src/rules/jitrules.jl | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/rules/jitrules.jl b/src/rules/jitrules.jl index 2387647dad..70aa7a5d5a 100644 --- a/src/rules/jitrules.jl +++ b/src/rules/jitrules.jl @@ -817,7 +817,11 @@ function push_if_not_ref( return darg end -@inline function push_inner(@nospecialize(vals), @nospecialize(arg), @nospecialize(darg)) +struct PushInnerStruct{Vals} + vals::Vals +end + +@inline function (v::PushInnerStruct)(@nospecialize(arg), @nospecialize(darg)) ty = Core.Typeof(arg) actreg = active_reg_nothrow(ty, Val(nothing)) if actreg == AnyState @@ -828,7 +832,7 @@ end darg = Base.inferencebarrier(darg) MixedDuplicated( arg, - push_if_not_ref(Val(reverse), vals, darg, ty)::Base.RefValue{ty}, + push_if_not_ref(Val(reverse), v.vals, darg, ty)::Base.RefValue{ty}, ) else Duplicated(arg, darg) @@ -841,7 +845,7 @@ end args, dargs, ) where {reverse} - map(Base.Fix1(push_inner, vals), args, dargs) + map(PushInnerStruct(vals), args, dargs) end @inline function iterate_unwrap_augfwd_batchdup( From 6ddbad617d91884f4da108a1214331e4c74996d5 Mon Sep 17 00:00:00 2001 From: William Moses Date: Mon, 30 Dec 2024 20:18:53 -0500 Subject: [PATCH 3/5] Update jitrules.jl --- src/rules/jitrules.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/rules/jitrules.jl b/src/rules/jitrules.jl index 70aa7a5d5a..b36accc2ca 100644 --- a/src/rules/jitrules.jl +++ b/src/rules/jitrules.jl @@ -817,11 +817,11 @@ function push_if_not_ref( return darg end -struct PushInnerStruct{Vals} +struct PushInnerStruct{reverse, Vals} vals::Vals end -@inline function (v::PushInnerStruct)(@nospecialize(arg), @nospecialize(darg)) +@inline function (v::PushInnerStruct{reverse})(@nospecialize(arg), @nospecialize(darg)) ty = Core.Typeof(arg) actreg = active_reg_nothrow(ty, Val(nothing)) if actreg == AnyState @@ -845,7 +845,7 @@ end args, dargs, ) where {reverse} - map(PushInnerStruct(vals), args, dargs) + map(PushInnerStruct{reverse, typeof(val)}(vals), args, dargs) end @inline function iterate_unwrap_augfwd_batchdup( From aa37f0cbfe03297b269f74af19f1d75edf178ad2 Mon Sep 17 00:00:00 2001 From: William Moses Date: Tue, 31 Dec 2024 00:33:10 -0500 Subject: [PATCH 4/5] Update jitrules.jl --- src/rules/jitrules.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/rules/jitrules.jl b/src/rules/jitrules.jl index b36accc2ca..dd7ffc8c63 100644 --- a/src/rules/jitrules.jl +++ b/src/rules/jitrules.jl @@ -845,7 +845,7 @@ end args, dargs, ) where {reverse} - map(PushInnerStruct{reverse, typeof(val)}(vals), args, dargs) + map(PushInnerStruct{reverse, typeof(vals)}(vals), args, dargs) end @inline function iterate_unwrap_augfwd_batchdup( From 3145f944a5b36f8da8b5a58590f271449e95a024 Mon Sep 17 00:00:00 2001 From: William Moses Date: Tue, 31 Dec 2024 10:29:29 -0500 Subject: [PATCH 5/5] Update jitrules.jl --- src/rules/jitrules.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/rules/jitrules.jl b/src/rules/jitrules.jl index dd7ffc8c63..724eb4fe7a 100644 --- a/src/rules/jitrules.jl +++ b/src/rules/jitrules.jl @@ -821,7 +821,7 @@ struct PushInnerStruct{reverse, Vals} vals::Vals end -@inline function (v::PushInnerStruct{reverse})(@nospecialize(arg), @nospecialize(darg)) +@inline function (v::PushInnerStruct{reverse})(@nospecialize(arg), @nospecialize(darg)) where reverse ty = Core.Typeof(arg) actreg = active_reg_nothrow(ty, Val(nothing)) if actreg == AnyState