array.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_ARRAY_H
#define FREE_TENSOR_ARRAY_H
 
#include <cstdint>
#include <vector>
 
#include <driver/device.h>
#include <tensor.h>
 
namespace freetensor {
 
struct ArrayCopy {
    Ref<Device> device_;
    uint8_t *ptr_ = nullptr;
    bool borrowed_ = false;
 
    ArrayCopy(const Ref<Device> &device, uint8_t *ptr, bool borrowed)
        : device_(device), ptr_(ptr), borrowed_(borrowed) {}
};
 
class Array {
    // Identical data on different devices.
    std::vector<ArrayCopy> ptrs_;
 
    // Temorary copy of the data used by AccessType::InputMutable. The data may
    // become different after the program runs.
    std::optional<ArrayCopy> tempPtr_;
 
    size_t size_ = 0, nElem_ = 0;
    std::vector<size_t> shape_;
    DataType dtype_;
    bool dontDropBorrow_, moved_;
 
  private:
    Array(const std::vector<size_t> &shape, DataType dtype,
          bool dontDropBorrow = false, bool moved = false);
 
  public:
    static Array moveFromRaw(void *ptr, const std::vector<size_t> &shape,
                             DataType dtype, const Ref<Device> &device);
 
    static Array borrowFromRaw(void *ptr, const std::vector<size_t> &shape,
                               DataType dtype, const Ref<Device> &device,
                               bool dontDropBorrow, bool moved);
 
    ~Array();
 
    Array(Array &&);
    Array(const Array &) = delete;
 
    Array &operator=(Array &&);
    Array &operator=(const Array &) = delete;
 
    size_t size() const { return size_; }
    size_t nElem() const { return nElem_; }
    const std::vector<size_t> &shape() const { return shape_; }
    DataType dtype() const { return dtype_; }
 
    bool dontDropBorrow() const { return dontDropBorrow_; }
    void setDontDropBorrow(bool flag) { dontDropBorrow_ = flag; }
 
    bool moved() const { return moved_; }
    void setMoved(bool flag) { moved_ = flag; }
 
    void *rawSharedTo(const Ref<Device> &device);
    void *rawMovedTo(const Ref<Device> &device);
    void *rawInitTo(const Ref<Device> &device);
    void *rawTemporarilyCopiedTo(const Ref<Device> &device);
 
    void makePrivateCopy();
};
 
} // namespace freetensor
 
#endif // FREE_TENSOR_ARRAY_H