1e67c90e2c
XLA is a compiler-based linear algebra execution engine that targets CPUs, GPUs and custom accelerators. XLA is still experimental; we are releasing it early to get the community involved. Change: 143990941
52 lines
1.6 KiB
Plaintext
52 lines
1.6 KiB
Plaintext
// 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 "third_party/eigen3/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);
|
|
}
|