Fix aten.amax lowering issue (#13381)

Summary:

There was an error when lowering amax around this line `input_tensor = self.get_tensor(input_node, node)` and the issue is that we're trying to permute the tensor inside node_visitors, op_node.meta[QCOM_AXIS_ORDER] is (0, 1), however, tensor.shape is (1, 980, 49).

Rollback Plan:

Differential Revision: D80187368

Author: cccclai

backends/qualcomm/_passes/layout_transform.py

@@ -175,6 +175,7 @@ def is_layout_agnostic(self, node: torch.fx.Node) -> bool:
             exir_ops.edge.aten.mean.dim,
             exir_ops.edge.aten.min.dim,
             exir_ops.edge.aten.sum.dim_IntList,
+            exir_ops.edge.aten.amax.default,
         }:
             # if dimemsion is not kept, we'll have no clue how to do layout transform
             if len(node.args) < 3 or not node.args[2]:

backends/qualcomm/tests/models.py

@@ -102,6 +102,24 @@ def forward(self, x):
         return torch.amax(x, dim=self.dim, keepdim=self.keepdim)
 
 
+class AMaxFollowingConv2D(torch.nn.Module):
+    def __init__(
+        self, in_channels, out_channels, kernel_size=3, dim=None, keepdim=False
+    ):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(
+            in_channels, out_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.dim = dim
+        self.keepdim = keepdim
+
+    def forward(self, x):
+        x = self.conv(
+            x
+        )  # Apply convolution (output shape: [batch, out_channels, H, W])
+        return torch.amax(x, dim=self.dim, keepdim=self.keepdim)
+
+
 class AMin(torch.nn.Module):
     def __init__(self, dim=None, keepdim=False):
         super().__init__()

backends/qualcomm/tests/test_qnn_delegate.py

@@ -134,6 +134,13 @@ def test_qnn_backend_amax(self):
             with self.subTest(i=i):
                 self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_amax_conv(self):
+        sample_input = (torch.randn(2, 3, 64, 64),)  # [batch, channels, height, width]
+        module = AMaxFollowingConv2D(  # noqa: F405
+            in_channels=3, out_channels=16, kernel_size=3, dim=-1, keepdim=False
+        )
+        self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_amin(self):
         modules = [AMin(dim=1, keepdim=False), AMin(dim=1, keepdim=True)]  # noqa: F405
         sample_input = (torch.randn(4, 4),)
@@ -1435,6 +1442,13 @@ def test_qnn_backend_amax(self):
                 module = self.get_qdq_module(module, sample_input)
                 self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_amax_conv(self):
+        sample_input = (torch.randn(2, 3, 64, 64),)  # [batch, channels, height, width]
+        module = AMaxFollowingConv2D(  # noqa: F405
+            in_channels=3, out_channels=16, kernel_size=3, dim=-1, keepdim=False
+        )
+        self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_amin(self):
         modules = [AMin(dim=1, keepdim=False), AMin(dim=1, keepdim=True)]  # noqa: F405
         sample_input = (torch.randn(4, 4),)
@@ -3418,7 +3432,6 @@ def test_qnn_backend_generate_optrace(self):
 
         for compiler_spec in compiler_specs:
             with tempfile.TemporaryDirectory() as tmp_dir:
-
                 edge_prog_mgr = to_edge_transform_and_lower_to_qnn(
                     module, sample_input, compiler_spec
                 ).to_executorch()