src/backends/cpu/tensor_factory.cpp - with_cpu - vuron/adept

src/backends/cpu/tensor_factory.cpp
103 строки3 KB
kolkir
Refactor tensor factory to get rid of temp objects
23 мар 2025, 11:53
23 мар 2025, 11:53f117bd9
#include <adept/backends/cpu/tensor_factory.hpp>
#include <adept/backends/cpu/tensorimpl.hpp>
#include <adept/types_dispatch.hpp>

#include <hwy/aligned_allocator.h>
#include <hwy/highway.h>

namespace adept::cpu {

namespace {

template <typename T>
class SIMDBuffer : public Buffer {
 public:
  SIMDBuffer(hwy::AlignedFreeUniquePtr<T[]> data, size_t size)
      : data_(std::move(data)), size_(size) {}
  ~SIMDBuffer() = default;
  SIMDBuffer(const SIMDBuffer&) = delete;
  SIMDBuffer& operator=(const SIMDBuffer&) = delete;
  SIMDBuffer(SIMDBuffer&&) = default;
  SIMDBuffer& operator=(SIMDBuffer&&) = default;

  void* ptr() const override { return data_.get(); }

  /**
   * @brief Return size of buffer in data type items, real size can be bugger due to alignemnt
   *
   * @return size_t size of buffer
   */
  size_t size() const override { return size_; }

 private:
  hwy::AlignedFreeUniquePtr<T[]> data_;
  size_t size_;
};

template <typename T>
requires std::integral<T> || std::floating_point<T>
class SIMDAllocator : public Allocator {
 public:
  buffer_ptr_t allocate(size_t size) {
    auto buffer = hwy::AllocateAligned<T>(size);
    auto mem_size = to_mem_size(size);
    return std::make_shared<SIMDBuffer<T>>(std::move(buffer), mem_size);
  }

  void free(buffer_ptr_t&& buffer) { std::move(buffer).reset(); }

  size_t to_mem_size(size_t numel) override { return numel * sizeof(T); }
};
}  // namespace

TensorFactory::TensorFactory() {
  auto make_mem_pool = [this]<typename T>() {
    mem_pools_[to_dtype<T>()] = std::make_shared<MemoryPool>(std::make_shared<SIMDAllocator<T>>());
  };
  for_each_data_type<float32_t, float64_t, int32_t, int8_t>(make_mem_pool);
}

std::shared_ptr<TensorImpl> TensorFactory::empty(const TensorProperties& props) {
  if (props.shape.empty()) {
    THROW_ERROR("Can't create tensor with empty shape ", props);
  }

  auto buffer = mem_pools_[props.dtype]->allocate(props.shape.numel());
  return std::make_shared<TensorImpl>(props, std::move(buffer));
}

std::shared_ptr<TensorImpl> TensorFactory::from_blob(const void* data,
                                                     const TensorProperties& props) {
  if (props.shape.empty()) {
    THROW_ERROR("Can't create tensor with empty shape ", props);
  }

  auto numel = props.shape.numel();
  auto buffer = mem_pools_[props.dtype]->allocate(numel);
  DISPATCH_TYPE(props.dtype,
                [&]() { hwy::CopyBytes(data, buffer->ptr(), numel * sizeof(scalar_t)); });
  return std::make_shared<TensorImpl>(props, std::move(buffer));
}

std::shared_ptr<TensorImpl> TensorFactory::clone(const DeviceTensor& tensor) {
  auto numel = tensor.properties().shape.numel();
  auto buffer = mem_pools_[tensor.properties().dtype]->allocate(numel);
  const auto& in_tensor = static_cast<const TensorImpl&>(tensor);
  DISPATCH_TYPE(tensor.properties().dtype, [&]() {
    hwy::CopyBytes(in_tensor.data(), buffer->ptr(), numel * sizeof(scalar_t));
  });
  return std::make_shared<TensorImpl>(tensor.properties(), std::move(buffer));
}

Tensor make_empty_tensor(const TensorProperties& props) {
  return Tensor(TensorFactory::instance().empty(props));
}

Tensor make_tensor_from_blob(const void* data, const TensorProperties& props) {
  return Tensor(TensorFactory::instance().from_blob(data, props));
}

Tensor clone_tensor(const Tensor& tensor) {
  return Tensor(TensorFactory::instance().clone(*tensor.impl()));
}
}  // namespace adept::cpu