Arm backend: Test int64 under the int64 extension

Author: abm-77

backends/arm/test/misc/test_int64.py

@@ -110,3 +110,57 @@ def test_int64_tosa_INT(test_data: Tuple):
         .to_executorch()
         .run_method_and_compare_outputs(inputs)
     )
+
+
+class Int64Add(torch.nn.Module):
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        return x + y
+
+
+class Int64Sub(torch.nn.Module):
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        return x - y
+
+
+class Int64Mul(torch.nn.Module):
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        return x * y
+
+
+class Int64GreaterScalar(torch.nn.Module):
+    """Compare against a scalar that a double cannot represent exactly."""
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        return x > 9007199254740993
+
+
+# Operands straddle the int32 range, so a narrowed lowering cannot reproduce
+# the reference result.
+_WIDE = torch.tensor([2**40, -(2**38), 2**31, -1], dtype=torch.int64)
+_WIDE_2 = torch.tensor([2**39, 2**37, 1, -(2**40)], dtype=torch.int64)
+_NEAR_2_53 = torch.tensor([2**53, 2**53 + 1, 2**53 + 2, -(2**53)], dtype=torch.int64)
+
+int64_extension_suite = {
+    "add": (Int64Add(), (_WIDE, _WIDE_2), "aten_add_tensor"),
+    "sub": (Int64Sub(), (_WIDE, _WIDE_2), "aten_sub_tensor"),
+    "mul": (Int64Mul(), (_WIDE, _WIDE_2), "aten_mul_tensor"),
+    "gt_scalar": (Int64GreaterScalar(), (_NEAR_2_53, _WIDE), "aten_gt_scalar"),
+}
+
+
+@common.parametrize("test_data", int64_extension_suite)
+def test_int64_extension_tosa_INT(test_data: Tuple):
+    """Int64 operands are delegated when the spec declares the extension."""
+    model, inputs, aten_op = test_data
+    (
+        ArmTester(
+            model,
+            inputs,
+            common.get_tosa_compile_spec("TOSA-1.1+INT+int64"),
+        )
+        .export()
+        .to_edge_transform_and_lower()
+        .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+        .check_not([aten_op])
+        .to_executorch()
+    )