xref: /aosp_15_r20/external/pytorch/torch/csrc/jit/mobile/model_tracer/BuildFeatureTracer.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1 #include <torch/csrc/jit/mobile/model_tracer/BuildFeatureTracer.h>
2 #include <mutex>
3 
4 namespace torch::jit::mobile {
BuildFeatureTracer()5 BuildFeatureTracer::BuildFeatureTracer() {
6   auto recorder_cb =
7       [](const at::RecordFunction& fn) -> std::unique_ptr<at::ObserverContext> {
8     std::string name = fn.name();
9     getBuildFeatures().withLock(
10         [&](BuildFeatureTracer::build_feature_type& build_features) {
11           build_features.insert(name);
12         });
13     return nullptr;
14   };
15 
16   handle_ =
17       at::addGlobalCallback(at::RecordFunctionCallback(recorder_cb)
18                                 .scopes({at::RecordScope::BUILD_FEATURE}));
19 }
20 
21 c10::Synchronized<BuildFeatureTracer::build_feature_type>& BuildFeatureTracer::
getBuildFeatures()22     getBuildFeatures() {
23   static c10::Synchronized<build_feature_type> build_features;
24   return build_features;
25 }
26 
27 } // namespace torch::jit::mobile
28