Qualcomm AI Engine Direct - Adapt to new IR capture flow (pytorch#2431)

Summary:
- Change existent IR capture flow (exir.capture) to torch.export.export
- Add custom decomposition table for mitigating maintaining effort
- Fix breakages encountered and make sure all tests passed as well

Pull Request resolved: pytorch#2431

Reviewed By: mergennachin

Differential Revision: D55353449

Pulled By: cccclai

fbshipit-source-id: aa2e27d0ae93aa62208fd03ec39b3891a70b954e

backends/qualcomm/builders/node_visitor.py

@@ -396,7 +396,8 @@ def register_node_visitor(visitor):
         and issubclass(visitor, NodeVisitor)
         and hasattr(visitor, "target")
     ), f"Illformed NodeVisitor subclass, can't register!, got: {visitor}"
-    _node_visitor_dict[visitor.target] = visitor
+    for target in visitor.target:
+        _node_visitor_dict[target] = visitor
 
 
 def generate_node_to_external_map(

backends/qualcomm/builders/op_add.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class Add(NodeVisitor):
-    target = "aten.add.Tensor"
+    target = ["aten.add.Tensor"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_avg_pool2d.py

@@ -16,7 +16,7 @@
 
 @register_node_visitor
 class AvgPool2d(NodeVisitor):
-    target = "aten.avg_pool2d.default"
+    target = ["aten.avg_pool2d.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_batch_norm.py

@@ -16,7 +16,7 @@
 
 @register_node_visitor
 class BatchNorm(NodeVisitor):
-    target = "aten._native_batch_norm_legit_no_training.default"
+    target = ["aten._native_batch_norm_legit_no_training.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_bmm.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class BMM(NodeVisitor):
-    target = "aten.bmm.default"
+    target = ["aten.bmm.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_cast.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class Cast(NodeVisitor):
-    target = "aten._to_copy.default"
+    target = ["aten._to_copy.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_cat.py

@@ -16,7 +16,7 @@
 
 @register_node_visitor
 class Cat(NodeVisitor):
-    target = "aten.cat.default"
+    target = ["aten.cat.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_ceil.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class Ceil(NodeVisitor):
-    target = "aten.ceil.default"
+    target = ["aten.ceil.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_clamp.py

@@ -16,7 +16,7 @@
 
 @register_node_visitor
 class Clamp(NodeVisitor):
-    target = "aten.clamp.default"
+    target = ["aten.clamp.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_conv2d.py

@@ -24,7 +24,7 @@
 
 @register_node_visitor
 class Conv2d(NodeVisitor):
-    target = "aten.convolution.default"
+    target = ["aten.convolution.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_depth_to_space.py

@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class DepthToSpaceVisitor(NodeVisitor):
-    target = "aten.pixel_shuffle.default"
+    target = ["aten.pixel_shuffle.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_dequantize.py

@@ -55,19 +55,19 @@ def define_node(
 
 @register_node_visitor
 class PerTensorDequantizeDefault(DequantizeOpBase):
-    target = "quantized_decomposed.dequantize_per_tensor.default"
+    target = ["quantized_decomposed.dequantize_per_tensor.default"]
 
 
 @register_node_visitor
 class PerTensorDequantizeTensor(DequantizeOpBase):
-    target = "quantized_decomposed.dequantize_per_tensor.tensor"
+    target = ["quantized_decomposed.dequantize_per_tensor.tensor"]
 
 
 @register_node_visitor
 class PerChannelDequantizeDefault(DequantizeOpBase):
-    target = "quantized_decomposed.dequantize_per_channel.default"
+    target = ["quantized_decomposed.dequantize_per_channel.default"]
 
 
 @register_node_visitor
 class PerChannelDequantizeTensor(DequantizeOpBase):
-    target = "quantized_decomposed.dequantize_per_channel.tensor"
+    target = ["quantized_decomposed.dequantize_per_channel.tensor"]

backends/qualcomm/builders/op_div.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class Div(NodeVisitor):
-    target = "aten.div.Tensor"
+    target = ["aten.div.Tensor"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_embedding.py

@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class Embedding(NodeVisitor):
-    target = "aten.embedding.default"
+    target = ["aten.embedding.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_expand.py

@@ -16,7 +16,7 @@
 
 @register_node_visitor
 class Expand(NodeVisitor):
-    target = "aten.expand_copy.default"
+    target = ["aten.expand_copy.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_gelu.py

@@ -16,7 +16,7 @@
 
 @register_node_visitor
 class GeluVisitor(NodeVisitor):
-    target = "aten.gelu.default"
+    target = ["aten.gelu.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_hardswish.py

@@ -16,7 +16,7 @@
 
 @register_node_visitor
 class HardSwishVisitor(NodeVisitor):
-    target = "aten.hardswish.default"
+    target = ["aten.hardswish.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_hardtanh.py

@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class HardTanhVisitor(NodeVisitor):
-    target = "aten.hardtanh.default"
+    target = ["aten.hardtanh.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_layer_norm.py

@@ -18,7 +18,7 @@
 
 @register_node_visitor
 class LayerNormVisitor(NodeVisitor):
-    target = "aten.native_layer_norm.default"
+    target = ["aten.native_layer_norm.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_linear.py

@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class LinearVisitor(NodeVisitor):
-    target = "aten.linear.default"
+    target = ["aten.linear.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_log_softmax.py

@@ -16,7 +16,7 @@
 
 @register_node_visitor
 class LogSoftmax(NodeVisitor):
-    target = "aten._log_softmax.default"
+    target = ["aten._log_softmax.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_matmul.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class Matmul(NodeVisitor):
-    target = "aten.matmul.default"
+    target = ["aten.matmul.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_max_pool2d.py

@@ -16,7 +16,7 @@
 
 @register_node_visitor
 class MaxPool2d(NodeVisitor):
-    target = "aten.max_pool2d_with_indices.default"
+    target = ["aten.max_pool2d_with_indices.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_mean_dim.py

@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class MeanDim(NodeVisitor):
-    target = "aten.mean.dim"
+    target = ["aten.mean.dim"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_mul.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class Mul(NodeVisitor):
-    target = "aten.mul.Tensor"
+    target = ["aten.mul.Tensor"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_pad.py

@@ -16,7 +16,7 @@
 
 @register_node_visitor
 class Pad(NodeVisitor):
-    target = "aten.constant_pad_nd.default"
+    target = ["aten.constant_pad_nd.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_pow.py

@@ -17,7 +17,7 @@
 # TODO Add more class Like PowTensorTensor if needed
 @register_node_visitor
 class PowTensorScalar(NodeVisitor):
-    target = "aten.pow.Tensor_Scalar"
+    target = ["aten.pow.Tensor_Scalar"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_quantize.py

@@ -61,9 +61,9 @@ def define_node(
 
 @register_node_visitor
 class PerTensorQuantize(QuantizeOpBase):
-    target = "quantized_decomposed.quantize_per_tensor.default"
+    target = ["quantized_decomposed.quantize_per_tensor.default"]
 
 
 @register_node_visitor
 class PerChannelQuantize(QuantizeOpBase):
-    target = "quantized_decomposed.quantize_per_channel.default"
+    target = ["quantized_decomposed.quantize_per_channel.default"]

backends/qualcomm/builders/op_relu.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class Relu(NodeVisitor):
-    target = "aten.relu.default"
+    target = ["aten.relu.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_reshape.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class Reshape(NodeVisitor):
-    target = "aten.view_copy.default"
+    target = ["aten.view_copy.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_rsqrt.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class Rsqrt(NodeVisitor):
-    target = "aten.rsqrt.default"
+    target = ["aten.rsqrt.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_select_copy.py

@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class SelectCopy(NodeVisitor):
-    target = "aten.select_copy.int"
+    target = ["aten.select_copy.int", "aten.select.int"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_sigmoid.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class Sigmoid(NodeVisitor):
-    target = "aten.sigmoid.default"
+    target = ["aten.sigmoid.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_skip_ops.py

@@ -35,7 +35,7 @@ class OpGetItem(OpSkipOps):
     do nothing if node is getitem
     """
 
-    target = "getitem"
+    target = ["getitem"]
 
     def define_node(
         self,

backends/qualcomm/builders/op_slice_copy.py

@@ -16,7 +16,7 @@
 
 @register_node_visitor
 class StrideSlice(NodeVisitor):
-    target = "aten.slice_copy.Tensor"
+    target = ["aten.slice_copy.Tensor"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_softmax.py

@@ -16,7 +16,7 @@
 
 @register_node_visitor
 class Softmax(NodeVisitor):
-    target = "aten._softmax.default"
+    target = ["aten._softmax.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_squeeze.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class Squeeze(NodeVisitor):
-    target = "aten.squeeze_copy.dims"
+    target = ["aten.squeeze_copy.dims", "aten.squeeze.dims"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_sub.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class Sub(NodeVisitor):
-    target = "aten.sub.Tensor"
+    target = ["aten.sub.Tensor"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_tanh.py

@@ -16,7 +16,7 @@
 
 @register_node_visitor
 class Tanh(NodeVisitor):
-    target = "aten.tanh.default"
+    target = ["aten.tanh.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_transpose.py

@@ -17,7 +17,7 @@
 
 @register_node_visitor
 class TransposeVisitor(NodeVisitor):
-    target = "aten.permute_copy.default"
+    target = ["aten.permute_copy.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_unsqueeze.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class Unsqueeze(NodeVisitor):
-    target = "aten.unsqueeze_copy.default"
+    target = ["aten.unsqueeze_copy.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_upsample_bilinear2d.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class ResizeBilinear(NodeVisitor):
-    target = "aten.upsample_bilinear2d.default"
+    target = ["aten.upsample_bilinear2d.default"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)