gather_test.cpp

#include <iostream>
#include <algorithm>
#include <random>
#include <numeric>

constexpr int Nx = 512;

typedef double Dtype;
// typedef int Dtype;
// typedef float Dtype;
// typedef int64_t Dtype;

void gather_test_ran(const int* idx,
                     const Dtype* val,
                     const int num) {
  Dtype sum = 0.0;
  for (int i = 0; i < num; i++) {
    const auto v = val[idx[i]];
    sum += v;
  }
  std::cout << sum << std::endl;
}

void gather_test_uni(const Dtype* val,
                     const int num) {
  Dtype sum = 0.0;
  for (int i = 0; i < num; i++) {
    const auto v = val[i * 100];
    sum += v;
  }
  std::cout << sum << std::endl;
}

void gather_test_uni_2d(const Dtype val[Nx][Nx],
                        const int num) {
  Dtype sum = 0.0;
  for (int i = 0; i < num; i++) {
    sum += val[0][i] * val[i][0];
  }
  std::cout << sum << std::endl;
}

void gather_test_matmulvec(const Dtype a[Nx][Nx],
                           const Dtype b[Nx][Nx],
                           Dtype c[Nx]) {
  for (int i = 0; i < Nx; i++) {
    for (int j = 0; j < Nx; j++) {
      c[i] += a[i][j] * b[j][0];
    }
  }
}

void gather_test_matmulmat(const Dtype a[Nx][Nx],
                           const Dtype b[Nx][Nx],
                           Dtype c[Nx][Nx]) {
  for (int i = 0; i < Nx; i++) {
    for (int j = 0; j < Nx; j++) {
      Dtype sum = 0;
      for (int k = 0; k < Nx; k++) {
        sum += a[i][k] * b[k][j];
      }
      c[i][j] = sum;
    }
  }
}

int idx[100];
Dtype val[10000], val2d[Nx][Nx], a[Nx][Nx], b[Nx][Nx], c[Nx], c_mat[Nx][Nx];

int main() {
  std::mt19937 mt;
  std::uniform_int_distribution<> ud(0, 99);
  std::generate(idx, idx + 100, [&mt, &ud](){return ud(mt);});
  std::generate(a[0], a[Nx], mt);
  std::generate(b[0], b[Nx], mt);
  std::fill(c, c + Nx, 0);
  std::fill(c_mat[0], c_mat[Nx], 0);
  std::iota(val, val + 100, 0.0);
  std::iota(val2d[0], val2d[100], 0.0);
  gather_test_ran(idx, val, 100);
  gather_test_uni(val, 100);
  gather_test_uni_2d(val2d, 100);
  gather_test_matmulvec(a, b, c);
  gather_test_matmulmat(a, b, c_mat);
}