Introduce dim order utils

Summary:
Common utils to work accross torch.memory_format and dim_order.

This is quite restrictive as of now, we can open it up for more usecases as we go.

Reviewed By: larryliu0820

Differential Revision: D47580149

fbshipit-source-id: de6017958a79b334e03dcce8a368cbde965078d6

Author: digantdesai

exir/TARGETS

@@ -255,3 +255,11 @@ python_library(
         "//caffe2:torch",
     ],
 )
+
+python_library(
+    name = "dim_order_utils",
+    srcs = ["dim_order_utils.py"],
+    deps = [
+        "//caffe2:torch",
+    ],
+)

exir/dim_order_utils.py

@@ -0,0 +1,57 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import List
+
+import torch
+
+"""
+Set of simple utilities for translating between torch.memory_format and dim_order
+"""
+
+
+def _get_contiguous_dim_order(ndim: int) -> List[int]:
+    if ndim <= 0:
+        raise AssertionError(f"Unsupported rank: {ndim}")
+
+    return list(range(ndim))
+
+
+def _get_channels_last_dim_order(ndim: int) -> List[int]:
+    if ndim == 4:
+        return [0, 2, 3, 1]
+
+    raise AssertionError(f"Unsupported rank: {ndim}")
+
+
+def get_memory_format(dim_order: List[int]) -> torch.memory_format:
+    """
+    Given a dim_order try to map it to torch.memory_format
+    """
+    if dim_order == _get_contiguous_dim_order(len(dim_order)):
+        return torch.contiguous_format
+    elif len(dim_order) == 4 and dim_order == _get_channels_last_dim_order(
+        len(dim_order)
+    ):
+        return torch.channels_last
+
+    raise AssertionError(
+        f"Failed to map a given dim_order: {dim_order} to a torch.memory_format"
+    )
+
+
+def get_dim_order(memory_format: torch.memory_format, ndim: int) -> List[int]:
+    """
+    Given a memory_format and a tensor rank, generate a dim_order
+    """
+    if memory_format == torch.contiguous_format:
+        return _get_contiguous_dim_order(ndim)
+    elif memory_format == torch.channels_last:
+        return _get_channels_last_dim_order(ndim)
+
+    raise AssertionError(
+        f"Failed to generate dim_order for a given memory format: {memory_format}"
+    )

exir/tests/TARGETS

@@ -420,3 +420,15 @@ python_unittest(
         "//executorch/exir/verification:verifier",
     ],
 )
+
+python_unittest(
+    name = "dim_order_utils",
+    srcs = [
+        "test_dim_order_utils.py",
+    ],
+    supports_static_listing = True,
+    deps = [
+        "//caffe2:torch",
+        "//executorch/exir:dim_order_utils",
+    ],
+)

exir/tests/test_dim_order_utils.py

@@ -0,0 +1,29 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+import unittest
+
+import torch
+from executorch.exir.dim_order_utils import get_dim_order, get_memory_format
+
+
+class TestDimOrderUtils(unittest.TestCase):
+    def test_get_memory_format(self) -> None:
+        mem_format = torch.contiguous_format
+        for ndim in range(1, 7):
+            dim_order = list(range(ndim))
+            self.assertEqual(mem_format, get_memory_format(dim_order))
+
+        mem_format = torch.channels_last
+        self.assertEqual(mem_format, get_memory_format([0, 2, 3, 1]))
+
+    def test_get_dim_order(self) -> None:
+        for ndim in range(1, 7):
+            self.assertEqual(
+                list(range(ndim)), get_dim_order(torch.contiguous_format, ndim)
+            )
+        self.assertEqual([0, 2, 3, 1], get_dim_order(torch.channels_last, 4))