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

src/backends/cpu/tensor_matmul.cpp
246 строк8 KB
kolkir
Add matmul tests
26 мар 2025, 00:49
26 мар 2025, 00:4927a44dc
#include <adept/backends/cpu/tensor_factory.hpp>
#include <adept/backends/cpu/tensorimpl.hpp>
#include <adept/irange.hpp>
#include <adept/types_dispatch.hpp>
#include <algorithm>
#include <functional>
#include <numeric>
#include <utility>
#include "adept/shape.hpp"
#include "adept/types.hpp"
#include "matmul.hpp"

namespace adept::cpu {

namespace {
std::shared_ptr<DeviceTensor> matmul_2d(const DeviceTensor& a,
                                        const Shape& a_shape,
                                        bool squeeze_prepend,
                                        const DeviceTensor& b,
                                        const Shape& b_shape,
                                        bool squeeze_append) {
  if (a_shape.dim(1) != b_shape.dim(0)) {
    THROW_ERROR("matmul incompatible dimenstions ", a.properties().shape, " ~ ",
                b.properties().shape);
  }

  Shape c_shape({a_shape.dim(0), b_shape.dim(1)});
  Shape real_c_shape = c_shape;

  if (squeeze_prepend) {
    real_c_shape = c_shape;
    real_c_shape.remove_dim(0);
  }
  if (squeeze_append) {
    real_c_shape = c_shape;
    real_c_shape.remove_dim(1);
  }

  auto new_tensor = Tensor::zero(
      {.shape = real_c_shape, .device = a.properties().device, .dtype = a.properties().dtype});

  DISPATCH_TYPE(a.properties().dtype, [&]() {
    const auto* a_data = a.const_data_ptr<scalar_t>();
    const auto* b_data = b.const_data_ptr<scalar_t>();
    auto* c_data = new_tensor.mutable_data_ptr<scalar_t>();
    adept::matmul(a_shape, a_data, b_shape, b_data, c_data);
  });
  return new_tensor.impl();
}

std::vector<index_t> get_batch_dims(const Shape& s) {
  std::vector<index_t> batch_dims;
  if (s.rang() > 2) {
    for (auto i : irange(s.rang() - 2)) {
      batch_dims.push_back(s.dim(i));
    }
  }
  return batch_dims;
}

std::pair<index_t, index_t> get_mat_dims(const Shape& s) {
  return {*std::next(s.dims().rbegin()), s.dims().back()};
}

std::shared_ptr<DeviceTensor> matmul_batched(const DeviceTensor& a,
                                             const Shape& a_shape,
                                             bool squeeze_prepend,
                                             const DeviceTensor& b,
                                             const Shape& b_shape,
                                             bool squeeze_append) {
  auto a_batch_dims = get_batch_dims(a_shape);
  index_t a_batches =
      std::reduce(a_batch_dims.begin(), a_batch_dims.end(), 1, std::multiplies<index_t>());

  auto b_batch_dims = get_batch_dims(b_shape);
  index_t b_batches =
      std::reduce(b_batch_dims.begin(), b_batch_dims.end(), 1, std::multiplies<index_t>());

  if (a_batch_dims.size() == b_batch_dims.size() && a_batches != b_batches) {
    THROW_ERROR("matmul wrong batch dimenstions ", a.properties().shape, " ~ ",
                b.properties().shape);
  } else if (a_batch_dims.size() > b_batch_dims.size()) {
    auto i = a_batch_dims.size() - 1;
    auto j = b_batch_dims.size() - 1;
    for ([[maybe_unused]] auto s : irange(b_batch_dims.size())) {
      if (a_batch_dims[i] != b_batch_dims[j]) {
        THROW_ERROR("matmul wrong batch dimenstions ", a.properties().shape, " ~ ",
                    b.properties().shape);
        --i;
        --j;
      }
    }
  } else {
    auto i = a_batch_dims.size() - 1;
    auto j = b_batch_dims.size() - 1;
    for ([[maybe_unused]] auto s : irange(a_batch_dims.size())) {
      if (a_batch_dims[i] != b_batch_dims[j]) {
        THROW_ERROR("matmul wrong batch dimenstions ", a.properties().shape, " ~ ",
                    b.properties().shape);
        --i;
        --j;
      }
    }
  }
  auto [min_batch_dim, max_batch_dim] = std::minmax(a_batches, b_batches);

  index_t a_m, a_k, b_k, b_n;
  std::tie(a_m, a_k) = get_mat_dims(a_shape);
  std::tie(b_k, b_n) = get_mat_dims(b_shape);

  if (a_k != b_k) {
    THROW_ERROR("matmul incompatible dimenstions ", a.properties().shape, " ~ ",
                b.properties().shape);
  }

  Shape a_mat_shape = {a_m, a_k};
  Shape b_mat_shape = {b_k, b_n};
  Shape c_shape = {max_batch_dim, a_m, b_n};
  coords_t real_new_dims;
  if (a_batches == max_batch_dim) {
    real_new_dims = a_batch_dims;
  } else {
    real_new_dims = b_batch_dims;
  }
  if (!squeeze_prepend || a_m != 1) {
    real_new_dims.push_back(a_m);
  }
  if (!squeeze_append || b_n != 1) {
    real_new_dims.push_back(b_n);
  }

  auto new_tensor = Tensor::zero(
      {.shape = real_new_dims, .device = a.properties().device, .dtype = a.properties().dtype});

  coords_t coords = {0, 0, 0};
  Indexer c_indexer(&c_shape);

  DISPATCH_TYPE(a.properties().dtype, [&]() {
    const auto* a_data = a.const_data_ptr<scalar_t>();
    const auto* b_data = b.const_data_ptr<scalar_t>();
    auto* c_data = new_tensor.mutable_data_ptr<scalar_t>();

    if (a_batch_dims.size() == b_batch_dims.size() && a_batches == b_batches) {
      // all same dims
      Shape a_batch_shape = {max_batch_dim, a_m, a_k};
      Shape b_batch_shape = {max_batch_dim, b_k, b_n};
      Indexer a_indexer(&a_batch_shape);
      Indexer b_indexer(&b_batch_shape);
      for (auto i : irange(max_batch_dim)) {
        coords[0] = i;
        auto a_offset = a_indexer.idxravel(coords);
        auto b_offset = b_indexer.idxravel(coords);
        auto c_offset = c_indexer.idxravel(coords);
        adept::matmul(a_mat_shape, a_data + a_offset, b_mat_shape, b_data + b_offset,
                      c_data + c_offset);
      }
    } else if (a_batch_dims.size() > b_batch_dims.size()) {
      Shape a_batch_shape = {max_batch_dim, a_m, a_k};
      Shape b_batch_shape = {min_batch_dim, b_k, b_n};
      Indexer a_indexer(&a_batch_shape);
      Indexer b_indexer(&b_batch_shape);
      for (auto i : irange(max_batch_dim / min_batch_dim)) {
        for (auto j : irange(min_batch_dim)) {
          coords[0] = i * min_batch_dim + j;
          auto a_offset = a_indexer.idxravel(coords);
          auto c_offset = c_indexer.idxravel(coords);

          coords[0] = j;
          auto b_offset = b_indexer.idxravel(coords);

          adept::matmul(a_mat_shape, a_data + a_offset, b_mat_shape, b_data + b_offset,
                        c_data + c_offset);
        }
      }
    } else {
      Shape a_batch_shape = {min_batch_dim, a_m, a_k};
      Shape b_batch_shape = {max_batch_dim, b_k, b_n};
      Indexer a_indexer(&a_batch_shape);
      Indexer b_indexer(&b_batch_shape);
      for (auto i : irange(max_batch_dim / min_batch_dim)) {
        for (auto j : irange(min_batch_dim)) {
          coords[0] = i * min_batch_dim + j;
          auto b_offset = b_indexer.idxravel(coords);
          auto c_offset = c_indexer.idxravel(coords);

          coords[0] = j;
          auto a_offset = a_indexer.idxravel(coords);

          adept::matmul(a_mat_shape, a_data + a_offset, b_mat_shape, b_data + b_offset,
                        c_data + c_offset);
        }
      }
    }
  });
  return new_tensor.impl();
}

}  // namespace

/* Broadcasting rules
If both arguments are 2-D they are multiplied like conventional matrices.

If either argument is N-D, N > 2, it is treated as a stack of matrices residing in the last two
indexes and broadcast accordingly.

If the first argument is 1-D, it is promoted to a matrix by prepending a 1 to its dimensions.
After matrix multiplication the prepended 1 is removed.

If the second argument is 1-D, it is promoted to a matrix by appending a 1 to its dimensions.
After matrix multiplication the appended 1 is removed.
*/
std::shared_ptr<DeviceTensor> TensorImpl::matmul(const DeviceTensor& other) const {
  bool self_dim1 = props_.shape.rang() == 1;
  bool other_dim1 = other.properties().shape.rang() == 1;

  bool self_batched = props_.shape.rang() > 2;
  bool other_batched = other.properties().shape.rang() > 2;

  bool has_batch_dim = self_batched || other_batched;

  bool squeeze_prepend = false;
  bool squeeze_append = false;

  Shape a_shape = props_.shape;
  Shape b_shape = other.properties().shape;

  // make matrices from 1D vectors
  if (self_dim1) {
    a_shape = Shape({1, a_shape.dim(0)});
    squeeze_prepend = true;
  }

  if (other_dim1) {
    b_shape = Shape({b_shape.dim(0), 1});
    squeeze_append = true;
  }

  if (has_batch_dim) {
    return matmul_batched(*this, a_shape, squeeze_prepend, other, b_shape, squeeze_append);

  } else {
    return matmul_2d(*this, a_shape, squeeze_prepend, other, b_shape, squeeze_append);
  }
}

}  // namespace adept::cpu