reset HEAD back, wip

Author: winskuo-quic

backends/qualcomm/quantizer/annotators.py

@@ -432,6 +432,10 @@ def annotate_ceil(node: Node, quantization_config: QuantizationConfig) -> None:
 def annotate_clamp(node: Node, quantization_config: QuantizationConfig) -> None:
     annotate_single_in_single_out(node, quantization_config)
 
+@register_annotator([torch.ops.aten.index_select.default])
+def annotate_index_select(node: Node, quantization_config: QuantizationConfig) -> None:
+    import pdb; pdb.set_trace()
+    annotate_single_in_single_out(node, quantization_config)
 
 @register_annotator([torch.ops.aten.floor.default])
 def annotate_floor(node: Node, quantization_config: QuantizationConfig) -> None:

backends/qualcomm/tests/models.py

@@ -813,7 +813,23 @@ def __init__(self):
     def forward(self, x):
         return torch.special.expm1(x)
 
+class Flip(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.dims = [0,2]
+
+    def forward(self, x):
+        return torch.flip(x, self.dims)
 
+class FlipDecomp(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.dims = [0,2]
+    def forward(self, x):
+        for dim in self.dims:
+            idx = torch.arange(x.size(dim) - 1, -1, -1, device=x.device)
+            x = torch.index_select(x, dim, idx)
+        return x
 class Floor(torch.nn.Module):
     def __init__(self):
         super().__init__()

backends/qualcomm/tests/test_qnn_delegate.py

@@ -2030,6 +2030,18 @@ def test_qnn_backend_expm1(self):
         module = self.get_qdq_module(module, sample_input)
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_flip(self):
+        sample_input = (torch.randn(3, 4, 5,6),)
+        # golden_module = Flip()
+        decomp_module = FlipDecomp()
+        decomp_module = self.get_qdq_module(decomp_module, sample_input)
+        self.lower_module_and_test_output(decomp_module, sample_input)
+        # golden_out = golden_module(sample_input)
+        # decomp_out = decomp_module(sample_input)
+        # torch.testing.assert_close(golden_out, decomp_out)
+        
+        
+    
     def test_qnn_backend_floor(self):
         sample_input = (torch.randn(3, 4),)
         module = Floor()  # noqa: F405