Add flag torch_jit_disable_warning_prints to allow disabling all warn…

…ings.warn (pytorch#49313)

Summary:
Adding a flag torch_jit_disable_warning_prints to optimize interpreter performance by suppressing (potentially large amount) of warnings.warn.

This is to work around TorchScript's warning behavior mismatch with Python. Python by default triggers a warning once per location but TorchScript doesn't support it. This causes same warning to trigger and print once per inference run, hurting performance.

Pull Request resolved: pytorch#49313

Reviewed By: SplitInfinity

Differential Revision: D25534274

Pulled By: gmagogsfm

fbshipit-source-id: eaeb57a335c3e6c7eb259671645db05d781e80a2

torch/csrc/jit/runtime/interpreter.cpp

@@ -891,6 +891,10 @@ struct CodeImpl {
   }
 
   void emitWarn(Node* node) {
+    if (FLAGS_torch_jit_disable_warning_prints) {
+      return;
+    }
+
     emitLoadInputs(node->inputs());
     int32_t idx = -1;
     if (node->hasAttribute(attr::warn_id)) {

torch/csrc/jit/runtime/interpreter.h

@@ -8,6 +8,8 @@
 #include <torch/csrc/WindowsTorchApiMacro.h>
 #include <torch/csrc/jit/frontend/source_range.h>
 
+C10_DECLARE_bool(torch_jit_disable_warning_prints);
+
 namespace at {
 class Tensor;
 CAFFE2_API void launch(std::function<void()> func);

torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp

@@ -35,6 +35,11 @@ C10_DEFINE_bool(
     true,
     "If this flag is set to false TorchScript will be using the legacy/original executor");
 
+C10_DEFINE_bool(
+    torch_jit_disable_warning_prints,
+    false,
+    "Disables warning.warn prints in TorchScript graph");
+
 constexpr size_t kDefaultNumProfiledRuns = 1;
 constexpr size_t kDefaultBailoutDepth = 20;