-
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
[RELAY] Port winograd ops to relay #2356
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.
I found an issue that alter_op_layout doesn't work with scalar
e.g. relay.add(relay.nn.conv2d(...), relay.const(1.0))
output of conv2d will fallback to original layout because the inferred layout of constant is undef
python/tvm/relay/build_module.py
Outdated
@@ -151,7 +151,8 @@ def optimize(func, params=None): | |||
func = ir_pass.combine_parallel_conv2d(func) | |||
|
|||
if cfg.pass_enabled("FoldConstant"): | |||
func = ir_pass.fold_constant(func) | |||
with _target.create("llvm"): |
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 new target is unnecessary because FoldConstant
always creates a new llvm target
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.
BTW, can you compile inception_v3 with relay?
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 can compile inception v3 without AlterOpLayout, otherwise I got type_infer.cc:314: the function is provided too many arguments (nullptr)
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.
Fixed
@@ -38,7 +38,7 @@ struct Conv2DAttrs : public tvm::AttrsNode<Conv2DAttrs> { | |||
IndexExpr channels; | |||
Array<IndexExpr> kernel_size; | |||
std::string data_layout; | |||
std::string weight_layout; | |||
std::string kernel_layout; |
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.
Purpose: rename weight_layout
to kernel_layout
to make relay consistent with nnvm and mxnet
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.
👍
3eaf5d0
to
78eaf0e
Compare
python/tvm/relay/build_module.py
Outdated
@@ -222,7 +225,8 @@ def build(func, | |||
cfg = BuildConfig.current | |||
|
|||
with tophub_context: | |||
func = optimize(func, params) | |||
with target: |
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.
May I ask why we need to have with target
here? Is it used by layout altering? It seems that constant folding is not target dependent because it always uses llvm as the target.
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.
@zhiics alter_op_layout relies on current target to query autotvm log
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.
@vinx13 I just saw that. Thanks for your quick response. Can we pass the target to alter layout? I am asking is because for heterogeneous compilation we will pass in multiple targets. It is probably not convenient to know which target to be with
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.
Ok, I passed target as an argument and use it explicitly for target-specific passes. Alter_op_layout calls functions similar to topi compute/topi schedule and will modify the graph, so I think it is not straightforward to port it to heterogeneous compilation.
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.
@merrymercy Thanks. It doesn't really solve the problem. But I think we can keep it like this first because it at least won't break homogeneous execution. I will think about it later in the heterogeneous pass.
src/relay/op/nn/convolution.cc
Outdated
// NOTE: Do not check weight shape here! | ||
// Different backend requires different layout to compute | ||
// the batch gemm stage in winograd efficiently, but we want to | ||
// make this NNVM symbol work for all backends. |
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 we update this comment?
|
||
for (auto new_arg : new_args) { | ||
// NOTE: do not support nested tuple |
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 we need to generically handle nested tuples, the handling of nested tuples has appeared in multiple places including the execution + optimization of AD.
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 leave it to later PRs..
@@ -523,9 +523,8 @@ def _callback(op): | |||
|
|||
##### REGISTER ALTER OP LAYOUT ##### | |||
@conv2d_alter_layout.register(["arm_cpu"]) | |||
def _alter_conv2d_layout_arm(attrs, inputs, tinfos): | |||
def _alter_conv2d_layout_arm(attrs, inputs, tinfos, F): |
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 we provide a documentation comment here with information about the parameters including F?
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.
added
Overall looks good, just some minor comments. |
76e9ac4
to
ebcf761
Compare
ebcf761
to
c397933
Compare
thanks @merrymercy . |
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.
AlterOpLayout changes LGTM.
Since this PR fixes some bugs that other PRs are pending on. I am going to merge this in. The proposed change of |
Thanks to @merrymercy @jroesch @vinx13 @srkreddy1238 @zhiics @yzhliu |
To make alter_op_layout in TOPI support both NNVM and Relay, a new argument
F
is added, which can be eitherrelay.op
ornnvm.sym
.Know issues: