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

ai_edge_torch/hlfb/mark_pattern/pattern.py

@@ -0,0 +1,260 @@
+# 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 dataclasses
+from typing import Any, Callable, Optional, Union
+
+import torch
+from torch.export.graph_signature import TensorArgument
+from torch.fx import Graph
+from torch.fx import GraphModule
+from torch.fx.passes.utils.matcher_utils import InternalMatch
+from torch.fx.passes.utils.matcher_utils import SubgraphMatcher
+
+from ai_edge_torch.hlfb.mark_pattern import passes
+
+
+def _are_equal(x: Any, y: Any) -> bool:
+  if type(x) != type(y):
+    return False
+  if type(x) in [int, str]:
+    return x == y
+  if isinstance(x, float):
+    rel_tol = 1e-07
+    abs_tol = 0.0
+    return abs(x - y) <= max(rel_tol * max(abs(x), abs(y)), abs_tol)
+  if isinstance(x, list):
+    if len(x) != len(y):
+      return False
+    return all([_are_equal(a, b) for a, b in zip(x, y)])
+
+  raise Exception(f"Cannot compare type: {type(x)}")
+
+
+@dataclasses.dataclass
+class ScalarAttrTracker:
+  """ScalarAttrTracker is used to track the occurrence of a pattern's
+  scalar arg/attr in the pattern decomposed graph. Since a scalar attr
+  to the pattern can be transformed and turned into a/some ops' scalar
+  arg in the decomposed graph, it would be hard to programmatically get
+  the attr value from the pattern match. With the tracker and tracking info,
+  we could target the position of the decomposed op's scalar arg derived
+  from the pattern arg/attr and retrieve the value from the InternalMatch.
+
+  Args:
+    name (str): name of the attr to track.
+    pattern_arg_pos (int): the index of the attr to track in the pattern's
+      export_args.
+    transform (Callable): the transform function used when targeting the
+      occurrence of the attr value in the decomposed graph. An attr value
+      may be transformed during the decomposition and appear as a derived
+      value.
+    inverse_transform (Callable): the inverse transform function that maps
+      the transformed value back to the original attr value.
+  """
+
+  attr_name: str
+  pattern_arg_pos: int
+  transform: Callable = lambda x: x
+  inverse_transform: Callable = lambda x: x
+  _source_targets: list[tuple[Any, Any]] = dataclasses.field(default_factory=list)
+
+  def track(self, *sources):
+    """Register magic values to track the (transformed) attr values in
+    the pattern decomposed graph.
+    """
+    for source in sources:
+      target = self.transform(source)
+      if not _are_equal(self.inverse_transform(target), source):
+        raise Exception(f"Invalid transform/inverse_transform for {self.attr_name}")
+      self._source_targets.append([source, target])
+    return self
+
+
+@dataclasses.dataclass
+class ScalarAttrLocation:
+  attr_name: str
+  node_name: str
+  pos: Union[int, str]
+  _tracker: ScalarAttrTracker
+
+  @property
+  def index(self):
+    return self.pos if isinstance(self.pos, int) else None
+
+  @property
+  def key(self):
+    return self.pos if isinstance(self.pos, str) else None
+
+
+def _find_scalar_attr(
+    pattern_module: torch.nn.Module, export_args: tuple[Any], tracker: ScalarAttrTracker
+) -> ScalarAttrLocation:
+  scalar_loc_intersections = None
+  for source, target in tracker._source_targets:
+    track_args = list(export_args)
+    track_args[tracker.pattern_arg_pos] = source
+    ep = torch.export.export(pattern_module, tuple(track_args))
+
+    scalar_locs = set()
+    nodes = ep.graph_module.graph.nodes
+    for n in nodes:
+      for arg_pos, arg in enumerate(n.args):
+        if type(arg) == type(target) and arg == target:
+          scalar_locs.add((n.name, arg_pos))
+      for attr, val in n.kwargs.items():
+        if type(val) == type(target) and val == target:
+          scalar_locs.add((n.name, attr))
+
+    if scalar_loc_intersections is None:
+      scalar_loc_intersections = scalar_locs
+    else:
+      scalar_loc_intersections = scalar_loc_intersections & scalar_locs
+
+    if not scalar_loc_intersections:
+      break
+
+  if not scalar_loc_intersections:
+    return None
+  # Choose any occurrence as the attr provider
+  node_name, pos = scalar_loc_intersections.pop()
+  return ScalarAttrLocation(tracker.attr_name, node_name, pos, _tracker=tracker)
+
+
+class Pattern:
+
+  def __init__(
+      self,
+      name: str,
+      module: Union[Callable, torch.nn.Module],
+      export_args: tuple[Any],
+      *,
+      attr_builder: Callable[
+          ["Pattern", GraphModule, InternalMatch], Optional[dict[str, Any]]
+      ] = None,
+      scalar_attr_trackers: list[ScalarAttrTracker] = None,
+  ):
+    """The PyTorch computation pattern to match against a model.
+
+    Args:
+      name (str): the name of the pattern. It would be propagated to
+        the `name` attr in StableHLO composite ops for the matched
+        model subgraphs in the lowering.
+      module (torch.nn.Module or Callable): the PyTorch computation.
+      export_args (tuple[Any]): the args used to export the pattern module
+        with torch.export.export. If export_args contains non-tensor
+        Python scalars, there must be a corresponding attr tracker
+        in `scalar_attr_trackers` for each scalar arg.
+      attr_builder (Callable[[Pattern, GraphModule, InternalMatch], Optional[dict[str, Any]]]):
+        the callable that produces the a scalar attrs dict, which would be
+        propagated to `attr` in StableHLO composite ops for the matched
+        model subgraphs in the lowering.
+      scalar_attr_trackers (list[ScalarAttrTracker]): the trackers
+        for scalar args in `export_args`, which are used to track
+        the attr occurrence(s) and retrieve their values from the
+        matched subgraph.
+    """
+    if not isinstance(module, torch.nn.Module):
+
+      class PatternModule(torch.nn.Module):
+
+        def __init__(self, func):
+          super().__init__()
+          self.func = func
+
+        def forward(self, *args, **kwargs):
+          return self.func(*args, **kwargs)
+
+      module = PatternModule(module).eval()
+
+    self.name = name
+    self.exported_program = torch.export.export(module, export_args)
+    self.graph_module = self.exported_program.graph_module
+    self.attr_builder = attr_builder
+    self._scalar_attr_trackers = scalar_attr_trackers if scalar_attr_trackers else []
+
+    # Sanitize graph_module for more precise pattern matching.
+    # The graph_module to match against this pattern should apply equivalent
+    # sanitization.
+    self.graph_module = passes.remove_clone_ops(self.graph_module)
+    self.graph_module = passes.remove_dangling_args(self.graph_module)
+
+    self._scalar_attr_locations = []
+    for tracker in self._scalar_attr_trackers:
+      self._scalar_attr_locations.append(
+          _find_scalar_attr(module, export_args, tracker)
+      )
+
+    # Builds list of ordered input and output nodes.
+    self.graph_nodes_map = {}
+    for node in self.graph_module.graph.nodes:
+      self.graph_nodes_map[node.name] = node
+
+    self.input_nodes = tuple(
+        self.graph_nodes_map[spec.arg.name]
+        for spec in self.exported_program.graph_signature.input_specs
+        if isinstance(spec.arg, TensorArgument)
+    )
+    self.output_nodes = tuple(
+        self.graph_nodes_map[spec.arg.name]
+        for spec in self.exported_program.graph_signature.output_specs
+    )
+
+  def register_attr_builder(self, attr_builder):
+    self.attr_builder = attr_builder
+    return attr_builder
+
+  def match(
+      self,
+      graph_module: GraphModule,
+  ) -> list[tuple[InternalMatch, dict[str, Any]]]:
+    matcher = SubgraphMatcher(
+        self.graph_module.graph,
+        match_output=False,
+        match_placeholder=False,
+        remove_overlapping_matches=True,
+        ignore_literals=True,
+    )
+    matches = matcher.match(graph_module.graph)
+
+    match_with_attrs = []
+    # Graph traversal must be done in the reverser order (from SubgraphMatcher).
+    for match in matches[::-1]:
+      if self.attr_builder is not None:
+        attrs = self.attr_builder(self, graph_module, match)
+      else:
+        attrs = {}
+
+      for loc in self._scalar_attr_locations:
+        attrs[loc.attr_name] = self._get_attr_value_from_pattern_match(match, loc)
+
+      attrs = attrs if attrs else None
+      match_with_attrs.append((match, attrs))
+    return match_with_attrs
+
+  def _get_attr_value_from_pattern_match(
+      self,
+      match: InternalMatch,
+      loc: ScalarAttrLocation,
+  ):
+    matched_val = None
+    for k, v in match.nodes_map.items():
+      if k.name == loc.node_name:
+        if loc.index:
+          matched_val = v.args[loc.index]
+        elif loc.key in v.kwargs.keys():
+          matched_val = v.kwargs[loc.key]
+    attr_val = loc._tracker.inverse_transform(matched_val)
+    return attr_val

