src/index_utils.cpp - with_cpu - vuron/adept

src/index_utils.cpp
115 строк3 KB
kolkir
Imlement dataset and dataloader python interfaces
05 мар 2025, 22:56
05 мар 2025, 22:56c63c6ed
#include <adept/index_utils.hpp>
#include <adept/irange.hpp>
#include <adept/print.hpp>

namespace adept {

Indexer::Indexer(const Shape* shape) : shape_{shape}, numel_(shape->numel()) {}

index_t Indexer::idxravel(const std::vector<index_t>& coords) const {
  if (coords.size() != shape_->rang()) {
    THROW_ERROR("idxravel failed: incompatible sizes ", coords.size(), " != ", shape_->rang());
  }
  index_t idx = 0;
  auto c = numel_;

  for (auto i : irange(shape_->rang())) {
    c /= shape_->dim(i);
    idx += coords[i] * c;
  }
  return idx;
}

std::vector<index_t> Indexer::idxunravel(index_t idx) const {
  if (idx >= numel_) {
    THROW_ERROR("idxunravel failed: index is too big ", idx, " >= ", numel_);
  }
  auto c = numel_;
  std::vector<index_t> coords;
  coords.reserve(shape_->rang());
  for (auto dim : shape_->dims()) {
    c /= dim;
    auto coord = idx / c;
    idx -= coord * c;
    coords.push_back(coord);
  }
  return coords;
}

std::optional<Shape> make_broadcast_shape(const Shape& a, const Shape& b) {
  auto ndims = std::max(a.rang(), b.rang());
  auto a_diff = ndims - a.rang();
  auto b_diff = ndims - b.rang();
  Shape new_shape(ndims);

  for (int64_t i = ndims - 1; i >= 0; --i) {
    index_t a_dim = 1;
    if (ndims > a.rang()) {
      auto a_i = i - static_cast<int64_t>(a_diff);
      if (a_i >= 0)
        a_dim = a.dim(a_i);
    } else {
      a_dim = a.dim(i);
    }

    index_t b_dim = 1;
    if (ndims > b.rang()) {
      auto b_i = i - static_cast<int64_t>(b_diff);
      if (b_i >= 0)
        b_dim = b.dim(b_i);
    } else {
      b_dim = b.dim(i);
    }

    if (a_dim != 1 && b_dim != 1 && a_dim != b_dim) {
      return std::nullopt;
    } else {
      new_shape.update_dim(i, std::max(a_dim, b_dim));
    }
  }
  return new_shape;
}

coords_t make_broadcast_strides(const Shape& target_shape, const Shape& original_shape) {
  auto ndims = target_shape.rang();
  auto diff = ndims - original_shape.rang();
  coords_t strides(ndims, 0);

  index_t stride = 1;
  for (int64_t i = ndims - 1; i >= 0; --i) {
    index_t dim = 1;
    if (ndims > original_shape.rang()) {
      auto a_i = i - static_cast<int64_t>(diff);
      if (a_i >= 0)
        dim = original_shape.dim(a_i);
    } else {
      dim = original_shape.dim(i);
    }

    if (dim != 1 && target_shape.dim(i) != dim) {
      THROW_ERROR("make_broadcast_strides filed: shapes are not broadcastable ", target_shape,
                  " ~ ", original_shape);
    } else if (dim == 1) {
      strides[i] = 0;
    } else {
      strides[i] = stride;
      stride *= target_shape.dim(i);
    }
  }

  return strides;
}

coords_t make_continuous_strides(const Shape& shape) {
  auto ndims = shape.rang();
  coords_t strides(ndims, 0);

  index_t stride = 1;
  for (int64_t i = ndims - 1; i >= 0; --i) {
    strides[i] = stride;
    stride *= shape.dim(i);
  }
  return strides;
}

}  // namespace adept