Skip to content

Commit 270b8af

Browse files
committed
Integrate NeutronAtenPassManager passes into pipeline
1 parent 0060581 commit 270b8af

File tree

1 file changed

+6
-9
lines changed

1 file changed

+6
-9
lines changed

backends/nxp/tests/executorch_pipeline.py

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,11 @@
66
import torch
77

88
from executorch import exir
9+
from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import (
10+
NeutronAtenPassManager,
11+
)
912
from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner
1013
from executorch.backends.nxp.nxp_backend import generate_neutron_compile_spec
11-
12-
# TODO (Robert Kalmar) Uncomment when NXP passes are ported to main
13-
# from executorch.backends.nxp.pytorch_passes.nxp_pytorch_pass_manager import NXPPyTorchPassManager
1414
from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
1515
from executorch.exir import (
1616
EdgeCompileConfig,
@@ -27,7 +27,7 @@ def _quantize_model(model, calibration_inputs: list[tuple[torch.Tensor]]):
2727
quantizer = NeutronQuantizer()
2828

2929
m = prepare_pt2e(model, quantizer)
30-
for _i, data in enumerate(calibration_inputs):
30+
for data in calibration_inputs:
3131
m(*data)
3232
m = convert_pt2e(m)
3333

@@ -48,12 +48,9 @@ def to_quantized_edge_program(
4848
model, example_input, strict=True
4949
)
5050

51-
# TODO(Robert Kalmar) uncoment when NXP passes are ported to main
5251
# Run pre-processing passes of the float32 aten dialect program.
53-
# pass_manager = NXPPyTorchPassManager(exir_program_aten)
54-
# pass_manager.run() # All passes by default.
55-
56-
exir_program_aten_module = exir_program_aten.module()
52+
pass_runner = NeutronAtenPassManager()
53+
exir_program_aten_module = pass_runner(exir_program_aten.module()).graph_module
5754
exir_program_aten__module_quant = _quantize_model(
5855
exir_program_aten_module, calibration_inputs
5956
)

0 commit comments

Comments
 (0)