diff --git a/backends/arm/_passes/conv1d_unsqueeze_pass.py b/backends/arm/_passes/conv1d_unsqueeze_pass.py index 591a14c455b..7bc86a4fcf2 100644 --- a/backends/arm/_passes/conv1d_unsqueeze_pass.py +++ b/backends/arm/_passes/conv1d_unsqueeze_pass.py @@ -7,6 +7,7 @@ from typing import Set, Type +from executorch.backends.arm._passes import ArmOpTargetedPass from executorch.backends.arm._passes.convert_squeezes_to_view import ( ConvertSqueezesToViewPass, ) @@ -15,10 +16,11 @@ from executorch.backends.transforms.convert_conv1d_to_conv2d_pass import ( ConvertConv1dToConv2dPass, ) +from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass -class Conv1dUnsqueezePass(ConvertConv1dToConv2dPass): +class Conv1dUnsqueezePass(ConvertConv1dToConv2dPass, ArmOpTargetedPass): """Arm wrapper for the shared Conv1d-to-Conv2d transform.""" _passes_required_after: Set[Type[ExportPass]] = { @@ -26,3 +28,4 @@ class Conv1dUnsqueezePass(ConvertConv1dToConv2dPass): RewriteConvPass, SizeAdjustInputPass, } + target_ops = (exir_ops.edge.aten.convolution.default,) diff --git a/backends/arm/test/ops/test_conv1d.py b/backends/arm/test/ops/test_conv1d.py index 5fcdcf50465..3a1713a9548 100644 --- a/backends/arm/test/ops/test_conv1d.py +++ b/backends/arm/test/ops/test_conv1d.py @@ -118,6 +118,11 @@ def forward(self, x): return x +class Conv1dWithLeakyReLU(Conv1d): + def forward(self, x): + return torch.nn.functional.leaky_relu(super().forward(x), negative_slope=0.0) + + conv1d_2_3x2x40_nobias = Conv1d( in_channels=2, out_channels=3, @@ -327,6 +332,18 @@ def test_convolution_1d_tosa_INT(test_data): pipeline.run() +def test_convolution_1d_with_leaky_relu_tosa_INT(): + model = Conv1dWithLeakyReLU() + pipeline = TosaPipelineINT[input_t]( + model, + model.get_inputs(), + [aten_op, "torch.ops.aten.leaky_relu.default"], + [exir_op, "executorch_exir_dialects_edge__ops_aten_leaky_relu_default"], + qtol=1, + ) + pipeline.run() + + @common.parametrize("test_data", test_data_INT) @common.XfailIfNoCorstone300 def test_convolution_1d_u55_INT(test_data):