Reducing allocations in arbitrary precision calculations with Julia

Question

I am new to Julia and have only written in python before which is most likely also reflected by my coding style. Unfortunately, my applications require very high numerical precision which is why I am resorting to BigFloat and BigInt in Julia. For some reason, this leads to a major increase in allocations besides the obvious increase in computational time. I am wondering where my code can be improved.

I have attached the current version of my code below with example parameters that allow for execution under 1s.

using Combinatorics

function gamma_mpnq(m::Int, p::Int, n::Int, q::Int)::BigFloat
    return sqrt(factorial(big(m)) * factorial(big(n))) / (big(2)^(p + q) * factorial(big(m - 2 * p)) * factorial(big(n - 2 * q)) * factorial(big(q)) * factorial(big(p)))
end

function general_binomial(alpha::Float64, k::Int)
    """
    Implementation of the generalized binomial.
    """
    _prod = BigFloat(1.0)
    for kk in 1:k
        _prod *= (alpha - kk + 1.)
    end
    return _prod / factorial(big(k))
end

function construct_M(nmax::Int, mmax::Int)::Array
    """
    Construct the matrix M containing all moments <C^n * C*^m> up to nmax and mmax.
    """
    _M = zeros(BigFloat, nmax, mmax)
    for mm in 1:nmax
        for nn in 1:mmax
            if nn == 1
                _M[nn, mm] = 1. / doublefactorial(BigInt(2*mm - 1))
            elseif mm == 1
                _M[nn, mm] = 1. / doublefactorial(BigInt(2*nn - 1))
            else
                _M[nn, mm] = (((nn-1) * _M[nn-1, mm] + (mm-1) * _M[nn, mm-1]) / (2. * (nn-mm)^2 + (nn-1) + (mm-1)))
            end
        end
    end
    return _M
end


function maclaurin_exp(n::Int, m::Int, p::Int, q::Int, M::Array, cutoff::Int)::BigFloat
    """
    Calculate the expectation value of ((1+C)/(1+C*))^((n-m)/2) * C^p * C*^q
    using the Maclaurin expansion and the matrix containing all moments M.
    """
    inner::BigFloat = 0.0
    @views for nn in 1:cutoff
        for mm in 1:cutoff
            inner += (general_binomial((n-m)/2., mm-1) * general_binomial((m-n)/2., nn-1) * M[q+nn, p+mm])
        end
    end
    return inner
end

function Hmn(m::Int, n::Int, M::Array, cutoff::Int)::Float64
    """
    Calculates a single entry of the Hmn matrix.
    """
    ppmax = div(m, 2)
    qqmax = div(n, 2)
    outer = 0.0
    for pp in 0:ppmax
        for qq in 0:qqmax
            outer += (gamma_mpnq(m, pp, n, qq) * maclaurin_exp(n, m, pp, qq, M, cutoff))
            # println(pp, " ", qq, " ", outer)
        end
    end
    return outer
end

dim = 128
cutoff = 64
M = construct_M(dim+cutoff, dim+cutoff)
@btime Hmn(12, 2, M, cutoff)

``` 

Answer 1

The following is a workaround for your speed and memory issue rather than a general review of your code.

Profiling is useful for investigating the source of performance issues, and Julia's standard library includes a profiling module (guide, doc).

Here's a basic use of it: (I'm using @eval just to interpolate in the construction of M)

using Profile
@eval @profile Hmn(12, 2, $(construct_M(192, 192)), 64)

julia> Profile.print()
Overhead ╎ [+additional indent] Count File:Line; Function
=========================================================
    ╎963  @Base/task.jl:356; (::IJulia.var"#15#18")()
...
    ╎    ╎   963  In[1]:61; Hmn(::Int64, ::Int64, ::Arr...
    ╎    ╎    1    In[1]:4; gamma_mpnq(::Int64, ::Int6...
...
  18╎    ╎    962  In[1]:46; maclaurin_exp(::Int64, ::I...
   1╎    ╎     1    In[1]:0; general_binomial(::Float6...
    ╎    ╎     14   In[1]:11; general_binomial(::Float6...
...
    ╎    ╎     786  In[1]:13; general_binomial(::Float6...
...
    ╎    ╎     97   In[1]:15; general_binomial(::Float6...
...

I've removed all the lines that point to code in external files. Of the 963 samples collected in Hmn, more than 800 were spent within calculations of general_binomial.

Looking at how general_binomial(alpha, k) is used, I noticed that it is called repeated with only a small number of different alpha and k values. Since it is a pure function (i.e., it does not interact with external state), basic memoization can be applied to speed it up.

We can use a Dict to cache the results (make sure it is defined with concrete key and value types). To allow for easy benchmarking, we'll define it as a global constant; for performance, it is important that global variables are constant (or otherwise type-constrained in code using it) so that type inference works (also, the compiled code can avoid an extra lookup).

Here's a simple way to do this. Rename general_binomial to _general_binomial, then create the Dict and redefine general_binomial to cache the results. Like so:

@isdefined(general_binomial) && Base.delete_method.(methods(general_binomial))

const GENERAL_BINOMIAL_CACHE = Dict{Tuple{Float64, Int}, BigFloat}()

function general_binomial(alpha, k)
    return get!(GENERAL_BINOMIAL_CACHE, (alpha, k)) do
        _general_binomial(alpha, k)
    end
end

Benchmarking shows a substantial improvement in speed and memory use using the parameters you provided.

Before:

julia> @benchmark Hmn(12, 2, $(construct_M(192, 192)), 64)
BenchmarkTools.Trial: 
  memory estimate:  435.29 MiB
  allocs estimate:  8543119
  --------------
  minimum time:     787.497 ms (5.88% GC)
  median time:      791.734 ms (5.83% GC)
  mean time:        792.690 ms (5.76% GC)
  maximum time:     802.070 ms (5.17% GC)
  --------------
  samples:          7
  evals/sample:     1

After: (the cache is cleared before each sample by empty!(GENERAL_BINOMIAL_CACHE))

julia> @benchmark Hmn(12, 2, $(construct_M(192, 192)), 64) evals=1 setup=empty!(GENERAL_BINOMIAL_CACHE)
BenchmarkTools.Trial: 
  memory estimate:  12.74 MiB
  allocs estimate:  239309
  --------------
  minimum time:     10.684 ms (0.00% GC)
  median time:      12.844 ms (0.00% GC)
  mean time:        14.181 ms (7.58% GC)
  maximum time:     32.631 ms (14.03% GC)
  --------------
  samples:          353
  evals/sample:     1