tests/index_utils_tests.cpp - with_cpu - vuron/adept

tests/index_utils_tests.cpp
263 строки6 KB
kolkir
Refactor loops
08 фев 2025, 13:56
08 фев 2025, 13:56b73a638
#include <adept/index_utils.hpp>
#include <adept/irange.hpp>

#include "catch.hpp"

using namespace adept;

TEST_CASE("Make broadcast shape n", "[index utils]") {
  Shape a{3, 3, 3, 3};
  Shape b{1, 3, 3, 3};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());
  REQUIRE(a == c);
}

TEST_CASE("Make broadcast shape c", "[index utils]") {
  Shape a{3, 3, 3, 3};
  Shape b{3, 1, 3, 3};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());
  REQUIRE(a == c);
}

TEST_CASE("Make broadcast shape h", "[index utils]") {
  Shape a{3, 3, 3, 3};
  Shape b{3, 3, 1, 3};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());
  REQUIRE(a == c);
}

TEST_CASE("Make broadcast shape w", "[index utils]") {
  Shape a{3, 3, 3, 3};
  Shape b{3, 3, 3, 1};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());
  REQUIRE(a == c);
}

TEST_CASE("Make broadcast shape 2d h", "[index utils]") {
  Shape a{3, 3};
  Shape b{1, 3};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());
  REQUIRE(a == c);
}

TEST_CASE("Make broadcast shape 2d w", "[index utils]") {
  Shape a{3, 3};
  Shape b{3, 1};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());
  REQUIRE(a == c);
}

TEST_CASE("Make broadcast shape mixed h", "[index utils]") {
  Shape a{3, 3, 3, 3};
  Shape b{1, 3, 1, 3};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());
  REQUIRE(a == c);
}

TEST_CASE("Make broadcast shape mixed w", "[index utils]") {
  Shape a{3, 3, 3, 3};
  Shape b{1, 3, 3, 1};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());
  REQUIRE(a == c);
}

TEST_CASE("Make broadcast shape fail c", "[index utils]") {
  Shape a{3, 3, 3, 3};
  Shape b{1, 2, 3, 1};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(!c.has_value());
}

TEST_CASE("Make broadcast shape fail n", "[index utils]") {
  Shape a{3, 3, 3, 3};
  Shape b{2, 1, 3, 1};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(!c.has_value());
}

TEST_CASE("Make broadcast shape fail h", "[index utils]") {
  Shape a{3, 3};
  Shape b{2, 1};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(!c.has_value());
}

TEST_CASE("Make broadcast shape promote h", "[index utils]") {
  Shape a{3, 3};
  Shape b{3};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());
  REQUIRE(a == c);
}

TEST_CASE("Make broadcast shape promote c", "[index utils]") {
  Shape a{3, 3, 3};
  Shape b{3, 1};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());
  REQUIRE(a == c);
}

TEST_CASE("Make broadcast shape mixed both", "[index utils]") {
  Shape a{3, 1, 3};
  Shape b{3, 1};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());
  REQUIRE(Shape{3, 3, 3} == c);
}

TEST_CASE("Make broadcast strides mixed both", "[index utils]") {
  Shape a{3, 1, 3};
  Shape b{3, 1};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());

  auto strides = make_broadcast_strides(c.value(), a);
  REQUIRE(strides.size() == a.rang());
  REQUIRE(strides[0] == 3);
  REQUIRE(strides[1] == 0);
  REQUIRE(strides[2] == 1);

  strides = make_broadcast_strides(c.value(), b);
  REQUIRE(strides.size() == a.rang());
  REQUIRE(strides[0] == 0);
  REQUIRE(strides[1] == 1);
  REQUIRE(strides[2] == 0);
}

TEST_CASE("Make broadcast strides w", "[index utils]") {
  Shape a{3, 3, 3, 3};
  Shape b{3, 1};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());

  auto strides = make_broadcast_strides(c.value(), a);
  REQUIRE(strides.size() == a.rang());
  REQUIRE(strides[0] == 27);
  REQUIRE(strides[1] == 9);
  REQUIRE(strides[2] == 3);
  REQUIRE(strides[3] == 1);

  strides = make_broadcast_strides(c.value(), b);
  REQUIRE(strides.size() == a.rang());
  REQUIRE(strides[0] == 0);
  REQUIRE(strides[1] == 0);
  REQUIRE(strides[2] == 1);
  REQUIRE(strides[3] == 0);
}

TEST_CASE("Strides loop to scalar", "[index utils]") {
  Shape a{3, 3, 3};
  Shape b{3, 1};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());
  auto target_shape = *c;

  auto a_strides = make_broadcast_strides(target_shape, a);
  auto b_strides = make_broadcast_strides(target_shape, b);

  auto row_based_dims = target_shape.dims();
  row_based_dims.resize(row_based_dims.size() - 1);

  Shape row_based_shape(row_based_dims);
  Indexer indexer(&row_based_shape);
  auto numel = row_based_shape.numel();
  auto row_count = 0;
  std::vector<index_t> a_row_start = {0, 3, 6, 9, 12, 15, 18, 21, 24};
  std::vector<index_t> b_row_start = {0, 1, 2, 0, 1, 2, 0, 1, 2};
  for (auto i : irange(numel)) {
    auto coords = indexer.idxunravel(i);

    index_t a_row = 0;
    size_t j = 0;
    for (; j < coords.size(); ++j) {
      a_row += coords[j] * a_strides[j];
    }
    REQUIRE(a_row == a_row_start[i]);
    REQUIRE(j == a_strides.size() - 1);
    REQUIRE(a_strides[j] != 0);

    index_t b_row = 0;
    j = 0;
    for (; j < coords.size(); ++j) {
      b_row += coords[j] * b_strides[j];
    }
    REQUIRE(b_row == b_row_start[i]);
    REQUIRE(j == b_strides.size() - 1);
    REQUIRE(b_strides[j] == 0);  // scalar

    ++row_count;
  }
  REQUIRE(row_count == 9);
}

TEST_CASE("Strides loop to vec", "[index utils]") {
  Shape a{3, 3, 3};
  Shape b{1, 3};

  auto c = make_broadcast_shape(a, b);
  REQUIRE(c.has_value());
  auto target_shape = *c;

  auto a_strides = make_broadcast_strides(target_shape, a);
  auto b_strides = make_broadcast_strides(target_shape, b);

  auto row_based_dims = target_shape.dims();
  row_based_dims.resize(row_based_dims.size() - 1);

  Shape row_based_shape(row_based_dims);
  Indexer indexer(&row_based_shape);
  auto numel = row_based_shape.numel();
  auto row_count = 0;
  std::vector<index_t> a_row_start = {0, 3, 6, 9, 12, 15, 18, 21, 24};
  std::vector<index_t> b_row_start = {0, 0, 0, 0, 0, 0, 0, 0, 0};
  for (auto i : irange(numel)) {
    auto coords = indexer.idxunravel(i);

    index_t a_row = 0;
    size_t j = 0;
    for (; j < coords.size(); ++j) {
      a_row += coords[j] * a_strides[j];
    }
    REQUIRE(a_row == a_row_start[i]);
    REQUIRE(j == a_strides.size() - 1);
    REQUIRE(a_strides[j] != 0);

    index_t b_row = 0;
    j = 0;
    for (; j < coords.size(); ++j) {
      b_row += coords[j] * b_strides[j];
    }
    REQUIRE(b_row == b_row_start[i]);
    REQUIRE(j == b_strides.size() - 1);
    REQUIRE(b_strides[j] != 0);  // vec

    ++row_count;
  }
  REQUIRE(row_count == 9);
}