ai-edge-torch-nightly 0.3.0.dev20250114__py3-none-any.whl

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_check.py

@@ -0,0 +1,258 @@
+# 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.
+# ==============================================================================
+"""Layout check for the optimized layout transposes pass."""
+
+import dataclasses
+import operator
+
+import ai_edge_torch
+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
+
+aten = torch.ops.aten
+
+__all__ = [
+    "is_4d",
+    "can_be_nhwc",
+    "must_be_nhwc",
+    "get_layout_sensitive_inputs",
+    "get_no_rewriter_nhwc_ops",
+]
+
+
+class LayoutSensitiveInputsGettersRegistry(OpFuncRegistry):
+
+  def __missing__(self, op):
+
+    def _default_getter(node: Node):
+      """Default layout sensitive inputs are all input nodes."""
+      return node.all_input_nodes
+
+    return _default_getter
+
+
+@dataclasses.dataclass
+class NHWCable:
+  can_be: bool
+  must_be: bool
+
+  def __bool__(self):
+    raise RuntimeError(
+        "Boolean value on NHWCable is disabled. Please call .can_be or .must_be"
+    )
+
+
+class NHWCableNodeCheckersRegistry(OpFuncRegistry):
+
+  def __init__(self):
+    self.no_rewriter_nhwc_ops = set()
+
+  def __missing__(self, op):
+
+    def _default_checker(node: Node):
+      """Default checker for most of the layout insensitive ops.
+
+      The node should be marked and rewritten to NHWC if:
+      1. The node output is a single 4-D tensor.
+      2. All layout sensitive input nodes (default all inputs) of this
+        node are all marked as NHWC.
+      3. All layout sensitive input nodes return 4-D tensors.
+      4. There exists a rewrite rule for this node (explicit registry
+        required for noop.)
+      """
+      nonlocal self
+      layout_sensitive_inputs = get_layout_sensitive_inputs(node)
+
+      can_be_nhwc = is_4d(node) and all_layout_sensitive_inputs_are_4d(node)
+      has_rewriter = layout_rewrite.has_nhwc_rewriter(node)
+
+      if can_be_nhwc and not has_rewriter:
+        self.no_rewriter_nhwc_ops.add(node.target)
+
+      return NHWCable(can_be_nhwc and has_rewriter, must_be=False)
+
+    return _default_checker
+
+
+nhwcable_node_checkers = NHWCableNodeCheckersRegistry()
+layout_sensitive_inputs_getters = LayoutSensitiveInputsGettersRegistry()
+
+
+def can_be_nhwc(node: Node):
+  return nhwcable_node_checkers[node.target](node).can_be
+
+
+def must_be_nhwc(node: Node):
+  return nhwcable_node_checkers[node.target](node).must_be
+
+
+def get_layout_sensitive_inputs(node: Node):
+  return layout_sensitive_inputs_getters[node.target](node)
+
+
+def get_no_rewriter_nhwc_ops():
+  """Debug only: get the ops that may be NHWC but not due to no rewriter registered."""
+  return nhwcable_node_checkers.no_rewriter_nhwc_ops
+
+
+def is_4d(node: Node):
+  val = node.meta.get("val")
+  if val is None:
+    return False
+
+  if isinstance(val, (list, tuple)) and val:
+    val = val[0]
+
+  if not hasattr(val, "shape"):
+    return False
+
+  return len(val.shape) == 4
+
+
+def all_layout_sensitive_inputs_are_4d(node: Node):
+  return all(is_4d(m) for m in get_layout_sensitive_inputs(node))
+
+
+# ==== Quantize ops (use default NHWC checker)
+
+
+@layout_sensitive_inputs_getters.register(
+    torch.ops.quantized_decomposed.dequantize_per_tensor
+)
+@layout_sensitive_inputs_getters.register(
+    torch.ops.quantized_decomposed.quantize_per_tensor
+)
+@layout_sensitive_inputs_getters.register(
+    torch.ops.quantized_decomposed.dequantize_per_channel
+)
+@layout_sensitive_inputs_getters.register(
+    torch.ops.quantized_decomposed.quantize_per_channel
+)
+def _qdq_layout_sensitive_inputs_getter(node: Node):
+  return [node.args[0]]
+
+
+# ==== Ops must be NHWC if possible
+
+
+@layout_sensitive_inputs_getters.register(aten.conv2d)
+@layout_sensitive_inputs_getters.register(aten.convolution)
+@layout_sensitive_inputs_getters.register(
+    aten._native_batch_norm_legit_no_training
+)
+@layout_sensitive_inputs_getters.register(aten.group_norm)
+@layout_sensitive_inputs_getters.register(aten.native_group_norm)
+def _first_arg_getter(node):
+  return [node.args[0]]
+
+
+# Note: default layout sensitive inputs are all inputs when not specified.
+@nhwcable_node_checkers.register(aten.max_pool2d)
+@nhwcable_node_checkers.register(aten.max_pool2d_with_indices)
+@nhwcable_node_checkers.register(aten.amax)
+@nhwcable_node_checkers.register(aten.avg_pool2d)
+@nhwcable_node_checkers.register(aten._prelu_kernel)
+@nhwcable_node_checkers.register(aten.upsample_bilinear2d)
+@nhwcable_node_checkers.register(aten.upsample_nearest2d)
+@nhwcable_node_checkers.register(aten._adaptive_avg_pool2d)
+@nhwcable_node_checkers.register(aten.conv2d)
+@nhwcable_node_checkers.register(aten.convolution)
+def _all_layout_sensitive_inputs_are_4d_checker(node: Node):
+  can_be = all_layout_sensitive_inputs_are_4d(node)
+  return NHWCable(can_be, must_be=can_be)
+
+
+@nhwcable_node_checkers.register(aten._native_batch_norm_legit_no_training)
+def _aten_norm_checker(node):
+  val = node.meta.get("val")
+  if (
+      not isinstance(val, (list, tuple))
+      or not val
+      or not hasattr(val[0], "shape")
+  ):
+    return NHWCable(can_be=False, must_be=False)
+  return NHWCable(can_be=len(val[0].shape) == 4, must_be=False)
+
+
+@nhwcable_node_checkers.register(aten.group_norm)
+def _aten_group_norm_checker(node):
+  val = node.meta.get("val")
+  if not hasattr(val, "shape"):
+    return NHWCable(can_be=False, must_be=False)
+
+  can_be = len(val.shape) == 4
+  must_be = can_be and ai_edge_torch.config.enable_group_norm_composite
+  return NHWCable(can_be=can_be, must_be=must_be)
+
+
+@nhwcable_node_checkers.register(aten.native_group_norm)
+def _aten_native_group_norm_checker(node):
+  val = node.meta.get("val")
+  if (
+      not isinstance(val, (list, tuple))
+      or not val
+      or not hasattr(val[0], "shape")
+  ):
+    return NHWCable(can_be=False, must_be=False)
+  if len(node.args) >= 3 and (
+      node.args[1] is not None or node.args[2] is not None
+  ):
+    # Disable NHWC rewriter due to precision issue with weight and bias.
+    # TODO(b/354780253): Re-enable NHWC rewriter with proper lowering.
+    return NHWCable(can_be=False, must_be=False)
+  return NHWCable(can_be=len(val[0].shape) == 4, must_be=False)
+
+
+# ==== Ops must be NCHW
+
+
+@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")
+@nhwcable_node_checkers.register(aten.view)
+@nhwcable_node_checkers.register(aten.unsqueeze_copy)
+@nhwcable_node_checkers.register(aten.expand)
+@nhwcable_node_checkers.register(aten.permute)
+@nhwcable_node_checkers.register(aten.as_strided)
+def _not_nhwc(node: Node):
+  return NHWCable(can_be=False, must_be=False)
+
+
+# ==== Others
+
+
+@layout_sensitive_inputs_getters.register(aten.index)
+@layout_sensitive_inputs_getters.register(aten._unsafe_index)
+def _aten_index_layout_sensitive_inputs_getter(node):
+  return [node.args[0]]
+
+
+@nhwcable_node_checkers.register(aten.index)
+@nhwcable_node_checkers.register(aten._unsafe_index)
+def _aten_index_checker(node):
+  layout_sensitive_inputs = get_layout_sensitive_inputs(node)
+  can_be = is_4d(node) and all_layout_sensitive_inputs_are_4d(node)
+  return NHWCable(can_be, must_be=False)
+
+
+@nhwcable_node_checkers.register(operator.getitem)
+def _getitem_checker(node):
+  src = node.args[0]
+  return nhwcable_node_checkers[src.target](src)

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_mark.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.
+# ==============================================================================
+"""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
+# partition.
+IS_NHWC_NODE = "OPTIMIZE_LAYOUT_TRANSPOSES_PASS__IS_NHWC_NODE"
+
+
+# Tag which is added to a node's meta to indicate that it is derived completely
+# from constant and/or weight tensor(s).
+IS_CONST_NODE = "OPTIMIZE_LAYOUT_TRANSPOSES_PASS__IS_CONST_NODE"
+
+
+def mark_as_nhwc_node(node: torch.fx.Node) -> None:
+  node.meta[IS_NHWC_NODE] = True
+
+
+def mark_as_nchw_node(node: torch.fx.Node) -> None:
+  node.meta[IS_NHWC_NODE] = False
+
+
+def is_nhwc_node(node: torch.fx.Node) -> bool:
+  return node.meta.get(IS_NHWC_NODE, False)
+
+
+def is_nchw_node(node: torch.fx.Node) -> bool:
+  return not is_nhwc_node(node)
+
+
+def mark_as_const_node(node: torch.fx.Node) -> None:
+  node.meta[IS_CONST_NODE] = True
+
+
+def is_const_node(node: torch.fx.Node) -> bool:
+  return node.meta.get(IS_CONST_NODE, False)

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+# ==============================================================================
+"""Layout partitioners."""
+
+from . import greedy
+from . import min_cut

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/greedy.py

@@ -0,0 +1,68 @@
+# 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.
+# ==============================================================================
+"""Greedy partitioning algorithm."""
+
+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.
+
+  Partition the graph module into NHWC and non-NHWC subgraphs and mark nodes in
+  the NHWC partitions.
+
+  Implements O(|V|) greedy partitioning algorithm.
+
+  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 not node.all_input_nodes:
+      # This node has no inputs so we don't need to change anything
+      continue
+
+    if layout_check.must_be_nhwc(node):
+      # If the node has must_be_nhwc equals true, mark this node as NHWC
+
+      layout_mark.mark_as_nhwc_node(node)
+    elif layout_check.can_be_nhwc(node):
+      # If the following conditions are all true, mark this node as NHWC
+      # - The node has can_be_nhwc equals true
+      # - Any of the node's layout sensitive inputs is marked as NHWC
+      # - All the node's layout sensitive inputs are 4D tensors
+
+      layout_sensitive_inputs = layout_check.get_layout_sensitive_inputs(node)
+
+      should_be_nhwc = any(
+          map(layout_mark.is_nhwc_node, layout_sensitive_inputs)
+      )
+      for input_node in layout_sensitive_inputs:
+        if not layout_mark.is_nhwc_node(input_node) and not layout_check.is_4d(
+            input_node
+        ):
+          should_be_nhwc = False
+
+      if should_be_nhwc:
+        layout_mark.mark_as_nhwc_node(node)
+
+  graph_module.recompile()
+  return graph_module

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitioners/min_cut.py

@@ -0,0 +1,216 @@
+# 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.
+# ==============================================================================
+"""Min cut solver for partitioning the graph module into NHWC and non-NHWC subgraphs."""
+
+import collections
+import dataclasses
+
+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
+
+
+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.
+  """
+  graph = graph_module.graph
+  n_nodes = len(graph.nodes)
+  n_edges = sum(len(n.users) for n in graph.nodes)
+
+  # According to the experiments our model set, |V| < 2000 can
+  # be partitioned generally in a reasonable time.
+  return n_nodes**2 * n_edges < 2000**3
+
+
+class MinCutSolver:
+  # A number that is large enough but can fit into int32 with all computations
+  # in the maximum flow.
+  INF_COST = 1 << 28
+
+  def __init__(self):
+    self._edges_map = collections.defaultdict(dict)
+    self._obj_to_node = {}
+    self._node_to_obj = {}
+    self._nodes_cnt = 0
+
+    self.source = self._next_nid()
+    self.sink = self._next_nid()
+
+  def _next_nid(self):
+    nid = self._nodes_cnt
+    self._nodes_cnt += 1
+    return nid
+
+  @property
+  def nodes(self):
+    return list(range(self._nodes_cnt))
+
+  @property
+  def edges_map(self):
+    return self._edges_map
+
+  @property
+  def edges(self):
+    return [
+        [n, m, cost]
+        for n, next_nodes in self._edges_map.items()
+        for m, cost in next_nodes.items()
+    ]
+
+  @property
+  def graph(self):
+    edges = np.array(self.edges)
+    return scipy.sparse.csr_matrix(
+        (
+            np.minimum(edges[:, 2], MinCutSolver.INF_COST),
+            (edges[:, 0], edges[:, 1]),
+        ),
+        shape=(self._nodes_cnt, self._nodes_cnt),
+        dtype=np.int32,
+    )
+
+  def get_nid(self, obj=None):
+    if obj is None:
+      return self._next_nid()
+
+    nid = self._obj_to_node.get(obj)
+    if nid is None:
+      nid = self._next_nid()
+
+      self._obj_to_node[obj] = nid
+      self._node_to_obj[nid] = obj
+    return nid
+
+  def get_obj(self, nid: int):
+    return self._node_to_obj.get(nid, None)
+
+  def add_edge(self, a_id: int, b_id: int, cost: int):
+    assert isinstance(cost, int)
+    self._edges_map[a_id][b_id] = cost
+
+  def solve(self):
+    flow = scipy.sparse.csgraph.maximum_flow(
+        self.graph, self.source, self.sink, method="dinic"
+    ).flow
+
+    # Max-flow min-cut theorem: find min-cuts in the residual network.
+    ds = scipy.cluster.hierarchy.DisjointSet(self.nodes)
+    for n, m, cost in self.edges:
+      if abs(flow[n, m]) < cost:
+        ds.merge(n, m)
+
+    residual_reachable_nodes = ds.subset(self.source)
+
+    cuts = set()
+    for n, m, cost in self.edges:
+      if n in residual_reachable_nodes and m not in residual_reachable_nodes:
+        cuts.add((n, m))
+
+    return cuts
+
+
+@dataclasses.dataclass(frozen=True)
+class MultiUsersDummyNode:
+  src: torch.fx.Node
+
+
+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.
+  """
+  graph = graph_module.graph
+
+  mc_solver = MinCutSolver()
+  for fx_node in graph.nodes:
+    if layout_mark.is_const_node(fx_node):
+      continue
+
+    nid = mc_solver.get_nid(fx_node)
+    if fx_node.op in ("placeholder", "output"):
+      # All inputs and outputs are not NHWCable nodes in the graph,
+      # connected to source S directly with inf cost to cut
+      mc_solver.add_edge(mc_solver.source, nid, cost=MinCutSolver.INF_COST)
+    elif not layout_check.can_be_nhwc(fx_node):
+      # All not NHWCable nodes are connected to source S directly,
+      # with inf cost to cut.
+      mc_solver.add_edge(mc_solver.source, nid, cost=MinCutSolver.INF_COST)
+    elif layout_check.must_be_nhwc(fx_node):
+      # All must be NHWC nodes are connected to sink T directly,
+      # with inf cost to cut
+      mc_solver.add_edge(nid, mc_solver.sink, cost=MinCutSolver.INF_COST)
+
+    cut_cost = 10  # set 10 to be a unit of cut cost
+    if fx_node.target in (torch.ops.aten.mean.default, torch.ops.aten.mean.dim):
+      # TFLite converter cannot fuse the lowering of (tpos-mean) but (mean-tpos)
+      # when it applies on the feature dimensions. Therefore decreasing the cut
+      # cost for aten.mean's out-going edges to favor having a cut (transpose)
+      # after the node than before when the number of transposes are equal.
+      # TODO: Remove this rule when converter has fuse rule for tpos-mean.
+      cut_cost = 9
+
+    if len(fx_node.users) > 1:
+      # If a node's (A1) output is used by multiple nodes (B1, B2, B3, ...),
+      # the cost to split A1 and Bs into different partitions would just be 1
+      # transpose. So we need to introduce a dummy node between A1 and Bs in the
+      # min-cut graph to reflect the fact that disconnecting them doesn't
+      # introduce multiple transposes.
+      dummy_nid = mc_solver.get_nid(MultiUsersDummyNode(fx_node))
+      mc_solver.add_edge(nid, dummy_nid, cost=cut_cost)
+      mc_solver.add_edge(dummy_nid, nid, cost=cut_cost)
+      nid = dummy_nid
+
+    for user in fx_node.users:
+      # All the other nodes and edges in the model graph are scattered
+      # and connected as is in the new graph, with 1 cost to cut an edge.
+      user_id = mc_solver.get_nid(user)
+      mc_solver.add_edge(nid, user_id, cost=cut_cost)
+      mc_solver.add_edge(user_id, nid, cost=cut_cost)
+
+  cuts = mc_solver.solve()
+
+  # Find nodes that is connected to sink after the min-cut and mark as NHWC.
+  ds = scipy.cluster.hierarchy.DisjointSet(mc_solver.nodes)
+  for n, m, cost in mc_solver.edges:
+    if (n, m) in cuts or (m, n) in cuts:
+      continue
+    ds.merge(n, m)
+  assert not ds.connected(mc_solver.source, mc_solver.sink)
+
+  for nid in mc_solver.nodes:
+    if ds.connected(nid, mc_solver.source):
+      continue
+
+    obj = mc_solver.get_obj(nid)
+    if obj is None:
+      continue
+    if isinstance(obj, MultiUsersDummyNode):
+      continue
+
+    assert isinstance(obj, torch.fx.Node)
+    layout_mark.mark_as_nhwc_node(obj)
+
+  graph_module.recompile()
+  return graph_module