Repository navigation
Conversation
The scatter_nd kernel applies each update element with a plain read-modify-write, so when indices repeat a target two units can read the same old value and one update is lost. For bool, or (Add/Max) and and (Mul/Min) have an absorbing value: the kernel now writes only that value (true for or, false for and) and skips the rest, so every concurrent write to a position stores the same value and the result no longer depends on scheduling. Numeric reductions keep the read-modify-write. Part of tracel-ai#6213.
There was a problem hiding this comment.
Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.
An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.
Codecov Report❌ Patch coverage is
❌ Your patch check has failed because the patch coverage (43.28%) is below the target coverage (80.00%). You can increase the patch coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## main #6215 +/- ##
==========================================
- Coverage 66.16% 66.16% -0.01%
==========================================
Files 878 878
Lines 127748 127806 +58
==========================================
+ Hits 84528 84557 +29
- Misses 43220 43249 +29 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Checklist
cargo run-checkshas been executed.changelog policy in CONTRIBUTING.md.
Related Issues/PRs
scatter_ndscatter_nd)Changes
scatter_nd_kernelruns one unit per update element and does a plain read-modify-write, so with repeated targets two units can read the same old value and one update is lost. On bool,Add/Max(or) andMul/Min(and) went through that path asMax/Min.Or and and have an absorbing value, so for bool the kernel now writes only that value and skips the rest (as suggested by @antimora in tracel-ai/burn-onnx#654):
truefor or,falsefor and. Every concurrent write to a position then stores the same value, so the result no longer depends on scheduling, without atomics. This is a#[comptime]ScatterNdWritemode on the existing kernel (Reducekeeps the read-modify-write for numeric data and boolAssign);bool_scatter_ndselects it.Testing
test_bool_scatter_nd_{or,and}_with_repeated_targets: 2048 updates on each of two positions, with a single update that differs from the rest, for all four reductions.--features metal) before the fix, both fail on every run (20/20): or loses the true (Position 0: 0 != 1), and loses the false (Position 1: 1 != 0).scatter/scatter_ndtest passes on Metal (77) and Flex (89).cargo run-checks: passes. One earlier run failed onf16::module::attention::test_attention_prompt_transformer_heads(abs diff 1.26e-3 vs tol 1e-3); that test fails 2 runs out of 3 of--test tensor_f16onmaintoo, so it's unrelated.AI Usage
I used Claude Code for most of the work: it reproduced the race on Metal, wrote the kernel change, the tests and this description, and ran the checks. I followed the work step by step and decided what to submit.