Fix segmentation fault.

- Reset session data after logging
- Add the return statement

Author: ysiraichi

test/test_dynamic_shapes_detector.py

@@ -5,7 +5,6 @@
 import test_utils
 import unittest
 
-
 # Processes a string, so that it can be used as the expected error regex.
 # Specifically, it does 3 things:
 #
@@ -23,8 +22,9 @@
 # regex wildcards, such as "*".
 ESCAPE_RE = re.compile(r"([\[\](){}])")
 
+
 def escape(s):
-    return ESCAPE_RE.sub(r"\\\1", textwrap.dedent(s[1:]))
+  return ESCAPE_RE.sub(r"\\\1", textwrap.dedent(s[1:]))
 
 
 class TestDynamicShapeDetector(test_utils.XlaTestCase):
@@ -85,7 +85,8 @@ def foo(x):
       return x + x
 
     inp1 = torch.rand(10, device=torch_xla.device())
-    self._run_and_compare(foo, args=(inp1,), max_different_graphs=max_different_graphs)
+    self._run_and_compare(
+        foo, args=(inp1,), max_different_graphs=max_different_graphs)
 
     expected_error_msg = escape(r"""
         Maximum number of different graphs allowed per function exceeded: 1
@@ -95,7 +96,8 @@ def foo(x):
 
     with self.assertRaisesRegex(RuntimeError, expected_error_msg):
       inp2 = torch.rand(5, device=torch_xla.device())
-      self._run_and_compare(foo, args=(inp2,), max_different_graphs=max_different_graphs)
+      self._run_and_compare(
+          foo, args=(inp2,), max_different_graphs=max_different_graphs)
 
   def test_graph_limit_exceeded_common_sequence_mismatch(self):
     # Test: catch graph limit exceeded error when the common sequence (i.e. compressed
@@ -117,7 +119,8 @@ def foo(x, step):
         return x * 5
 
     inp = torch.rand(10, device=torch_xla.device())
-    self._run_and_compare(foo, args=(inp, 0), max_different_graphs=max_different_graphs)
+    self._run_and_compare(
+        foo, args=(inp, 0), max_different_graphs=max_different_graphs)
 
     expected_error_msg = escape(r"""
         Maximum number of different graphs allowed per function exceeded: 1
@@ -126,7 +129,8 @@ def foo(x, step):
     """)
 
     with self.assertRaisesRegex(RuntimeError, expected_error_msg):
-      self._run_and_compare(foo, args=(inp, 2), max_different_graphs=max_different_graphs)
+      self._run_and_compare(
+          foo, args=(inp, 2), max_different_graphs=max_different_graphs)
 
   def test_graph_limit_exceeded_children_mismatch(self):
     # Test: catch graph limit exceeded error when the expected child of the trie
@@ -154,8 +158,10 @@ def foo(x, step):
       return r / 3
 
     inp = torch.rand(10, device=torch_xla.device())
-    self._run_and_compare(foo, args=(inp, 0), max_different_graphs=max_different_graphs)
-    self._run_and_compare(foo, args=(inp, 1), max_different_graphs=max_different_graphs)
+    self._run_and_compare(
+        foo, args=(inp, 0), max_different_graphs=max_different_graphs)
+    self._run_and_compare(
+        foo, args=(inp, 1), max_different_graphs=max_different_graphs)
 
     expected_error_msg = escape(r"""
         Maximum number of different graphs allowed per function exceeded: 2
@@ -166,7 +172,8 @@ def foo(x, step):
     """)
 
     with self.assertRaisesRegex(RuntimeError, expected_error_msg):
-      self._run_and_compare(foo, args=(inp, 2), max_different_graphs=max_different_graphs)
+      self._run_and_compare(
+          foo, args=(inp, 2), max_different_graphs=max_different_graphs)
 
   def test_graph_limit_exceeded_common_sequence_early_stop(self):
     # Test: catch graph limit exceeded error when the graph ends unexpectedly in
@@ -188,7 +195,8 @@ def foo(x, mul=False):
         return r
 
     inp = torch.rand(10, device=torch_xla.device())
-    self._run_and_compare(foo, args=(inp, True), max_different_graphs=max_different_graphs)
+    self._run_and_compare(
+        foo, args=(inp, True), max_different_graphs=max_different_graphs)
 
     expected_error_msg = escape(r"""
         Maximum number of different graphs allowed per function exceeded: 1
@@ -224,8 +232,10 @@ def foo(x, step):
       return r
 
     inp = torch.rand(10, device=torch_xla.device())
-    self._run_and_compare(foo, args=(inp, 0), max_different_graphs=max_different_graphs)
-    self._run_and_compare(foo, args=(inp, 1), max_different_graphs=max_different_graphs)
+    self._run_and_compare(
+        foo, args=(inp, 0), max_different_graphs=max_different_graphs)
+    self._run_and_compare(
+        foo, args=(inp, 1), max_different_graphs=max_different_graphs)
 
     expected_error_msg = escape(r"""
         Maximum number of different graphs allowed per function exceeded: 2
@@ -236,7 +246,8 @@ def foo(x, step):
     """)
 
     with self.assertRaisesRegex(RuntimeError, expected_error_msg):
-      self._run_and_compare(foo, args=(inp, 2), max_different_graphs=max_different_graphs)
+      self._run_and_compare(
+          foo, args=(inp, 2), max_different_graphs=max_different_graphs)
 
 
 if __name__ == "__main__":

torch_xla/csrc/dynamic_shape_detector.cpp

@@ -9,7 +9,8 @@ namespace torch_xla {
 // Maximum number of allowed graphs per function (i.e. session).
 static std::size_t max_different_graphs = 1;
 
-TrieNode::TrieNode(const TrieValue& value, bool is_graph_boundary) : TrieNode() {
+TrieNode::TrieNode(const TrieValue& value, bool is_graph_boundary)
+    : TrieNode() {
   common_sequence_.push_back(value);
   is_graph_boundary_ = is_graph_boundary;
 }
@@ -115,6 +116,7 @@ TrieBuilder TrieNode::AddValue(TrieValue value, std::size_t matched,
   // Maybe split the current node into: prefix (before matched) and suffix
   // (after matched).
   bool did_split = MaybeSplitAt(matched);
+  TF_VLOG(5) << "MaybeSplitAt(" << matched << "): " << did_split;
 
   // Create a new node that contains only the given value.
   std::unique_ptr<TrieNode> node = std::make_unique<TrieNode>(value);
@@ -132,6 +134,7 @@ TrieBuilder TrieNode::AddValue(TrieValue value, std::size_t matched,
     is_graph_boundary_ = false;
   }
 
+  TF_VLOG(5) << "Added value: " << value.str << " (" << value.hash << ")";
   return {children_[value.hash].get(), 1};
 }
 
@@ -144,6 +147,8 @@ bool TrieNode::MaybeSplitAt(std::size_t matched) {
       common_sequence.subspan(0, /*len=*/matched);
   absl::Span<const TrieValue> suffix = common_sequence.subspan(matched);
 
+  bool did_split = false;
+
   // A split only occurs if suffix is not empty.
   if (!suffix.empty()) {
     std::unique_ptr<TrieNode> suffix_node =
@@ -159,10 +164,13 @@ bool TrieNode::MaybeSplitAt(std::size_t matched) {
     TF_VLOG(5) << "Split node " << children_[suffix.front().hash].get()
                << " at position " << matched << ": " << suffix.front().str
                << " (" << suffix.front().hash << ")";
+
+    did_split = true;
   }
 
   // This node's common_sequence_ will be whatever the prefix was.
-  common_sequence_ = std::vector<TrieValue>{prefix.begin(), prefix.end()};
+  common_sequence_.erase(common_sequence_.begin() + matched, common_sequence_.end());
+  return did_split;
 }
 
 DynamicShapeDetector* DynamicShapeDetector::Get() {
@@ -209,8 +217,8 @@ void DynamicShapeDetector::EndSession() {
       TF_VLOG(5) << "Created new graph.";
     }
 
-    ResetSession();
     TF_VLOG(5) << "Ended session: " << current_session_->name_;
+    ResetSession();
   } catch (const std::exception& e) {
     // MarkGraphBoundary might raise an exception if AllowNewGraph() is false.
     // Catch it here, so that we can correctly end the session.

torch_xla/csrc/init_python_bindings.cpp

@@ -2619,9 +2619,7 @@ void InitXlaModuleBindings(py::module m) {
           DynamicShapeDetector::SetMaxDifferentGraphs(max_different_graphs);
         });
   m.def("_dynamic_shape_detector_get_max_different_graphs",
-        []() {
-          return DynamicShapeDetector::GetMaxDifferentGraphs();
-        });
+        []() { return DynamicShapeDetector::GetMaxDifferentGraphs(); });
   m.def("_replace_xla_tensor",
         [](at::Tensor& self, const at::Tensor& source) -> at::Tensor& {
           return XLANativeFunctions::set_(self, source);