TastyPancakes/Flux.jl
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

more robust loop lift

Browse Source
This commit is contained in:
Mike J Innes 2016-11-15 20:45:24 +00:00
parent 27554494b7
commit 46aa456bd1
2 changed files with 10 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/compiler/code.jl
View File

@ -4,7 +4,7 @@ export @net, @ml
function process_func(ex, params = []) function process_func(ex, params = [])
@capture(shortdef(ex), (args__,) -> body_) @capture(shortdef(ex), (args__,) -> body_)
body = @> body MacroTools.flatten block liftloops(params) graphm DataFlow.il body = @> body MacroTools.flatten block liftloops graphm DataFlow.il
body = mapconst(x -> x in params ? :(self.$x) : x, body) body = mapconst(x -> x in params ? :(self.$x) : x, body)
return args, body return args, body
end end

12
src/compiler/loops.jl
View File

@ -10,12 +10,18 @@ Offset(name, n) = Offset(name, n, nothing)
Base.:-(o::Offset) = Offset(o.name, -o.n, o.default) Base.:-(o::Offset) = Offset(o.name, -o.n, o.default)
function liftloops(ex, params) function liftloops(ex)
ex = DataFlow.normedges(ex) ex = DataFlow.normedges(ex)
MacroTools.postwalk(ex) do ex decls = Dict()
ex = MacroTools.postwalk(ex) do ex
@capture(ex, x_{n_}) || return ex @capture(ex, x_{n_}) || return ex
:($(Offset(x,n))($x)) haskey(decls, (x,n)) && return namify(decls[(x,n)])
@gensym edge
decls[(x,n)] = :($edge = $(Offset(x,n))($x))
edge
end end
prepend!(ex.args, collect(values(decls)))
ex
end end
function hasloops(model) function hasloops(model)