tensor.h

Go to the documentation of this file.

#ifndef ATEN_CORE_TENSOR_H_
#define ATEN_CORE_TENSOR_H_
 
#include <base/core/allocator.h>
#include <ATen/core/tensor_types.h>
#include <ATen/core/tensor_shape.h>
#include <ATen/core/tensor_buffer.h>
#include <ATen/core/tensor_accessor.h>
 
#include <ATen/kernels/memory.h>
 
#include <base/macros/macros.h>
 
// TODO: 
// 1. add log system
// 2. add exception system
// 3. refact cmake system, use cmake parant scope to construct the libraries
 
namespace ct = container;
 
namespace container {
 
class Tensor {
  public:
    
    Tensor();
 
    explicit Tensor(DataType data_type);
 
    Tensor(DataType data_type, const TensorShape& shape);
 
    Tensor(DataType data_type, DeviceType device, const TensorShape& shape);
 
    Tensor(base::core::Allocator* a, DataType data_type, DeviceType device, const TensorShape& shape);
 
    Tensor(const Tensor& other);
 
    Tensor(Tensor&& other) noexcept;
    
    ~Tensor();
    template <typename T> 
    Tensor(std::initializer_list<T> values, DeviceType device = DeviceType::CpuDevice) :
        Tensor(DataTypeToEnum<T>::value, device, TensorShape({static_cast<int64_t>(values.size())})) {
        TEMPLATE_ALL_2(this->data_type_, this->device_,
            kernels::synchronize_memory<T, DEVICE_, DEVICE_CPU>()(
                this->data<T>(), values.begin(), this->NumElements()))
    }
 
    DataType data_type() const;
 
    DeviceType device_type() const;
 
    const TensorShape& shape() const;
 
    int64_t NumElements() const;
 
    void* data() const;
 
    template <typename T>
    T* data() const {
        if ((std::is_same<T, float>::value && data_type_ != DataType::DT_FLOAT) ||
            (std::is_same<T, int>::value && data_type_ != DataType::DT_INT) ||
            (std::is_same<T, int64_t>::value && data_type_ != DataType::DT_INT64) ||
            (std::is_same<T, double>::value && data_type_ != DataType::DT_DOUBLE) ||
            (std::is_same<T, std::complex<float>>::value && data_type_ != DataType::DT_COMPLEX) ||
            (std::is_same<T, std::complex<double>>::value && data_type_ != DataType::DT_COMPLEX_DOUBLE))
        {
            std::cerr << "Tensor data type does not match requested type." << std::endl;
            exit(EXIT_FAILURE);
        }
        return buffer_->base<T>();
    }
 
    static size_t SizeOfType(DataType data_type) {
        switch (data_type) {
            case DataType::DT_FLOAT:
                return sizeof(float);
            case DataType::DT_INT:
                return sizeof(int32_t);
            case DataType::DT_INT64:
                return sizeof(int64_t);
            case DataType::DT_DOUBLE:
                return sizeof(double);
            case DataType::DT_COMPLEX:
                return sizeof(std::complex<float>);
            case DataType::DT_COMPLEX_DOUBLE:
                return sizeof(std::complex<double>);
            default:
                std::cerr << "Unsupported data type!" << std::endl;
                exit(EXIT_FAILURE);
        }
    }
 
    const TensorBuffer& buffer() const;
 
    template <typename DEVICE>
    Tensor to_device() const {
        if (this->device_ == DeviceTypeToEnum<DEVICE>::value) {
            return *this;
        }
        // Create output tensor on device
        Tensor output(this->data_type_, DeviceTypeToEnum<DEVICE>::value, this->shape_);
 
        // Copy data to a specified device
        // TODO: move the memory operator into the tensor_buff class.
        TEMPLATE_ALL_2(this->data_type_, this->device_,
                   kernels::synchronize_memory<T_, DEVICE, DEVICE_>()(
                           output.data<T_>(), this->data<T_>(), this->NumElements()))
 
        return output;
    }
 
    template <typename DEVICE, typename T>
    void copy_from_device(const T* data, int64_t num_elements = -1) {
        if (num_elements == -1) {
            num_elements = this->NumElements();
        }
        REQUIRES_OK(this->shape_.NumElements() >= num_elements,
                    "The number of elements of the input data must match the number of elements of the tensor.")
 
        TEMPLATE_CZ_2(this->data_type_, this->device_,
                   kernels::cast_memory<T_, T, DEVICE_, DEVICE>()(
                           this->data<T_>(), data, num_elements))
    }
 
    template <typename T>
    Tensor cast() const {
        // Create output tensor on device
        Tensor output(DataTypeToEnum<T>::value, this->device_, this->shape_);
 
        // TODO: error handle of cast memory
        // TODO: move the memory operator into the tensor_buff class.
        // Copy data to a specified device
        TEMPLATE_CZ_2(this->data_type_, this->device_,
                   kernels::cast_memory<T, T_, DEVICE_, DEVICE_>()(
                           output.data<T>(), this->data<T_>(), this->NumElements()))
 
        return output;
    }
 
