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 #ifndef TENSORFLOW_COMPILER_JIT_TESTS_AUTO_CLUSTERING_TEST_HELPER_H_ 16 #define TENSORFLOW_COMPILER_JIT_TESTS_AUTO_CLUSTERING_TEST_HELPER_H_ 17 18 #include "tensorflow/compiler/xla/statusor.h" 19 #include "tensorflow/core/common_runtime/graph_constructor.h" 20 #include "tensorflow/core/platform/test.h" 21 #include "tensorflow/core/platform/test_benchmark.h" 22 23 namespace tensorflow { 24 // Helper to write integration tests and benchmarks for the auto-clustering pass 25 // pipeline. These tests run auto-clustering on a graphdef and compare a 26 // summary of the auto-clustering decisions with a "golden" summary. 27 // 28 // To create a new test from an TF workload first run the workload with the 29 // following environment variables set: 30 // 31 // TF_DUMP_GRAPH_PREFIX=<some temporary directory> 32 // TF_XLA_FLAGS="--tf_xla_clustering_debug" 33 // 34 // If auto-clustering is enabled this should produce files named 35 // before_mark_for_compilation_<N>.pbtxt in the temporary directory. As the 36 // file name suggests, these are graphdefs that have been dumped right before 37 // the mark_for_compilation pass. There should be one 38 // before_mark_for_compilation_<N>.pbtxt for every TF graph that was 39 // auto-clustered, out of which usually only one is the "main" graph that's 40 // running training/inference. 41 // 42 // Copy the pbtxt for that "main" graph to tensorflow/compiler/jit/tests/ 43 // (i.e. this directory) and create a corresponding empty .golden_summary file. 44 // Add the .pbtxt and .golden_summary files to the "data" section of the cc_test 45 // rule for :auto_clustering_test and then see the comment on update_golden on 46 // how to auto-generate the .golden_summary file. 47 48 class AutoClusteringTest : public ::testing::Test { 49 protected: 50 Status RunAutoClusteringTestWithPbtxt( 51 absl::string_view pbtxt_file_path, 52 absl::string_view golden_summary_file_path); 53 Status RunAutoClusteringTestWithGzippedPbtxt( 54 absl::string_view gzipped_pbtxt_file_path, 55 absl::string_view golden_summary_file_path); 56 57 private: 58 Status RunAutoClusteringTestImpl(GraphDef graphdef, 59 absl::string_view golden_summary_file_path); 60 }; 61 62 #if defined(PLATFORM_GOOGLE) 63 // Reads the GraphDef stored in graph_def_path (which must be a pbtxt file) and 64 // benchmarks MarkForCompilationPass on this graphdef. 65 Status BenchmarkMarkForCompilation(absl::string_view graph_def_path, 66 benchmark::State& state); 67 #endif // PLATFORM_GOOGLE 68 69 } // namespace tensorflow 70 71 #endif // TENSORFLOW_COMPILER_JIT_TESTS_AUTO_CLUSTERING_TEST_HELPER_H_ 72