[nativert] Move TensorMeta to pytorch core (#152475)

Summary:
Torch Native Runtime RFC: https://github.com/pytorch/rfcs/pull/72

This diff moves `TensorMeta.cpp` and `TensorMeta.h` to PyTorch core under `torch/nativert/graph/`

Existing `torch::_export::TensorMeta` in `torch/csrc/utils/generated_serialization_types.h` is auto-generated from the export serde schema and therefore only containing the most basic serializable types. We need the newly added `TensorMeta.cpp` to deserialize the metadata into a in-memory class with c10 types so that it can be consumed by the runtime later.

Test Plan:

Added test under `test/cpp/nativert/test_tensor_meta.cpp`

Differential Revision: D73820548

Pull Request resolved: https://github.com/pytorch/pytorch/pull/152475
Approved by: https://github.com/albanD

build_variables.bzl

@@ -587,6 +587,11 @@ jit_sources_full = [
 
 libtorch_core_jit_sources = sorted(jit_sources_full)
 
+
+libtorch_nativert_sources = [
+    "torch/nativert/graph/TensorMeta.cpp",
+]
+
 torch_mobile_tracer_sources = [
     "torch/csrc/jit/mobile/model_tracer/tracer.cpp",
     "torch/csrc/jit/mobile/model_tracer/TensorUtils.cpp",
@@ -619,7 +624,7 @@ libtorch_lite_cmake_sources = sorted(
     torch_mobile_core,
 )
 
-libtorch_cmake_sources = libtorch_core_sources + libtorch_core_jit_sources
+libtorch_cmake_sources = libtorch_core_sources + libtorch_core_jit_sources + libtorch_nativert_sources
 
 libtorch_extra_sources = libtorch_core_jit_sources + [
     "torch/csrc/autograd/TraceTypeManual.cpp",
@@ -659,7 +664,7 @@ libtorch_extra_sources = libtorch_core_jit_sources + [
 
 def libtorch_sources(gencode_pattern = ":generate-code[{}]"):
     return (
-        libtorch_generated_sources(gencode_pattern) + libtorch_core_sources + libtorch_distributed_sources + libtorch_extra_sources
+        libtorch_generated_sources(gencode_pattern) + libtorch_core_sources + libtorch_distributed_sources + libtorch_extra_sources + libtorch_nativert_sources
     )
 
 libtorch_cuda_core_sources = [

caffe2/CMakeLists.txt

@@ -1319,6 +1319,7 @@ if(BUILD_TEST)
     )
   else()
     add_subdirectory(${TORCH_ROOT}/test/cpp/jit ${CMAKE_BINARY_DIR}/test_jit)
+    add_subdirectory(${TORCH_ROOT}/test/cpp/nativert ${CMAKE_BINARY_DIR}/test_nativert)
     add_subdirectory(${TORCH_ROOT}/test/inductor ${CMAKE_BINARY_DIR}/test_inductor)
     add_subdirectory(
       ${TORCH_ROOT}/test/cpp/tensorexpr

test/cpp/nativert/CMakeLists.txt

@@ -0,0 +1,40 @@
+set(NATIVERT_TEST_ROOT ${TORCH_ROOT}/test/cpp/nativert)
+
+# Build the cpp gtest binary containing the cpp-only tests.
+set(NATIVERT_TEST_SRCS
+  ${NATIVERT_TEST_ROOT}/test_tensor_meta.cpp
+  ${TORCH_ROOT}/torch/nativert/graph/TensorMeta.cpp
+)
+
+add_executable(test_nativert
+  ${TORCH_ROOT}/test/cpp/common/main.cpp
+  ${NATIVERT_TEST_SRCS}
+)
+
+# TODO temporary until we can delete the old gtest polyfills.
+target_compile_definitions(test_nativert PRIVATE USE_GTEST)
+
+set(NATIVERT_TEST_DEPENDENCIES torch gtest)
+
+target_link_libraries(test_nativert PRIVATE ${NATIVERT_TEST_DEPENDENCIES})
+target_include_directories(test_nativert PRIVATE ${ATen_CPU_INCLUDE})
+
+if(USE_CUDA)
+  target_compile_definitions(test_nativert PRIVATE USE_CUDA)
+elseif(USE_ROCM)
+  target_link_libraries(test_nativert PRIVATE
+    hiprtc::hiprtc
+    hip::amdhip64
+    ${TORCH_CUDA_LIBRARIES})
+
+  target_compile_definitions(test_nativert PRIVATE USE_ROCM)
+endif()
+
+if(INSTALL_TEST)
+  set_target_properties(test_nativert PROPERTIES INSTALL_RPATH "${CMAKE_INSTALL_RPATH}:${_rpath_portable_origin}/../lib")
+  install(TARGETS test_nativert DESTINATION bin)
+  # Install PDB files for MSVC builds
+  if(MSVC AND BUILD_SHARED_LIBS)
+    install(FILES $<TARGET_PDB_FILE:test_nativert> DESTINATION bin OPTIONAL)
+  endif()
+endif()

test/cpp/nativert/test_tensor_meta.cpp

@@ -0,0 +1,62 @@
+#include <gtest/gtest.h>
+#include <torch/nativert/graph/TensorMeta.h>
+
+namespace torch::nativert {
+TEST(TensorMetaTest, ScalarTypeConversion) {
+  EXPECT_EQ(
+      convertJsonScalarType(torch::_export::ScalarType::FLOAT),
+      c10::ScalarType::Float);
+  EXPECT_EQ(
+      convertJsonScalarType(torch::_export::ScalarType::INT),
+      c10::ScalarType::Int);
+  EXPECT_EQ(
+      convertJsonScalarType(torch::_export::ScalarType::HALF),
+      c10::ScalarType::Half);
+  EXPECT_EQ(
+      convertJsonScalarType(torch::_export::ScalarType::COMPLEXHALF),
+      c10::ScalarType::ComplexHalf);
+  EXPECT_EQ(
+      convertJsonScalarType(torch::_export::ScalarType::BFLOAT16),
+      c10::ScalarType::BFloat16);
+  EXPECT_THROW(
+      convertJsonScalarType(static_cast<torch::_export::ScalarType>(100)),
+      c10::Error);
+}
+TEST(TensorMetaTest, MemoryFormatConversion) {
+  EXPECT_EQ(
+      convertJsonMemoryFormat(torch::_export::MemoryFormat::ContiguousFormat),
+      c10::MemoryFormat::Contiguous);
+  EXPECT_EQ(
+      convertJsonMemoryFormat(torch::_export::MemoryFormat::ChannelsLast),
+      c10::MemoryFormat::ChannelsLast);
+  EXPECT_EQ(
+      convertJsonMemoryFormat(torch::_export::MemoryFormat::PreserveFormat),
+      c10::MemoryFormat::Preserve);
+  EXPECT_THROW(
+      convertJsonMemoryFormat(static_cast<torch::_export::MemoryFormat>(100)),
+      c10::Error);
+}
+
+TEST(TensorMetaTest, LayoutConversion) {
+  EXPECT_EQ(
+      convertJsonLayout(torch::_export::Layout::Strided), c10::Layout::Strided);
+  EXPECT_EQ(
+      convertJsonLayout(torch::_export::Layout::SparseCsr),
+      c10::Layout::SparseCsr);
+  EXPECT_EQ(
+      convertJsonLayout(torch::_export::Layout::_mkldnn), c10::Layout::Mkldnn);
+  EXPECT_THROW(
+      convertJsonLayout(static_cast<torch::_export::Layout>(100)), c10::Error);
+}
+TEST(TensorMetaTest, DeviceConversion) {
+  torch::_export::Device cpu_device;
+  cpu_device.set_type("cpu");
+  EXPECT_EQ(convertJsonDevice(cpu_device), c10::Device(c10::DeviceType::CPU));
+  torch::_export::Device cuda_device;
+  cuda_device.set_type("cuda");
+  cuda_device.set_index(0);
+  EXPECT_EQ(
+      convertJsonDevice(cuda_device), c10::Device(c10::DeviceType::CUDA, 0));
+}
+
+} // namespace torch::nativert

torch/nativert/graph/TensorMeta.cpp

@@ -0,0 +1,142 @@
+#include <torch/nativert/graph/TensorMeta.h>
+
+#include <c10/util/Logging.h>
+
+namespace torch::nativert {
+
+c10::ScalarType convertJsonScalarType(
+    const torch::_export::ScalarType& scalarType) {
+  switch (scalarType) {
+    case torch::_export::ScalarType::UNKNOWN:
+      TORCH_CHECK(false, "scalar type is not properly set");
+    case torch::_export::ScalarType::BYTE:
+      return c10::ScalarType::Byte;
+    case torch::_export::ScalarType::CHAR:
+      return c10::ScalarType::Char;
+    case torch::_export::ScalarType::SHORT:
+      return c10::ScalarType::Short;
+    case torch::_export::ScalarType::INT:
+      return c10::ScalarType::Int;
+    case torch::_export::ScalarType::LONG:
+      return c10::ScalarType::Long;
+    case torch::_export::ScalarType::HALF:
+      return c10::ScalarType::Half;
+    case torch::_export::ScalarType::FLOAT:
+      return c10::ScalarType::Float;
+    case torch::_export::ScalarType::DOUBLE:
+      return c10::ScalarType::Double;
+    case torch::_export::ScalarType::COMPLEXHALF:
+      return c10::ScalarType::ComplexHalf;
+    case torch::_export::ScalarType::COMPLEXFLOAT:
+      return c10::ScalarType::ComplexFloat;
+    case torch::_export::ScalarType::COMPLEXDOUBLE:
+      return c10::ScalarType::ComplexDouble;
+    case torch::_export::ScalarType::BOOL:
+      return c10::ScalarType::Bool;
+    case torch::_export::ScalarType::BFLOAT16:
+      return c10::ScalarType::BFloat16;
+    case torch::_export::ScalarType::UINT16:
+      return c10::ScalarType::UInt16;
+    case torch::_export::ScalarType::FLOAT8E4M3FN:
+      return c10::ScalarType::Float8_e4m3fn;
+    case torch::_export::ScalarType::FLOAT8E5M2:
+      return c10::ScalarType::Float8_e5m2;
+    default:
+      TORCH_CHECK(false, "unknown scalar type", static_cast<int>(scalarType));
+  }
+}
+
+c10::MemoryFormat convertJsonMemoryFormat(
+    const torch::_export::MemoryFormat& memoryFormat) {
+  switch (memoryFormat) {
+    case torch::_export::MemoryFormat::Unknown:
+      TORCH_CHECK(false, "got unknown scalar type");
+    case torch::_export::MemoryFormat::ContiguousFormat:
+      return c10::MemoryFormat::Contiguous;
+    case torch::_export::MemoryFormat::ChannelsLast:
+      return c10::MemoryFormat::ChannelsLast;
+    case torch::_export::MemoryFormat::ChannelsLast3d:
+      return c10::MemoryFormat::ChannelsLast3d;
+    case torch::_export::MemoryFormat::PreserveFormat:
+      return c10::MemoryFormat::Preserve;
+    default:
+      TORCH_CHECK(
+          false, "unknown memory format", static_cast<int>(memoryFormat));
+  }
+}
+
+c10::Layout convertJsonLayout(const torch::_export::Layout& layout) {
+  switch (layout) {
+    case torch::_export::Layout::Unknown:
+      TORCH_CHECK(false, "got unknown layout");
+    case torch::_export::Layout::SparseCoo:
+      // TODO is this the right translation
+      return c10::Layout::Sparse;
+    case torch::_export::Layout::SparseCsr:
+      return c10::Layout::SparseCsr;
+    case torch::_export::Layout::SparseCsc:
+      return c10::Layout::SparseCsc;
+    case torch::_export::Layout::SparseBsr:
+      return c10::Layout::SparseBsr;
+    case torch::_export::Layout::SparseBsc:
+      return c10::Layout::SparseBsc;
+    case torch::_export::Layout::_mkldnn:
+      return c10::Layout::Mkldnn;
+    case torch::_export::Layout::Strided:
+      return c10::Layout::Strided;
+    default:
+      TORCH_CHECK(false, "unknown layout", static_cast<int>(layout));
+  }
+}
+
+c10::Device convertJsonDevice(const torch::_export::Device& device) {
+  c10::Device d(device.get_type());
+  if (auto index = device.get_index()) {
+    d.set_index(*index);
+  }
+  return d;
+}
+
+TensorMeta::TensorMeta(const torch::_export::TensorMeta& tensorMeta)
+    : device_(convertJsonDevice(tensorMeta.get_device())) {
+  dtype_ = convertJsonScalarType(tensorMeta.get_dtype());
+  layout_ = convertJsonLayout(tensorMeta.get_layout());
+  requiresGrad_ = tensorMeta.get_requires_grad();
+
+  if (tensorMeta.get_storage_offset().tag() ==
+      torch::_export::SymInt::Tag::AS_INT) {
+    storage_offset_ = tensorMeta.get_storage_offset().get_as_int();
+  } else {
+    CHECK(false) << "SymInt not supported yet";
+  }
+
+  for (const auto& size : tensorMeta.get_sizes()) {
+    if (size.tag() == torch::_export::SymInt::Tag::AS_INT) {
+      int64_t val = size.get_as_int();
+      sizes_.emplace_back(val);
+      numel_ *= val;
+    } else if (size.tag() == torch::_export::SymInt::Tag::AS_EXPR) {
+      // TODO: it's still unclear how SymInt shape should be used in runtime
+      // One potential use cases is for verifing inputs shape matches constrain
+      // This would require unpacking the serialized constrain, which is NYI
+      //
+      // For the time being, we just set the symbolic dim to -1
+      hasSymbolicShape_ = true;
+      sizes_.emplace_back(-1);
+      numel_ = -1;
+    }
+  }
+
+  for (const auto& stride : tensorMeta.get_strides()) {
+    if (stride.tag() == torch::_export::SymInt::Tag::AS_INT) {
+      strides_.emplace_back(stride.get_as_int());
+    } else if (stride.tag() == torch::_export::SymInt::Tag::AS_EXPR) {
+      // TODO: it's still unclear how SymInt shape should be used in runtime
+      // Setting symbolic shape to -1 for now
+      hasSymbolicShape_ = true;
+      strides_.emplace_back(-1);
+    }
+  }
+}
+
+} // namespace torch::nativert

torch/nativert/graph/TensorMeta.h

@@ -0,0 +1,92 @@
+#pragma once
+
+#include <c10/core/Device.h>
+#include <c10/util/Logging.h>
+
+#include <c10/core/MemoryFormat.h>
+#include <c10/core/ScalarType.h>
+#include <c10/core/TensorOptions.h>
+#include "c10/core/Layout.h"
+#include <c10/util/ArrayRef.h>
+
+#include <torch/csrc/utils/generated_serialization_types.h>
+
+namespace torch::nativert {
+
+c10::ScalarType convertJsonScalarType(
+    const torch::_export::ScalarType& scalarType);
+c10::MemoryFormat convertJsonMemoryFormat(
+    const torch::_export::MemoryFormat& memoryFormat);
+c10::Layout convertJsonLayout(const torch::_export::Layout& layout);
+c10::Device convertJsonDevice(const torch::_export::Device& device);
+
+class TensorMeta {
+ public:
+  explicit TensorMeta(const torch::_export::TensorMeta& tensorMeta);
+
+  c10::IntArrayRef sizes() const {
+    CHECK(!hasSymbolicShape_) << "TensorMeta has symbolic shape";
+    return sizes_;
+  }
+
+  c10::IntArrayRef strides() const {
+    CHECK(!hasSymbolicShape_) << "TensorMeta has symbolic shape";
+    return strides_;
+  }
+
+  c10::Layout layout() const {
+    return layout_;
+  }
+
+  c10::ScalarType dtype() const {
+    return dtype_;
+  }
+
+  bool requires_grad() const {
+    return requiresGrad_;
+  }
+
+  int64_t storage_offset() const {
+    return storage_offset_;
+  }
+
+  int64_t dim() const {
+    return sizes_.size();
+  }
+
+  int64_t numel() const {
+    CHECK(!hasSymbolicShape_) << "TensorMeta has symbolic shape";
+    return numel_;
+  }
+
+  c10::Device device() const {
+    return device_;
+  }
+
+  c10::TensorOptions asTensorOptions() const {
+    return c10::TensorOptions().dtype(dtype_).layout(layout_).requires_grad(
+        requiresGrad_);
+  }
+
+  // NYI
+  // c10::SymIntArrayRef sym_sizes() const {}
+  // c10::SymIntArrayRef sym_strides() const {}
+  // c10::SymInt sym_storage_offset() const {}
+  // c10::SymInt sym_numel() const {}
+
+ private:
+  bool hasSymbolicShape_ = false;
+
+  std::vector<int64_t> sizes_;
+  std::vector<int64_t> strides_;
+  int64_t storage_offset_ = 0;
+  int64_t numel_ = 1;
+
+  c10::ScalarType dtype_;
+  c10::Layout layout_;
+  bool requiresGrad_;
+
+  c10::Device device_;
+};
+
+} // namespace torch::nativert