pytorch基于intrusive_ptr_target实现的核心数据结构介绍

最新推荐文章于 2024-08-09 10:25:03 发布

来自底层的仰望

最新推荐文章于 2024-08-09 10:25:03 发布

阅读量356

点赞数 3

分类专栏： pytorch 文章标签： c++ pytorch 深度学习人工智能

本文链接：https://blog.csdn.net/weixin_43969854/article/details/129348171

版权

pytorch 专栏收录该内容

2 篇文章 0 订阅

订阅专栏

文章详细介绍了C++中的一种智能指针intrusive_ptr的实现，包括其内部的引用计数机制、友元函数、构造函数等。同时，它在PyTorch中的具体应用，如TensorImpl和StorageImpl类的使用，展示了如何通过intrusive_ptr管理对象生命周期。此外，还提到了MaybeOwnedTraits特化类和weak_intrusive_ptr类的功能和注意事项。

摘要由CSDN通过智能技术生成

1 指针类实现

1 intrusive_ptr_target 类

1）该类主要构建了两个计数器，且支持拷贝和移动构造。但是对于拷贝和移动构造没有执行任何动作，不改变原有计数器的数值。

2）声明了两个友元类和友元函数，其中友元函数不是类函数，还是特殊作用域下的函数。

3）当其他类继承这个类时，将自动维护两个计数器，而不需要在intrusive ptr中维护计数器。

从目前来看，intrusive_ptr想对于shared_ptr而言，其对应的引用计数的方式被修改了，并且要求对于T是继承于intrusive_ptr_target.

主要私有变量：
  mutable std::atomic<size_t> refcount_;
  mutable std::atomic<size_t> weakcount_;
友元：
  template <typename T, typename NullType>
  friend class intrusive_ptr;
  friend inline void raw::intrusive_ptr::incref(intrusive_ptr_target* self);

  template <typename T, typename NullType>
  friend class weak_intrusive_ptr;
  friend inline void raw::weak_intrusive_ptr::incref(
      intrusive_ptr_target* self);

2 intrusive_ptr类

1）私有函数reset_()看，只有当类中维护的指针对象的refcount_和weakcount_都减为0时，才会出发delete对象动作。(在该类的处理中，只会操作refcount_, 不会操作weakcount_)

2）建议通过make_intrusive()，reclaim() and weak_intrusive_ptr::lock()方式创建intrusive_ptr。

2.1）make_intrusive()，调用私有构造函数构造intrusive_ptr，增加refcount_和weakcount_计数。

2.2）reclaim()不增加计数器，需要特别注意。（如果需要特别考虑，可以使用unsafe_reclaim_from_nonowning）

2.3）weak_intrusive_ptr::lock()

3）release()操作不会减少refcount_，所以该操作只能把获取的指针重新放回intrusive_ptr，才是安全的。

4）unsafe_steal_from_new可以调用私有构造函数构造intrusive_ptr，增加refcount_和weakcount_计数。（如果使用裸指针对象创建intrusive_ptr，debug模式会检查refcount_和weakcount_必须要为0）

5）对应的operator比较等操作，都是通过包裹对应T定义的opertor方法进行处理。

6）私有函数retain_()只会操作refcount_。

特别说明1：

  intrusive_ptr& operator=(intrusive_ptr&& rhs) & noexcept {
    return operator=<TTarget, NullType>(std::move(rhs)); // 调用模板函数operator=，其通过编译器static_assert函数判断，from和ttarget是可转换的类型，从而此处直接用目标类型，而不是rhs的类型检查获取。
  }

  template <class From, class FromNullType>
  intrusive_ptr& operator=(intrusive_ptr<From, FromNullType>&& rhs) & noexcept {
    static_assert(
        std::is_convertible<From*, TTarget*>::value,
        "Type mismatch. intrusive_ptr move assignment got pointer of wrong type.");
    intrusive_ptr tmp = std::move(rhs);
    swap(tmp);
    return *this;
  }

特别说明2：

  template <class TTarget,class NullType = detail::intrusive_target_default_null_type<TTarget>>
  class intrusive_ptr final {  
    ...
    template <class TTarget2, class NullType2>  // 可以访问特化实现对象的私有成员变量
    friend class intrusive_ptr;
    ...
  }

3 MaybeOwnedTraits特化类MaybeOwnedTraits<c10::intrusive_ptr<T>>

