Flux.jl/src/compiler/loops.jl

79 lines
1.8 KiB
Julia
Raw Normal View History

2016-10-26 14:37:30 +00:00
export unroll
2016-11-07 19:44:51 +00:00
type Offset
2016-08-29 14:17:54 +00:00
name::Symbol
2016-11-07 19:44:51 +00:00
n::Int
default::Nullable{Param}
2016-08-29 14:17:54 +00:00
end
2016-11-07 19:44:51 +00:00
Offset(name, n) = Offset(name, n, nothing)
2016-11-07 19:44:51 +00:00
function liftloops(ex, params)
2016-10-31 12:38:18 +00:00
ex = DataFlow.normedges(ex)
2016-11-07 19:44:51 +00:00
MacroTools.postwalk(ex) do ex
@capture(ex, x_{n_}) || return ex
:($(Offset(x,n))($x))
2016-08-29 14:17:54 +00:00
end
end
2016-08-31 01:37:53 +00:00
function hasloops(model)
g = graph(model)
g == nothing && return false
iscyclic(g) && return true
result = false
map(m -> hasloops(m) && (result = true), g)
return result
end
function atomise(model)
postwalk(graph(model)) do v
hasloops(value(v)) || return v
spliceinputs(atomise(value(v)), inputs(v)...)
end
end
2016-08-29 14:17:54 +00:00
2016-11-07 19:44:51 +00:00
# hiddeninput(n) = vertex(Split(n), inputnode(1))
#
# function unroll!(delay::IVertex, n)
# prewalk!(delay[1]) do v
# v === delay ? hiddeninput(n) : v
# end
# end
#
# function break!(g::IVertex)
# g = bumpinputs(g)
# loops = []
# defaults = []
# g = prewalk!(g) do v
# isa(value(v), Offset) || return v
# n = length(loops)+1
# push!(loops, unroll!(v, n))
# push!(defaults, get(value(v).default))
# hiddeninput(n)
# end
# cse(group(group(loops...), g)), defaults
# end
#
# function unrollgraph(model, n)
# graph, defaults = break!(atomise(model))
# outputs = [spliceinputs(graph, group([constant(splitnode(inputnode(1),i)) for i = 1:length(defaults)]...),
# splitnode(inputnode(2), 1))]
# for i = 2:n
# push!(outputs, spliceinputs(graph, outputs[end][1], splitnode(inputnode(2), i)))
# end
# state = outputs[end][1]
# outputs = map(x -> x[2], outputs)
# (@> group(state, group(outputs...)) detuple), map(x->x.x, defaults)
# end
2016-10-25 20:10:04 +00:00
2016-10-26 10:57:03 +00:00
type Unrolled <: Model
model
graph::IVertex{Any}
2016-10-29 23:19:57 +00:00
state::Vector{Any}
2016-10-26 10:57:03 +00:00
steps::Int
end
graph(u::Unrolled) = u.graph
2016-10-28 20:17:48 +00:00
unroll(model, n) = Unrolled(model, unrollgraph(model, n)..., n)