Tensor.h

Go to the documentation of this file.

#pragma once
#include <cstddef>
#include <vector>
#include <deepity/backend/IComputeBackend.h>
#include <deepity/backend/DeviceType.h>

 
namespace Deep
{
    class Tensor
    {
    public:
        Tensor(IComputeBackend *backend, DeviceType device, size_t numFloats)
            : backend(backend), device(device), size(numFloats)
        {
            data = backend->Allocate(numFloats);
            backend->Zero(data, numFloats);
        }

        Tensor(IComputeBackend *backend, DeviceType device, const float *hostData, size_t numFloats)
            : backend(backend), device(device), size(numFloats)
        {
            data = backend->Allocate(numFloats);
            backend->CopyFromHost(data, hostData, numFloats);
        }

        Tensor(IComputeBackend *backend, DeviceType device, const std::vector<float> &hostData)
            : Tensor(backend, device, hostData.data(), hostData.size())
        {
        }

        ~Tensor()
        {
            if (data != nullptr)
            {
                backend->Free(data);
            }
        }
 
        // Move-only -- see file-level note.
        Tensor(const Tensor &) = delete;
        Tensor &operator=(const Tensor &) = delete;

        Tensor(Tensor &&other) noexcept
            : backend(other.backend), device(other.device), data(other.data), size(other.size)
        {
            other.data = nullptr;
            other.size = 0;
        }

        Tensor &operator=(Tensor &&other) noexcept
        {
            if (this != &other)
            {
                if (data != nullptr)
                {
                    backend->Free(data);
                }
                backend = other.backend;
                device = other.device;
                data = other.data;
                size = other.size;
                other.data = nullptr;
                other.size = 0;
            }
            return *this;
        }

        float *Data() noexcept { return data; }
        const float *Data() const noexcept { return data; }

        size_t Size() const noexcept { return size; }
        DeviceType Device() const noexcept { return device; }
        IComputeBackend *Backend() const noexcept { return backend; }

        void CopyToHost(float *hostDst) const noexcept
        {
            backend->CopyToHost(hostDst, data, size);
        }

        void CopyToHost(std::vector<float> &hostDst) const
        {
            hostDst.resize(size);
            CopyToHost(hostDst.data());
        }
 
    private:
        IComputeBackend *backend = nullptr; // non-owning; must outlive this Tensor
        DeviceType device;
        float *data = nullptr;
        size_t size = 0;
    };
}