src/backends/cpu/tensor_arithmetics.cpp - master - vuron/adept

src/backends/cpu/tensor_arithmetics.cpp
335 строк12 KB
kolkir
Device checks and exchange
06 апр 2025, 15:28
06 апр 2025, 15:2873b23e8
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE "../src/backends/cpu/tensor_arithmetics.cpp"
#include <hwy/foreach_target.h>

#include <hwy/highway.h>
#include <hwy/per_target.h>
#include "arithmetics-inl.hpp"
#include "logic-inl.hpp"

#if HWY_ONCE

#include <adept/backends/cpu/tensor_factory.hpp>
#include <adept/backends/cpu/tensorimpl.hpp>
#include <adept/threading.hpp>
#include <adept/types_dispatch.hpp>
#include "arithmetics.hpp"  // should go after arithmetics-inl.hpp
#include "logic.hpp"        // should go after logic-inl.hpp
#include "loop.h"

namespace adept::cpu {

TensorImpl::TensorImpl(const TensorProperties& props, pool_buffer_ptr_t data_buffer)
    : props_(props),
      indexer_{&props_.shape},
      data_buffer_(std::move(data_buffer)),
      is_aligned_(is_shape_aligned_impl()) {}

TensorImpl::~TensorImpl() {}

TensorImpl::TensorImpl(TensorImpl&& other)
    : props_(std::move(other.props_)),
      indexer_(std::move(other.indexer_)),
      data_buffer_(std::move(other.data_buffer_)),
      is_aligned_(other.is_aligned_){};

TensorImpl& TensorImpl::operator=(TensorImpl&& other) {
  if (&other != this) {
    props_ = std::move(other.props_);
    indexer_ = std::move(other.indexer_);
    data_buffer_ = std::move(other.data_buffer_);
    is_aligned_ = other.is_aligned_;
  }
  return *this;
}

const void* TensorImpl::host_data() const {
  return data();
}

bool TensorImpl::is_same_backend(backend_t backend, index_t) const {
  return backend_t::SIMD == backend;
}

void TensorImpl::check_dev_type_compatibility(const DeviceTensor& other) const {
  if (props_.device != other.properties().device)
    THROW_ERROR("SIMD Tensor op failed: incompatible devices ", props_.device,
                " != ", other.properties().device);

  if (props_.dtype != other.properties().dtype)
    THROW_ERROR("SIMD Tensor op failed: incompatible data types!");
}

bool TensorImpl::is_shape_aligned() const {
  return is_aligned_;
}

bool TensorImpl::is_shape_aligned_impl() const {
  for (auto dim : props_.shape.dims()) {
    if (dim % hwy::VectorBytes() != 0) {
      return false;
    }
  }
  return true;
}

std::shared_ptr<DeviceTensor> TensorImpl::squeeze(index_t dim) const {
  // the mempool and factory is not needed - buffer will be tha same
  auto new_props = props_;
  if (new_props.shape.dim(dim) != 1) {
    THROW_ERROR("Tensor::squeeze failed the given dim is not 1, shape=", new_props.shape,
                ", dim=", dim);
  } else {
    new_props.shape.remove_dim(dim);
  }
  return std::make_shared<TensorImpl>(new_props, data_buffer_);
}

std::shared_ptr<DeviceTensor> TensorImpl::unsqueeze(index_t dim) const {
  // the mempool and factory is not needed - buffer will be tha same
  auto new_props = props_;
  new_props.shape.add_dim(dim, 1);
  return std::make_shared<TensorImpl>(new_props, data_buffer_);
}

float32_t TensorImpl::get_float32_at(const coords_t& coords) const {
  return get<float32_t>(coords);
}
float64_t TensorImpl::get_float64_at(const coords_t& coords) const {
  return get<float64_t>(coords);
}
int32_t TensorImpl::get_int32_at(const coords_t& coords) const {
  return get<int32_t>(coords);
}
int8_t TensorImpl::get_int8_at(const coords_t& coords) const {
  return get<int8_t>(coords);
}

template <typename Op>
void broadcast_arithmetic_call(TensorImpl& self, const TensorImpl& other) {
  self.check_dev_type_compatibility(other);
  bool aligned = self.is_aligned_ && other.is_aligned_;

  DISPATCH_TYPE(self.props_.dtype, [&]() {
    auto* x_array = self.mutable_data_ptr<scalar_t>();
    const auto* y_array = other.const_data_ptr<scalar_t>();
    if (other.properties().shape.numel() == 1) {
      // scalar broadcast
      auto scalar = static_cast<float32_t>(y_array[0]);
      parallel_for<scalar_t>(0, self.props_.shape.numel(), [&](auto begin, auto end) {
        Op::apply(x_array + begin, scalar, end - begin, aligned);
      });
    } else if (self.props_.shape == other.props_.shape) {
      // same shape
      parallel_for<scalar_t>(0, self.props_.shape.numel(), [&](auto begin, auto end) {
        Op::apply(x_array + begin, y_array + begin, end - begin, aligned);
      });
    } else {
      auto target_shape = make_broadcast_shape(self.props_.shape, other.props_.shape);
      if (!target_shape || self.props_.shape.rang() < target_shape->rang()) {
        THROW_ERROR("Tensors shapes are not broadcastable ", self.props_.shape, " ~ ",
                    other.props_.shape);
      }
      auto self_strides = make_broadcast_strides(*target_shape, self.props_.shape);
      auto other_strides = make_broadcast_strides(*target_shape, other.props_.shape);
      if (self_strides.back() == other_strides.back() && self_strides.back() == 0) {
        // special case - last dim is column vector
        target_shape->remove_dim(target_shape->rang() - 1);
        self_strides.pop_back();
        other_strides.pop_back();
        (row_loop<Op, scalar_t>(*target_shape, x_array, self_strides, y_array, other_strides,
                                aligned));
      } else {
        (row_loop<Op, scalar_t>(*target_shape, x_array, self_strides, y_array, other_strides,
                                aligned));
      }
    }
  });
}

void TensorImpl::add(const DeviceTensor& other) {
  const auto& other_tensor = static_cast<const TensorImpl&>(other);
  broadcast_arithmetic_call<detail::add>(*this, other_tensor);
}

void TensorImpl::add(float32_t scalar) {
  DISPATCH_TYPE(props_.dtype, [&]() {
    auto* x_array = mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, props_.shape.numel(), [&](auto begin, auto end) {
      detail::add::apply(x_array + begin, scalar, end - begin, is_aligned_);
    });
  });
}

