tensor.h

Go to the documentation of this file.

#pragma once

#include <vector>
#include "types.h"
#include "memory/memorymanager.h"
#include "tensor_internal.h"
#include "utilities_internal.h"

#if defined(_HAS_CUDA_)
#include "cudnn.h"
#endif

namespace magmadnn {

/* Default values for tensors.
   Initialize to CPU 0 if not indicated otherwise.
   And don't fill the tensor on creation unless specified.
*/
const memory_t TENSOR_DEFAULT_MEM_TYPE = HOST;
const device_t TENSOR_DEFAULT_DEVICE_ID = (device_t) 0;
const tensor_fill_t TENSOR_DEFAULT_FILL_TYPE = NONE;
const tensor_filler_t<float> TENSOR_DEFAULT_FILLER = { TENSOR_DEFAULT_FILL_TYPE, {} };


template <typename T>
class Tensor {
public:

    Tensor(std::vector<unsigned int> shape);

    Tensor(std::vector<unsigned int> shape, memory_t mem_type);

    Tensor(std::vector<unsigned int> shape, memory_t mem_type, device_t device_id);

    Tensor(std::vector<unsigned int> shape, tensor_filler_t<T> filler);

    Tensor(std::vector<unsigned int> shape, tensor_filler_t<T> filler, memory_t mem_type);
    
    Tensor(std::vector<unsigned int> shape, tensor_filler_t<T> filler, memory_t mem_type, device_t device_id);

    ~Tensor();


    magmadnn_error_t copy_from(const Tensor<T>& src, unsigned int begin_idx, unsigned int size);

    magmadnn_error_t copy_from(const Tensor<T>& src);


    magmadnn_error_t copy_from(const Tensor<T>& src, const std::vector<unsigned int>& dims);


    T get(const std::vector<int>& idx) const;

    T get(const std::vector<unsigned int>& idx) const;

    T get(const std::initializer_list<int>& idx) const;

    T get(const std::initializer_list<unsigned int>& idx) const;

    T get(unsigned int flattened_idx) const;

    const T operator[](unsigned int idx) const;

    const T operator[](const std::initializer_list<unsigned int>& idx);

    void set(const std::vector<int>& idx, T val);

    void set(const std::vector<unsigned int>& idx, T val);

    void set(const std::initializer_list<int>& idx, T val);

    void set(const std::initializer_list<unsigned int>& idx, T val);

    void set(unsigned int flattened_idx, T val);    
    

    MemoryManager<T>* get_memory_manager() const { return this->mem_manager; }

    std::vector<unsigned int> get_shape() const { return this->shape; }

    unsigned int get_shape(unsigned int idx) const;

    unsigned int get_size() const { return this->size; }

    T* get_ptr() { return this->mem_manager->get_ptr(); }

    memory_t get_memory_type() const { return this->mem_type; }

    device_t get_device_id() const { return this->device_id; }

    #if defined(_HAS_CUDA_)
    cudnnTensorDescriptor_t get_cudnn_tensor_descriptor() { return desc; }
    #endif

    void reshape(const std::vector<unsigned int>& dims);

    void squeeze();

    void unsqueeze(unsigned int dim = 0);

private:
    void init(std::vector<unsigned int>& shape, tensor_filler_t<T> filler, memory_t mem_type, device_t device_id);
    unsigned int get_flattened_index(const std::vector<unsigned int>& idx) const;
    unsigned int get_flattened_index_old(const std::vector<unsigned int>& idx) const;

    /* device specific code */
    #if defined(_HAS_CUDA_)
    void init_cudnn_descriptor();
    void free_cudnn_descriptor();

    cudnnTensorDescriptor_t desc;
    #endif

    MemoryManager<T> *mem_manager;  /* allocated by init */
    
    std::vector<unsigned int> shape;    /* tensor axes (shape) */
    std::vector<unsigned int> strides;  /* axis strides in memory */
    unsigned int size;      /* total number of elements in tensor */
    memory_t mem_type;      /* the type of memory to use for this tensor */
    device_t device_id;     /* device number i.e. gpu0 or cpu1 */

};

typedef Tensor<int> tensori_t;
typedef Tensor<float> tensorf_t;
typedef Tensor<double> tensord_t;

} // namespace magmadnn