-
Notifications
You must be signed in to change notification settings - Fork 0
Majel: From Allocation to DArray
Allocation在Majel中是对一块内存空间的表示，其基本定义为：
class Allocation {
public:
Allocation();
Allocation(size_t size);
Allocation(size_t size, Place place);
// Creates a non-owned allocation (an allocation not owned by the Majel
// memory allocator); non-owned allocations are not cleaned up in the
// destructor.
Allocation(void* ptr, size_t size, Place place);
~Allocation();
//No copying!
Allocation(const Allocation&) = delete;
//No assigning!
Allocation& operator=(const Allocation&) = delete;
void* ptr() const;
void* end() const;
Place place() const;
size_t size() const;
private:
bool owned_;
void* ptr_;
size_t size_;
Place place_;
};Allocation本身也负责了内存的申请和释放工作，上面的定义中前三个构造函数都会去申请内存。成员变量owned_用来记录这块这块内存是否由Allocation本身申请，如果为true，则在Allocation析构的时候会销毁对应的内存。
Buffer在Majel中表示一块数据，是Array的基类，也是Array的完整定义的一部分。Buffer与Allocation可以是多对一的关系，即多个Buffer指向同一块Allocation。其定义为：
class Buffer {
public:
Buffer() : external_address_(nullptr), allocation_(std::make_shared<Allocation>(0)) { }
Buffer(void* address) : external_address_(address), allocation_(std::make_shared<Allocation>(0)) { }
Buffer(void* address, Place p) : external_address_(address), allocation_(std::make_shared<Allocation>(0, p)) { }
Buffer(std::shared_ptr<Allocation> allocation) :external_address_(nullptr), allocation_(allocation) {}
public:
void* get_address() const {
if (allocation_->ptr() == nullptr) {
return external_address_;
}
return allocation_->ptr();
}
Place get_place() const { return allocation_->place(); }
std::shared_ptr<Allocation> data() const {
return allocation_;
}
private:
void* external_address_;
std::shared_ptr<Allocation> allocation_;
};可以看到Buffer支持两种类型的内存表示，一是直接的裸指针void* external_address_，二是Allocation的shared_ptr指针。显然，在第一种表示中Buffer不会负责对应内存的管理，第二种会由allocation_负责管理。
在Buffer的构造过程中，如果传入的参数是裸指针，则会用external_address_表示内存，同时allocation_中的ptr_被置为nullptr。如果传入的参数是Allocation对象，则external_address_被置为nullptr。
换言之，external_address_和allocation_有且只会有一个有效。
Reference是一个类模板，表示对某一个数据的引用，该数据块可以在内存上，也可以在显存上，Reference提供了统一的访问方式。
template<typename T>
struct Reference {
PlacedPointer<T> ptr;
std::shared_ptr<Allocation> alloc;
};PlacedPointer<T>存在的目的是为内存和显存上的数据提供统一的访问方式。其内部封装了一个Place对象和一个T*指针，并提供了get()和put()函数：
template<typename T>
struct PlacedPointer {
Place place;
T* ptr;
PlacedPointer(Place p, T* pt) : place(p), ptr(pt) {}
T get() const {
return boost::apply_visitor(PlacedGetter<T>(ptr), place);
}
void put(const T& value) {
return boost::apply_visitor(PlacedPutter<T>(ptr, value),
place);
}
};PlacedGetter<T>和PlacedPutter<T>均继承自boost::static_visitor<T>，其内部根据Place选择不同的访问方式。如果是CpuPlace，则直接访问内存；如果是GpuPlace，则进行内存和显存之间的拷贝，并访问内存中的对应位置。
alloc表示上文的ptr所指向的数据所属的Allocation。关于它和ptr的关系的更多信息将在后面的Array中有所体现。
1、构造函数：
Reference(PlacedPointer<T> p,
std::shared_ptr<Allocation> a) : ptr(p),
alloc(a){}2、类型转换函数
分别提供了从Reference<T>类型向T类型和其他可以从T转化过去的类型之间的转换方法：
// This allows Reference<T> to convert to T
operator T() const {
return ptr.get();
}
// This allows Reference<T> to convert to any
// type that is convertible from T
template<typename U,
typename = typename std::enable_if<
std::is_convertible<T, U>::value>::type>
operator U() const {
return U(ptr.get());
}3、=和==的运算符重载
Reference<T>& operator=(const T& other) {
ptr.put(other);
return *this;
}
Reference<T>& operator=(const Reference<T>& other) {
T value = other;
*this = value;
return *this;
}
bool operator==(const T& other) const {
return ptr.get() == other;
}
bool operator==(const Reference<T>& other) const {
return ptr.get() == other.ptr.get();
}可以看到=和==的重载函数本质上都是在通过成员变量ptr进行操作。
定义了上面这些转换方法和算符重载后，Reference<T>对象就可以像C++中原生的引用那样使用，直接用=和T类型的变量进行相互赋值。
Array是一个类模板，继承自Buffer，用来表示一个固定dimension的array（或者说tensor）。
template<typename T, int D>
class Array : public Buffer {
Dim<D> size_;
Dim<D> stride_;
T* ptr_;
};Array主要提供了多种构造函数，当构造函数传入的参数中包含std::shared_ptr<Allocation>对象时，该对象会用来初始化作为基类的Buffer；如果不包含，那么会在构造函数中新申请一个Allocation，再用来初始化Buffer。
从Array的构造函数上来看，作为Array的基类的Buffer必须通过std::shared_ptr<Allocation> allocation_来管理内存。
成员变量Array::ptr_表示这个Array所表示的数据在内存上的开始位置。一般情况下与Buffer::allocation_::ptr_指向同一位置，但因为可能有多个Array基于同一个Buffer::allocation_，且其中的一些是另一些的切片，因此Array::ptr_也可能指向Buffer::allocation_所管理内存的中间某个位置。
注意在Array这个层面上才出现了stride的概念，因此可以认为在Array之下的所有概念（Buffer、Allocation）中，内存都是连续的。Array需要stride的概念是因为它可以进行切片操作（详见后面的全局函数部分）。
Array提供的成员函数中比较重要的有下面几个：
[]算符重载
majel::Reference<T> operator[](const Dim<D>& idx) {
scheduler::synchronize(*this);
T* location = index(idx);
return majel::Reference<T>(majel::PlacedPointer<T>(place(), location),
data());
}
T operator[](const Dim<D>& idx) const {
scheduler::synchronize(*this);
T* location = index(idx);
return majel::Reference<T>(majel::PlacedPointer<T>(place(), location),
data());
}返回array中某一下标位置上的元素引用或者值。返回过程分为三步：
- 根据下标计算出指针偏移量。
- 用偏移后的指针和
Place构建PalcedPointer，再用这个PlacedPointer和Buffer::allocation_构建元素的Reference。 - 直接返回这个
Refernece或是转换成值本身在返回。
raw_ptr()
T* raw_ptr() const {
return ptr_;
}直接返回成员变量ptr_。
1、元素取值、赋值
template<typename T, int D>
T get(const Array<T, D>& arr, const Dim<D>& idx) {
return arr[idx];
}
template<typename T, int D>
void set(Array<T, D>& arr, const Dim<D>& idx, const T& value) {
arr[idx] = value;
}如果数据在显存上，那么执行过程会涉及到内存和显存的互相拷贝，可能比较耗时。
2、改变Array的形状
改变Array的形状并不会改变元素的个数，也不涉及到内存的重新分配，因此要求改变前后各维度值相乘的结果一致。
template<typename T, int OldD, int NewD>
Array<T, NewD> reshape(const Array<T, OldD>& x, Dim<NewD> new_size) {
if (!contiguous(x.size(), x.stride())) {
throw std::invalid_argument("Sorry, reshaping non-contiguous Arrays is not currently implemented.");
}
if (x.numel() != product(new_size)) {
throw std::invalid_argument("Sorry, reshaping Arrays must preserve the number of elements.");
}
return Array<T, NewD>(x.data(), new_size, contiguous_strides(new_size), x.raw_ptr());
}
/// Flatten an array to one dimension
template<typename T, int D>
Array<T, 1> flatten(const Array<T, D>& x) {
return reshape(x, Dim<1>(x.numel()));
}函数返回了一个新的Array，这个新Array和老的Array共用同一段内存（Allocation）。
3、切片操作
从老的Array中截取出一块生成新的Array，并且支持各个维度上的间隔切片。
template<typename T, int D>
Array<T, D> slice(const Array<T, D>& x, DDim start, DDim stop, DDim step) {
Dim<D> size;
Dim<D> stride;
std::tie(size, stride) = detail::slice_math<D>(start,
stop,
step);
stride = stride * x.stride();
return Array<T, D>(x.data(),
size,
stride,
x.index(boost::get<Dim<D>>(start)));
}其中detail::slice_math函数用来检查参数合法性并生成新的size和stride。返回的新Array和老Array共用内存，但是ptr_可能不再指向这段内存的开头，而是中间的某个位置。
4、从std::vector构造一维Array
可以作为一个Array构造过程的例子来看。
template<typename T>
Array<T, 1>
make_array(const std::vector<T>& input, Place place) {
std::shared_ptr<majel::Allocation> alloc =
std::make_shared<majel::Allocation>(sizeof(T) * input.size(), place);
T *ptr = static_cast<T*>(alloc->ptr());
if (is_gpu_place(place)) {
gpu::detail::memcpy_sync(ptr, input.data(), sizeof(T) * input.size(),
cudaMemcpyHostToDevice);
} else if (is_cpu_place(place)) {
memcpy(ptr, input.data(), sizeof(T) * input.size());
}
return Array<T, 1>(alloc, make_dim(input.size()));
}如上所述，Array是一个模板类，有两个模板参数，分别是元素类型T和维度D，这意味着不同数据类型和维度的Array是完全不同的class。实际使用中，经常会遇到Array的维度无法在编译期间就确定的情况，这时就需要一个统一的、可以承接所有类型Array的接口，这就是DArray。
typedef boost::variant<
Array<float, 1>,
Array<float, 2>,
Array<float, 3>,
Array<float, 4>,
Array<double, 1>,
Array<double, 2>,
Array<double, 3>,
Array<double, 4>,
Array<float16, 1>,
Array<float16, 2>,
Array<float16, 3>,
Array<float16, 4> > DArrayVar;
}
struct DArray {
DArrayVar var;
DArray();
template<typename T, int D>
DArray(Array<T, D> in) : var(in) {}
template<typename T>
DArray& operator=(T in) {
var = in;
return *this;
}
const DValue operator[](const DDim&) const;
DReference operator[](const DDim&);
template<typename Visitor>
typename Visitor::result_type
apply_visitor(Visitor& visitor) {
return var.apply_visitor(visitor);
}
template<typename Visitor>
typename Visitor::result_type
apply_visitor(Visitor& visitor) const {
return var.apply_visitor(visitor);
}
};DArray的核心是一个boost::variant类型的成员变量var，var被设定为可以接受数据类型为float、double、float16、维度为1-4的Array。
DArray实现了名为apply_visitor的成员函数，用来将外部的boost::static_visitor传递作用到var上。这样一来，我们就可以在外部定义各种各样针对var的visitor，并在外部直接以boost::apply_visitor(visitor, DArray)的形式使用。换言之，经过这样的定义，虽然看上去传递进去的参数是DArray，但其实真正传给visitor的参数是DArray的var成员，即Array。（阅读boost源码可以发现，boost::apply_visitor(visitor, v)函数内部其实是就调用了v.apply_visitor(visitor)）。
DArray的成员函数除了apply_visitor()以外只有对=和[]的算符重载。代码为：
template<typename T>
DArray& operator=(T in) {
var = in;
return *this;
}
const DValue DArray::operator[](const DDim& idx) const {
return boost::apply_visitor(DevalueVisitor(),
var, idx);
}
DReference DArray::operator[](const DDim& idx) {
return boost::apply_visitor(DereferenceVisitor(),
var, idx);
}其中主要是用到了DevalueVisitor和DereferenceVisitor这两个可调用对象，它们的定义为：
struct DereferenceVisitor
: public boost::static_visitor<DReference> {
template<typename T, int D, int E>
DReference operator()(Array<T, D> a, const Dim<E>& d) const {
std::stringstream ss;
ss << "Index dimension does not match array in DReferenceVisitor, type is: " <<
a.get_type_and_dim() << " is a gpu place: " << is_gpu_place(a.place()) <<
" size: " << a.size() << " and dim is " << d.dimensions << "D with contents: " <<
d;
throw std::invalid_argument(ss.str());
}
template<typename T, int D>
DReference operator()(Array<T, D> a, const Dim<D>& d) const {
return a[d];
}
};
struct DevalueVisitor
: public boost::static_visitor<DValue> {
template<typename T, int D, int E>
DValue operator()(const Array<T, D>& a, const Dim<E>& d) const {
std::stringstream ss;
ss << "Index dimension does not match array in DevalueVisitor, type is: " <<
a.get_type_and_dim() << " is a gpu place: " << is_gpu_place(a.place()) <<
" size: " << a.size() << " and dim is " << d.dimensions << "D with contents: " <<
d;
throw std::invalid_argument(ss.str());
}
template<typename T, int D>
DValue operator()(const Array<T, D>& a, const Dim<D>& d) const {
return a[d];
}
};两个visitor分别用于返回引用和值，可以看到在每个visitor中，operator()有两种函数模板定义，上面一种针对异常情况，下面一种针对正常情况。如果Dim和Array有一致的维度，则走下面一种operator()模板定义，特化出的函数直接调用Array的[]返回值或者引用；如果两者维度不一致，意味着参数不合法，走上面一种operator()模板定义，特化出的函数专门用于报错。
TODO：此处有个问题，既然对于维度不一致的处理方法就是报错，那可否直接不写上面一种operator()模板定义？这样应该在编译期就可以报错了，而现在的报错还是在运行期。
上面的代码中还用到了DValue和DReference这两个概念。与DArray和DDim类似，它们分别是对不同类型的值和Refernece的封装，使用的也是boost::variant。其中DValue的定义为：
typedef boost::variant<float, double, float16> DValue;元素访问
DValue get(const DArray& arr, const DDim& idx) {
return arr[idx];
}
void set(DArray& arr, const DDim& idx, const DValue& value) {
arr[idx] = value;
}内部调用DArray::[]来实现。
返回DArray::var的成员变量
DArray::var是一个Array，有一些全局函数用来返回它的size_、stride_等成员变量。
构造DArray
Majel提供了一系列的全局函数来构造DArray，这些函数名字都叫make_darray()。对DArray的构造其实就是对其中存放Array成员变量var的构造。构造过程同样通过visitor进行：
struct DArrayConstructor
: public boost::static_visitor<DArray> {
std::shared_ptr<Allocation> alloc;
DDim stride;
DArrayConstructor(std::shared_ptr<Allocation> alloc_,
DDim stride_) : alloc(alloc_), stride(stride_) {}
template<int i, typename T>
typename std::enable_if<(i < 5), DArray>::type
operator()(Dim<i> size, T type) {
Dim<i> s_stride = boost::get<Dim<i>>(stride);
return Array<T, i>(alloc, size, s_stride);
}
template<int i, typename T>
typename std::enable_if<(i >= 5), DArray>::type
operator()(Dim<i> dims, T type) {
throw std::invalid_argument("DArrays are limited to 4 dimensions");
}
};
DArray make_darray(std::shared_ptr<Allocation> alloc,
DDim dims,
DDim strides,
DValue type) {
DArrayConstructor ctor(alloc, strides);
return boost::apply_visitor(ctor, dims, type);
}有这么几点需要注意一下：
-
visitor并不负责内存的管理，Allocation需要从外面作为初始化参数传递给visitor。如果有需要，Allocation的申请可以在make_darray()中调用boost::apply_visitor()之前进行：
DArray make_darray(DDim dims, DDim strides, DValue type, Place place) { size_t size = product(dims) * boost::apply_visitor(DSizeOf(), type); auto alloc = std::make_shared<Allocation>(size, place); return make_darray(alloc, dims, strides, type); }
- 虽然定义的返回类型是
DArray（boost::static_visitor<DArray>），但实际return的是一个Array。这里会自动调用DArray对于Array类型的=重载，将返回结果封装为DArray。 -
make_darray()有一个DValue type参数，该最终参数作为T type传递给visitor的operator()。然而，在整个过程中type都只是用于推导确定Array的数据类型，其本身的值完全没有被使用。作为由用户指定的值，这么设计是不是会给用户一种会将生成的DArray初始化成这个值的错觉？
在使用中，经常会需要在Array和DArray之间相互转换，下面对此进行说明。
DArray的定义中重载了operator[]:
template<typename T>
DArray& operator=(T in) {
var = in;
return *this;
}var就是DArray中用于存放各种类型Array的boost::variant成员变量。通过特化模板参数T，该等号操作符可以接受各种类型的Array，并将参数赋值给var，完成Array向DArray的转化（或者说“封装”）。
这种转化过程在Majel代码中大量隐式地被调用。例如在很多函数中，虽然声明的返回值是DArray，但实际上return的是一个Array对象，这时候转化就会被自动调用。
Majel重载了boost::get()函数，用来获取DArray中封装的Array对象：
namespace boost {
template<typename T>
T get(const majel::DArray& in) {
return boost::get<T>(in.var);
}
}其实大部分时候，我们的需求并不是真的将DArray转化为Array，而是希望能够访问和使用DArray中封装的Array。Majel提供了一种基于boost::apply_visitor()的统一实现，来让用户可以自己定义对封装的Array的操作，并在全局函数中加以调用。
这种实现的关键在于，为DArray类实现apply_visitor()成员函数：
template<typename Visitor>
typename Visitor::result_type
apply_visitor(Visitor& visitor) {
return var.apply_visitor(visitor);
}
template<typename Visitor>
typename Visitor::result_type
apply_visitor(Visitor& visitor) const {
return var.apply_visitor(visitor);
}为了进一步解释这种实现的原理，我们需要先了解一下boost::apply_visitor()内部的实现原理。假设我们已经定义了一个visitor和一个boost::variant类型的变量v，那么用visitor对v进行访问和操作的通用方法是：
boost::apply_visitor(visitor, v);事实上在boost内部，这个调用会被转化成下面的形式：
v.apply_visitor(visitor);只要变量v有apply_visitor()这个成员函数，这个调用就能成功。boost::variant内部已经定义实现了自己的apply_visitor()，这个apply_visitor()的具体行为是：调用传进来的visitor的operator()，并将自身所存储的对象作为参数传进去。
这意味着对于一般的类来说，只要我们定义一下它的apply_visitor()成员函数，就可以使用boost::apply_visitor()和定义好的visitor来访问它，并且可以在apply_visitor()中订制我们自己的处理visitor的逻辑。DArray正是利用了这样的做法。
假设有一个名为darray的DArray对象，另外有针对Array进行某种操作的array_visitor，那么如果想对darray中封装的Array对象执行这种操作，只需要直接：
boost::apply_visitor(array_visitor, darray);就可以。注意传入的参数可以直接是darray，无需将封装的Array从中取出。原理如下：
-
boost::apply_visitor(array_visitor, darray)会进一步调用darray.apply_visitor(array_visitor) - 根据上文贴出的
DArray类型的成员函数apply_visitor()的定义，接下来会继续调用var.apply_visitor(array_visitor) -
var是一个boost::variant变量，它的成员函数apply_visitor()的具体行为是：调用传进来的array_visitor的operator()，并将自身所存储的Array作为参数传进去。 -
Array被传给operator()，执行操作。
Majel的DDim和Dim之间同样使用了这样的技巧。如果将Dim和Array这类对象称为“静态对象”、将DDim和DArray这类称为“动态对象”、将访问和操作静态对象的visitor称为“静态visitor”，那么我们可以得出如下结论：
在Majel风格的代码实现中，我们可以直接写出boost::apply_visitor(静态visitor, 动态对象)这样的代码，静态visitor会自动向动态对象内部传递，并最终作用于其中封装的静态对象。