Using Threads.@spawn with BLAS.set_num_threads

General Usage
1

Hi,

I’m implementing a frequency sweep algorithm. In each iteration, two reduced models are evaluated on a grid of candidate points. Since each candidate frequency is evaluated independently of the others, I can evaluate them in different tasks.

The issue is that the evaluations involve BLAS operations. I ran some tests, and in general, when I increase the number of tasks into which i split the candidates, I get slightly better perfomance when I set the number of BLAS threads to 1.

The pseudocode would look something like this:


function with_BLAS_num_threads(f::F,n::Integer=1) where F
    previous=BLAS.get_num_threads()
    previous==n && return f()
    try
        BLAS.set_num_threads(n)
        return f()
    finally
        BLAS.set_num_threads(previous)
    end
end

# I can use ChunkSplitters.jl
function pseudo_errors!(errors, candidates, E_r1, A_r1, B_r1, C_r1, E_r2, A_r2, B_r2, C_r2, Δf, nthreads)
    n = length(candidates)
    chunk_size = ceil(Int, n / nthreads)
    tasks = map(1:nthreads) do t
        Threads.@spawn begin
            # allocate matrices
            H1 = # ...
            chunk = (t-1)*chunk_size+1 : min(t*chunk_size, n)
            for idx in chunk
                fi = candidates[idx]
                # compute s1 and s2
                eval_reduced_model!(H1, args1...) #
                eval_reduced_model!(H2, args2...)
                errors[idx] = # something
            end
        end
    end
    foreach(wait, tasks)
    return errors
end


with_BLAS_num_threads(1) do
    pseudo_errors!(errors, candidates, E_r1, A_r1, B_r1, C_r1, E_r2, A_r2, B_r2, C_r2, Δf, nthreads)
end

My main interest isn’t performance, since I can benchmark that and determine whether the trade off is worthwhile. My main concern is whether there are any issues with this type of pattern.

  • Are there any side effects or issues that I might be overlooking?
  • Is there a more convenient or suitable pattern for this type of problem?

I’m interested in hearing how people would approach this in Julia, and whether there are any unexpected issues with this pattern that I should be aware of.

2

The pattern seems fine. Just a couple of comments.

There’s a function cld which you can use like chunk_size = cld(n, nthreads) to avoid possible problems with rounding (e.g. n/nthreads = 6.000000000001)

Your last chunk may end up very small. There’s a simple package ChunkSplitters.jl which balances the chunk sizes better. An alternative is taking indices one by one:

idxs = Atomic{Int}(1)
...
    @spawn begin
    ...
    while (idx = atomic_add!(idxs, 1)) <= n)
        ...
    end

Or have the chunk_size a variable:

while (lastidx = atomic_add!(idxs, chunk_size)) 
    for idx in lastidx-chunk_size+1:lastidx
        ...
    end

It’s also possible to decrease the chunk size as one approaches the end, but that requires the chunk_size to be atomic as well, and some care must be taken.