Convenience APIs for script objects (#20226)

suo · facebook-github-bot · commit ce033485ebb4 · 2019-05-10T13:03:58.000-07:00
Summary: Pull Request resolved: #20226 ghimport-source-id: 22937d72e35ec4eba38019284a368453089fe3eb Differential Revision: D15243625 Pulled By: suo fbshipit-source-id: 5e9fb773da244f9ef201dba524155c3b19b2b4e0
diff --git a/aten/src/ATen/core/ivalue.cpp b/aten/src/ATen/core/ivalue.cpp
@@ -113,6 +113,16 @@ std::string ivalue::Object::name() const {
   return this->type_->qualname();
 }
 
+IValue ivalue::Object::getAttr(const std::string& name) const {
+  const size_t slot = type_->getAttributeSlot(name);
+  return getSlot(slot);
+}
+
+void ivalue::Object::setAttr(const std::string& name, IValue v) {
+  const size_t slot = type_->getAttributeSlot(name);
+  setSlot(slot, std::move(v));
+}
+
 void ivalue::Object::resizeObject(size_t slot) {
   AT_ASSERT(slot < type()->numAttributes());
   slots_.resize(type()->numAttributes());
diff --git a/aten/src/ATen/core/ivalue.h b/aten/src/ATen/core/ivalue.h
@@ -12,6 +12,13 @@
 
 #include <ATen/core/Tensor.h>
 
+namespace torch {
+namespace jit {
+namespace script {
+struct Function;
+}
+} // namespace jit
+} // namespace torch
 namespace c10 {
 struct IValue;
 struct ClassType;
@@ -695,6 +702,14 @@ struct C10_EXPORT ivalue::Object final : c10::intrusive_ptr_target {
     return c10::make_intrusive<Object>(std::move(type), numSlots);
   }
 
+  /**
+   * Slot API.
+   *
+   * Attributes are stored as a simple vector so that lookups are fast at
+   * runtime. A "slot" is just an index into that vector, which can be computed
+   * statically if you have access to the class type. Use this API if you are
+   * writing compiler stuff.
+   */
   void setSlot(size_t slot, IValue v) {
     if (slot >= slots_.size()) {
       // for module types, it is possible that the members of the class have
@@ -709,6 +724,19 @@ struct C10_EXPORT ivalue::Object final : c10::intrusive_ptr_target {
     return slots_.at(slot);
   }
 
+  /**
+   * Attribute API.
+   *
+   * Wrappers around the slot stuff so that users can access attributes
+   * directly. Use this API if you are a user.
+   *
+   * Note: Unlike in Python, TorchScript must make a distinction between
+   * attributes (which are IValues) and methods (which are Methods). If you
+   * want a method, use `obj.type()->getMethod()`
+   */
+  IValue getAttr(const std::string& name) const;
+  void setAttr(const std::string& name, IValue v);
+
   std::string name() const;
 
   const std::vector<IValue>& slots() const {
diff --git a/aten/src/ATen/core/type.cpp b/aten/src/ATen/core/type.cpp
@@ -156,6 +156,8 @@ TypePtr incompleteInferTypeFrom(const IValue& value) {
     return TupleType::create(fmap(value.toTuple()->elements(), incompleteInferTypeFrom));
   } else if (value.isDevice()) {
     return DeviceObjType::get();
+  } else if (value.isObject()) {
+    return value.toObject()->type();
   }
   AT_ERROR("Type cannot be accurately recovered from this IValue.");
 }
diff --git a/test/cpp/jit/test.cpp b/test/cpp/jit/test.cpp
@@ -80,7 +80,8 @@ namespace jit {
   _(SubgraphMatching)              \
   _(ModuleDefine)                  \
   _(QualifiedName)                 \
-  _(ClassImport)
+  _(ClassImport)                   \
+  _(ScriptObject)
 
 #define TH_FORALL_TESTS_CUDA(_) \
   _(ArgumentSpec)               \
diff --git a/test/cpp/jit/test_class_import.h b/test/cpp/jit/test_class_import.h
@@ -1,8 +1,11 @@
 #pragma once
 
-#include <ATen/core/qualified_name.h>
 #include <test/cpp/jit/test_base.h>
+#include <test/cpp/jit/test_utils.h>
+
+#include <ATen/core/qualified_name.h>
 #include <torch/csrc/jit/import_source.h>
+#include <torch/torch.h>
 
 namespace torch {
 namespace jit {
@@ -13,10 +16,12 @@ op_version_set = 1
 class FooNestedTest:
     def __init__(self, y):
         self.y = y
+
 class FooNestedTest2:
     def __init__(self, y):
         self.y = y
         self.nested = __torch__.FooNestedTest(y)
+
 class FooTest:
     def __init__(self, x):
         self.class_attr = __torch__.FooNestedTest(x)
@@ -58,6 +63,37 @@ void testClassImport() {
   ASSERT_FALSE(c);
 }
 
+void testScriptObject() {
+  Module m1;
+  Module m2;
+  std::vector<at::Tensor> constantTable;
+  import_libs(
+      m1.class_compilation_unit(),
+      "__torch__",
+      classSrcs1,
+      constantTable,
+      nullptr);
+  import_libs(
+      m2.class_compilation_unit(),
+      "__torch__",
+      classSrcs2,
+      constantTable,
+      nullptr);
+
+  // Incorrect arguments for constructor should throw
+  c10::QualifiedName base("__torch__");
+  ASSERT_ANY_THROW(m1.create_class(c10::QualifiedName(base, "FooTest"), {1}));
+  auto x = torch::ones({2, 3});
+  auto obj = m2.create_class(c10::QualifiedName(base, "FooTest"), x).toObject();
+  auto dx = obj->getAttr("dx");
+  ASSERT_TRUE(test::almostEqual(x, dx.toTensor()));
+
+  auto new_x = torch::rand({2, 3});
+  obj->setAttr("dx", new_x);
+  auto new_dx = obj->getAttr("dx");
+  ASSERT_TRUE(test::almostEqual(new_x, new_dx.toTensor()));
+}
+
 } // namespace script
 } // namespace jit
 } // namespace torch
diff --git a/torch/csrc/jit/script/module.cpp b/torch/csrc/jit/script/module.cpp
@@ -1,12 +1,12 @@
-#include <torch/csrc/jit/script/module.h>
 #include <c10/util/Exception.h>
+#include <torch/csrc/autograd/generated/variable_factories.h>
 #include <torch/csrc/jit/export.h>
 #include <torch/csrc/jit/operator.h>
 #include <torch/csrc/jit/passes/dead_code_elimination.h>
 #include <torch/csrc/jit/script/compiler.h>
 #include <torch/csrc/jit/script/error_report.h>
+#include <torch/csrc/jit/script/module.h>
 #include <torch/csrc/jit/script/schema_matching.h>
-#include <torch/csrc/autograd/generated/variable_factories.h>
 
 namespace torch {
 namespace jit {
@@ -106,15 +106,15 @@ void module_state_to(
     const c10::optional<at::Device>& device,
     const c10::optional<at::ScalarType>& dtype,
     bool non_blocking) {
-    // Need to access the `at::Tensor` as a `Variable` here.
-    autograd::Variable variable = s.value().toTensor();
-    at::Tensor data = variable.data();
-    // Use the data's original device or dtype if not supplied here.
-    auto new_data = data.to(
-        device.value_or(data.device()),
-        dtype.value_or(data.scalar_type()),
-        non_blocking);
-    variable.set_data(new_data);
+  // Need to access the `at::Tensor` as a `Variable` here.
+  autograd::Variable variable = s.value().toTensor();
+  at::Tensor data = variable.data();
+  // Use the data's original device or dtype if not supplied here.
+  auto new_data = data.to(
+      device.value_or(data.device()),
+      dtype.value_or(data.scalar_type()),
+      non_blocking);
+  variable.set_data(new_data);
 }
 
 void Module::to_impl(
@@ -301,7 +301,6 @@ void Module::copy_into(
   }
 }
 
-
 void Module::clone_method(
     const Module& orig,
     const std::string& name,
@@ -348,10 +347,42 @@ void Module::clone_method(const Module& orig, const std::string& name) {
 }
 
 void Module::train(bool on) {
- for (auto& submod : get_modules()) {
-   submod->train(on);
- }
- register_buffer("training", torch::tensor(on ? 1 : 0, at::kLong));
+  for (auto& submod : get_modules()) {
+    submod->train(on);
+  }
+  register_buffer("training", torch::tensor(on ? 1 : 0, at::kLong));
+}
+
+IValue Module::create_class(const c10::QualifiedName& name, Stack stack) const {
+  // Classes live in the top-level compilation unit.
+  if (parent_) {
+    return parent_->create_class(name, std::move(stack));
+  }
+
+  // Look up the class
+  const auto classType =
+      class_compilation_unit().get_class(c10::QualifiedName(name));
+  if (!classType) {
+    AT_ERROR(
+        "Could not find class with name: '",
+        name.qualifiedName(),
+        "' in module.");
+  }
+
+  // Create a bare object with correct number of slots
+  const size_t numAttrs = classType->numAttributes();
+  auto obj = c10::ivalue::Object::create(classType, numAttrs);
+
+  // Invoke the `__init__()` of the class with the arguments provided.
+  Stack stackWithSelf = {obj};
+  for (auto& arg : stack) {
+    stackWithSelf.push_back(std::move(arg));
+  }
+  // Note: following Python, `__init__()` modifies its first parameter in-place
+  // and returns nothing.
+  classType->getMethod("__init__")->operator()(std::move(stackWithSelf));
+
+  return obj;
 }
 
 } // namespace script
diff --git a/torch/csrc/jit/script/module.h b/torch/csrc/jit/script/module.h
@@ -442,6 +442,14 @@ struct TORCH_API Module {
   // so that C++ users can easily add methods
   void define(const std::string& src, const ResolverPtr& resolver = nullptr);
 
+  template <typename... Types>
+  IValue create_class(const c10::QualifiedName& name, Types&&... args)
+      const {
+    return create_class(name, {IValue(std::forward<Types>(args))...});
+  }
+
+  IValue create_class(const c10::QualifiedName& name, Stack stack) const;
+
  private:
   std::pair<std::shared_ptr<Function>, std::vector<Slot>>
   lower_first_class_method(Function* fn);

Original file line number	Diff line number	Diff line change
`@@ -156,6 +156,8 @@ TypePtr incompleteInferTypeFrom(const IValue& value) {`
`156`	`156`	`return TupleType::create(fmap(value.toTuple()->elements(), incompleteInferTypeFrom));`
`157`	`157`	`} else if (value.isDevice()) {`
`158`	`158`	`return DeviceObjType::get();`
	`159`	`+ } else if (value.isObject()) {`
	`160`	`+ return value.toObject()->type();`
`159`	`161`	`}`
`160`	`162`	`AT_ERROR("Type cannot be accurately recovered from this IValue.");`
`161`	`163`	`}`