defer if yield results type setting (ie use the set_type api)

Author: makslevental

mlir_utils/__init__.py

@@ -1,6 +1,7 @@
-from ._configuration.configuration import alias_upstream_bindings
 import atexit
 
+from ._configuration.configuration import alias_upstream_bindings
+
 if alias_upstream_bindings():
     from mlir import ir
 

mlir_utils/_configuration/configuration.py

@@ -5,9 +5,10 @@
 from base64 import urlsafe_b64encode
 from importlib.metadata import distribution, packages_distributions
 from importlib.resources import files
-from importlib.resources.readers import MultiplexedPath
 from pathlib import Path
 
+from importlib.resources.readers import MultiplexedPath
+
 from .module_alias_map import get_meta_path_insertion_index, AliasedModuleFinder
 
 __MLIR_PYTHON_PACKAGE_PREFIX__ = "__MLIR_PYTHON_PACKAGE_PREFIX__"

mlir_utils/ast/canonicalize.py

@@ -1,67 +1,43 @@
 import ast
-import logging
-
+import enum
 import inspect
+import logging
 import types
 from abc import ABC
+from opcode import opmap
 from types import CodeType
-from typing import Type, Optional
 
 import libcst as cst
 from bytecode import ConcreteBytecode
 from libcst.matchers import MatcherDecoratableTransformer
-from libcst.metadata import ParentNodeProvider
 
 from mlir_utils.ast.util import get_module_cst, copy_func
 
 logger = logging.getLogger(__name__)
 
 
-class FuncIdentTypeTable(cst.CSTVisitor):
-    METADATA_DEPENDENCIES = (ParentNodeProvider,)
-
-    def __init__(self, f):
-        super().__init__()
-        self.ident_type: dict[str, Type] = {}
-        module_cst = get_module_cst(f)
-        wrapper = cst.MetadataWrapper(module_cst)
-        wrapper.visit(self)
-
-    def visit_Annotation(self, node: cst.Annotation) -> Optional[bool]:
-        parent = self.get_metadata(ParentNodeProvider, node)
-        if isinstance(node.annotation, (cst.Tuple, cst.List)):
-            self.ident_type[parent.target.value] = [
-                e.value.value for e in node.annotation.elements
-            ]
-        else:
-            self.ident_type[parent.target.value] = [node.annotation.value]
-
-    def __getitem__(self, ident):
-        return self.ident_type[ident]
-
-
 class Transformer(MatcherDecoratableTransformer):
-    def __init__(self, context, func_sym_table: FuncIdentTypeTable):
+    def __init__(self, context):
         super().__init__()
         self.context = context
-        self.func_sym_table = func_sym_table
 
 
 class StrictTransformer(Transformer):
     def visit_FunctionDef(self, node: cst.FunctionDef):
         return False
 
 
-def transform_cst(f, transformers: list[type(StrictTransformer)] = None):
+def transform_cst(
+    f, transformers: list[type(Transformer) | type(StrictTransformer)] = None
+):
     if transformers is None:
         return f
 
     module_cst = get_module_cst(f)
-    func_sym_table = FuncIdentTypeTable(f)
     context = types.SimpleNamespace()
     for transformer in transformers:
         func_node = module_cst.body[0]
-        replace = transformer(context, func_sym_table)
+        replace = transformer(context)
         new_func = func_node._visit_and_replace_children(replace)
         module_cst = module_cst.deep_replace(func_node, new_func)
 
@@ -79,6 +55,23 @@ def transform_cst(f, transformers: list[type(StrictTransformer)] = None):
     return copy_func(f, new_f_code_o)
 
 
+# this is like this because i couldn't figure out how to subclass
+# Enum and simultaneously pass in opmap
+OpCode = enum.Enum("OpCode", opmap)
+
+
+def to_int(self: OpCode):
+    return self.value
+
+
+def to_str(self: OpCode):
+    return self.name
+
+
+setattr(OpCode, "__int__", to_int)
+setattr(OpCode, "__str__", to_str)
+
+
 class BytecodePatcher(ABC):
     def __init__(self, context=None):
         self.context = context
@@ -109,10 +102,10 @@ def bytecode_patchers(self) -> list[BytecodePatcher]:
         pass
 
 
-def canonicalize(*, with_: Canonicalizer):
+def canonicalize(*, using: Canonicalizer):
     def wrapper(f):
-        f = transform_cst(f, with_.cst_transformers)
-        f = patch_bytecode(f, with_.bytecode_patchers)
+        f = transform_cst(f, using.cst_transformers)
+        f = patch_bytecode(f, using.bytecode_patchers)
         return f
 
     return wrapper

mlir_utils/dialects/ext/arith.py

@@ -1,5 +1,5 @@
 from copy import deepcopy
-from functools import lru_cache, partialmethod, cached_property
+from functools import partialmethod, cached_property
 from typing import Union, Optional
 
 import numpy as np
@@ -143,7 +143,7 @@ def __call__(cls, *args, **kwargs):
             # which by default (through the Python buffer protocol) does not copy;
             # see mlir/lib/Bindings/Python/IRAttributes.cpp#L556
             arg_copy = deepcopy(arg)
-            val = constant(arg, dtype).result
+            return constant(arg_copy, dtype)
         else:
             raise NotImplementedError(f"{cls.__name__} doesn't support wrapping {arg}.")
 
@@ -155,7 +155,7 @@ def __call__(cls, *args, **kwargs):
         # the Python object protocol; first an object is new'ed and then
         # it is init'ed. Note we pass arg_copy here in case a subclass wants to
         # inspect the literal.
-        cls.__init__(cls_obj, val, arg_copy)
+        cls.__init__(cls_obj, val)
         return cls_obj
 
 
@@ -276,19 +276,12 @@ class ArithValue(Value, metaclass=ArithValueMeta):
                           Value.__init__
     """
 
-    def __init__(
-        self,
-        val,
-        arg: Optional[Union[int, float, bool, np.ndarray]] = None,
-    ):
-        self.__arg = arg
+    def __init__(self, val):
         super().__init__(val)
 
-    # @lru_cache(maxsize=1)
     def __str__(self):
         return f"{self.__class__.__name__}({self.get_name()}, {self.type})"
 
-    # @lru_cache(maxsize=1)
     def __repr__(self):
         return str(self)
 

mlir_utils/dialects/ext/func.py

@@ -1,5 +1,4 @@
 import inspect
-from functools import wraps, partial
 
 from mlir.dialects.func import FuncOp, ReturnOp, CallOp
 from mlir.ir import (