fix upstream binding aliasing yet again

Author: makslevental

README.md

@@ -1,10 +1,39 @@
 # mlir-utils
 
+The missing pieces (as far as boilerplate reduction goes) of the upstream MLIR python bindings.
+
+## Install
+
+This package is meant to work in concert with the upstream bindings.
+Practically speaking that means you need to have *some* package installed that includes mlir python bindings.
+In addition, you have to do one of two things to **configure this package** (after installing it):
+
+1. `$ configure-mlir-utils -y <MLIR_PYTHON_PACKAGE_PREFIX>`, where `MLIR_PYTHON_PACKAGE_PREFIX` is (as it says) the
+   package prefix for your chosen upstream bindings. So for example, for `torch-mlir`, you would
+   execute `configure-mlir-utils torch_mlir`, since `torch-mlir`'s bindings are the root of the `torch-mlir` python
+   package. **When in doubt about this prefix**, it is everything up until `ir` (e.g., as
+   in `from torch_mlir import ir`).
+2. `$ export MLIR_PYTHON_PACKAGE_PREFIX=<MLIR_PYTHON_PACKAGE_PREFIX>`, i.e., you can set this string as an environment
+   variable each time you use this package. Note, in this case, if you want to make use of the "op trampolines", you
+   still need to run `generate_trampolines.generate_all_upstream_trampolines()` by hand at least once.
+
+There is a convenience `extra_requires` in `pyproject.toml` such that you can do this:
+
+```shell
+$ pip install .[mlir] -f https://github.com/makslevental/mlir-wheels/releases/expanded_assets/latest
+$ configure-mlir-utils mlir
+```
+
+# Examples
+
+Check out the [tests](tests).
+
 ## Dev
 
 ```shell
 # you need setuptools >= 64 for build_editable
 pip install setuptools -U
-pip install -e ".[mlir-test]" \
-  -f https://github.com/makslevental/mlir-wheels/releases/expanded_assets/latest 
+pip install -e .[torch-mlir-test] \
+   -f https://github.com/makslevental/mlir-wheels/releases/expanded_assets/latest \
+   -f https://llvm.github.io/torch-mlir/package-index/
 ```

examples/demo.py

@@ -0,0 +1,58 @@
+import inspect
+from functools import wraps
+
+from mlir_utils.context import mlir_mod_ctx
+from mlir_utils.dialects.generate_trampolines import generate_all_upstream_trampolines
+from mlir_utils.dialects.memref import alloca_scope, return_
+from mlir_utils.dialects.transform import foreach, yield_
+from mlir_utils.dialects import gpu
+from mlir_utils.dialects.ext import func
+
+
+from mlir_utils.dialects.util import constant
+
+generate_all_upstream_trampolines()
+# # generate_all_upstream_trampolines()
+# from mlir.dialects.scf import WhileOp
+# from mlir.ir import InsertionPoint
+#
+#
+# # from mlir_utils.dialects.scf import execute_region, yield_
+#
+#
+# # def doublewrap(f):
+# #     """
+# #     a decorator decorator, allowing the decorator to be used as:
+# #     @decorator(with, arguments, and=kwargs)
+# #     or
+# #     @decorator
+# #     """
+# #
+# #     @wraps(f)
+# #     def new_dec(*args, **kwargs):
+# #         if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
+# #             # actual decorated function
+# #             return f(args[0])
+# #         else:
+# #             # decorator arguments
+# #             return lambda realf: f(realf, *args, **kwargs)
+# #
+# #     return new_dec
+# #
+#
+#
+# with mlir_mod_ctx() as ctx:
+#     one = constant(1)
+#
+#     @func.func
+#     def demo_fun1():
+#         one = constant(1)
+#         return
+#
+#     demo_fun1()
+#     ctx.module.operation.verify()
+#
+#
+# print(ctx.module)
+# ctx.module.operation.verify()
+# print(ctx.module)

examples/throwaway.py

