add generated files to manifest (for deletion during uninstall)

Author: makslevental

.github/workflows/test.yml

@@ -38,7 +38,7 @@ jobs:
           pip install .[torch-mlir-test] \
             -f https://github.com/makslevental/mlir-wheels/releases/expanded_assets/latest \
             -f https://llvm.github.io/torch-mlir/package-index/
-          pushd /tmp && python -m mlir_utils.__configuration -y torch_mlir && popd
+          pushd /tmp && python -m mlir_utils._configuration -y torch_mlir && popd
 
       - name: Test
         shell: bash

examples/throwaway.py

@@ -1,8 +1,16 @@
 import inspect
+import sys
 from functools import wraps
+from pathlib import Path
 
 from mlir_utils.context import mlir_mod_ctx
-from mlir_utils.dialects.generate_trampolines import generate_all_upstream_trampolines
+from mlir_utils._configuration.generate_trampolines import (
+    generate_all_upstream_trampolines,
+)
+from mlir_utils._configuration.configuration import _add_file_to_sources_txt_file
+
+# _add_file_to_sources_txt_file(Path("_configuration/__MLIR_PYTHON_PACKAGE_PREFIX__"))
+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
@@ -11,7 +19,6 @@
 
 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
@@ -56,3 +63,4 @@
 # print(ctx.module)
 # ctx.module.operation.verify()
 # print(ctx.module)
+# from importlib.resources import files

mlir_utils/__configuration/__init__.py

@@ -1,3 +1,3 @@
-from .__configuration import load_upstream_bindings
+from ._configuration.configuration import alias_upstream_bindings
 
-load_upstream_bindings()
+alias_upstream_bindings()

mlir_utils/__configuration/__main__.py

@@ -0,0 +1,2 @@
+from . import configuration
+from . import generate_trampolines

mlir_utils/__configuration/configuration.py

@@ -0,0 +1,6 @@
+from .configuration import configure_host_bindings
+
+# noinspection PyUnresolvedReferences
+from .generate_trampolines import generate_trampolines
+
+configure_host_bindings()

mlir_utils/__init__.py

@@ -0,0 +1,115 @@
+import argparse
+import hashlib
+import os
+import sys
+from base64 import urlsafe_b64encode
+from importlib.metadata import distribution, packages_distributions
+from importlib.resources import files
+from pathlib import Path
+
+from .module_alias_map import get_meta_path_insertion_index, AliasedModuleFinder
+
+__MLIR_PYTHON_PACKAGE_PREFIX__ = "__MLIR_PYTHON_PACKAGE_PREFIX__"
+PACKAGE = __package__.split(".")[0]
+PACKAGE_ROOT_PATH = files(PACKAGE)
+DIST = distribution(packages_distributions()[PACKAGE][0])
+MLIR_PYTHON_PACKAGE_PREFIX_TOKEN_PATH = (
+    Path(__file__).parent / __MLIR_PYTHON_PACKAGE_PREFIX__
+).absolute()
+
+
+def _add_file_to_sources_txt_file(file_path: Path):
+    assert file_path.exists(), f"file being added doesn't exist at {file_path}"
+    relative_file_path = Path(PACKAGE) / file_path.relative_to(PACKAGE_ROOT_PATH)
+    if DIST._read_files_egginfo() is not None:
+        with open(DIST._path / "SOURCES.txt", "a") as sources_file:
+            sources_file.write(f"\n{relative_file_path}")
+    if DIST._read_files_distinfo():
+        with open(file_path, "rb") as file, open(
+            DIST._path / "RECORD", "a"
+        ) as sources_file:
+            # https://packaging.python.org/en/latest/specifications/recording-installed-packages/#the-record-file
+            m = hashlib.sha256()
+            file = file.read()
+            m.update(file)
+            encoded = urlsafe_b64encode(m.digest())
+            sources_file.write(
+                f"{relative_file_path},sha256={encoded[:-1].decode()},{len(file)}\n"
+            )
+    else:
+        raise RuntimeError("unsupported distribution scheme; please file a bug!")
+
+
+def _get_mlir_package_prefix():
+    mlir_python_package_prefix = None
+    if MLIR_PYTHON_PACKAGE_PREFIX_TOKEN_PATH.exists():
+        with open(MLIR_PYTHON_PACKAGE_PREFIX_TOKEN_PATH) as f:
+            mlir_python_package_prefix = f.read().strip()
+
+    if os.getenv("MLIR_PYTHON_PACKAGE_PREFIX"):
+        mlir_python_package_prefix = os.getenv("MLIR_PYTHON_PACKAGE_PREFIX")
+
+    return mlir_python_package_prefix
+
+
+def alias_upstream_bindings():
+    if mlir_python_package_prefix := _get_mlir_package_prefix():
+        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"
+        )
+
+
+def configure_host_bindings():
+    parser = argparse.ArgumentParser(
+        prog="configure-mlir-utils",
+        description="Configure mlir-utils",
+    )
+    parser.add_argument("-y", "--yes", action="store_true", default=False)
+    parser.add_argument("mlir_python_package_prefix")
+    args = parser.parse_args()
+    mlir_python_package_prefix = args.mlir_python_package_prefix
+    mlir_python_package_prefix = (
+        mlir_python_package_prefix.replace("'", "").replace('"', "").strip()
+    )
+
+    if current_mlir_python_package_prefix := _get_mlir_package_prefix():
+        print(
+            f'mlir_python_package_prefix has already been set to "{current_mlir_python_package_prefix}"'
+        )
+        if not args.yes:
+            answer = input("do you want to reset? [y/n]: ")
+            if answer.lower() not in {"1", "true", "yes", "y"}:
+                return
+
+    if not args.yes:
+        answer = input(f"new {mlir_python_package_prefix=}; continue? [y/n]: ")
+        if answer.lower() not in {"1", "true", "yes", "y"}:
+            return
+    else:
+        print(f"new {mlir_python_package_prefix=}")
+
+    # check if valid package/module
+    try:
+        _host_bindings_mlir = __import__(f"{mlir_python_package_prefix}._mlir_libs")
+    except (ImportError, ModuleNotFoundError) as e:
+        print(f"couldn't import {mlir_python_package_prefix=} due to: {e}")
+        raise e
+
+    with open(MLIR_PYTHON_PACKAGE_PREFIX_TOKEN_PATH, "w") as f:
+        f.write(mlir_python_package_prefix)
+
+    _add_file_to_sources_txt_file(MLIR_PYTHON_PACKAGE_PREFIX_TOKEN_PATH)
+
+    alias_upstream_bindings()
+
+    from .generate_trampolines import generate_all_upstream_trampolines
+
+    generate_all_upstream_trampolines()