Update XNN delegate to handle non-decomposed upsample ops

backends/xnnpack/_passes/convert_to_upsample_bilinear2d.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import torch
 from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass
 from executorch.backends.xnnpack.partition.graphs import bilinear_2d
@@ -23,6 +25,10 @@ def create_upsample_bilinear_2d(
         align_corners: bool,
     ):
         output = internal_match.returning_nodes[0]
+        if output.target == exir_ops.edge.aten.upsample_bilinear2d.vec:
+            # Op was not decomposed, do nothing
+            return
+
         output_shape = output.meta["val"].shape
         output_h = output_shape[-2]
         output_w = output_shape[-1]

backends/xnnpack/test/ops/test_bilinear2d.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import unittest
 
 import torch
@@ -131,11 +133,17 @@ def test_fp32_bilinear2d_dynamic_bilinear2d_not_partitioned(self):
                 3: torch.export.Dim("w", min=1, max=12),
             }
         }
-        (
+        artifact_str = str(
             Tester(self.StaticResizeBilinear2dModule(), example_inputs)
             .export(Export(dynamic_shapes))
             .to_edge_transform_and_lower()
-            # NOTE The decomposition is partially delegated. This will need to be replaced
-            # with the aten upsample op once decomp is removed.
-            .check("executorch_exir_dialects_edge__ops_aten_index_Tensor")
+            .get_artifact()
+            .exported_program()
+        )
+        # NOTE The decomposition can be partially delegated. This will need to be replaced
+        # with the aten upsample op once decomp is removed.
+        self.assertTrue(
+            "executorch_exir_dialects_edge__ops_aten_index_Tensor" in artifact_str
+            or "executorch_exir_dialects_edge__ops_aten_upsample_bilinear2d_vec"
+            in artifact_str
         )

exir/emit/test/test_emit.py

@@ -642,7 +642,7 @@ class M(torch.nn.Module):
             def forward(self, x):
                 return torch.nn.functional.interpolate(x, scale_factor=2)
 
-        x = (torch.randn(1, 1, 2, 2),)
+        x = (torch.randn(1, 1, 2, 2, 2),)
         program = (
             to_edge(export(M(), x, strict=True)).to_executorch().executorch_program
         )