-
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
[TIR] add support for multi-blocking layout and their transformation #9996
Conversation
Nice, is it possible to remove custom layout transform in tvm/python/tvm/topi/x86/conv2d_alter_op.py Lines 171 to 181 in d1dafbd
|
@masahi Thanks for your timely comment. |
I'm more than happy to see those ad-hoc transform removed, can you do that? I'd also like to remove |
Yeah, it's a pleasure to make the code clean and clear. I'll search the code for workarounds about multi-blocking and refine them. |
Sounds great!
I think the one we already discussed is the only instance of such a workaround. |
* \note this function does eager constant folding for | ||
* shape types(int32, int64) when possible. | ||
*/ | ||
TVM_DLL PrimExpr shapediv(PrimExpr a, PrimExpr b, Span span = Span()); |
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.
the name shapediv could be confusing given indexdiv used the other case, how about nonneg_ceildiv?
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 think it's a kind of symmetry, as indexdiv
is an alias of floordiv
to prevent access out-of-boundary, and shapediv
is an alias of ceildiv
to prevent the shrink of a Tensor.
If this is confusing, I prefer to remove indexdiv
and shapediv
since they are just aliases of floordiv
and ceildiv
now, or we can keep them and add check codes for non-negative then change their names to nonneg_floordiv/ceildiv.
@yangulei Any update? I'm doing some VNNI stuff, I want to transform like
|
Yeah I think that's possible. If I understand correctly, the point is to have 16 x 4 loop in the inner-most axis, so something like |
Yes, the transformation from "NK" to "NK16n4k" is well supported, even without this PR. |
d023ec1
to
aab406b
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.
LGTM modulo one clarifying question
…pache#9996) * add ceildiv and shapediv * add boundary checking in layout_transform * support multi-blocking and shape padding * refine the log for shape transform * add test for multi-blocking layout transform * delete unwanted comments * remove workaround * fix lint errors
Main works in this PR including:
With those enhancement, a tensor with layout="OIHW" could be transformed to a tensor with layout="OHWI16i64o2i".