Fix methodinstance usage and backedges #2199

wsmoses · 2024-12-14T21:56:41Z

Hopefully now that I understand backedges, this fixes pending issues.

Simultaneously this gets rid of a bunch of inner generated functions so hopefully compile time comes down.

Time to first derivative (post precompile) is now:

julia> @time autodiff(Forward, sin, Duplicated(1.0, 2.0))
  0.000014 seconds
(1.0806046117362795,)

julia> @time autodiff(Forward, sin, Duplicated(1.0, 2.0))
  0.000003 seconds
(1.0806046117362795,)

julia> @time autodiff(Forward, sin, Duplicated(1.0, 2.0))
  0.000002 seconds
(1.0806046117362795,)

github-actions · 2024-12-14T22:07:59Z

Benchmark Results

	main	`4dbca9f`...	main/4dbca9f3a41ab7...
basics/overhead	4.34 ± 0.01 ns	4.34 ± 0.92 ns	1
time_to_load	1.15 ± 0.0054 s	1.15 ± 0.0089 s	0.997

Benchmark Plots

A plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR.
Go to "Actions"->"Benchmark a pull request"->[the most recent run]->"Artifacts" (at the bottom).

vchuravy · 2024-12-16T12:58:15Z

src/compiler.jl

@@ -1258,6 +1272,8 @@ Create the methodinstance pair, and lookup the primal return type.
    @nospecialize(TT::Type),
    world::Union{UInt,Nothing} = nothing,
 )
+
+fdsafdsafsa


vchuravy · 2024-12-16T12:59:51Z

src/compiler.jl

+        fwd_sig = Tuple{typeof(EnzymeRules.forward), <:EnzymeRules.FwdConfig, <:Enzyme.EnzymeCore.Annotation, Type{<:Enzyme.EnzymeCore.Annotation},Vararg{Enzyme.EnzymeCore.Annotation}}
+        push!(edges, ccall(:jl_method_table_for, Any, (Any,), fwd_sig)::Core.MethodTable)
+        push!(edges, fwd_sig)
    else
-        push!(edges, GPUCompiler.methodinstance(typeof(Compiler.Interpreter.rule_backedge_holder), Tuple{typeof(EnzymeRules.augmented_primal)}, world))
+        rev_sig = Tuple{typeof(EnzymeRules.augmented_primal), <:EnzymeRules.RevConfig, <:Enzyme.EnzymeCore.Annotation, Type{<:Enzyme.EnzymeCore.Annotation},Vararg{Enzyme.EnzymeCore.Annotation}}
+        push!(edges, ccall(:jl_method_table_for, Any, (Any,), rev_sig)::Core.MethodTable)
+        push!(edges, rev_sig)
+
+        rev_sig = Tuple{typeof(EnzymeRules.reverse), <:EnzymeRules.RevConfig, <:Enzyme.EnzymeCore.Annotation, Union{Type{<:Enzyme.EnzymeCore.Annotation}, Enzyme.EnzymeCore.Active}, Any, Vararg{Enzyme.EnzymeCore.Annotation}}
+        push!(edges, ccall(:jl_method_table_for, Any, (Any,), rev_sig)::Core.MethodTable)
+        push!(edges, rev_sig)
+    end
+
+    ina_sig = Tuple{typeof(EnzymeRules.inactive), Vararg{Any}}
+    push!(edges, ccall(:jl_method_table_for, Any, (Any,), ina_sig)::Core.MethodTable)
+    push!(edges, ina_sig)


This does look much better.

vchuravy · 2024-12-16T13:00:48Z

src/compiler/interpreter.jl

@@ -110,7 +94,7 @@ function rule_backedge_holder_generator(world::UInt, source, self, ft::Type)
    new_ci.slotflags = UInt8[0x00 for i = 1:2]

    # return the codegen world age
-    push!(new_ci.code, Core.Compiler.ReturnNode(0))
+    push!(new_ci.code, Core.Compiler.ReturnNode(world))


Uhm this, on the other hand, looks horribly illegal.

vchuravy · 2024-12-16T13:01:59Z

src/compiler/interpreter.jl

@@ -198,6 +176,37 @@ function EnzymeInterpreter(
    else
        InferenceParams(; unoptimize_throw_blocks=false)
    end
+
+    @static if HAS_INTEGRATED_CACHE


I was also contemplating that Enzyme for 1.11 could simply not use the integrated cache...
But we might have removed some interface...

vchuravy · 2024-12-16T13:02:39Z

src/compiler/tfunc.jl

@@ -2,30 +2,32 @@ import EnzymeCore: Annotation
 import EnzymeCore.EnzymeRules: FwdConfig, RevConfig, forward, augmented_primal, inactive, _annotate_tt

 function has_frule_from_sig(@nospecialize(interp::Core.Compiler.AbstractInterpreter),
-    @nospecialize(TT), sv::Core.Compiler.AbsIntState)::Bool
+    @nospecialize(TT::Type), sv::Core.Compiler.AbsIntState, partialedge::Bool=true)::Bool


What is a partialedge? It seems not used?

wsmoses added 2 commits December 14, 2024 16:56

Fix methodinstance usage and backedges

9e50cc5

fixup

2d3f43b

wsmoses force-pushed the metinst branch from 90130b0 to 2d3f43b Compare December 14, 2024 22:01

fix

67e6626

wsmoses force-pushed the metinst branch from 939ddc0 to 67e6626 Compare December 14, 2024 22:05

fix

cdf9223

wsmoses force-pushed the metinst branch from 6913d41 to cdf9223 Compare December 14, 2024 22:18

wsmoses added 13 commits December 14, 2024 17:39

fix

339316f

fix

a5e5ae4

fix

36632e4

final fix

258c5ed

fix

501b9e8

more tracking

1dc2182

fix

07eead4

fix

060626a

fix

10c8379

fix

5030f93

edges

08d44a9

fix

4f6f452

fix

4dbca9f

wsmoses merged commit ee21090 into main Dec 15, 2024
18 of 29 checks passed

wsmoses deleted the metinst branch December 15, 2024 06:47

vchuravy reviewed Dec 16, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix methodinstance usage and backedges #2199

Fix methodinstance usage and backedges #2199

wsmoses commented Dec 14, 2024 •

edited

Loading

github-actions bot commented Dec 14, 2024 •

edited

Loading

vchuravy Dec 16, 2024

vchuravy Dec 16, 2024

vchuravy Dec 16, 2024

vchuravy Dec 16, 2024

vchuravy Dec 16, 2024

Fix methodinstance usage and backedges #2199

Fix methodinstance usage and backedges #2199

Conversation

wsmoses commented Dec 14, 2024 • edited Loading

github-actions bot commented Dec 14, 2024 • edited Loading

Benchmark Results

Benchmark Plots

vchuravy Dec 16, 2024

Choose a reason for hiding this comment

vchuravy Dec 16, 2024

Choose a reason for hiding this comment

vchuravy Dec 16, 2024

Choose a reason for hiding this comment

vchuravy Dec 16, 2024

Choose a reason for hiding this comment

vchuravy Dec 16, 2024

Choose a reason for hiding this comment

wsmoses commented Dec 14, 2024 •

edited

Loading

github-actions bot commented Dec 14, 2024 •

edited

Loading