NXP backend: Create NeutronAtenPassManager with initial BatchNorm fusing passes #10579
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/10579
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 5d69743 with merge base d7030aa ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
|
@pytorchbot label "module: nxp" "release notes: nxp" |
pytorch-bot
bot
added
module: nxp
|
|
||
| bn_node = node | ||
|
|
||
| if not _is_linear(bn_node.args[0]): |
There was a problem hiding this comment.
here and above, you can be defensive and also make sure args len == 1. there is nothing stopping someone from adding a skip connection after conv/linear and before batchnorm, or some other op. Can't say i've seen it in practice though
There was a problem hiding this comment.
Added check for number of users of Conv/Linear op.
|
|
||
| import torch | ||
|
|
||
| from executorch.backends.nxp.aten_passes.fuse_batch_norm_with_conv_pass import ( |
There was a problem hiding this comment.
nit: can become unwieldly since there are a few other permutations of "batch norm" fusion, maybe just one file for fuse_batch_norm would be good.
There was a problem hiding this comment.
Agree with you. Maybe we can create subdirectories in backend/transforms based on what dialect/ir is pass targeting?
|
LGTM but @digantdesai thoughts on the aten pass setup? |
skywall
changed the title
skywall
force-pushed
the
feature/nxf93343/neutron-batch-norm-fusion-passes
branch
from
d1913ae to
5d69743
Compare
|
I have also disabled some tests that require changes from #10258 . We'll update PR that will be merged last. |
| program = torch.export.export_for_training(module, example_input, strict=True) | ||
| og_module = program.module() | ||
|
|
||
| pm = NeutronAtenPassManager() |
There was a problem hiding this comment.
curious why does it have to be before export? Typically this is done in neutron.preprocess(). The reason is, it is not easy to run these nodes if Neutron partitioner fails to pick them up. There are several examples of this in Arm, XNNPACK etc.
| class FuseBatchNormWithLinearPass(PassBase): | ||
| """The executorch batch normalization carries out the following computation [1]. | ||
|
|
||
| (x - mean) / (var + eps) * W + B |
There was a problem hiding this comment.
| (x - mean) / (var + eps) * W + B | |
| (x - mean) / sqrt(var + eps) * W + B |
| from torch.fx import GraphModule, Node | ||
| from torch.fx.passes.infra.pass_base import PassBase, PassResult | ||
| from torch.nn.parameter import Parameter | ||
| from torch.nn.utils import fuse_conv_bn_weights |
There was a problem hiding this comment.
glad you are using this as is
| conv_args.append(None) | ||
|
|
||
| weight_attr_name = conv_weight_node.target | ||
| _assign_attr( |
There was a problem hiding this comment.
you can do this without using internal methods, and insert new params and delete old params https://github.com/pytorch/executorch/blob/main/backends/transforms/utils.py#L65 (assuming you are doing this after export)
| assert np.allclose(out1, out2, atol=1.2e-7) | ||
|
|
||
|
|
||
| # TODO - enable within PR #10196 |
Summary
Add NeutronAtenPassManager that by default runs passes for optimal model conversion. Initially, only passes for BatchNorm + Conv/Linear fusing are added. I haven't intentionally moved passes into backend/transforms because added passes are executed on Aten ops and not Edge ops. Please let me know if I should follow different approach.
This pass manager is planned to be included in Neutron backend's pipeline - https://github.com/pytorch/executorch/pull/10196/files#diff-b3b0844670da322bb4c245ae440e8c71378d631bb94334148322cdf758369baaR51
cc @digantdesai @JakeStevens @robert-kalmar