ai-edge-torch-nightly 0.1.dev202405131930__py3-none-any.whl

ai_edge_torch/debug/culprit.py

@@ -0,0 +1,423 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import contextlib
+import copy
+import dataclasses
+import functools
+import io
+import operator
+import os
+import sys
+from typing import Any, Generator, List, Optional, Tuple
+
+from functorch.compile import minifier as fx_minifier
+import torch
+from torch._functorch import aot_autograd
+import torch.utils._pytree as pytree
+
+import ai_edge_torch
+from ai_edge_torch.debug import utils
+
+_torch_float_dtypes = {
+    torch.float32,
+    torch.float,
+    torch.float64,
+    torch.double,
+    torch.float16,
+    torch.half,
+    torch.bfloat16,
+}
+_torch_int_dtypes = {
+    torch.uint8,
+    torch.int8,
+    torch.int16,
+    torch.short,
+    torch.int32,
+    torch.int,
+    torch.int64,
+    torch.long,
+}
+
+_fx_op_runner = {
+    "call_function": lambda target, args, kwargs: target(*args, **kwargs),
+    "call_method": lambda target, args, kwargs: getattr(args[0], target)(
+        *args[1:], **kwargs
+    ),
+}
+
+_CULPRIT_GRAPH_MODULE_NAME = "CulpritGraphModule"
+
+
+def _get_shape_str(t: torch.Tensor):
+  return f"({', '.join(map(str, t.shape))},)"
+
+
+def _tensor_to_random_tensor_call(t: torch.Tensor):
+  shape_str = _get_shape_str(t)
+  if t.dtype in _torch_float_dtypes:
+    return f"torch.randn({shape_str}, dtype={t.dtype})"
+  elif t.dtype in _torch_int_dtypes:
+    return f"torch.randint(0, 10, {shape_str}, dtype={t.dtype})"
+  elif t.dtype == torch.bool:
+    return f"torch.randint(0, 2, {shape_str}, dtype={t.dtype})"
+  else:
+    raise ValueError(f"Unsupported dtype: {t.dtype}")
+
+
+def _tensor_to_buffer(t: torch.Tensor):
+  buff = io.BytesIO()
+  torch.save(t, buff)
+  buff.seek(0)
+  return buff.read()
+
+
+@dataclasses.dataclass
+class Culprit:
+  graph_module: torch.fx.GraphModule
+  inputs: Tuple[Any]
+  _runtime_errors: bool
+
+  @property
+  def graph(self) -> torch.fx.Graph:
+    return self.graph_module.graph
+
+  @graph.setter
+  def graph(self, fx_g: torch.fx.Graph):
+    self.graph_module.graph = fx_g
+
+  @property
+  def stack_traces(self) -> List[str]:
+    stack_traces = set()
+    for node in self.graph.nodes:
+      if node.op.startswith("call_") and "stack_trace" in node.meta:
+        stack_traces.add(node.meta["stack_trace"])
+    return list(stack_traces)
+
+  def print_readable(self, print_output=True):
+    """Print the Python code for culprit graph module and sample args.
+
+    Args:
+      print_output: bool - If true, prints the code to stdout. Otherwise returns
+        the code in a str.
+    """
+    # TODO (b/321263453): Support Python code gen with sample arg tensor values.
+    random_inputs = True
+
+    graph_module_code = self.graph_module.print_readable(print_output=False).rstrip()
+
+    input_strs = []
+    for value in self.inputs:
+      if torch.is_tensor(value):
+        if not random_inputs:
+          input_strs.append(f"# size={_get_shape_str(value)}, dtype={value.dtype}")
+          input_strs.append(f"torch.load(io.BytesIO({_tensor_to_buffer(value)})),")
+        else:
+          input_strs.append(_tensor_to_random_tensor_call(value) + ",")
+      else:
+        input_strs.append(str(value) + ",")
+
+    inputs_code = (
+        "_args = (\n" + "\n".join([" " * 4 + code for code in input_strs]) + "\n)"
+    )
+
+    code = graph_module_code + "\n\n" + inputs_code
+    if print_output:
+      print(code)
+    else:
+      return code
+
+  def print_code(self, print_output=True):
+    """Print the Python code for culprit graph module, sample args, and AI
+    Edge Torch conversion that will fail with the error.
+
+    Args:
+      print_output: bool - If true, prints the code to stdout. Otherwise returns
+        the code in a str.
+    """
+    definitions = self.print_readable(print_output=False)
+    code = (
+        "import torch\n"
+        + "from torch import device\n"
+        + "import ai_edge_torch\n\n"
+        + definitions
+        + f"\n\n_edge_model = ai_edge_torch.convert({_CULPRIT_GRAPH_MODULE_NAME}().eval(), _args)\n"
+    )
+    if self._runtime_errors:
+      code += "_edge_model(*_args)\n"
+
+    if print_output:
+      print(code)
+    else:
+      return code
+
+  @property
+  def code(self):
+    return self.print_code(print_output=False)
+
+  def __repr__(self):
+    return self.print_readable(print_output=False)
+
+  def __str__(self):
+    return self.print_readable(print_output=False)
+
+
+def _normalize_getitem_nodes(fx_gm: torch.fx.GraphModule):
+  """
+  This function turns all operator getitem nodes in ExportedProgram FX graph to
+  new nodes composed of "computation + getitem". The normalization duplicates
+  some computations in the graph but would make the graph more friendly for
+  partitioning in FX minifier.
+  """
+
+  fx_gm = copy.deepcopy(fx_gm)
+  graph = fx_gm.graph
+  for n in graph.nodes:
+    if n.target != operator.getitem:
+      continue
+
+    src_n, key = n.args
+    if src_n.op not in _fx_op_runner:
+      continue
+
+    runner = _fx_op_runner.get(src_n.op)
+
+    with graph.inserting_after(n):
+      new_n = graph.call_function(
+          lambda src_target, key, args, kwargs: operator.getitem(
+              runner(src_target, args, kwargs), key
+          ),
+          (src_n.target, key, src_n.args, src_n.kwargs),
+      )
+      n.replace_all_uses_with(new_n)
+
+  graph.eliminate_dead_code()
+  fx_gm.graph = graph
+  return fx_gm
+
+
+def _erase_unused_inputs(fx_gm: torch.fx.GraphModule, inputs: Tuple[torch.Tensor]):
+  fx_gm = copy.deepcopy(fx_gm)
+  inputs = tuple(inputs)
+  args = fx_gm.graph.process_inputs(*inputs)
+  args_iter = iter(args)
+
+  graph = fx_gm.graph
+  new_inputs = []
+  for n in graph.nodes:
+    if n.op == "placeholder":
+      if n.target.startswith("*"):
+        new_inputs += list(args_iter)
+      elif len(n.users) > 0:
+        new_inputs.append(next(args_iter))
+      else:
+        graph.erase_node(n)
+        next(args_iter)
+  new_inputs = tuple(new_inputs)
+  fx_gm.graph = graph
+  return fx_gm, new_inputs
+
+
+def _lift_dead_ops_to_outputs(fx_gm: torch.fx.GraphModule):
+  fx_gm = copy.deepcopy(fx_gm)
+
+  new_outputs = []
+  graph = fx_gm.graph
+  nodes = list(graph.nodes)
+  assert nodes[-1].op == "output" and sum(n.op == "output" for n in nodes) == 1
+  for node in nodes:
+    if node.op not in ("placeholder", "output") and len(node.users) == 0:
+      new_outputs.append(node)
+
+  output_node = nodes[-1]
+  # FX output node returns the first arg as is.
+  # ref: https://github.com/pytorch/pytorch/blob/1a578df57cc0f417f671634e564c62ef5d9a97e2/torch/fx/interpreter.py#L337
+  new_outputs, _ = pytree.tree_flatten([new_outputs, output_node.args[0]])
+  output_node.update_arg(0, tuple(new_outputs))
+
+  fx_gm.graph = graph
+  return fx_gm
+
+
+def _erase_trivial_outputs(fx_gm: torch.fx.GraphModule):
+  """Remove output nodes directly connected to an input node."""
+  fx_gm = copy.deepcopy(fx_gm)
+
+  graph = fx_gm.graph
+  nodes = list(graph.nodes)
+  assert nodes[-1].op == "output" and sum(n.op == "output" for n in nodes) == 1
+  output_node = nodes[-1]
+
+  outputs, _ = pytree.tree_flatten(output_node.args[0])
+  new_outputs = [output for output in outputs if output.op != "placeholder"]
+  output_node.update_arg(0, tuple(new_outputs))
+
+  fx_gm.recompile()
+  return fx_gm
+
+
+def _erase_sub_gm_from_gm(
+    fx_gm: torch.fx.GraphModule,
+    fx_inputs: Tuple[torch.Tensor],
+    sub_gm: torch.fx.GraphModule,
+    sub_inputs: Tuple[torch.Tensor],
+):
+  fx_gm = copy.deepcopy(fx_gm)
+  fx_inputs = list(fx_inputs)
+
+  class EraseNodeInterpreter(torch.fx.Interpreter):
+
+    def run_node(self, node):
+      nonlocal fx_gm, fx_inputs
+      res = super().run_node(node)
+      if node.op not in ("placeholder", "output"):
+        to_erase = next(m for m in fx_gm.graph.nodes if m.name == node.name)
+        # Raise the output (tensor) of the erased node to be an input of
+        # the new model graph. Some raised inputs may become unused later
+        # when all the users are within the erased subgraph, those inputs
+        # will be removed by the followed `_erase_unused_inputs` pass.
+        with fx_gm.graph.inserting_before(to_erase):
+          new_input = fx_gm.graph.placeholder(node.name + "__value")
+        to_erase.replace_all_uses_with(new_input)
+
+        fx_gm.graph.erase_node(to_erase)
+        fx_inputs.append(res)
+      return res
+
+  interpreter = EraseNodeInterpreter(sub_gm)
+  interpreter.run(*sub_inputs)
+
+  fx_gm.graph.lint()
+  fx_gm.recompile()
+
+  # Ops prior to the erased subgraph may be dangling. Lift them as outputs.
+  fx_gm = _lift_dead_ops_to_outputs(fx_gm)
+  fx_gm = _erase_trivial_outputs(fx_gm)
+  fx_gm, fx_inputs = _erase_unused_inputs(fx_gm, fx_inputs)
+
+  fx_gm.graph.lint()
+  fx_gm.recompile()
+  return fx_gm, fx_inputs
+
+
+def _normalize_minified_fx_gm(fx_gm: torch.fx.GraphModule, inputs: Tuple[torch.Tensor]):
+  fx_gm, inputs = _erase_unused_inputs(fx_gm, inputs)
+  fx_gm = _lift_dead_ops_to_outputs(fx_gm)
+  fx_gm, _ = aot_autograd.aot_export_module(fx_gm, inputs, trace_joint=False)
+  fx_gm.__class__.__name__ = _CULPRIT_GRAPH_MODULE_NAME
+  return fx_gm, inputs
+
+
+def _fx_minifier_checker(fx_gm, inputs, runtime_errors=False):
+  fx_gm, inputs = _normalize_minified_fx_gm(fx_gm, inputs)
+
+  trivial_aten_ops = {
+      torch.ops.aten.view,
+      torch.ops.aten.view.default,
+  }
+  if all(
+      node.op in ("placeholder", "output") or node.target in trivial_aten_ops
+      for node in fx_gm.graph.nodes
+  ):
+    return False
+
+  try:
+    edge_model = ai_edge_torch.convert(fx_gm.eval(), inputs)
+    if runtime_errors:
+      edge_model(*inputs)
+  except Exception as err:
+    return True
+  return False
+
+
+def find_culprits(
+    torch_model: torch.nn.Module,
+    args: Tuple[Any],
+    max_granularity: Optional[int] = None,
+    runtime_errors: bool = False,
+    *,
+    enable_fx_minifier_logging: bool = False,
+) -> Generator[Culprit, None, None]:
+  """Finds culprits in the AI Edge Torch model conversion.
+
+  Args:
+    torch_model: model to export and save
+    args: A set of args to trace the model with, i.e.
+      torch_model(*args) must run
+    max_granularity - FX minifier arg. The maximum granularity (number of nodes)
+      in the returned ATen FX subgraph of the culprit.
+    runtime_errors: If true, find culprits for Python runtime errors
+      with converted model.
+    enable_fx_minifier_logging: If true, allows the underlying FX minifier to log
+      the progress.
+  """
+
+  try:
+    ep = torch.export.export(torch_model, args)
+  except Exception as err:
+    raise ValueError(
+        "Your model is not exportable by torch.export.export. Please modify your model to be torch-exportable first."
+    ) from err
+
+  fx_gm, fx_inputs = utils.exported_program_to_fx_graph_module_and_inputs(ep)
+  fx_gm = _normalize_getitem_nodes(fx_gm)
+
+  fx_minifier_checker = functools.partial(
+      _fx_minifier_checker, runtime_errors=runtime_errors
+  )
+
+  # HACK: temporarily disable XLA_HLO_DEBUG so that fx_minifier won't dump
+  # intermediate stablehlo files to storage.
+  # https://github.com/pytorch/pytorch/blob/main/torch/_functorch/fx_minifier.py#L440
+  @contextlib.contextmanager
+  def disable_xla_hlo_debug():
+    xla_hlo_debug_value = None
+    if "XLA_HLO_DEBUG" in os.environ:
+      xla_hlo_debug_value = os.environ["XLA_HLO_DEBUG"]
+      del os.environ["XLA_HLO_DEBUG"]
+
+    try:
+      yield None
+    finally:
+      if xla_hlo_debug_value is not None:
+        os.environ["XLA_HLO_DEBUG"] = xla_hlo_debug_value
+
+  found_culprits_num = 0
+  while True:
+    try:
+      with disable_xla_hlo_debug(), open(os.devnull, "w") as devnull:
+        with contextlib.nullcontext() if enable_fx_minifier_logging else utils.redirect_stdio(
+            stdout=devnull,
+            stderr=devnull,
+        ):
+          raw_min_fx_gm, raw_min_inputs = fx_minifier(
+              fx_gm,
+              fx_inputs,
+              fx_minifier_checker,
+              max_granularity=max_granularity,
+          )
+
+      min_fx_gm, min_inputs = _normalize_minified_fx_gm(raw_min_fx_gm, raw_min_inputs)
+      found_culprits_num += 1
+      yield Culprit(min_fx_gm, min_inputs, _runtime_errors=runtime_errors)
+
+      fx_gm, fx_inputs = _erase_sub_gm_from_gm(
+          fx_gm, fx_inputs, raw_min_fx_gm, raw_min_inputs
+      )
+
+    except RuntimeError as e:
+      if str(e) == "Input graph did not fail the tester" and found_culprits_num > 0:
+        break
+      raise e

