[jit] Refcounting spot fixes (#65346)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/65346

Tidying up the top sources of reference count decrements seen during static runtime startup.
ghstack-source-id: 140027349

Test Plan:
CI

perf now shows under 2% time spend in ~__shared_count instead of about 5%.

Reviewed By: suo

Differential Revision: D31057277

fbshipit-source-id: 9a16daf2e655fda80d4ec21290b30f02ba63d8da

aten/src/ATen/core/jit_type.h

@@ -1960,7 +1960,7 @@ struct TORCH_API ClassType : public NamedType {
       std::vector<std::string> unresolved_class_attributes = {});
 
   bool operator==(const Type& rhs) const override {
-    if (auto user_rhs = rhs.cast<ClassType>()) {
+    if (auto user_rhs = rhs.castRaw<ClassType>()) {
       const auto& lhs_name = name().value();
       const auto& rhs_name = user_rhs->name().value();
 

c10/util/MaybeOwned.h

@@ -8,46 +8,50 @@
 
 namespace c10 {
 
-/// Traits class describing how to borrow from T.  As a synopsis, here
-/// is how we might implement borrowing from an arbitrary type T using
-/// a raw pointer to const:
-///
-/// template <typename T>
-/// struct MaybeOwnedTraits {
-///   using owned_type = T;
-///   using borrow_type = const T*;
-///
-///   static borrow_type createBorrow(const owned_type& from) {
-///     return &from;
-///   }
-///
-///   static void assignBorrow(borrow_type& lhs, borrow_type rhs) {
-///     lhs = rhs;
-///   }
-///
-///   static void destroyBorrow(borrow_type& toDestroy) {}
-///
-///   static const owned_type& referenceFromBorrow(const borrow_type& borrow) {
-///     return *borrow;
-///   }
-///
-///   static const owned_type* pointerFromBorrow(const borrow_type& borrow) {
-///     return borrow;
-///   }
-///
-///   static bool debugBorrowIsValid(const borrow_type& borrow) {
-///     return borrow != nullptr;
-///   }
-/// };
-///
-/// (This implementation is not in use because MaybeOwned is an unsafe
-/// abstraction and we don't want to encourage it widely, just for
-/// skipping reference counting in critical paths.)
-///
-/// For examples that are in use, see intrusive_ptr.h and TensorBody.h.
+/// MaybeOwnedTraits<T> describes how to borrow from T.  Here is how we
+/// can implement borrowing from an arbitrary type T using a raw
+/// pointer to const:
+template <typename T>
+struct MaybeOwnedTraitsGenericImpl {
+  using owned_type = T;
+  using borrow_type = const T*;
+
+  static borrow_type createBorrow(const owned_type& from) {
+    return &from;
+  }
+
+  static void assignBorrow(borrow_type& lhs, borrow_type rhs) {
+    lhs = rhs;
+  }
+
+  static void destroyBorrow(borrow_type& toDestroy) {}
+
+  static const owned_type& referenceFromBorrow(const borrow_type& borrow) {
+    return *borrow;
+  }
+
+  static const owned_type* pointerFromBorrow(const borrow_type& borrow) {
+    return borrow;
+  }
+
+  static bool debugBorrowIsValid(const borrow_type& borrow) {
+    return borrow != nullptr;
+  }
+};
+
+/// It is possible to eliminate the extra layer of indirection for
+/// borrows for some types that we control. For examples, see
+/// intrusive_ptr.h and TensorBody.h.
+
 template <typename T>
 struct MaybeOwnedTraits;
 
+// Explicitly enable MaybeOwned<shared_ptr<T>>, rather than allowing
+// MaybeOwned to be used for any type right away.
+template <typename T>
+struct MaybeOwnedTraits<std::shared_ptr<T>>
+    : public MaybeOwnedTraitsGenericImpl<std::shared_ptr<T>> {};
+
 /// A smart pointer around either a borrowed or owned T. When
 /// constructed with borrowed(), the caller MUST ensure that the
 /// borrowed-from argument outlives this MaybeOwned<T>. Compare to

torch/csrc/jit/frontend/tree_views.h

@@ -459,6 +459,9 @@ struct ClassDef : public TreeView {
   explicit ClassDef(const TreeRef& tree) : TreeView(tree) {
     tree->match(TK_CLASS_DEF);
   }
+  explicit ClassDef(TreeRef&& tree) : TreeView(std::move(tree)) {
+    tree_->match(TK_CLASS_DEF);
+  }
   ClassDef withName(std::string new_name) const {
     auto new_ident = Ident::create(name().range(), std::move(new_name));
     return create(range(), new_ident, superclass(), body());

torch/csrc/jit/ir/type_hashing.cpp

@@ -11,7 +11,7 @@ namespace jit {
 
 namespace {
 size_t hashType(const c10::ConstTypePtr& type) {
-  if (auto named_type = type->cast<ClassType>()) {
+  if (auto named_type = type->castRaw<ClassType>()) {
     return get_hash(named_type->name().value());
   }
   auto hashes = fmap(type->containedTypes(), [](const TypePtr& elem) {

torch/csrc/jit/serialization/import_source.cpp

@@ -133,7 +133,7 @@ TypePtr SourceImporterImpl::findNamedType(const QualifiedName& name) {
   parseSourceIfNeeded(name.prefix());
   auto it = to_be_defined_.find(name);
   if (it != to_be_defined_.end() && it->second->kind() == TK_CLASS_DEF) {
-    ClassDef cd(it->second);
+    ClassDef cd(std::move(it->second));
     to_be_defined_.erase(it);
     importNamedType(name.prefix(), cd);
   }