1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "absl/strings/str_cat.h"
17 #include "tensorflow/compiler/jit/tests/auto_clustering_test_helper.h"
18 #include "tensorflow/core/lib/core/status_test_util.h"
19
20 namespace tensorflow {
21 namespace {
22 class AutoClusteringTestImpl : public AutoClusteringTest {
23 protected:
24 // Test auto-clustering with a proto text file ${key}.pbtxt.
RunAutoClusteringTestWithPbtxt(absl::string_view key)25 Status RunAutoClusteringTestWithPbtxt(absl::string_view key) {
26 string file_name_without_extension =
27 absl::StrCat(testing::TensorFlowSrcRoot(), "/compiler/jit/tests/", key);
28
29 return AutoClusteringTest::RunAutoClusteringTestWithPbtxt(
30 absl::StrCat(file_name_without_extension, ".pbtxt"),
31 absl::StrCat(file_name_without_extension, ".golden_summary"));
32 }
33
34 // Test auto-clustering with a gzipped proto text file ${key}.pbtxt.gz.
RunAutoClusteringTestWithGzippedPbtxt(absl::string_view key)35 Status RunAutoClusteringTestWithGzippedPbtxt(absl::string_view key) {
36 string file_name_without_extension =
37 absl::StrCat(testing::TensorFlowSrcRoot(), "/compiler/jit/tests/", key);
38
39 return AutoClusteringTest::RunAutoClusteringTestWithGzippedPbtxt(
40 absl::StrCat(file_name_without_extension, ".pbtxt.gz"),
41 absl::StrCat(file_name_without_extension, ".golden_summary"));
42 }
43 };
44
TEST_F(AutoClusteringTestImpl,KerasImagenetMain)45 TEST_F(AutoClusteringTestImpl, KerasImagenetMain) {
46 // Generated from
47 //
48 // TARGET_PATH=tensorflow_models/official/legacy/image_classification \
49 // bazel run -c opt --config=cuda ${TARGET_PATH}:resnet_imagenet_main \
50 // -- --skip_eval --num_gpus=1 --dtype=fp16 --batch_size=192 \
51 // --train_steps=210 --enable_xla --enable_eager=true
52 //
53 // At CL 245846452
54 TF_ASSERT_OK(RunAutoClusteringTestWithPbtxt("keras_imagenet_main"));
55 }
56
TEST_F(AutoClusteringTestImpl,KerasImagenetMainGraphMode)57 TEST_F(AutoClusteringTestImpl, KerasImagenetMainGraphMode) {
58 // Generated from
59 //
60 // TARGET_PATH=tensorflow_models/official/legacy/image_classification \
61 // bazel run -c opt --config=cuda ${TARGET_PATH}:resnet_imagenet_main \
62 // -- --use_synthetic_data --num_gpus=1 --batch_size=117 --train_steps=600 \
63 // --skip_eval=True --logtostderr --enable_xla
64 TF_ASSERT_OK(
65 RunAutoClusteringTestWithPbtxt("keras_imagenet_main_graph_mode"));
66 }
67
TEST_F(AutoClusteringTestImpl,OpenSeq2SeqGNMT)68 TEST_F(AutoClusteringTestImpl, OpenSeq2SeqGNMT) {
69 // Model is from https://github.com/NVIDIA/OpenSeq2Seq.
70 // Generated from
71 //
72 // python run.py \
73 // --config_file=example_configs/text2text/en-de/en-de-gnmt-like-4GPUs.py \
74 // --use_xla_jit
75 TF_ASSERT_OK(
76 RunAutoClusteringTestWithGzippedPbtxt("opens2s_gnmt_mixed_precision"));
77 }
78
79 #if defined(PLATFORM_GOOGLE)
BenchmarkHelper(absl::string_view key,benchmark::State & state)80 Status BenchmarkHelper(absl::string_view key, benchmark::State& state) {
81 return BenchmarkMarkForCompilation(
82 absl::StrCat(testing::TensorFlowSrcRoot(), "/compiler/jit/tests/", key,
83 ".pbtxt"),
84 state);
85 }
86
BM_MarkForCompilationPass_KerasImagenetMain(benchmark::State & state)87 void BM_MarkForCompilationPass_KerasImagenetMain(benchmark::State& state) {
88 TF_CHECK_OK(BenchmarkHelper("keras_imagenet_main", state));
89 }
90
91 BENCHMARK(BM_MarkForCompilationPass_KerasImagenetMain);
92 #endif // PLATFORM_GOOGLE
93
94 } // namespace
95 } // namespace tensorflow
96