6
6
import torch
7
7
8
8
from executorch import exir
9
+ from executorch .backends .nxp .aten_passes .neutron_aten_pass_manager import (
10
+ NeutronAtenPassManager ,
11
+ )
9
12
from executorch .backends .nxp .neutron_partitioner import NeutronPartitioner
10
13
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
14
14
from executorch .backends .nxp .quantizer .neutron_quantizer import NeutronQuantizer
15
15
from executorch .exir import (
16
16
EdgeCompileConfig ,
@@ -27,7 +27,7 @@ def _quantize_model(model, calibration_inputs: list[tuple[torch.Tensor]]):
27
27
quantizer = NeutronQuantizer ()
28
28
29
29
m = prepare_pt2e (model , quantizer )
30
- for _i , data in enumerate ( calibration_inputs ) :
30
+ for data in calibration_inputs :
31
31
m (* data )
32
32
m = convert_pt2e (m )
33
33
@@ -48,12 +48,9 @@ def to_quantized_edge_program(
48
48
model , example_input , strict = True
49
49
)
50
50
51
- # TODO(Robert Kalmar) uncoment when NXP passes are ported to main
52
51
# 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
57
54
exir_program_aten__module_quant = _quantize_model (
58
55
exir_program_aten_module , calibration_inputs
59
56
)
0 commit comments