Make sure each warnings.warn only executes once inside TorchScript.

aten/src/ATen/core/interned_strings.h

@@ -360,7 +360,8 @@ namespace c10 {
   _(attr, scope)                     \
   _(attr, keepdims)                  \
   _(attr, cache_id)                  \
-  _(attr, new_axis)
+  _(attr, new_axis)                  \
+  _(attr, warn_id)
 #else
 #define FORALL_NS_SYMBOLS(_) \
   _(namespaces, prim)              \

test/jit/test_warn.py

@@ -0,0 +1,165 @@
+import os
+import sys
+import io
+
+import torch
+import warnings
+from contextlib import redirect_stderr
+from torch.testing import FileCheck
+
+# Make the helper files in test/ importable
+pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
+sys.path.append(pytorch_test_dir)
+from torch.testing._internal.jit_utils import JitTestCase
+
+if __name__ == '__main__':
+    raise RuntimeError("This test file is not meant to be run directly, use:\n\n"
+                       "\tpython test/test_jit.py TESTNAME\n\n"
+                       "instead.")
+
+
+class TestWarn(JitTestCase):
+    def test_warn(self):
+        @torch.jit.script
+        def fn():
+            warnings.warn("I am warning you")
+
+        f = io.StringIO()
+        with redirect_stderr(f):
+            fn()
+
+        FileCheck() \
+            .check_count(
+                str="UserWarning: I am warning you",
+                count=1,
+                exactly=True) \
+            .run(f.getvalue())
+
+    def test_warn_only_once(self):
+        @torch.jit.script
+        def fn():
+            for _ in range(10):
+                warnings.warn("I am warning you")
+
+        f = io.StringIO()
+        with redirect_stderr(f):
+            fn()
+
+        FileCheck() \
+            .check_count(
+                str="UserWarning: I am warning you",
+                count=1,
+                exactly=True) \
+            .run(f.getvalue())
+
+    def test_warn_only_once_in_loop_func(self):
+        def w():
+            warnings.warn("I am warning you")
+
+        @torch.jit.script
+        def fn():
+            for _ in range(10):
+                w()
+
+        f = io.StringIO()
+        with redirect_stderr(f):
+            fn()
+
+        FileCheck() \
+            .check_count(
+                str="UserWarning: I am warning you",
+                count=1,
+                exactly=True) \
+            .run(f.getvalue())
+
+    def test_warn_once_per_func(self):
+        def w1():
+            warnings.warn("I am warning you")
+
+        def w2():
+            warnings.warn("I am warning you")
+
+        @torch.jit.script
+        def fn():
+            w1()
+            w2()
+
+        f = io.StringIO()
+        with redirect_stderr(f):
+            fn()
+
+        FileCheck() \
+            .check_count(
+                str="UserWarning: I am warning you",
+                count=2,
+                exactly=True) \
+            .run(f.getvalue())
+
+    def test_warn_once_per_func_in_loop(self):
+        def w1():
+            warnings.warn("I am warning you")
+
+        def w2():
+            warnings.warn("I am warning you")
+
+        @torch.jit.script
+        def fn():
+            for _ in range(10):
+                w1()
+                w2()
+
+        f = io.StringIO()
+        with redirect_stderr(f):
+            fn()
+
+        FileCheck() \
+            .check_count(
+                str="UserWarning: I am warning you",
+                count=2,
+                exactly=True) \
+            .run(f.getvalue())
+
+    def test_warn_multiple_calls_multiple_warnings(self):
+        @torch.jit.script
+        def fn():
+            warnings.warn("I am warning you")
+
+        f = io.StringIO()
+        with redirect_stderr(f):
+            fn()
+            fn()
+
+        FileCheck() \
+            .check_count(
+                str="UserWarning: I am warning you",
+                count=2,
+                exactly=True) \
+            .run(f.getvalue())
+
+    def test_warn_multiple_calls_same_func_diff_stack(self):
+        def warn(caller: str):
+            warnings.warn("I am warning you from " + caller)
+
+        @torch.jit.script
+        def foo():
+            warn("foo")
+
+        @torch.jit.script
+        def bar():
+            warn("bar")
+
+        f = io.StringIO()
+        with redirect_stderr(f):
+            foo()
+            bar()
+
+        FileCheck() \
+            .check_count(
+                str="UserWarning: I am warning you from foo",
+                count=1,
+                exactly=True) \
+            .check_count(
+                str="UserWarning: I am warning you from bar",
+                count=1,
+                exactly=True) \
+            .run(f.getvalue())

test/test_jit.py

@@ -32,6 +32,7 @@
 from jit.test_enum import TestEnum  # noqa: F401
 from jit.test_profiler import TestProfiler  # noqa: F401
 from jit.test_slice import TestSlice  # noqa: F401
+from jit.test_warn import TestWarn  # noqa: F401
 
 # Torch
 from torch import Tensor

tools/build_variables.bzl

@@ -148,6 +148,7 @@ core_sources_full = [
     "torch/csrc/jit/ir/scope.cpp",
     "torch/csrc/jit/ir/subgraph_matcher.cpp",
     "torch/csrc/jit/jit_log.cpp",
+    "torch/csrc/jit/passes/annotate_warns.cpp",
     "torch/csrc/jit/passes/bailout_graph.cpp",
     "torch/csrc/jit/passes/batch_mm.cpp",
     "torch/csrc/jit/passes/canonicalize.cpp",

torch/csrc/jit/frontend/ir_emitter.cpp

@@ -8,6 +8,7 @@
 #include <torch/csrc/jit/frontend/schema_matching.h>
 #include <torch/csrc/jit/frontend/script_type_parser.h>
 #include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/passes/annotate_warns.h>
 #include <torch/csrc/jit/passes/canonicalize.h>
 #include <torch/csrc/jit/passes/constant_pooling.h>
 #include <torch/csrc/jit/passes/constant_propagation.h>
@@ -4090,6 +4091,10 @@ void runCleanupPasses(std::shared_ptr<Graph>& to_clean) {
 
   // For jitter
   CanonicalizeOutputs(to_clean);
+
+  // Annotate aten::warns so that each has its unique ID. This enables us to
+  // mimic Python behavior of only emitting each warning only once.
+  AnnotateWarns(to_clean);
 }
 
 // we consider _N where N is a number, to be a non-meaningful name

torch/csrc/jit/passes/annotate_warns.cpp

@@ -0,0 +1,29 @@
+#include <torch/csrc/jit/passes/annotate_warns.h>
+
+#include <atomic>
+
+namespace torch {
+namespace jit {
+
+void AnnotateWarns(Block* b) {
+  static std::atomic<int64_t> idx(0);
+  for (Node* n : b->nodes()) {
+    for (Block* child_b : n->blocks()) {
+      AnnotateWarns(child_b);
+    }
+
+    if (n->kind() != aten::warn) {
+      continue;
+    }
+
+    n->i_(attr::warn_id, idx);
+    idx++;
+  }
+}
+
+void AnnotateWarns(const std::shared_ptr<Graph>& graph) {
+  AnnotateWarns(graph->block());
+}
+
+} // namespace jit
+} // namespace torch

torch/csrc/jit/passes/annotate_warns.h

@@ -0,0 +1,11 @@
+#pragma once
+
+#include <torch/csrc/jit/ir/ir.h>
+
+namespace torch {
+namespace jit {
+
+TORCH_API void AnnotateWarns(const std::shared_ptr<Graph>& graph);
+
+} // namespace jit
+} // namespace torch

torch/csrc/jit/runtime/instruction.h

@@ -52,7 +52,7 @@ namespace jit {
   _(ISINSTANCE, "TI") /* check object is one of  types[X:X+N]  */              \
   _(TUPLE_SLICE, "II") /* slice tup[X:(X+N)] */                                \
   _(FORK, "CN") /* launch a thread to run code entry x with N inputs  */       \
-  _(WARN, "") /* emit a warning with line information */                       \
+  _(WARN, "I") /* emit a warning with line information */                      \
   _(ENTER, "EN") /* enter scope of a contextmanager */                         \
   _(EXIT, "EX") /* exit the last entered contextmanager */
 

torch/csrc/jit/runtime/interpreter.cpp

@@ -428,6 +428,16 @@ struct TLSCurrentInterpreterGuard {
   InterpreterStateImpl* prev_state_;
 };
 
+template <class Ttarget, class Tsource>
+Ttarget safe_narrow_cast(Tsource v) {
+  Ttarget res = static_cast<Ttarget>(v);
+  // Casting it back to check whether it overflew.
+  if (static_cast<Tsource>(res) != v) {
+    throw std::runtime_error("safe_narrow_cast<>() failed due to overflow");
+  }
+  return res;
+}
+
 struct CodeImpl {
   friend struct InterpreterState;
   std::vector<Instruction> instructions_;
@@ -535,7 +545,10 @@ struct CodeImpl {
   }
 
   void insertInstruction(OpCode op, int64_t X = 0, uint64_t N = 0) {
-    instructions_.emplace_back(op, X, N);
+    instructions_.emplace_back(
+        op,
+        safe_narrow_cast<int32_t, int64_t>(X),
+        safe_narrow_cast<int16_t, int64_t>(N));
     instructions_source_.emplace_back(current_node_);
 
     // check that we didn't accidentally emit nodes out of topological order
@@ -873,7 +886,11 @@ struct CodeImpl {
 
   void emitWarn(Node* node) {
     emitLoadInputs(node->inputs());
-    insertInstruction(WARN);
+    int32_t idx = -1;
+    if (node->hasAttribute(attr::warn_id)) {
+      idx = static_cast<int32_t>(node->i(attr::warn_id));
+    }
+    insertInstruction(WARN, idx);
   }
 
   void emitEnter(Node* node) {
@@ -1017,6 +1034,22 @@ struct InterpreterStateImpl : c10::intrusive_ptr_target {
   }
 
  private:
+  struct WarnedNodes {
+   public:
+    // Inserts idx into warned_nodes_, returns a boolean indicates whether
+    // insertion actually happened (idx wasn't originally in the set).
+    bool insert(int32_t idx) {
+      std::unique_lock<std::mutex> lock(mutex_);
+      return warned_nodes_.insert(idx).second;
+    }
+
+   private:
+    std::mutex mutex_;
+    std::unordered_set<int32_t> warned_nodes_;
+  };
+
+  WarnedNodes warned_nodes_;
+
   // if we need to suspend, where do we reset the stack?
   // answer: to where it was when we were called, not
   // including any inputs to this function
@@ -1487,21 +1520,35 @@ struct InterpreterStateImpl : c10::intrusive_ptr_target {
             ++frame.pc;
           } break;
           case WARN: {
-            Node* node = frame.function->instructions_source_.at(frame.pc);
+            // Keeps track of which WARN instruction has been executed before,
+            // we only want to execute each WARN once to match default Python
+            // warning behavior.
+            bool need_warn = true;
+            if (inst.X != -1) {
+              need_warn = warned_nodes_.insert(inst.X);
+            }
+
+            Node* node =
+                frames.back().function->instructions_source_.at(frame.pc);
             auto range = node->sourceRange().source();
             if (range->filename()) {
-              auto line = range->starting_line_no() +
-                  range->lineno_for_offset(node->sourceRange().start());
               drop(stack, 1);
-              c10::SourceLocation location{
-                  "", range->filename()->c_str(), uint32_t(line)};
-              // Sends the warning to the warning handler with the
-              // "verbatim" flag. This flag ensures the warning handler
-              // will print the exception as configured.
-              c10::Warning::warn(
-                  location, pop(stack).toStringRef(), /*verbatim=*/true);
+              const auto msg = pop(stack).toStringRef();
+              if (need_warn) {
+                auto line = range->starting_line_no() +
+                    range->lineno_for_offset(node->sourceRange().start());
+                c10::SourceLocation location{
+                    "", range->filename()->c_str(), uint32_t(line)};
+                // Sends the warning to the warning handler with the
+                // "verbatim" flag. This flag ensures the warning handler
+                // will print the exception as configured.
+                c10::Warning::warn(location, msg, /*verbatim=*/true);
+              }
             } else {
-              TORCH_WARN(pop(stack).toStringRef());
+              const auto msg = pop(stack).toStringRef();
+              if (need_warn) {
+                TORCH_WARN(msg);
+              }
             }
             ++frame.pc;
           } break;