Add stream_loads and stream_stores scheduling directives#9207
Add stream_loads and stream_stores scheduling directives#9207alexreinking wants to merge 2 commits into
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. |
| src.set_host_alignment(32); | ||
| dst.output_buffer().set_host_alignment(32); | ||
|
|
||
| dst.vectorize(x, 32, TailStrategy::GuardWithIf); |
There was a problem hiding this comment.
You'd apply it here then.
There was a problem hiding this comment.
Perhaps two versions: with streaming stores and loads, and without, just as informative output when someone runs this test?
There was a problem hiding this comment.
We don't really show "wrong" schedules in the other performance tests, do we?
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #9207 +/- ##
=======================================
Coverage ? 69.76%
=======================================
Files ? 255
Lines ? 78828
Branches ? 18843
=======================================
Hits ? 54997
Misses ? 18267
Partials ? 5564 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
alexreinking
force-pushed
the
alexreinking/non-temporal
branch
from
e0d2142 to
4fe8903
Compare
| } else if (op->is_intrinsic(Call::stream_store_fence)) { | ||
| // The C backend does not implement non-temporal loads/stores. | ||
| rhs << "0"; |
There was a problem hiding this comment.
Something to discuss: what do we want to do here. Warn? Silently drop (this)?
There was a problem hiding this comment.
I vote warn. We're ignoring a scheduling directive.
alexreinking
force-pushed
the
alexreinking/non-temporal
branch
from
d16165a to
9d1398d
Compare
Co-authored-by: GPT 5.6 Sol <noreply@openai.com> Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
alexreinking
force-pushed
the
alexreinking/non-temporal
branch
from
538dbc8 to
057fec5
Compare
| * resulting Load node(s) and then discards it -- it should not survive | ||
| * past that point. */ | ||
| struct StreamingLoads : public StmtNode<StreamingLoads> { | ||
| std::optional<std::vector<std::string>> names; |
There was a problem hiding this comment.
Is an empty vector meaningful, or could it be used instead of nullopt?
There was a problem hiding this comment.
It is meaningful:
- nullopt = stream from all loads except self
- vector (any size) = stream from this set of loads
Therefore, an empty vector means "stream nothing". That's a valid thing to say if you're metaprogramming the schedule and some combination of parameters leads to that decision.
There was a problem hiding this comment.
But isn't "stream nothing" representable by just not having the IR node at all? Or might you have a stream-nothing context inside of a stream-some-stuff context? Actually, what does it even mean to have nested StreamingLoad nodes? Does the inner one override the outer, or is it the union?
There was a problem hiding this comment.
It should not be possible to nest them. A Provide node is immediately surrounded by StreamingLoads/StreamingStores/Atomic in that order, with no opportunity to nest any of them. I chose that implementation over complicating the Provide and especially Call nodes.
Yes, it's possible to represent it by just removing the IR node. But that's true of constant predicates, extent-1 RDoms, etc.
There was a problem hiding this comment.
Maybe add the nesting rule to the comment. Agree that empty vector as stream-nothing is reasonable. Does this imply that in .stream_loads({foo}).stream_loads({}), the second overrides the first, as opposed to being a no-op? I guess it might show up in code like:
for (auto f : {a, b, c, d, e}) {
f.compute_at(...).tile(...).vectorize(...).stream_loads(...);
}
// This specific func shouldn't actually stream any loads
c.stream_loads({})
There was a problem hiding this comment.
Yeah, I think a second declaration would override a first.
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