Skip to content

Interpreter fixes (prerequisite to CUDA support)#364

Closed
wsmoses wants to merge 28 commits intomainfrom
interp
Closed

Interpreter fixes (prerequisite to CUDA support)#364
wsmoses wants to merge 28 commits intomainfrom
interp

Conversation

@wsmoses
Copy link
Copy Markdown
Member

@wsmoses wsmoses commented Dec 11, 2024

and also should fix the nested gradient overlay, and ideally improve compile times (but not necessarily time to first @ compile ]

github-actions[bot]
github-actions bot reviewed Dec 11, 2024
View reviewed changes
Copy link
Copy Markdown
Contributor

@github-actions github-actions bot left a comment

Choose a reason for hiding this comment

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

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

JuliaFormatter

[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/utils.jl

Lines 286 to 288 in ff729ce

mi = ccall(:jl_specializations_get_linfo, Ref{Core.MethodInstance},
(Any, Any, Any), match.method, match.spec_types, match.sparams)

[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/utils.jl

Line 290 in ff729ce

frame = Core.Compiler.InferenceState(result, #=cache_mode=#:local, interp)

[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/utils.jl

Lines 307 to 346 in ff729ce

opt = Core.Compiler.OptimizationState(frame, interp)
caller = frame.result
@static if VERSION < v"1.11-"
ir = Core.Compiler.run_passes(opt.src, opt, caller)
else
ir = Core.Compiler.run_passes_ipo_safe(opt.src, opt, caller)
Core.Compiler.ipo_dataflow_analysis!(interp, ir, caller)
end
# Rewrite type unstable calls to recurse into call_with_reactant to ensure
# they continue to use our interpreter. Reset the derived return type
# to Any if our interpreter would change the return type of any result.
# Also rewrite invoke (type stable call) to be :call, since otherwise apparently
# screws up type inference after this (TODO this should be fixed).
any_changed = false
for (i, inst) in enumerate(ir.stmts)
@static if VERSION < v"1.11"
changed, next = rewrite_inst(inst[:inst], ir)
Core.Compiler.setindex!(ir.stmts[i], next, :inst)
else
changed, next = rewrite_inst(inst[:stmt], ir)
Core.Compiler.setindex!(ir.stmts[i], next, :stmt)
end
if changed
any_changed = true
Core.Compiler.setindex!(ir.stmts[i], Any, :type)
end
end
Core.Compiler.finish(interp, opt, ir, caller)
src = Core.Compiler.ir_to_codeinf!(opt)
# Julia hits various internal errors trying to re-perform type inference
# on type infered code (that we undo inference of), if there is no type unstable
# code to be rewritten, just use the default methodinstance (still using our methodtable),
# to improve compatibility as these bugs are fixed upstream.
if !any_changed
src = Core.Compiler.retrieve_code_info(mi, world)
@show "post non change", src
end

[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/utils.jl

Line 362 in ff729ce

code_info.slotnames = Any[:call_with_reactant, REDUB_ARGUMENTS_NAME, code_info.slotnames...]

[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/utils.jl

Line 372 in ff729ce

[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/utils.jl

Line 389 in ff729ce

actual_argument = Expr(:call, Core.GlobalRef(Core, :getfield), overdub_args_slot, offset)

[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/utils.jl

Lines 394 to 396 in ff729ce

#push!(overdubbed_code, actual_argument)
push!(fn_args, Core.SSAValue(length(overdubbed_code)))

[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/utils.jl

Lines 405 to 406 in ff729ce

pop!(overdubbed_codelocs)
pop!(fn_args)

[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/utils.jl

Line 410 in ff729ce

push!(overdubbed_code, Expr(:call, Core.GlobalRef(Core, :getfield), overdub_args_slot, offset - 1))

[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/utils.jl

Lines 415 to 417 in ff729ce

push!(overdubbed_code, Expr(:(=), Core.SlotNumber(n_method_args + n_prepended_slots), trailing_arguments))
push!(overdubbed_codelocs, code_info.codelocs[1])
push!(fn_args, Core.SSAValue(length(overdubbed_code)))

[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/utils.jl

Lines 423 to 424 in ff729ce

arg_partially_inline!(code_info.code, fn_args, method.sig, Any[static_params...],
n_prepended_slots, n_prepended_slots, length(overdubbed_code), :propagate)

[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/src/utils.jl

Line 445 in ff729ce

$(Expr(:meta, :generated, call_with_reactant_generator))

[JuliaFormatter] reported by reviewdog 🐶

Reactant.jl/test/cuda.jl

Line 21 in ff729ce

@show @code_hlo optimize=false square!(A)

@wsmoses wsmoses requested review from gbaraldi and vchuravy December 11, 2024 04:41
@wsmoses @github-actions
Apply suggestions from code review
3bd5608 
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
github-actions[bot]
github-actions bot reviewed Dec 11, 2024
View reviewed changes
@wsmoses
Copy link
Copy Markdown
Member Author

wsmoses commented Dec 11, 2024

Currently I want to ensure this doesn't break anything.

Then I will potentially revamp the Enzyme overload to use more of this (and already this should hopefully fix the use of the Enzyme overload in type unstable code), as well as CUDA and perhaps pythoncall stuff

@wsmoses wsmoses mentioned this pull request Dec 11, 2024
Closed
github-actions[bot]
github-actions bot reviewed Dec 11, 2024
View reviewed changes
@wsmoses
Copy link
Copy Markdown
Member Author

wsmoses commented Dec 11, 2024

....and naturally I segfault julia. @aviatesk @vtjnash @gbaraldi @vchuravy @topolarity do you have any thoughts here?

@wsmoses @github-actions
Apply suggestions from code review
4807a79 
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
@wsmoses wsmoses mentioned this pull request Dec 11, 2024
Merged
mofeing
mofeing reviewed Dec 11, 2024
View reviewed changes
src/Interpreter.jl

Base.Experimental.@MethodTable(REACTANT_METHOD_TABLE)

function var"@reactant_override"(__source__::LineNumberNode, __module__::Module, def)
Copy link
Copy Markdown
Collaborator

@mofeing mofeing Dec 11, 2024

Choose a reason for hiding this comment

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

now that we're adding it, should we rename it to @reactant_overlay? because it's just a partial macro to @overlay

@mofeing
Copy link
Copy Markdown
Collaborator

mofeing commented Dec 11, 2024 •
edited
Loading

I think we're putting too much into "utils.jl"? Maybe this code makes more sense in "Compiler.jl"?

Also, with these fixes we'll be able to write diff rules from the Julia side?

@wsmoses
Copy link
Copy Markdown
Member Author

wsmoses commented Dec 14, 2024

Closing in favor of #365

@wsmoses wsmoses closed this Dec 14, 2024
@giordano giordano deleted the interp branch January 17, 2025 19:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@mofeing mofeing mofeing left review comments

@github-actions github-actions[bot] github-actions[bot] left review comments

@avik-pal avik-pal Awaiting requested review from avik-pal

@Pangoraw Pangoraw Awaiting requested review from Pangoraw

@vchuravy vchuravy Awaiting requested review from vchuravy

@gbaraldi gbaraldi Awaiting requested review from gbaraldi

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@wsmoses @mofeing @jumerckx