[eudsl-python-extras] name mangle generics "instantiations" (#207)

Author: makslevental

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

@@ -203,12 +203,17 @@ def from_object(cls, obj) -> Self:
         return cls.from_address(address(obj))
 
 
+# https://github.com/python/cpython/blob/9648eed33f5fca0fcb8802fe0be8d35907bc33e3/Objects/typevarobject.c#L19-L31
 class PyTypeVarObject(Structure):
     _fields_ = _py_object_fields + [
         ("ob_size", c_ssize_t),
         ("name", _Ptr[PyObject]),
-        # not sure why but this is the only thing that works but that's fine because it's the only thing we need
         ("bound", _Ptr[PyObject]),
+        ("evaluate_bound", _Ptr[PyObject]),
+        ("constraints", _Ptr[PyObject]),
+        ("evaluate_constraints", _Ptr[PyObject]),
+        ("default_value", _Ptr[PyObject]),
+        ("evaluate_default", _Ptr[PyObject]),
     ]
 
     @classmethod

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

@@ -4,13 +4,15 @@
 import ast
 import functools
 import inspect
+import io
 import types
 from opcode import opmap
 from textwrap import dedent
 from typing import Dict
 
 from bytecode import ConcreteBytecode
 from cloudpickle import cloudpickle
+from ...ir import Type
 
 
 def set_lineno(node, n=1):
@@ -117,6 +119,17 @@ def replace_closure(code, new_closure: Dict):
     return new_code, closure
 
 
+def unpickle_mlir_type(v):
+    return Type.parse(v)
+
+
+class MLIRTypePickler(cloudpickle.Pickler):
+    def reducer_override(self, obj):
+        if isinstance(obj, Type):
+            return unpickle_mlir_type, (str(obj),)
+        return super().reducer_override(obj)
+
+
 # Based on http://stackoverflow.com/a/6528148/190597 (Glenn Maynard);
 # potentially more complete approach https://stackoverflow.com/a/56901529/9045206
 def copy_func(f, new_closure: Dict = None):
@@ -125,7 +138,10 @@ def copy_func(f, new_closure: Dict = None):
     else:
         # see https://github.com/cloudpipe/cloudpickle/blob/f111f7ab6d302e9b1e2a568d0e4c574895db6a6e/cloudpickle/cloudpickle.py#L813
         # for how this trick is accomplished (dill and pickle both fail to pickle eg generic typevars)
-        closure = cloudpickle.loads(cloudpickle.dumps(f.__closure__))
+        with io.BytesIO() as file:
+            cp = MLIRTypePickler(file)
+            cp.dump(f.__closure__)
+            closure = cloudpickle.loads(file.getvalue())
         code = f.__code__
 
     g = types.FunctionType(

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

@@ -9,7 +9,7 @@
 from typing import Optional, List, Union, TypeVar
 
 from ..ast.util import copy_func
-from ..ast.py_type import PyTypeVarObject
+from ..ast.py_type import PyTypeVarObject, _Ptr, PyObject
 from ..meta import op_region_builder
 from .. import types as T
 from ..util import get_user_code_loc, make_maybe_no_args_decorator
@@ -141,6 +141,7 @@ def prep_func_types(sig, return_types):
 class ReifiedTypeParams:
     name: str
     val: object
+    type: Optional[type]
 
 
 class FuncBase:
@@ -232,7 +233,7 @@ def emit(self, *call_args, decl=False, force=False) -> FuncOp:
                     if self.generics is not None:
                         for t in self.generics:
                             if not isinstance(t, ReifiedTypeParams):
-                                raise RuntimeError(f"{t=} must reified")
+                                raise RuntimeError(f"{t=} must be reified")
                             locals[t.name] = t.val
                     for i, v in enumerate(input_types):
                         if isinstance(v, TypeVar):
@@ -308,35 +309,63 @@ def __getitem__(self, item):
         # this also copies the function so that the original body_builder remains "generic" (via its closure)
         body_builder = copy_func(self.body_builder)
         reified_type_params = []
-        # dumb but whatever
+
+        # For "generics" (i.e. typevars) which are dependent on previous generics (identified by the fact that they have vals in their own closures),
+        # we collect all such previous generics along with the concrete vals (into already_reified_type_params) and then
+        # evaluate the typevars in the fully-populated closure. Note, in order to get the unevaled typevar bound and default value
+        # we access them in the PyTypeVarObject C struct itself instead of the API that python provides.
         already_reified_type_params = {}
+
+        def maybe_eval_type_data_closure_vals(unevaled_type_data: _Ptr[PyObject]):
+            assert type(unevaled_type_data) == _Ptr[PyObject]
+            unevaled_type_data = unevaled_type_data.contents.into_object()
+            cvrs = inspect.getclosurevars(unevaled_type_data).nonlocals
+            if len(cvrs):
+                for k, v in cvrs.items():
+                    if not isinstance(v, TypeVar):
+                        continue
+                    if k not in already_reified_type_params:
+                        raise RuntimeError(
+                            f"typevar {k} not reified prior to evaluating dependent typevar {t}"
+                        )
+                    cvrs[k] = already_reified_type_params[k]
+                unevaled_type_data = copy_func(unevaled_type_data, cvrs)
+            return unevaled_type_data()
+
         generics = copy.deepcopy(self.generics)
         for i, t in enumerate(generics):
+            type_var_default = None
             if sys.version_info >= (3, 12):
-                type_var_bound = PyTypeVarObject.from_object(t).bound
+                type_var = PyTypeVarObject.from_object(t)
+                type_var_bound = type_var.bound
+                if sys.version_info >= (3, 13) and t.has_default():
+                    type_var_default = type_var.default_value
             else:
                 type_var_bound = t.__bound__
-            if type_var_bound:
+
+            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 = lambda: type_var_bound
+                if sys.version_info >= (3, 12):
+                    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 {t} must have concrete val")
+                if isinstance(item[i], Type):
+                    type_var_bound = "type"
                 else:
-                    type_var_bound = type_var_bound.contents.into_object()
-                    cvrs = inspect.getclosurevars(type_var_bound).nonlocals
-                    if len(cvrs):
-                        for k, v in cvrs.items():
-                            if not isinstance(v, TypeVar):
-                                continue
-                            if k not in already_reified_type_params:
-                                raise RuntimeError(
-                                    f"typevar {k} not reified prior to evaluating dependent typevar {t}"
-                                )
-                            cvrs[k] = already_reified_type_params[k]
-                        type_var_bound = copy_func(type_var_bound, cvrs)
-                r = ReifiedTypeParams(t.__name__, type_var_bound())
+                    type_var_bound = type(item[i]).__name__
+
+            if bool(type_var_default):
+                type_var_default = maybe_eval_type_data_closure_vals(type_var_default)
+                type_var_bound = type(type_var_default).__name__
+                val = type_var_default
             else:
-                r = ReifiedTypeParams(t.__name__, item[i])
+                if i >= len(item):
+                    raise RuntimeError(f"generic {t} must have concrete val")
+                val = item[i]
+
+            r = ReifiedTypeParams(t.__name__, val, type_var_bound)
 
             reified_type_params.append(r)
             already_reified_type_params[r.name] = r.val
@@ -357,6 +386,11 @@ def __getitem__(self, item):
             self.call_op_ctor,
             return_types=self.return_types,
             sym_visibility=self.sym_visibility,
+            sym_name=(
+                self.func_name
+                + "_"
+                + "_".join([f"{r.type}_{r.val}" for r in reified_type_params])
+            ),
             arg_attrs=self.arg_attrs,
             res_attrs=self.res_attrs,
             func_attrs=self.func_attrs,

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

@@ -300,8 +300,10 @@ def negf(I, O, *, loc=None, ip=None):
     return linalg.negf(I, loc=loc, ip=ip, outs=[O])
 
 
-def pooling_nchw_max(I, K, O, *, loc=None, ip=None):
-    return linalg.pooling_nchw_max(I, K, loc=loc, ip=ip, outs=[O])
+def pooling_nchw_max(I, K, O, *, strides, dilations, loc=None, ip=None):
+    return linalg.pooling_nchw_max(
+        I, K, strides=strides, dilations=dilations, loc=loc, ip=ip, outs=[O]
+    )
 
 
 def pooling_nchw_sum(I, K, O, *, loc=None, ip=None):

projects/eudsl-python-extras/mlir/extras/runtime/refbackend.py

@@ -178,6 +178,9 @@ def invoke(*args):
 
         return invoke
 
+    def __getitem__(self, item):
+        return getattr(self, item)
+
 
 # A return consumer is a trampoline to a python function that will store/capture the return from the return;
 # this is done because you can't return structs etc from C APIs.