Julia: tipos de função e desempenho

Sep 28 2020

Há alguma maneira de Julia generalizar um padrão como o seguinte?

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

Isso calcula a raiz quadrada de cada elemento do vetor e soma tudo. É muito mais rápido do que as versões "ingênuas" com compreensão de array ou map, e também não aloca memória adicional:

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

Infelizmente, a versão genérica "óbvia" tem um desempenho terrível de wrt:

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

Respostas

7 BogumiłKamiński Sep 28 2020 at 18:18

O motivo é que x -> sqrt(x)é definido como uma nova função anônima a cada chamada a compute_sum2, portanto, isso causa uma nova compilação toda vez que você a chama.

Se você definir antes mesmo, por exemplo, assim:

julia> f = x -> sqrt(x)

então você tem:

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

Observe que uma abordagem natural seria definir uma função com um nome como este:

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

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

Você pode ver que x -> sqrt(x)define uma nova função anônima cada vez que é encontrada quando você escreve, por exemplo:

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

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

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

Observe que isso seria diferente se uma função anônima fosse definida em um corpo de função:

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"

e você verá que, desta vez, a função anônima é sempre a mesma.

7 JakobNissen Sep 28 2020 at 18:40

Além da excelente resposta por Bogumil, eu gostaria apenas de acrescentar que uma maneira muito conveniente de generalizar isso é usar a função de programação funcional normal como map, reduce, fold, etc.

Neste caso, você está fazendo uma maptransformação (ou seja sqrt) e uma redução (ou seja +), então você também pode obter o resultado com mapreduce(sqrt, +, xs). Essencialmente, isso não tem sobrecarga e é comparável a um loop manual no desempenho.

Se você tem uma série realmente complicada de transformações, pode obter um desempenho ideal e ainda usar uma função usando o pacote Transducers.jl.

3 phipsgabler Sep 28 2020 at 19:35

Bogumił já respondeu a parte sobre o tipo de função. Quero salientar que sua implementação já está funcionando da forma mais eficiente possível, se comparada corretamente , mas pode ser substituída por funções integradas equivalentes:

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

E é sempre melhor usar uma função não lambda equivalente a eta, se possível: em fvez de x -> f(x). Especialmente para funções integradas, uma vez que às vezes são despachadas.

1 StefanKarpinski Sep 28 2020 at 21:00

As outras respostas são bastante abrangentes, mas gostaria de salientar que você pode simplesmente omitir os colchetes de sum([sqrt(x) for x in xs])e obter a versão mais rápida de todas:

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

Observe também que no meu computador o Julia master compute_sumé a maneira mais lenta e não a mais rápida de somar esses números.