Add support for dynamo

Author: JackCaoG

test/dynamo/test_dynamo_graph_dump.py

@@ -32,10 +32,14 @@ def test_dump_graph_with_dynamo_execution(self):
     xla_y = torch.tensor(200.0).to(device)
     res_xla_dynamo = self.fn_simple_dynamo(xla_x, xla_y)
     with open(save_file, 'rb') as f:
-      current_line = sum(1 for line in f)
+      lines = f.readlines()
+      current_line = len(lines)
+      self.assertIn('Graph Hash:', lines[-5].decode())
     with open(save_file, 'rb') as f:
       res_xla_dynamo_2 = self.fn_simple_dynamo(xla_x, xla_y)
-      new_line = sum(1 for line in f)
+      lines = f.readlines()
+      new_line = len(lines)
+      self.assertIn('Graph Hash:', lines[-5].decode())
     self.assertGreater(new_line, current_line)
 
 

torch_xla/csrc/xla_graph_executor.cpp

@@ -146,8 +146,11 @@ void XLAGraphExecutor::DeviceContextArena::SaveGraphAsString(
       "";
   if (should_save_graph &&
       hash_to_graph_map.find(hash) == hash_to_graph_map.end()) {
-    hash_to_graph_map[hash] =
-        DebugUtil::GetTensorsGraphInfo(tensors, indices, format);
+    std::stringstream ss;
+    ss << DebugUtil::GetTensorsGraphInfo(tensors, indices, format);
+    ss << "Graph Hash: " << torch::lazy::HashToString(hash)
+       << "\n\n## END_GRAPH\n\n";
+    hash_to_graph_map[hash] = ss.str();
   }
 }