juliaallocationin-place

How to make my julia in-place function wrapper work without allocating new memory?

I am having a specific problem in julia, I cannot seem to figure out. I broke down my code to some minimal working example to demonstrate the problem:

using BenchmarkTools

function f!(u::Vector)
    u[1] = 5im
    u[2] = 6
end

function g!(u::Vector)
    u = u[1:2] + u[3:4] * 1im
    f!(u)
end 

# first version, works as expected
f!(ComplexF64.(Vector([1,2]))) # for jit compilation
u = ComplexF64.(Vector([1,2]))
println(u)
allo = @allocated f!(u)
println("Allocations : $allo bytes")
println(u)

# second try, not working
g!(Float64.(Vector([1,2,3,4]))) # for jit compilation
u = Float64.(Vector([1,2,3,4]))
println(u)
allo = @allocated g!(u)
println("Allocations : $allo bytes")
println(u)

The function f! takes a 2d-Vector (note the vector must be complex to not throw an error) and simply changes the first entry to the pure imaginary number 5im and the second entry to 6.

Now I have a second function g! which is basically a wrapper for f, which allows for a Float64 vector with twice as many entries. Its use is to convert the vector to ComplexF64 and update its entries in-place via f!.

I specifically want that both functions f! and g! do not allocate additional memory. Although technically a Float64 vector with 4 entries should use the same amount of memory as a ComplexF64 vector with 2 entries, I am not sure if this is possible to do in julia.

The output to the above code is:

ComplexF64[1.0 + 0.0im, 2.0 + 0.0im]
Allocations : 0 bytes
ComplexF64[0.0 + 5.0im, 6.0 + 0.0im]
[1.0, 2.0, 3.0, 4.0]
Allocations : 352 bytes
[1.0, 2.0, 3.0, 4.0]

Problem: The output is not what I expected. While f! works just fine, applying g! doesn't produce the output I want, most likely because while locally converting u to a complex vector f! will just update this variable in local scope not modifying the original vector. I am rather new to julia and I have not figured out a way to make this wrapper work. Also I am not sure if this is possible without allocating new memory.

Is there a way to achieve with g! the same as with f!, i.e. update the entries of u in-place but without allocating additional memory?

Thanks in advance!

Edit: To have an idea what I am actually trying to achieve, see the following

function f!(u::Vector{ComplexF64})
    # takes as input a complex vector of length N
    # u = a + bim (where a and b are float64 vectors)
    # ... modify u in place
    return
end

function g!(u::Vector{Float64})
    # takes as input a float vector of length 2*N
    # u = (a,b) (where a and b are float64 vectors of length N, representing real and imag part)
    # modify u = (a,b) without new memory allocation to achieve the same as calling f!(a + bim)
    # More specifically: a = real(f(a + bim)), b = imag(f(a + bim))
    return
end

Hence, for given f!, which modifies a complex vector in place, g! should be the split of f! into real and imaginary parts (if possible without additional memory allocation).

Solution

  • Yes, you can do this, but you need to make several changes to your code:

    function f!(u::AbstractVector{<:Complex})
    u[1] = 5im
    u[2] = 6
    return u 
end

    As you already guessed, writing u = ... doesn't actually modify u , it just assigns the label to whatever is on the right hand side, and the right hand side happens to be a newly allocated array, since slices, like u[1:2] allocates. In fact the right hand side creates several allocations.

    Instead, you can turn a real array into a complex one by reinterpreting the underlying data, like below. Here, v now refers to the same memory, just with a different interpretation of what those data mean. In this case, it gives you a complex vector of half the length of the real input. Then you pass it to f! .

    function g!(u::AbstractVector{T}) where {T<:Real}
    v = reinterpret(Complex{T}, u)
    return f!(v)
end

    Example:

    julia> x = [1,2,3,4]
4-element Vector{Int64}:
 1
 2
 3
 4

julia> g!(x)
2-element reinterpret(Complex{Int64}, ::Vector{Int64}):
 0 + 5im
 6 + 0im

    Note that if you print x , it looks like this

    julia> x
4-element Vector{Int64}:
 0
 5
 6
 0

    It is only the reinterpreted wrapper that has the right type.

    BTW, since you are concerned about avoiding allocations, here are some additional remarks: