From 73d91204ddbc812387e680c33148ef85b4c6b344 Mon Sep 17 00:00:00 2001 From: Stella Laurenzo Date: Wed, 30 Oct 2024 18:53:46 -0700 Subject: [PATCH] Add ignored file (?) --- compiler/bindings/python/iree/build/main.py | 244 ++++++++++++++++++++ 1 file changed, 244 insertions(+) create mode 100644 compiler/bindings/python/iree/build/main.py diff --git a/compiler/bindings/python/iree/build/main.py b/compiler/bindings/python/iree/build/main.py new file mode 100644 index 000000000000..272533c848d6 --- /dev/null +++ b/compiler/bindings/python/iree/build/main.py @@ -0,0 +1,244 @@ +# Copyright 2024 The IREE Authors +# +# Licensed under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +from typing import Any, IO + +import argparse +import importlib +import importlib.util +from pathlib import Path +import sys + +from iree.build.executor import BuildEntrypoint, Entrypoint, Executor + +__all__ = [ + "iree_build_main", + "load_build_module", +] + + +def iree_build_main( + module="__main__", + args: list[str] | None = None, + stdout: IO | None = None, + stderr: IO | None = None, +): + """Make a build module invoke iree.build on itself when run. + + Typically, if you have a module that declares build entrypoints, you will + add a stanza at the end: + + .. code-block:: python + from iree.build import * + + if __name__ == "__main__": + iree_build_main() + """ + main = CliMain(module=module, args=args, stdout=stdout, stderr=stderr) + main.run() + + +def load_build_module(module_path: Path | str): + """Loads a build module by path, evaling and returning it.""" + spec = importlib.util.spec_from_file_location("__iree_build__", module_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +class CliMain: + """Composes command line programs.""" + + def __init__( + self, + *, + args: list[str] | None = None, + module=None, + stdout: IO | None = None, + stderr: IO | None = None, + ): + self.stdout = stdout if stdout is not None else sys.stdout + self.stderr = stderr if stderr is not None else sys.stderr + if args is None: + args = sys.argv[1:] + if module is not None and isinstance(module, str): + module = __import__(module) + module = module + + p = self.arg_parser = argparse.ArgumentParser( + description="IREE program build driver" + ) + if module is None: + args, self.top_module = self._resolve_module_arguments(args) + else: + self.top_module = ModuleWrapper(module) + + p.add_argument( + "--output-dir", + type=Path, + default=Path.cwd(), + help="Output directory for the build tree (defaults to current directory)", + ) + + cmd_group_desc = p.add_argument_group( + title="Build command", + description="Selects a build sub-command to invoke (default '--build')", + ) + cmd_group = cmd_group_desc.add_mutually_exclusive_group() + cmd_group.add_argument( + "--build", + dest="command", + action="store_const", + const=self.build_command, + help="Executes build actions", + ) + cmd_group.add_argument( + "--list", + dest="command", + action="store_const", + const=self.list_command, + help="Lists top level build actions", + ) + + cmd_group.add_argument( + "--list-all", + dest="command", + action="store_const", + const=self.list_all_command, + help="Lists all build actions", + ) + + p.add_argument( + "action_path", + nargs="*", + help="Paths of actions to build (default to top-level actions)", + ) + + self._define_action_arguments(p) + self.args = self.arg_parser.parse_args(args) + + def abort(self): + sys.exit(1) + + def _define_action_arguments(self, p: argparse.ArgumentParser): + user_group = p.add_argument_group("Action defined options") + for ep in self.top_module.entrypoints.values(): + for cl_arg in ep.cl_args(): + cl_arg.define_arg(user_group) + + def _resolve_module_arguments( + self, args: list[str] + ) -> tuple[list[str], "ModuleWrapper"]: + p = argparse.ArgumentParser( + add_help=False, + usage="python -m iree.build [-m] build_module [... additional module specific options ...]", + prog="python -m iree.build", + ) + # Invoked as a standalone tool: need the user to specify the + # module. + p.add_argument( + "-m", + dest="parse_as_module", + action="store_true", + help="Interpret the build definitions argument as a module (vs a file)", + ) + p.add_argument( + "build_module", + help="The Python file or module from which to load build definitions", + ) + + bootstrap_args, rem_args = p.parse_known_args(args) + # Resolve from arguments. + is_module = bootstrap_args.parse_as_module or _is_module_like_str( + bootstrap_args.build_module + ) + if is_module: + try: + top_module = ModuleWrapper.load_module(bootstrap_args.build_module) + except ModuleNotFoundError as e: + print( + f"ERROR: Module '{bootstrap_args.build_module}' not found: {e}", + file=self.stderr, + ) + self.abort() + else: + top_module = ModuleWrapper.load_py_file(bootstrap_args.build_module) + return rem_args, top_module + + def _create_executor(self) -> Executor: + executor = Executor(self.args.output_dir, self.args, stderr=self.stderr) + executor.analyze(*self.top_module.entrypoints.values()) + return executor + + def run(self): + command = self.args.command + if command is None: + command = self.build_command + command() + + def build_command(self): + executor = self._create_executor() + + if not self.args.action_path: + # Default to all. + build_actions = list(executor.entrypoints) + else: + # Look up each requested and add it. + build_actions = [] + for action_path in self.args.action_path: + try: + build_actions.append(executor.all[action_path]) + except KeyError: + all_paths = "\n".join(executor.all.keys()) + print( + f"ERROR: Action '{action_path}' not found. Available: \n{all_paths}", + file=self.stderr, + ) + self.abort() + executor.build(*build_actions) + + for build_action in build_actions: + if isinstance(build_action, BuildEntrypoint): + for output in build_action.outputs: + print(f"{output.get_fs_path()}", file=self.stdout) + + def list_command(self): + executor = self._create_executor() + for ep in executor.entrypoints: + print(ep.path, file=self.stdout) + + def list_all_command(self): + executor = self._create_executor() + for name in executor.all.keys(): + if name: + print(name, file=self.stdout) + + +class ModuleWrapper: + """Wraps a raw, loaded module with access to discovered details.""" + + def __init__(self, mod): + self.mod = mod + self.entrypoints = self._collect_entrypoints() + + @staticmethod + def load_module(module_name: str) -> "ModuleWrapper": + return ModuleWrapper(importlib.import_module(module_name)) + + @staticmethod + def load_py_file(module_path: Path | str) -> "ModuleWrapper": + return ModuleWrapper(load_build_module(str(module_path))) + + def _collect_entrypoints(self) -> dict[str, Entrypoint]: + results: dict[str, Entrypoint] = {} + for attr_name, attr_value in self.mod.__dict__.items(): + if isinstance(attr_value, Entrypoint): + results[attr_name] = attr_value + return results + + +def _is_module_like_str(s: str) -> bool: + return "/" not in s and "\\" not in s and not s.endswith(".py")