[eudsl-python-extras] drop support for generics for py < 3.12 (#236)

After much consternation and flaggelation, I've decided we
can't/shouldn't support generics for python < 3.12. The context/reason
is that I thought I had robust support for just plain `TypeVar`s as
globals (and in the closure) but recently I've realized that that
support is not robust.

In the weeds: the "generic" kernel instantiation effectively mutates the
`TypeVar` and thus when the `TypeVar` is a global, multiple
instantiations of the same kernel (with different concrete params) will
race (on that global `TypeVar`). `TypeVar`s in the closure don't suffer
from this because we copy the `body_builder` and its closure when
instantiating the kernel. In 3.12 (when using the brackets `def fun[M,
N, ...]` syntax) the `TypeVar`s are always in the closure[^1]. I tried
to figure out somehow to move globals into closures but failed and
decided it's not worth the effort anyway.

Note, the added test (`test_generics`) is gated behind 3.13 because
generics default value assignment is a 3.13 feature and the 3.12 parser
will choke on it, but everything else works in 3.12.

[^1]: That's actually a CPython implementation detail - there's a
compiler transform which rewrites the function into a closed over
function with those TypeVars in its closure.

Author: makslevental

.github/workflows/build_test_release_eudsl_python_extras.yml

@@ -64,6 +64,7 @@ jobs:
 
       - name: "Build eudsl-python-extras sdist"
         run: |
+          
           SHA_SHORT="$(git rev-parse --short HEAD)"
           WHEEL_VERSION="$(date +'%Y%m%d.%H%M')+$SHA_SHORT"
           pushd projects/eudsl-python-extras 
@@ -175,7 +176,14 @@ jobs:
         run: python -m pip install dist/eudsl_python_extras*.tar.gz
 
       - name: "Test eudsl-python-extras"
-        run: python -m pytest projects/eudsl-python-extras/tests
+        run: |
+          
+          IGNORE=""
+          if [[ $(python -c "print(__import__('sys').version_info < (3, 13))") == "True" ]]; then
+            IGNORE="--ignore projects/eudsl-python-extras/tests/dialect/test_generics.py"
+          fi
+          
+          python -m pytest projects/eudsl-python-extras/tests $IGNORE
 
   release-eudsl-python-extras:
 

projects/eudsl-python-extras/mlir/extras/ast/util.py

@@ -12,6 +12,7 @@
 
 from bytecode import ConcreteBytecode
 from cloudpickle import cloudpickle
+
 from ...ir import Type
 
 

projects/eudsl-python-extras/mlir/extras/dialects/func.py

@@ -1,17 +1,16 @@
 # Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-import copy
 import inspect
 import sys
 from dataclasses import dataclass
 from functools import update_wrapper
 from typing import Optional, List, Union, TypeVar
 
-from ..ast.util import copy_func, copy_object
+from .. import types as T
 from ..ast.py_type import PyTypeVarObject, _Ptr, PyObject
+from ..ast.util import copy_func
 from ..meta import op_region_builder
-from .. import types as T
 from ..util import get_user_code_loc, make_maybe_no_args_decorator
 from ...dialects._ods_common import get_op_result_or_op_results
 from ...dialects.func import *
@@ -27,7 +26,6 @@
     Value,
 )
 
-
 _call = call
 
 
