caffe2源码分析之core之存储篇----storage

core模块的数据存储与表示主要分为以下几个模块,后面逐一分析。
storage
tensor
blob
qtensor
上面模块的包含关系为blob大多数时候包含tensor指针, tensor包含storage.
这里开始介绍数据的存储与表示,storage:

struct C10_API StorageImpl final : public c10::intrusive_ptr_target {  //继承intrusive_ptr_target可以使用intrusive_ptr智能指针
private:
/*TypeMeta 类是一个数据类型的封装,可以作为blob或者tensor等的container,另外还包含了itme的size,type name等数据。*/
  caffe2::TypeMeta data_type_; 
  /*可以参考源码分析的base章节*/
  DataPtr data_ptr_;
  int64_t numel_;
  bool resizable_;
  // Identifies that Storage was received from another process and doesn't have
  // local to process cuda memory allocation
  bool received_cuda_;
  Allocator* allocator_;
};

下面分析全部的源码:

#pragma once

#include <c10/core/Allocator.h>
#include <c10/core/ScalarType.h>

#include <c10/util/intrusive_ptr.h>

namespace c10 {
/*继承自intrusive_ptr_target 是为了可以使用intrusive_ptr智能指针*/
struct C10_API StorageImpl final : public c10::intrusive_ptr_target {
 public:
 /*构造函数*/
  StorageImpl(
      caffe2::TypeMeta data_type,
      int64_t numel,    //item number,就是有几个data item。
      at::DataPtr data_ptr,
      at::Allocator* allocator,
      bool resizable)
      : data_type_(data_type),
        data_ptr_(std::move(data_ptr)), //这里因为data_ptr里面包含unique_ptr,所以需要move,不能copy。
        numel_(numel),
        resizable_(resizable),
        received_cuda_(false),
        allocator_(allocator) {
    if (resizable) {
    /*asseration, 也就是如果resizable为true,同时allocator为null的话,则返回error。所以从这里看出,
    如果需要resize的话,则需要利用allocator来分配内存*/
      AT_ASSERTM( 
          allocator_, "For resizable storage, allocator must be provided");
    }
    if (numel > 0) {
    /*如果numel大于0,说明包含meta data, 则data_type也必须有确定格式*/
      if (data_type_.id() == caffe2::TypeIdentifier::uninitialized()) {
        AT_ERROR(
            "Constructing a storage with meta of unknown type and non-zero numel");
      }
    }
  }

  StorageImpl(
      caffe2::TypeMeta data_type,
      int64_t numel,
      at::Allocator* allocator,
      bool resizable)
      : StorageImpl(   //借用了上面的构造函数。
            data_type,
            numel,
            allocator->allocate(data_type.itemsize() * numel),  //用传入的alloctor来分配data_ptr.
            allocator,
            resizable) {}
/*从下面这些运算符重载可以看出,只能接受右值引用,不接受左值,因为data_ptr包含unique_ptr,所以不能copy,只能move*/
  StorageImpl& operator=(StorageImpl&& other) = default;
  StorageImpl& operator=(const StorageImpl&) = delete;
  StorageImpl() = delete;
  StorageImpl(StorageImpl&& other) = default;
  StorageImpl(const StorageImpl&) = delete;
  ~StorageImpl() = default;

  void reset() {  //释放data_ptr, 将item num 设置为0
    data_ptr_.clear();
    numel_ = 0;
  }
/*这是个模板函数,返回bool来表明是不是该type。*/
  template <typename T>
  inline bool IsType() const {
    return data_type_.Match<T>();
  }

  template <typename T>
  inline T* data() const {
    auto data_type = caffe2::TypeMeta::Make<T>();
    if (dtype() != data_type) {
      AT_ERROR(
          "Attempt to access StorageImpl having data type ",
          dtype(),
          " as data type ",
          data_type);
    }
    return unsafe_data<T>();
  }

  template <typename T>
  inline T* unsafe_data() const {
    return static_cast<T*>(this->data_ptr_.get());
  }

  void release_resources() override {
    data_ptr_.clear();
  }

  size_t itemsize() const {
    return data_type_.itemsize();
  }