void TensorImpl::sub(const DeviceTensor& other) {
  const auto& other_tensor = static_cast<const TensorImpl&>(other);
  broadcast_arithmetic_call<detail::sub>(*this, other_tensor);
}

void TensorImpl::sub(float32_t scalar) {
  DISPATCH_TYPE(props_.dtype, [&]() {
    auto* x_array = mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, props_.shape.numel(), [&](auto begin, auto end) {
      detail::sub::apply(x_array + begin, scalar, end - begin, is_aligned_);
    });
  });
}

void TensorImpl::mul(const DeviceTensor& other) {
  const auto& other_tensor = static_cast<const TensorImpl&>(other);
  broadcast_arithmetic_call<detail::mul>(*this, other_tensor);
}

void TensorImpl::mul(float32_t scalar) {
  DISPATCH_TYPE(props_.dtype, [&]() {
    auto* x_array = mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, props_.shape.numel(), [&](auto begin, auto end) {
      detail::mul::apply(x_array + begin, scalar, end - begin, is_aligned_);
    });
  });
};

void TensorImpl::div(const DeviceTensor& other) {
  const auto& other_tensor = static_cast<const TensorImpl&>(other);
  broadcast_arithmetic_call<detail::div>(*this, other_tensor);
}

void TensorImpl::div(float32_t scalar) {
  DISPATCH_TYPE(props_.dtype, [&]() {
    auto* x_array = mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, props_.shape.numel(), [&](auto begin, auto end) {
      detail::div::apply(x_array + begin, scalar, end - begin, is_aligned_);
    });
  });
}

namespace {
template <typename Op>
std::shared_ptr<DeviceTensor> tensor_logic_call(const TensorImpl& self, float32_t scalar) {
  auto new_tensor = TensorFactory::instance().empty(self.properties());
  DISPATCH_TYPE(self.properties().dtype, [&]() {
    const auto* x_array = self.const_data_ptr<scalar_t>();
    auto* z_array = new_tensor->mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, self.properties().shape.numel(), [&](auto begin, auto end) {
      Op::apply(x_array + begin, scalar, z_array + begin, end - begin, self.is_shape_aligned());
    });
  });
  return new_tensor;
}
}  // namespace

std::shared_ptr<DeviceTensor> TensorImpl::gt(float32_t scalar) const {
  return tensor_logic_call<detail::gt>(*this, scalar);
}

std::shared_ptr<DeviceTensor> TensorImpl::ge(float32_t scalar) const {
  return tensor_logic_call<detail::ge>(*this, scalar);
}

