// Generated by the tf_library build rule.  DO NOT EDIT!
//
// This file contains the main function and logic for benchmarking code
// generated by tfcompile.  All tokens of the form `{{TFCOMPILE_*}}` must be
// rewritten to real values before this file can be compiled.
//
//    TFCOMPILE_HEADER    : Path to the header file generated by tfcompile.
//    TFCOMPILE_CPP_CLASS : Name of the C++ class generated by tfcompile.
//
// The tf_library bazel macro in tfcompile.bzl performs the token rewriting, and
// generates a cc_binary rule for you.

// These macros must be defined before eigen files are included.
#define EIGEN_USE_THREADS
#define EIGEN_USE_CUSTOM_THREAD_POOL

// clang-format off
#include "{{TFCOMPILE_HEADER}}"  // NOLINT(whitespace/braces)
// clang-format on

#include "tensorflow/compiler/aot/benchmark.h"
#include "unsupported/Eigen/CXX11/Tensor"

// Macros that expand to tokens based on the entry point name.
// clang-format off
#define CPP_CLASS {{TFCOMPILE_CPP_CLASS}}  // NOLINT(whitespace/braces)
// clang-format on

namespace tensorflow {
namespace tfcompile {

int Main(int argc, char** argv) {
  Eigen::ThreadPool pool(1 /* num_threads */);
  Eigen::ThreadPoolDevice device(&pool, pool.NumThreads());

  CPP_CLASS computation;
  computation.set_thread_pool(&device);

  benchmark::Options options;
  benchmark::Stats stats;
  benchmark::Benchmark(options, [&] { computation.Run(); }, &stats);
  benchmark::DumpStatsToStdout(stats);
  return 0;
}

}  // namespace tfcompile
}  // namespace tensorflow

int main(int argc, char** argv) {
  return tensorflow::tfcompile::Main(argc, argv);
}