ai-edge-torch-nightly 0.2.0.dev20240806__py3-none-any.whl → 0.2.0.dev20240807__py3-none-any.whl

ai_edge_torch/{convert → _convert}/fx_passes/build_interpolate_composite_pass.py

@@ -12,19 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Build interpolate composite pass."""
 
 import functools
 
-from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassBase
-from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassResult  # NOQA
+from ai_edge_torch._convert.fx_passes._pass_base import ExportedProgramPassBase
+from ai_edge_torch._convert.fx_passes._pass_base import ExportedProgramPassResult  # NOQA
 from ai_edge_torch.hlfb import mark_pattern
+from ai_edge_torch.hlfb.mark_pattern import pattern as pattern_module
 import torch
 
 # For torch nightly released after mid June 2024,
 # torch.nn.functional.interpolate no longer gets exported into decomposed graph
-# but single aten op torch.ops.aten.upsample_nearest2d.vec/torch.ops.aten.upsample_bilinear2d.vec.
-# This behavior would our pattern matching based composite builder.
-# It requires the pattern and model graph to get decomposed first for backward compatibility.
+# but a single aten op:
+# torch.ops.aten.upsample_nearest2d.vec/torch.ops.aten.upsample_bilinear2d.vec.
+# This would interefere with our pattern matching based composite builder.
+# Here we register the now missing decompositions first.
 _INTERPOLATE_DECOMPOSITIONS = torch._decomp.get_decompositions([
     torch.ops.aten.upsample_bilinear2d.vec,
     torch.ops.aten.upsample_nearest2d.vec,
@@ -33,7 +36,7 @@ _INTERPOLATE_DECOMPOSITIONS = torch._decomp.get_decompositions([
 
 @functools.cache
 def _get_upsample_bilinear2d_pattern():
-  pattern = mark_pattern.Pattern(
+  pattern = pattern_module.Pattern(
       "odml.upsample_bilinear2d",
       lambda x: torch.nn.functional.interpolate(
           x, scale_factor=2, mode="bilinear", align_corners=False
@@ -56,7 +59,7 @@ def _get_upsample_bilinear2d_pattern():
 
 @functools.cache
 def _get_upsample_bilinear2d_align_corners_pattern():
-  pattern = mark_pattern.Pattern(
+  pattern = pattern_module.Pattern(
       "odml.upsample_bilinear2d",
       lambda x: torch.nn.functional.interpolate(
           x, scale_factor=2, mode="bilinear", align_corners=True
@@ -79,7 +82,7 @@ def _get_upsample_bilinear2d_align_corners_pattern():
 
 @functools.cache
 def _get_interpolate_nearest2d_pattern():
-  pattern = mark_pattern.Pattern(
+  pattern = pattern_module.Pattern(
       "tfl.resize_nearest_neighbor",
       lambda x: torch.nn.functional.interpolate(
           x, scale_factor=2, mode="nearest"

ai_edge_torch/{convert → _convert}/fx_passes/canonicalize_pass.py

@@ -13,8 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassBase
-from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassResult  # NOQA
+from ai_edge_torch._convert.fx_passes import _pass_base
 import torch
 from torch.export import ExportedProgram
 
@@ -29,8 +28,8 @@ _dummy_decomp_table = {
 }
 
 
-class CanonicalizePass(ExportedProgramPassBase):
+class CanonicalizePass(_pass_base.ExportedProgramPassBase):
 
   def call(self, exported_program: ExportedProgram):
     exported_program = exported_program.run_decompositions(_dummy_decomp_table)
-    return ExportedProgramPassResult(exported_program, True)
+    return _pass_base.ExportedProgramPassResult(exported_program, True)

ai_edge_torch/{convert → _convert}/fx_passes/inject_mlir_debuginfo_pass.py

@@ -13,11 +13,11 @@
 # limitations under the License.
 # ==============================================================================
 
+from ai_edge_torch import lowertools
 import torch
 from torch.fx.passes.infra.pass_base import PassBase
 from torch.fx.passes.infra.pass_base import PassResult
 import torch.utils._pytree as pytree
-import torch_xla.experimental.xla_mlir_debuginfo  # Import required to register torch.ops.xla.write_mlir_debuginfo
 
 
 def _get_mlir_debuginfo(node: torch.fx.Node):
@@ -54,7 +54,7 @@ def _wrap_call_function_node_with_debuginfo_writer(node: torch.fx.GraphModule):
     outputs = target(*args, **kwargs)
     outputs = pytree.tree_map_only(
         torch.Tensor,
-        lambda x: torch.ops.xla.write_mlir_debuginfo(x, debuginfo),
+        lambda x: lowertools.write_mlir_debuginfo_op(x, debuginfo),
         outputs,
     )
     return outputs

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+# ==============================================================================
+
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass.pass_body import OptimizeLayoutTransposesPass  # NOQA

ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_check.py

@@ -12,13 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Layout check for the optimized layout transposes pass."""
+
 import dataclasses
 import operator
 
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_mark  # NOQA
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_rewrite  # NOQA
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import utils  # NOQA
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass.op_func_registry import OpFuncRegistry  # NOQA
+from ai_edge_torch import lowertools
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_rewrite
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import utils
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass.op_func_registry import OpFuncRegistry
 import torch
 from torch.fx import Node
 
