# Multi-threaded code requiring scratch space

**URL:** <https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997>\
**Category:** Performance\
**Created:** [February 9, 2024, 5:48pm UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997 "2024-02-09T17:48:15Z")\
**Posts on this page:** 18\
**Page:** 1

<div class="post-metadata">

**Author:** ![PGS62](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pgs62/32/207030_2.png) [@PGS62](https://discourse.julialang.org/u/PGS62)\
**Post date:** [February 9, 2024, 5:48pm UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/1 "2024-02-09T17:48:15Z")

</div>

I think it’s true that to get decent performance from multi-threaded code, it’s a good idea to minimise the amount of allocation done by each task. One way to do this can be calling functions that accept a “scratch” argument, allowing the same memory to be re-used multiple times i.e. less allocation.

So I have code that looks something like this:

```julia
using Base.Threads: nthreads, @threads, threadid

function bugged_code(input_matrix::AbstractMatrix)

    output_matrix = similar(input_matrix)
    n = some_large_number

    scratch_args = [fill(1.0, n) for _ in 1:nthreads()]

    @threads for i in eachindex(input_matrix)
        scratch = scratch_args[threadid()]
        output_matrix[i] = some_operator(input_matrix[i], scratch)
    end
    return (output_matrix)
end

function some_operator(x, scratch::AbstractVector)
    #calculate y without allocation, thanks to scratch argument
    return (y)
end

```

My hope was that each scratch vector would never be in use by more than one task simultaneously. But on reading [Thread local state…](https://julialang.org/blog/2023/07/PSA-dont-use-threadid/#fixing_buggy_code_which_uses_this_pattern) I see that the code is probably bugged, with a “quickfix, not recommended longterm” being to use the static scheduler i.e. replace `@threads for` with `@threads :static for`

Maybe that’s OK, but my `some_operator` calls `sort!` which might well one day become multi-threaded in which case bad stuff (“… even deadlock”) might happen.

What’s the solution? How should I ensure that each executing task has sole access to dedicated scratch space?

---

<div class="post-metadata">

**Author:** ![Oscar\_Smith](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/oscar_smith/32/25343_2.png) [@Oscar\_Smith](https://discourse.julialang.org/u/Oscar_Smith)\
**Post date:** [February 9, 2024, 6:02pm UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/2 "2024-02-09T18:02:26Z")

</div>

You can use `task_local_storage`.

---

<div class="post-metadata">

**Author:** ![adienes](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/adienes/32/37459_2.png) [@adienes](https://discourse.julialang.org/u/adienes)\
**Post date:** [February 9, 2024, 6:08pm UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/3 "2024-02-09T18:08:45Z")

</div>

this might work

```julia
using OhMyThreads: TaskLocalValue, tmap

function bugged_code(input_matrix::AbstractMatrix)
    n = some_large_number
    scratch_space = TaskLocalValue{Vector{Float64}}(() -> fill(1.0, n))

    tmap(input_matrix) do x
        some_operator(x, scratch_space[])
    end
end

```

---

<div class="post-metadata">

**Author:** ![PGS62](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pgs62/32/207030_2.png) [@PGS62](https://discourse.julialang.org/u/PGS62)\
**Post date:** [February 9, 2024, 8:19pm UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/4 "2024-02-09T20:19:43Z")

</div>

Thanks to both @Oscar_Smith and @adienes for those replies, which I’m sure are pointing me in the right direction.

I’ve amended the code to use `TaskLocalValue` with all tests passing, but allocations are rather higher, and performance a bit lower, than my existing “quickfix, not recommended” approach. So I need to experiment some more.

---

<div class="post-metadata">

**Author:** ![adienes](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/adienes/32/37459_2.png) [@adienes](https://discourse.julialang.org/u/adienes)\
**Post date:** [February 9, 2024, 8:32pm UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/5 "2024-02-09T20:32:24Z")

</div>

> but allocations are rather higher, and performance a bit lower, than my existing “quickfix, not recommended” approach

I think this “shouldn’t” be the case. if you share a more complete example (and benchmark) it may be possible to help look more deeply into the performance difference.

---

<div class="post-metadata">

**Author:** ![lmiq](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lmiq/32/18314_2.png) [@lmiq](https://discourse.julialang.org/u/lmiq)\
**Post date:** [February 10, 2024, 12:25am UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/6 "2024-02-10T00:25:51Z")

</div>

You can use ChunkSplitters.jl for that.

---

<div class="post-metadata">

**Author:** ![carstenbauer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/carstenbauer/32/4981_2.png) [@carstenbauer](https://discourse.julialang.org/u/carstenbauer)\
**Post date:** [February 10, 2024, 8:27am UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/7 "2024-02-10T08:27:55Z")

</div>

> **[Task-Local Storage · OhMyThreads.jl](https://juliafolds2.github.io/OhMyThreads.jl/stable/examples/tls/tls/)**
>
> Documentation for OhMyThreads.jl.

---

<div class="post-metadata">

**Author:** ![carstenbauer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/carstenbauer/32/4981_2.png) [@carstenbauer](https://discourse.julialang.org/u/carstenbauer)\
**Post date:** [February 10, 2024, 10:23am UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/8 "2024-02-10T10:23:01Z")

</div>

> [@PGS62](#):
>
> I’ve amended the code to use `TaskLocalValue` with all tests passing, but allocations are rather higher, and performance a bit lower, than my existing “quickfix, not recommended” approach. So I need to experiment some more.

As @adienes said, a MWE would be nice. Could also just be dynamic vs static scheduling.

You could try `scheduler=StaticScheduler()`, e.g.

```julia
using OhMyThreads: TaskLocalValue, tmap

function bugged_code(input_matrix::AbstractMatrix)
    n = some_large_number
    scratch_space = TaskLocalValue{Vector{Float64}}(() -> fill(1.0, n))

    tmap(input_matrix; scheduler=StaticScheduler()) do x
        some_operator(x, scratch_space[])
    end
end

```

Or `scheduler=DynamicScheduler(; nchunks=nthreads())`.

---

<div class="post-metadata">

**Author:** ![PGS62](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pgs62/32/207030_2.png) [@PGS62](https://discourse.julialang.org/u/PGS62)\
**Post date:** [February 10, 2024, 10:53am UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/9 "2024-02-10T10:53:15Z")

</div>

Thanks for offering to “look more deeply”!

Quick background: I have (an unregistered) package [KendallTau](https://github.com/PGS62/KendallTau.jl) for calculating [kendall correlation](https://en.wikipedia.org/wiki/Kendall_rank_correlation_coefficient) via a multi-threaded function `corkendall` that I hope will replace the single-threaded function of the same name in StatsBase (mostly my code from 2021).

As of yesterday, the **main** branch of `KendallTau` uses the “quickfix” (`threadid` and static scheduler) and the **TLV\_Experiment** branch uses `OhMyThreads.TaskLocalValue`. The relevant function is `_corkendall` in `src/corkendall.jl`. See [here](https://github.com/PGS62/KendallTau.jl/blob/71c0af1c376429ac29e3e7d59ab107852e7934cf/src/corkendall.jl#L49C1-L84C4).

Here’s the comparison (on an ancient 4-core desktop):

```julia
julia> x=randn(1000,1000);

julia> res_qf=corkendall(x);#compile

julia> @time res_qf=corkendall(x);
  6.501597 seconds (1.58 k allocations: 15.884 MiB)

#Change branch from main to TLV_Experiment, and restart Julia (not sure if Revise plays ball when branch switching).

julia> using KendallTau

julia> x=randn(1000,1000);

julia> res_tlv=corkendall(x);#compile

julia> @time res_tlv=corkendall(x);
  6.595984 seconds (1.52 M allocations: 93.466 MiB, 1.70% gc time, 0.12% compilation time)

```

So, contrary to what I wrote yesterday, the speed of the two approaches looks about the same, but the number of allocations has gone up by a factor of 1,000 and the allocated memory by a factor of about 6. OTOH `StatsBase.corkendall` allocates way more:

```julia
julia> using StatsBase

julia> x=randn(1000,1000);

julia> res_sb = StatsBase.corkendall(x);#compile

julia> @time res_sb = StatsBase.corkendall(x);
 25.176168 seconds (3.00 M allocations: 17.090 GiB, 2.28% gc time)

```

I confess I have yet to fire up Julia with `--track-allocation` to investigate, but if you did have some insight as to how the code on the TLV\_Experiment branch could be improved then that would be wonderful.

---

<div class="post-metadata">

**Author:** ![mkoculak](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mkoculak/32/28310_2.png) [@mkoculak](https://discourse.julialang.org/u/mkoculak)\
**Post date:** [February 10, 2024, 11:26am UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/10 "2024-02-10T11:26:43Z")

</div>

Not sure if this would be the source of those allocations, but as you can see in the examples above (as well as in your original code) scratches are created outside of the threaded loop and here you have a new set with each iteration.

---

<div class="post-metadata">

**Author:** ![PGS62](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pgs62/32/207030_2.png) [@PGS62](https://discourse.julialang.org/u/PGS62)\
**Post date:** [February 11, 2024, 5:38pm UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/11 "2024-02-11T17:38:58Z")

</div>

That was sloppy - thanks for pointing it out. With the mistake [corrected](https://github.com/PGS62/KendallTau.jl/blob/28f57b2a331727cce1d8606d547a18b0294ecf68/src/corkendall.jl#L49C1-L84C4):

```julia
julia> res_tlv = KendallTau.corkendall(x);#compile

julia> @time res_tlv = KendallTau.corkendall(x);
  6.973501 seconds (1.50 M allocations: 38.920 MiB)

```

which continues to compare not so well with the “quickfix” code on the main branch:

```julia
julia> res_qf=corkendall(x);#compile

julia> @time res_qf=corkendall(x);
  6.501597 seconds (1.58 k allocations: 15.884 MiB)

```

---

<div class="post-metadata">

**Author:** ![PGS62](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pgs62/32/207030_2.png) [@PGS62](https://discourse.julialang.org/u/PGS62)\
**Post date:** [February 11, 2024, 8:36pm UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/12 "2024-02-11T20:36:14Z")

</div>

I’ve now figured out the cause of the extra allocations in the code on branch TLV\_Experiment. It’s that a `TaskLocalValue` is an untyped location so when getting data I need to specify the type as [here](https://github.com/PGS62/KendallTau.jl/blob/28bebe25c47f2776bc380a748ba660d68805cb46/src/corkendall.jl#L71C1-L77C49).

But for my use case I think it would be simpler to not use OhMyThreads, but instead use `Base.task_local_storage` (as suggested by @Oscar_Smith) as [here](https://github.com/PGS62/KendallTau.jl/blob/f999efa51af29abb0fba26f7d2b1a880422c3cad/src/corkendall.jl#L47C1-L92C4).

---

<div class="post-metadata">

**Author:** ![carstenbauer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/carstenbauer/32/4981_2.png) [@carstenbauer](https://discourse.julialang.org/u/carstenbauer)\
**Post date:** [February 11, 2024, 9:49pm UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/13 "2024-02-11T21:49:47Z")

</div>

> [@PGS62](#):
>
> But for my use case I think it would be simpler to not use OhMyThreads, but instead use `Base.task_local_storage` (as suggested by @Oscar_Smith) as [here](https://github.com/PGS62/KendallTau.jl/blob/f999efa51af29abb0fba26f7d2b1a880422c3cad/src/corkendall.jl#L47C1-L92C4).

Note that `TaskLocalValue` is a slim wrapper around `Base.task_local_storage`. See [here](https://github.com/vchuravy/TaskLocalValues.jl/blob/main/src/TaskLocalValues.jl).

---

<div class="post-metadata">

**Author:** ![carstenbauer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/carstenbauer/32/4981_2.png) [@carstenbauer](https://discourse.julialang.org/u/carstenbauer)\
**Post date:** [February 11, 2024, 9:52pm UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/14 "2024-02-11T21:52:52Z")

</div>

> [@PGS62](#):
>
> I’ve now figured out the cause of the extra allocations in the code on branch TLV\_Experiment. It’s that a `TaskLocalValue` is an untyped location so when getting data I need to specify the type as [here](https://github.com/PGS62/KendallTau.jl/blob/28bebe25c47f2776bc380a748ba660d68805cb46/src/corkendall.jl#L71C1-L77C49).

That shouldn’t be necessary. You specify the type of the value explicitly in the constructor and [`getindex` has a type assertion](https://github.com/vchuravy/TaskLocalValues.jl/blob/main/src/TaskLocalValues.jl#L14).

---

<div class="post-metadata">

**Author:** ![PGS62](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pgs62/32/207030_2.png) [@PGS62](https://discourse.julialang.org/u/PGS62)\
**Post date:** [February 12, 2024, 10:27am UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/15 "2024-02-12T10:27:31Z")

</div>

Ah, I suspected that `OhMyThreads.TaskLocalValue` might be simply a wrapper around `Base.task_local_storage` though I hadn’t dug in to check - thanks for confirming.

I’m a bit perplexed as to what was going on in my tests of `TaskLocalValue` yesterday (done at home, today I’m at my office) since I can’t replicate type assertions being necessary, just as you wrote above.

By the way, the documentation for OhMyThreads is excellent!

---

<div class="post-metadata">

**Author:** ![carstenbauer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/carstenbauer/32/4981_2.png) [@carstenbauer](https://discourse.julialang.org/u/carstenbauer)\
**Post date:** [February 12, 2024, 10:37am UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/16 "2024-02-12T10:37:06Z")

</div>

> [@PGS62](#):
>
> I’m a bit perplexed as to what was going on in my tests of `TaskLocalValue` yesterday (done at home, today I’m at my office) since I can’t replicate type assertions being necessary, just as you wrote above.

Great, otherwise I would have been really confused 🙂

> [@PGS62](#):
>
> By the way, the documentation for OhMyThreads is excellent!

Thanks. The package and its documentation are brand new. Please file issues/PRs if you have suggestions for how to improve the documentation or also new features that you would like to see.

---

<div class="post-metadata">

**Author:** ![PGS62](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pgs62/32/207030_2.png) [@PGS62](https://discourse.julialang.org/u/PGS62)\
**Post date:** [February 12, 2024, 12:14pm UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/17 "2024-02-12T12:14:18Z")

</div>

In case you are interested, I’ve managed to replicate the circumstance in which allocations _do_ happen (by looking at git history of what I actually did yesterday!)

Allocations don’t happen with [this](https://github.com/PGS62/KendallTau.jl/blob/28bebe25c47f2776bc380a748ba660d68805cb46/src/corkendall.jl#L49C1-L92C4) code, but they do happen with [this](https://github.com/PGS62/KendallTau.jl/blob/28f57b2a331727cce1d8606d547a18b0294ecf68/src/corkendall.jl#L49C1-L84C4) code, which you could install via:

`]add https://github.com/PGS62/KendallTau#TLV_with_allocations`

I then see this:

```julia
julia> using KendallTau
Precompiling KendallTau
  1 dependency successfully precompiled in 2 seconds. 52 already precompiled.

julia> x=randn(1000,1000);res=corkendall(x);

julia> @time corkendall(x);
  1.685487 seconds (1.50 M allocations: 39.410 MiB)

```

So 1.5 million allocations as opposed to the “good state” of 1.5 thousand.

---

<div class="post-metadata">

**Author:** ![PGS62](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pgs62/32/207030_2.png) [@PGS62](https://discourse.julialang.org/u/PGS62)\
**Post date:** [February 14, 2024, 12:40pm UTC](https://discourse.julialang.org/t/multi-threaded-code-requiring-scratch-space/109997/18 "2024-02-14T12:40:37Z")

</div>

Thanks to @lmiq for the suggestion that I look at ChunkSplitters.jl An aspect of the problem that I haven’t paid too much attention to is that when calculating `corkendall(x)` (as opposed to `corkendall(x,y)`) the workload is naturally unbalanced as one calculates only one half of a symmetric matrix. ChunkSplitters helps with unbalanced workloads.