@@ -0,0 +1,85 @@
+import importlib
+import importlib.util
+import sys
+from importlib.abc import Loader, MetaPathFinder
+from importlib.machinery import ModuleSpec, PathFinder
+from types import ModuleType
+from typing import Mapping, Optional, Sequence, Union
+
+# stolen from https://github.com/dagster-io/dagster/blob/master/python_modules/dagster/dagster/_module_alias_map.py
+
+# The AliasedModuleFinder should be inserted in front of the built-in PathFinder.
+def get_meta_path_insertion_index() -> int:
+    for i in range(len(sys.meta_path)):
+        finder = sys.meta_path[i]
+        if isinstance(finder, type) and issubclass(finder, PathFinder):
+            return i
+    raise Exception(
+        "Could not find the built-in PathFinder in sys.meta_path-- cannot insert the"
+        " AliasedModuleFinder"
+    )
+
+
+# Key reference to understand the load process:
+#   https://docs.python.org/3/reference/import.html#loading
+
+
+# While it is possible to override `Loader.create_module` to simply return the base module, this
+# is undesirable because the import system modifies the module's metadata attributes after
+# creation and outside of our control. This means that, if we simply had:
+#
+#     def create_module(self, spec):
+#         return importlib.import_module(self.base_spec_name)
+#
+# The returned base module would have its name etc modified, e.g. the already-loaded
+# `dagster._core` would be renamed to alias `dagster.core`. To avoid this, we let the import system
+# generate a module using default logic, then simply discard this module in `exec_module` (the final
+# step), where it is passed in. This is the point at which we swap in the base module, which we
+# obtain through `importlib.import_module` (like a standard import statetment, this will simply
+# returned the cached module from `sys.modules` if it has already been loaded). The swap is done by
+# simply replacing the dummy module (already stored in `sys.modules` outside of our control) with
+# the imported base module.
+class AliasedModuleLoader(Loader):
+    def __init__(self, alias: str, base_spec: ModuleSpec):
+        self.alias = alias
+        self.base_spec = base_spec
+
+    def exec_module(self, _module: ModuleType) -> None:
+        base_module = importlib.import_module(self.base_spec.name)
+        sys.modules[self.alias] = base_module
+
+    def module_repr(self, module: ModuleType) -> str:
+        assert self.base_spec.loader
+        return self.base_spec.loader.module_repr(module)
+
+
+class AliasedModuleFinder(MetaPathFinder):
+    def __init__(self, alias_map: Mapping[str, str]):
+        self.alias_map = alias_map
+
+    def find_spec(
+        self,
+        fullname: str,
+        _path: Optional[Sequence[Union[bytes, str]]] = None,
+        _target: Optional[ModuleType] = None,
+    ) -> Optional[ModuleSpec]:
+        head = next(
+            (
+                k
+                for k in self.alias_map.keys()
+                if fullname == k or fullname.startswith(k + ".")
+            ),
+            None,
+        )
+        if head is not None:
+            base_name = self.alias_map[head] + fullname[len(head) :]
+            base_spec = importlib.util.find_spec(base_name)
+            assert base_spec, f"Could not find module spec for {base_name}."
+            return ModuleSpec(
+                fullname,
+                AliasedModuleLoader(fullname, base_spec),
+                origin=base_spec.origin,
+                is_package=base_spec.submodule_search_locations is not None,
+            )
+        else:
+            return None

mlir_utils/__configuration/_module_alias_map.py

@@ -1,22 +1,16 @@
 import argparse
-import importlib
 import os
-import pkgutil
 import sys
 from pathlib import Path
 
+from mlir_utils.__configuration._module_alias_map import (
+    get_meta_path_insertion_index,
+    AliasedModuleFinder,
+)
 
-__MLIR_PYTHON_PACKAGE_PREFIX__ = None
 THIS_DIR = Path(__file__).resolve().parent
 MLIR_PYTHON_PACKAGE_PREFIX_FILE_PATH = THIS_DIR / "__MLIR_PYTHON_PACKAGE_PREFIX__"
-
-
-def import_submodules(package_name):
-    package = sys.modules[package_name]
-    return {
-        name: importlib.import_module(package_name + "." + name)
-        for loader, name, is_pkg in pkgutil.walk_packages(package.__path__)
-    }
+__MLIR_PYTHON_PACKAGE_PREFIX__ = None
 
 
 def load_upstream_bindings():
@@ -28,35 +22,18 @@ def load_upstream_bindings():
 
     if os.getenv("MLIR_PYTHON_PACKAGE_PREFIX"):
         __MLIR_PYTHON_PACKAGE_PREFIX__ = os.getenv("MLIR_PYTHON_PACKAGE_PREFIX")
-
     if __MLIR_PYTHON_PACKAGE_PREFIX__ is not None:
