add BLAS.with_num_threads

[johnnychen94]

I added this BLAS.with_num_threads as a convenient wrapper, but I'm not sure if we want it live in Base because it's not thread-safe.

In some cases, multi-threading could happen at a more efficient higher level instead of at the BLAS level. Thus we may want to temporarily disable BLAS threads during the julia threads execution.

As an example, SVD is one typical case: if there are a bunch of SVD operations and if we can do embarrassingly parallel at a higher level, then it's better to disable the BLAS threads.

using LinearAlgebra

function make_test_data(n)
    U = rand(n, n)
    S = rand(n)
    V = rand(n, n)
    U * Diagonal(S) * V'
end

function svd_solver(X, k=5)
    U, S, Vt = svd(X)
    S[k:end] .= 0
    U * Diagonal(S) * Vt
end

dataset = [make_test_data(64) for _ in 1:64]

function svd_with_blas_threads(dataset)
    losses = zeros(size(dataset))
    Threads.@threads for i in eachindex(dataset)
        X = dataset[i]
        losses[i] = norm(X - svd_solver(X))
    end
    return losses
end

function svd_without_blas_threads(dataset)
    losses = zeros(size(dataset))
    BLAS.with_num_threads(1) do 
        Threads.@threads for i in eachindex(dataset)
            X = dataset[i]
            losses[i] = norm(X - svd_solver(X))
        end
    end
    return losses
end

Threads.nthreads() # 8
BLAS.get_num_threads() # 8

@btime svd_with_blas_threads(dataset); # 15.149 ms (1194 allocations: 18.55 MiB)
@btime svd_without_blas_threads(dataset);# 10.411 ms (1194 allocations: 18.55 MiB)

Would need to add a news entry if approved.

[ViralBShah]

@KristofferC @dkarrasch @chriselrod @YingboMa Thoughts on this?

I generally don't like code and APIs that code in numbers of threads, but this is a reality we have to work with today. Since this is an unexported API, I feel comfortable merging.

[ViralBShah]

@johnnychen94 Any ideas why the doctests are failing?

[johnnychen94]

@johnnychen94 Any ideas why the doctests are failing?

That seems to be noise. After rebasing on the current master the test passes now.

[chriselrod]

@KristofferC @dkarrasch @chriselrod @YingboMa Thoughts on this?

I generally don't like code and APIs that code in numbers of threads, but this is a reality we have to work with today. Since this is an unexported API, I feel comfortable merging.

Would you prefer they're integrated with Julia's threading and the scheduler handles it?

I think this is an improvement, but it'd be concerning if we want to break it later (e.g. after/if this merges).

[ViralBShah]

I think this is an improvement, but it'd be concerning if we want to break it later (e.g. after/if this merges).

That PR has not seen any real work since it was created. I believe this PR would be compatible with that, in any case.

[DilumAluthge]

Also, if we make the additions in this PR experimental and not part of the public API, then we can break it (and even remove it) in a future minor release.

[johnnychen94]

Should I only tag with_num_threads as experimental or also to set_num_threads?

[ViralBShah]

Only with_num_threads needs to be marked experimental.

[ViralBShah]

Let's rebase this and get it merged.

[KristofferC]

What is this test testing?

[johnnychen94]

This is used to simulate the case that when there exists a task with with_num_threads, then all threads are affected until that task is finished.

[KristofferC]

Suggested change

	retval = f()
	return f()

and fixup the rest?

[johnnychen94]

Sorry, I must have messed up with rebasing on a computer that has different local git history. This is pointed out in #41785 (comment) already.
Thanks for catching this again!

[vchuravy]

I am really not a fan of this. Mutating global state like this makes this effectively unusable from a library writers perspective.

Another question that came to mind is if this is portable across blas vendors? Or is this OpenBLAS specific.

[vchuravy]

This execution order is not guaranteed. Use a barrier instead to guarantee the ordering of events

[vchuravy]

Unnecessary catch

[chriselrod]

Another question that came to mind is if this is portable across blas vendors? Or is this OpenBLAS specific.

BLAS.get_num_threads and BLAS.set_num_threads are supported by MKL as well.

[johnnychen94]

I am really not a fan of this.

Okay. Should I fix the comments or just close this PR? I'm okay with having this on my local projects. I propose this PR as I feel like there might be other people wanting it.

[johnnychen94]

This isn't a very good solution so let me just close this one.

[ViralBShah]

Yes, I have been thinking that it isn't a good idea to change global state in base Julia like this.