xref: /aosp_15_r20/external/pytorch/torch/csrc/jit/backends/backend_debug_info.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #ifndef BUILD_LITE_INTERPRETER
4 #include <torch/csrc/jit/backends/backend_debug_handler.h>
5 #endif
6 #include <torch/custom_class.h>
7 
8 namespace torch {
9 namespace jit {
10 
11 constexpr static auto kBackendUtilsNamespace = "backendutils";
12 constexpr static auto kBackendDebugInfoClass = "BackendDebugInfo";
13 
14 #ifndef BUILD_LITE_INTERPRETER
15 /*
16  * Custom class for holding debug information in lowered modules, intended
17  * purely for keeping this information to be later serialized outside of the
18  * lowered module itself.
19  * Its usage pattern is:
20  * 1. LoweredModule declares an instance of this class in __backend_debug_info
21  * 2. During serialization, __backend_debug_info is used to obtain the debug
22  *    information.
23  * 3. The contents of LoweredModule.__backend_debug_info are not serialized
24  *    within the LoweredModule itself.
25  */
26 class TORCH_API PyTorchBackendDebugInfo : public torch::CustomClassHolder {
27  public:
28   PyTorchBackendDebugInfo() = default;
29 
getDebugInfoMap()30   std::optional<BackendDebugInfoMapType>& getDebugInfoMap() {
31     return debug_info_map_;
32   }
33 
setDebugInfoMap(BackendDebugInfoMapType && debug_info_map)34   void setDebugInfoMap(BackendDebugInfoMapType&& debug_info_map) {
35     debug_info_map_ = std::move(debug_info_map);
36   }
37 
38  private:
39   std::optional<BackendDebugInfoMapType> debug_info_map_;
40 };
41 
42 #else
43 
44 /*
45  * Dummy instance exists for the following reason:
46  * __backend_debug_info is of type BackendDebugInfo which is a torchbind'
47  * class backed by cpp class PyTorchBackendDebugInfo.
48  * PyTorchBackendDebugInfo, depends on ir.h., scope.h, source_range etc.
49  * We dont include this on lite interpreter side. Thus on lite interpreter side
50  * we cannot have valid definition of PyTorchBackendDebugInfo. However we do not
51  * need valid instance of __backend_debug_info in lite interpreter anyway as we
52  * dont serialize this info as part of LowerdModule as mentioned ealrier.
53  * However since LoweredModule has registered attribute of __backend_debug_info
54  * we still need to make sure that BackendDebugInfo is registered with
55  * TorchScript. However in this instance it does not have to be backed by
56  * PyTorchBackendDebugInfo, so we create a dummy PyTorchBackendDebugInfoDummy
57  * just for this purpose.
58  */
59 class PyTorchBackendDebugInfoDummy : public torch::CustomClassHolder {
60  public:
61   PyTorchBackendDebugInfoDummy() = default;
62 };
63 #endif
64 } // namespace jit
65 } // namespace torch
66