ai-edge-torch-nightly 0.2.0.dev20240714__py3-none-any.whl → 0.3.0.dev20240926__py3-none-any.whl

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

@@ -12,22 +12,22 @@
 # 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
+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
 
-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
-
 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):
@@ -39,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
 
 
@@ -55,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 = {}
 
@@ -200,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
 
 
@@ -213,25 +213,28 @@ 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
 
   node.target = nhwc_op
 
 
+@rewriters.register(aten.conv2d)
 @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):
+  def conv_nhwc(input, weight, bias=None, *args, **kwargs):
     nonlocal op
     nhwc_bias = None
     if bias is not None and len(bias.shape) == 1:
@@ -287,7 +290,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 = []
@@ -329,7 +332,12 @@ def _aten__native_batch_norm_legit_no_training(node):
   def batch_norm(input, weight, bias, running_mean, running_var, momentum, eps):
     a = input - running_mean
     b = torch.sqrt(running_var + eps)
-    return a / b * weight + bias, None, None
+    out = a / b
+    if weight is not None:
+      out = out * weight
+    if bias is not None:
+      out = out + bias
+    return out, None, None
 
   node.target = batch_norm
 
@@ -349,7 +357,12 @@ def _aten_native_group_norm(node):
   ):
     input_reshaped = torch.reshape(
         input,
-        [batch_size, flattened_inner_size, num_groups, num_channels // num_groups],
+        [
+            batch_size,
+            flattened_inner_size,
+            num_groups,
+            num_channels // num_groups,
+        ],
     )
     reduction_dims = [1, 3]
 

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

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-import torch
+"""Op function registry for the optimized layout transposes pass."""
 
-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 import utils
 
 
 class OpFuncRegistry(dict):

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

@@ -12,28 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-import os
-from typing import Optional, Tuple, Union
+"""Optimize layout transposes pass."""
 
+import operator
+import os
+from typing import Union
+
+from ai_edge_torch import fx_pass_base
+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
-
-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
 
 TransposeFunc = Union[utils.tensor_to_nchw, utils.tensor_to_nhwc]
 
 
-class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
+class OptimizeLayoutTransposesPass(fx_pass_base.ExportedProgramPassBase):
 
   def get_source_meta(self, node: torch.fx.Node):
     keys = ["stack_trace", "nn_module_stack", "source_fn_stack", "from_node"]
@@ -51,8 +48,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 +83,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
@@ -96,7 +93,7 @@ class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
 
     q_args = input_q.args[1:]
     q_kwargs = input_q.kwargs
-    q_op, dq_op = self.get_paired_q_dq_ops(input_q.target)
+    q_op, dq_op = utils.get_paired_q_dq_ops(input_q.target)
     with graph.inserting_before(target):
       # Q and DQ inserted here may required updating the `axis` arg when they
       # are per_channel ops. However, instead of updating here, the nodes would
@@ -207,7 +204,8 @@ class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
 
     if not layout_check.is_4d(input_node):
       raise AssertionError(
-          f"Attempting to convert non-NHWC compatible node to NHWC: {input_node}"
+          "Attempting to convert non-NHWC compatible node to NHWC:"
+          f" {input_node}"
       )
 
     # Assign target node's source meta to the to_NHWC node, because the transpose
@@ -249,7 +247,9 @@ class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
 
     for node in graph.nodes:
       has_input_nodes = len(node.all_input_nodes) > 0
-      all_inputs_are_const = all(map(layout_mark.is_const_node, node.all_input_nodes))
+      all_inputs_are_const = all(
+          map(layout_mark.is_const_node, node.all_input_nodes)
+      )
       if (
           node.name in non_user_input_names
           or (has_input_nodes and all_inputs_are_const)
@@ -261,7 +261,9 @@ class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
     self.mark_const_nodes(exported_program)
 
     graph_module = exported_program.graph_module
-    partitioner = os.environ.get("AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER", None)
+    partitioner = os.environ.get(
+        "AIEDGETORCH_LAYOUT_OPTIMIZE_PARTITIONER", None
+    )
     if partitioner == "MINCUT":
       graph_module = layout_partitioners.min_cut.partition(graph_module)
     elif partitioner == "GREEDY":
@@ -274,6 +276,14 @@ class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
         graph_module = layout_partitioners.greedy.partition(graph_module)
 
     graph = graph_module.graph
+    for node in list(graph.nodes):
+      if node.target == operator.getitem:
+        # force the layout mark of a getitem node to follow its producer.
+        if layout_mark.is_nchw_node(node.args[0]):
+          layout_mark.mark_as_nchw_node(node)
+        else:
+          layout_mark.mark_as_nhwc_node(node)
+
     for node in list(graph.nodes):
       if layout_mark.is_nhwc_node(node):
         for input_node in layout_check.get_layout_sensitive_inputs(node):
@@ -290,4 +300,4 @@ class OptimizeLayoutTransposesPass(ExportedProgramPassBase):
     # Mark const node again for debugging
     self.mark_const_nodes(exported_program)
 
-    return ExportedProgramPassResult(exported_program, True)
+    return fx_pass_base.ExportedProgramPassResult(exported_program, True)

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,66 @@
+# 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
+
+from ai_edge_torch import lowertools
+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 = lowertools.flat_dict_names(
+        kwargs_spec.children_specs, kwargs_spec.context
+    )
+    names.extend(kwargs_names)
+    return names
+
+  @property
+  def flat_args(self) -> tuple[Any]:
+    args, kwargs = self._normalized_sample_args_kwargs
+    return tuple([*args, *kwargs.values()])