Julia: function types and performance

user_7756364 0
Viewed 269

Is there a way in Julia to generalise a pattern like the following?

function compute_sum(xs::Vector{Float64})
    res = 0
    for i in 1:length(xs)
        res += sqrt(xs[i])
    end
    res
end

This computes the square-root of each vector element and then sums everything. It is much faster than the "naive" versions with array comprehension or map, and also doesn't allocate additional memory:

xs = rand(1000)

julia> @time compute_sum(xs)
  0.000004 seconds
676.8372556762225

julia> @time sum([sqrt(x) for x in xs])
  0.000013 seconds (3 allocations: 7.969 KiB)
676.837255676223

julia> @time sum(map(sqrt, xs))
  0.000013 seconds (3 allocations: 7.969 KiB)
676.837255676223

Unfortunately the "obvious" generic version is terrible wrt performance:

function compute_sum2(xs::Vector{Float64}, fn::Function)
    res = 0
    for i in 1:length(xs)
        res += fn(xs[i])
    end
    res
end

julia> @time compute_sum2(xs, x -> sqrt(x))
  0.013537 seconds (19.34 k allocations: 1.011 MiB)
676.8372556762225
julia performance
4 Answers

The reason is that x -> sqrt(x) is defined as a new anonymous function with each call to compute_sum2, so this causes new compilation every time you call it.

If you define it before even e.g. like this:

julia> f = x -> sqrt(x)

then you have:

julia> @time compute_sum2(xs, f) # here you pay compilation cost
  0.010053 seconds (19.46 k allocations: 1.064 MiB)
665.2469135020949

julia> @time compute_sum2(xs, f) # here you have already compiled everything
  0.000003 seconds (1 allocation: 16 bytes)
665.2469135020949

Note that a natural approach would be to define a function with a name like this:

julia> g(x) = sqrt(x)
g (generic function with 1 method)

julia> @time compute_sum2(xs, g)
  0.000002 seconds
665.2469135020949

You can see that x -> sqrt(x) defines a fresh anonymous function each time it is encountered when you write e.g.:

julia> typeof(x -> sqrt(x))
var"#3#4"

julia> typeof(x -> sqrt(x))
var"#5#6"

julia> typeof(x -> sqrt(x))
var"#7#8"

Note that this would be different if an anonymous function would be defined in a function body:

julia> h() = typeof(x -> sqrt(x))
h (generic function with 2 methods)

julia> h()
var"#11#12"

julia> h()
var"#11#12"

julia> h()
var"#11#12"

and you see that this time the anonymous function is the same every time.

user_12695670

In addition to the excellent response by Bogumil, I would just like to add that a very convenient way of generalizing this is to use the normal functional programming function like map, reduce, fold, etc.

In this case, you're doing a map transformation (namely sqrt) and a reduce (namely +), so you can also achieve the result with mapreduce(sqrt, +, xs). This has essentially no overhead and is comparable to a manual loop in performance.

If you have a really complicated series of transformations, you can get optimal performance and still use a function using the Transducers.jl package.

user_109926670

Bogumił has already answered the part about the function type. I want to point out that your implementation is already working as efficiently as possible, if benchmarked properly, but can be replaced by equivalent built-in functions:

julia> @btime compute_sum($xs)
  2.149 μs (0 allocations: 0 bytes)
661.6571623823567

julia> @btime sum(sqrt, $xs)
  2.149 μs (0 allocations: 0 bytes)
661.6571623823567

julia> @btime compute_sum2($xs, sqrt)
  2.149 μs (0 allocations: 0 bytes)
661.6571623823567

julia> @btime mapreduce(sqrt, +, $xs)
  2.149 μs (0 allocations: 0 bytes)
661.6571623823567

And it is always better to use an eta-equivalent non-lambda function, if possible: f instead of x -> f(x). Especially for builtin functions, since they are sometimes dispatched on.

user_13462760

The other answers are pretty comprehensive, but I wanted to point out that you can just omit the square brackets from sum([sqrt(x) for x in xs]) and get the fastest version of all:

julia> using BenchmarkTools

julia> @btime compute_sum($xs)
  1.779 μs (0 allocations: 0 bytes)
679.0943275393031

julia> @btime sum([sqrt(x) for x in $xs])
  1.626 μs (1 allocation: 7.94 KiB)
679.0943275393028

julia> @btime sum(map(sqrt, $xs))
  1.628 μs (1 allocation: 7.94 KiB)
679.0943275393028

julia> @btime sum(sqrt(x) for x in $xs)
  1.337 μs (0 allocations: 0 bytes)
679.0943275393031

Also note that on my computer on Julia master compute_sum is the slowest, not the fastest way to sum these numbers.

user_6592480
Related