  size_t capacity() const {
    return numel_ * itemsize();
  }

  int64_t numel() const {
    return numel_;
  };

  // TODO: remove later
  void set_numel(int64_t numel) {
    numel_ = numel;
  };

  bool resizable() const {
    return resizable_;
  };

  at::DataPtr& data_ptr() {
    return data_ptr_;
  };

  const at::DataPtr& data_ptr() const {
    return data_ptr_;
  };

  // Returns the previous data_ptr 
  at::DataPtr set_data_ptr(at::DataPtr&& data_ptr) {
    std::swap(data_ptr_, data_ptr);
    return std::move(data_ptr);
  };
/*下面注释说明,此接口一定要慎重使用,因为一旦变化type,里面很多参数要同时变化。*/
  // XXX: TERRIBLE! DONT USE UNLESS YOU HAVE TO! AND EVEN THEN DONT, JUST DONT!
  // Setting the data_type will require you to audit many other parts of the
  // struct again to make sure it's still valid.
  void set_dtype(const caffe2::TypeMeta& data_type) {
    int64_t capacity = numel_ * data_type_.itemsize();
    data_type_ = data_type;
    numel_ = capacity / data_type_.itemsize();
  }

  // TODO: Return const ptr eventually if possible
  void* data() {
    return data_ptr_.get();
  }

  void* data() const {
    return data_ptr_.get();
  }

  at::DeviceType device_type() const {
    return data_ptr_.device().type();
  }

  at::Allocator* allocator() {
    return allocator_;
  }

  const caffe2::TypeMeta& dtype() const {
    return data_type_;
  }

  const at::Allocator* allocator() const {
    return allocator_;
  };

  // You generally shouldn't use this method, but it is occasionally
  // useful if you want to override how a tensor will be reallocated,
  // after it was already allocated (and its initial allocator was
  // set)
  void set_allocator(at::Allocator* allocator) {
    allocator_ = allocator;
  }

  Device device() const {
    return data_ptr_.device();
  }

  void set_resizable(bool resizable) {
    if (resizable) {
      // We need an allocator to be resizable
      AT_ASSERT(allocator_);
    }
    resizable_ = resizable;
  }

  /**
   * Can only be called when use_count is 1
   */
  void UniqueStorageShareExternalPointer(
      void* src,
      const caffe2::TypeMeta& data_type,
      size_t capacity,
      DeleterFnPtr d = nullptr) {
    UniqueStorageShareExternalPointer(
        at::DataPtr(src, src, d, data_ptr_.device()), data_type, capacity);
  }

  /**
   * Can only be called when use_count is 1
   */
  void UniqueStorageShareExternalPointer(
      at::DataPtr&& data_ptr,
      const caffe2::TypeMeta& data_type,
      size_t capacity) {
    data_type_ = data_type;
    // TODO: Use CAFFE_ENFORCE_WITH_CALLER equivalent
    // For now causes lots of redefine issues if caffe2/core/logging.h is used
    if (data_type_.id() == caffe2::TypeIdentifier::uninitialized()) {
      AT_ERROR(
          "To share with a raw external pointer you need to have meta "
          "already set.");
    }
    data_ptr_ = std::move(data_ptr);
    // NOTE: data_type might change and so it's also possible that capacity
    // might not be divisible by itemsize. There is no way for us to keep track
    // of the exact capacity if we're not explicity storing is. More conrectely
    // capacity() might not return the value that was set here, if itemsize does
    // not evenly divide it.
    numel_ = capacity / data_type_.itemsize();
    allocator_ = nullptr;
    resizable_ = false;
  }

  // This method can be used only after storage construction and cannot be used
  // to modify storage status
  void set_received_cuda(bool received_cuda) {
    received_cuda_ = received_cuda;
  }

  bool received_cuda() {
    return received_cuda_;
  }

 private:
  caffe2::TypeMeta data_type_;
  DataPtr data_ptr_;
  int64_t numel_;
  bool resizable_;
  // Identifies that Storage was received from another process and doesn't have
  // local to process cuda memory allocation
  bool received_cuda_;
  Allocator* allocator_;
};
} // namespace c10

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值