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

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py

@@ -0,0 +1,400 @@
+# 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 operator
+
+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):
+
+  def __missing__(self, op):
+    def _rewriter(node):
+      raise RuntimeError(f"NHWC node rewriter not found: {str(node)}")
+
+    return _rewriter
+
+
+rewriters = NHWCNodeRewritersRegistry()
+
+
+def rewrite_nhwc_node(node: Node):
+  if not layout_mark.is_nhwc_node(node):
+    return
+
+  rewriters[node.target](node)
+
+
+def has_nhwc_rewriter(node: Node):
+  return node.target in rewriters
+
+
+# ======= Quantize ops
+
+
+@rewriters.register(torch.ops.quantized_decomposed.dequantize_per_tensor)
+@rewriters.register(torch.ops.quantized_decomposed.quantize_per_tensor)
+def noop(node: 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):
+  new_args = []
+  new_kwargs = {}
+
+  def axis_nchw_to_nhwc(axis: int):
+    axis = axis if axis >= 0 else 4 + axis
+    return {3: 2, 2: 1, 1: 3}.get(axis, axis)
+
+  for arg, spec in zip(node.args, op._schema.arguments):
+    if spec.name == "axis":
+      new_args.append(axis_nchw_to_nhwc(arg))
+    else:
+      new_args.append(arg)
+
+  for spec in op._schema.arguments[len(node.args) :]:
+    if spec.name not in node.kwargs:
+      continue
+
+    if spec.name == "axis":
+      new_kwargs[spec.name] = axis_nchw_to_nhwc(node.kwargs[spec.name])
+    else:
+      new_kwargs[spec.name] = node.kwargs[spec.name]
+
+  node.args = tuple(new_args)
+  node.kwargs = new_kwargs
+
+
+# ======= Noop ops (layout insensitive ops)
+
+
+@rewriters.register(utils.tensor_to_nhwc)
+@rewriters.register(utils.tensor_to_nchw)
+@rewriters.register(operator.getitem)
+@rewriters.register("output")
+@rewriters.register(aten.add.Tensor)
+@rewriters.register(aten.add.Scalar)
+@rewriters.register(aten.atan2.default)
+@rewriters.register(aten.atan2.out)
+@rewriters.register(aten.bitwise_and.Tensor)
+@rewriters.register(aten.bitwise_and.Scalar)
+@rewriters.register(aten.bitwise_or.Tensor)
+@rewriters.register(aten.bitwise_or.Scalar)
+@rewriters.register(aten.bitwise_xor.Tensor)
+@rewriters.register(aten.bitwise_xor.Scalar)
+@rewriters.register(aten.div.Tensor)
+@rewriters.register(aten.div.Scalar)
+@rewriters.register(aten.div.Tensor_mode)
+@rewriters.register(aten.div.Scalar_mode)
+@rewriters.register(aten.fmod.Tensor)
+@rewriters.register(aten.fmod.Scalar)
+@rewriters.register(aten.mul.Tensor)
+@rewriters.register(aten.mul.Scalar)
+@rewriters.register(aten.remainder.Tensor)
+@rewriters.register(aten.remainder.Scalar)
+@rewriters.register(aten.sub.Tensor)
+@rewriters.register(aten.sub.Scalar)
+@rewriters.register(aten.eq.Tensor)
+@rewriters.register(aten.eq.Scalar)
+@rewriters.register(aten.ne.Tensor)
+@rewriters.register(aten.ne.Scalar)
+@rewriters.register(aten.le.Tensor)
+@rewriters.register(aten.le.Scalar)
+@rewriters.register(aten.ge.Tensor)
+@rewriters.register(aten.ge.Scalar)
+@rewriters.register(aten.gt.Tensor)
+@rewriters.register(aten.gt.Scalar)
+@rewriters.register(aten.lt.Tensor)
+@rewriters.register(aten.lt.Scalar)
+@rewriters.register(aten.maximum.default)
+@rewriters.register(aten.minimum.default)
+@rewriters.register(aten.mean.default)
+@rewriters.register(aten.prod.default)
+@rewriters.register(aten.abs.default)
+@rewriters.register(aten.acos.default)
+@rewriters.register(aten.acosh.default)
+@rewriters.register(aten.asin.default)
+@rewriters.register(aten.asinh.default)
+@rewriters.register(aten.atan.default)
+@rewriters.register(aten.atanh.default)
+@rewriters.register(aten.bitwise_not.default)
+@rewriters.register(aten.ceil.default)
+@rewriters.register(aten.clamp.default)
+@rewriters.register(aten.clamp.Tensor)
+@rewriters.register(aten.cos.default)
+@rewriters.register(aten.cosh.default)
+@rewriters.register(aten.erf.default)
+@rewriters.register(aten.exp.default)
+@rewriters.register(aten.expm1.default)
+@rewriters.register(aten.floor.default)
+@rewriters.register(aten.log.default)
+@rewriters.register(aten.log10.default)
+@rewriters.register(aten.log1p.default)
+@rewriters.register(aten.log2.default)
+@rewriters.register(aten.isnan.default)
+@rewriters.register(aten.neg.default)
+@rewriters.register(aten.pow.Tensor_Tensor)
+@rewriters.register(aten.pow.Tensor_Scalar)
+@rewriters.register(aten.pow.Scalar)
+@rewriters.register(aten.reciprocal.default)
+@rewriters.register(aten.round.default)
+@rewriters.register(aten.rsqrt.default)
+@rewriters.register(aten.sigmoid.default)
+@rewriters.register(aten.sign.default)
+@rewriters.register(aten.sin.default)
+@rewriters.register(aten.sinh.default)
+@rewriters.register(aten.sqrt.default)
+@rewriters.register(aten.tan.default)
+@rewriters.register(aten.tanh.default)
+@rewriters.register(aten.trunc.default)
+@rewriters.register(aten.nonzero.default)
+@rewriters.register(aten.copy.default)
+@rewriters.register(aten.mm.default)
+@rewriters.register(aten.fill.Scalar)
+@rewriters.register(aten.col2im.default)
+@rewriters.register(aten.addmm.default)
+@rewriters.register(aten.gelu.default)
+@rewriters.register(aten.hardtanh.default)
+@rewriters.register(aten.leaky_relu.default)
+@rewriters.register(aten.relu.default)
+@rewriters.register(aten.arange.start_step)
+@rewriters.register(aten.isinf.default)
+@rewriters.register(aten.logical_and.default)
+@rewriters.register(aten.logical_not.default)
+@rewriters.register(aten.logical_or.default)
+@rewriters.register(aten.logical_xor.default)
+@rewriters.register(aten.where.self)
+@rewriters.register(aten.clone.default)
+@rewriters.register(aten.any.default)
+@rewriters.register(aten.repeat.default)
+@rewriters.register(aten.alias.default)
+@rewriters.register(aten._pdist_forward.default)
+@rewriters.register(aten._cdist_forward.default)
+@rewriters.register(aten.bmm.default)
+@rewriters.register(aten.hardswish)
+@rewriters.register(aten.hardsigmoid)
+@rewriters.register(aten._to_copy)
+@rewriters.register(aten._prelu_kernel)
+@rewriters.register(aten.softplus)
+@rewriters.register(aten.silu)
+def noop(node: Node):
+  pass
+
+
+# ======= Add transposes before and after NCHW-only ops (T-aten-T)
+
+
+@rewriters.register(aten.upsample_bilinear2d)
+@rewriters.register(aten.upsample_nearest2d)
+@rewriters.register(aten.max_pool2d)
+@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):
+  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)
+    )
+    return res
+
+  node.target = nhwc_op
+
+
+@rewriters.register(aten.convolution)
+def _aten_convolution_rewriter(node: Node):
+  op = node.target
+
+  def conv_nhwc(input, weight, bias, *args, **kwargs):
+    nonlocal op
+    nhwc_bias = None
+    if bias is not None and len(bias.shape) == 1:
+      nhwc_bias = bias
+      bias = None
+
+    input = utils.tensor_to_nchw(input)
+    res = pytree.tree_map_only(
+        torch.Tensor,
+        utils.tensor_to_nhwc,
+        op(input, weight, bias, *args, **kwargs),
+    )
+
+    if nhwc_bias is not None:
+      res += nhwc_bias
+    return res
+
+  node.target = conv_nhwc
+
+
+# ======= Rewrite dim attribute(s)
+
+
+@rewriters.register(aten._softmax.default)
+@rewriters.register(aten.select.int)
+@rewriters.register(aten.slice.Tensor)
+@rewriters.register(aten.sum.dim_IntList)
+@rewriters.register(aten.mean.dim)
+@rewriters.register(aten.prod.dim_int)
+@rewriters.register(aten.var.dim)
+@rewriters.register(aten.var.correction)
+@rewriters.register(aten.slice_scatter.default)
+@rewriters.register(aten.diagonal.default)
+@rewriters.register(aten.select_scatter.default)
+@rewriters.register(aten.sym_size.int)
+@rewriters.register(aten.sym_stride.int)
+@rewriters.register(aten._log_softmax.default)
+@rewriters.register(aten.split_with_sizes.default)
+@rewriters.register(aten.squeeze.dim)
+@rewriters.register(aten.squeeze.dims)
+@rewriters.register(aten.scatter.value)
+@rewriters.register(aten.scatter.src)
+@rewriters.register(aten.scatter_add.default)
+@rewriters.register(aten.scatter_reduce.two)
+@rewriters.register(aten.any.dim)
+@rewriters.register(aten.any.dims)
+@rewriters.register(aten.flip.default)
+@rewriters.register(aten.index_select.default)
+@rewriters.register(aten.cumsum.default)
+@rewriters.register(aten.max.dim)
+@rewriters.register(aten.min.dim)
+@rewriters.register(aten.gather.default)
+@rewriters.register(aten.sort.default)
+@rewriters.register(aten.topk.default)
+@rewriters.register(aten.cat.default)
+def dim_attr_rewriter(node: Node):
+  op = node.target
+
+  new_args = []
+  new_kwargs = {}
+
+  def dims_nchw_to_nhwc(dims: list[int]):
+    def convert(dim: int):
+      dim = dim if dim >= 0 else 4 + dim
+      return {3: 2, 2: 1, 1: 3}.get(dim, dim)
+
+    dims = pytree.tree_map_only(int, convert, dims)
+    dims = pytree.tree_map_only(torch.SymInt, convert, dims)
+    return dims
+
+  for arg, spec in zip(node.args, op._schema.arguments):
+    if spec.name.startswith("dim"):
+      new_args.append(dims_nchw_to_nhwc(arg))
+    else:
+      new_args.append(arg)
+
+  for spec in op._schema.arguments[len(node.args) :]:
+    if spec.name not in node.kwargs:
+      continue
+
+    if spec.name.startswith("dim"):
+      new_kwargs[spec.name] = dims_nchw_to_nhwc(node.kwargs[spec.name])
+    else:
+      new_kwargs[spec.name] = node.kwargs[spec.name]
+
+  node.args = tuple(new_args)
+  node.kwargs = new_kwargs
+
+
+# ======= Others
+
+
+@rewriters.register(aten._native_batch_norm_legit_no_training.default)
+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
+
+  node.target = batch_norm
+
+
+@rewriters.register(aten.native_group_norm.default)
+def _aten_native_group_norm(node):
+
+  def native_group_norm(
+      input,
+      weight,
+      bias,
+      batch_size: int,
+      num_channels: int,
+      flattened_inner_size: int,
+      num_groups: int,
+      eps: float,
+  ):
+    input_reshaped = torch.reshape(
+        input,
+        [batch_size, flattened_inner_size, num_groups, num_channels // num_groups],
+    )
+    reduction_dims = [1, 3]
+
+    biased_var, mean = torch.var_mean(
+        input_reshaped, dim=reduction_dims, unbiased=False, keepdim=True
+    )
+    rstd = torch.rsqrt(biased_var + eps)
+
+    out = (input_reshaped - mean) * rstd
+    out = torch.reshape(out, input.shape)
+
+    if weight is not None:
+      out = out * weight
+    if bias is not None:
+      out = out + bias
+
+    mean = torch.squeeze(mean, reduction_dims)
+    rstd = torch.squeeze(rstd, reduction_dims)
+
+    return out, mean, rstd
+
+  node.target = native_group_norm
+
+
+@rewriters.register(aten.index)
+@rewriters.register(aten._unsafe_index)
+def _aten_index(node):
+  op = node.target
+
+  def index_nhwc(x, indices=[], *args, **kwargs):
+    nonlocal op
+    indices = list(indices)
+    if len(indices) < 4:
+      indices += [None] * (4 - len(indices))
+
+    indices[1:4] = indices[2], indices[3], indices[1]
+    return op(x, indices, *args, **kwargs)
+
+  node.target = index_nhwc
+
+
+@rewriters.register(aten.reflection_pad2d.default)
+def _aten_reflection_pad2d(node):
+  def reflection_pad2d_nhwc(x, padding):
+    padding = [0, 0] + padding
+    return torch.nn.functional.pad(x, padding, mode="reflect")
+
+  node.target = reflection_pad2d_nhwc

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/op_func_registry.py

@@ -0,0 +1,30 @@
+# 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.optimize_layout_transposes_pass import utils  # NOQA
+
+
+class OpFuncRegistry(dict):
+
+  def register(self, op):
+    ops = utils.flatten_torch_op_overloads(op)
+
+    def inner(func):
+      for op in ops:
+        self[op] = func
+      return func
+
+    return inner

ai_edge_torch/convert/fx_passes/optimize_layout_transposes_pass/pass_body.py

@@ -0,0 +1,286 @@
+# 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 os
+from typing import Optional, Tuple, Union
+
+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):
+
+  def get_source_meta(self, node: torch.fx.Node):
+    keys = ["stack_trace", "nn_module_stack", "source_fn_stack", "from_node"]
+    meta = {}
+    for key in keys:
+      if key in node.meta:
+        meta[key] = node.meta[key]
+    return meta
+
+  def insert_t_q_dq(
+      self,
+      graph: torch.fx.Graph,
+      input_dq: torch.fx.Node,
+      target: torch.fx.Node,
+      transpose_func: TransposeFunc,
+      transpose_node_meta: dict,
+  ) -> list[torch.fx.Node]:
+    """
+    original:
+        input_dq -> target
+    insert the node as:
+        input_dq -> (T q dq) -> target
+    """
+    assert utils.is_dq_node(input_dq)
+
+    q_args = input_dq.args[1:]
+    q_kwargs = input_dq.kwargs
+    q_op, dq_op = utils.get_paired_q_dq_ops(input_dq.target)
+    with graph.inserting_before(target):
+      t = graph.call_function(transpose_func, (input_dq,))
+      # 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
+      # be marked as NHWC/NCHW and applied rewriters after partitioning.
+      q = graph.call_function(q_op, (t,) + q_args, q_kwargs)
+      dq = graph.call_function(dq_op, (q,) + q_args, q_kwargs)
+
+    input_dq.meta = transpose_node_meta
+    t.meta = transpose_node_meta
+    q.meta = transpose_node_meta
+    dq.meta = self.get_source_meta(target)
+
+    target.replace_input_with(input_dq, dq)
+    return [t, q, dq]
+
+  def insert_dq_t_q(
+      self,
+      graph: torch.fx.Graph,
+      input_q: torch.fx.Node,
+      target: torch.fx.Node,
+      transpose_func: TransposeFunc,
+      transpose_node_meta: dict,
+  ) -> list[torch.fx.Node]:
+    """
+    original:
+        input_q -> target
+    insert the node as:
+        input_q -> (dq T q) -> target
+    """
+    assert utils.is_q_node(input_q)
+
+    q_args = input_q.args[1:]
+    q_kwargs = input_q.kwargs
+    q_op, dq_op = self.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
+      # be marked as NHWC/NCHW and applied rewriters after partitioning.
+      dq = graph.call_function(dq_op, (input_q,) + q_args, q_kwargs)
+      t = graph.call_function(transpose_func, (dq,))
+      q = graph.call_function(q_op, (t,) + q_args, q_kwargs)
+
+    dq.meta = transpose_node_meta
+    t.meta = transpose_node_meta
+    q.meta = transpose_node_meta
+
+    target.replace_input_with(input_q, q)
+    return [dq, t, q]
+
+  def insert_layout_transpose(
+      self,
+      graph: torch.fx.Graph,
+      input_node: torch.fx.Node,
+      target_node: torch.fx.Node,
+      transpose_func: TransposeFunc,
+      transpose_node_meta: dict,
+  ) -> None:
+    assert transpose_func in (utils.tensor_to_nchw, utils.tensor_to_nhwc)
+
+    # new_nodes only contains Q/DQ/Transpose nodes, which are all SISO.
+    # Insertion order input nodes -> output nodes
+    new_nodes = []
+
+    # Constraint Q2: the NHWC partition's entry and exit must not be output
+    # edges of Q/DQ ops that are connected to a constant/weight tensor.
+    while layout_mark.is_const_node(input_node) and (
+        utils.is_dq_node(input_node) or utils.is_q_node(input_node)
+    ):
+      with graph.inserting_before(target_node):
+        new_input_node = graph.node_copy(input_node)
+
+      target_node.replace_input_with(input_node, new_input_node)
+
+      new_nodes = [new_input_node] + new_nodes
+      input_node, target_node = new_input_node.args[0], new_input_node
+
+    if utils.is_q_node(input_node):
+      # Constraint Q3: when the entry and exit is right after a q op (occur after a (dq-op-q)
+      # triplet), the transpose must be added as a quantized transpose in (dq-T-q)
+      # input_q -> (dq T q) -> target
+      new_nodes = (
+          self.insert_dq_t_q(
+              graph,
+              input_node,
+              target_node,
+              transpose_func,
+              transpose_node_meta,
+          )
+          + new_nodes
+      )
+    elif utils.is_dq_node(input_node):
+      # Constraint Q1: the NHWC partition's entry and exit cannot be edges
+      # within (dq-op-q) triplet.
+      # input_dq -> (T q dq) -> target
+      new_nodes = (
+          self.insert_t_q_dq(
+              graph,
+              input_node,
+              target_node,
+              transpose_func,
+              transpose_node_meta,
+          )
+          + new_nodes
+      )
+    else:
+      # input -> target
+      with graph.inserting_before(target_node):
+        t = graph.call_function(transpose_func, (input_node,))
+      t.meta = transpose_node_meta
+      target_node.replace_input_with(input_node, t)
+      new_nodes = [t] + new_nodes
+
+    # Mark new nodes as NCHW or NHWC
+    # For all nodes before the transpose, mark it as input_marker
+    # For all nodes after the transpose (incl. transpose), mark it as output_marker
+    if transpose_func == utils.tensor_to_nchw:
+      input_marker, target_marker = (
+          layout_mark.mark_as_nhwc_node,
+          layout_mark.mark_as_nchw_node,
+      )
+    else:
+      input_marker, target_marker = (
+          layout_mark.mark_as_nchw_node,
+          layout_mark.mark_as_nhwc_node,
+      )
+
+    marker = input_marker
+    for node in new_nodes:
+      if node.target == transpose_func:
+        marker = target_marker
+      marker(node)
+    assert marker == target_marker
+
+  def input_to_nhwc(
+      self,
+      graph: torch.fx.Graph,
+      input_node: torch.fx.Node,
+      target_node: torch.fx.Node,
+  ) -> None:
+    if layout_mark.is_nhwc_node(input_node):
+      return
+
+    if not layout_check.is_4d(input_node):
+      raise AssertionError(
+          f"Attempting to convert non-NHWC compatible node to NHWC: {input_node}"
+      )
+
+    # Assign target node's source meta to the to_NHWC node, because the transpose
+    # is added for the existence of target node.
+    self.insert_layout_transpose(
+        graph,
+        input_node,
+        target_node,
+        utils.tensor_to_nhwc,
+        self.get_source_meta(target_node),
+    )
+
+  def input_to_nchw(
+      self,
+      graph: torch.fx.Graph,
+      input_node: torch.fx.Node,
+      target_node: torch.fx.Node,
+  ) -> None:
+    if layout_mark.is_nchw_node(input_node):
+      return
+
+    self.insert_layout_transpose(
+        graph,
+        input_node,
+        target_node,
+        utils.tensor_to_nchw,
+        self.get_source_meta(input_node),
+    )
+
+  def mark_const_nodes(self, exported_program: torch.export.ExportedProgram):
+    graph_module = exported_program.graph_module
+    graph = graph_module.graph
+
+    input_specs = exported_program.graph_signature.input_specs
+    non_user_input_names = set()
+    for spec in input_specs:
+      if spec.kind != torch.export.graph_signature.InputKind.USER_INPUT:
+        non_user_input_names.add(spec.arg.name)
+
+    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))
+      if (
+          node.name in non_user_input_names
+          or (has_input_nodes and all_inputs_are_const)
+          or (node.op != "placeholder" and not has_input_nodes)
+      ):
+        layout_mark.mark_as_const_node(node)
+
+  def call(self, exported_program: torch.export.ExportedProgram):
+    self.mark_const_nodes(exported_program)
+
+    graph_module = exported_program.graph_module
+    if os.environ.get("AIEDGETORCH_LAYOUT_OPTIMIZE_USE_MINCUT_PARTITIONER"):
+      graph_module = layout_partitioners.min_cut.partition(graph_module)
+    else:
+      graph_module = layout_partitioners.greedy.partition(graph_module)
+
+    graph = graph_module.graph
+    for node in list(graph.nodes):
+      if layout_mark.is_nhwc_node(node):
+        for input_node in layout_check.get_layout_sensitive_inputs(node):
+          self.input_to_nhwc(graph, input_node, node)
+        layout_rewrite.rewrite_nhwc_node(node)
+      else:
+        for input_node in layout_check.get_layout_sensitive_inputs(node):
+          # Note: for non-4D tensors input_to_nchw is always noop.
+          self.input_to_nchw(graph, input_node, node)
+
+    graph_module.graph.eliminate_dead_code()
+    graph_module.recompile()
+    graph_module.graph.lint()
+    # Mark const node again for debugging
+    self.mark_const_nodes(exported_program)
+
+    return ExportedProgramPassResult(exported_program, True)