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

ai_edge_torch/convert/fx_passes/__init__.py

@@ -0,0 +1,59 @@
+# 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 typing import Sequence, Union
+
+from torch.export import ExportedProgram
+from torch.fx.passes.infra.pass_manager import pass_result_wrapper
+import torch.utils._pytree as pytree
+
+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 FxPassBase
+from ai_edge_torch.convert.fx_passes._pass_base import FxPassResult
+from ai_edge_torch.convert.fx_passes.build_aten_composite_pass import BuildAtenCompositePass  # NOQA
+from ai_edge_torch.convert.fx_passes.build_upsample_bilinear2d_composite_pass import BuildUpsampleBilinear2DCompositePass  # NOQA
+from ai_edge_torch.convert.fx_passes.canonicalize_pass import CanonicalizePass
+from ai_edge_torch.convert.fx_passes.inject_mlir_debuginfo_pass import InjectMlirDebuginfoPass  # NOQA
+from ai_edge_torch.convert.fx_passes.optimize_layout_transposes_pass import OptimizeLayoutTransposesPass  # NOQA
+
+
+# TODO(cnchan): make a PassManager class.
+def run_passes(
+    exported_program: ExportedProgram,
+    passes: Sequence[Union[ExportedProgramPassBase, FxPassBase]],
+) -> ExportedProgram:
+  passes, _ = pytree.tree_flatten(passes)
+  for pass_ in passes:
+    if not isinstance(pass_, ExportedProgramPassBase):
+      pass_ = pass_result_wrapper(pass_)
+    if isinstance(pass_, ExportedProgramPassBase):
+      exported_program = pass_(exported_program).exported_program
+    else:
+      gm = exported_program.graph_module
+      gm, modified = pass_(gm)
+      if modified and gm is not exported_program.graph_module:
+        exported_program = ExportedProgram(
+            root=gm,
+            graph=gm.graph,
+            graph_signature=exported_program.graph_signature,
+            state_dict=exported_program.state_dict,
+            range_constraints=exported_program.range_constraints,
+            module_call_graph=exported_program.module_call_graph,
+            example_inputs=exported_program.example_inputs,
+            verifier=exported_program.verifier,
+            constants=exported_program.constants,
+        )
+  return exported_program

ai_edge_torch/convert/fx_passes/_pass_base.py

