Disabling allocations

General Usage Performance
7

For Array allocations, or ones arising from mutable structures you know about, you can write a custom Cassette.jl pass that throws an error if you try to allocate an Array or the known type. E.g.

using Cassette
Cassette.@context AlloCatcher
function Cassette.overdub(::AlloCatcher, ::Type{<:Array}, ::UndefInitializer, dims...)
    error("Tried to allocate an array")
end

Then if you had e.g.

bad_add!(dst, a, b)  = dst .= a  + b # forgot to broadcast a .+ b
good_add!(dst, a, b) = dst .= a .+ b 

dst, a, b = rand(5), rand(5), rand(5)

# runs like normal:
Cassette.overdub(AlloCatcher(), good_add!, dst, a, b) 

# errors, and you get the stack trace to the allocation:
Cassette.overdub(AlloCatcher(), bad_add!, dst, a, b)

FWIW this is exactly what AutoPreallocation.jl does except instead of erroring, it saves the allocated array for you so you can reuse it the next time you call the same function.

10 Likes