Add global boolean for controlling whether to record concrete shapes …

…or not (pytorch#101043)

Summary:
Pull Request resolved: pytorch#101043

We don't think the performance impact of recording concrete shapes is significant; but it's good to have a knob for turning it off quickly in case it has a large performance impact.

Test Plan:
Ran D45681838. It prints the state of that "concrete inputs" boolean. I ran it before and after canarying a change to `pytorch/kineto:pytorch_record_concrete_inputs`; before, it returns true; after, it returns false.

Note that D45681838 had to add `service` on the main function. That's because we need to `initFacebook` in order to use jks.

Reviewed By: aaronenyeshi

Differential Revision: D45680162

fbshipit-source-id: 680780ad9c947c0c60d09b3ba9904490730e96cb

torch/csrc/profiler/collection.cpp

@@ -1213,6 +1213,21 @@ RecordQueue::getRecords(
   return {out, std::move(trace)};
 }
 
+namespace {
+std::atomic<bool>& record_concrete_inputs_enabled() {
+  static std::atomic<bool> val{true};
+  return val;
+}
+} // namespace
+
+bool get_record_concrete_inputs_enabled() {
+  return record_concrete_inputs_enabled();
+}
+
+void set_record_concrete_inputs_enabled(bool val) {
+  record_concrete_inputs_enabled() = val;
+}
+
 } // namespace impl
 } // namespace profiler
 } // namespace torch

torch/csrc/profiler/collection.h

@@ -604,6 +604,9 @@ class TORCH_API RecordQueue {
   std::unique_ptr<python_tracer::PythonTracerBase> python_tracer_;
 };
 
+TORCH_API bool get_record_concrete_inputs_enabled();
+TORCH_API void set_record_concrete_inputs_enabled(bool);
+
 } // namespace impl
 } // namespace profiler
 } // namespace torch