tensor.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <string>
#include <stdexcept>
#include <initializer_list>
#include <vector>
 
#include "pulsatrix/assert.hpp"
#include "pulsatrix/device_backend.hpp"
#include "pulsatrix/host_guard.hpp"
#include "pulsatrix/shape.hpp"
 
namespace pulsatrix {
 
class Tensor {
public:
    Tensor(Shape shape, DeviceBackend* backend);
 
    Tensor(Shape shape, DeviceBackend* backend, DeviceType device);
 
    Tensor(Shape shape, DeviceBackend* backend, std::initializer_list<float> values, DeviceType device);
 
    Tensor(Shape shape, DeviceBackend* backend, std::initializer_list<float> values);
 
    Tensor(Shape shape, DeviceBackend* backend, const std::vector<float>& values, DeviceType device);
 
    Tensor(Shape shape, DeviceBackend* backend, const std::vector<float>& values);
 
    ~Tensor();
 
    [[nodiscard]] static Tensor Stack(const std::vector<Tensor>& tensors, DeviceBackend* backend);
 
    Tensor(const Tensor& other);
    Tensor& operator=(const Tensor& other);
    Tensor(Tensor&& other) noexcept;
    Tensor& operator=(Tensor&& other) noexcept;
 
    [[nodiscard]] const Shape& shape() const { return shape_; }
 
    [[nodiscard]] int64_t numel() const { return shape_.numel(); }
 
    [[nodiscard]] int64_t rank() const { return shape_.rank(); }
 
    [[nodiscard]] DeviceType device() const { return device_; }
 
    [[nodiscard]] DeviceBackend* backend() const { return backend_; }
 
    [[nodiscard]] bool requires_grad() const { return requires_grad_; }
 
    void set_requires_grad(bool requires_grad) { requires_grad_ = requires_grad; }
 
    [[nodiscard]] float read_element(int64_t flat_index) const;
 
    [[nodiscard]] std::vector<float> to_host_vector() const;
 
    void write_element(int64_t flat_index, float value);
 
    [[nodiscard]] const float* data() const { return data_; }
 
    [[nodiscard]] float* data() { return data_; }
 
    [[nodiscard]] float& at(std::initializer_list<int64_t> index);
 
    [[nodiscard]] const float& at(std::initializer_list<int64_t> index) const;
 
    [[nodiscard]] float& operator[](int64_t flat_index) {
        PULSATRIX_ASSERT(flat_index >= 0 && flat_index < numel());
        return data_[flat_index];
    }
 
    [[nodiscard]] const float& operator[](int64_t flat_index) const {
        PULSATRIX_ASSERT(flat_index >= 0 && flat_index < numel());
        return data_[flat_index];
    }
 
    Tensor& fill(float value);
 
    Tensor& accumulate(const Tensor& other);
 
    Tensor& reshape(Shape new_shape);
 
    Tensor& to(DeviceType target);
 
    Tensor& to(DeviceType target, DeviceBackend* target_backend);
 
private:
    [[nodiscard]] int64_t flat_index_of(std::initializer_list<int64_t> index) const;
 
    float* data_;
    Shape shape_;
    DeviceBackend* backend_;
    DeviceType device_;
    bool requires_grad_ = true;
};
 
inline void require_device(const Tensor& t, DeviceType expected, const char* where) {
    if (t.device() != expected) {
        auto name = [](DeviceType d) {
            switch (d) {
                case DeviceType::Cpu:
                    return "Cpu";
                case DeviceType::Cuda:
                    return "Cuda";
                case DeviceType::Hip:
                    return "Hip";
            }
            return "unknown";
        };
        throw std::invalid_argument(std::string(where) + ": tensor is on " + name(t.device()) +
                                    " but this computes on " + name(expected) + "; move it first with Tensor::to()");
    }
}
 
}  // namespace pulsatrix