Assign vector to slices and allocation

I am trying to optimize my code by reducing all unnecessary memory allocation. In particular I have a some assignment I need to repeat many many times. Hence I pre-allocate all necessary vectors and matrices in order to reduce GC time and allocation time in general. Here I have a working example


function test_assign()
    N = 100
    A = rand(N,N,N)
    B = zeros(N)
    C = zeros(N,N)
    assign_1!(A,B)
    assign_2!(A,C)
end

function assign_1!(A,B)
    N,N,N = size(A)
    for i=1:N
        @views B[i] = test_scalar_field(A[i,:,1])
    end
end

function assign_2!(A,C)
    N,N,N = size(A)
    for i=1:N
        y = @view A[i,:,1]
        x = @view C[i,:]
        x .= test_vector_field(y)
    end
end

function test_scalar_field(x)
    return 10*norm(x)^2
end

function test_vector_field(x)
    return 10*2*x
end

If I do

julia> @benchmark assign_1!(A,B)
BenchmarkTools.Trial: 
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     5.355 μs (0.00% GC)
  median time:      5.391 μs (0.00% GC)
  mean time:        5.586 μs (0.00% GC)
  maximum time:     78.947 μs (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     6

so everything is good here where I assign scalar results. While when I assign vectors

@benchmark assign_2!(A,C)
BenchmarkTools.Trial: 
  memory estimate:  87.50 KiB
  allocs estimate:  100
  --------------
  minimum time:     21.977 μs (0.00% GC)
  median time:      45.901 μs (0.00% GC)
  mean time:        50.889 μs (13.14% GC)
  maximum time:     22.403 ms (99.77% GC)
  --------------
  samples:          10000
  evals/sample:     1

I haven’t yet found a way to remove those N allocation that apparently are unnecessary since everything is already there in memory. It should be possible somehow. As a rule of thumb I usually think about allocations in the following way:
Are they really mandatory? Can I write a Fortran90 code with less allocation (where you explicitly allocate) ?
If the answer is YES, then there should be a way

x .= test_vector_field.(y)

Or, maybe test_vector_field must be itself a function that mutates data in-place, i.e. test_vector_field!(x, y) overwrites x with some values based on y.

this doesn’t really makes much sense. It works in this specific examples where test_vector_field acts only component wise without mixing the components. But think at something like test_vector_field(x) = A*x where the function is really vector -> vector and not the vector version of something that acts on scalar -> scalar written in vector form.

This is perhaps the way to go, and actually achieves what I need. It is utterly boring to write everything in that way but at least it has 0 allocations if x and y are already pre-allocated.
Thanks for the idea!

Somewhat better looking, maybe:

function assign_2!(A,C)
    N,N,N = size(A)
    buf = zeros(eltype(A), N)
    vector_field(y) = test_vector_field!(buf, y)
    for i=1:N
        y = @view A[i,:,1]
        x = @view C[i,:]
        x .= vector_field(y)
    end
end

Otherwise, there is just no way semantically to avoid allocation if you need a non-destructive vector -> vector function in Julia.
A macro transforming x = f(args...) into f(x, args...) for such purposes must be possible, though.

If x is a vector, this necessarily allocates a new vector, there’s no way around this. If test_vector_field works on scalars, you dot both function call and assignment, otherwise, you just have to go with test_vector_field!.

For in-place matrix-vector product you anyway have to use mul!(y, A, x).

Perhaps with the exception when x is a static array? Depending on the context that would not allocate anything.

julia> function f(x)
         10*x
       end

julia> let 
         x = SVector{3,Float64}(rand(3))
         @allocated f(x)
       end
0

I am mentioning this because effectively you have taught me this a few days ago, and that allowed me to write some parts of a code of mine much clearer (without carrying the ! and preallocating some vectors obsessively everywhere, which was my programming habit).

Well, yes and no. It would create a new vector, but there is no performance cost to it, so it doesn’t matter. So I guess it doesn’t allocate a vector, but it does create it.

:+1:

1 Like

Here’s a macro for inplace operations:

macro in(code)
  code.head == :(=) || error("assignment expected")
  lhs, expr = code.args
  expr.head == :call || error("function call expected")
  # append `!` to the function name
  expr.args[1] = Symbol(String(expr.args[1]) * "!")
  # insert `lhs` as the first argument to the function
  insert!(expr.args, 2, lhs)
  esc(expr)
end

For instance, as can be verified with @macroexpand,

@in lhs = f(g(x), y, h(z) + k(w))

gets translated to

f!(lhs, g(x), y, h(z) + k(w))

I do not advocate the use of this macro.

2 Likes