I am currently working on a way to memoise functions using fully typed containers (see this discussion for the general concept). I have a working implementation (sorry, still very barebone ATM), that seems to largely do what it is supposed to do and has already noticeably improved the runtime of my simulation.
However, under specific conditions the call to the memoised function seems to break: The internal cache is generated with the correct types and the function runs and returns a value, but the cache does not get modified. Mystifyingly the problem disappears as soon as I wrap the function call into a closure (i.e. use (x->myfun(x, v2, v3))(v1) instead of myfun(v1, v2, v3) or mark it as @noinline.
My suspicion is that the cache that is being created in the generating part of the generated function is somehow getting lost and not correctly stored in the global register of the package. I’m not sure, though, if this is a compiler bug or if what I’m trying to do is outside of the specifications for generated functions.
I will try to find a MWE that reproduces the effect outside of my somewhat baroque simulation code, but maybe someone has some insights on this in the meantime.
I have an MWE now, at least minimal in the sense that it only uses the package.
using TypedMemo
const data = rand(1000)
@cached Dict function check_cached(data, x, y)
x*y
end
function calc(data)
s = 0.0
for a in 1:10, b in 1:10
# this works:
# @noinline s += check_cached(data, a, b)
# this doesn't
s += check_cached(data, a, b)
end
s
end
println(calc(data))
println(get_all_caches(check_cached))
The two lines in the for loop should be identical, but only the one with @noinline does actually fill the cache.
Interestingly, if I simplify further by removing the data argument neither version works:
# for some reason this never works
@cached Dict function check_cached2(x, y)
x*y
end
function calc2(data)
s = 0.0
for a in 1:10, b in 1:10
# this doesn't work:
@noinline s += check_cached2(a, b)
# this doesn't either
# s += check_cached2(a, b)
end
s
end
println(calc2(data))
println(get_all_caches(check_cached2))
For reference, this is what @macroexpand spits out for the first version:
quote
#= /home/martin/Science/src/TypedMemo.jl/src/TypedMemo.jl:155 =#
function var"##check_cached#316"(data, x, y; )
#= In[9]:1 =#
#= In[9]:2 =#
x * y
end
#= /home/martin/Science/src/TypedMemo.jl/src/TypedMemo.jl:157 =#
function check_cached(data, x, y)
#= /home/martin/Science/src/TypedMemo.jl/src/TypedMemo.jl:157 =#
if $(Expr(:generated))
local var"#41###tmp#317" = begin
#= /home/martin/Science/src/TypedMemo.jl/src/TypedMemo.jl:157 =#
#= /home/martin/Science/src/TypedMemo.jl/src/TypedMemo.jl:158 =#
var"#42#arg_types" = TypedMemo.Tuple{data, x, y}
#= /home/martin/Science/src/TypedMemo.jl/src/TypedMemo.jl:159 =#
var"#43#ret_type" = (TypedMemo.Core).Compiler.return_type(var"##check_cached#316", var"#42#arg_types")
#= /home/martin/Science/src/TypedMemo.jl/src/TypedMemo.jl:161 =#
var"#44#cache" = TypedMemo.Dict{TypedMemo.Tuple{data, x, y}, var"#43#ret_type"}()
#= /home/martin/Science/src/TypedMemo.jl/src/TypedMemo.jl:162 =#
(TypedMemo.get!(TypedMemo.caches, check_cached, TypedMemo.IdDict{TypedMemo.Any, TypedMemo.Any}()))[(data, x, y)] = var"#44#cache"
#= /home/martin/Science/src/TypedMemo.jl/src/TypedMemo.jl:164 =#
Core._expr(:block, $(QuoteNode(:(#= /home/martin/Science/src/TypedMemo.jl/src/TypedMemo.jl:146 =#))), Core._expr(:call, :get!, $(Expr(:copyast, :($(QuoteNode(:(TypedMemo.Closure(var"##check_cached#316", (data, x, y)))))))), var"#44#cache", $(Expr(:copyast, :($(QuoteNode(:((data, x, y)))))))))
end
if var"#41###tmp#317" isa Core.CodeInfo
#= expr.jl:849 =#
return var"#41###tmp#317"
else
#= expr.jl:849 =#
var"#41###tmp#317"
end
else
$(Expr(:meta, :generated_only))
return
end
end
end
And here’s a completely standalone MWE:
const caches = Dict{Any, Any}()
const data = rand(1000)
struct Closure{F,A} <: Function
f::F
args::A
end
(x::Closure)() = x.f(x.args...)
kernel(data, x, y) = x * y
@generated function check_cached(data, x, y)
ret_type = Core.Compiler.return_type(kernel, Tuple{data, x, y})
global caches
cache = Dict{Tuple{data, x, y}, ret_type}()
caches[check_cached] = cache
:(get!(Closure(kernel, (data, x, y)), $cache, (data, x, y)))
end
function calc(data)
s = 0.0
for a in 1:10, b in 1:10
# this works:
# @noinline s += check_cached(data, a, b)
# this doesn't
s += check_cached(data, a, b)
end
s
end
println(calc(data))
println(caches)
As before this only works (in the sense that the cache gets filled) if I put @noinline before the call in the for loop.
I found the same thing when I played around with this stuff the other day. IIRC it’s a world age thing. Moving the definition of the kernel function to after the memorization function fixes it.
My version which puts return_type outside of the generated function also works (because it’s calculated at the same world age as the memorized function).
Do you mean after the generated function? That would contradict the official guidelines, though, as they state that a GF can only call functions that have been defined before.
@noinlineing the generated function also seems to work.
I wonder, though, whether all of this is supposed to work in the first place. The documentation on generated functions is annoyingly vague (starting with the fact that it doesn’t make clear where it talks about the generating and where about the generated code), but I wouldn’t be surprised if all of this relies on non-guaranteed behaviour.
Anyway, I might just file a bug report and see what the language developers have to say about it.
It’s official, this is not meant to work. Pity.