[nativert] Move Weights to PyTorch core (pytorch#155156)

Summary:
Pull Request resolved: pytorch#155156

Moves Weights class to PyTorch core
Torch Native Runtime RFC: pytorch/rfcs#72
 README: https://github.com/pytorch/pytorch/blob/main/torch/nativert/OVERVIEW.md

Test Plan: buck2 run mode/dev-nosan caffe2/test/cpp/nativert:weights_test

Differential Revision: D75973156

Author: yiming0416

build_variables.bzl

@@ -596,6 +596,7 @@ libtorch_nativert_sources = [
     "torch/nativert/graph/TensorMeta.cpp",
     "torch/nativert/executor/Placement.cpp",
     "torch/nativert/executor/PlacementUtils.cpp",
+    "torch/nativert/executor/Weights.cpp",
     "torch/nativert/common/FileUtil.cpp",
 ]
 

test/cpp/nativert/CMakeLists.txt

@@ -9,6 +9,7 @@ set(NATIVERT_TEST_SRCS
   ${TORCH_ROOT}/torch/nativert/graph/Graph.cpp
   ${TORCH_ROOT}/torch/nativert/graph/GraphSignature.cpp
   ${TORCH_ROOT}/torch/nativert/executor/PlacementUtils.cpp
+  ${TORCH_ROOT}/torch/nativert/executor/Weights.cpp
   ${TORCH_ROOT}/torch/nativert/common/FileUtil.cpp
 )
 

test/cpp/nativert/test_weights.cpp

@@ -0,0 +1,92 @@
+#include <gtest/gtest.h>
+#include <torch/csrc/jit/serialization/pickle.h>
+#include <torch/custom_class.h>
+#include <torch/torch.h>
+#include <memory>
+
+#include <torch/nativert/executor/Placement.h>
+#include <torch/nativert/executor/Weights.h>
+#include <torch/nativert/graph/Graph.h>
+
+namespace torch::nativert {
+class WeightsTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    static constexpr std::string_view source =
+        R"(graph(%foo, %bar, %baz):
+%o1, %o2 = aten.foo(self=%foo, target=%bar, alpha=0.1)
+return(%o2, %baz)
+)";
+    graph = stringToGraph(source);
+    placement = std::make_unique<Placement>(c10::Device(c10::DeviceType::CPU));
+  }
+  std::shared_ptr<Graph> graph;
+  std::unique_ptr<Placement> placement;
+};
+TEST_F(WeightsTest, ConstructEmptyStateDict) {
+  std::unordered_map<std::string, c10::IValue> stateDict;
+  Weights weights(graph.get(), stateDict, *placement);
+  // Check that weights are initialized correctly
+  EXPECT_TRUE(weights.parameters().empty());
+  EXPECT_TRUE(weights.buffers().empty());
+  EXPECT_FALSE(weights.contains("non_existent_weight"));
+}
+TEST_F(WeightsTest, SetAndGetValue) {
+  std::unordered_map<std::string, c10::IValue> stateDict;
+  Weights weights(graph.get(), stateDict, *placement);
+  at::Tensor tensor = at::ones({2, 2});
+  weights.setValue("added_weight", tensor);
+  EXPECT_TRUE(weights.contains("added_weight"));
+  EXPECT_EQ(weights.at("added_weight").sizes(), tensor.sizes());
+}
+
+} // namespace torch::nativert
+
+using namespace ::testing;
+struct ContainsTensorDict : torch::CustomClassHolder {
+  explicit ContainsTensorDict(at::Tensor t) : t_(t) {}
+
+  explicit ContainsTensorDict(c10::Dict<std::string, at::Tensor> dict) {
+    t_ = dict.at(std::string("init_tensor"));
+  }
+
+  c10::Dict<std::string, at::Tensor> serialize() const {
+    c10::Dict<std::string, at::Tensor> dict;
+    dict.insert(std::string("init_tensor"), t_);
+    return dict;
+  }
+
+  at::Tensor t_;
+};
+
+static auto reg =
+    torch::class_<ContainsTensorDict>("testing", "ContainsTensorDict")
+        .def(torch::init<at::Tensor>())
+        .def_pickle(
+            // __getstate__
+            [](const c10::intrusive_ptr<ContainsTensorDict>& self)
+                -> c10::Dict<std::string, at::Tensor> {
+              return self->serialize();
+            },
+            // __setstate__
+            [](c10::Dict<std::string, at::Tensor> data)
+                -> c10::intrusive_ptr<ContainsTensorDict> {
+              return c10::make_intrusive<ContainsTensorDict>(std::move(data));
+            });
+
+TEST(CustomWeightsTest, TestCustomObjWithContainedTensor) {
+  // Save
+  auto customObj =
+      c10::make_intrusive<ContainsTensorDict>(torch::tensor({1, 2, 3}));
+  const auto bytes = torch::jit::pickle_save(c10::IValue(std::move(customObj)));
+
+  // Load
+  const auto loadedCustomObj =
+      torch::jit::pickle_load_obj(std::string{bytes.begin(), bytes.end()});
+  EXPECT_TRUE(loadedCustomObj.isObject());
+  EXPECT_EQ(
+      loadedCustomObj.to<c10::intrusive_ptr<ContainsTensorDict>>()
+          ->t_[0]
+          .item<int>(),
+      1);
+}