Adding Support for preserve_format when converting Evalues (#2242)

Summary:
Pull Request resolved: #2242

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

Reviewed By: dbort

Differential Revision: D54288165

fbshipit-source-id: 3fc2772ecadb9d65e16d3b23ee7ae02fb7a9259c

Author: Jack-Khuu

exir/emit/_emitter.py

@@ -463,14 +463,15 @@ def _constant_to_evalue(  # noqa: C901
             return EValue(Int(layout_enum(val)))
 
         if isinstance(val, torch.memory_format):
-            if val != torch.contiguous_format:
+            try:
+                return EValue(Int(memory_format_enum(val)))
+            except KeyError:
                 raise InternalError(
                     self._emit_node_specific_error(
                         self.node,
-                        "Non contiguous tensors are not supported in ExecuTorch",
+                        f"Tensor has a memory_format that is unsupported in ExecuTorch: {val}",
                     )
                 )
-            return EValue(Int(memory_format_enum(val)))
 
         if isinstance(val, torch.Tensor):
             raise ExportError(

exir/emit/test/test_emit.py

@@ -21,6 +21,7 @@
 from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult
 from executorch.exir.dialects._ops import ops as exir_ops
 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
@@ -838,6 +839,48 @@ def _compare_execution_plans(
             else:
                 self.assertEqual(single_val, merged_val)
 
+    def test_emit_memory_format_valid(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")
+
+    def test_emit_memory_format_invalid(self) -> None:
+        class SimpleLinear(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 = SimpleLinear()
+        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]
 

runtime/core/portable_type/tensor_options.h

@@ -14,17 +14,21 @@ namespace torch {
 namespace executor {
 
 /**
- * Tensor data memory format. This concept only exists for compatibility
- * with ATen.
+ * Tensor data memory formats supported by ExecuTorch. This concept only exists
+ * for compatibility with ATen; use dim_order to describe non-contiguous
+ * layouts.
  */
 enum class MemoryFormat : int8_t {
   /**
-   * Row-major contiguous data format.
-   *
-   * This is the only format supported by ExecuTorch. Use dim orders to
-   * describe other layouts.
+   * Row-major contiguous data.
+   */
+  Contiguous = 0,
+  /**
+   * Output tensor format should remain the same as the input tensor format.
+   * E.g. if the input tensor is in channels_last format, operator output
+   * should be in channels_last format.
    */
-  Contiguous,
+  Preserve = 1,
 };
 
 /**
@@ -39,7 +43,7 @@ enum class Layout : int8_t {
    *
    * This is the only layout supported by ExecuTorch.
    */
-  Strided,
+  Strided = 0,
 };
 } // namespace executor
 } // namespace torch