-
Notifications
You must be signed in to change notification settings - Fork 3.5k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[TensorIR][M1c] LowerInitBlock #7806
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
minor nits, overall LGTM
3ff3602
to
87d5a58
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Overall LGTM.
One suggestion: Since we have defined a clear stage (i.e., M1c) for this series of PRs, it would be better for all PR submitters to use the unified PR titles (for exampe: [TensorIR][M1c] ...) so that they can be easily queried.
return Block(n); | ||
} | ||
|
||
static Stmt RealizeInitBlock(const Stmt& init, const Array<IterVar>& iter_vars) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Could you clarify the meaning of the word "realize" here?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What about "concretize"? My understanding is that the init-block is something "conceptual" for initializing "zero point" of a reduction, and this pass is used to make it concrete, lowering to "if-then-else" statement
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I got your point, and I feel "LowerInitBlock" itself is the best name to describe what it does, so maybe the simplest way is manually inlining this function so that we won't be bothered by its name, or just give it a general helper name to let people know this is just the implementation of lowering init block.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This helper is something a bit standalone (making the init block to If-Then-Else), and inlining this function might make the logic in the function above a bit vague, so what about we just call it "DoLowering"?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
DoLowering sounds good to me.
|
||
|
||
def test_lower_reduction(): | ||
origin_mod = WithInit() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just a note here: per discussion with @Hzfengsy, we need to remove "()" here in future PRs, so that the script parser API is consistent across PrimFunc and IRModule
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Co-authored-by: Junru Shao <junrushao1994@gmail.com> Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
Thanks @comaniac @junrushao1994 for reviewing. Thanks @Hzfengsy |
Co-authored-by: Junru Shao <junrushao1994@gmail.com> Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
Co-authored-by: Junru Shao <junrushao1994@gmail.com> Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
Co-authored-by: Junru Shao <junrushao1994@gmail.com> Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
Co-authored-by: Junru Shao <junrushao1994@gmail.com> Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
Co-authored-by: Junru Shao <junrushao1994@gmail.com> Co-authored-by: Ruihang Lai <lairuihangdongdong@qq.com>
This PR is the first part of the stage M1c, TensorIR upstreaming plan (#7527) on lowering init block.
We just use
if
branch to replace the init block. e.g.Before
After
Please note that we only provide the default behavior for lowering the init part, which is not performent. We will provide primitives for users to manually "decompose" init part.
cc @tqchen @junrushao1994 @jroesch @comaniac @jcf94