Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

proper support for interpreting opaque closures #593

Merged
merged 6 commits into from
Sep 10, 2024
Merged

Conversation

simeonschaub
Copy link
Collaborator

@simeonschaub simeonschaub commented Oct 18, 2023

fixes #537

@codecov
Copy link

codecov bot commented Oct 18, 2023

Codecov Report

Attention: 4 lines in your changes are missing coverage. Please review.

Comparison is base (7beca92) 86.54% compared to head (ae687a0) 84.15%.
Report is 1 commits behind head on master.

❗ Current head ae687a0 differs from pull request most recent head ca7a028. Consider uploading reports for the commit ca7a028 to get more accurate results

Files Patch % Lines
src/builtins.jl 33.33% 2 Missing ⚠️
src/construct.jl 66.66% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##           master     #593      +/-   ##
==========================================
- Coverage   86.54%   84.15%   -2.40%     
==========================================
  Files          12       12              
  Lines        2542     2404     -138     
==========================================
- Hits         2200     2023     -177     
- Misses        342      381      +39     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@simeonschaub
Copy link
Collaborator Author

@oxinabox Any chance you could give this a try some time for your Diffractor use case? It's a bit hard to rely on CI currently due to #580

@oxinabox
Copy link
Contributor

Sorry yeah, I was going to then i fixed the bug I was needing to debug.
It is in my TODOs, I get a ping from slack each day reminding me

@oxinabox
Copy link
Contributor

oxinabox commented Nov 6, 2023

Finally tied this and I get an error

1|debug> nc
ERROR: optimization of inferred code not implemented
Stacktrace:
 [1] (::DAECompiler.JITOpaqueClosure{Tuple{Vector{Float64}, Nothing, Float64}, DAECompiler.var"#163#168"{TransformedIRODESystem, Vector{Int64}, Vector{Int64}, DAECompiler.DebugConfig}, Core.OpaqueClosure{Tuple{Vector{Float64}, Nothing, Float64}, Tuple{Vector{Union{Nothing, Float64}}, Vector{Union{Nothing, Float64}}}}})(args::Tuple{Vector{Float64}, Vector{Float64}, Float64})
...

and the debugger exits

@simeonschaub
Copy link
Collaborator Author

Ok, let's track this down! Do you have a reproducer or is this dependent on internal code?

@oxinabox
Copy link
Contributor

oxinabox commented Nov 6, 2023

Its dependent on internal code.
Let me see if I can make a reproducer that isn't

@oxinabox
Copy link
Contributor

oxinabox commented Nov 6, 2023

this reproduces it. It could probably be minimized more though

using Debugger

function get_toplevel_mi_from_ir(ir, _module::Module)
    mi = ccall(:jl_new_method_instance_uninit, Ref{Core.MethodInstance}, ());
    mi.specTypes = Tuple{ir.argtypes...}
    mi.def = _module
    return mi
end

