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
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
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!
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.
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,