Move MemoryFormat/Layout to headeronly

ghstack-source-id: 9acfa60
Pull Request resolved: #168034

Author: janeyx99

c10/core/Layout.h

@@ -3,30 +3,9 @@
 #include <c10/core/Backend.h>
 #include <c10/util/Exception.h>
 
-#include <cstdint>
-#include <ostream>
+#include <torch/headeronly/core/Layout.h>
 
 namespace c10 {
-enum class Layout : int8_t {
-  Strided,
-  Sparse,
-  SparseCsr,
-  Mkldnn,
-  SparseCsc,
-  SparseBsr,
-  SparseBsc,
-  Jagged,
-  NumOptions
-};
-
-constexpr auto kStrided = Layout::Strided;
-constexpr auto kSparse = Layout::Sparse;
-constexpr auto kSparseCsr = Layout::SparseCsr;
-constexpr auto kMkldnn = Layout::Mkldnn;
-constexpr auto kSparseCsc = Layout::SparseCsc;
-constexpr auto kSparseBsr = Layout::SparseBsr;
-constexpr auto kSparseBsc = Layout::SparseBsc;
-constexpr auto kJagged = Layout::Jagged;
 
 inline Layout layout_from_backend(Backend backend) {
   C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wswitch-enum")
@@ -56,28 +35,4 @@ inline Layout layout_from_backend(Backend backend) {
   C10_DIAGNOSTIC_POP()
 }
 
-inline std::ostream& operator<<(std::ostream& stream, at::Layout layout) {
-  switch (layout) {
-    case at::kStrided:
-      return stream << "Strided";
-    case at::kSparse:
-      return stream << "Sparse";
-    case at::kSparseCsr:
-      return stream << "SparseCsr";
-    case at::kSparseCsc:
-      return stream << "SparseCsc";
-    case at::kSparseBsr:
-      return stream << "SparseBsr";
-    case at::kSparseBsc:
-      return stream << "SparseBsc";
-    case at::kMkldnn:
-      return stream << "Mkldnn";
-    case at::kJagged:
-      return stream << "Jagged";
-    case Layout::NumOptions:
-    default:
-      TORCH_CHECK(false, "Unknown layout");
-  }
-}
-
 } // namespace c10

c10/core/MemoryFormat.h

@@ -3,64 +3,18 @@
 #include <c10/util/ArrayRef.h>
 #include <c10/util/Exception.h>
 
+#include <torch/headeronly/core/MemoryFormat.h>
+
 #include <cstdint>
-#include <ostream>
 #include <vector>
 
-// Memory format is not the property of a Tensor. It is the way to tell an
-// operator how the result should be organized in memory and nothing more. That
-// means memory format should never be used as return value for any tensor state
-// interrogation functions (internally and externally).
-//
-// Possible options are:
-//  Preserve:
-//    If any of the input tensors is in channels_last format, operator output
-//    should be in channels_last format
-//
-//  Contiguous:
-//    Regardless of input tensors format, the output should be contiguous
-//    Tensor.
-//
-//  ChannelsLast:
-//    Regardless of input tensors format, the output should be in channels_last
-//    format.
-
 namespace c10 {
-enum class MemoryFormat : int8_t {
-  Contiguous,
-  Preserve,
-  ChannelsLast,
-  ChannelsLast3d,
-  NumOptions
-};
 
 // If you are seeing this, it means that this call site was not checked if
 // the memory format could be preserved, and it was switched to old default
 // behaviour of contiguous
 #define LEGACY_CONTIGUOUS_MEMORY_FORMAT c10::get_contiguous_memory_format()
 
-inline MemoryFormat get_contiguous_memory_format() {
-  return MemoryFormat::Contiguous;
-}
-
-inline std::ostream& operator<<(
-    std::ostream& stream,
-    at::MemoryFormat memory_format) {
-  switch (memory_format) {
-    case MemoryFormat::Preserve:
-      return stream << "Preserve";
-    case MemoryFormat::Contiguous:
-      return stream << "Contiguous";
-    case MemoryFormat::ChannelsLast:
-      return stream << "ChannelsLast";
-    case MemoryFormat::ChannelsLast3d:
-      return stream << "ChannelsLast3d";
-    case MemoryFormat::NumOptions:
-    default:
-      TORCH_CHECK(false, "Unknown memory format ", memory_format);
-  }
-}
-
 // Note: Hardcoded the channel last stride indices here to get better
 // performance
 template <typename T>

test/cpp/aoti_abi_check/CMakeLists.txt

@@ -16,8 +16,10 @@ set(AOTI_ABI_CHECK_TEST_SRCS
   ${AOTI_ABI_CHECK_TEST_ROOT}/test_dtype.cpp
   ${AOTI_ABI_CHECK_TEST_ROOT}/test_exception.cpp
   ${AOTI_ABI_CHECK_TEST_ROOT}/test_headeronlyarrayref.cpp
+  ${AOTI_ABI_CHECK_TEST_ROOT}/test_layout.cpp
   ${AOTI_ABI_CHECK_TEST_ROOT}/test_macros.cpp
   ${AOTI_ABI_CHECK_TEST_ROOT}/test_math.cpp
+  ${AOTI_ABI_CHECK_TEST_ROOT}/test_memoryformat.cpp
   ${AOTI_ABI_CHECK_TEST_ROOT}/test_metaprogramming.cpp
   ${AOTI_ABI_CHECK_TEST_ROOT}/test_rand.cpp
   ${AOTI_ABI_CHECK_TEST_ROOT}/test_scalartype.cpp

test/cpp/aoti_abi_check/test_layout.cpp

@@ -0,0 +1,65 @@
+#include <gtest/gtest.h>
+
+#include <torch/headeronly/core/Layout.h>
+
+TEST(TestLayout, TestLayout) {
+  using torch::headeronly::Layout;
+  constexpr Layout expected_layouts[] = {
+      torch::headeronly::kStrided,
+      torch::headeronly::kSparse,
+      torch::headeronly::kSparseCsr,
+      torch::headeronly::kMkldnn,
+      torch::headeronly::kSparseCsc,
+      torch::headeronly::kSparseBsr,
+      torch::headeronly::kSparseBsc,
+      torch::headeronly::kJagged,
+  };
+  for (int8_t i = 0; i < static_cast<int8_t>(Layout::NumOptions); i++) {
+    EXPECT_EQ(static_cast<Layout>(i), expected_layouts[i]);
+  }
+}
+
+TEST(TestLayout, operator_left_shift) {
+  using torch::headeronly::Layout;
+
+  {
+    std::stringstream ss;
+    ss << Layout::Strided;
+    EXPECT_EQ(ss.str(), "Strided");
+  }
+  {
+    std::stringstream ss;
+    ss << Layout::Sparse;
+    EXPECT_EQ(ss.str(), "Sparse");
+  }
+  {
+    std::stringstream ss;
+    ss << Layout::SparseCsr;
+    EXPECT_EQ(ss.str(), "SparseCsr");
+  }
+  {
+    std::stringstream ss;
+    ss << Layout::SparseCsc;
+    EXPECT_EQ(ss.str(), "SparseCsc");
+  }
+  {
+    std::stringstream ss;
+    ss << Layout::SparseBsr;
+    EXPECT_EQ(ss.str(), "SparseBsr");
+  }
+  {
+    std::stringstream ss;
+    ss << Layout::SparseBsc;
+    EXPECT_EQ(ss.str(), "SparseBsc");
+  }
+  {
+    std::stringstream ss;
+    ss << Layout::Mkldnn;
+    EXPECT_EQ(ss.str(), "Mkldnn");
+  }
+  {
+    std::stringstream ss;
+    ss << Layout::Jagged;
+    EXPECT_EQ(ss.str(), "Jagged");
+  }
+}

test/cpp/aoti_abi_check/test_memoryformat.cpp

@@ -0,0 +1,48 @@
+#include <gtest/gtest.h>
+
+#include <torch/headeronly/core/MemoryFormat.h>
+
+TEST(TestMemoryFormat, TestMemoryFormat) {
+  using torch::headeronly::MemoryFormat;
+  constexpr MemoryFormat expected_memory_formats[] = {
+      MemoryFormat::Contiguous,
+      MemoryFormat::Preserve,
+      MemoryFormat::ChannelsLast,
+      MemoryFormat::ChannelsLast3d,
+  };
+  for (int8_t i = 0; i < static_cast<int8_t>(MemoryFormat::NumOptions); i++) {
+    EXPECT_EQ(static_cast<MemoryFormat>(i), expected_memory_formats[i]);
+  }
+}
+
+TEST(TestMemoryFormat, get_contiguous_memory_format) {
+  using torch::headeronly::get_contiguous_memory_format;
+  using torch::headeronly::MemoryFormat;
+
+  EXPECT_EQ(get_contiguous_memory_format(), MemoryFormat::Contiguous);
+}
+
+TEST(TestMemoryFormat, operator_left_shift) {
+  using torch::headeronly::MemoryFormat;
+
+  {
+    std::stringstream ss;
+    ss << MemoryFormat::Preserve;
+    EXPECT_EQ(ss.str(), "Preserve");
+  }
+  {
+    std::stringstream ss;
+    ss << MemoryFormat::Contiguous;
+    EXPECT_EQ(ss.str(), "Contiguous");
+  }
+  {
+    std::stringstream ss;
+    ss << MemoryFormat::ChannelsLast;
+    EXPECT_EQ(ss.str(), "ChannelsLast");
+  }
+  {
+    std::stringstream ss;
+    ss << MemoryFormat::ChannelsLast3d;
+    EXPECT_EQ(ss.str(), "ChannelsLast3d");
+  }
+}

test/cpp_extensions/libtorch_agnostic_2_10_extension/libtorch_agnostic_2_10/csrc/my_empty.cpp

@@ -10,14 +10,16 @@ using torch::stable::Tensor;
 Tensor my_empty(
     torch::headeronly::HeaderOnlyArrayRef<int64_t> size,
     std::optional<torch::headeronly::ScalarType> dtype,
+    std::optional<torch::headeronly::Layout> layout,
     std::optional<torch::stable::Device> device,
-    std::optional<bool> pin_memory) {
-  return empty(size, dtype, device, pin_memory);
+    std::optional<bool> pin_memory,
+    std::optional<torch::headeronly::MemoryFormat> memory_format) {
+  return empty(size, dtype, layout, device, pin_memory, memory_format);
 }
 
 STABLE_TORCH_LIBRARY_FRAGMENT(libtorch_agnostic_2_10, m) {
   m.def(
-      "my_empty(int[] size, ScalarType? dtype=None, Device? device=None, bool? pin_memory=None) -> Tensor");
+      "my_empty(int[] size, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor");
 }
 
 STABLE_TORCH_LIBRARY_IMPL(libtorch_agnostic_2_10, CompositeExplicitAutograd, m) {

test/cpp_extensions/libtorch_agnostic_2_10_extension/libtorch_agnostic_2_10/ops.py

@@ -156,20 +156,24 @@ def test_get_num_threads() -> int:
     return torch.ops.libtorch_agnostic_2_10.test_get_num_threads.default()
 
 
-def my_empty(size, dtype=None, device=None, pin_memory=None) -> Tensor:
+def my_empty(
+    size, dtype=None, layout=None, device=None, pin_memory=None, memory_format=None
+) -> Tensor:
     """
-    Creates an empty tensor with the specified size, dtype, device, and pin_memory.
+    Creates an empty tensor with the specified size, dtype, layout, device, pin_memory, and memory_format.
 
     Args:
         size: list[int] - size of the tensor to create
         dtype: ScalarType or None - data type of the tensor
+        layout: Layout or None - layout of the tensor
         device: Device or None - device on which to create the tensor
         pin_memory: bool or None - whether to use pinned memory
+        memory_format: MemoryFormat or None - memory format of the tensor
 
     Returns: Tensor - an uninitialized tensor with the specified properties
     """
     return torch.ops.libtorch_agnostic_2_10.my_empty.default(
-        size, dtype, device, pin_memory
+        size, dtype, layout, device, pin_memory, memory_format
     )