Julia - Is there a way to make this in-place assignment function look better

user_8408902 0
Viewed 751

Consider this excerpt of code:

q = zeros(5,5);
x = rand(5,5);
function testFunc!(n, q, x)
    for i = 1:n
        q .= x.^2;
    end
end

I initialize q and update the values in q through an in-place assignment function. No matter how I choose n, the number of allocations is constant since I'm not creating anything new.

However, is there a way to write a function called testFunc2 that looks like this:

q = zeros(5,5)
x = rand(5,5);
q = testFunc2(n, x)

such that the memory allocation is invariant to n.

I know that this looks silly, but in my actual code, I have numerous variables that are similar to q and I have a loop that updates them each iteration. I don't want the code to occupy excessive memory. If I write a function that looks like testFunc2(n, x, q1, q2, q3, q4), it would look very cumbersome in terms of the syntax. Therefore, I want to be able to have q1, q2, q3, q4 = testFunc2(n, x) without having to worry about memory.

julia allocation
2 Answers

If I understood correctly, the simplest way is to use a closure:

let 
   q1 = similar(q)
   q2 = similar(q)
   global function testFunc2(n, x)
       for i = 1:n
          q1 .= x.^2
          q2 .= x.^2
       end
       return q1, q2
   end
end

julia> @btime testFunc2(10, $x)
  786.000 ns (0 allocations: 0 bytes)

julia> @btime testFunc2(100, $x)
  7.476 μs (0 allocations: 0 bytes)

As pointed by @Liso in the comment below, if we don't attempt to share the same memory between calls, the copy should be returned e.g. return copy(q1), copy(q2) instead of return q1, q2, then we'll still get constant memory allocation:

@btime testFunc2(10, $x)
  821.684 ns (3 allocations: 704 bytes)

@btime testFunc2(100, $x)
  7.326 μs (3 allocations: 704 bytes)

Anyway, choosing which method depends on corresponding use cases, functor is a more flexible choice.

user_45945320

I'm not sure if q1, q2, q3, q4 = testFunc2(n, x) looks much better than testFunc2!(n, x, q1, q2, q3, q4), but you can do that type of transformation with a macro like this:

macro wrap(ex)
    outputs = ex.args[1].args
    ex.args[2].args[1] = Symbol(ex.args[2].args[1], "!")
    append!(ex.args[2].args, outputs)
    esc(ex)
end

which expands

@wrap q1, q2, q3, q4 = testFunc2(n, x)

into

(q1, q2, q3, q4) = testFunc2!(n, x, q1, q2, q3, q4)
user_16574450
Related