src/backends/cpu/loop.h - master - vuron/adept

src/backends/cpu/loop.h
66 строк2 KB
kolkir
Refactor simd algos to use map reduce
16 мар 2025, 23:09
16 мар 2025, 23:095f536e1
#pragma once

#include <adept/index_utils.hpp>
#include <adept/shape.hpp>
#include <adept/threading.hpp>
#include <adept/types.hpp>

namespace adept {

template <typename Op, typename DataType>
void row_loop(const Shape& target_shape,
              DataType* self_array,
              const coords_t& self_strides,
              const DataType* other_array,
              const coords_t& other_strides,
              bool aligned) {
  Shape row_based_shape(target_shape);
  row_based_shape.remove_dim(row_based_shape.rang() - 1);

  Indexer indexer(&row_based_shape);
  auto rows = row_based_shape.numel();
  parallel_for<DataType>(0, rows, target_shape.dims().back(), [&](auto begin, auto end) {
    for (index_t i = 0; i < rows; ++i) {
      auto coords = indexer.idxunravel(i);

      index_t self_row_start = 0;
      size_t j = 0;
      for (; j < coords.size(); ++j) {
        self_row_start += coords[j] * self_strides[j];
      }
      bool self_vec = true;
      if (self_strides[j] == 0) {
        self_vec = false;
      }

      index_t other_row_start = 0;
      j = 0;
      for (; j < coords.size(); ++j) {
        other_row_start += coords[j] * other_strides[j];
      }
      bool other_vec = true;
      if (other_strides[j] == 0) {
        other_vec = false;
      }

      if (self_vec && other_vec) {
        Op::apply(self_array + self_row_start, other_array + other_row_start,
                  target_shape.dims().back(), aligned);
      } else if (self_vec && !other_vec) {
        Op::apply(self_array + self_row_start, (other_array + other_row_start)[0],
                  target_shape.dims().back(), aligned);
      }
      // This branch makes sense for the case with new tensor output
      // else if (!self_vec && other_vec) {
      //   Func::call(other_array + other_row_start, (self_array + self_row_start)[0],
      //              target_shape.dims.back(), aligned);
      // }
      else {
        // shouldn't happened - self should have fully defined shape
        THROW_ERROR("Broadcast failed - self tensor has not fully defined shape");
      }
    }
  });
}

}  // namespace adept