-        _mlir = sys.modules["mlir"] = __import__(
-            __MLIR_PYTHON_PACKAGE_PREFIX__, globals(), locals(), fromlist=["*"]
+        sys.meta_path.insert(
+            get_meta_path_insertion_index(),
+            AliasedModuleFinder({"mlir": __MLIR_PYTHON_PACKAGE_PREFIX__}),
+        )
+    elif not (
+        sys.argv[0].endswith("configure-mlir-utils")
+        or ("-m" in sys.orig_argv and "mlir_utils.__configuration" in sys.orig_argv)
+    ):
+        raise Exception(
+            "mlir-utils not configured and MLIR_PYTHON_PACKAGE_PREFIX env variable not set"
         )
-        for submod in ["ir", "dialects", "_mlir_libs"]:
-            sys.modules[f"mlir.{submod}"] = __import__(
-                f"{__MLIR_PYTHON_PACKAGE_PREFIX__}.{submod}",
-                globals(),
-                locals(),
-                fromlist=["*"],
-            )
-        mlir_modules = {}
-        for name, mod in sys.modules.items():
-            if name.startswith(__MLIR_PYTHON_PACKAGE_PREFIX__ + "."):
-                mlir_name = (
-                    "mlir." + name[len(__MLIR_PYTHON_PACKAGE_PREFIX__ + ".") + 1 :]
-                )
-                mlir_modules[mlir_name] = mod
-        sys.modules.update(mlir_modules)
-
-    else:
-        if not (
-            sys.argv[0].endswith("configure-mlir-utils")
-            or ("-m" in sys.orig_argv and "mlir_utils.__configuration" in sys.orig_argv)
-        ):
-            raise Exception(
-                "mlir-utils not configured and MLIR_PYTHON_PACKAGE_PREFIX env variable not set"
-            )
 
 
 def configure_host_bindings():

mlir_utils/__configuration/configuration.py

@@ -2,7 +2,7 @@
 from dataclasses import dataclass
 from typing import Optional
 
-import mlir
+import mlir.ir
 
 
 @dataclass

mlir_utils/context.py

@@ -1,8 +1,14 @@
 import inspect
 from functools import wraps
 
-from mlir.dialects.func import FuncOp, ReturnOp
-from mlir.ir import InsertionPoint, FunctionType, StringAttr, TypeAttr
+from mlir.dialects.func import FuncOp, ReturnOp, CallOp
+from mlir.ir import (
+    InsertionPoint,
+    FunctionType,
+    StringAttr,
+    TypeAttr,
+    FlatSymbolRefAttr,
+)
 
 from mlir_utils.dialects.util import (
     get_result_or_results,
@@ -16,18 +22,20 @@ def func(sym_visibility=None, arg_attrs=None, res_attrs=None, loc=None, ip=None)
 
     def builder_wrapper(body_builder):
         @wraps(body_builder)
-        def wrapper(*args):
+        def wrapper(*call_args):
             sig = inspect.signature(body_builder)
             implicit_return = sig.return_annotation is inspect._empty
+            input_types = [a.type for a in call_args]
             function_type = TypeAttr.get(
                 FunctionType.get(
-                    inputs=[a.type for a in args],
+                    inputs=input_types,
                     results=[] if implicit_return else sig.return_annotation,
                 )
             )
             # FuncOp is extended but we do really want the base
-            op = FuncOp.__base__(
-                body_builder.__name__,
+            func_name = body_builder.__name__
+            func_op = FuncOp.__base__(
+                func_name,
                 function_type,
                 sym_visibility=StringAttr.get(str(sym_visibility))
                 if sym_visibility is not None
@@ -37,19 +45,30 @@ def wrapper(*args):
                 loc=loc,
                 ip=ip,
             )
-            op.regions[0].blocks.append(*[a.type for a in args])
-            with InsertionPoint(op.regions[0].blocks[0]):
-                r = get_result_or_results(
-                    body_builder(*op.regions[0].blocks[0].arguments)
+            func_op.regions[0].blocks.append(*[a.type for a in call_args])
+            with InsertionPoint(func_op.regions[0].blocks[0]):
+                results = get_result_or_results(
+                    body_builder(*func_op.regions[0].blocks[0].arguments)
                 )
-                if r is not None:
-                    if isinstance(r, (tuple, list)):
-                        ReturnOp(list(r))
+                if results is not None:
+                    if isinstance(results, (tuple, list)):
+                        results = list(results)
                     else:
-                        ReturnOp([r])
+                        results = [results]
                 else:
-                    ReturnOp([])
-                return r
+                    results = []
+                ReturnOp(results)
+            # Recompute the function type.
+            return_types = [v.type for v in results]
+            function_type = FunctionType.get(inputs=input_types, results=return_types)
+            func_op.attributes["function_type"] = TypeAttr.get(function_type)
+
+            call_op = CallOp(
+                [r.type for r in results], FlatSymbolRefAttr.get(func_name), call_args
+            )
+            if results is None:
+                return None
+            return get_result_or_results(call_op)
 
         # wrapper.op = op
         return wrapper