function infer_ir!(ir, interp::Core.Compiler.AbstractInterpreter, mi::Core.Compiler.MethodInstance)
    method_info = Core.Compiler.MethodInfo(#=propagate_inbounds=#true, nothing)
    min_world = world = Core.Compiler.get_world_counter(interp)
    max_world = Base.get_world_counter()
    irsv = Core.Compiler.IRInterpretationState(interp, method_info, ir, mi, ir.argtypes, world, min_world, max_world)
    rt = Core.Compiler._ir_abstract_constant_propagation(interp, irsv)
    return ir
end


# add overloads from Core.Compiler into Base
Base.iterate(compact::Core.Compiler.IncrementalCompact, state) = Core.Compiler.iterate(compact, state)
Base.iterate(compact::Core.Compiler.IncrementalCompact) = Core.Compiler.iterate(compact)
Base.getindex(c::Core.Compiler.IncrementalCompact, args...) = Core.Compiler.getindex(c, args...)
Base.setindex!(c::Core.Compiler.IncrementalCompact, args...) = Core.Compiler.setindex!(c, args...)
Base.setindex!(i::Core.Compiler.Instruction, args...) = Core.Compiler.setindex!(i, args...)


###################################
# Demo

function foo(x)
    a = sin(x+pi/2)
    b = cos(x)
    return a - b
end



input_ir = first(only(Base.code_ircode(foo, Tuple{Float64})))
ir = Core.Compiler.copy(input_ir)
empty!(ir.argtypes)
push!(ir.argtypes, Tuple{})  # the function object itself
push!(ir.argtypes, Float64)  # x 
compact = Core.Compiler.IncrementalCompact(ir)
for ((_, idx), inst) in compact
    ssa = Core.SSAValue(idx)
    if Meta.isexpr(inst, :invoke)
        # we can insert nodes, lets print the function objects
        Core.Compiler.insert_node_here!(
            compact,
            Core.Compiler.NewInstruction(
                Expr(:call, println, inst.args[2]),
                Any, # type
                Core.Compiler.NoCallInfo(), # call info
                Int32(1), # line
                Core.Compiler.IR_FLAG_REFINED  # flag
            )
        )
        compact[ssa][:inst] = Expr(:call, inst.args[2:end]...)
        compact[ssa][:type] = Any
        compact[ssa][:flag] |= Core.Compiler.IR_FLAG_REFINED        
    end
    
end
ir = Core.Compiler.finish(compact)
ir = Core.Compiler.compact!(ir)
interp = Core.Compiler.NativeInterpreter()
mi = get_toplevel_mi_from_ir(ir, @__MODULE__);
ir = infer_ir!(ir, interp, mi)
inline_state = Core.Compiler.InliningState(interp)
ir = Core.Compiler.ssa_inlining_pass!(ir, inline_state, #=propagate_inbounds=#true)
ir = Core.Compiler.compact!(ir)
ir = Core.Compiler.sroa_pass!(ir, inline_state)
ir = Core.Compiler.adce_pass!(ir, inline_state)
ir = Core.Compiler.compact!(ir)
f1 = Core.OpaqueClosure(ir; do_compile=true)
f1(1.2)

@enter f1(1.2)

@simeonschaub
Copy link
Collaborator Author

I'm assuming this is latest Julia nightly?

@oxinabox
Copy link
Contributor

oxinabox commented Nov 6, 2023

11 days old nightly, but I assume it reproduces on nightly as well

@simeonschaub
Copy link
Collaborator Author

Ok, I think I have an idea what's happening here now. So you're putting optimized IR into an opaque closure, but unfortunately JuliaInterpreter can't handle that rn (at least in the general case) because it would have to deal with PhiNodes and such.

So what we'd need is either implement handling for interpreting post-opt IR (basically continuing #491) or an inverse for slot2reg I think (Maybe someone already implemented that?). For now, we should probably call the compiled oc if we discover the IR has been optimized so we don't error but that probably doesn't help you much then.

Not confident on all this stuff though, maybe @aviatesk or @Keno have some ideas?

@aviatesk
Copy link
Member

aviatesk commented Nov 6, 2023

We have ir_to_codeinf!, which converts back IRCode to CodeInfo, although it does preserve post-opt IR elements like PhiNode. AFAIR there is no utilities like an inverse of slot2reg. We would need to implement a new JuliaInterpreter pass for interpreting them..

@simeonschaub
Copy link
Collaborator Author

Yeah, I think the code is saved as a CodeInfo already, but that's what I was afraid of. Thanks for confirming!

@simeonschaub
Copy link
Collaborator Author

@aviatesk Does this look reasonable to you?

@simeonschaub
Copy link
Collaborator Author

Bump

@simeonschaub simeonschaub merged commit 958bd73 into master Sep 10, 2024
9 of 10 checks passed
@simeonschaub simeonschaub deleted the sds/opaque branch September 10, 2024 14:37
@testset "opaque closures" begin
g(x) = 3x
f = Base.Experimental.@opaque x -> g(x)
@test @interpret f(4) == 12
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
@test @interpret f(4) == 12
@test (@interpret f(4)) == 12

Otherwise this interprets the call to ==.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, good catch! I guess it doesn't really matter whether we also interpret the == here, but I can make a follow-up PR to make this more consistent.

@aviatesk
Copy link
Member

This looks very reasonable. Thanks so much for tackling this and sorry for being very late to give a review.

simeonschaub added a commit that referenced this pull request Sep 11, 2024
ref #593 (comment)

Co-authored-by: Shuhei Kadowaki <aviatesk@gmail.com>
@simeonschaub simeonschaub mentioned this pull request Sep 11, 2024
simeonschaub added a commit that referenced this pull request Sep 11, 2024
ref #593 (comment)

Co-authored-by: Shuhei Kadowaki <aviatesk@gmail.com>
aviatesk added a commit that referenced this pull request Sep 11, 2024
ref #593 (comment)

Co-authored-by: Shuhei Kadowaki <aviatesk@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Opaque Closures not supported.
3 participants