    void zero();
 
    void reshape(TensorShape shape);
 
    Tensor shaped(const TensorShape& shape) const;
 
    Tensor slice(const std::vector<int>& start, const std::vector<int>& size) const;
 
    void resize(const TensorShape& new_shape);
 
    // TODO: Delete this function, and use a singleton allocator class.
    static base::core::Allocator* GetAllocator(DeviceType device);
 
    template <typename T, typename... Indices>
    T& get_value(Indices... indices) const {
        if (sizeof...(Indices) != shape_.ndim()) {
            throw std::invalid_argument("Incorrect number of indices.");
        }
 
        // Calculate the linear index corresponding to the given indices
        size_t linearIndex = calculateLinearIndex(indices...);
 
        // Access the element at the calculated linear index
        return *reinterpret_cast<T*>(data<T>() + linearIndex);
    }
 
    template <typename T>
    T* inner_most_ptr(const int &index) const {
        if (shape_.ndim() > 2) {
            throw std::invalid_argument("Invalid call, inner_most_ptr only support tensor rank <= 2!");
        }
        if (index > shape_.dim_size(static_cast<int>(shape_.ndim() - 2))) {
            throw std::invalid_argument("Invalid index, index of the inner-most must less than the inner-most shape size!");
        }
        if (shape_.ndim() == 1) {
            return data<T>() + index;
        }
        return data<T>() + index * shape_.dim_size(static_cast<int>(shape_.ndim()) - 1);
    }
 
 
    bool operator==(const Tensor& other) const;
 
    Tensor& operator=(const Tensor& other);
 
    Tensor& operator=(Tensor&& other) noexcept;
 
    bool CopyFrom(const Tensor& other);
 
    bool CopyFrom(const Tensor& other, const TensorShape& shape);
 
    bool AllocateFrom(const Tensor& other, const TensorShape& shape);
 
    template <typename T, size_t N, typename index_t = int64_t>
    TensorAccessor<T, N, index_t> accessor() const& {
        // Check if the tensor's rank is greater than 0
        static_assert(N > 0, 
            "Accessor is used to access the data of a tensor with rank > 0, for scalars use *data<T>()");
        // Check if the rank of the tensor matches the rank of the accessor
        REQUIRES_OK(this->shape_.ndim() == N, 
            "The rank of the tensor must match the rank of the accessor.")
        // Create and return a TensorAccessor object
        return TensorAccessor<T, N, index_t>(this->data<T>(), this->shape_.dims().data(), this->shape_.strides().data());
    }
    template<typename T, size_t N, typename index_t = int>
    TensorAccessor<T, N, index_t> accessor() && = delete;
 
    void sync(const Tensor& rhs);
 
    Tensor operator[] (const int& index) const;
 
    explicit operator bool() const {
        return this->NumElements() > 0;
    }
 
    template<typename T>
    void set_value(T value) {
        TEMPLATE_ALL_2(this->data_type_, this->device_,
            kernels::set_memory<T, DEVICE_>()(this->data<T>(), value, this->NumElements()))
    }
 
protected:
 
    DataType data_type_;
 
    DeviceType device_;
 
    TensorShape shape_;
 
    TensorBuffer* buffer_{};
 
    template <typename... Indices>
    size_t calculateLinearIndex(Indices... indices) const {
        size_t stride = 1;
        size_t linearIndex = 0;
        size_t indexArray[] = { static_cast<size_t>(indices)... };
 
        for (int ii = static_cast<int>(shape_.ndim()) - 1; ii >= 0; --ii) {
            linearIndex += indexArray[ii] * stride;
            stride *= shape_.dim_size(ii);
        }
        return linearIndex;
    }
 
    // This function is used to copy data and properties from another Tensor instance, 'other', into the current Tensor instance.
    // The 'shape' parameter specifies the new shape for the current Tensor.
    inline void CopyFromInternal(const Tensor& other, const TensorShape& shape) {
        // Copy the data type and device from the 'other' Tensor.
        data_type_ = other.data_type_;
        device_ = other.device_;
        // Set the shape of the current Tensor to the provided 'shape'.
        shape_ = shape;
        // Check if the buffer of the current Tensor is different from the buffer of the 'other' Tensor.
        if (buffer_ != other.buffer_) {
            // If the current Tensor has a buffer, decrease its reference count.
            // Note this could indicate a deleted of current buffer_
            if (buffer_) buffer_->unref();
            // Assign the buffer of the 'other' Tensor to the current Tensor's buffer.
            buffer_ = other.buffer_;
            // Increase the reference count of the buffer to indicate shared ownership.
            if (buffer_) buffer_->ref();
        }
    }
 
};
 
std::ostream& operator<<(std::ostream& os, const Tensor& tensor);
 
} // namespace container
 
#endif // ATEN_CORE_TENSOR_H_