@@ -175,7 +173,7 @@ def __init__(
         self.call_op_ctor = call_op_ctor
         self.arg_attrs = arg_attrs
         self.res_attrs = res_attrs
-        self.generics = copy_object(generics)
+        self.generics = generics
         self.loc = loc
         self.ip = ip
         self._func_op = None
@@ -323,33 +321,24 @@ def maybe_eval_type_data_closure_vals(unevaled_type_data: _Ptr[PyObject]):
                         continue
                     if k not in already_reified_type_params:
                         raise RuntimeError(
-                            f"typevar {k} not reified prior to evaluating dependent typevar {tvar}"
+                            f"typevar {k} not reified prior to evaluating dependent typevar {v}"
                         )
                     cvrs[k] = already_reified_type_params[k]
                 unevaled_type_data = copy_func(unevaled_type_data, cvrs)
             return unevaled_type_data()
 
-        generics = copy_object(self.generics)
-        for i, tvar in enumerate(generics):
-            if not isinstance(tvar, TypeVar):
-                raise RuntimeError(
-                    f"{i}th generic has probably already been reified as {tvar}; if you're using a global tvar for the generic, "
-                    f"you should use a unique one for each generic function."
-                )
+        for i, tvar in enumerate(self.generics):
+            if tvar.__name__ in body_builder.__globals__:
+                raise RuntimeError("global typevars for generics are not supported")
+
             type_var_default = None
-            if sys.version_info >= (3, 12):
-                type_var = PyTypeVarObject.from_object(tvar)
-                type_var_bound = type_var.bound
-                if sys.version_info >= (3, 13) and tvar.has_default():
-                    type_var_default = type_var.default_value
-            else:
-                type_var_bound = tvar.__bound__
+            type_var = PyTypeVarObject.from_object(tvar)
+            type_var_bound = type_var.bound
+            if sys.version_info >= (3, 13) and tvar.has_default():
+                type_var_default = type_var.default_value
 
             if bool(type_var_bound):
-                # before 3.12 typevar was just a python class
-                # https://github.com/python/cpython/blob/3.11/Lib/typing.py#L966
-                if sys.version_info >= (3, 12):
-                    type_var_bound = maybe_eval_type_data_closure_vals(type_var_bound)
+                type_var_bound = maybe_eval_type_data_closure_vals(type_var_bound)
             elif not bool(type_var_default):
                 if i >= len(item):
                     raise RuntimeError(f"generic {tvar=} must have concrete val")
@@ -372,15 +361,13 @@ def maybe_eval_type_data_closure_vals(unevaled_type_data: _Ptr[PyObject]):
             reified_type_params.append(r)
             already_reified_type_params[r.name] = r.val
 
-            # replace the tvar in body_builder's global context with the reified val
-            if tvar.__name__ in body_builder.__globals__:
-                body_builder.__globals__[tvar.__name__] = r.val
-            # replace the tvar in body_builder's closure with the reified val
+            # only in the closure if used in the body
             if r.name in body_builder.__code__.co_freevars:
                 free_i = body_builder.__code__.co_freevars.index(r.name)
-                assert (
-                    body_builder.__closure__[free_i].cell_contents == tvar
-                ), "typevars don't match"
+                if body_builder.__closure__[free_i].cell_contents != tvar:
+                    raise RuntimeError(
+                        f"typevars don't match: {id(body_builder.__closure__[free_i].cell_contents)=}, {id(tvar)=}"
+                    )
                 body_builder.__closure__[free_i].cell_contents = r.val
 
         name_mangled_generics = []
@@ -419,12 +406,9 @@ def func(
     func_attrs=None,
     function_type=None,
     emit=False,
-    generics=None,
     loc=None,
     ip=None,
 ) -> FuncBase:
-    if generics is None and hasattr(f, "__type_params__") and f.__type_params__:
-        generics = f.__type_params__
     func_ = FuncBase(
         body_builder=f,
         func_op_ctor=FuncOp.__base__,
@@ -436,7 +420,7 @@ def func(
         res_attrs=res_attrs,
         func_attrs=func_attrs,
         function_type=function_type,
-        generics=generics,
+        generics=getattr(f, "__type_params__", None),
         loc=loc,
         ip=ip,
     )

projects/eudsl-python-extras/mlir/extras/dialects/gpu.py

@@ -5,7 +5,6 @@
 from functools import partial
 from typing import Any, List, Optional, Tuple, Union
 
-
 from .func import FuncBase
 from .. import types as T
 from ..meta import (
@@ -25,6 +24,10 @@
     get_op_result_or_op_results,
 )
 from ...dialects.gpu import *
+
+del constant
+# constant needs to be below gpu import because it needs to shadow upstream's arith.constant
+# noinspection PyUnusedImports
 from .arith import constant
 from ...ir import (
     ArrayAttr,
@@ -439,13 +442,10 @@ def func(
     res_attrs=None,
     func_attrs=None,
     emit=False,
-    generics=None,
     loc=None,
     ip=None,
     emit_grid=False,
 ) -> Grid:
-    if generics is None and hasattr(f, "__type_params__") and f.__type_params__:
-        generics = f.__type_params__
     func_ = GPUFunc(
         body_builder=f,
         func_op_ctor=GPUFuncOp,
@@ -455,7 +455,7 @@ def func(
         arg_attrs=arg_attrs,
         res_attrs=res_attrs,
         func_attrs=func_attrs,
-        generics=generics,
+        generics=getattr(f, "__type_params__", None),
         loc=loc,
         ip=ip,
     )