Make task sticky before entering threaded region

[tkf]

When we have t = @spawn @threads for ..., the task t that was at the primary thread when scheduling @threads for may be migrated away to a non-primary thread. This makes me wonder if we need to set t.sticky = true before calling jl_enter_threaded_region to ensure that jl_exit_threaded_region is called at the primary thread. I haven't encountered a bug due to this but I'm opening as a PR to see if my concern is valid.

[tkf]

Actually, this is (accidentally?) not a problem ATM since all tasks in @threads for are sticky and so schedule'ing them sets current_task().sticky = true. But making it explicit like this may be better? Or, I think it's better to put some comments in the code at least. (Edit: No, this is not the case since @threads for sets the tid of the child tasks before schedule. So, I think something like this PR is required.)

[vchuravy]

I don't understand jl_enter_threaded_region well enough. But yes I agree that this seems needed in the context of migration and I question the usefulness of threaded_region in the context of migration.

So maybe we should remove it, and support IO on all threads even outside a threaded_region.

[vtjnash]

I don't think this is necessary, or desirable. The calling Task can freely migrate, as all of the worker tasks will be pinned to 1:n, and this Task will simply sleep until all of the workers are done.

[tkf]

@vtjnash How do we balance the increments and decrements of _threadedregion without disabling the task migration?

[vtjnash]

_threadedregion is a simple global counter

[tkf]

Uh, sorry, I just realized that I've been reading the code from #43418 (which motivated me to look into the threaded region business) where you added if (jl_atomic_load_relaxed(&jl_current_task->tid) != 0) checks #43418 (comment)

julia/src/threading.c

Lines 538 to 549 in 2a6a8ab

	JL_DLLEXPORT void jl_enter_threaded_region(void)
	{
	if (jl_atomic_load_relaxed(&jl_current_task->tid) != 0)
	return;
	jl_atomic_fetch_add(&_threadedregion, 1);
	}
	JL_DLLEXPORT void jl_exit_threaded_region(void)
	{
	if (jl_atomic_load_relaxed(&jl_current_task->tid) != 0)
	return;
	jl_atomic_fetch_add(&_threadedregion, -1);

I can demonstrate the scenario that I mentioned in the OP with

function demo(n)
    ok = false
    nprimaries = 0
    niters = 0
    post = zeros(Int, Threads.nthreads())
    for outer niters in 1:n
        t = Threads.@spawn begin
            @assert !current_task().sticky
            Threads.threadid() == 1 || return nothing
            Threads.@threads :static for i in 1:2
                if i == 1
                    t0 = time_ns() + 1000
                    while t0 > time_ns()
                    end
                end
            end
            @assert !current_task().sticky
            return Threads.threadid()
        end
        y = fetch(t)
        if y !== nothing
            nprimaries += 1
            tid = y::Int
            post[tid] += 1
        end
        if ccall(:jl_in_threaded_region, Cint, ()) != 0
            ok = true
            break
        end
    end
    return (; ok, nprimaries, niters, post)
end

With Julia from the master branch, I always get demo(_).ok == false. But with #43418 I get

julia> ccall(:jl_in_threaded_region, Cint, ())
0

julia> demo(10000)
(ok = true, nprimaries = 2, niters = 39, post = [1, 1])

julia> ccall(:jl_in_threaded_region, Cint, ())
1

julia> Base.GIT_VERSION_INFO
Base.GitVersionInfo("2a6a8ab0b2ac5c2c6826c8895cc650ea05ab2159", "2a6a8ab0b2", "improve-threads", 1135, "2021-12-13 19:46 UTC", false, 9, 1.639413997e9, "", "")

i.e., now the threaded region is permanently enabled.

I'll change the merge target of this PR to #43418.

[tkf]

Maybe a better name is

Suggested change

	function prefer_threaded_region(f)
	function prefer_io(f)

?

I also wonder if it makes sense to expose this as a public API with the caveat that it is just a hint and f -> f() is a valid implementation of this. It then maybe even makes sense to not do this in @threads for by default to help CPU-bound (almost no I/O) use cases since people can recover the old behavior with prefer_io() do; @threads for ... end; end.

[tkf]

It seems like a better fix is coming #43418 (comment)