深入理解TVM：内存分配器

一、从make_object开始

TVM的内存分配器用来给前文Object家族中介绍的Object系列类分配内存和构造对象，同时提供了make_object这个helper function来完成这个工作，make_object在TVM中用的非常多，随便举个例子来说，比如src/te/tensor.cc中Tensor类的一个构造函数中就通过make_object来构造了一个TensorNode的对象：

Tensor::Tensor(Array<PrimExpr> shape, DataType dtype, Operation op,
               int value_index) {
  auto n = make_object<TensorNode>();
  n->shape = std::move(shape);
  n->dtype = dtype;
  n->op = op;
  n->value_index = value_index;
  data_ = std::move(n);
}

Tensor类继承自DataProducer，DataProducer继承自ObjectRef，上面Tensor构造函数中make_object构造的对象在完成相关初始化之后，最终交由data_管理，如下面代码所示，data_定义在Tensor的顶层基类ObjectRef中，可以看作是一个指向Object对象的指针，Object、ObjectPtr、ObjectRef三者的关系已经在前文Object家族中讲过了：

class ObjectRef {
 protected:
  ObjectPtr<Object> data_;
};

前面是一个make_object这个helper function使用的具体例子，下面来看make_object具体的定义，在include/tvm/runtime/memory.h中：

template <typename T, typename... Args>
inline ObjectPtr<T> make_object(Args&&... args) {
  return SimpleObjAllocator().make_object<T>(std::forward<Args>(args)...);
}

可以看到，make_object内部调用了SimpleObjAllocator的相关函数，SimpleObjAllocator就是本文要说的内存分配器。

make_object只是构造单个的Object，其实除了make_object之外，还有另外一个helper function make_inplace_array_object用于构造Object数组，同样是使用了SimpleObjAllocator这个内存分配器，下面是它的定义：

template <typename ArrayType, typename ElemType, typename... Args>
inline ObjectPtr<ArrayType> make_inplace_array_object(size_t num_elems, Args&&... args) {
  return SimpleObjAllocator().make_inplace_array<ArrayType, ElemType>(
                num_elems, std::forward<Args>(args)...);
}

下面详细分析下SimpleObjAllocator这个内存分配器。

二、ObjAllocatorBase

SimpleObjAllocator继承自ObjAllocatorBase，两者的类关系定义如下：

template <typename Derived> class ObjAllocatorBase {};
class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {};

这里用到了C++的一种编程技巧CRTP，它既可以实现静态多态，又可以复用代码，CRTP在TVM有多处应用，其它的如AttrRegistry也用到了CRTP，这里不详细说了，以后抽空单独写篇文章来详细分析CRTP。

先看基类ObjAllocatorBase，它有两个成员函数，一个是构造单个Object的make_object，一个是构造Object数组的make_inplace_array：

template <typename Derived> 
class ObjAllocatorBase {
  template <typename T, typename... Args>
  ObjectPtr<T> make_object(Args&&... args) {
    using Handler = typename Derived::template Handler<T>;
    T* ptr = Handler::New(static_cast<Derived*>(this), std::forward<Args>(args)...);
    ptr->type_index_ = T::RuntimeTypeIndex();
    ptr->deleter_ = Handler::Deleter();
    return ObjectPtr<T>(ptr);
  }

  template <typename ArrayType, typename ElemType, typename... Args>
  ObjectPtr<ArrayType> make_inplace_array(size_t num_elems, Args&&... args) {
    using Handler = typename Derived::template ArrayHandler<ArrayType, ElemType>;
    ArrayType* ptr = Handler::New(static_cast<Derived*>(this), num_elems, std::forward<Args>(args)...);
    ptr->type_index_ = ArrayType::RuntimeTypeIndex();
    ptr->deleter_ = Handler::Deleter();
    return ObjectPtr<ArrayType>(ptr);
  }
};

从上面的定义中可以看出，make_object和make_inplace_array的处理流程相同，都是通过静态多态的方法调用了Derived这个子类定义的New函数来构造对象，同时把Derived这个子类中定义的Deleter这个删除器赋值给Object中定义的deleter_变量中，用于析构对象的时候用，下面是deleter_在Object中的定义：

class TVM_DLL Object {
 public:
  typedef void (*FDeleter)(Object* self);
 protected:
  FDeleter deleter_ = nullptr;
};

