check for failing case

Author: makslevental

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

@@ -38,9 +38,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
 
 
@@ -92,7 +92,7 @@ def replace_closure(code, new_closure: Dict):
     LOAD_DEREF = opmap["LOAD_DEREF"]
 
     # get the orig localplus that will be loaded from by the orig bytecode LOAD_DEREF arg_i
-    localsplus, localsplus_name_to_idx = get_localsplus_name_to_idx(code)
+    localsplus, _localsplus_name_to_idx = get_localsplus_name_to_idx(code)
 
     # closure vars go into co_freevars
     new_code = code.replace(co_freevars=tuple(new_closure.keys()))
@@ -167,9 +167,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

@@ -298,6 +298,8 @@ def __call__(self, *call_args):
         return call(self.emit(*call_args), call_args)
 
     def __getitem__(self, item):
+        if not isinstance(item, tuple):
+            item = (item,)
         if self.generics is None:
             raise RuntimeError(
                 "using a generic call requires the func be generic (i.e., have type_params)"
@@ -322,22 +324,26 @@ 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 {t}"
+                            f"typevar {k} not reified prior to evaluating dependent typevar {tvar}"
                         )
                     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):
+        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, you should give it a unique name."
+                )
             type_var_default = None
             if sys.version_info >= (3, 12):
-                type_var = PyTypeVarObject.from_object(t)
+                type_var = PyTypeVarObject.from_object(tvar)
                 type_var_bound = type_var.bound
-                if sys.version_info >= (3, 13) and t.has_default():
+                if sys.version_info >= (3, 13) and tvar.has_default():
                     type_var_default = type_var.default_value
             else:
-                type_var_bound = t.__bound__
+                type_var_bound = tvar.__bound__
 
             if bool(type_var_bound):
                 # before 3.12 typevar was just a python class
@@ -346,7 +352,7 @@ def maybe_eval_type_data_closure_vals(unevaled_type_data: _Ptr[PyObject]):
                     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")
+                    raise RuntimeError(f"generic {tvar=} must have concrete val")
                 if isinstance(item[i], Type):
                     type_var_bound = "type"
                 else:
@@ -358,29 +364,31 @@ def maybe_eval_type_data_closure_vals(unevaled_type_data: _Ptr[PyObject]):
                 val = type_var_default
             else:
                 if i >= len(item):
-                    raise RuntimeError(f"generic {t} must have concrete val")
+                    raise RuntimeError(f"generic {tvar=} must have concrete val")
                 val = item[i]
 
-            r = ReifiedTypeParams(t.__name__, val, type_var_bound)
+            r = ReifiedTypeParams(tvar.__name__, val, type_var_bound)
 
             reified_type_params.append(r)
             already_reified_type_params[r.name] = r.val
 
-            if t.__name__ in body_builder.__globals__:
-                body_builder.__globals__[t.__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
             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 == t
+                    body_builder.__closure__[free_i].cell_contents == tvar
                 ), "typevars don't match"
                 body_builder.__closure__[free_i].cell_contents = r.val
 
         name_mangled_generics = []
         for r in reified_type_params:
-            t, v = r.type, r.val
+            tvar, v = r.type, r.val
             if callable(v):
                 v = v.__name__
-            name_mangled_generics.append(f"{t}_{v}")
+            name_mangled_generics.append(f"{tvar}_{v}")
 
         return FuncBase(
             body_builder,

projects/eudsl-python-extras/tests/dialect/test_func.py

@@ -232,9 +232,9 @@ def mat_product_kernel2():
         one = arith.constant(1, T.f32())
         two = _op(one, one)
 
-    mat_product_kernel1[arith.maximumf,].emit()
-    mat_product_kernel2[arith.minimumf,].emit()
-    mat_product_kernel2[arith.maximumf,].emit()
+    mat_product_kernel1[arith.maximumf].emit()
+    mat_product_kernel2[arith.minimumf].emit()
+    mat_product_kernel2[arith.maximumf].emit()
 
     # CHECK:  func.func @mat_product_kernel1_function_maximumf() {
     # CHECK:    %cst = arith.constant 1.000000e+00 : f32
@@ -265,7 +265,7 @@ def test_global_closures(ctx: MLIRContext):
     def _generic_pool2d_scf(a: T.f32(), b: T.f32()):
         _op(a, b)
 
-    _maxpool2d_scf = _generic_pool2d_scf[arith.maximumf,]
+    _maxpool2d_scf = _generic_pool2d_scf[arith.maximumf]
 
     # _op = TypeVar("_op")
 
@@ -276,7 +276,11 @@ def _generic_pool3d_scf(
     ):
         _op(a, b)
 
-    _maxpool3d_scf = _generic_pool3d_scf[arith.maximumf,]
+    with pytest.raises(
+        RuntimeError,
+        match="0th generic has probably already been reified as <function maximumf at .*?>; if you're using a global tvar for the generic, you should give it a unique name.",
+    ):
+        _maxpool3d_scf = _generic_pool3d_scf[arith.maximumf]
 
 
 def test_generics_with_canonicalizations(ctx: MLIRContext):