1 | #pragma once |
2 | |
3 | #include <iostream> |
4 | #include <string> |
5 | #include <vector> |
6 | |
7 | #include <torch/csrc/lazy/core/tensor.h> |
8 | |
9 | namespace torch { |
10 | namespace lazy { |
11 | |
12 | TORCH_API std::function<std::vector<SourceLocation>()>& |
13 | GetPythonFramesFunction(); |
14 | |
15 | TORCH_API std::string GetFirstUserFrameInPython(); |
16 | |
17 | class TORCH_API DebugUtil { |
18 | public: |
19 | enum GraphFormat { |
20 | kText, |
21 | kDot, |
22 | kBackend, |
23 | }; |
24 | |
25 | static GraphFormat GetDefaultGraphFormat(); |
26 | |
27 | // Dumps the current Python frame and the IR Graph whose roots are the IR |
28 | // values held at the tensors. If indices is not nullptr, it selects the |
29 | // indices of the tensors whose graph will be emitted. |
30 | static std::string GetTensorsGraphInfo( |
31 | c10::ArrayRef<torch::lazy::LazyTensorPtr> tensors, |
32 | const std::vector<size_t>* indices, |
33 | GraphFormat format = GetDefaultGraphFormat()); |
34 | |
35 | // If the environment variable LTC_SAVE_TENSORS_FILE is set to the proper |
36 | // output path, an instance of the report returned by GetTensorsGraphInfo() is |
37 | // saved. |
38 | static void SaveTensorsGraphInfo( |
39 | const char* name, |
40 | c10::ArrayRef<torch::lazy::LazyTensorPtr> tensors, |
41 | const std::vector<size_t>* indices, |
42 | GraphFormat format = GetDefaultGraphFormat()); |
43 | |
44 | static bool ExperimentEnabled(const std::string& name); |
45 | }; |
46 | |
47 | } // namespace lazy |
48 | } // namespace torch |
49 | |