[eudsl-python-extras] fix more func stuff

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,16 @@ 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: |
+          
+          DESELECT=""
+          if [[ $(python -c "print(__import__('sys').version_info < (3, 13))") == "True" ]]; then
+            DESELECT="--deselect projects/eudsl-python-extras/tests/dialect/test_generics.py"
+          fi
+          
+          echo $DESELECT
+          
+          python -m pytest projects/eudsl-python-extras/tests $DESELECT
 
   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
 
 
@@ -38,9 +39,9 @@ def ast_call(name, args=None, keywords=None):
 def get_module_cst(f):
     f_src = dedent(inspect.getsource(f))
     tree = ast.parse(f_src)
-    assert isinstance(tree.body[0], ast.FunctionDef), (
-        f"unexpected ast node {tree.body[0]}"
-    )
+    assert isinstance(
+        tree.body[0], ast.FunctionDef
+    ), f"unexpected ast node {tree.body[0]}"
     return tree
 
 
@@ -167,9 +168,9 @@ def copy_func(f, new_closure: Dict = None):
 
 def append_hidden_node(node_body, new_node):
     last_statement = node_body[-1]
-    assert last_statement.end_lineno is not None, (
-        f"last_statement {ast.unparse(last_statement)} must have end_lineno"
-    )
+    assert (
+        last_statement.end_lineno is not None
+    ), f"last_statement {ast.unparse(last_statement)} must have end_lineno"
     new_node = ast.fix_missing_locations(
         set_lineno(new_node, last_statement.end_lineno)
     )

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

@@ -7,10 +7,10 @@
 from functools import update_wrapper
 from typing import Optional, List, Union, TypeVar
 
-from ..ast.util import copy_func
+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 *
@@ -26,7 +26,6 @@
     Value,
 )
 
-
 _call = call
 
 

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,
@@ -108,9 +111,9 @@ def z(cls):
 
 def thread_id():
     return (
-        block_dim.x * block_dim.y * thread_idx.z
-        + block_dim.x * thread_idx.y
-        + thread_idx.x
+            block_dim.x * block_dim.y * thread_idx.z
+            + block_dim.x * thread_idx.y
+            + thread_idx.x
     )
 
 
@@ -126,7 +129,7 @@ def set_container_module(module):
 
 @register_attribute_builder("DeviceMappingArrayAttr")
 def get_device_mapping_array_attr(
-    mapping: List[Attribute], context: Optional[Context] = None
+        mapping: List[Attribute], context: Optional[Context] = None
 ) -> ArrayAttr:
     if context is None:
         context = Context.current
@@ -177,7 +180,7 @@ def smem_space(int=False):
 @_cext.register_operation(_Dialect, replace=True)
 class GPUModuleOp(GPUModuleOp):
     def __init__(
-        self, sym_name, targets: Optional[List[Attribute]] = None, *, loc=None, ip=None
+            self, sym_name, targets: Optional[List[Attribute]] = None, *, loc=None, ip=None
     ):
         if targets is None:
             targets = []
