Defining object destructor in c10 (#21984)

zdevito · facebook-github-bot · commit f6aac413911c · 2019-06-19T16:27:40.000-07:00
Summary: Pull Request resolved: #21984 ghimport-source-id: 5767592e37ed388422eed5639f8ba0722aec66e2 Test Plan: Imported from OSS Differential Revision: D15906530 Pulled By: zdevito fbshipit-source-id: bec8c8b0b5b9dcc2e8fc69b5031fcfa6bb22d54e
diff --git a/aten/src/ATen/core/ivalue_inl.h b/aten/src/ATen/core/ivalue_inl.h
@@ -277,14 +277,19 @@ struct C10_EXPORT ivalue::Future final : c10::intrusive_ptr_target {
 // User-defined object.
 struct C10_EXPORT ivalue::Object final : c10::intrusive_ptr_target {
  public:
-  Object(std::shared_ptr<ClassType> type, size_t numSlots) : type_(std::move(type)) {
+  // temporary way to break cyclic dependencies in modules by forcing the deletion
+  // of functions when the module object is destructed
+  typedef void (*OnDelete)(ivalue::Object*);
+  Object(std::shared_ptr<ClassType> type, size_t numSlots, OnDelete on_delete)
+      : type_(std::move(type)), on_delete_(on_delete) {
     slots_.resize(numSlots);
   }
 
   static c10::intrusive_ptr<Object> create(
       std::shared_ptr<ClassType> type,
-      size_t numSlots) {
-    return c10::make_intrusive<Object>(std::move(type), numSlots);
+      size_t numSlots,
+      OnDelete on_delete = nullptr) {
+    return c10::make_intrusive<Object>(std::move(type), numSlots, on_delete);
   }
 
   /**
@@ -337,6 +342,7 @@ struct C10_EXPORT ivalue::Object final : c10::intrusive_ptr_target {
   void resizeObject(size_t slot);
   std::shared_ptr<ClassType> type_;
   std::vector<IValue> slots_;
+  OnDelete on_delete_;
 };
 
 std::vector<std::pair<IValue, IValue>> iterationOrder(const c10::DictPtr<IValue, IValue>& dict);
diff --git a/aten/src/ATen/core/type.cpp b/aten/src/ATen/core/type.cpp
@@ -5,11 +5,13 @@
 #include <c10/macros/Macros.h>
 namespace c10 {
 
-#ifdef C10_ANDROID
 namespace ivalue {
-Object::~Object() {}
+Object::~Object() {
+  if (on_delete_) {
+    on_delete_(this);
+  }
+}
 } // namespace ivalue
-#endif
 
 std::ostream& operator<<(std::ostream & out, const Type & t) {
   if(auto value = t.cast<CompleteTensorType>()) {
diff --git a/torch/csrc/jit/script/class_type.cpp b/torch/csrc/jit/script/class_type.cpp
@@ -7,16 +7,6 @@ namespace c10 {
 // This file exists because we need to reference module.h, which we can't from
 // c10. Sigh...
 
-#ifndef C10_ANDROID
-namespace ivalue {
-Object::~Object() {
-  if (type_->is_module()) {
-    type_->compilation_unit()->drop_all_functions();
-  }
-}
-} // namespace ivalue
-#endif
-
 std::shared_ptr<Function> ClassType::getMethod(const std::string& name) const {
   return compilation_unit_->find_function(name);
 }
diff --git a/torch/csrc/jit/script/module.h b/torch/csrc/jit/script/module.h
@@ -104,6 +104,9 @@ struct TORCH_API Method {
   // first-class function in class_compilation_unit()
   std::shared_ptr<Function> function_;
 };
+static void clearMethods(c10::ivalue::Object* self) {
+  self->type()->compilation_unit()->drop_all_functions();
+}
 
 struct TORCH_API Module {
   Module(std::string class_name)
@@ -112,7 +115,8 @@ struct TORCH_API Module {
                 QualifiedName(std::move(class_name)),
                 std::make_shared<CompilationUnit>(),
                 /*is_module=*/true),
-            0)) {}
+            0,
+            clearMethods)) {}
   Module() : Module("$Module") {}
   Module(ModulePtr module_value) : module_value_(std::move(module_value)) {}
   ~Module() {}
