ai-edge-torch-nightly 0.2.0.dev20240617__py3-none-any.whl → 0.2.0.dev20240625__py3-none-any.whl

ai_edge_torch/convert/conversion.py

@@ -30,6 +30,7 @@ from ai_edge_torch.convert.fx_passes import CanonicalizePass
 from ai_edge_torch.convert.fx_passes import InjectMlirDebuginfoPass
 from ai_edge_torch.convert.fx_passes import OptimizeLayoutTransposesPass
 from ai_edge_torch.convert.fx_passes import run_passes
+from ai_edge_torch.generative.fx_passes import run_generative_passes
 from ai_edge_torch.quantize import quant_config as qcfg
 
 os.environ["EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM"] = "1"
@@ -38,6 +39,7 @@ os.environ["EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM"] = "1"
 def _run_convert_passes(
     exported_program: ExportedProgram,
 ) -> ExportedProgram:
+  exported_program = run_generative_passes(exported_program)
   return run_passes(
       exported_program,
       [

ai_edge_torch/debug/__init__.py

@@ -13,4 +13,5 @@
 # limitations under the License.
 # ==============================================================================
 
+from .culprit import _search_model
 from .culprit import find_culprits

ai_edge_torch/debug/culprit.py

@@ -21,7 +21,7 @@ import io
 import operator
 import os
 import sys
-from typing import Any, Generator, List, Optional, Tuple
+from typing import Any, Callable, Generator, List, Optional, Tuple, Union
 
 from functorch.compile import minifier as fx_minifier
 import torch
@@ -85,10 +85,9 @@ def _tensor_to_buffer(t: torch.Tensor):
 
 
 @dataclasses.dataclass
-class Culprit:
+class SearchResult:
   graph_module: torch.fx.GraphModule
   inputs: Tuple[Any]
-  _runtime_errors: bool
 
   @property
   def graph(self) -> torch.fx.Graph:
@@ -98,6 +97,11 @@ class Culprit:
   def graph(self, fx_g: torch.fx.Graph):
     self.graph_module.graph = fx_g
 
+
+@dataclasses.dataclass
+class Culprit(SearchResult):
+  _runtime_errors: bool
+
   @property
   def stack_traces(self) -> List[str]:
     stack_traces = set()
@@ -342,42 +346,42 @@ def _fx_minifier_checker(fx_gm, inputs, runtime_errors=False):
   return False
 
 
-def find_culprits(
-    torch_model: torch.nn.Module,
-    args: Tuple[Any],
-    max_granularity: Optional[int] = None,
-    runtime_errors: bool = False,
+def _search_model(
+    predicate_f: Callable[[torch.fx.GraphModule, List[Any]], bool],
+    model: Union[torch.export.ExportedProgram, torch.nn.Module],
+    export_args: Tuple[Any] = None,
     *,
+    max_granularity: Optional[int] = None,
     enable_fx_minifier_logging: bool = False,
-) -> Generator[Culprit, None, None]:
-  """Finds culprits in the AI Edge Torch model conversion.
+) -> Generator[SearchResult, None, None]:
+  """Finds subgraphs in the torch model that satify a certain predicate function provided by the users.
 
   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
+    predicate_f: a predicate function the users specify.
+      It takes a FX (sub)graph and the inputs to this graph,
+      return True if the graph satisfies the predicate,
+      return False otherwise.
+    model: model in which to search subgraph.
+    export_args: A set of args to trace the model with,
+      i.e. 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.
+    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
+  if isinstance(model, torch.nn.Module):
+    try:
+      ep = torch.export.export(model, export_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
+  else:
+    ep = model
 
   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
@@ -405,13 +409,13 @@ def find_culprits(
           raw_min_fx_gm, raw_min_inputs = fx_minifier(
               fx_gm,
               fx_inputs,
-              fx_minifier_checker,
+              predicate_f,
               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)
+      yield SearchResult(min_fx_gm, min_inputs)
 
       fx_gm, fx_inputs = _erase_sub_gm_from_gm(
           fx_gm, fx_inputs, raw_min_fx_gm, raw_min_inputs
@@ -421,3 +425,40 @@ def find_culprits(
       if str(e) == "Input graph did not fail the tester" and found_culprits_num > 0:
         break
       raise e
+
+
+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.
+  """
+
+  fx_minifier_checker = functools.partial(
+      _fx_minifier_checker, runtime_errors=runtime_errors
+  )
+
+  for search_result in _search_model(
+      fx_minifier_checker,
+      torch_model,
+      args,
+      max_granularity=max_granularity,
+      enable_fx_minifier_logging=enable_fx_minifier_logging,
+  ):
+    yield Culprit(
+        search_result.graph_module, search_result.inputs, _runtime_errors=runtime_errors
+    )

ai_edge_torch/debug/test/test_search_model.py

@@ -0,0 +1,50 @@
+# 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 unittest
+
+import torch
+
+from ai_edge_torch.debug import _search_model
+
+
+class TestSearchModel(unittest.TestCase):
+
+  def test_search_model_with_ops(self):
+    class MultipleOpsModel(torch.nn.Module):
+
+      def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        sub_0 = x - 1
+        add_0 = y + 1
+        mul_0 = x * y
+        add_1 = sub_0 + add_0
+        mul_1 = add_0 * mul_0
+        sub_1 = add_1 - mul_1
+        return sub_1
+
+    model = MultipleOpsModel().eval()
+    args = (torch.rand(10), torch.rand(10))
+
+    def find_subgraph_with_sub(fx_gm, inputs):
+      return torch.ops.aten.sub.Tensor in [n.target for n in fx_gm.graph.nodes]
+
+    results = list(_search_model(find_subgraph_with_sub, model, args))
+    self.assertEqual(len(results), 2)
+    self.assertIn(torch.ops.aten.sub.Tensor, [n.target for n in results[0].graph.nodes])
+
+
+if __name__ == "__main__":
+  unittest.main()

ai_edge_torch/generative/fx_passes/__init__.py

@@ -0,0 +1,31 @@
+# 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 torch
+
+from ai_edge_torch.convert.fx_passes import CanonicalizePass
+from ai_edge_torch.convert.fx_passes import run_passes
+from ai_edge_torch.generative.fx_passes.remove_sdpa_zero_mask_pass import RemoveSDPACompositeZeroMaskPass  # NOQA
+
+
+def run_generative_passes(
+    exported_program: torch.export.ExportedProgram,
+) -> torch.export.ExportedProgram:
+  return run_passes(
+      exported_program,
+      [
+          RemoveSDPACompositeZeroMaskPass(),
+          CanonicalizePass(),
+      ],
+  )

ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py

@@ -0,0 +1,47 @@
+# 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 torch
+
+from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassBase
+from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassResult  # NOQA
+
+
+class RemoveSDPACompositeZeroMaskPass(ExportedProgramPassBase):
+
+  def is_zero_tensor_node(self, node: torch.fx.Node):
+    return node.target == torch.ops.aten.zeros.default
+
+  def call(self, exported_program: torch.export.ExportedProgram):
+    graph = exported_program.graph_module.graph
+    for node in graph.nodes:
+      if not (
+          node.op == "call_function"
+          and node.target == torch.ops.xla.mark_tensor.default
+      ):
+        continue
+
+      source, name, io_position, id, is_input = node.args[:5]
+      # Composite info:
+      # - name: odml.scaled_dot_product_attention
+      # - inputs: q, k, v, mask
+      if name == "odml.scaled_dot_product_attention" and is_input and io_position == 3:
+        if self.is_zero_tensor_node(source):
+          # Remove the mark_tensor call on the mask input by
+          # replacing the target with an identity function.
+          node.target = lambda *args, **kwargs: args[0]
+
+    exported_program.graph_module.graph.lint()
+    exported_program.graph_module.recompile()
+    return ExportedProgramPassResult(exported_program, True)

{ai_edge_torch_nightly-0.2.0.dev20240617.dist-info → ai_edge_torch_nightly-0.2.0.dev20240625.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.2.0.dev20240617
+Version: 0.2.0.dev20240625
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.2.0.dev20240617.dist-info → ai_edge_torch_nightly-0.2.0.dev20240625.dist-info}/RECORD

@@ -1,7 +1,7 @@
 ai_edge_torch/__init__.py,sha256=FPMmuFU3pyMREtjB_san1fy_0PFtAsgA0VZfOYvDrb4,1008
 ai_edge_torch/model.py,sha256=kmcgELjsYl8YzF8nUF6P7q4i8MWS-pLGpfsy-yTUXmE,4243
 ai_edge_torch/convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/convert/conversion.py,sha256=GN2Js232u_5Y118wg3qIfEoYewxbxLl3TpSnO6osi8c,4029
+ai_edge_torch/convert/conversion.py,sha256=8K8jQuaCjlUWoj7jiimxp_zpN6mYThLOcQ858UDcYnE,4159
 ai_edge_torch/convert/conversion_utils.py,sha256=9BqCL38DErv1vEVGtT3BIJVhdwZjw2EQ-_m5UpvVVYE,11294
 ai_edge_torch/convert/converter.py,sha256=bjj5TV5_g4sGyuSh8ThEDydlNMqhkGSY4SzXK6vwhqI,6927
 ai_edge_torch/convert/fx_passes/__init__.py,sha256=EPs4PSIDLuRH5EBETi6deaOvaaf_Q4xD3_9NVcR7x8o,2810
@@ -24,11 +24,12 @@ ai_edge_torch/convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrK
 ai_edge_torch/convert/test/test_convert.py,sha256=2qPmmGqnfV_o1gfsSdjGq3-JR1b323ligiy5MdAv9NA,8021
 ai_edge_torch/convert/test/test_convert_composites.py,sha256=_Ojc-H6GOS5s8ek3_8eRBL_AiCs-k3srziPJ2R4Ulrg,7255
 ai_edge_torch/convert/test/test_convert_multisig.py,sha256=kMaGnHe9ylfyU68qCifYcaGwJqyejKz--QQt9jS2oUA,4537
-ai_edge_torch/debug/__init__.py,sha256=TKvmnjVk3asvYcVh6C-LPr6srgAF_nppSAupWEXqwPY,707
-ai_edge_torch/debug/culprit.py,sha256=vklaxBUfINdo44OsH7csILK70N41gEThCGchGEfbTZw,12789
+ai_edge_torch/debug/__init__.py,sha256=N05Mmvi41KgSuK0JhuMejERESgP8QekiGdp9_PEyuKU,742
+ai_edge_torch/debug/culprit.py,sha256=urtCKPXORPvn6oyDxDSCSjgvngUnjjcsUMwAOeIl15E,14236
 ai_edge_torch/debug/utils.py,sha256=hjVmQVVl1dKxEF0D6KB4a3ouQ3wBkTsebOX2YsUObZM,1430
 ai_edge_torch/debug/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/debug/test/test_culprit.py,sha256=9An_n9p_RWTAYdHYTCO-__EJlbnjclCDo8tDhOzMlwk,3731
+ai_edge_torch/debug/test/test_search_model.py,sha256=0guAEon5cvwBpPXk6J0wVOKj7TXMDaiuomEEQmHgO5o,1590
 ai_edge_torch/experimental/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -64,6 +65,8 @@ ai_edge_torch/generative/examples/test_models/toy_model_with_kv_cache.py,sha256=
 ai_edge_torch/generative/examples/tiny_llama/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_edge_torch/generative/examples/tiny_llama/convert_to_tflite.py,sha256=E4I5OlC4zyl5cxiiu7uTED-zcwYRu210lP1zuT3xLBE,2566
 ai_edge_torch/generative/examples/tiny_llama/tiny_llama.py,sha256=IFRLPG9wz_aLl_zV_6CETCjSM03ukA6bZqqyDLVACuw,5651
+ai_edge_torch/generative/fx_passes/__init__.py,sha256=aXvYiaHDvETIrh0Q9DDZA_ZBiazGk80DT6nt7lLtC1o,1172
+ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py,sha256=IehLwFNwa0C9fnk1pmNmyfuAwwWbuwdyKy46BSqNVdI,1948
 ai_edge_torch/generative/layers/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/attention.py,sha256=AW0Qo3uOIe6p1rJNJ6zR_r4fqL2y-6QJHh0yUd-5Yb0,11966
 ai_edge_torch/generative/layers/attention_utils.py,sha256=hXhuyKblPPxKIRzlAf1YNlwHgpbj-6nReRLhRHELx5k,6350
@@ -109,8 +112,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=eO9Ra160ITjQSyRBEGy6nNIVH3gYacSWDd
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=EIyKz-HY70DguWuSrJal8LpYXQ5ZSEUf3ZrVl7jikFM,4286
-ai_edge_torch_nightly-0.2.0.dev20240617.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240617.dist-info/METADATA,sha256=Z9rUO2CabVbBpydpRk8OxNlwK4yznGCb2QHGlJhqRsM,1748
-ai_edge_torch_nightly-0.2.0.dev20240617.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240617.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240617.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240625.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240625.dist-info/METADATA,sha256=BMjo3IdVwNBpXCIGJC01y1buq5mPWWPyAJV1LFvxsJw,1748
+ai_edge_torch_nightly-0.2.0.dev20240625.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240625.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240625.dist-info/RECORD,,