tensor_impl.hpp

#pragma once
 
#include <iostream>
#include <numeric>
#include <vector>
#include "data_type.hpp"
#include "device_type.hpp"
#include "dispatch_type.hpp"
#include "lamppp/common/assert.hpp"
#include "lamppp/tensor/align_utils.hpp"
#include "lamppp/tensor/native/memory_ops.hpp"
#include "lamppp/tensor/storage.hpp"
#include "scalar.hpp"
 
namespace lmp::tensor {
 
 
class TensorImpl {
 public:
  template <typename T>
  explicit TensorImpl(const std::vector<T>& data,
                      const std::vector<size_t>& shape, DeviceType device,
                      DataType dtype)
      : data_(LMP_DISPATCH_ALL_TYPES(
            dtype,
            [&] { return Storage(data.size() * sizeof(scalar_t), device); })),
        shape_(shape),
        type_(dtype),
        strides_(std::vector<detail::stride_t>(shape.size())),
        numel_(shape.empty() ? 0
                             : std::accumulate(shape.begin(), shape.end(), 1,
                                               std::multiplies<>())) {
    LMP_CHECK(data.size() == numel_) <<
              "Size mismatch, product of shape must equal num elements";
    DataType src_dtype = TypeMeta<T>::value;
    ops::copy_stub()(DeviceType::CPU, device, data.data(),
                                data_.data(), numel_, src_dtype, type_);
    update_strides();
  }
  explicit TensorImpl(Storage storage, const std::vector<size_t>& shape,
                      DataType dtype);
 
  void* data() const noexcept;
  DataType type() const noexcept;
  DeviceType device() const noexcept;
  const std::vector<size_t>& shape() const noexcept;
  const std::vector<detail::stride_t>& strides() const noexcept;
  size_t numel() const noexcept;
 
  TensorImpl reshape(std::vector<size_t> new_shape);
  TensorImpl squeeze(size_t dim);
  TensorImpl expand_dims(size_t dim);
  Scalar index(const std::vector<size_t>& idx);
 
  void copy(const TensorImpl& other) const;
  void fill(Scalar item) const;
  void print(std::ostream& os) const;
 
 private:
  friend class Tensor;
 
  void update_strides();
 
  DataType type_;
  Storage data_;
  size_t numel_;
  std::vector<size_t> shape_;
  std::vector<detail::stride_t> strides_;
};
 
}  // namespace lmp::tensor