llvm/mlir/python/mlir/dialects/linalg/opdsl/dump_oplib.py

#!/usr/bin/which python
# Command line tool to load an oplib module and dump all of the operations
# it contains in some format.
"""Loads one or more modules containing op definitions and dumps them.

The dump format can be:

* `--dump_format=yaml` (default)
* `--dump_format=repr`

Positional arguments are interpreted as module names (optionally, relative to
this module). Loose module files can be specified via `--file <filepath>`.

Sample usage:
  # Dump the YAML op definitions for the core named ops (as in the dialect
  # source tree).
  python -m mlir.dialects.linalg.opdsl.dump_oplib .ops.core_named_ops

Note: YAML output is emitted in "document list" format with each operation
as its own "document". Practically, this means that each operation (or group
of composite ops) is emitted with a "---" preceding it, which can be useful
for testing.
"""

import argparse
import importlib

from .lang import *
from .lang.config import *
from .lang.yaml_helper import *


def create_arg_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description="Dump an oplib in various formats")
    p.add_argument(
        "modules", metavar="M", type=str, nargs="*", help="Op module to dump"
    )
    p.add_argument(
        "--file", metavar="F", type=str, nargs="*", help="Python op file to dump"
    )
    p.add_argument(
        "--format",
        type=str,
        dest="format",
        default="yaml",
        choices=("yaml", "repr"),
        help="Format in which to dump",
    )
    return p


def load_module_from_file(module_name, file_path):
    spec = importlib.util.spec_from_file_location(module_name, file_path)
    m = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(m)
    return m


def main(args):
    # Load all configs.
    configs = []
    modules = []
    for module_name in args.modules:
        modules.append(
            importlib.import_module(module_name, package="mlir.dialects.linalg.opdsl")
        )
    for i, file_path in enumerate(args.file or []):
        modules.append(load_module_from_file(f"_mlir_eval_oplib{i}", file_path))
    for m in modules:
        for attr_name, value in m.__dict__.items():
            # TODO: This class layering is awkward.
            if isinstance(value, DefinedOpCallable):
                try:
                    linalg_config = LinalgOpConfig.from_linalg_op_def(value.op_def)
                except Exception as e:
                    raise ValueError(
                        f"Could not create LinalgOpConfig from {value.op_def}"
                    ) from e
                configs.extend(linalg_config)

    # Print.
    if args.format == "yaml":
        print(yaml_dump_all(configs))
    elif args.format == "repr":
        for config in configs:
            print(repr(config))


if __name__ == "__main__":
    main(create_arg_parser().parse_args())