World backedge holder #2183

wsmoses · 2024-12-07T00:10:05Z

No description provided.

wsmoses · 2024-12-07T01:01:33Z

@aviatesk @vchuravy @gbaraldi so the tuple backedge doesn't work. What's the right way to add a generic backedge?

wsmoses · 2024-12-07T01:14:38Z

@vchuravy this really seems to want a MI: https://github.com/JuliaLang/julia/pull/32237/files

wsmoses · 2024-12-07T01:15:07Z

@vtjnash idk if you have any thoughts here either

github-actions · 2024-12-07T05:39:16Z

Benchmark Results

	main	`d8f8a90`...	main/d8f8a90616dc9a...
basics/overhead	4.64 ± 0.01 ns	4.64 ± 0.01 ns	1
time_to_load	1.14 ± 0.015 s	1.13 ± 0.0014 s	1.01

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-10T05:50:37Z

src/compiler.jl

@@ -5554,7 +5554,22 @@ function thunk_generator(world::UInt, source::LineNumberNode, @nospecialize(FA::
    # new_ci.min_world = min_world[]
    new_ci.min_world = world
    new_ci.max_world = max_world[]
-    new_ci.edges = Core.MethodInstance[mi]
+
+    edges = Core.MethodInstance[mi]


edges = Any[mi]

vchuravy · 2024-12-10T05:54:55Z

@vchuravy this really seems to want a MI: https://github.com/JuliaLang/julia/pull/32237/files

Yeah, Jameson's comment there is correct. This should allow edges more diverse than just the mi.
as implemented correctly in the runtime https://github.com/JuliaLang/julia/blob/7192df7e900a6e4562e1670f544b8be8878ee359/src/method.c#L847-L863

vchuravy · 2024-12-10T05:58:30Z

src/compiler/interpreter.jl

+begin
+    # Forward-rule catch all
+    fwd_rule_be = GPUCompiler.methodinstance(typeof(rule_backedge_holder), Tuple{typeof(EnzymeRules.forward)})
+    # Reverse-rule catch all
+    rev_rule_be = GPUCompiler.methodinstance(typeof(rule_backedge_holder), Tuple{typeof(EnzymeRules.augmented_primal)})
+    # Inactive-rule catch all
+    ina_rule_be = GPUCompiler.methodinstance(typeof(rule_backedge_holder), Tuple{typeof(EnzymeRules.inactive)})
+    # All other derivative-related catch all (just for autodiff, not inference), including inactive_noinl, noalias, and inactive_type
+    gen_rule_be = GPUCompiler.methodinstance(typeof(rule_backedge_holder), Tuple{Val{0}})
+
+
+    fwd_sig = Tuple{typeof(EnzymeRules.forward), <:EnzymeRules.FwdConfig, <:Enzyme.EnzymeCore.Annotation, Type{<:Enzyme.EnzymeCore.Annotation},Vararg{Enzyme.EnzymeCore.Annotation}}
+    EnzymeRules.add_mt_backedge!(fwd_rule_be, ccall(:jl_method_table_for, Any, (Any,), fwd_sig)::Core.MethodTable, fwd_sig)
+
+    rev_sig = Tuple{typeof(EnzymeRules.augmented_primal), <:EnzymeRules.RevConfig, <:Enzyme.EnzymeCore.Annotation, Type{<:Enzyme.EnzymeCore.Annotation},Vararg{Enzyme.EnzymeCore.Annotation}}
+    EnzymeRules.add_mt_backedge!(rev_rule_be, ccall(:jl_method_table_for, Any, (Any,), rev_sig)::Core.MethodTable, rev_sig)
+
+
+    for ina_sig in (
+        Tuple{typeof(EnzymeRules.inactive), Vararg{Any}},
+    )
+        EnzymeRules.add_mt_backedge!(ina_rule_be, ccall(:jl_method_table_for, Any, (Any,), ina_sig)::Core.MethodTable, ina_sig)
+    end
+
+    for gen_sig in (
+        Tuple{typeof(EnzymeRules.inactive_noinl), Vararg{Any}},
+        Tuple{typeof(EnzymeRules.noalias), Vararg{Any}},
+        Tuple{typeof(EnzymeRules.inactive_type), Type},
+    )
+        EnzymeRules.add_mt_backedge!(gen_rule_be, ccall(:jl_method_table_for, Any, (Any,), gen_sig)::Core.MethodTable, gen_sig)
+    end
+end


What does this code aim to do? IIUC backedges are only triggered once and after that pruned. So this would only survive one invalidation.

ah so just move this inside our interp then? Basically this does the split we discussed, since I failed to add the generic methodtable backedge earlier =/. (I genuinely tried to add it there and failed, so you're going to have to help me add a non methodinstance backedge)

so all functions get a backedge to rule_backedge_holder. And rule_backedge_holder gets a backedge from .

vchuravy · 2024-12-10T06:01:06Z

src/compiler/interpreter.jl

+        if interp.forward_rules
+            Core.Compiler.add_backedge!(sv, GPUCompiler.methodinstance(typeof(Enzyme.Compiler.Interpreter.rule_backedge_holder), Tuple{typeof(EnzymeRules.forward)}, interp.world)::Core.MethodInstance)
+            Enzyme.Compiler.Interpreter.rule_backedge_holder(Base.inferencebarrier(EnzymeRules.forward))
+        end
+        if interp.reverse_rules
+            Core.Compiler.add_backedge!(sv, GPUCompiler.methodinstance(typeof(Enzyme.Compiler.Interpreter.rule_backedge_holder), Tuple{typeof(EnzymeRules.augmented_primal)}, interp.world)::Core.MethodInstance)
+            Enzyme.Compiler.Interpreter.rule_backedge_holder(Base.inferencebarrier(EnzymeRules.augmented_primal))
+        end
+        Core.Compiler.add_backedge!(sv, GPUCompiler.methodinstance(typeof(Enzyme.Compiler.Interpreter.rule_backedge_holder), Tuple{typeof(EnzymeRules.inactive)}, interp.world)::Core.MethodInstance)
+        Enzyme.Compiler.Interpreter.rule_backedge_holder(Base.inferencebarrier(typeof(EnzymeRules.inactive)))


This should have had a comment that a mi->mi edge is cheaper than a mi->mt one.

And generally within inference the right thing is to use something like abstract_call. Instead of using the hacky version of GPUCompiler.methodinstance

yeah sorry I don't know the absint well enough to know how to set that up, what's the right way to do?

vchuravy · 2024-12-10T06:01:25Z

test/ruleinvalidation.jl

+@static if VERSION < v"1.11-"
+    @test_broken autodiff(Forward, call_issue696, Duplicated(1.0, 1.0))[1] ≈ 0.0
+else
+    @test autodiff(Forward, call_issue696, Duplicated(1.0, 1.0))[1] ≈ 0.0
+end


Yeah that tells me we are doing something wrong again.

yeah...........but at least not stupid compile time , but yea.....

vchuravy · 2024-12-10T06:03:16Z

src/compiler/interpreter.jl

            callinfo = NoInlineCallInfo(callinfo, atype, :inactive)
        else
-            # 2. Check if rule is defined
-            has_rule = get!(interp.rules_cache, specTypes) do
-                if interp.forward_rules && has_frule_from_sig(interp, specTypes, sv)


Probably should have then deleted the code in tfuncs.jl since this was the sole user.

World backedge holder

c568216

wsmoses requested a review from vchuravy December 7, 2024 00:10

fix

278207e

wsmoses and others added 17 commits December 6, 2024 19:18

fixup

15e64d8

Update compiler.jl

ea1c69b

Update compiler.jl

159d33b

Update interpreter.jl

9918798

Update interpreter.jl

952f277

Update interpreter.jl

c7ef6b6

Update interpreter.jl

ec76838

Update interpreter.jl

e680d1b

Update interpreter.jl

d64d0c9

try2

30ef92a

nothing works rip

9a1d79b

more test

df9ed1c

hn

4846fdb

keep trying

da76eb2

more

ca722d7

fix

4292697

fix

78d41eb

wsmoses added 3 commits December 6, 2024 23:40

isapplic

76f9195

fix2

04b169c

mark broken

d8f8a90

wsmoses merged commit 8e10a0a into main Dec 7, 2024
21 of 30 checks passed

wsmoses deleted the rbh branch December 7, 2024 17:05

vchuravy reviewed Dec 10, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

World backedge holder #2183

World backedge holder #2183

wsmoses commented Dec 7, 2024

wsmoses commented Dec 7, 2024

wsmoses commented Dec 7, 2024

wsmoses commented Dec 7, 2024

github-actions bot commented Dec 7, 2024 •

edited

Loading

vchuravy Dec 10, 2024

vchuravy commented Dec 10, 2024

vchuravy Dec 10, 2024

wsmoses Dec 11, 2024

vchuravy Dec 10, 2024

vchuravy Dec 10, 2024

wsmoses Dec 11, 2024

vchuravy Dec 10, 2024

wsmoses Dec 11, 2024

vchuravy Dec 10, 2024

World backedge holder #2183

World backedge holder #2183

Conversation

wsmoses commented Dec 7, 2024

wsmoses commented Dec 7, 2024

wsmoses commented Dec 7, 2024

wsmoses commented Dec 7, 2024

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

Benchmark Results

Benchmark Plots

Choose a reason for hiding this comment

vchuravy commented Dec 10, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot commented Dec 7, 2024 •

edited

Loading