Skip to content

[WIP] automatic control flow #1754

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.

Sign up for GitHub

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+3,018 −978
Open

[WIP] automatic control flow #1754

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 0 additions & 104 deletions src/Interpreter.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,107 +4,3 @@
# - https://github.com/JuliaLang/julia/blob/v1.10.4/test/compiler/newinterp.jl#L9

const CC = Core.Compiler
using Enzyme

import Core.Compiler:
AbstractInterpreter,
abstract_call,
abstract_call_known,
ArgInfo,
StmtInfo,
AbsIntState,
get_max_methods,
CallMeta,
Effects,
NoCallInfo,
MethodResultPure

Base.Experimental.@MethodTable(REACTANT_METHOD_TABLE)

function var"@reactant_overlay"(__source__::LineNumberNode, __module__::Module, def)
return Base.Experimental.var"@overlay"(
__source__, __module__, :(Reactant.REACTANT_METHOD_TABLE), def
)
end

function set_reactant_abi(
interp,
@nospecialize(f),
arginfo::ArgInfo,
si::StmtInfo,
sv::AbsIntState,
max_methods::Int=get_max_methods(interp, f, sv),
)
(; fargs, argtypes) = arginfo

if f === ReactantCore.within_compile
if length(argtypes) != 1
@static if VERSION < v"1.11.0-"
return CallMeta(Union{}, Effects(), NoCallInfo())
else
return CallMeta(Union{}, Union{}, Effects(), NoCallInfo())
end
end
@static if VERSION < v"1.11.0-"
return CallMeta(
Core.Const(true), Core.Compiler.EFFECTS_TOTAL, MethodResultPure()
)
else
return CallMeta(
Core.Const(true), Union{}, Core.Compiler.EFFECTS_TOTAL, MethodResultPure()
)
end
end

# Improve inference by considering call_with_reactant as having the same results as
# the original call
if f === call_with_reactant
arginfo2 = ArgInfo(fargs isa Nothing ? nothing : fargs[2:end], argtypes[2:end])
return abstract_call(interp, arginfo2::ArgInfo, si, sv, max_methods)
end

return Base.@invoke abstract_call_known(
interp::AbstractInterpreter,
f::Any,
arginfo::ArgInfo,
si::StmtInfo,
sv::AbsIntState,
max_methods::Int,
)
end

@static if Enzyme.GPUCompiler.HAS_INTEGRATED_CACHE
struct ReactantCacheToken end

function ReactantInterpreter(; world::UInt=Base.get_world_counter())
return Enzyme.Compiler.Interpreter.EnzymeInterpreter(
ReactantCacheToken(),
REACTANT_METHOD_TABLE,
world,
false, #=forward_rules=#
false, #=reverse_rules=#
false, #=inactive_rules=#
false, #=broadcast_rewrite=#
false, #=within_autodiff_rewrite=#
set_reactant_abi,
)
end
else
const REACTANT_CACHE = Enzyme.GPUCompiler.CodeCache()

function ReactantInterpreter(;
world::UInt=Base.get_world_counter(), code_cache=REACTANT_CACHE
)
return Enzyme.Compiler.Interpreter.EnzymeInterpreter(
REACTANT_CACHE,
REACTANT_METHOD_TABLE,
world,
false, #=forward_rules=#
false, #=reverse_rules=#
false, #=inactive_rules=#
false, #=broadcast_rewrite=#
false, #=within_autodiff_rewrite=#
set_reactant_abi,
)
end
end
Loading
Loading