ai_edge_torch/hlfb/test/__init__.py

@@ -0,0 +1,14 @@
+# 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.
+# ==============================================================================

ai_edge_torch/hlfb/test/test_mark_pattern.py

@@ -0,0 +1,133 @@
+# 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 unittest
+
+import torch
+import torch_xla
+
+from ai_edge_torch.hlfb import mark_pattern
+
+
+def _export_stablehlo_mlir(model, args=None):
+  if not isinstance(model, torch.export.ExportedProgram):
+    ep = torch.export.export(model, args)
+  else:
+    ep = model
+  stablehlo_gm = torch_xla.stablehlo.exported_program_to_stablehlo(ep)
+  return stablehlo_gm.get_stablehlo_text()
+
+
+class TestMarkPattern(unittest.TestCase):
+
+  def test_mark_pattern(self):
+
+    class TestModel(torch.nn.Module):
+
+      def forward(self, x):
+        return x * x + x + x
+
+    pattern = mark_pattern.Pattern(
+        "test.add",
+        lambda a, b: a + b,
+        export_args=(torch.rand(2, 2), torch.rand(2, 2)),
+    )
+
+    model = TestModel().eval()
+    args = (torch.rand(20, 20),)
+    exported_program = torch.export.export(model, args)
+    mark_pattern.mark_pattern(exported_program.graph_module, pattern)
+    mlir = _export_stablehlo_mlir(exported_program)
+
+    self.assertEqual(mlir.count('stablehlo.composite "test.add"'), 2)
+
+  def test_mark_pattern_with_attr_builder(self):
+    class TestModel(torch.nn.Module):
+
+      def forward(self, x):
+        return x * x * x + x - x * x + x
+
+    pattern = mark_pattern.Pattern(
+        "test.add",
+        lambda a, b: a + b,
+        export_args=(torch.rand(2, 2), torch.rand(2, 2)),
+        attr_builder=lambda *args: {"alias": "test.test_add"},
+    )
+
+    model = TestModel().eval()
+    args = (torch.rand(20, 20),)
+    exported_program = torch.export.export(model, args)
+    mark_pattern.mark_pattern(exported_program.graph_module, pattern)
+    mlir = _export_stablehlo_mlir(exported_program)
+
+    self.assertEqual(mlir.count('stablehlo.composite "test.add"'), 2)
+    self.assertEqual(mlir.count('composite_attributes = {alias = "test.test_add"}'), 2)
+
+  def test_mark_pattern_with_scalar_attr_tracker(self):
+    class TestModel(torch.nn.Module):
+
+      def forward(self, x):
+        r = x
+        for idx in range(5):
+          r = torch.nn.LogSoftmax(dim=idx % 2)(r) * x
+        return r
+
+    pattern = mark_pattern.Pattern(
+        "test.log_softmax",
+        lambda x, dim: torch.nn.functional.log_softmax(x, dim=dim),
+        export_args=(torch.rand(10, 10, 10), 1),
+        scalar_attr_trackers=[
+            mark_pattern.ScalarAttrTracker("dim", pattern_arg_pos=1)
+            .track(0)
+            .track(1)
+            .track(2),
+        ],
+    )
+
+    model = TestModel().eval()
+    args = (torch.rand(10, 10),)
+    exported_program = torch.export.export(model, args)
+    mark_pattern.mark_pattern(exported_program.graph_module, pattern)
+    mlir = _export_stablehlo_mlir(exported_program)
+
+    self.assertEqual(mlir.count('stablehlo.composite "test.log_softmax"'), 5)
+    self.assertEqual(mlir.count("composite_attributes = {dim = 0 : i64}"), 3)
+    self.assertEqual(mlir.count("composite_attributes = {dim = 1 : i64}"), 2)
+
+  def test_mark_tangent_model_and_pattern_input(self):
+    class TestModel(torch.nn.Module):
+
+      def forward(self, x, y):
+        z = torch.ops.aten.relu(x)
+        z = z + y
+        return z
+
+    pattern = mark_pattern.Pattern(
+        "test.relu",
+        lambda x: torch.ops.aten.relu(x),
+        export_args=(torch.rand(2, 2),),
+    )
+
+    model = TestModel().eval()
+    args = (torch.rand(20, 20), torch.rand(20, 20))
+    exported_program = torch.export.export(model, args)
+    mark_pattern.mark_pattern(exported_program.graph_module, pattern)
+    mlir = _export_stablehlo_mlir(exported_program)
+
+    self.assertEqual(mlir.count('stablehlo.composite "test.relu'), 1)
+
+
+if __name__ == "__main__":
+  unittest.main()

