Add pass to replace some ops with channels last variants.#20788
Add pass to replace some ops with channels last variants.#20788MartinPavella wants to merge 3 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20788
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 2 Unrelated FailuresAs of commit 7e23e2d with merge base a19d1ba ( NEW FAILURE - The following job has failed:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
BROKEN TRUNK - The following job failed but was present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
|
This PR needs a
|
|
@AdrianLundell this is the first part of #20094. I am planning on introducing the second part (a pass which makes the graph contiguous) separately. |
| _NHWC_TO_NCHW_PERM: list[int] = [0, 3, 1, 2] | ||
|
|
||
|
|
||
| def _requires_rank(rank: int) -> Callable[[torch.fx.Node], bool]: |
There was a problem hiding this comment.
Why not allow (CHW,) -> (HWC) as well?
There was a problem hiding this comment.
Good point, I missed that use case.
I will make the relevant changes.
There was a problem hiding this comment.
For convolution, unsqueeze and squeeze operators are inserted during lowering to edge, which explicitly add/remove the batch size of 1, and the exisitng code with 4D permutations works correctly.
For avg_pool2d, this doesn't happen.
I have decided to update this pass to also insert the [un]squeeze for implicit batch cases, which should solve all related issues while remaining consistent with convolution.
|
Great to see the collaboration on this, just one question! |
rascani
left a comment
rascani
left a comment
There was a problem hiding this comment.
Can you add BUCK support in the targets.bzl?
runtime.python_library(
name = "replace_ops_with_channels_last_variants",
srcs = [
"replace_ops_with_channels_last_variants.py",
],
visibility = [
"//executorch/backends/...",
],
deps = [
"//caffe2:torch",
":channels_last_ops",
"//executorch/exir:pass_base",
],
)
runtime.python_test(
name = "test_replace_ops_with_channels_last_variants",
srcs = [
"test/test_replace_ops_with_channels_last_variants.py",
],
deps = [
"//caffe2:torch",
":channels_last_ops",
":replace_ops_with_channels_last_variants",
"//executorch/exir:lib",
"fbsource//third-party/pypi/pytest:pytest",
],
)
|
|
||
| args = list(node.args) | ||
|
|
||
| with graph.inserting_before(node): |
There was a problem hiding this comment.
Might want to consider using the super().call_operator() style of rewrites, which eliminates having to deal with meta & graph cleanup. Not a blocker though.
def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
self._modified = False
result = super().call(graph_module)
return PassResult(result.graph_module, self._modified)
def call_operator(self, op, args, kwargs, meta: NodeMetadata) -> ProxyValue:
spec = self.op_map.get(op)
if spec is None:
return super().call_operator(op, args, kwargs, meta)
val = meta["val"]
contiguous_dim_order = tuple(range(len(val.shape)))
if val.dim_order() != contiguous_dim_order or (
spec.filter_fn is not None and not spec.filter_fn(val)
):
return super().call_operator(op, args, kwargs, meta)
new_args = list(args)
for idx in spec.input_indices:
new_args[idx] = super().call_operator(
exir_ops.edge.channels_last.permute_copy.default,
(new_args[idx], _NCHW_TO_NHWC_PERM),
{},
meta,
)
nhwc_out = super().call_operator(spec.target, tuple(new_args), kwargs, meta)
self._modified = True
return super().call_operator(
exir_ops.edge.channels_last.permute_copy.default,
(nhwc_out, _NHWC_TO_NCHW_PERM),
{},
meta,
)
There was a problem hiding this comment.
Thank you for the suggestion. After consideration, I prefer your solution. I have updated the implementation.
There was a problem hiding this comment.
Update: When I tried to add support for max_pool2d_with_indices, I realized that the call_operator() style won't work. This is because permute and squeeze operators must be added after the getitem nodes of the max_pool2d_with_indices, so the resulting sugbraph would have multiple distinct outputs (which as far as I can tell is incompatible with the call_operator() pass style).
Therefore, I have switched back to the original approach.
|
Implemented changes based on reviews. |
When replacing aten.avg_pool2d` with `channels_last.avg_pool2d`, some attributes of the aten operator can be left out, and must be replaced by default values. This can either be done in the "to channels last" transformation pass, or in the definition of the operator. Adding the default attributes of the aten operator to the channels_last operator seems cleaner to me.
The pass currently supports `convolution` and `avg_pool2d`, as those are the only 2 operators with a channels last variant. The pass can be easily extended to more ops in the future.
7a3b692 to
7e23e2d
Compare
AdrianLundell
left a comment
AdrianLundell
left a comment
There was a problem hiding this comment.
Nice, it is a bit unfortunate that there are these arbitrary differences between ops which must be handled everywhere but not much to do about it.
|
@MartinPavella - Would you mind taking care of the EasyCLA? This is part of the migration of ExecuTorch to the PyTorch/Linux Foundation. |
3 participants
Summary
Add a transformation pass to replace some operators with their channels last variant. Currently only
convolutionandavg_pool2dare implemented, as those are the only operators with a channels last variant. In the future, the pass can be easily extended to support more operators.Test plan
Tested by
pytest backends/transforms/test/test_replace_ops_with_channels_last_variants.py.