@@ -205,7 +207,7 @@ def _aten_native_group_norm_checker(node):
 # ==== Ops must be NCHW
 
 
-@nhwcable_node_checkers.register(torch.ops.xla.mark_tensor)
+@nhwcable_node_checkers.register(lowertools.mark_tensor_op)
 @nhwcable_node_checkers.register(utils.tensor_to_nchw)
 @nhwcable_node_checkers.register(utils.tensor_to_nhwc)
 @nhwcable_node_checkers.register("output")

ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_mark.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Layout mark for the optimized layout transposes pass."""
+
 import torch
 
 # Tag which is added to a node's meta to indicate that is is part of the NHWC

ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_partitioners/__init__.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Layout partitioners."""
 
 from . import greedy
 from . import min_cut

ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py

@@ -12,23 +12,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Greedy partitioning algorithm."""
 
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_check  # NOQA
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_mark  # NOQA
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_check
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_mark
 import torch
 
 
 def partition(graph_module: torch.fx.GraphModule):
-  """Partition the graph module into NHWC and non-NHWC subgraphs, and mark
-  nodes in the NHWC partitions.
+  """Partition the graph module into NHWC and non-NHWC subgraphs.
+
+  Partition the graph module into NHWC and non-NHWC subgraphs and mark nodes in
+  the NHWC partitions.
 
   Implements O(|V|) greedy partitioning algorithm.
-  See go/pytorch-layout-transpose-optimization for more details.
+
+  Args:
+    graph_module: The graph module to be partitioned.
+
+  Returns:
+    The partitioned graph module.
   """
   graph = graph_module.graph
 
   for node in list(graph.nodes):
-    if len(node.all_input_nodes) == 0:
+    if not node.all_input_nodes:
       # This node has no inputs so we don't need to change anything
       continue
 

ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py

@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Min cut solver for partitioning the graph module into NHWC and non-NHWC subgraphs."""
 
 import collections
 import dataclasses
-import itertools
 
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_check  # NOQA
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_mark  # NOQA
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_check  # NOQA
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_mark  # NOQA
 import numpy as np
 import scipy
 import torch
@@ -26,13 +26,12 @@ import torch
 
 def can_partition(graph_module: torch.fx.GraphModule):
   """Returns true if the input graph_module can be partitioned by min cut solver
+
   in a reasonable time.
 
   The min cut solver implements O(|V|^2|E|) Dinic's algorithm, which may
   take a long time to complete for large graph module. This function determines
   whether the graph module can be partitioned by the graph module size.
-
-  See go/pytorch-layout-transpose-optimization for more details.
   """
   graph = graph_module.graph
   n_nodes = len(graph.nodes)
@@ -137,10 +136,10 @@ class MultiUsersDummyNode:
 
 def partition(graph_module: torch.fx.GraphModule):
   """Partition the graph module into NHWC and non-NHWC subgraphs, and mark
+
   nodes in the NHWC partitions.
 
   Implements O(|V|^2|E|) min-cut (optimal) partitioning algorithm.
-  See go/pytorch-layout-transpose-optimization for more details.
   """
   graph = graph_module.graph
 

ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py

@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Layout rewrite for the optimized layout transposes pass."""
+
 import operator
 
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_mark  # NOQA
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import utils  # NOQA
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass.op_func_registry import OpFuncRegistry  # NOQA
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_mark
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import op_func_registry
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import utils
 import torch
