Add regression test for Julia 1.12 custom-rule calling convention rewrite (#3196)

* Initial plan

* Add callconv regression test

* fix

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: Billy Moses <wmoses@google.com>

Author: Copilot

src/rules/customrules.jl

@@ -736,22 +736,47 @@ function enzyme_custom_setup_args(
                     @assert ival !== nothing
 
                     for idx = 1:width
+                        local_shadow_root = if roots_ival !== nothing
+                            (width == 1) ? roots_ival : extract_value!(B, roots_ival, idx - 1)
+                        end
+
                         if !is_constant_value(gutils, op)
                             ev =
                                 (width == 1) ? ptr_val : extract_value!(B, ptr_val, idx - 1)
-                            ld = load!(B, iarty, ev, "rules_shadow_load")
-                            metadata(ld)["enzyme_mustcache"] = MDNode(LLVM.Metadata[])
+                            
+                            ld = if uncache_arg
+                                if !reverse
+                                    ld0 = load!(B, iarty, ev, "rules_shadow_load")
+                                    metadata(ld0)["enzyme_mustcache"] = MDNode(LLVM.Metadata[])
+                                    if roots_op != nothing
+                                        if uncacheable[arg.arg_i + 1] != 0
+                                            ld0 = recombine_value!(B, ld0, local_shadow_root)
+                                        else
+                                            ld0 = nullify_rooted_values!(B, ld0)
+                                        end
+                                    end
+                                    push!(byval_tapes, ld0)
+                                    ld0
+                                else
+                                    @assert tape isa LLVM.Value
+                                    ld0 = extract_value!(B, tape, length(byval_tapes), "shadow_roots_op_extract_v1_")
+                                    @assert value_type(ld0) == iarty
+                                    push!(byval_tapes, ld0)
+                                    ld0
+                                end
+                            else
+                                ld0 = load!(B, iarty, ev, "rules_shadow_load")
+                                metadata(ld0)["enzyme_mustcache"] = MDNode(LLVM.Metadata[])
+                                ld0
+                            end
+
                             ival = (width == 1) ? ld : insert_value!(B, ival, ld, idx - 1)
                             @assert ival !== nothing
                         else
                             ival = (width == 1) ? ptr_val : insert_value!(B, ival, ptr_val, idx - 1)
                             @assert ival !== nothing
                         end
 
-                        local_shadow_root = if roots_ival !== nothing
-                            (width == 1) ? roots_ival : extract_value!(B, roots_ival, idx - 1)
-                        end
-
                         if shadow_roots !== nothing
 
                             for r = 1:n_primal_roots

test/callconv.jl

@@ -1,4 +1,5 @@
 using Enzyme, Test
+using Enzyme: EnzymeRules
 
 @noinline function force_stup(A)
     A11 = A[];
@@ -141,3 +142,113 @@ g_test(p) = sum(Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), inn
     @test dp ≈ [3.18, 2.68, 2.82]
 end
 
+abstract type AbstractDomainCallConv end
+
+mutable struct InnerPlanCallConv
+    b::Float64
+end
+
+struct MyPlanCallConv
+    a::Int64
+    b::Float64
+    c::Int32
+    d::Int32
+    plan::InnerPlanCallConv
+    phases::Vector{Float64}
+    indices::Tuple{Vector{Int64}, Vector{Int64}}
+    h::Int64
+end
+
+struct MyDomainCallConv <: AbstractDomainCallConv
+    plan::MyPlanCallConv
+end
+
+@noinline forward_plan_callconv(g::AbstractDomainCallConv) = getfield(g, :plan)
+EnzymeRules.inactive(::typeof(forward_plan_callconv), args...) = nothing
+
+@noinline getplan_callconv(p::MyPlanCallConv) = getfield(p, :plan)
+EnzymeRules.inactive(::typeof(getplan_callconv), args...) = nothing
+
+@noinline getindices_callconv(p::MyPlanCallConv) = getfield(p, :indices)
+EnzymeRules.inactive(::typeof(getindices_callconv), args...) = nothing
+
+@noinline function my_nuft_callconv!(out, A, b)
+    out .= b .* A.b
+    return nothing
+end
+
+function EnzymeRules.augmented_primal(
+    config::EnzymeRules.RevConfigWidth,
+    func::Const{typeof(my_nuft_callconv!)},
+    ::Type{<:Const},
+    out::EnzymeRules.Annotation,
+    A::EnzymeRules.Annotation,
+    b::EnzymeRules.Annotation,
+)
+    primal = EnzymeRules.needs_primal(config) ? out.val : nothing
+    shadow = EnzymeRules.needs_shadow(config) ? out.dval : nothing
+    func.val(out.val, A.val, b.val)
+    return EnzymeRules.AugmentedReturn(primal, shadow, nothing)
+end
+
+function EnzymeRules.reverse(
+    config::EnzymeRules.RevConfigWidth,
+    ::Const{typeof(my_nuft_callconv!)},
+    ::Type{RT},
+    tape,
+    out::EnzymeRules.Annotation,
+    A::EnzymeRules.Annotation,
+    b::EnzymeRules.Annotation,
+) where {RT}
+    b.dval .+= out.dval .* A.val.b
+    fill!(out.dval, 0)
+    return (nothing, nothing, nothing)
+end
+
+function applyphases_callconv!(vis, phases)
+    for i in eachindex(vis, phases)
+        vis[i] *= phases[i]
+    end
+    return vis
+end
+
+@inline function applyft_callconv(p, img)
+    vis = similar(img)
+    plan = getplan_callconv(p)
+    iminds, visinds = getindices_callconv(p)
+    for i in eachindex(iminds, visinds)
+        imind = iminds[i]
+        visind = visinds[i]
+        vis_view = @view(vis[visind:visind])
+        img_view = @view(img[imind:imind])
+        my_nuft_callconv!(vis_view, plan, img_view)
+    end
+    applyphases_callconv!(vis, p.phases)
+    return vis
+end
+
+@noinline function visibilitymap_numeric_callconv(grid::AbstractDomainCallConv, img::Vector{Float64})
+    return applyft_callconv(forward_plan_callconv(grid), img)
+end
+
+@noinline function foo_callconv(grid::AbstractDomainCallConv, img)
+    return sum(visibilitymap_numeric_callconv(grid, img))
+end
+
+@testset "Custom rule calling conv rewrite" begin
+    inner = InnerPlanCallConv(2.0)
+    plan = MyPlanCallConv(1, 2.0, 3, 4, inner, [2.0, 3.0, 4.0], ([1, 2, 3], [1, 2, 3]), 7)
+    grid = MyDomainCallConv(plan)
+    img = [1.0, 2.0, 3.0]
+    dimg = zeros(3)
+
+    Enzyme.autodiff(
+        Enzyme.set_runtime_activity(Enzyme.Reverse),
+        foo_callconv,
+        Active,
+        Const(grid),
+        Duplicated(img, dimg),
+    )
+
+    @test dimg == [4.0, 6.0, 8.0]
+end