content: import argparse
  import logging
  import re
  from typing import Any, Dict, List, NamedTuple, Optional, Tuple
  
  import pydot
  
  _LOGGER: logging.Logger = logging.getLogger(__name__)
  """
  log_file is generated by tensorrt verbose logger during building engine.
  profile_file is generated by tensorrt profiler.
  
  Curretnly we support processing multiple logs in one log_file, which
  would generate multiple dot graphs. However, multiple engine profiles are not
  supported.
  
  Usage:
      python torch_tensorrt.fx/tools/engine_layer_visualize.py --log_file aaa --profile_file bbb
  """
  
  parser = argparse.ArgumentParser()
  parser.add_argument(
      "--log_file",
      type=str,
      default="",
      help="TensorRT VERBOSE logging when building engines.",
  )
  parser.add_argument(
      "--profile_file",
      type=str,
      default="",
      help="TensorRT execution context profiler output.",
  )
  args = parser.parse_args()
  
  
  class LayerInfo(NamedTuple):
      kernel_name: str
      layer_name: str
      tactic: str
      input_names: Optional[List[str]]
      input_types: Optional[List[str]]
      output_name: str
      output_type: str
      time: str
  
      @classmethod
      def from_string(cls, string, tactic_names, layer_times=None):
          input_names = []
          input_types = []
          kernel_name, layer_name, tactic, inputs, output_name, output_type = re.findall(
              "Layer\\((.+)\\): (.+), Tactic: (-?\\d+), (.+)? -> (.+)\\[(.+)\\]", string
          )[0]
  
          if kernel_name != "Constant":
              inputs = re.findall(
                  "[, ]*(.+?)\\[([Half|Float|Int8]+\\(\\d[,\\d]*\\))\\]", inputs
              )
              for input_name, input_type in inputs:
                  input_names.append(input_name)
                  input_types.append(input_type)
  
              if layer_name in tactic_names:
                  kernel_name = tactic_names[layer_name]
          else:
              input_names = input_types = None  # type:ignore[assignment]
  
          return cls(
              kernel_name,
              layer_name,
              tactic,
              input_names,
              input_types,
              output_name,
              output_type,
              layer_times[layer_name] if layer_times else "NA",
          )
  
  
  def build_node(layer):
      layer_name = layer.layer_name.replace("|", "\\|")
      label = f"{{{layer_name}|kernel: {layer.kernel_name}\\l|tactic: {layer.tactic}\\l|time: {layer.time}\\l}}"
      label = label.replace(">", "\\>")
      return pydot.Node(layer.layer_name, label=label, **style)
  
  
  def build_edge(layer, graph, reformat_layers, output_name2node, layer_name2node):
      if layer.input_names is None:
          return
  
      for input_name, input_type in zip(layer.input_names, layer.input_types):
          if input_name not in output_name2node:
              if input_name in reformat_layers:
                  from_node = pydot.Node(
                      input_name,
                      label="{reformatter|kernel: Reformat\\l|tactic: 0\\l}",
                      **style,
                  )
                  graph.add_node(from_node)
                  if reformat_layers[input_name][0] in output_name2node:
                      graph.add_edge(
                          pydot.Edge(
                              output_name2node[reformat_layers[input_name][0]],
                              from_node,
                              label=f"{reformat_layers[input_name][0]}\\l{reformat_layers[input_name][1]}\\l",
                          )
                      )
              else:
                  _LOGGER.info(f"Missing node {input_name}")
                  from_node = input_name
          else:
              from_node = output_name2node[input_name]
  
          edge_name = input_name.replace(">", "\\>")
          graph.add_edge(
              pydot.Edge(
                  from_node,
                  layer_name2node[layer.layer_name],
                  label=f"{edge_name}\\l{input_type}\\l",
              )
          )
  
  
  if args.profile_file != "":
      layer_times = {}
      with open(args.profile_file) as f:
          times = f.readlines()
  
      for t in times:
          t = t.strip("\n").split(": ")  # type: ignore[assignment]
          layer_times[": ".join(t[:-1])] = t[-1]
  else:
      layer_times = None  # type: ignore[assignment]
  
  if args.log_file != "":
      with open(args.log_file) as f:
          lines = f.readlines()
  
      graphs = []
      layers = []
      reformat_layers: Dict[str, Tuple[str, str]] = {}
      tactic_names: Dict[str, str] = {}
      layer_info_start = False
      tactic_name_start = False
  
      for line in lines:
          line = line.strip("\n")
  
          if layer_info_start:
              if "Layer(" in line:
                  layers.append(LayerInfo.from_string(line, tactic_names, layer_times))
              else:
                  layer_info_start = False
                  graphs.append((layers, reformat_layers))
                  layers = []
                  reformat_layers = {}
                  tactic_names = {}
  
          if tactic_name_start and "Set Tactic Name:" in line:
              layer_name, kernel_name, _ = re.findall(
                  "VERBOSE: (.*) Set Tactic Name: (.*) Tactic: (.*)$", line
              )[0]
              tactic_names[layer_name] = kernel_name
  
          # Some reformat layers aren't displayed in Engine Layer Information
          if "Adding reformat layer" in line:
              output_name, input_name, from_type, to_type = re.findall(
                  "reformat layer: (.+) \\((.+)\\) from (.+) to (.+)", line
              )[0]
              reformat_layers[output_name] = (input_name, from_type)
  
          if "Total Activation Memory:" in line:
              tactic_name_start = True
  
          if "Engine Layer Information" in line:
              layer_info_start = True
              tactic_name_start = False
  
      style = {
          "shape": "record",
          "fillcolor": "Salmon",
          "style": '"filled,rounded"',
          "fontcolor": "#000000",
      }
  
      dot_graphs: List[Any] = []
      i = 0
      for layers, reformat_layers in graphs:
          output_name2node = {}
          layer_name2node = {}
          dot_graph = pydot.Dot("Layer Graph")
  
          for layer in layers:
              node = build_node(layer)
              dot_graph.add_node(node)
              output_name2node[layer.output_name] = node
              layer_name2node[layer.layer_name] = node
  
          for layer in layers:
              build_edge(
                  layer, dot_graph, reformat_layers, output_name2node, layer_name2node
              )
  
          dot_graph.write_raw(f"EngineLayers_{i}.dot")
          i += 1
  
  if args.profile_file != "":
      est_reformat_time = 0.0
      est_total_time = 0.0
  
      for layer in layers:
          if layer.kernel_name == "Reformat":
              est_reformat_time += float(layer.time[:-2])
          est_total_time += float(layer.time[:-2])
  
      _LOGGER.info(f"Time Cost on Reformatting: {est_reformat_time} ms")
      _LOGGER.info(f"Total Time Cost: {est_total_time} ms")
  