device.hpp

#pragma once
 
#include "tensor/shape.hpp"
#include "tensor/tensorimpl.hpp"
#include "utils/scopeguard.hpp"
 
namespace dl {
    class Device {
    private:
        thread_local static Device const* defaultDevice;
 
        Device(const Device& other) = delete;
        Device(Device&& other) = delete;
        Device& operator=(const Device& other) = delete;
        Device& operator=(Device&& other) = delete;
 
    protected:
        Device() = default;
 
    public:
        virtual ~Device() = default;
 
        static Device const& cpu;
 
        virtual TensorPtr empty(Shape shape, bool requiresGrad = false) const noexcept = 0;
        virtual TensorPtr zeros(Shape shape, bool requiresGrad = false) const noexcept = 0;
        virtual TensorPtr ones(Shape shape, bool requiresGrad = false) const noexcept = 0;
        virtual TensorPtr rand(Shape shape, bool requiresGrad = false) const noexcept = 0;
        virtual TensorPtr constant(int value, bool requiresGrad = false) const noexcept = 0;
        virtual TensorPtr constant(float value, bool requiresGrad = false) const noexcept = 0;
        virtual TensorPtr constant(double value, bool requiresGrad = false) const noexcept = 0;
        virtual TensorPtr constant(InitializerTensor<float>&& value, bool requiresGrad = false) const noexcept = 0;
        virtual TensorPtr fromBytesFP32(const char* buffer, size_t bufsize, Shape shape) const noexcept = 0;
 
        template <typename T>
        void setDefaultFloatTensorType();
 
        static const Device& getDefault() noexcept { return *defaultDevice; }
 
        static auto changeDefaultDevice(Device& device) {
            auto oldDefault = defaultDevice;
            defaultDevice = &device;
            return dl::utils::ScopeGuard([oldDefault] { Device::defaultDevice = oldDefault; });
        }
    };
 
    template <>
    void Device::setDefaultFloatTensorType<float>();
    template <>
    void Device::setDefaultFloatTensorType<double>();
 
    inline TensorPtr empty(const Shape& size, Device const& device = Device::getDefault()) {
        return device.empty(size);
    }
    inline TensorPtr zeros(const Shape& size, Device const& device = Device::getDefault()) {
        return device.zeros(size);
    }
    inline TensorPtr ones(const Shape& size, Device const& device = Device::getDefault()) { return device.ones(size); }
    inline TensorPtr rand(const Shape& size, Device const& device = Device::getDefault()) { return device.rand(size); }
    TensorPtr zeros_like(const TensorPtr& tensor, Device const& device = Device::getDefault());
    TensorPtr ones_like(const TensorPtr& tensor, Device const& device = Device::getDefault());
    TensorPtr rand_like(const TensorPtr& tensor, Device const& device = Device::getDefault());
    inline TensorPtr constant(int value, Device const& device = Device::getDefault()) { return device.constant(value); }
    inline TensorPtr constant(float value, Device const& device = Device::getDefault()) {
        return device.constant(value);
    }
    inline TensorPtr constant(double value, Device const& device = Device::getDefault()) {
        return device.constant(value);
    }
    inline TensorPtr constant(InitializerTensor<float>&& value, Device const& device = Device::getDefault()) {
        return device.constant(std::move(value));
    }
 
    template <typename T>
    TensorPtr
    fromBytes(const char* buffer, size_t buflen, const Shape& shape, Device const& device = Device::getDefault());
    template <>
    inline TensorPtr fromBytes<float>(const char* buffer, size_t buflen, const Shape& shape, Device const& device) {
        return device.fromBytesFP32(buffer, buflen, shape);
    }
 
    TensorPtr clone(const TensorPtr& tensor);
} // namespace dl