vuron/adept: Платформа для обучения и вывода нейронных сетей, с API приближенным к PyTorch.

src/backends/cpu/stack.cpp

46 строк1 KB

kolkir

revert compiler compatibility to gcc12

03 мар 2025, 23:52

03 мар 2025, 23:5226a9274

#include <adept/backends/cpu/stack.hpp>
#include <adept/backends/cpu/tensor_factory.hpp>
#include <adept/backends/cpu/tensorimpl.hpp>
#include <adept/exception.hpp>
#include <adept/print.hpp>
#include <adept/types_dispatch.hpp>

#include <hwy/highway.h>

namespace adept::cpu {

Tensor stack_tensors(const std::vector<Tensor>& tensors) {
  if (!tensors.empty()) {
    auto props = tensors[0].properties();
    for (index_t i = 0; i < tensors.size(); ++i) {
      if (tensors[i].impl()->properties().device != props.device ||
          tensors[i].impl()->properties().dtype != props.dtype ||
          tensors[i].impl()->properties().shape != props.shape) {
        THROW_ERROR("stack failed: incompatible tensor properties ", props, " ~ ",
                    tensors[i].impl()->properties());
      }
    }

    // make new shape with extra dimsion
    auto result_props = tensors[0].properties();
    result_props.shape.add_dim(0, tensors.size());

    // make result tensot
    auto result = Tensor::empty(result_props);

    // copy data
    auto tensor_size = tensors[0].properties().shape.numel();
    DISPATCH_TYPE(result_props.dtype, [&]() {
      auto buffer_ptr = result.mutable_data_ptr<scalar_t>();
      for (auto& t : tensors) {
        hwy::CopyBytes(t.const_data_ptr<scalar_t>(), buffer_ptr, tensor_size * sizeof(scalar_t));
        buffer_ptr += tensor_size;
      }
    });

    return result;
  }
  return Tensor(nullptr);
}

}  // namespace adept::cpu