@@ -188,8 +191,8 @@ def __init__(
         sym_name = (
             sym_name
             if (
-                issubclass(type(sym_name), Attribute)
-                or not AttrBuilder.contains("SymbolNameAttr")
+                    issubclass(type(sym_name), Attribute)
+                    or not AttrBuilder.contains("SymbolNameAttr")
             )
             else AttrBuilder.get("SymbolNameAttr")(sym_name, context=_ods_context)
         )
@@ -227,17 +230,17 @@ def __prepare__(cls, name, bases, **kwargs):
 
 class GPUFuncOp(GPUFuncOp_):
     def __init__(
-        self,
-        sym_name,
-        function_type,
-        *,
-        sym_visibility=None,
-        arg_attrs=None,
-        res_attrs=None,
-        workgroup_attrib_attrs=None,
-        private_attrib_attrs=None,
-        loc=None,
-        ip=None,
+            self,
+            sym_name,
+            function_type,
+            *,
+            sym_visibility=None,
+            arg_attrs=None,
+            res_attrs=None,
+            workgroup_attrib_attrs=None,
+            private_attrib_attrs=None,
+            loc=None,
+            ip=None,
     ):
         super().__init__(
             function_type=function_type,
@@ -255,17 +258,17 @@ def __init__(
             self.operation.attributes["arg_attrs"] = (
                 arg_attrs
                 if (
-                    isinstance(arg_attrs, Attribute)
-                    or not AttrBuilder.contains("DictArrayAttr")
+                        isinstance(arg_attrs, Attribute)
+                        or not AttrBuilder.contains("DictArrayAttr")
                 )
                 else AttrBuilder.get("DictArrayAttr")(arg_attrs, context=_ods_context)
             )
         if res_attrs is not None:
             self.operation.attributes["res_attrs"] = (
                 res_attrs
                 if (
-                    isinstance(res_attrs, Attribute)
-                    or not AttrBuilder.contains("DictArrayAttr")
+                        isinstance(res_attrs, Attribute)
+                        or not AttrBuilder.contains("DictArrayAttr")
                 )
                 else AttrBuilder.get("DictArrayAttr")(res_attrs, context=_ods_context)
             )
@@ -274,23 +277,23 @@ def __init__(
             self.operation.attributes["sym_visibility"] = (
                 sym_visibility
                 if (
-                    issubclass(type(sym_visibility), Attribute)
-                    or not AttrBuilder.contains("StrAttr")
+                        issubclass(type(sym_visibility), Attribute)
+                        or not AttrBuilder.contains("StrAttr")
                 )
                 else AttrBuilder.get("StrAttr")(sym_visibility, context=_ods_context)
             )
 
 
 class LaunchOp(LaunchOp):
     def __init__(
-        self,
-        grid_size: Tuple[Any, Any, Any],
-        block_size: Tuple[Any, Any, Any],
-        async_dependencies=None,
-        dynamic_shared_memory_size: Optional[Value] = None,
-        *,
-        loc=None,
-        ip=None,
+            self,
+            grid_size: Tuple[Any, Any, Any],
+            block_size: Tuple[Any, Any, Any],
+            async_dependencies=None,
+            dynamic_shared_memory_size: Optional[Value] = None,
+            *,
+            loc=None,
+            ip=None,
     ):
         _ods_context = get_default_loc_context(loc)
         if async_dependencies is None:
@@ -309,13 +312,13 @@ def __init__(
 
 
 def launch_(
-    grid_size: Tuple[Any, Any, Any],
-    block_size: Tuple[Any, Any, Any],
-    async_dependencies=None,
-    dynamic_shared_memory_size: Optional[Value] = None,
-    *,
-    loc=None,
-    ip=None,
+        grid_size: Tuple[Any, Any, Any],
+        block_size: Tuple[Any, Any, Any],
+        async_dependencies=None,
+        dynamic_shared_memory_size: Optional[Value] = None,
+        *,
+        loc=None,
+        ip=None,
 ):
     for size in [grid_size, block_size]:
         for i, s in enumerate(size):
@@ -337,17 +340,17 @@ def launch_(
 
 class LaunchFuncOp(LaunchFuncOp):
     def __init__(
-        self,
-        kernel: List[str],
-        grid_size: Tuple[Any, Any, Any],
-        block_size: Tuple[Any, Any, Any],
-        kernel_operands: List[Value] = None,
-        async_dependencies=None,
-        dynamic_shared_memory_size: Optional[Value] = None,
-        async_object=None,
-        *,
-        loc=None,
-        ip=None,
+            self,
+            kernel: List[str],
+            grid_size: Tuple[Any, Any, Any],
+            block_size: Tuple[Any, Any, Any],
+            kernel_operands: List[Value] = None,
+            async_dependencies=None,
+            dynamic_shared_memory_size: Optional[Value] = None,
+            async_object=None,
+            *,
+            loc=None,
+            ip=None,
     ):
         _ods_context = get_default_loc_context(loc)
         if async_dependencies is None:
@@ -370,15 +373,15 @@ def __init__(
 
 class GPUFunc(FuncBase):
     def __call__(
-        self,
-        *kernel_operands: List[Value],
-        grid_size: Tuple[Any, Any, Any],
-        block_size: Tuple[Any, Any, Any],
-        async_dependencies=None,
-        dynamic_shared_memory_size: Optional[Value] = None,
-        stream=None,
-        loc=None,
-        ip=None,
+            self,
+            *kernel_operands: List[Value],
+            grid_size: Tuple[Any, Any, Any],
+            block_size: Tuple[Any, Any, Any],
+            async_dependencies=None,
+            dynamic_shared_memory_size: Optional[Value] = None,
+            stream=None,
+            loc=None,
+            ip=None,
     ):
         for size in [grid_size, block_size]:
             for i, s in enumerate(size):
@@ -432,16 +435,16 @@ def __call__(self, *args, **kwargs):
 
 @make_maybe_no_args_decorator
 def func(
-    f,
-    *,
-    sym_visibility=None,
-    arg_attrs=None,
-    res_attrs=None,
-    func_attrs=None,
-    emit=False,
-    loc=None,
-    ip=None,
-    emit_grid=False,
+        f,
+        *,
+        sym_visibility=None,
+        arg_attrs=None,
+        res_attrs=None,
+        func_attrs=None,
+        emit=False,
+        loc=None,
+        ip=None,
+        emit_grid=False,
 ) -> Grid:
     func_ = GPUFunc(
         body_builder=f,
@@ -490,14 +493,14 @@ def wait(async_dependencies: Optional[List[Value]] = None, *, loc=None, ip=None)
 
 
 def alloc(
-    sizes: Union[int, Value],
-    element_type: Type = None,
-    async_dependencies=None,
-    dynamic_sizes=None,
-    symbol_operands=None,
-    host_shared=None,
-    loc=None,
-    ip=None,
+        sizes: Union[int, Value],
+        element_type: Type = None,
+        async_dependencies=None,
+        dynamic_sizes=None,
+        symbol_operands=None,
+        host_shared=None,
+        loc=None,
+        ip=None,
 ):
     if symbol_operands is None:
         symbol_operands = []