move modules to ops (#257)

Summary:
Pull Request resolved: #257

Decrease the supported module surface. For ops that are canonical (not decomposed) and single op partitions, we can always partition just using the operator rather than the module (It is difficult to maintain every variant of the op's module, and if the op is used in a larger module like multihead attention then it is impossible to partition using source_fn)

In the future, we will want to only support module which require recomposition or some module level recomposition

Reviewed By: manuelcandales

Differential Revision: D49111303

fbshipit-source-id: eb46aa51cd129200fa3ecd45be34675ab0fa0a0e

Author: mcr229

backends/xnnpack/partition/configs.py

@@ -24,6 +24,7 @@
     exir_ops.edge.aten.upsample_bilinear2d.default,
     exir_ops.edge.aten.mean.dim,
     exir_ops.edge.aten.max.dim,
+    exir_ops.edge.aten.max_pool2d_with_indices.default,
     exir_ops.edge.aten.hardtanh.default,
     exir_ops.edge.aten.sqrt.default,
     exir_ops.edge.aten.ceil.default,
@@ -33,40 +34,49 @@
     exir_ops.edge.aten.abs.default,
     exir_ops.edge.aten._prelu_kernel.default,
     exir_ops.edge.aten.slice_copy.Tensor,
+    exir_ops.edge.aten.relu.default,
+    exir_ops.edge.aten.hardtanh.default,
+    exir_ops.edge.aten.permute_copy.default,
+    exir_ops.edge.aten.sigmoid.default,
+    exir_ops.edge.aten._softmax.default,
+    exir_ops.edge.aten.cat.default,
+    exir_ops.edge.aten.elu.default,
+    exir_ops.edge.aten.avg_pool2d.default,
+    exir_ops.edge.aten.leaky_relu.default,
 ]
 
 SUPPORTED_MODULES = [
     torch.nn.Conv1d,
     # TODO(T161981984) recomposed hardswish into a single node
-    torch.nn.Hardswish,
-    torch.nn.Hardsigmoid,
-    torch.nn.Conv2d,
-    torch.nn.ReLU,
-    torch.nn.Sigmoid,
-    torch.nn.Softmax,
-    torch.nn.BatchNorm1d,
+    torch.nn.Hardswish,  # we need to recompose
+    torch.nn.Hardsigmoid,  # we can handle decomposition
     torch.nn.BatchNorm2d,
+    torch.nn.BatchNorm1d,
+    torch.nn.Conv2d,
     torch.nn.Linear,
     torch.nn.functional.linear,
-    torch.nn.Hardtanh,
-    torch.nn.MaxPool2d,
-    torch.nn.LeakyReLU,
-    torch.nn.ELU,
-    torch.nn.AvgPool2d,
     torch.nn.PReLU,  # Without this, the PReLU weight becomes not a get_attr
-    torch.cat,
-    torch.concat,
-    torch.concatenate,
 ]
 
 # TODO delete this and should use SUPPORTED_OPS instead once we align fp32 and quant support
 SUPPORTED_QUANT_OPS = [
     exir_ops.edge.aten.add.Tensor,
+    exir_ops.edge.aten.clamp.default,
+    exir_ops.edge.aten.relu.default,
     exir_ops.edge.aten.sub.Tensor,
     exir_ops.edge.aten.mul.Tensor,
     exir_ops.edge.aten.mean.dim,
-    exir_ops.edge.aten.hardtanh.default,  # TODO - which one module or op or both?
+    exir_ops.edge.aten.hardtanh.default,
     exir_ops.edge.aten.slice_copy.Tensor,
+    exir_ops.edge.aten.permute_copy.default,
+    exir_ops.edge.aten.hardtanh.default,
+    exir_ops.edge.aten.mean.dim,
+    exir_ops.edge.aten.cat.default,
+    exir_ops.edge.aten.max_pool2d_with_indices.default,
+    exir_ops.edge.aten.max_pool2d.default,
+    exir_ops.edge.aten.constant_pad_nd.default,
+    exir_ops.edge.aten.elu.default,
+    exir_ops.edge.aten.leaky_relu.default,
 ]
 
 SUPPORTED_IMPLICIT_Q_DQ_OP_NAMES_SET = {
@@ -75,7 +85,6 @@
         SUPPORTED_QUANT_OPS
         + [
             exir_ops.edge.aten._to_copy.default,
-            exir_ops.edge.aten.max_pool2d.default,
             exir_ops.edge.aten.linear.default,
         ]
     )
@@ -88,37 +97,18 @@
 
 # TODO delete this and should use SUPPORTED_MODULES instead once we align fp32 and quant support
 SUPPORTED_QUANT_MODULES = [
-    torch.clamp,
-    torch.mean,
-    torch.permute,
-    torch.permute_copy,
-    torch.cat,
-    torch.concat,
-    torch.concatenate,
     torch.nn.Linear,
     torch.nn.functional.linear,
     # TODO - T158982884
     # torch.ao.nn.quantized.reference.modules.linear.Linear,
-    torch.nn.MaxPool2d,
     torch.nn.Conv1d,
     torch.nn.functional.conv1d,
     torch.ao.nn.quantized.reference.modules.conv.Conv1d,
     torch.nn.Conv2d,
     torch.nn.functional.conv2d,
-    torch.nn.functional.pad,
-    torch.nn.functional.elu,
     torch.ao.nn.quantized.reference.modules.conv.Conv2d,
     torch.nn.BatchNorm1d,
     torch.nn.BatchNorm2d,
-    torch.nn.ConstantPad2d,
-    torch.nn.ELU,
-    torch.nn.Hardtanh,
-    torch.nn.ReLU,
-    torch.nn.functional.relu,
-    torch.nn.functional.relu_,
-    torch.nn.functional.leaky_relu,
-    torch.nn.functional.leaky_relu_,
-    torch.nn.LeakyReLU,
 ]
 
 SUPPORTED_IMPLICIT_Q_DQ_MODULES_SET = set(SUPPORTED_QUANT_MODULES)

backends/xnnpack/test/test_xnnpack_quantized.py

@@ -199,8 +199,6 @@ def forward(self, x):
 
         self.quantize_and_test_model(LeakyReLUModule(), example_inputs)
 
-    # TODO(T158652796)
-    @unittest.expectedFailure
     def test_xnnpack_leaky_relu2(self):
         example_inputs = (torch.randn(1, 3, 3),)