ai_edge_torch/debug/test/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

ai_edge_torch/debug/test/test_culprit.py

@@ -0,0 +1,133 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+
+import ast
+import io
+import sys
+import unittest
+
+import torch
+
+from ai_edge_torch.debug import find_culprits
+
+_test_culprit_lib = torch.library.Library("test_culprit", "DEF")
+
+_test_culprit_lib.define("non_lowerable_op(Tensor x) -> Tensor")
+
+
+@torch.library.impl(_test_culprit_lib, "non_lowerable_op", "CompositeExplicitAutograd")
+def non_lowerable_op(x):
+  if x.max() > 10.0:
+    return x + 1.0
+  return x
+
+
+@torch.library.impl(_test_culprit_lib, "non_lowerable_op", "Meta")
+def non_lowerable_op_meta(x):
+  return torch.empty_like(x)
+
+
+class BadModel(torch.nn.Module):
+
+  def forward(self, x):
+    x = x + 1
+    x = torch.ops.test_culprit.non_lowerable_op.default(x)
+    return x
+
+
+class TestCulprit(unittest.TestCase):
+
+  def test_find_culprits(self):
+    model = BadModel().eval()
+    args = (torch.rand(10),)
+
+    culprits = list(find_culprits(model, args))
+    self.assertEqual(len(culprits), 1)
+    self.assertIn(
+        torch.ops.test_culprit.non_lowerable_op.default,
+        [n.target for n in culprits[0].graph.nodes],
+    )
+
+  def test_valid_culprit_readable(self):
+    model = BadModel().eval()
+    args = (torch.rand(10),)
+
+    culprits = list(find_culprits(model, args))
+    self.assertEqual(len(culprits), 1)
+
+    code = culprits[0].print_readable(print_output=False)
+
+    # The code should be a valid Python code
+    ast.parse(code)
+
+  def test_valid_culprit_code(self):
+    model = BadModel().eval()
+    args = (torch.rand(10),)
+
+    culprits = list(find_culprits(model, args))
+    self.assertEqual(len(culprits), 1)
+
+    code = culprits[0].print_code(print_output=False)
+
+    # The code should be a valid Python code
+    ast.parse(code)
+
+  def test_find_multiple_culprits(self):
+    class MultiBadOpsModel(torch.nn.Module):
+
+      def forward(self, x):
+        x = x + 1
+        a = torch.ops.test_culprit.non_lowerable_op.default(x)
+        b = torch.ops.test_culprit.non_lowerable_op.default(x)
+        c = a + b
+        d = torch.ops.test_culprit.non_lowerable_op.default(c)
+        return d
+
+    model = MultiBadOpsModel().eval()
+    args = (torch.rand(10),)
+
+    culprits = list(find_culprits(model, args))
+    self.assertEqual(len(culprits), 3)
+    for culprit in culprits:
+      self.assertIn(
+          torch.ops.test_culprit.non_lowerable_op.default,
+          [n.target for n in culprit.graph.nodes],
+      )
+
+  def test_find_culprits_with_trivial_inputs_outputs(self):
+
+    class MultiBadOpsModel(torch.nn.Module):
+
+      def forward(self, x, y, z):
+        x = x + 1
+        a = torch.ops.test_culprit.non_lowerable_op.default(x)
+        b = torch.ops.test_culprit.non_lowerable_op.default(y)
+        return a, b, x, y, a, b
+
+    model = MultiBadOpsModel().eval()
+    args = (torch.rand(10), torch.rand(10), torch.rand(10))
+
+    culprits = list(find_culprits(model, args))
+    self.assertEqual(len(culprits), 2)
+    for culprit in culprits:
+      self.assertIn(
+          torch.ops.test_culprit.non_lowerable_op.default,
+          [n.target for n in culprit.graph.nodes],
+      )
+
+
+if __name__ == "__main__":
+  unittest.main()

ai_edge_torch/debug/utils.py

@@ -0,0 +1,48 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import contextlib
+import sys
+
+import torch
+from torch.export.graph_signature import InputKind
+import torch.fx._pytree as fx_pytree
+from torch.utils import _pytree as pytree
+
+
+def exported_program_to_fx_graph_module_and_inputs(ep: torch.export.ExportedProgram):
+  fx_gm = ep.graph_module
+  fx_inputs = pytree.tree_map(
+      torch.tensor, ep._graph_module_flat_inputs(*ep.example_inputs)
+  )
+  return fx_gm, fx_inputs
+
+
+@contextlib.contextmanager
+def redirect_stdio(stdout, stderr):
+  old_stdout = sys.stdout
+  old_stderr = sys.stderr
+
+  old_stdout.flush()
+  old_stderr.flush()
+
+  sys.stdout = stdout
+  sys.stderr = stderr
+  try:
+    yield stdout, stderr
+  finally:
+    stdout.flush()
+    stderr.flush()
+    sys.stdout = old_stdout
+    sys.stderr = old_stderr

ai_edge_torch/experimental/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

ai_edge_torch/generative/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

ai_edge_torch/generative/examples/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

ai_edge_torch/generative/examples/gemma/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================