-from torch.fx import Node
 import torch.utils._pytree as pytree
 
 aten = torch.ops.aten
@@ -26,7 +27,7 @@ aten = torch.ops.aten
 __all__ = ["rewrite_nhwc_node", "has_nhwc_rewriter"]
 
 
-class NHWCNodeRewritersRegistry(OpFuncRegistry):
+class NHWCNodeRewritersRegistry(op_func_registry.OpFuncRegistry):
 
   def __missing__(self, op):
     def _rewriter(node):
@@ -38,14 +39,14 @@ class NHWCNodeRewritersRegistry(OpFuncRegistry):
 rewriters = NHWCNodeRewritersRegistry()
 
 
-def rewrite_nhwc_node(node: Node):
+def rewrite_nhwc_node(node: torch.fx.Node):
   if not layout_mark.is_nhwc_node(node):
     return
 
   rewriters[node.target](node)
 
 
-def has_nhwc_rewriter(node: Node):
+def has_nhwc_rewriter(node: torch.fx.Node):
   return node.target in rewriters
 
 
@@ -54,13 +55,13 @@ def has_nhwc_rewriter(node: Node):
 
 @rewriters.register(torch.ops.quantized_decomposed.dequantize_per_tensor)
 @rewriters.register(torch.ops.quantized_decomposed.quantize_per_tensor)
-def noop(node: Node):
+def noop(node: torch.fx.Node):
   pass
 
 
 @rewriters.register(torch.ops.quantized_decomposed.dequantize_per_channel)
 @rewriters.register(torch.ops.quantized_decomposed.quantize_per_channel)
-def _qdq_per_channel_rewriter(node: Node):
+def _qdq_per_channel_rewriter(node: torch.fx.Node):
   new_args = []
   new_kwargs = {}
 
@@ -199,7 +200,7 @@ def _qdq_per_channel_rewriter(node: Node):
 @rewriters.register(aten._prelu_kernel)
 @rewriters.register(aten.softplus)
 @rewriters.register(aten.silu)
-def noop(node: Node):
+def noop(node: torch.fx.Node):
   pass
 
 
@@ -212,14 +213,16 @@ def noop(node: Node):
 @rewriters.register(aten.max_pool2d_with_indices)
 @rewriters.register(aten.avg_pool2d)
 @rewriters.register(aten._adaptive_avg_pool2d.default)
-def transpose_first_arg_rewriter(node: Node):
+def transpose_first_arg_rewriter(node: torch.fx.Node):
   op = node.target
 
   def nhwc_op(x, *args, **kwargs):
     nonlocal op
     x = utils.tensor_to_nchw(x)
     res = pytree.tree_map_only(
-        torch.Tensor, utils.tensor_to_nhwc, op(x, *args, **kwargs)
+        torch.Tensor,
+        utils.tensor_to_nhwc,
+        op(x, *args, **kwargs),
     )
     return res
 
@@ -227,7 +230,7 @@ def transpose_first_arg_rewriter(node: Node):
 
 
 @rewriters.register(aten.convolution)
-def _aten_convolution_rewriter(node: Node):
+def _aten_convolution_rewriter(node: torch.fx.Node):
   op = node.target
 
   def conv_nhwc(input, weight, bias, *args, **kwargs):
@@ -286,7 +289,7 @@ def _aten_convolution_rewriter(node: Node):
 @rewriters.register(aten.sort.default)
 @rewriters.register(aten.topk.default)
 @rewriters.register(aten.cat.default)
-def dim_attr_rewriter(node: Node):
+def dim_attr_rewriter(node: torch.fx.Node):
   op = node.target
 
   new_args = []

ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/op_func_registry.py

@@ -12,8 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import utils  # NOQA
-import torch
+"""Op function registry for the optimized layout transposes pass."""
+
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import utils
 
 
 class OpFuncRegistry(dict):

ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/pass_body.py

