vector_sum_profile.cc

#include <algorithm>
#include <cassert>
#include <chrono>
#include <random>
#include <vector>
#include "cc_benchmarks.h"

void gen_rand(std::vector<int32_t>& v, size_t length) {
  std::random_device rd;
  std::mt19937 gen(rd());
  std::uniform_int_distribution<> dis(0, 1);

  v.resize(length);
  for (size_t i = 0; i < length; ++i) {
    v[i] = dis(gen);
  }
}

class Inputs {
public:
  Inputs(const size_t invoke_count,
         const size_t vector_size,
         const size_t batch_size,
         const int32_t* data,
         int32_t* batch_sum_data)
      : invoke_count(invoke_count)
      , vector_size(vector_size)
      , batch_size(batch_size)
      , data(data)
      , batch_sum_data(batch_sum_data) {
    assert(vector_size % batch_size == 0);
    sequence_length = vector_size / batch_size;
  }

  const size_t invoke_count;
  const size_t vector_size;
  const size_t batch_size;
  size_t sequence_length;
  const int32_t* data;
  int32_t* batch_sum_data;
};

void test_std_if_block(const Inputs& i) {
  for (int b = 0; b < i.batch_size; ++b) {
    i.batch_sum_data[b] = static_cast<int32_t>(std::count_if(
      i.data + (static_cast<int64_t>(b) * i.sequence_length),
      i.data + (static_cast<int64_t>(b) * i.sequence_length) + i.sequence_length,
      [](int v) { return v == 1; }));
  }
}

void test_std_if_block2(const Inputs& i) {
  for (int b = 0; b < i.batch_size; ++b) {
    const int32_t* start = i.data + (static_cast<int64_t>(b) * i.sequence_length);
    const int32_t* end = i.data + (static_cast<int64_t>(b) * i.sequence_length) + i.sequence_length;
    i.batch_sum_data[b] = static_cast<int32_t>(std::count_if(start, end, [](int v) { return v == 1; }));
  }
}

void test_sum_vector(const Inputs& i) {
  for (int b = 0; b < i.batch_size; ++b) {
    const int32_t* cur_data = i.data + (static_cast<int64_t>(b) * i.sequence_length);
    int32_t cur_sum = 0;
    for (int s = 0; s < i.sequence_length; ++s) {
      cur_sum += cur_data[s];
    }
    i.batch_sum_data[b] = cur_sum;
  }
}

template <typename F>
void profile(const Inputs& inputs, const char* name, F&& func) {
  auto sum = 0;
  for (size_t i = 0; i < inputs.invoke_count; ++i) {
    auto start = std::chrono::high_resolution_clock::now();
    func();
    auto end = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::nanoseconds>(end - start);
    sum += duration.count();
  }
  std::cerr << name << " duration (ns): " << sum / inputs.invoke_count << std::endl;
}

int main() {
  size_t vector_size = 6;
  size_t batch_size = 3;

  std::vector<int32_t> vector;
  gen_rand(vector, vector_size);

  std::vector<int32_t> batch_sum;
  batch_sum.reserve(batch_size);

  Inputs inputs(/*invoke_count=*/10000, vector_size, batch_size, vector.data(), batch_sum.data());

  profile(inputs, "std::count_if", [=]() { test_std_if_block(inputs); });
  profile(inputs, "std::count_if2", [=]() { test_std_if_block2(inputs); });
  profile(inputs, "sum vector", [=]() { test_sum_vector(inputs); });

  return 0;
}