blob: a4df6ed1cff1b5315a3ad76e190e4ce8bc087cba (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
|
// 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);
}
|