std::shared_ptr<DeviceTensor> TensorImpl::lt(float32_t scalar) const {
  return tensor_logic_call<detail::lt>(*this, scalar);
}

std::shared_ptr<DeviceTensor> TensorImpl::le(float32_t scalar) const {
  return tensor_logic_call<detail::le>(*this, scalar);
}

void TensorImpl::neg() {
  DISPATCH_TYPE(props_.dtype, [&]() {
    auto* x_array = mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, props_.shape.numel(), [&](auto begin, auto end) {
      detail::neg::apply<scalar_t>(x_array + begin, end - begin, is_aligned_);
    });
  });
}

void TensorImpl::exp() {
  DISPATCH_FLOATING_TYPES(props_.dtype, [&]() {
    auto* x_array = mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, props_.shape.numel(), [&](auto begin, auto end) {
      detail::exp::apply<scalar_t>(x_array + begin, end - begin, is_aligned_);
    });
  });
}

void TensorImpl::sqrt() {
  DISPATCH_FLOATING_TYPES(props_.dtype, [&]() {
    auto* x_array = mutable_data_ptr<scalar_t>();
    parallel_for<scalar_t>(0, props_.shape.numel(), [&](auto begin, auto end) {
      detail::sqrt::apply<scalar_t>(x_array + begin, end - begin, is_aligned_);
    });
  });
}

std::shared_ptr<DeviceTensor> TensorImpl::sum_div(index_t divider) const {
  auto new_tensor = TensorFactory::instance().empty(
      TensorProperties{.shape = {1, 1}, .device = props_.device, .dtype = props_.dtype});
  DISPATCH_TYPE(props_.dtype, [&]() {
    const auto* x_array = const_data_ptr<scalar_t>();
    auto* z_array = new_tensor->mutable_data_ptr<scalar_t>();
    z_array[0] = parallel_reduce<scalar_t>(
        0, props_.shape.numel(), scalar_t{0},
        [&](auto begin, auto end, auto init) {
          scalar_t partial_sum = init;
          detail::sum::apply<scalar_t>(x_array + begin, partial_sum, end - begin, is_aligned_);
          return partial_sum;
        },
        [](scalar_t a, scalar_t b) { return a + b; });
    z_array[0] /= divider;
  });
  return new_tensor;
}

std::shared_ptr<DeviceTensor> TensorImpl::sum_dim0() const {
  Shape new_shape(props_.shape);
  new_shape.remove_dim(0);
  auto new_tensor =
      Tensor::zero({.shape = new_shape, .device = props_.device, .dtype = props_.dtype});

  auto batch_size = new_tensor.properties().shape.numel();
  DISPATCH_TYPE(props_.dtype, [&]() {
    const auto* x_array = const_data_ptr<scalar_t>();
    auto* z_array = new_tensor.mutable_data_ptr<scalar_t>();
    coords_t coords(props_.shape.rang(), 0);
    // loop over batches
    for (index_t i = 0; i < props_.shape.dim(0); ++i) {
      coords[0] = i;
      auto dim0_offset = indexer_.idxravel(coords);
      // sum all batch elemets to result(accumulator)
      parallel_for<scalar_t>(0, batch_size, [&](auto begin, auto end) {
        detail::add::apply(z_array + begin, x_array + dim0_offset + begin, end - begin,
                           is_aligned_);
      });
    }
  });
  return new_tensor.impl();
}

std::shared_ptr<DeviceTensor> TensorImpl::max() const {
  auto new_tensor = TensorFactory::instance().empty(
      TensorProperties{.shape = {1, 1}, .device = props_.device, .dtype = props_.dtype});
  DISPATCH_TYPE(props_.dtype, [&]() {
    const auto* x_array = const_data_ptr<scalar_t>();
    auto* z_array = new_tensor->mutable_data_ptr<scalar_t>();
    z_array[0] = parallel_reduce<scalar_t>(
        0, props_.shape.numel(), scalar_t{0},
        [&](auto begin, auto end, auto init) {
          scalar_t partial_sum = init;
          detail::max::apply<scalar_t>(x_array + begin, partial_sum, end - begin, is_aligned_);
          return partial_sum;
        },
        [](scalar_t a, scalar_t b) { return a + b; });
  });
  return new_tensor;
}

std::shared_ptr<DeviceTensor> TensorImpl::sum() const {
  return sum_div(1);
}

std::shared_ptr<DeviceTensor> TensorImpl::mean() const {
  return sum_div(props_.shape.numel());
}

}  // namespace adept::cpu
#endif  // HWY_ONCE