1 #pragma once 2 3 #include <ATen/record_function.h> 4 #include <c10/util/Synchronized.h> 5 #include <map> 6 #include <set> 7 #include <string> 8 9 namespace torch::jit::mobile { 10 11 /* The BuildFeatureTracer class handles the attachment and removal of a 12 * recording callback that traces the invocation of code that handles executing 13 * generic build features. 14 * 15 * You can get the set of used build features using 16 * getBuildFeatures(). 17 * 18 * Note: This class is not thread safe or re-entrant, and should not be used 19 * across multiple threads of execution. 20 * 21 */ 22 struct BuildFeatureTracer final { 23 at::CallbackHandle handle_; 24 /* These are the custom class names (constant 25 * character string) which shows up in code. 26 */ 27 typedef std::set<std::string> build_feature_type; 28 29 BuildFeatureTracer(); 30 static c10::Synchronized<build_feature_type>& getBuildFeatures(); 31 ~BuildFeatureTracerfinal32 ~BuildFeatureTracer() { 33 at::removeCallback(handle_); 34 } 35 }; 36 37 } // namespace torch::jit::mobile 38