三、SimpleObjAllocator

SimpleObjAllocator中主要实现了基类中调用的New和Deleter函数，先看使用make_object构造单个Object时调用的New的实现：

class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
  template <typename T> class Handler {
    using StorageType = typename std::aligned_storage<sizeof(T), alignof(T)>::type;
    template <typename... Args>
    static T* New(SimpleObjAllocator*, Args&&... args) {
      StorageType* data = new StorageType();
      new (data) T(std::forward<Args>(args)...);
      return reinterpret_cast<T*>(data);
    }
  };
};

看到这里，所有的一切就真相大白了，make_object调用的New通过标准库的new来分配空间，然后再通过placement new来在分配的空间上构造对象，这里可能大家会有个疑问，就是为什么要用placement new来单独构造对象，TVM code base里有下面这段解释，我没做实验验证这个解释，大家有做过实验的可以留言告知一下：

// NOTE2: Use placement new to allocate
// This is used to get rid of warning when deleting a virtual
// class with non-virtual destructor.

再看make_object中给Object的deleter_变量赋值时调用的子类的Deleter函数的定义：

class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
  template <typename T> class Handler {
    static Object::FDeleter Deleter() { return Deleter_; }
    static void Deleter_(Object* objptr) {
      T* tptr = static_cast<T*>(objptr);
      tptr->T::~T();
      delete reinterpret_cast<StorageType*>(tptr);
    }
  };
};

上面代码可以看到，删除器中先调用了析构函数，然后使用标准库的delete来释放内存。关于new/delete和inplacement new/delete，前文《内存管理：new and delete》有详细介绍，大家感兴趣可以点进去看看。

上面是make_object使用的New和Deleter，make_inplace_array所使用的New和Deleter的实现原理相同，只不过在分配和释放内存的时候，调用的是标准库的operator new/delete的数组版本，下面是make_inplace_array所使用的New和Deleter的关键代码：

class SimpleObjAllocator : public ObjAllocatorBase<SimpleObjAllocator> {
  template <typename ArrayType, typename ElemType> class ArrayHandler {
    using StorageType = typename std::aligned_storage<sizeof(ArrayType), alignof(ArrayType)>::type;

    template <typename... Args>
    static ArrayType* New(SimpleObjAllocator*, size_t num_elems, Args&&... args) {
      size_t unit = sizeof(StorageType);
      size_t requested_size = num_elems * sizeof(ElemType) + sizeof(ArrayType);
      size_t num_storage_slots = (requested_size + unit - 1) / unit;
      StorageType* data = new StorageType[num_storage_slots];
      new (data) ArrayType(std::forward<Args>(args)...);
      return reinterpret_cast<ArrayType*>(data);
    }

    static Object::FDeleter Deleter() { return Deleter_; }
    static void Deleter_(Object* objptr) {
      ArrayType* tptr = static_cast<ArrayType*>(objptr);
      tptr->ArrayType::~ArrayType();
      StorageType* p = reinterpret_cast<StorageType*>(tptr);
      delete[] p;
    }
  };
};

make_inplace_array所使用的New中有一点需要注意，它所申请的空间构成是下面这样子的：

图1

从图1可以看出，make_inplace_array申请的内存开头是Array作为header，后面紧跟着n个elements的空间，申请的时候把sizeof(array)+n*sizeof(element)+sizeof(padding)的空间作为num_storage_slots*sizeof(array)来申请，然后使用placement new来构造array header。

对于Deleter_函数，也有一点尤其需要注意，里面只调用了数组类型对象的析构函数，并没有调用里面所放元素的析构函数，需要user自己来手动在自己实现的Array中来调用存放元素的析构函数。Deleter_函数最后释放了数组对象和所有元素所占的内存空间。

四、summary and refenrence

TVM的内存分配器的实现很简单，但是它是整个TVM目前为Object分配对象的基石，非常重要值得一说，总结起来主要有下面几点：

• 使用基类+继承类，可以非常方便的扩展新的内存分配器，目前只实现了一个SimpleObjAllocator

• 使用CRTP实现静态多态，理论上来讲提高了内存分配的效率，因为如果使用动态多态的话，显然虚函数的调用开销要大于正常函数
  

• 既支持构造单个Object，也支持构造Object数组

本文参考的资料有下面这些：

• TVM code base：https://github.com/apache/tvm/