@@ -12,23 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Optimize layout transposes pass."""
+
 import operator
 import os
-from typing import Optional, Tuple, Union
-
-from ai_edge_torch.convert.fx_passes import ExportedProgramPassBase
-from ai_edge_torch.convert.fx_passes import ExportedProgramPassResult
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_check  # NOQA
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_mark  # NOQA
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_partitioners  # NOQA
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import layout_rewrite  # NOQA
-from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import utils  # NOQA
+from typing import Union
+
+from ai_edge_torch._convert.fx_passes import ExportedProgramPassBase
+from ai_edge_torch._convert.fx_passes import ExportedProgramPassResult
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_check  # NOQA
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_mark  # NOQA
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_partitioners  # NOQA
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_rewrite  # NOQA
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import utils  # NOQA
 import torch
 import torch.ao.quantization.quantize_pt2e
-from torch.export import ExportedProgram
-from torch.fx import GraphModule
-from torch.fx import Node
-import torch.utils._pytree as pytree
 
 TransposeFunc = Union[utils.tensor_to_nchw, utils.tensor_to_nhwc]
 
@@ -51,8 +49,8 @@ class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
       transpose_func: TransposeFunc,
       transpose_node_meta: dict,
   ) -> list[torch.fx.Node]:
-    """
-    original:
+    """original:
+
         input_dq -> target
     insert the node as:
         input_dq -> (T q dq) -> target
@@ -86,8 +84,8 @@ class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
       transpose_func: TransposeFunc,
       transpose_node_meta: dict,
   ) -> list[torch.fx.Node]:
-    """
-    original:
+    """original:
+
         input_q -> target
     insert the node as:
         input_q -> (dq T q) -> target

ai_edge_torch/{convert → _convert}/fx_passes/optimize_layout_transposes_pass/utils.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""Utils for the optimized layout transposes pass."""
+
 from typing import Callable
 
 import torch

ai_edge_torch/_convert/signature.py

@@ -0,0 +1,100 @@
+# 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 dataclasses
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.utils._pytree as pytree
+
+
+@dataclasses.dataclass
+class Signature:
+  name: str
+  module: torch.nn.Module
+  sample_args: tuple[torch.Tensor]
+  sample_kwargs: dict[str, torch.Tensor]
+  dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None
+
+  @property
+  def _normalized_sample_args_kwargs(self):
+    args, kwargs = self.sample_args, self.sample_kwargs
+    if args is not None:
+      if not isinstance(args, tuple):
+        # TODO(b/352584188): Check value types
+        raise ValueError("sample_args must be a tuple of torch tensors.")
+    if kwargs is not None:
+      if not isinstance(kwargs, dict) or not all(
+          isinstance(key, str) for key in kwargs.keys()
+      ):
+        # TODO(b/352584188): Check value types
+        raise ValueError("sample_kwargs must be a dict of string to tensor.")
+    args = args if args is not None else tuple()
+    kwargs = kwargs if kwargs is not None else {}
+    return args, kwargs
+
+  @property
+  def flat_arg_names(self) -> list[str]:
+    spec = pytree.tree_flatten(self._normalized_sample_args_kwargs)[1]
+    args_spec, kwargs_spec = spec.children_specs
+    names = []
+    for i in range(args_spec.num_leaves):
+      names.append(f"args_{i}")
+
+    kwargs_names = self._flat_kwarg_names(
+        kwargs_spec.children_specs, kwargs_spec.context
+    )
+    names.extend(kwargs_names)
+    return names
+
+  def _flat_kwarg_names(self, specs, context) -> List[str]:
+    flat_names = []
+    if context is None:
+      for i, spec in enumerate(specs):
+        if spec.children_specs:
+          flat_names.extend([
+              f"{i}_{name}"
+              for name in self._flat_kwarg_names(
+                  spec.children_specs, spec.context
+              )
+          ])
+        else:
+          flat_names.append(f"{i}")
+    else:
+      flat_ctx = self._flatten_list(context)
+      for prefix, spec in zip(flat_ctx, specs):
+        leaf_flat_names = self._flat_kwarg_names(
+            spec.children_specs, spec.context
+        )
+        if leaf_flat_names:
+          flat_names.extend([f"{prefix}_{name}" for name in leaf_flat_names])
+        else:
+          flat_names.append(prefix)
+
+    return flat_names
+
+  def _flatten_list(self, l: List) -> List:
+    flattened = []
+    for item in l:
+      if isinstance(item, list):
+        flattened.extend(self._flatten_list(item))
+      else:
+        flattened.append(item)
+    return flattened
+
+  @property
+  def flat_args(self) -> tuple[Any]:
+    args, kwargs = self._normalized_sample_args_kwargs
+    return tuple([*args, *kwargs.values()])