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

src/backends/cpu/tensor_mat_tr.cpp
62 строки2 KB
kolkir
Fix training issues
06 мар 2025, 23:54
06 мар 2025, 23:54efe26e8
#include <adept/backends/cpu/tensor_factory.hpp>
#include <adept/backends/cpu/tensorimpl.hpp>
#include <adept/irange.hpp>
#include <adept/threading.hpp>
#include <adept/types_dispatch.hpp>

#include "transpose.hpp"

namespace adept::cpu {

std::shared_ptr<DeviceTensor> TensorImpl::transpose2d() const {
  if (props_.shape.rang() > 3) {
    THROW_ERROR("transpose op called for the tensor ndims > 3: ", props_.shape);
  }

  if (props_.shape.rang() == 1) {
    return TensorFactory::instance().clone(*this);
  }

  // if we have ndims == 3 the first dim will be used as batch dimension

  index_t batch_dim = 0;
  if (props_.shape.rang() == 3) {
    batch_dim = props_.shape.dim(0);
  }

  Shape src_batched_shape = props_.shape;
  Shape new_shape = props_.shape;
  Shape dst_batched_shape;
  if (batch_dim == 0) {
    new_shape.swap_dims(0, 1);
    batch_dim = 1;
    src_batched_shape = Shape{batch_dim, props_.shape.dim(0), props_.shape.dim(1)};
    dst_batched_shape = Shape{batch_dim, new_shape.dim(0), new_shape.dim(1)};
  } else {
    new_shape.swap_dims(1, 2);
    dst_batched_shape = new_shape;
  }

  auto new_tensor = TensorFactory::instance().empty(
      TensorProperties{.shape = new_shape, .device = props_.device, .dtype = props_.dtype});

  Indexer src_indexer(&src_batched_shape);
  Indexer dst_indexer(&dst_batched_shape);

  coords_t batch_pos = {0, 0, 0};
  DISPATCH_TYPE(props_.dtype, [&]() {
    parallel_for<scalar_t>(0, batch_dim, batch_dim, [&](auto begin, auto end) {
      for (auto n : irange(begin, end)) {
        batch_pos[0] = n;
        auto src_batch_offset = src_indexer.idxravel(batch_pos);
        const auto* src = const_data_ptr<scalar_t>() + src_batch_offset;
        auto dst_batch_offset = dst_indexer.idxravel(batch_pos);
        auto* dst = new_tensor->mutable_data_ptr<scalar_t>() + dst_batch_offset;
        transpose_mat(src_batched_shape.dim(1), src_batched_shape.dim(2), src, dst);
      }
    });
  });
  return new_tensor;
}

}  // namespace adept::cpu