benchmark.cpp

const int kRuns = 1000;
const int kWarmUp = 100000;

void benchmark() {
  using clock = std::chrono::high_resolution_clock;
  using duration = std::chrono::duration<double, std::micro>;

  auto x = make_variable(CPU(kFloat).randn({10, 10}), true);
  backward(x.sum());

  size_t u = 0;

  for (int i = 0; i < kWarmUp; ++i) {
    x.set_requires_grad(i % 2);
    // u += x.requires_grad();
    // u += x.grad().numel();
  }

  float total = 0, min = 1000000000000, max = -1;
  for (int i = 0; i < kRuns; ++i) {
    const auto start = clock::now();

    x.set_requires_grad(i % 2);
    // u += x.requires_grad();
    // u += x.grad().numel();

    const auto end = clock::now();
    const auto elapsed = std::chrono::duration_cast<duration>(end - start);
    const auto e = elapsed.count();
    total += e;
    if (e > max) {
      max = e;
    } else if (e < min) {
      min = e;
    }
  }

  std::cout << "average time: " << (total / kRuns) << " | min time: " << min
            << " | max time: " << max << std::endl;
}