xref: /llvm-project/mlir/python/mlir/dialects/linalg/opdsl/dump_oplib.py (revision f9008e6366c2496b1ca1785b891d5578174ad63e) 
1#!/usr/bin/which python
2# Command line tool to load an oplib module and dump all of the operations
3# it contains in some format.
4"""Loads one or more modules containing op definitions and dumps them.
5
6The dump format can be:
7
8* `--dump_format=yaml` (default)
9* `--dump_format=repr`
10
11Positional arguments are interpreted as module names (optionally, relative to
12this module). Loose module files can be specified via `--file <filepath>`.
13
14Sample usage:
15  # Dump the YAML op definitions for the core named ops (as in the dialect
16  # source tree).
17  python -m mlir.dialects.linalg.opdsl.dump_oplib .ops.core_named_ops
18
19Note: YAML output is emitted in "document list" format with each operation
20as its own "document". Practically, this means that each operation (or group
21of composite ops) is emitted with a "---" preceding it, which can be useful
22for testing.
23"""
24
25import argparse
26import importlib
27
28from .lang import *
29from .lang.config import *
30from .lang.yaml_helper import *
31
32
33def create_arg_parser() -> argparse.ArgumentParser:
34    p = argparse.ArgumentParser(description="Dump an oplib in various formats")
35    p.add_argument(
36        "modules", metavar="M", type=str, nargs="*", help="Op module to dump"
37    )
38    p.add_argument(
39        "--file", metavar="F", type=str, nargs="*", help="Python op file to dump"
40    )
41    p.add_argument(
42        "--format",
43        type=str,
44        dest="format",
45        default="yaml",
46        choices=("yaml", "repr"),
47        help="Format in which to dump",
48    )
49    return p
50
51
52def load_module_from_file(module_name, file_path):
53    spec = importlib.util.spec_from_file_location(module_name, file_path)
54    m = importlib.util.module_from_spec(spec)
55    spec.loader.exec_module(m)
56    return m
57
58
59def main(args):
60    # Load all configs.
61    configs = []
62    modules = []
63    for module_name in args.modules:
64        modules.append(
65            importlib.import_module(module_name, package="mlir.dialects.linalg.opdsl")
66        )
67    for i, file_path in enumerate(args.file or []):
68        modules.append(load_module_from_file(f"_mlir_eval_oplib{i}", file_path))
69    for m in modules:
70        for attr_name, value in m.__dict__.items():
71            # TODO: This class layering is awkward.
72            if isinstance(value, DefinedOpCallable):
73                try:
74                    linalg_config = LinalgOpConfig.from_linalg_op_def(value.op_def)
75                except Exception as e:
76                    raise ValueError(
77                        f"Could not create LinalgOpConfig from {value.op_def}"
78                    ) from e
79                configs.extend(linalg_config)
80
81    # Print.
82    if args.format == "yaml":
83        print(yaml_dump_all(configs))
84    elif args.format == "repr":
85        for config in configs:
86            print(repr(config))
87
88
89if __name__ == "__main__":
90    main(create_arg_parser().parse_args())
91