297 lines
8.9 KiB
C++
297 lines
8.9 KiB
C++
|
|
|
|
#pragma once
|
|
|
|
#include <stddef.h>
|
|
#include <iostream>
|
|
#include <string>
|
|
#include <vector>
|
|
#include <utility>
|
|
#include <string.h>
|
|
#include <memory>
|
|
|
|
namespace dynasty {
|
|
namespace common {
|
|
struct BCHW {
|
|
uint32_t b_;
|
|
uint32_t c_;
|
|
uint32_t h_;
|
|
uint32_t w_;
|
|
BCHW(): b_(0),c_(0),h_(0),w_(0) {};
|
|
BCHW(const std::vector<int32_t>& shape) {
|
|
b_ = shape.size() > 0 ? shape.at(0) : 1;
|
|
c_ = shape.size() > 1 ? shape.at(1) : 1;
|
|
h_ = shape.size() > 2 ? shape.at(2) : 1;
|
|
w_ = shape.size() > 3 ? shape.at(3) : 1;
|
|
}
|
|
size_t getBatchSize() const {
|
|
return c_*h_*w_;
|
|
}
|
|
size_t getChannelSize() const {
|
|
return h_*w_;
|
|
}
|
|
size_t getSize() const {
|
|
return b_*c_*h_*w_;
|
|
}
|
|
};
|
|
//--------------------------------------------------------------------
|
|
// Tensor Shape
|
|
//--------------------------------------------------------------------
|
|
class TensorShape : private std::vector<int32_t> {
|
|
// We use negative numbers for unknown symbolic dimension. Each negative
|
|
// number represents a unique symbolic dimension.
|
|
// Private inheritance is used to prevent ambiguity of element versus dimension size
|
|
public:
|
|
TensorShape() = default;
|
|
|
|
TensorShape(const TensorShape& /*other*/) = default;
|
|
TensorShape& operator=(const TensorShape& /*other*/) = default;
|
|
|
|
TensorShape(TensorShape&& /*other*/) = default;
|
|
TensorShape& operator=(TensorShape&& /*other*/) = default;
|
|
|
|
TensorShape(const std::vector<int32_t>& dims) : std::vector<int32_t>(dims) {}
|
|
|
|
TensorShape(std::vector<int32_t>&& dims) : std::vector<int32_t>(std::move(dims)) {}
|
|
|
|
TensorShape(const std::initializer_list<int32_t>& dims) : std::vector<int32_t>(dims) {}
|
|
|
|
TensorShape(const std::vector<int32_t>& dims, size_t start, size_t end){
|
|
assign(dims.begin() + start, dims.begin() + end);
|
|
}
|
|
|
|
/**
|
|
Return the dimension specified by <idx>.
|
|
*/
|
|
const int32_t& operator[](size_t idx) const {
|
|
return std::vector<int32_t>::operator[](static_cast<int>(idx));
|
|
}
|
|
|
|
int32_t& operator[](size_t idx) {
|
|
return std::vector<int32_t>::operator[](static_cast<int>(idx));
|
|
}
|
|
|
|
bool operator==(const TensorShape& other) const noexcept {
|
|
auto thisVector = static_cast<const std::vector<int32_t>*>(this);
|
|
auto otherVector = static_cast<const std::vector<int32_t>*>(&other);
|
|
return *thisVector == *otherVector;
|
|
}
|
|
|
|
bool operator!=(const TensorShape& other) const noexcept {
|
|
return !(*this == other);
|
|
}
|
|
|
|
size_t NumDimensions() const noexcept {
|
|
return size();
|
|
}
|
|
|
|
/**
|
|
Copy dims into an array with given size
|
|
*/
|
|
void CopyDims(int32_t* dims, size_t num_dims) const {
|
|
memcpy(dims, data(), sizeof(value_type) * std::min(num_dims, NumDimensions()));
|
|
}
|
|
|
|
/**
|
|
Return underlying vector representation.
|
|
*/
|
|
const std::vector<int32_t>& GetDims() const { return *this; }
|
|
|
|
/**
|
|
* Return the total number of elements. Returns 1 for an empty (rank 0) TensorShape.
|
|
*
|
|
* May return -1
|
|
*/
|
|
int32_t Size() const{
|
|
size_t arraySize = size();
|
|
return SizeHelper(0, arraySize);
|
|
}
|
|
|
|
/**
|
|
Return the total number of elements up to the specified dimension.
|
|
If the dimension interval is empty (dimension == 0), return 1.
|
|
@param dimension Return size up to this dimension. Value must be between 0 and this->NumDimensions(), inclusive.
|
|
*/
|
|
int32_t SizeToDimension(size_t dimension) const{
|
|
const size_t num_dims = size();
|
|
if ( dimension > num_dims) {
|
|
printf("SizeToDimension: incorrect dimension\n");
|
|
exit(-1);
|
|
}
|
|
return SizeHelper(0, dimension);
|
|
|
|
}
|
|
|
|
/**
|
|
Return the total number of elements from the specified dimension to the end of the tensor shape.
|
|
If the dimension interval is empty (dimension == this->NumDimensions()), return 1.
|
|
@param dimension Return size from this dimension to the end. Value must be between 0 and this->NumDimensions(),
|
|
inclusive.
|
|
*/
|
|
int32_t SizeFromDimension(size_t dimension) const{
|
|
const size_t num_dims = size();
|
|
if ( dimension > num_dims) {
|
|
printf("SizeToDimension: incorrect dimension\n");
|
|
exit(-1);
|
|
}
|
|
return SizeHelper(dimension, num_dims);
|
|
|
|
}
|
|
|
|
std::string ToString() const {
|
|
std::string result;
|
|
|
|
result.append("{");
|
|
bool first = true;
|
|
for (auto dim : (*this)) {
|
|
if (!first) {
|
|
result.append(",");
|
|
}
|
|
|
|
result.append(std::to_string(dim));
|
|
first = false;
|
|
}
|
|
result.append("}");
|
|
|
|
return result;
|
|
}
|
|
|
|
/**
|
|
Calculate size between start and end.
|
|
Assumes start and end are between 0 and this->NumDimensions(), inclusive, and that
|
|
start < end.
|
|
*/
|
|
int32_t SizeHelper(size_t start, size_t end) const{
|
|
int32_t size = 1; // this is used to calculate the size, which is used for memory allocations, so validate no overflow
|
|
for (size_t i = start; i < end; i++) {
|
|
size *= (*this)[i];
|
|
}
|
|
return size;
|
|
}
|
|
/**
|
|
empty shape or 1D shape (1) is regarded as scalar tensor
|
|
*/
|
|
bool IsScalar() const {
|
|
size_t len = size();
|
|
return len == 0 || (len == 1 && operator[](0) == 1);
|
|
}
|
|
|
|
BCHW GetBCHW() const{
|
|
return BCHW(GetDims());
|
|
}
|
|
};
|
|
|
|
|
|
|
|
|
|
//---------------------------------------------------------------------------------------
|
|
// Tensor
|
|
// . Placeholder for a piece of memory, with additional shape information.
|
|
// . It's abstract class with virtual fun: host(), device()
|
|
//---------------------------------------------------------------------------------------
|
|
using Deletor = void(*)(void*);
|
|
|
|
class Tensor : public std::enable_shared_from_this<Tensor> {
|
|
public:
|
|
Tensor() = default;
|
|
Tensor(const TensorShape& shape, void * p_data, Deletor deletor) : p_data_(p_data),shape_(shape), deletor_(deletor){}
|
|
virtual ~Tensor() {
|
|
if(deletor_) deletor_(p_data_);
|
|
}
|
|
|
|
// Force to use of pointer after it's created
|
|
Tensor(const Tensor& other) = delete ;
|
|
Tensor(Tensor&& other) = delete;
|
|
Tensor& operator=(const Tensor& other) =delete;// copy assignment
|
|
Tensor& operator=(Tensor&& other) =delete; // move assignment
|
|
/**
|
|
Returns the shape of the tensor.
|
|
*/
|
|
const TensorShape& Shape() const noexcept { return shape_; }
|
|
|
|
/**
|
|
May return nullptr if tensor size is zero
|
|
*/
|
|
template<typename T>
|
|
T* MutableData() {
|
|
return reinterpret_cast<T*>(p_data_);
|
|
}
|
|
|
|
template<typename T>
|
|
std::pair<T*, size_t> MutableDataAndSize() {
|
|
return std::make_pair<T *, size_t>(MutableData<T>(), Size());
|
|
}
|
|
|
|
template<typename T>
|
|
const T* Data() const {
|
|
return reinterpret_cast<const T*>(p_data_);
|
|
}
|
|
|
|
template<typename T>
|
|
std::pair<const T*, size_t> DataAndSize() const {
|
|
return std::make_pair<const T *, size_t>(Data<T>(), Size());
|
|
}
|
|
|
|
size_t Size() const { return shape_.Size();}
|
|
const std::vector<int32_t>& GetShapeVector() const { return shape_.GetDims();}
|
|
|
|
virtual std::shared_ptr<Tensor> host() =0;
|
|
virtual std::shared_ptr<Tensor> device() =0;
|
|
|
|
protected:
|
|
|
|
void* p_data_ {};
|
|
TensorShape shape_ {};
|
|
Deletor deletor_{};
|
|
};
|
|
|
|
//----------------------------------------------------------------------------------
|
|
// CPU Tensor
|
|
//-----------------------------------------------------------------------------------
|
|
template <typename T>
|
|
class CPUTensor : public Tensor{
|
|
public:
|
|
CPUTensor() = default;
|
|
CPUTensor(const TensorShape& shape, void * p_data, Deletor deletor=nullptr) : Tensor(shape, p_data, deletor) {}
|
|
virtual ~CPUTensor() = default;
|
|
|
|
std::shared_ptr<Tensor> host() override;
|
|
std::shared_ptr<Tensor> device() override;
|
|
|
|
};
|
|
|
|
//----------------------------------------------------------------------------------
|
|
// CUDA Tensor
|
|
//-----------------------------------------------------------------------------------
|
|
template <typename T>
|
|
class CUDATensor : public Tensor {
|
|
public:
|
|
CUDATensor() = default;
|
|
CUDATensor(const TensorShape& shape, void * p_data, Deletor deletor=nullptr) : Tensor(shape, p_data, deletor) {}
|
|
virtual ~CUDATensor() = default;
|
|
|
|
std::shared_ptr<Tensor> host() override;
|
|
std::shared_ptr<Tensor> device() override;
|
|
|
|
};
|
|
|
|
//----------------------------------------------------------------------------------
|
|
// OpenCL Tensor
|
|
//-----------------------------------------------------------------------------------
|
|
template <typename T>
|
|
class OpenCLTensor : public Tensor{
|
|
public:
|
|
OpenCLTensor() = default;
|
|
OpenCLTensor(const TensorShape& shape, void * p_data, Deletor deletor=nullptr) : Tensor(shape, p_data, deletor) {}
|
|
virtual ~OpenCLTensor() = default;
|
|
|
|
std::shared_ptr<Tensor> host() override;
|
|
std::shared_ptr<Tensor> device() override;
|
|
|
|
};
|
|
|
|
|
|
|
|
} // namespace common
|
|
} // namespace dynasty
|