Move calling convention validation to KernelSignature constructor.

So that it is performed regardless of name mangling.

Signed-off-by: Greg Bonik <gbonik@nvidia.com>

Author: gbonik

src/cuda/tile/compilation/_name_mangling.py

@@ -16,20 +16,6 @@
 from .._cext import CallingConvention
 
 
-def cconv_require_tuple_constraint(cconv: CallingConvention, cursor: "_Cursor | None" = None):
-    if cconv.version < 2:
-        msg = (f"Tuple parameters ('T' constraint) are not supported by calling convention"
-               f" {cconv.name}; version >= 2 is required")
-        raise cursor.make_error(msg) if cursor is not None else ValueError(msg)
-
-
-def cconv_require_static_shape(cconv: CallingConvention, cursor: "_Cursor | None" = None):
-    if cconv.version < 2:
-        msg = (f"Static array shapes ('s' predicate) are not supported by calling convention"
-               f" {cconv.name}; version >= 2 is required")
-        raise cursor.make_error(msg) if cursor is not None else ValueError(msg)
-
-
 def mangle_kernel_name(function_name: str,
                        kernel_signature: KernelSignature) -> str:
     alias_group_map, alias_group_names = _map_alias_groups(kernel_signature.parameters)
@@ -158,7 +144,6 @@ def _mangle_constraint(p: ParameterConstraint, alias_group_map: dict[str, int],
         assert isinstance(p.element, ArrayConstraint)
         return "L" + _mangle_list_constraint(p, alias_group_map, cconv)
     elif isinstance(p, TupleConstraint):
-        cconv_require_tuple_constraint(cconv)
         return "T" + _mangle_tuple_constraint(p, alias_group_map, cconv)
     elif isinstance(p, ScalarConstraint):
         return "S" + _mangle_dtype(p.dtype)
@@ -186,7 +171,6 @@ def _demangle_constraint(cursor: _Cursor,
     elif c == "L":
         return _demangle_list_constraint(cursor, alias_group_demangler, cconv)
     elif c == "T":
-        cconv_require_tuple_constraint(cconv, orig_cursor)
         return _demangle_tuple_constraint(cursor, alias_group_demangler, cconv)
     elif c == "S":
         dtype = _demangle_dtype(cursor)
@@ -206,8 +190,6 @@ def _demangle_constraint(cursor: _Cursor,
 def _mangle_array_constraint(a: ArrayConstraint,
                              alias_group_map: dict[str, int],
                              cconv: CallingConvention) -> str:
-    if any(v is not None for v in a.shape_constant):
-        cconv_require_static_shape(cconv)
     ret = f"{a.ndim}{_mangle_dtype(a.dtype)}"
 
     # NOTE: since we encode axis masks as hex, letters a-f can't be used for predicates
@@ -273,10 +255,8 @@ def _demangle_array_constraint(cursor: _Cursor,
             raise mask_cursor.make_error(f"Axis mask {axis_mask:x} has more bits"
                                          f" ({axis_mask.bit_length()}) than array ndim ({ndim})")
 
-        s_cursor = cursor.clone()
         axis_shape_constant = None
         if cursor.read("s") is not None:
-            cconv_require_static_shape(cconv, s_cursor)
             axis_shape_constant = _demangle_signed_int(cursor)
 
         axis_shape_div_by = 1

src/cuda/tile/compilation/_signature.py

@@ -330,6 +330,8 @@ def __init__(self,
 
         parameters = tuple(_to_constraint(c) for c in parameters)
         _validate_alias_groups(parameters)
+        for x in parameters:
+            _validate_constraint_support(x, calling_convention)
 
         object.__setattr__(self, "parameters", parameters)
         object.__setattr__(self, "calling_convention", calling_convention)
@@ -504,3 +506,24 @@ def _remove_redundant_divisibility_constraints(static_values: tuple[int, ...],
                                  f" which is not divisible by {ret[i]}")
             ret[i] = 1
     return tuple(ret)
+
+
+def _validate_constraint_support(constraint: ParameterConstraint, cconv: CallingConvention):
+    if isinstance(constraint, ScalarConstraint):
+        pass
+    elif isinstance(constraint, ArrayConstraint):
+        if any(x is not None for x in constraint.shape_constant) and cconv.version < 2:
+            raise ValueError(f"Static array shapes are not supported by calling convention"
+                             f" {cconv.name}; version >= 2 is required")
+    elif isinstance(constraint, ListConstraint):
+        _validate_constraint_support(constraint.element, cconv)
+    elif isinstance(constraint, TupleConstraint):
+        if cconv.version < 2:
+            raise ValueError(f"Tuple parameters are not supported by calling convention"
+                             f" {cconv.name}; version >= 2 is required")
+        for x in constraint.items:
+            _validate_constraint_support(x, cconv)
+    elif isinstance(constraint, ConstantConstraint):
+        pass
+    else:
+        raise TypeError(f"Unexpected constraint type: {type(constraint)}")

test/test_export_compat.py

@@ -3,10 +3,12 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import ctypes
+import re
 from ctypes import (c_char_p, POINTER, c_void_p, c_int, c_uint64, pointer, CFUNCTYPE, c_uint,
                     c_int32, c_float, byref)
 from io import BytesIO
 
+import pytest
 import torch.cuda
 
 import cuda.tile as ct
@@ -201,3 +203,24 @@ def test_export_compat_cutile_python_v2():
     out = torch.zeros((), dtype=torch.int32, device="cuda")
     _call_kernel(io.getvalue(), "kernel_2_Kt2_T2Si32Si32_T1I10_A0i32", ((3, 7), out), is_v2=True)
     assert out.item() == 20
+
+
+def test_static_shape_with_v1_raises():
+    cconv = ct.compilation.CallingConvention.cutile_python_v1()
+    expected_message = re.escape("Static array shapes are not supported by calling convention"
+                                 " cutile_python_v1; version >= 2 is required")
+    with pytest.raises(ValueError, match=expected_message):
+        ct.compilation.KernelSignature(
+            [ct.compilation.ArrayConstraint(
+                ct.float32, 1, index_dtype=ct.int32, stride_lower_bound_incl=0,
+                alias_groups=(), may_alias_internally=False,
+                shape_constant=(8,))],
+            cconv)
+
+
+def test_tuple_with_v1_raises():
+    cconv = ct.compilation.CallingConvention.cutile_python_v1()
+    expected_message = re.escape("Tuple parameters are not supported by calling convention"
+                                 " cutile_python_v1; version >= 2 is required")
+    with pytest.raises(ValueError, match=expected_message):
+        ct.compilation.KernelSignature([(ct.compilation.ScalarConstraint(ct.int32),)], cconv)

test/test_name_mangling.py

@@ -269,30 +269,12 @@ def test_name_mangling_cutile_python_v2(parameters, expected_suffix):
     assert demangled_name == func_name
 
 
-def test_mangle_tuple_with_v1_raises():
-    cconv = CallingConvention.cutile_python_v1()
-    sig = KernelSignature([TupleConstraint([ScalarConstraint(int32)])], cconv)
-    with pytest.raises(ValueError, match="version >= 2"):
-        mangle_kernel_name("my_kernel", sig)
-
-
 def test_demangle_tuple_with_v1_raises():
     symbol = "my_kernel_Kt1_T1Si32"
     with pytest.raises(ValueError, match="version >= 2"):
         demangle_kernel_name(symbol)
 
 
-def test_mangle_static_shape_with_v1_raises():
-    cconv = CallingConvention.cutile_python_v1()
-    sig = KernelSignature(
-        [ArrayConstraint(float32, 1, index_dtype=int32, stride_lower_bound_incl=0,
-                         alias_groups=(), may_alias_internally=False,
-                         shape_constant=(8,))],
-        cconv)
-    with pytest.raises(ValueError, match="version >= 2"):
-        mangle_kernel_name("my_kernel", sig)
-
-
 def test_demangle_static_shape_with_v1_raises():
     symbol = "my_kernel_Kt1_A1f32_1s8l0"
     with pytest.raises(ValueError, match="version >= 2"):