1 /* Copyright 2020 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 #ifndef TENSORFLOW_COMPILER_XLA_PYTHON_PPROF_PROFILE_BUILDER_H_ 17 #define TENSORFLOW_COMPILER_XLA_PYTHON_PPROF_PROFILE_BUILDER_H_ 18 19 #include <string> 20 #include <utility> 21 22 #include "absl/container/flat_hash_map.h" 23 #include "pybind11/pybind11.h" 24 #include "tensorflow/compiler/xla/statusor.h" 25 #include "tensorflow/core/profiler/profile.pb.h" 26 27 namespace xla { 28 29 // Helper class for building pprof::Profile profiles. 30 class PprofProfileBuilder { 31 public: 32 PprofProfileBuilder(); profile()33 tensorflow::tfprof::pprof::Profile& profile() { return profile_; } 34 35 // Adds or returns the ID of `s` in the table. 36 int StringId(const std::string& s); 37 38 // Adds or returns the ID of a function. 39 int FunctionId(PyCodeObject* code); 40 41 // Adds or returns the ID of a code location. 42 int LocationId(PyCodeObject* code, int instruction); 43 44 private: 45 tensorflow::tfprof::pprof::Profile profile_; 46 47 absl::flat_hash_map<std::string, int> strings_; 48 absl::flat_hash_map<PyCodeObject*, int> functions_; 49 absl::flat_hash_map<std::pair<PyCodeObject*, int>, int> locations_; 50 }; 51 52 // Converts the JSON representation of a pprof profile protocol buffer into 53 // a serialized protocol buffer. We want to allow Python code to construct pprof 54 // protocol buffers, but we don't want to export the generated protocol buffer 55 // bindings for Python because they cause conflicts between multiple Python 56 // extensions that contain the same protocol buffer message. Instead, we accept 57 // a JSON representation from Python and use this function to serialize it to 58 // a uncompressed binary protocol buffer. 59 StatusOr<pybind11::bytes> JsonToPprofProfile(std::string json); 60 61 // The reverse, useful for testing. 62 StatusOr<std::string> PprofProfileToJson(pybind11::bytes binary_proto); 63 64 } // namespace xla 65 66 #endif // TENSORFLOW_COMPILER_XLA_PYTHON_PPROF_PROFILE_BUILDER_H_ 67