xref: /aosp_15_r20/external/executorch/backends/xnnpack/partition/configs.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
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
7import torch
8from executorch.exir.dialects._ops import ops as exir_ops
9
10"""
11** How to incorporate a new op into the XNNPACK Partitioner? **
12
13[1] When the new edge op being added is direct descendent of a core-aten op,
14and is also supported* by XNNPACK, prefer partitioning it via SUPPORTED_OPS
15mechanism e.g. torch.add
16
17[2] When the new op being added is not a core-aten op,
18
19[2.1] If the original torch op is supported* by XNNPACK, prefer partitioning it
20via SUPPORTED_MODULES. This will require "recomposing" the op before lowering
21it to XNNPACK e.g. torch.nn.Linear. Make sure to include all variants of the
22modules in the SUPPORTED_MODULES list.
23
24[2.2] If the original torch op is not supported by XNNPACK, then it is assumed
25that out of all the decomposed core-aten ops, SUPPORTED_OPS will be lowered to
26XNNPACK.
27
28* - Supported fully or partially. The partial support does not mean only few
29ops from the decomposition but means only some variants of the op "modes"
30possible with the arg combinations.
31"""
32
33SUPPORTED_OPS = [
34    exir_ops.edge.aten.div.Tensor,
35    exir_ops.edge.aten.add.Tensor,
36    exir_ops.edge.aten.clamp.default,
37    exir_ops.edge.aten.sub.Tensor,
38    exir_ops.edge.aten.floor.default,
39    exir_ops.edge.aten.maximum.default,
40    exir_ops.edge.aten.minimum.default,
41    exir_ops.edge.aten.mul.Tensor,
42    exir_ops.edge.aten.constant_pad_nd.default,
43    exir_ops.edge.aten.upsample_bilinear2d.default,
44    exir_ops.edge.aten.mean.dim,
45    exir_ops.edge.aten.max.dim,
46    exir_ops.edge.aten.max_pool2d_with_indices.default,
47    exir_ops.edge.aten.hardtanh.default,
48    exir_ops.edge.aten.sqrt.default,
49    exir_ops.edge.aten.ceil.default,
50    exir_ops.edge.aten.hardswish.default,
51    exir_ops.edge.aten.neg.default,
52    exir_ops.edge.aten.pow.Tensor_Scalar,
53    exir_ops.edge.aten.abs.default,
54    exir_ops.edge.aten._prelu_kernel.default,
55    exir_ops.edge.aten.slice_copy.Tensor,
56    exir_ops.edge.aten.relu.default,
57    exir_ops.edge.aten.hardtanh.default,
58    exir_ops.edge.aten.permute_copy.default,
59    exir_ops.edge.aten.sigmoid.default,
60    exir_ops.edge.aten._softmax.default,
61    exir_ops.edge.aten.cat.default,
62    exir_ops.edge.aten.elu.default,
63    exir_ops.edge.aten.avg_pool2d.default,
64    exir_ops.edge.aten.leaky_relu.default,
65    exir_ops.edge.aten.addmm.default,  # TODO(T163877189) add constraint for addmm
66]
67
68SUPPORTED_MODULES = [
69    torch.nn.Conv1d,
70    # TODO(T161981984) recomposed hardswish into a single node
71    torch.nn.Hardswish,  # we need to recompose
72    torch.nn.Hardsigmoid,  # we can handle decomposition
73    torch.nn.BatchNorm2d,
74    torch.nn.BatchNorm1d,
75    torch.nn.Conv2d,
76    torch.nn.Linear,
77    torch.nn.functional.linear,
78    torch.nn.PReLU,  # Without this, the PReLU weight becomes not a get_attr
79]
80
81# TODO delete this and should use SUPPORTED_OPS instead once we align fp32 and quant support
82SUPPORTED_QUANT_OPS = [
83    exir_ops.edge.aten.add.Tensor,
84    exir_ops.edge.aten.clamp.default,
85    exir_ops.edge.aten.relu.default,
86    exir_ops.edge.aten.sub.Tensor,
87    exir_ops.edge.aten.mul.Tensor,
88    exir_ops.edge.aten.mean.dim,
89    exir_ops.edge.aten.hardtanh.default,
90    exir_ops.edge.aten.slice_copy.Tensor,
91    exir_ops.edge.aten.permute_copy.default,
92    exir_ops.edge.aten.hardtanh.default,
93    exir_ops.edge.aten.mean.dim,
94    exir_ops.edge.aten.cat.default,
95    exir_ops.edge.aten.max_pool2d_with_indices.default,
96    exir_ops.edge.aten.max_pool2d.default,
97    exir_ops.edge.aten.constant_pad_nd.default,
98    exir_ops.edge.aten.elu.default,
99    exir_ops.edge.aten.t_copy.default,
100    exir_ops.edge.aten.leaky_relu.default,
101    exir_ops.edge.aten.addmm.default,  # TODO(T163877189) add constraint for addmm
102]
103
104# This set is used to determine if an op is a supported Quantized Op. This is
105# used to determine whether a quantization op is implicit or explicit.
106SUPPORTED_IMPLICIT_Q_DQ_OP_NAMES_SET = {
107    op.name()
108    for op in (
109        SUPPORTED_QUANT_OPS
110        + [
111            exir_ops.edge.aten._to_copy.default,
112            exir_ops.edge.aten.linear.default,
113            exir_ops.edge.aten.convolution.default,
114        ]
115    )
116}
117
118UNSUPPORTED_QUANT_MODULES = [
119    torch.nn.Hardswish,
120    torch.nn.Hardsigmoid,
121]
122
123# TODO delete this and should use SUPPORTED_MODULES instead once we align fp32 and quant support
124SUPPORTED_QUANT_MODULES = [
125    torch.nn.Linear,
126    torch.nn.functional.linear,
127    # TODO - T158982884
128    # torch.ao.nn.quantized.reference.modules.linear.Linear,
129    torch.nn.Conv1d,
130    torch.nn.functional.conv1d,
131    torch.ao.nn.quantized.reference.modules.conv.Conv1d,
132    torch.nn.Conv2d,
133    torch.nn.functional.conv2d,
134    torch.ao.nn.quantized.reference.modules.conv.Conv2d,
135    torch.nn.BatchNorm1d,
136    torch.nn.BatchNorm2d,
137]
138
139SUPPORTED_IMPLICIT_Q_DQ_MODULES_SET = set(SUPPORTED_QUANT_MODULES)
140
141# Modules which support dynamic quantization
142# These already support dynamic shape.
143SUPPORTED_DYN_QUANT_LINEAR_MODULES = [
144    torch.nn.Linear,
145    torch.nn.functional.linear,
146]
147
148SUPPORTED_DYN_QUANT_MODULES = SUPPORTED_DYN_QUANT_LINEAR_MODULES
149
150# XNNPACK supports majority of shape dynamism, however some ops are
151# explicitly static, so we maintain a set here to exclude them from
152# dynamic shape support.
153STATIC_OPS = [
154    exir_ops.edge.aten.cat.default,
155    exir_ops.edge.aten.slice_copy.Tensor,
156]
157
158STATIC_MODULES = []
159