Specialize AOT sizes and strides as int32 (#155)

* Specialize AOT sizes and strides as `int32`

* Add a test case for AOT `int32` overflow checks

Author: voltjia

src/ninetoothed/aot.py

@@ -81,10 +81,19 @@ def _find_tensor_by_source_name(tensors, name):
     variant_specs = _enumerate_variant_specs(
         launch_arg_names, tensors, _find_tensor_by_source_name
     )
+    _, tensor_ndims, _ = _per_tensor_dim_options(
+        launch_arg_names, tensors, _find_tensor_by_source_name
+    )
 
     output_contents = {}
 
-    for variant_suffix, divisibility_spec, contiguity_spec in variant_specs:
+    for (
+        variant_suffix,
+        divisibility_spec,
+        contiguity_spec,
+        size_type,
+        stride_type,
+    ) in variant_specs:
         variant_outputs = _build_variant(
             source_file,
             kernel_func,
@@ -99,11 +108,13 @@ def _find_tensor_by_source_name(tensors, name):
             num_stages=num_stages,
             divisibility_spec=divisibility_spec,
             contiguity_spec=contiguity_spec,
+            size_type=size_type,
+            stride_type=stride_type,
         )
         output_contents.update(variant_outputs)
 
     dispatcher_source, dispatcher_header = _generate_dispatcher(
-        kernel_name, launch_arg_names, variant_specs
+        kernel_name, launch_arg_names, variant_specs, tensor_ndims
     )
 
     output_contents[f"{kernel_name}.cpp"] = dispatcher_source
@@ -112,7 +123,7 @@ def _find_tensor_by_source_name(tensors, name):
     return output_contents
 
 
-def _generate_dispatcher(kernel_name, launch_arg_names, variant_specs):
+def _generate_dispatcher(kernel_name, launch_arg_names, variant_specs, tensor_ndims):
     tensor_params = ", ".join(f"NineToothedTensor {name}" for name in launch_arg_names)
     signature_params = (
         f"NineToothedStream stream, {tensor_params}"
@@ -138,14 +149,30 @@ def _generate_dispatcher(kernel_name, launch_arg_names, variant_specs):
     externs = []
     branches = []
 
-    for variant_suffix, divisibility_spec, contiguity_spec in variant_specs:
+    fallback_call = None
+
+    for (
+        variant_suffix,
+        divisibility_spec,
+        contiguity_spec,
+        size_type,
+        stride_type,
+    ) in variant_specs:
         variant_name = f"launch_{kernel_name}_{variant_suffix}"
         externs.append(
             f'extern "C" NineToothedResult {variant_name}({signature_params});'
         )
 
         call = f"return {variant_name}({call_args});"
 
+        if (size_type, stride_type) == (
+            ninetoothed.dtype.int64,
+            ninetoothed.dtype.int64,
+        ):
+            fallback_call = call
+
+            continue
+
         checks = tuple(
             f"{name}.shape[{dim}] % 16 == 0" for name, dim in divisibility_spec
         ) + tuple(f"{name}.strides[{dim}] == 1" for name, dim in contiguity_spec)
@@ -155,11 +182,23 @@ def _generate_dispatcher(kernel_name, launch_arg_names, variant_specs):
         else:
             branches.append(f"{_INDENTATION}{call}")
 
+    prelude_lines = []
+
+    if fallback_call is not None and launch_arg_names:
+        overflow_terms = _overflow_terms(launch_arg_names, tensor_ndims)
+
+        if overflow_terms:
+            prelude_lines.append(
+                f"{_INDENTATION}if ({' || '.join(overflow_terms)}) {fallback_call}"
+            )
+
+    body_lines = prelude_lines + branches
+
     source = (
         f'#include "{_HEADER_PATH}"\n\n'
         + "\n".join(externs)
         + f'\n\nextern "C" {signature} {{\n'
-        + "\n".join(branches)
+        + "\n".join(body_lines)
         + "\n}\n"
     )
 
@@ -181,6 +220,8 @@ def _build_variant(
     num_stages,
     divisibility_spec,
     contiguity_spec,
+    size_type=ninetoothed.dtype.int32,
+    stride_type=ninetoothed.dtype.int32,
 ):
     divisibility_set = {
         (naming.remove_prefixes(name), dim) for name, dim in divisibility_spec
@@ -211,9 +252,9 @@ def _build_variant(
             bare_source_name = naming.remove_prefixes(source_name)
 
             if (bare_source_name, dim_index) in divisibility_set:
-                param_types.append(f"{ninetoothed.dtype.int64}:16")
+                param_types.append(f"{size_type}:16")
             else:
-                param_types.append(ninetoothed.dtype.int64)
+                param_types.append(size_type)
         elif match := Tensor.stride_pattern().fullmatch(param):
             source_name = match.group(1)
             dim_index = int(match.group(3))
@@ -224,7 +265,7 @@ def _build_variant(
                 constexpr_param_indices.append(len(param_types) - 1)
                 constexpr_strides.append((source_name, dim_index))
             else:
-                param_types.append(f"{ninetoothed.dtype.int64}:16")
+                param_types.append(f"{stride_type}:16")
         else:
             source_name = param
             tensor = find_tensor(tensors, source_name)
@@ -331,15 +372,28 @@ def _spec_from_combo(combo):
     for divisibility_spec in dim_specs:
         for contiguity_spec in dim_specs:
             suffix = _variant_suffix(
-                divisibility_spec, contiguity_spec, launch_arg_names, tensor_ndims
+                divisibility_spec,
+                contiguity_spec,
+                launch_arg_names,
+                tensor_ndims,
+                size_type=ninetoothed.dtype.int32,
+                stride_type=ninetoothed.dtype.int32,
+            )
+            specs.append(
+                (
+                    suffix,
+                    divisibility_spec,
+                    contiguity_spec,
+                    ninetoothed.dtype.int32,
+                    ninetoothed.dtype.int32,
+                )
             )
-            specs.append((suffix, divisibility_spec, contiguity_spec))
 
     def _num_innermost(spec):
         return sum(1 for name, dim in spec if innermost_dims.get(name) == dim)
 
     def _specificity(entry):
-        _, divisibility_spec, contiguity_spec = entry
+        _, divisibility_spec, contiguity_spec, _, _ = entry
 
         return (
             -len(divisibility_spec),
@@ -350,6 +404,24 @@ def _specificity(entry):
 
     specs.sort(key=_specificity)
 
+    fallback_suffix = _variant_suffix(
+        (),
+        (),
+        launch_arg_names,
+        tensor_ndims,
+        size_type=ninetoothed.dtype.int64,
+        stride_type=ninetoothed.dtype.int64,
+    )
+    specs.append(
+        (
+            fallback_suffix,
+            (),
+            (),
+            ninetoothed.dtype.int64,
+            ninetoothed.dtype.int64,
+        )
+    )
+
     return specs
 
 
@@ -381,15 +453,40 @@ def _per_tensor_dim_options(launch_arg_names, tensors, find_tensor):
     return per_tensor_dims, tensor_ndims, innermost_dims
 
 
-def _variant_suffix(divisibility_spec, contiguity_spec, launch_arg_names, tensor_ndims):
+def _overflow_terms(launch_arg_names, tensor_ndims):
+    int32_min = -(2**31)
+    int32_max = 2**31 - 1
+
+    return tuple(
+        term
+        for name, ndim in zip(launch_arg_names, tensor_ndims)
+        for dim in range(ndim)
+        for term in (
+            f"{name}.shape[{dim}] > {int32_max}ULL",
+            f"{name}.strides[{dim}] > {int32_max}LL",
+            f"{name}.strides[{dim}] < {int32_min}LL",
+        )
+    )
+
+
+def _variant_suffix(
+    divisibility_spec,
+    contiguity_spec,
+    launch_arg_names,
+    tensor_ndims,
+    size_type=ninetoothed.dtype.int32,
+    stride_type=ninetoothed.dtype.int32,
+):
     divisibility_part = _divisibility_suffix(
         divisibility_spec, launch_arg_names, tensor_ndims
     )
     contiguity_part = _contiguity_suffix(
         contiguity_spec, launch_arg_names, tensor_ndims
     )
 
-    return f"{divisibility_part}_{contiguity_part}"
+    return (
+        f"{divisibility_part}_{contiguity_part}_size_{size_type}_stride_{stride_type}"
+    )
 
 
 def _divisibility_suffix(divisibility_spec, launch_arg_names, tensor_ndims):

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