[TOPI] Allow conv definition to have custom kernel layout #11936
Conversation
tkonolige
left a comment
tkonolige
left a comment
There was a problem hiding this comment.
Looks pretty good, but I don't see any tests for different kernel_layouts. Could you add some?
|
@tkonolige There is no good ways to test numerical results for arbitrary combinations of layouts due to lack of schedule. Therefore, I added some tests to check |
| out_dtype="float32", | ||
| ), | ||
| ], | ||
| [A, W, topi.nn.conv2d_nhwc(A, W, stride, padding, dilation, "float32")], |
There was a problem hiding this comment.
topi.nn.conv2d_nhwc just calls topi.nn.conv2d under the hood, so this test isn't really checking anything.
There was a problem hiding this comment.
topi.nn.conv2d_nhwc is already covered by existing test cases, but I didn't find any test cases invoke topi.nn.conv2d so we just need to check they are the same
Under the hood, both topi.nn.conv2d_nhwc and topi.nn.conv2d calls topi.nn.conv and cover the code path with explicit kernel layout.
|
@vinx13 I think you can use |
vinx13
force-pushed
the
feat/conv_generic_kernel_layout
branch
3 times, most recently
from
c8daa85 to
1f1b428
Compare
python/tvm/topi/nn/conv2d.py
Outdated
|
|
||
| def conv2d(input, filter, strides, padding, dilation, layout="NCHW", out_dtype=None): | ||
| def conv2d( | ||
| input, filter, strides, padding, dilation, data_layout="NCHC", kernel_layout="", out_dtype=None |
There was a problem hiding this comment.
maybe data_layout="NCHW" ?
|
|
||
| kernel_layout : Optional[str] | ||
| layout of kernel. If unspecified, use default layout inferred from data_layout. "OHWI" if | ||
| data_layout == "NCHW", "HWIO" if data_layout == "NHWC". |
There was a problem hiding this comment.
A little bit confused to me, shouldn’t "OIHW" & "NCHW" a default pair? I can also see this pattern at line 248.
tkonolige
left a comment
tkonolige
left a comment
There was a problem hiding this comment.
Looks good! Thanks for adding the tests!
vinx13
force-pushed
the
feat/conv_generic_kernel_layout
branch
from
cf01ebb to
249be49
Compare
vinx13
force-pushed
the
feat/conv_generic_kernel_layout
branch
from
249be49 to
a9eb3d4
Compare
* [TOPI] Allow conv definition to have custom kernel layout * add tests * fix * fix
* [TOPI] Allow conv definition to have custom kernel layout * add tests * fix * fix
4 participants
This PR added an optional parameter
kernel_layouttoconv. This allows arbitrary data and kernel layout combination to be lowered from Relay, and scheduled by auto / meta schedule.cc @junrushao1994 @masahi @tkonolige