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

move cache logic

Browse Source
This commit is contained in:
Mike J Innes 2018-02-27 21:41:03 +00:00
parent 2c74976602
commit bdb8aae107
3 changed files with 14 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
src/jit/JIT.jl
View File

@ -1,7 +1,6 @@
module JIT
include("shapes.jl")
include("inplace.jl")
include("trace.jl")
include("lib.jl")

11
src/jit/inplace.jl
View File

@ -1,11 +0,0 @@
mutable struct Cached{F,A}
f::F
buffer::A
end
function (c::Cached)(args...)
sh = shape(c.f, shape(args)...)
bytes(sh) > length(c.buffer) && (c.buffer = similar(c.buffer, bytes(sh)))
y = restructure(sh, c.buffer)
inplace!(c.f, y, args...)
end

14
src/jit/shapes.jl
View File

@ -35,3 +35,17 @@ function restructure(sh::Shape{T}, buf::Vector{UInt8}) where T
buf = unsafe_wrap(Array, pointer(buf), sizeof(sh))
reshape(reinterpret(T, buf), size(sh))
end
# Execution with caches
mutable struct Cached{F,A}
f::F
buffer::A
end
function (c::Cached)(args...)
sh = shape(c.f, shape(args)...)
bytes(sh) > length(c.buffer) && (c.buffer = similar(c.buffer, bytes(sh)))
y = restructure(sh, c.buffer)
inplace!(c.f, y, args...)
end