【问题分析】Tensorflow Unit Test与Benchmark Test

 

概述

Tensorflow中的native程序的测试用例分为两类，gtest写的unit test与BENCHMARK宏定义的Benchmark Test，本文旨在分析它们是如何被框架调用执行。

 

实例

举个例子：Kernel cast_op的测试用例cast_op_test 中用TEST_XXX的宏定义了可以运行在gtest框架下的测试用例

#define TEST_CAST(in, out)                                              \
  TEST_F(CastOpTest, TestCast##_##in##_##out) { CheckCast(); } \
  TEST_F(CastOpTest, TestCast2##_##in##_##out) { CheckCast(true); }

#define TEST_ALL_CASTS_FROM(in) \
  TEST_CAST(in, uint8);         \
  ...

TEST_ALL_CASTS_FROM(uint8)
...

 

Kernel random_op的测试用例 random_op_test用BENCHMARK的宏定义了Tensorflow的Benchmark 测试用例 

#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  test::Benchmark("gpu", Cast(num)).Run(iters);
#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
#ifdef TENSORFLOW_USE_SYCL
  test::Benchmark("sycl", Cast(num)).Run(iters);
#endif  // TENSORFLOW_USE_SYCL
}
BENCHMARK(BM_gpu_float_int64)->Arg(64 << 10)->Arg(32 << 20);

测试执行命令行 

#   # for CPU tests
#   $ bazel test --config opt //third_party/tensorflow/core/kernels:my_op_test
#   # for GPU benchmarks
#   $ bazel run --config opt --config=cuda //third_party/tensorflow/core/kernels:my_op_test_gpu -- --benchmarks=..

 

原理分析 

BENCHMARK 宏的实现原理

 

BENCHMARK宏的实现原理如上图所示，可以看到BENCHMARK宏注册的测试用例需要有人去调用Run/RunBenchmarks才能触发执行，包括TEST_XXX宏注册的gtest用例也需要有人去调用RUN_ALL_TESTS才能触发执行，但是在上面Kernel的test 代码中并没有看到调用Run/RunBenchmarks/RUN_ALL_TESTS

根据Google提供的测试用例运行命令，怀疑还有其它模块和上面的Kernel test代码一起编译连接，实现了对Run/RunBenchmarks/RUN_ALL_TESTS的调用

#   # for CPU tests
#   $ bazel test --config opt //third_party/tensorflow/core/kernels:my_op_test
#   # for GPU benchmarks
#   $ bazel run --config opt --config=cuda //third_party/tensorflow/core/kernels:my_op_test_gpu -- --benchmarks=..

 

//tensorflow/core/platform/test_main.cc

GTEST_API_ int main(int argc, char** argv) {
  std::cout << "Running main() from test_main.cc\n";

  tensorflow::testing::InstallStacktraceHandler();
  testing::InitGoogleTest(&argc, argv);
  for (int i = 1; i < argc; i++) {
    if (absl::StartsWith(argv[i], "--benchmarks=")) {
      const char* pattern = argv[i] + strlen("--benchmarks=");
      tensorflow::testing::Benchmark::Run(pattern);
      return 0;
    }
  }
  return RUN_ALL_TESTS();
}

奥秘就在//tensorflow/core/platform/test_main.cc，它提供了main函数可以根据命令参数中是否带有--benchmarks来决定是否调用Run触发Benchmark测试用例，还是调用RUN_ALL_TESTS触发gtest测试用例，而test_main.cc自己被编译为test_main的library

# Main program for tests
cc_library(
    name = "test_main",
    testonly = 1,
    srcs = ["//tensorflow/core/platform:test_main.cc"],
...
)

 

tf_cuda_cc_test(
    name = "cast_op_test",
    size = "small",
    srcs = ["cast_op_test.cc"],
    deps = [
...
        "//tensorflow/core:test_main",
...


tf_cuda_cc_test(
    name = "random_op_test",
    size = "small",
    srcs = ["random_op_test.cc"],
    deps = [
...
        "//tensorflow/core:test_main",
...

kernel中的cast_op_test和random_op_test的编译都依赖了test_main lib，所以虽然它们的代码中没有显示调用Run/RunBenchmarks/RUN_ALL_TESTS，但是在test_main中提供了main函数会根据运行参数调用Run/RunBenchmarks/RUN_ALL_TESTS 执行gtest或者Benchmark测试用例，其它的op test case也是同样的原理得到执行