Adding Support for preserve_format when converting Evalues (#2242)

Summary:

This diff adds support for `torch.preserve_format` when casting constants to Evalues.
---

Differential Revision: D54288165

Author: Jack-Khuu

exir/emit/_emitter.py

@@ -462,11 +462,11 @@ def _constant_to_evalue(  # noqa: C901
             return EValue(Int(layout_enum(val)))
 
         if isinstance(val, torch.memory_format):
-            if val != torch.contiguous_format:
+            if val != torch.contiguous_format and val != torch.preserve_format:
                 raise InternalError(
                     self._emit_node_specific_error(
                         self.node,
-                        "Non contiguous tensors are not supported in ExecuTorch",
+                        "Only Tensors that have a contiguous or preserve_format memory_format are supported in ExecuTorch",
                     )
                 )
             return EValue(Int(memory_format_enum(val)))

exir/emit/test/test_emit.py

@@ -20,6 +20,7 @@
 from executorch.exir.backend.backend_api import to_backend
 from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult
 from executorch.exir.emit import emit_program  # noqa
+from executorch.exir.error import InternalError
 from executorch.exir.passes.constant_prop_pass import constant_prop_pass
 from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass
 from executorch.exir.print_program import pretty_print, print_program  # noqa
@@ -836,6 +837,47 @@ def _compare_execution_plans(
             else:
                 self.assertEqual(single_val, merged_val)
 
+    def test_emit_memory_format(self) -> None:
+        class SimpleLinear(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                contiguous = x.to(
+                    dtype=torch.float32, memory_format=torch.contiguous_format
+                )
+                preserve = x.to(
+                    dtype=torch.float32, memory_format=torch.preserve_format
+                )
+                return contiguous + preserve
+
+        # Should succeed at exporting model with legal memory format (contiguous, preserve)
+        model = SimpleLinear()
+        inputs = (torch.ones(10, 5),)
+        try:
+            to_edge(
+                export(model, inputs),
+                compile_config=exir.EdgeCompileConfig(_check_ir_validity=False),
+            ).to_executorch()
+        except:
+            self.fail("Failed to export model with legal memory format")
+
+        class SimpleLinear2(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return x.to(dtype=torch.float32, memory_format=torch.channels_last)
+
+        #  Failure Expected when exporting model with illegal memory format (channels_last)
+        model = SimpleLinear2()
+        inputs = (torch.ones(10, 5, 2, 1),)
+        with self.assertRaises(InternalError):
+            to_edge(
+                export(model, inputs),
+                compile_config=exir.EdgeCompileConfig(_check_ir_validity=False),
+            ).to_executorch()
+
     def test_emit_multiple_entry_points(self) -> None:
         class SimpleLinear(torch.nn.Module):
             def __init__(self) -> None:

exir/tensor.py

@@ -231,6 +231,7 @@ def memory_format_enum(memory_format: torch.memory_format) -> int:
     )
     table = {
         torch.contiguous_format: 0,
+        torch.preserve_format: 1,
     }
     return table[memory_format]