diff --git a/backends/arm/_passes/arm_pass.py b/backends/arm/_passes/arm_pass.py index 0d48eab0cc0..019846975b7 100644 --- a/backends/arm/_passes/arm_pass.py +++ b/backends/arm/_passes/arm_pass.py @@ -87,9 +87,23 @@ def get_name(pass_) -> str: def call_operator(self, op, args, kwargs, meta, updated: Optional[bool] = False): ops_without_quantized_fake_kernel = { + exir_ops.edge.aten.atan.default, exir_ops.edge.aten.bmm.default, + exir_ops.edge.aten.cos.default, + exir_ops.edge.aten.erf.default, + exir_ops.edge.aten.exp.default, exir_ops.edge.aten.leaky_relu.default, + exir_ops.edge.aten.log.default, + exir_ops.edge.aten.log1p.default, + exir_ops.edge.aten.pow.Tensor_Scalar, + exir_ops.edge.aten.reciprocal.default, + exir_ops.edge.aten.remainder.Scalar, + exir_ops.edge.aten.rsqrt.default, + exir_ops.edge.aten.sigmoid.default, + exir_ops.edge.aten.sin.default, + exir_ops.edge.aten.sinh.default, exir_ops.edge.aten.softplus.default, + exir_ops.edge.aten.tanh.default, } if ( op in ops_without_quantized_fake_kernel diff --git a/backends/arm/_passes/replace_scalar_with_tensor_pass.py b/backends/arm/_passes/replace_scalar_with_tensor_pass.py index 5928911f93f..bf667aec11a 100644 --- a/backends/arm/_passes/replace_scalar_with_tensor_pass.py +++ b/backends/arm/_passes/replace_scalar_with_tensor_pass.py @@ -123,18 +123,16 @@ def call_operator(self, op, args, kwargs, meta): if included_ops == {}: raise ValueError("Profile must support at least INT or FP") - if op in TableOps.included_ops(): - # Do not handle quantized table ops; forward unchanged. - input_qparams = meta.data.get("input_qparams", {}) - output_qparams = meta.data.get("output_qparams", {}) - if len(input_qparams) > 0 and len(output_qparams) > 0: - return self._call_quantized_op_without_fake_kernel( - op, args, kwargs, meta - ) + # Skip quantized table ops - they will be turned to tables regardless. + if op in TableOps.included_ops() and self._is_quantized_meta(meta): + return ArmPass.call_operator(self, op, args, kwargs, meta) + # Apply ReplaceScalarWithTensorArgPass to included ops. if op in included_ops: - # Include this op based on the current profile. return super().call_operator(op, args, kwargs, meta) - else: - # Do not handle; forward unchanged. - return ArmPass.call_operator(self, op, args, kwargs, meta) + # Due to MRO, we need to call ExportPass's call_operator to not end up in ReplaceScalarWithTensorArgPass + if op in _all_ops: + return ExportPass.call_operator(self, op, args, kwargs, meta) + + # Otherwise, we can call the normal ArmPass call_operator + return ArmPass.call_operator(self, op, args, kwargs, meta)