@@ -0,0 +1,49 @@
+# 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 abc
+from collections import namedtuple
+
+import torch
+from torch.export import ExportedProgram
+from torch.fx.passes.infra.pass_base import PassBase as FxPassBase
+from torch.fx.passes.infra.pass_base import PassResult as FxPassResult
+
+
+class ExportedProgramPassResult(
+    namedtuple("ExportedProgramPassResult", ["exported_program", "modified"])
+):
+
+  def __new__(cls, exported_program, modified):
+    return super().__new__(cls, exported_program, modified)
+
+
+class ExportedProgramPassBase(abc.ABC):
+
+  def __call__(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
+    self.requires(exported_program)
+    res = self.call(exported_program)
+    self.ensures(exported_program)
+    return res
+
+  @abc.abstractmethod
+  def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
+    pass
+
+  def requires(self, exported_program: ExportedProgram) -> None:
+    pass
+
+  def ensures(self, exported_program: ExportedProgram) -> None:
+    pass

ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py

@@ -0,0 +1,192 @@
+# 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 copy
+import functools
+from typing import Any, Callable
+
+import torch
+from torch.fx import GraphModule
+from torch.fx import Node
+from torch.fx.passes.infra.pass_base import PassBase
+from torch.fx.passes.infra.pass_base import PassResult
+import torch.utils._pytree as pytree
+
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+
+_composite_builders: dict[Callable, Callable[[GraphModule, Node], None]] = {}
+
+
+def _register_composite_builder(op):
+  def inner(func):
+    if isinstance(op, torch._ops.OpOverloadPacket):
+      for overload in v.overloads():
+        _composite_builders[getattr(v, overload)] = func
+    else:
+      _composite_builders[op] = func
+    return func
+
+  return inner
+
+
+def _tree_map_to_composite_attr_values(values, *, stringify_incompatible_values=True):
+
+  def convert(value):
+    nonlocal stringify_incompatible_values
+    if value is None:
+      return "py_None"
+    if isinstance(value, (str, int, float, bool)):
+      return value
+
+    if stringify_incompatible_values:
+      return str(value)
+    return value
+
+  return pytree.tree_map(convert, values)
+
+
+class TorchOpArgumentsMapper:
+
+  def __init__(self, op):
+    if isinstance(op, torch._ops.OpOverloadPacket):
+      op = op.default
+
+    assert hasattr(op, "_schema")
+    self.op = op
+    self.arg_specs = [(spec.name, spec.default_value) for spec in op._schema.arguments]
+
+  def get_full_kwargs(self, args, kwargs=None) -> dict[str, Any]:
+    """Inspect the op's schema and extract all its args and kwargs
+    into one single kwargs dict, with default values for those
+    unspecified args and kwargs.
+    """
+    full_kwargs = {**(kwargs or {})}
+
+    for arg, (name, default_value) in zip(args, self.arg_specs):
+      full_kwargs[name] = arg
+
+    for name, default_value in self.arg_specs[len(args) :]:
+      if name not in full_kwargs:
+        full_kwargs[name] = default_value
+
+    return full_kwargs
+
+
+@_register_composite_builder(torch.ops.aten.hardswish.default)
+def _aten_hardswish(gm: GraphModule, node: Node):
+  op = node.target
+
+  def hardswish(self: torch.Tensor):
+    nonlocal op
+    builder = StableHLOCompositeBuilder("aten.hardswish.default")
+    self = builder.mark_inputs(self)
+    output = op(self)
+    output = builder.mark_outputs(output)
+    return output
+
+  node.target = hardswish
+
+
+@_register_composite_builder(torch.ops.aten.avg_pool2d.default)
+def _aten_avg_pool2d(gm: GraphModule, node: Node):
+  op = node.target
+  args_mapper = TorchOpArgumentsMapper(op)
+
+  def avg_pool2d(*args, **kwargs):
+    nonlocal op, args_mapper
+
+    full_kwargs = args_mapper.get_full_kwargs(args, kwargs)
+
+    def is_same_padding(
+        input_shape: list[int],
+        kernel_size: list[int],
+        stride: list[int],
+        padding: list[int],
+    ):
+      for dim_input_size, dim_kernel_size, dim_stride, dim_padding in zip(
+          input_shape, kernel_size, stride, padding
+      ):
+        dim_output_size = int((dim_input_size + dim_stride - 1) / dim_stride)
+        padding_needed = max(
+            0, (dim_output_size - 1) * dim_stride + dim_kernel_size - dim_input_size
+        )
+        if padding_needed % 2 != 0:
+          return False
+
+        if padding_needed // 2 != dim_padding:
+          return False
+      return True
+
+    def is_valid_padding(padding: list[int]):
+      return not any(padding)
+
+    # We prefer to avoid passing empty arrays to composite attributes
+    # as they will be lowered to an ArrayAttr so canonicalizing according
+    # to the default behaviour here.
+    if not full_kwargs["stride"]:
+      full_kwargs["stride"] = full_kwargs["kernel_size"]
+
+    # Only wrap in a composite when the underlying converter can handle it.
+    # TODO We should be able to remove this if the converter can inline composites when it can not handle them.
+
+    # We don't cover any cases where ceil_mode is True or divisor_override is set.
+    if full_kwargs["ceil_mode"] or full_kwargs["divisor_override"] is not None:
+      return op(*args, **kwargs)
+
+    # We also can not cover a case where count_include_pad is False but the padding is custom.
+    if (
+        not full_kwargs["count_include_pad"]
+        and not is_valid_padding(full_kwargs["padding"])
+        and not is_same_padding(
+            list(full_kwargs["self"].shape)[2:],
+            full_kwargs["kernel_size"],
+            full_kwargs["stride"],
+            full_kwargs["padding"],
+        )
+    ):
+      return op(*args, **kwargs)
+
+    builder = StableHLOCompositeBuilder(
+        "aten.avg_pool2d.default",
+        attr=_tree_map_to_composite_attr_values(
+            {
+                "kernel_size": full_kwargs["kernel_size"],
+                "stride": full_kwargs["stride"],
+                "padding": full_kwargs["padding"],
+                "ceil_mode": full_kwargs["ceil_mode"],
+                "count_include_pad": full_kwargs["count_include_pad"],
+                "divisor_override": full_kwargs["divisor_override"],
+            }
+        ),
+    )
+
+    full_kwargs["self"] = builder.mark_inputs(full_kwargs["self"])
+    output = op(**full_kwargs)
+    output = builder.mark_outputs(output)
+    return output
+
+  node.target = avg_pool2d
+
+
+class BuildAtenCompositePass(PassBase):
+
+  def call(self, graph_module: GraphModule):
+    for node in graph_module.graph.nodes:
+      if node.target in _composite_builders:
+        _composite_builders[node.target](graph_module, node)
+
+    graph_module.graph.lint()
+    graph_module.recompile()
+    return PassResult(graph_module, True)

ai_edge_torch/convert/fx_passes/build_upsample_bilinear2d_composite_pass.py

@@ -0,0 +1,84 @@
+# 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 functools
+
+import torch
+
+from ai_edge_torch.convert.fx_passes import FxPassBase
+from ai_edge_torch.convert.fx_passes import FxPassResult
+from ai_edge_torch.hlfb import mark_pattern
+
+
+@functools.cache
+def _get_upsample_bilinear2d_pattern():
+  pattern = mark_pattern.Pattern(
+      "odml.upsample_bilinear2d",
+      lambda x: torch.nn.functional.interpolate(
+          x, scale_factor=2, mode="bilinear", align_corners=False
+      ),
+      export_args=(torch.rand(1, 3, 100, 100),),
+  )
+
+  @pattern.register_attr_builder
+  def attr_builder(pattern, graph_module, internal_match):
+    output = internal_match.returning_nodes[0]
+    output_h, output_w = output.meta["val"].shape[-2:]
+    return {
+        "output": (int(output_h), int(output_w)),
+        "align_corners": False,
+    }
+
+  return pattern
+
+
+@functools.cache
+def _get_upsample_bilinear2d_align_corners_pattern():
+  pattern = mark_pattern.Pattern(
+      "odml.upsample_bilinear2d",
+      lambda x: torch.nn.functional.interpolate(
+          x, scale_factor=2, mode="bilinear", align_corners=True
+      ),
+      export_args=(torch.rand(1, 3, 100, 100),),
+  )
+
+  @pattern.register_attr_builder
+  def attr_builder(graph_module, pattern, internal_match):
+    output = internal_match.returning_nodes[0]
+    output_h, output_w = output.meta["val"].shape[-2:]
+    return {
+        "output": (int(output_h), int(output_w)),
+        "align_corners": True,
+    }
+
+  return pattern
+
+
+class BuildUpsampleBilinear2DCompositePass(FxPassBase):
+
+  def __init__(self):
+    super().__init__()
+    self._patterns = [
+        _get_upsample_bilinear2d_pattern(),
+        _get_upsample_bilinear2d_align_corners_pattern(),
+    ]
+
+  def call(self, graph_module: torch.fx.GraphModule):
+    for pattern in self._patterns:
+      graph_module = mark_pattern.mark_pattern(graph_module, pattern)
+
+    graph_module.graph.lint()
+    graph_module.recompile()
+    return FxPassResult(graph_module, True)

ai_edge_torch/convert/fx_passes/canonicalize_pass.py

@@ -0,0 +1,37 @@
+# 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 torch.export import ExportedProgram
+
+from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassBase
+from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassResult  # NOQA
+
+# A dummy decomp table for running ExportedProgram.run_decompositions without
+# any op decompositions but just aot_export_module. Due to the check in
+# run_decompositions, if None or an empty dict is passed as decomp_table,
+# it will run the default aten-coreaten decompositions. Therefore a non-empty
+# dummy decomp table is needed.
+# Ref: https://github.com/pytorch/pytorch/blob/db895ace1d36726e64781774f53b3d3098206116/torch/export/exported_program.py#L543
+_dummy_decomp_table = {
+    torch._ops.OperatorBase(): lambda: None,
+}
+
+
+class CanonicalizePass(ExportedProgramPassBase):
+
+  def call(self, exported_program: ExportedProgram):
+    exported_program = exported_program.run_decompositions(_dummy_decomp_table)
+    return ExportedProgramPassResult(exported_program, True)

ai_edge_torch/convert/fx_passes/inject_mlir_debuginfo_pass.py

@@ -0,0 +1,73 @@
+# 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 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):
+  def class_fullname(cls):
+    module = cls.__module__
+    if module == "builtins":
+      return cls.__qualname__
+    return module + "." + cls.__qualname__
+
+  def get_hierarchy(node: torch.fx.Node):
+    nn_module_stack = node.meta.get("nn_module_stack", {})
+    layers = []
+    for name, layer in nn_module_stack.values():
+      iid = ("_" + name.split(".")[-1]) if name else ""
+      layer_str = layer if isinstance(layer, str) else class_fullname(layer)
+      layers.append(layer_str + iid)
+
+    hierachy_str = "/".join(layers) + ";"
+    return hierachy_str
+
+  # TODO(yijieyang): Encode aten op and attrs.
+  return get_hierarchy(node)
+
+
+def _wrap_call_function_node_with_debuginfo_writer(node: torch.fx.GraphModule):
+  if not node.op.startswith("call_function"):
+    return
+
+  target = node.target
+  debuginfo = _get_mlir_debuginfo(node)
+
+  def debuginfo_writer(*args, **kwargs):
+    nonlocal target, debuginfo
+    outputs = target(*args, **kwargs)
+    outputs = pytree.tree_map_only(
+        torch.Tensor,
+        lambda x: torch.ops.xla.write_mlir_debuginfo(x, debuginfo),
+        outputs,
+    )
+    return outputs
+
+  node.target = debuginfo_writer
+
+
+class InjectMlirDebuginfoPass(PassBase):
+
+  def call(self, graph_module: torch.fx.GraphModule):
+    for node in graph_module.graph.nodes:
+      _wrap_call_function_node_with_debuginfo_writer(node)
+
+    graph_module.graph.lint()
+    graph_module.recompile()
+    return PassResult(graph_module, True)

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