pytorch

1
#include <torch/csrc/jit/mobile/model_tracer/BuildFeatureTracer.h>
2
#include <mutex>
3

4
namespace torch {
5
namespace jit {
6
namespace mobile {
7
BuildFeatureTracer::BuildFeatureTracer() {
8
  auto recorder_cb =
9
      [](const at::RecordFunction& fn) -> std::unique_ptr<at::ObserverContext> {
10
    std::string name = fn.name();
11
    getBuildFeatures().withLock(
12
        [&](BuildFeatureTracer::build_feature_type& build_features) {
13
          build_features.insert(name);
14
        });
15
    return nullptr;
16
  };
17

18
  handle_ =
19
      at::addGlobalCallback(at::RecordFunctionCallback(recorder_cb)
20
                                .scopes({at::RecordScope::BUILD_FEATURE}));
21
}
22

23
c10::Synchronized<BuildFeatureTracer::build_feature_type>& BuildFeatureTracer::
24
    getBuildFeatures() {
25
  static c10::Synchronized<build_feature_type> build_features;
26
  return build_features;
27
}
28

29
} // namespace mobile
30
} // namespace jit
31
} // namespace torch
32