Disable polyester threads if in an outer @batch loop

[tetrakai1]

The readme gives this example, which works great:

disable_polyester_threads() do
    @threads for i in 1:n
        f()
    end
end

For my use case, I would like to disable an inner @batch in f() if the outer loop uses @batch. Eg,:

disable_polyester_threads() do
    @batch for i in 1:n
        f()
    end
end

Is there an easy way to do this? Thanks.

[chriselrod]

This should probably roughly be what you observe.

[tetrakai1]

Ha, yes. It does indeed "work" as advertised, but the outer @batch is also disabled.

[chriselrod]

Sorry, I meant that

@batch for i in 1:n
    @batch for j in 1:m # should be single threaded
        f()
    end
end

you shouldn't need disable_polyester_threads.

[tetrakai1]

Thanks, maybe its a bug then. For me it freezes with some cores stuck at 100%. My situation is a bit more complicated, but didn't think it would matter. Schematically:

function playreps!(sims)
    @batch for i in eachindex(sims)
        playseason!(sims[i])
    end
    return nothing
end

function playseason!(lg_data)
    @unpack_LeagueData lg_data
    for wk in eachindex(schedule)
        playgames!(lg_data, schedule[wk])
    end
    lgrank!(lg_table, Val(length(lg_table)))

    return nothing
end

function playgames!(lg_data, wk_sched)
    @unpack_LeagueData lg_data
    @batch for idx in eachrow(wk_sched)
    # for idx in eachrow(wk_sched)
        idx1 = idx[1]
        idx2 = idx[2]

        playgame!(tv, idx1, idx2)
        comms = (tv[idx1].comm, tv[idx2].comm)

        update_lgtble!(lg_table, comms, idx1, idx2)
        for j in 1:2
            tv[idx[j]].roster  = update_roster(tv[idx[j]].roster, comms[j])
            tv[idx[j]].teamsht = update_teamsheet(tv[idx[j]].roster; tactic = "N")
        end
    end
    return nothing
end

for i in 1:nsteps
    playreps!(sims)
end

There are different modules and so on involved though. Unless you can identify the issue from looking at that, I'll try to figure out a MWE.

Edit:
My code with the above scheme of nested @batch loops eventually froze my AWS instance. But locally, it runs 3x faster than if I remove it from the inner (playgames) loop.

Neither case seems to be the expected behavior. I'll look into it.

[tetrakai1]

I got something close to a MWE. Try this using eg julia --threads=32 (more than nreps threads), and watch htop thread usage:

using Polyester

###################
# Inner Functions #
###################
# Single-threaded
function getvec(x, i)
   return add2vec(x[i])
end

function add2vec(x)
    for i in eachindex(x)
        x[i] = x[i] + 1
    end
    return x
end

# Multi-threaded
function getvec_batch(x, i)
   return add2vec_batch(x[i])
end

function add2vec_batch(x)
    @batch for i in eachindex(x)
        x[i] = x[i] + 1
    end
    return x
end

###################
# Outer Functions #
###################
# Always multi-threaded
function addall(nreps)
    x  = [rand(100000) for _ in 1:nreps]
    @batch for i in 1:nreps
        x[i] = getvec(x, i)
    end
    return x
end

function addall_batch(nreps)
    x  = [rand(100000) for _ in 1:nreps]
    @batch for i in 1:nreps
        x[i] = getvec_batch(x, i)
    end
    return x
end

# Compile
addall(10);
addall_batch(10);

########
# Test #
########
N     = 10000
nreps = 3
# Inner single-threaded
init_time = time();
for i in 1:N
    res = addall(nreps)
end
elapsed1 = time() - init_time

# Inner multi-threaded
init_time = time();
for i in 1:N
    res = addall_batch(nreps)
end
elapsed2 = time() - init_time

I see three threads used for the "Inner single-threaded" case, but many more for the "Inner multi-threaded" case. My understanding is those should be the same because the inner @batch should be disabled.