JuliaSIMD / JuliaSIMD/Polyester.jl

multiple loops only uses 2 threads

Open
#24 3 comments 1 reaction 0 assignees View on GitHub
Dominant language
Julia
Stars
285
Forks
20
PR merge metrics
No merged PRs in 30d

Description

I have three loops and because the outermost loop and innermost loop has a relationship that their summation is roughly a constant, thus I need to divide the thread evenly on these three loops, but somehow `@batch` only uses two cores

the implementation

```julia
using Test
using LoopVectorization
using LinearAlgebra
using BQCESubroutine
using CheapThreads
using YaoLocations: YaoLocations, plain, Locations, CtrlLocations, merge_locations
using BQCESubroutine: ctrl_offset, bcomspace, bsubspace, broutine!, log2dim
using ThreadingUtilities
using BenchmarkTools

function broutine2x2!(S::AbstractMatrix{Complex{T}}, U::AbstractMatrix, locs::Locations) where T
plain_locs = plain(locs)::Int
step_1 = 1 << (plain(locs) - 1)
step_2 = 1 << plain(locs)

U11 = U[1, 1]; U12 = U[1, 2];
U21 = U[2, 1]; U22 = U[2, 2];
@inbounds @batch for j in 0:step_2:size(S, 2)-step_1, b in 1:size(S, 1), i in j:j+step_1-1
ST1 = U11 * S[b, i+1] + U12 * S[b, i+step_1+1]
ST2 = U21 * S[b, i+1] + U22 * S[b, i+step_1+1]

S[b, i+1] = ST1
S[b, i+step_1+1] = ST2
end
return S
end
```

it works fine when the outer loop is large enough, but if you change `locs` to make the outer loop small, e.g

```julia
N = 20
S = rand(ComplexF64, 1, 1 << N)
U = rand(ComplexF64, 2, 2)
locs = Locations(19)
@benchmark broutine2x2!(S, U, locs) setup=(S=$(copy(S)))
```

it only uses two cores in this case. since `@batch` doesn't error like `Threads.@threads` on multiple loops, I assume it actually handles multiple loops? but not sure how to make this work.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.