逻辑上使用该类，需要非常谨慎。需要显示地销毁对应接口获得的intrusive_ptr。

1）createBorrow和assignBorrow不改变引用计数，调用reclaim进行操作。

2）destroyBorrow不改变引用基数，只是将自身置为nullptr，调用release()操作。

3）源文件路径： c10/util/MaybeOwned.h

4 weak_intrusive_ptr类

1）私有函数retain_和reset_只会操作weakcount_计数器。

2）通过intrusive_ptr构造weak_intrusive_ptr对象时，只改变weakcount_计数器。其他自身的拷贝和移动构造，常规操作。

3）release()和reclaim()函数都不改变weakcount_计数器。

代码路经：

c10/util/intrusive_ptr.h

c10/util/intrusive_ptr.cpp

2 使用案例

c10::intrusive_ptr 的初始化需要 intrusive_ptr_target 或者其子类。
TensorImpl 和 StorageImpl 两个类分别为intrusive_ptr_target 的子类，然后StorageImpl 主要负责 tensor 的实际物理内存相关的操作，设置空间配置器，获取数据指针，以及占用物理空间大小等； Storage 仅仅是对 StorageImpl 直接包了一下，直接调用的是 StorageImpl 的相关成员函数。TensorImpl 是 Tensor 类实现的主要依赖类，，其初始化就需要依赖 Storage 类，所以上面说：Tensor = TensorImpl + StorgaeImpl。

上述文字及图片截图链接：[Pytorch 源码阅读] —— Tensor C++相关实现_c++ tensor_Chris_zhangrx的博客-CSDN博客

1 StorageImpl类

代码路经： c10/core/StorageImpl.h

1）其继承于intrusive_ptr_target，其构造的时候，默认调用intrusive_ptr_target的默认构造，增加其refcount_和weakcount_的引用计数。

2）其中主要的私有变量为： DataPtr， SymInt， Allocator*， resizable_。

2.1）DataPtr结构体主要为一个智能指针 + 设备信息;

2.1.1）智能指针为UniqueVoidPtr，其主要包含一个ptr和unique_ptr(ptr + deleter function); ----为什么这么设计？

2.1.2）设备信息就是设备类型和设备index。

2.2）SymInt就是一个long数据类型的包装。

2.3）resizable_ 可否被resize？

3）支持移动构造，不支持拷贝构造。不支持默认构造，默认的稀构。

4）存在接口可以直接通过裸指针构造，UniqueStorageShareExternalPointer。

2）Storage类

代码路经：c10/core/Storage.h

1）protected成员变量为c10::intrusive_ptr<StorageImpl> storage_impl_。

2）本质就是对storage_impl的再次封装。但由于其包裹的是storage_impl的智能指针，其可以做更多的事情，比如拷贝构造，空构造，拷贝赋值等处理。

3）TensorImpl类

代码路经： c10/core/TensorImpl.h

1）不支持移动和复制构造，也不支持移动和赋值操作;

2）只支持storage的右值构造，而storage主要是一个intrusive_ptr包裹的stroage_impl; 或者dispatch_key的一个构造函数。其他均不支持。

3）新增了police的控制策略，用户可以自定义size，stride等metadata的行为。

辅助类介绍：

1）SizesAndStrides：用于存储tensor的size和stride信息。当size<=5时，存储在栈空间inlineStorage_，当超过5时，存储在堆outOfLineStorage_上。（大小分别都是2×size_）

2）VariableVersion：主要是一个intrusive_ptr包裹的atomic计数器。其实用来记录saved variable是否被修改。当进行inplace操作或者set_data或者.data的操作后，其计数器会增长。

3）AutogradMetaInterface：父类定义，后续AutogradMeta对其进行真实的实现。这么处理的原因是AutogradMeta定义在单独的编译单元，TensorImpl无法对其进行实例化。故通过父类接口完成对于子类实现的调用。

4）AutogradMetaFactory：父类定义，后续ConcreteAutogradMetaFactory对其进行真实的实现。其中ConcreteAutogradMetaFactory通过注册类AutogradMetaFactoryRegisterer，完成ConcreteAutogradMetaFactory对象到AutogradMetaFactory类中注册，后续调用接口都是基于ConcreteAutogradMetaFactory的对象完成。调用impl::GetAutogradMetaFactory()→make()获取AutogradMeta的对象指针。

具体关系如下：