Add a test case for AOT int32 overflow checks

Author: voltjia

tests/test_aot.py

@@ -5,6 +5,7 @@
 import torch.nn.functional as F
 
 import ninetoothed
+import ninetoothed.aot
 import ninetoothed.generation
 import tests.test_addmm as addmm
 import tests.test_attention as attention
@@ -342,6 +343,19 @@ def _application(input, scale, output):
     assert torch.allclose(output, expected)
 
 
+def test_overflow_terms():
+    terms = ninetoothed.aot._overflow_terms(("input", "scale"), (2, 0))
+
+    assert terms == (
+        "input.shape[0] > 2147483647ULL",
+        "input.strides[0] > 2147483647LL",
+        "input.strides[0] < -2147483648LL",
+        "input.shape[1] > 2147483647ULL",
+        "input.strides[1] > 2147483647LL",
+        "input.strides[1] < -2147483648LL",
+    )
+
+
 def _generate_kernel_name_suffix():
     count = _generate_kernel_name_suffix._kernel_count
     _generate_kernel_name_suffix._kernel_count += 1