Disabling allocations

Is there a way of disabling allocations altogether within the body of a function? I optimized my code to avoid all allocations, but they keep creeping back in with changes and are annoying to locate. I would much rather have a compile-time error.

1 Like

I don’t think there is a way. I would also like to add that there is nothing inherently wrong with allocations. At some point, you need data in memory for the CPU to do its thing. Are the allocations affecting your runtime/performance? If your code is type stable (check using @code_warntype), I wouldn’t worry about small allocations.

2 Likes

There are sometimes good reasons for avoiding allocations in specific parts of a program, usually for performance purposes.

You can measure allocations with @allocated and add a corresponding unit test.

2 Likes

I’m building a raytracer and allocations are indicative of a mistake (dynamic dispatch, accidental allocation, etc.) which degrades performance. Obviously the program as a whole will allocate memory but the actual render loop really should not. From my experience so far its fairly easy to screw up and accidentally cause things like dynamic dispatch (in some cases the reason can be pretty insidious, for instance a type annotation which is not specific enough). @code_warntype works great but isn’t recursive. This makes locating the source of an allocation very painful.

@allocated works great but actually locating the source of the allocation can get pretty tricky.

Have you tried this:?https://docs.julialang.org/en/v1/manual/profile/#Memory-allocation-analysis

Also I use @allocated like this:

a = @allocated begin
    Block to test
end; if a > 0 println(a) end

It is not hard to find sources of allocations moving the test around.

5 Likes

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.

9 Likes

Yes, this is what I do currently. I’m looking for a way of pinpointing the cause of the allocation without manually moving the @allocated test into every code branch.

This is really cool. What does UndefInitializer do in this example?

Perhaps the manual is not clear. Check this example:

This is the code (file name here: test.jl):

struct A
  x
end

function test(n,x)
  y = Vector{A}(undef,n)
  for i in 1:n
    y[i] = A(i*x)
  end
  y
end

Run julia with:

julia --track-allocation=user

Within Julia, do:

julia> using Profile

julia> include("./test.jl")
test (generic function with 1 method)

julia> test(10,rand()); # gets compiled

julia> Profile.clear_malloc_data() # clear allocations

julia> test(10,rand());

Exit Julia, this will generate a file test.jl.XXX.mem (extension .mem), which, in this case, contains:

        -
        - struct A
        -   x
        - end
        -
        - function test(n,x)
      160   y = Vector{A}(undef,n)
        0   for i in 1:n
      160     y[i] = A(i*x)
        -   end
        0   y
        - end

Where the lines with non-zero numbers are the lines where allocations occur.

5 Likes

Basically every array allocation at some point necessarily hits a constructor like Array{T}(undef, M, N) and that overdub rule catches that and rewrites it into an error.

Thanks, I didn’t understand that Array was always called with undef even for literals. I’ve looked a bit more into Cassette, it’s some pretty cool stuff!

Thanks, this works great. Although I still wish I could just disable the darn thing!

Note that arrays are not the only possible source of allocations though, there’s also dynamic dispatch, boxing and mutable structs and such.

In theory, the new version 1.6 compiler stuff could be used to catch all allocations / dynamism and throw an error though.

2 Likes

Yes please!

Yes, accidental dynamic dispatch is my main source of problems. I’d love to hear about the 1.6 compiler changes in question; if you could give me some pointers I would like to try and implement this feature!

I also tried to rewrite a raytracer serval months ago, tinsel. Currently, I can render the coffee cup example without scene bvh acceleration.
I use julia --track-allocation=user to find unnecessary memory allocation and the Cthulhu.jl to find type unstable. Unfortunately, unless single dispatch/typed function pointer is possible in Julia, some allocation caused by dynamic dispatch are impossible to elimate.
Another problem is the stack-allocated value. In a raytracer, it’s common to use a stack-allocated temporary value(for example, hit information) and pass its reference to other functions for further process. It’s impossible to do this in Julia, since the value might escape, which leads to allocation. Currently I have to preallocate all the intermediate values and passing them to the raytracers, which is quite fallible. Really hope there’s a way to force the value allocating on stacks.

That’s an impressive piece of work (I never got that far with my path tracer). But… it seems to be written in C++ and not Julia? :face_with_raised_eyebrow:

I believe that tinsel is the original whereas @ChenNingCong was referring to a port of the former.

Doh, I missed he said “rewrite” and not “write”, sorry! :grinning_face_with_smiling_eyes:

1 Like