ai_edge_torch/hlfb/test/test_stablehlo_composite_builder.py

@@ -0,0 +1,270 @@
+# 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 math
+import unittest
+
+import torch
+import torch.nn.functional as F
+import torch_xla
+
+from ai_edge_torch.hlfb import StableHLOCompositeBuilder
+
+
+def _export_stablehlo_mlir(model, args):
+  ep = torch.export.export(model, args)
+  stablehlo_gm = torch_xla.stablehlo.exported_program_to_stablehlo(ep)
+  return stablehlo_gm.get_stablehlo_text()
+
+
+class TestStableHLOCompositeBuilder(unittest.TestCase):
+
+  def test_build_composite(self):
+    class SampleModel(torch.nn.Module):
+
+      def forward(self, x):
+        builder = StableHLOCompositeBuilder(name="test.plus_two")
+        y = x + 1
+        y = builder.mark_inputs(y)
+        z = y + 2
+        z = builder.mark_outputs(z)
+        return z
+
+    mlir = _export_stablehlo_mlir(SampleModel().eval(), (torch.rand((2, 2)),))
+    self.assertEqual(mlir.count('stablehlo.composite "test.plus_two"'), 1)
+
+  def test_build_multiple_composites(self):
+    class SampleModel(torch.nn.Module):
+
+      def plus_one(self, x: torch.Tensor):
+        builder = StableHLOCompositeBuilder("test.plus_one")
+        x = builder.mark_inputs(x)
+        y = x + 1
+        y = builder.mark_outputs(y)
+        return y
+
+      def plus_two(self, x: torch.Tensor):
+        builder = StableHLOCompositeBuilder("test.plus_two")
+        x = builder.mark_inputs(x)
+        y = x + 2
+        y = builder.mark_outputs(y)
+        return y
+
+      def forward(self, x):
+        x = self.plus_two(x)
+        x = x + 3
+        x = self.plus_one(x)
+        x = x + 4
+        x = self.plus_two(x)
+        return x
+
+    mlir = _export_stablehlo_mlir(SampleModel().eval(), (torch.rand((2, 2)),))
+    self.assertEqual(mlir.count('stablehlo.composite "test.plus_one"'), 1)
+    self.assertEqual(mlir.count('stablehlo.composite "test.plus_two"'), 2)
+
+  def test_build_composite_with_attr(self):
+    class SampleModel(torch.nn.Module):
+
+      def __init__(self):
+        super().__init__()
+
+      def log_softmax(self, x: torch.Tensor, dim: int):
+        builder = StableHLOCompositeBuilder(name="test.log_softmax", attr={"dim": dim})
+        x = builder.mark_inputs(x)
+        y = torch.nn.functional.log_softmax(x, dim=dim)
+        y = builder.mark_outputs(y)
+        return y
+
+      def forward(self, x):
+        x = x + 1
+        x = self.log_softmax(x, 0)
+        x = self.log_softmax(x, 1)
+        return x
+
+    mlir = _export_stablehlo_mlir(SampleModel().eval(), (torch.rand((2, 2)),))
+    self.assertEqual(mlir.count('stablehlo.composite "test.log_softmax"'), 2)
+    self.assertEqual(mlir.count("composite_attributes = {dim = 0 : i64}"), 1)
+    self.assertEqual(mlir.count("composite_attributes = {dim = 1 : i64}"), 1)
+
+  def test_build_composite_with_mix_type_attrs(self):
+    class SampleModel(torch.nn.Module):
+
+      def __init__(self):
+        super().__init__()
+
+      def log_softmax(self, x: torch.Tensor, dim: int):
+        builder = StableHLOCompositeBuilder(
+            name="test.log_softmax",
+            attr={
+                "dim": dim,
+                "source": "torch.nn",
+                "version": 1.0,
+            },
+        )
+        x = builder.mark_inputs(x)
+        y = torch.nn.functional.log_softmax(x, dim=dim)
+        y = builder.mark_outputs(y)
+        return y
+
+      def forward(self, x):
+        x = x + 1
+        x = self.log_softmax(x, 0)
+        return x
+
+    mlir = _export_stablehlo_mlir(SampleModel().eval(), (torch.rand((2, 2)),))
+    self.assertEqual(mlir.count('stablehlo.composite "test.log_softmax"'), 1)
+    self.assertEqual(
+        mlir.count(
+            'composite_attributes = {dim = 0 : i64, source = "torch.nn", version = 1.000000e+00 : f32}'
+        ),
+        1,
+    )
+
+  def test_sdpa_composite(self):
+    class SDPAModel(torch.nn.Module):
+
+      def scaled_dot_product_attention(
+          self,
+          q: torch.Tensor,
+          k: torch.Tensor,
+          v: torch.Tensor,
+          head_size: int,
+          mask: torch.Tensor,
+      ):
+        builder = StableHLOCompositeBuilder("test.scaled_dot_product_attention")
+        q, k, v, mask = builder.mark_inputs(q, k, v, mask)
+
+        scale = 1.0 / math.sqrt(head_size)
+
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        y = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=mask,
+            dropout_p=0.0,
+            is_causal=mask is None,
+            scale=scale,
+        )
+        result = y.transpose(1, 2)
+        result = builder.mark_outputs(result)
+        return result
+
+      def forward(self, q, k, v, mask):
+        x = self.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            8,
+            mask,
+        )
+        return x
+
+    query = torch.rand(1, 1, 32, 4)
+    key = torch.rand(1, 500, 1, 4)
+    value = torch.rand(1, 500, 1, 4)
+    mask = torch.rand(1, 1, 1, 500)
+
+    mlir = _export_stablehlo_mlir(
+        SDPAModel().eval(),
+        (query, key, value, mask),
+    )
+    self.assertEqual(
+        mlir.count('stablehlo.composite "test.scaled_dot_product_attention"'), 1
+    )
+
+  def test_sdpa_composite_with_attr(self):
+    class SDPAModel(torch.nn.Module):
+
+      def scaled_dot_product_attention(
+          self,
+          q: torch.Tensor,
+          k: torch.Tensor,
+          v: torch.Tensor,
+          head_size: int,
+          include_captanh: bool,
+      ):
+        builder = StableHLOCompositeBuilder(
+            name="test.scaled_dot_product_attention",
+            attr={"include_captanh": include_captanh},
+        )
+        q, k, v = builder.mark_inputs(q, k, v)
+
+        scale = 1.0 / math.sqrt(head_size)
+
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        y = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=None,
+            dropout_p=0.0,
+            is_causal=True,
+            scale=scale,
+        )
+        result = y.transpose(1, 2)
+        result = builder.mark_outputs(result)
+        return result
+
+      def forward(self, q, k, v):
+        x = self.scaled_dot_product_attention(q, k, v, 8, True)
+        y = self.scaled_dot_product_attention(q, k, v, 8, False)
+        return x + y
+
+    query = torch.rand(1, 1, 32, 4)
+    key = torch.rand(1, 500, 1, 4)
+    value = torch.rand(1, 500, 1, 4)
+    mlir = _export_stablehlo_mlir(
+        SDPAModel().eval(),
+        (query, key, value),
+    )
+    self.assertEqual(
+        mlir.count('stablehlo.composite "test.scaled_dot_product_attention"'), 2
+    )
+    self.assertEqual(mlir.count("composite_attributes = {include_captanh = true}"), 1)
+    self.assertEqual(mlir.count("composite_attributes = {include_captanh = false}"), 1)
+
+  def test_build_composite_with_multiple_inputs_outputs(self):
+    class SampleModel(torch.nn.Module):
+
+      def mimo_sample(self, a, b, c):
+        builder = StableHLOCompositeBuilder(name="test.mimo_sample")
+
+        a, b, c = builder.mark_inputs(a, b, c)
+        x = a + b + c
+        y = (a - b) * x
+        z = (c + 1.0) * a
+        x, y, z = builder.mark_outputs(x, y, z)
+
+        result = x + y * z
+        return result
+
+      def forward(self, a, b, c):
+        x = self.mimo_sample(a, b, c)
+        x = self.mimo_sample(a, b, x)
+        x = self.mimo_sample(x, x, c)
+        return x
+
+    mlir = _export_stablehlo_mlir(
+        SampleModel().eval(), (torch.rand(2), torch.rand(2), torch.rand(2))
+    )
+    self.assertEqual(mlir.count('stablehlo.composite "test.mimo_sample"'), 3)
+
+
+if __name__ == "__main__":
+  unittest.main()