Add stream_loads and stream_stores scheduling directives#9207
Conversation
|
This implementation was super straightforward, but also super messy. It would be nice if we had an annotations system so that places that just forward those annotations don't need to be updated for every new feature. |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #9207 +/- ##
=======================================
Coverage ? 69.64%
=======================================
Files ? 255
Lines ? 78880
Branches ? 18859
=======================================
Hits ? 54936
Misses ? 18303
Partials ? 5641 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
e0d2142 to
4fe8903
Compare
538dbc8 to
057fec5
Compare
| // after the whole (specialized) production is equivalent to, | ||
| // and simpler than, fencing inside each branch individually. | ||
| if (any_specialization_requests_streaming(def)) { | ||
| Expr fence = Call::make(Int(32), Call::stream_store_fence, {}, Call::Intrinsic); |
There was a problem hiding this comment.
So I guess the StreamingStores Stmt doesn't actually imply code then, and is purely communication to future lowering passes?
There was a problem hiding this comment.
It implies a fence. However, rather than place a fence right then-and-there, this implementation places it after the specialization branch (if any) for the stage. This is sound because stream_stores only applies to stages with only pure Vars and RDoms. So that point assuredly post-dominates all such stores and pre-dominates any potential loads. There might be a later point in some programs with the same soundness property (especially when streaming out from the last stage), but that analysis is deferred to future work. This is safe and not too inefficient.
| Stmt Simplify::visit(const StreamingStore *op) { | ||
| Stmt body = mutate(op->body); | ||
| if (is_no_op(body)) { | ||
| return Evaluate::make(0); |
There was a problem hiding this comment.
I just did what Atomic did above?
There was a problem hiding this comment.
That actually looks correct since it drops a side-effect-free evaluation. The is_no_op predicate is true iff body is either:
- Undefined (i.e.
Stmt()) - An
Evaluatenode whosevaluesatisfiesis_pure
So Evaluate(0) is more defined than (1) and no more work than (2). Seems fine to me.
abadams
left a comment
abadams
left a comment
There was a problem hiding this comment.
Approving in advance so you're not blocked on me, but I did leave a few comments
Co-authored-by: GPT 5.6 Sol <noreply@openai.com> Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
a9baf07 to
44d6168
Compare
f044531 to
1c5f8ba
Compare
3 participants
Implements the
stream_loadsandstream_storesdirectives as proposed in #9206.stream_loads- ensure that loads this stage makes are streamed. It has two modes.stream_loads()-- stream all direct loads after inlining, excluding self-loads.g.stream_loads({f1, ..., fN})-- stream loads from the named funcs after inlining. Namingghere is auser_error.stream_stores- ensure that stores to this func are streamed on a per-stage basis. A stage can have this specialization only if all of its RVars are pure (meaning that no two stores can alias). If any stage-specialization requests streaming, a fence is placed after the specialization block. This is vacuously true for pure funcs, RDom-less updates, and stages without specializations.This fully recovers performance of
memcpyon my Apple M3.Fixes #9206
Checklist