Program Listing for File tensor_allocator.h¶
↰ Return to documentation for file (src/tensors/tensor_allocator.h
)
#pragma once
#include <deque>
#include <set>
#include "common/definitions.h"
#include "tensors/allocator.h"
#include "tensors/tensor.h"
namespace marian {
class TensorAllocator {
private:
const size_t CHUNK = 128;
const size_t MBYTE = 1024 * 1024;
const size_t GROW = CHUNK * MBYTE;
const size_t ALIGN = 256;
Ptr<Backend> backend_;
Ptr<Allocator> allocator_;
public:
TensorAllocator(Ptr<Backend> backend)
: backend_(backend),
allocator_(New<Allocator>(backend_->getDeviceId(), 0, GROW, ALIGN)) {}
TensorAllocator(Ptr<Backend> backend, Ptr<Device> device)
: backend_(backend),
allocator_(New<Allocator>(backend_->getDeviceId(), device, 0, GROW, ALIGN)) {}
~TensorAllocator() { clear(); }
void throwAtReallocation(bool throwRealloc) {
allocator_->throwAtReallocation(throwRealloc);
}
void reserve(size_t bytes = 0) {
auto mult = bytes / GROW + 1;
LOG(info,
"[memory] Extending reserved space to {} MB (device {})",
mult * CHUNK,
allocator_->getDeviceId());
allocator_->reserve(mult * GROW);
}
void reserveExact(const std::vector<size_t>& bytes) {
size_t total = 0;
for(auto part : bytes)
total += allocator_->alignedSize(part);
reserveExact(total);
}
void reserveExact(size_t bytes = 0) {
size_t mbytes = bytes / MBYTE;
if(mbytes == 0) {
LOG(info,
"[memory] Reserving {} B, device {}",
bytes,
allocator_->getDeviceId());
} else {
LOG(info,
"[memory] Reserving {} MB, device {}",
mbytes,
allocator_->getDeviceId());
}
allocator_->reserve(bytes);
}
void clear() { allocator_->clear(); }
size_t capacity(Shape shape, Type type = Type::float32) {
return allocator_->capacity<char>(requiredBytes(shape, type));
}
void allocate(/*out*/ Tensor& t, Shape shape, Type type = Type::float32) {
if(!t || t->shape() != shape) {
auto mem = allocator_->alloc(requiredBytes(shape, type));
t = Tensor(TensorBase::New(mem, shape, type, backend_));
}
}
void free(const Tensor& t) { allocator_->free(t->memory()); }
Tensor asTensor(Type type = Type::float32) {
auto mem = allocator_->memory();
auto size = mem->size() / sizeOf(type);
return TensorBase::New(mem, Shape({1, (int)size}), type, backend_);
}
size_t size(Type type = Type::float32) { return allocator_->size() / sizeOf(type); }
Ptr<Allocator> allocator() { return allocator_; }
};
} // namespace marian