Parallelize contiguous f32 elementwise ops over the barrier pool - #45
Parallelize contiguous f32 elementwise ops over the barrier pool#45DrJesseGlass wants to merge 4 commits into
Conversation
Thread the contiguous f32 unary/binary elementwise ops (SiLU, SwiGLU mul, residual adds, scaling) across the barrier pool at the unary_impl/binary_impl dispatch point, split into one disjoint output range per worker. Bit-identical to the serial path; gated by a size threshold so small tensors stay serial. Recovers multi-thread prefill scaling where these single-threaded elementwise ops were an Amdahl bottleneck (matmuls scale ~Nx, these did not). Knobs: CANDLE_PAR_ELEMWISE=0 disables; CANDLE_PAR_ELEMWISE_MIN sets the element threshold (default 16384).
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e6c45d3b51
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let pool = crate::utils::barrier_pool(); | ||
| let n_total = pool.n_workers() + 1; | ||
| let per = len.div_ceil(n_total); | ||
| pool.execute(|tid| { |
There was a problem hiding this comment.
Fall back for nested barrier-pool calls
When a caller is already inside candle_core::utils::barrier_pool().execute (the pool is public and used directly by candle-nn kernels), this default-on fast path calls pool.execute again for any large contiguous f32 elementwise op. BarrierPool::execute holds its call_lock until all workers finish, so a worker or the main participant that reaches this nested call blocks on the lock while the outer call waits for that participant, deadlocking; detect reentrancy and use the serial fallback in that case.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f1586ea894
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| d[i] = s[i] * s[i] - 0.5; // silu-like nonlinearity stand-in | ||
| } | ||
| } | ||
| let par = par_unary_vec_f32(&src, &layout, fv).expect("should parallelize"); |
There was a problem hiding this comment.
Make elementwise tests ignore opt-out env vars
When the documented CANDLE_PAR_ELEMWISE=0 knob is set, par_unary_vec_f32 returns None, so this expect panics even though the configuration is supported; the binary test below has the same problem, and I verified CANDLE_PAR_ELEMWISE=0 cargo test -p candle-core par_elemwise_tests --lib --quiet fails on both. Consider scoping the env vars for these tests or asserting conditionally so valid serial-forcing runs do not break the test suite.
Useful? React with 👍 / 👎.
|
@codex review |
|
Codex Review: Didn't find any major issues. 👍 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
Thread the contiguous f32 unary/binary elementwise ops (SiLU, SwiGLU mul, residual adds, scaling) across the barrier pool at the unary_impl/binary_impl dispatch point, split into one disjoint output range per worker. Bit-identical to the serial path; gated by a size threshold so small tensors stay serial.
Recovers multi-thread prefill scaling where these single-threaded elementwise ops were an Amdahl bottleneck (matmuls scale ~Nx, these did not). Knobs: CANDLE_PAR_ELEMWISE=0 disables; CANDLE_PAR_ELEMWISE_MIN sets the element threshold (default 16384).