Arm backend: Assert delegation in the int64 tests

Author: abm-77

backends/arm/test/misc/test_int64.py

@@ -128,25 +128,30 @@ def forward(self, x: torch.Tensor, y: torch.Tensor):
 
 
 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 > 0
+        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(), (torch.arange(-5, 5), torch.arange(0, 10))),
-    "gt_scalar": (
-        Int64GreaterScalar(),
-        (torch.arange(-5, 5), torch.arange(0, 10)),
-    ),
-    "sub": (Int64Sub(), (torch.arange(-5, 5), torch.arange(0, 10))),
-    "mul": (Int64Mul(), (torch.arange(-5, 5), torch.arange(0, 10))),
+    "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):
-    """Lowering keeps int64 operands when the spec declares the extension."""
-    model, inputs = test_data
+    """Int64 operands are delegated when the spec declares the extension."""
+    model, inputs, aten_op = test_data
     (
         ArmTester(
             model,
@@ -155,5 +160,7 @@ def test_int64_extension_tosa_INT(test_data: Tuple):
         )
         .export()
         .to_edge_transform_and_lower()
+        .check_count({"torch.ops.higher_order.executorch_call_delegate": 1})
+        .check_not([aten_op])
         .to_executorch()
     )