Skip to content

Commit 1490af9

Browse files
ethanng72facebook-github-bot
ethanng72
authored and
facebook-github-bot
committed
Create decompose_ops.py and test_decompose_ops.py (#11299)
Summary: Create new class and test suite for passes that decompose an op into a equivalent series of simpler ops Test Plan: Imported from GitHub, without a `Test Plan:` line. Rollback Plan: Differential Revision: D75826474 Pulled By: ethanng72
1 parent 450d1f9 commit 1490af9

File tree

5 files changed

+251
-116
lines changed

5 files changed

+251
-116
lines changed

‎backends/cadence/aot/TARGETS‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,24 @@ python_library(
276276
],
277277
)
278278

279+
python_library(
280+
name = "decompose_ops",
281+
srcs = [
282+
"decompose_ops.py",
283+
],
284+
typing = True,
285+
deps = [
286+
":pass_utils",
287+
"//caffe2:torch",
288+
"//executorch/backends/cadence/aot:pass_utils",
289+
"//executorch/exir:pass_base",
290+
"//executorch/exir/dialects:lib",
291+
"//executorch/exir/dialects/edge:lib",
292+
"//executorch/exir/passes:spec_prop_pass",
293+
],
294+
)
295+
296+
279297
python_unittest(
280298
name = "test_graph_builder",
281299
srcs = [
@@ -314,6 +332,27 @@ python_unittest(
314332
],
315333
)
316334

335+
python_unittest(
336+
name = "test_decompose_ops_passes",
337+
srcs = [
338+
"tests/test_decompose_ops_passes.py",
339+
],
340+
supports_static_listing = False,
341+
typing = True,
342+
deps = [
343+
"fbsource//third-party/pypi/parameterized:parameterized",
344+
":compiler",
345+
":decompose_ops",
346+
"//caffe2:torch",
347+
"//executorch/backends/cadence/aot:compiler",
348+
"//executorch/backends/cadence/aot:graph_builder",
349+
"//executorch/backends/cadence/aot:pass_utils",
350+
"//executorch/exir:pass_base",
351+
"//executorch/exir/dialects:lib",
352+
"//executorch/exir/passes:lib",
353+
],
354+
)
355+
317356
python_unittest(
318357
name = "test_fusion_ops_passes",
319358
srcs = [

‎backends/cadence/aot/decompose_ops.py‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
# Copyright 2025 Arm Limited and/or its affiliates.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
7+
8+
9+
# This file contains all the functions that decompose one op into simpler ops in the
10+
# graph. The functions decomposing ops for models deployed with Jarvis are grouped
11+
# together in class 'DecomposeOpsInGraph'. Some examples of functions in the class are
12+
# 1. functions that decompose an ATen gelu op into an equivalent series of simpler ops
13+
14+
# pyre-unsafe
15+
16+
from typing import Dict, Tuple
17+
18+
from executorch.backends.cadence.aot.pass_utils import (
19+
CadencePassAttribute,
20+
register_cadence_pass,
21+
)
22+
from executorch.exir.dialects._ops import ops as exir_ops
23+
from executorch.exir.dialects.edge._ops import EdgeOpOverload
24+
from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue
25+
from torch.fx.node import Argument
26+
27+
# A map to represent ops that:
28+
# (a) are functionally equivalent wrt. Jarvis; and
29+
# (b) have identical arguments
30+
# An op whose target is 'key' in this dict can be replaced by the functionally euivalent
31+
# op whose target is 'value'. The replacement would just involve changing the op target.
32+
functionally_equivalent_op_targets: Dict[EdgeOpOverload, EdgeOpOverload] = {
33+
exir_ops.edge.aten.relu_.default: exir_ops.edge.aten.relu.default,
34+
exir_ops.edge.aten.unsafe_split.Tensor: exir_ops.edge.aten.split_copy.Tensor,
35+
}
36+
37+
38+
@register_cadence_pass(CadencePassAttribute(opt_level=0))
39+
class DecomposeAtenApproxGeluPass(ExportPass):
40+
"""
41+
Decompose the aten gelu op with an approximate arg to a series of simpler ops
42+
"""
43+
44+
def call_operator(
45+
self,
46+
op,
47+
args: Tuple[Argument, ...],
48+
kwargs: Dict[str, Argument],
49+
meta: NodeMetadata,
50+
) -> ProxyValue:
51+
# compute the approximate gelu (0.7978845608028654 is sqrt(2 / pi))
52+
# as 0.5 * x * (1 + torch.tanh(0.7978845608028654 * ( x + 0.044715 * x^3)))
53+
54+
# Get 0.5 * x
55+
half = super().call_operator(
56+
exir_ops.edge.aten.mul.Tensor,
57+
(args[0], 0.5),
58+
{},
59+
meta,
60+
)
61+
62+
scaled = super().call_operator(
63+
exir_ops.edge.aten.mul.Tensor,
64+
(args[0], 0.044715),
65+
{},
66+
meta,
67+
)
68+
69+
# Get x^2 (note that we use mul.Tensor twice instead of pow.Tensor because
70+
# it is much more efficient on DSP backends)
71+
scaled_square = super().call_operator(
72+
exir_ops.edge.aten.mul.Tensor,
73+
(scaled, args[0]),
74+
{},
75+
meta,
76+
)
77+
78+
# Get x^3
79+
scaled_cubed = super().call_operator(
80+
exir_ops.edge.aten.mul.Tensor,
81+
(scaled_square, args[0]),
82+
{},
83+
meta,
84+
)
85+
86+
# Get x + 0.044715 * x^3
87+
inner_sum = super().call_operator(
88+
exir_ops.edge.aten.add.Tensor,
89+
(scaled_cubed, args[0]),
90+
{},
91+
meta,
92+
)
93+
94+
# Get 0.7978845608028654 * ( x + 0.044715 * x^3)
95+
scaled_sum = super().call_operator(
96+
exir_ops.edge.aten.mul.Tensor,
97+
(inner_sum, 0.7978845608028654),
98+
{},
99+
meta,
100+
)
101+
102+
# Get torch.tanh(0.7978845608028654 * ( x + 0.044715 * x^3))
103+
tanh = super().call_operator(
104+
exir_ops.edge.aten.tanh.default,
105+
(scaled_sum,),
106+
{},
107+
meta,
108+
)
109+
110+
# Get 1 + torch.tanh(0.79788456 * ( x + 0.044715 * x^3))
111+
# TODO(): Check why this is not working properly with integer values (e.g. 1 instead of 1.)
112+
outer_sum = super().call_operator(
113+
exir_ops.edge.aten.add.Tensor,
114+
(tanh, 1.0),
115+
{},
116+
meta,
117+
)
118+
119+
# Return the final result
120+
return super().call_operator(
121+
exir_ops.edge.aten.mul.Tensor,
122+
(half, outer_sum),
123+
{},
124+
meta,
125+
)
126+
127+
128+
# This class encapsulates all the functions that decompose one op in the graph.
129+
class CadenceDecomposeOpsInGraph:
130+
passes = [
131+
DecomposeAtenApproxGeluPass,
132+
]

‎backends/cadence/aot/replace_ops.py‎

Lines changed: 1 addition & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -2078,89 +2078,11 @@ def call_operator(
20782078
kwargs: Dict[str, Argument],
20792079
meta: NodeMetadata,
20802080
) -> ProxyValue:
2081-
if "approximate" not in kwargs:
2082-
return super().call_operator(op, args, kwargs, meta)
2083-
20842081
if op not in {
20852082
exir_ops.edge.aten.gelu.default,
20862083
}:
20872084
return super().call_operator(op, args, kwargs, meta)
2088-
2089-
# compute the approximate gelu (0.7978845608028654 is sqrt(2 / pi))
2090-
# as 0.5 * x * (1 + torch.tanh(0.7978845608028654 * ( x + 0.044715 * x^3)))
2091-
2092-
# Get 0.5 * x
2093-
half = super().call_operator(
2094-
exir_ops.edge.aten.mul.Tensor,
2095-
(args[0], 0.5),
2096-
{},
2097-
meta,
2098-
)
2099-
2100-
scaled = super().call_operator(
2101-
exir_ops.edge.aten.mul.Tensor,
2102-
(args[0], 0.044715),
2103-
{},
2104-
meta,
2105-
)
2106-
2107-
# Get x^2 (note that we use mul.Tensor twice instead of pow.Tensor because
2108-
# it is much more efficient on DSP backends)
2109-
scaled_square = super().call_operator(
2110-
exir_ops.edge.aten.mul.Tensor,
2111-
(scaled, args[0]),
2112-
{},
2113-
meta,
2114-
)
2115-
2116-
# Get x^3
2117-
scaled_cubed = super().call_operator(
2118-
exir_ops.edge.aten.mul.Tensor,
2119-
(scaled_square, args[0]),
2120-
{},
2121-
meta,
2122-
)
2123-
2124-
# Get x + 0.044715 * x^3
2125-
inner_sum = super().call_operator(
2126-
exir_ops.edge.aten.add.Tensor,
2127-
(scaled_cubed, args[0]),
2128-
{},
2129-
meta,
2130-
)
2131-
2132-
# Get 0.7978845608028654 * ( x + 0.044715 * x^3)
2133-
scaled_sum = super().call_operator(
2134-
exir_ops.edge.aten.mul.Tensor,
2135-
(inner_sum, 0.7978845608028654),
2136-
{},
2137-
meta,
2138-
)
2139-
2140-
# Get torch.tanh(0.7978845608028654 * ( x + 0.044715 * x^3))
2141-
tanh = super().call_operator(
2142-
exir_ops.edge.aten.tanh.default,
2143-
(scaled_sum,),
2144-
{},
2145-
meta,
2146-
)
2147-
2148-
# Get 1 + torch.tanh(0.79788456 * ( x + 0.044715 * x^3))
2149-
# TODO(): Check why this is not working properly with integer values (e.g. 1 instead of 1.)
2150-
outer_sum = super().call_operator(
2151-
exir_ops.edge.aten.add.Tensor,
2152-
(tanh, 1.0),
2153-
{},
2154-
meta,
2155-
)
2156-
2157-
# Retunr the final result
2158-
return super().call_operator(
2159-
exir_ops.edge.aten.mul.Tensor,
2160-
(half, outer_sum),
2161-
{},
2162-
meta,
2163-
)
2085+
return super().call_operator(op, args, kwargs, meta)
21642086

21652087

21662088
# Adapted from fbcode/pyspeech/opt_passes/replace_ops.py

‎backends/cadence/aot/tests/test_decompose_ops_passes.py‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
# pyre-unsafe
8+
9+
import unittest
10+
from typing import Any, Callable, List, Tuple, Union
11+
12+
import torch
13+
from executorch.backends.cadence.aot.decompose_ops import DecomposeAtenApproxGeluPass
14+
from executorch.backends.cadence.aot.graph_builder import single_op_builder
15+
from executorch.backends.cadence.aot.pass_utils import count_node
16+
from executorch.exir.dialects._ops import ops as exir_ops
17+
from executorch.exir.pass_base import ExportPass
18+
19+
20+
class TestDecomposeOpsPasses(unittest.TestCase):
21+
def assertTargetCountEqual(
22+
self,
23+
graph_module: torch.fx.GraphModule,
24+
target: Union[Callable[..., Any], str],
25+
expected_count: int,
26+
):
27+
"""Helper function to check the number of nodes with a given target."""
28+
actual_count = count_node(graph_module, target)
29+
self.assertEqual(
30+
actual_count,
31+
expected_count,
32+
f"{target} count mismatch for graph {graph_module}",
33+
)
34+
35+
def assertTargetCountsEqual(
36+
self,
37+
graph_module: torch.fx.GraphModule,
38+
targets_and_counts: List[Tuple[Union[Callable[..., Any], str], int]],
39+
):
40+
"""Helper function to check the number of nodes of all types for a given target."""
41+
for target, expected_count in targets_and_counts:
42+
self.assertTargetCountEqual(graph_module, target, expected_count)
43+
44+
def test_decompose_aten_approximate_gelu(self):
45+
inputs = torch.randn(2, 1, 64)
46+
47+
gm = single_op_builder(
48+
placeholders=(inputs,),
49+
op=exir_ops.edge.aten.gelu.default,
50+
args=(inputs,),
51+
kwargs={"approximate": "tanh"},
52+
)
53+
gm = ExportPass().call(gm).graph_module
54+
55+
p = DecomposeAtenApproxGeluPass()
56+
graph_after_passes = p.call(gm).graph_module
57+
58+
# Assert that aten.gelu op was decomposed
59+
self.assertEqual(
60+
count_node(
61+
graph_after_passes,
62+
exir_ops.edge.aten.gelu.default,
63+
),
64+
0,
65+
)
66+
67+
# The decomposition should have one tanh, 2 add and 6 mul
68+
self.assertEqual(
69+
count_node(graph_after_passes, exir_ops.edge.aten.tanh.default),
70+
1,
71+
)
72+
self.assertEqual(
73+
count_node(graph_after_passes, exir_ops.edge.aten.add.Tensor),
74+
2,
75+
)
76+
self.assertEqual(
77+
count_node(graph_after_passes, exir_ops.edge.aten.mul.Tensor),
78+
6,
79+
)

0 commit comments

Comments
 (0)