[TensorIR][M2a] Compute-At #8943
Conversation
junrushao
requested review from
Hzfengsy,
MarisaKirisame,
ZihengJiang,
anijain2305,
areusch,
comaniac,
jroesch,
jwfromm,
kparzysz-quic,
masahi,
mbrookhart,
merrymercy,
slyubomirsky,
tqchen,
vinx13,
wweic,
yzhliu and
zhiics
as code owners
This PR is part of the TensorIR upstreaming effort (#7527), which adds the following schedule primitives: * `compute-at` * `reverse-compute-at` Co-authored-by: Bohan Hou <[email protected]> Co-authored-by: Ruihang Lai <[email protected]> Co-authored-by: Hongyi Jin <[email protected]> Co-authored-by: Wuwei Lin <[email protected]> Co-authored-by: Siyuan Feng <[email protected]>
MasterJH5574
left a comment
MasterJH5574
left a comment
There was a problem hiding this comment.
Leave the implementation and changes of state.cc to tomorrow :-)
| bool CalculateAffineFlag(const ScheduleState& self, const StmtSRef& block_sref) { | ||
| if (block_sref->parent == nullptr) { | ||
| return true; | ||
| } |
There was a problem hiding this comment.
I have a question for a long period of time: if a block has no IterVar (i.e., opaque block), does it have affine bindings?
There was a problem hiding this comment.
I don't actually have a good answer to this question. CC: @Hzfengsy
There was a problem hiding this comment.
That's a good question. I think affine flag for opaque block makes no sense to me, no matter it's affine or not.
Blocks with affine bindings have more schedule abilities than other blocks, while opaque blocks have no schedule abilities.
|
@MasterJH5574 I fixed the wording issue on "compact dataflow" in the doc. Please take another look. Thanks! |
|
I notice that boolean template parameters are widely used in this PR. Just curious that, what's the edge of boolean template parameters over normal boolean parameters? |
|
Just finished a whole review of the entire PR. Thanks @junrushao1994 for the tireless work! Please take a look at my comments. I believe that this PR can be merged very soon :-) |
|
@MasterJH5574 So the template thing here basically does compile time code generation to ensure it is zero overhead |
|
@MasterJH5574 Would you mind taking another look? Also I didn't find the docs you mention that I didn't change, so it would be great if you can just use "suggest" on those lines you want me to rephrase - would be much easier for me. |
|
@MasterJH5574 Thanks for the careful review! Without your help, I cannot find the spot I mistakes one primitive's doc for the other lol. Thanks a lot! |
Hzfengsy
left a comment
Hzfengsy
left a comment
There was a problem hiding this comment.
LGTM. Thanks @junrushao1994.
MasterJH5574
left a comment
MasterJH5574
left a comment
There was a problem hiding this comment.
Thanks @junrushao1994! I think this PR is in a really good state now :-)
|
There's another issue that I want @junrushao1994 to take a look at. I've informed him privately. So we may wait for his response first. |
This PR is part of the TensorIR upstreaming effort (apache#7527), which adds the following schedule primitives: * `compute-at` * `reverse-compute-at` Co-authored-by: Bohan Hou <[email protected]> Co-authored-by: Ruihang Lai <[email protected]> Co-authored-by: Hongyi Jin <[email protected]> Co-authored-by: Wuwei Lin <[email protected]> Co-authored-by: Siyuan Feng <[email protected]>
This PR is part of the TensorIR upstreaming effort (apache#7527), which adds the following schedule primitives: * `compute-at` * `reverse-compute-at` Co-authored-by: Bohan Hou <[email protected]> Co-authored-by: Ruihang Lai <[email protected]> Co-authored-by: Hongyi Jin <[email protected]> Co-authored-by: Wuwei Lin <[email protected]> Co-authored-by: Siyuan Feng <[email protected]>
4 participants
This PR is part of the TensorIR upstreaming effort (#7527), which adds the following schedule primitives:
compute-atreverse-compute-atCo-authored-by: Bohan Hou [email protected]
Co-authored-by: Ruihang Lai [email protected]
Co-authored-by: Hongyi Jin [email protected]
Co-authored-by: Wuwei Lin [email protected]
Co-authored-by: Siyuan Feng [email protected]