diff --git a/scripts/infini_ops_plugin_registry.py b/scripts/infini_ops_plugin_registry.py index 09778771b..dbb15e20a 100644 --- a/scripts/infini_ops_plugin_registry.py +++ b/scripts/infini_ops_plugin_registry.py @@ -205,13 +205,18 @@ def _append_unique(values, new_values): values.append(value) -def load_plugin_registry(plugin_root, requested_plugins): +def load_plugin_manifests(plugin_root): plugin_root = pathlib.Path(plugin_root) - manifests = { + + return { path.parent.name: _load_manifest(path) for path in sorted(plugin_root.glob("*/plugin.json")) } + +def load_plugin_registry(plugin_root, requested_plugins): + manifests = load_plugin_manifests(plugin_root) + ordered = [] visiting = [] visited = set() diff --git a/scripts/infini_ops_plugin_test_matrix.py b/scripts/infini_ops_plugin_test_matrix.py new file mode 100644 index 000000000..4a5fd5aa2 --- /dev/null +++ b/scripts/infini_ops_plugin_test_matrix.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 +import argparse +import json +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) + +import infini_ops_plugin_registry + +CORE_PATHS = { + "CMakeLists.txt", + ".github/ci_config.yml", + "docs/plugin_contract.md", + "scripts/generate_wrappers.py", + "scripts/infini_ops_plugin_registry.py", + "scripts/infini_ops_plugin_test_matrix.py", + "src/CMakeLists.txt", +} + +CORE_PREFIXES = ( + ".github/workflows/", + "cmake/", + "include/", +) + + +def _normalize_path(path): + return pathlib.PurePosixPath(str(path).replace("\\", "/")).as_posix() + + +def _matches(path, root): + return path == root or path.startswith(f"{root}/") + + +def _device_plugins(manifests): + return [name for name, manifest in manifests.items() if manifest["kind"] == "device"] + + +def _dependent_plugins(manifests, plugin_name): + return [ + name + for name, manifest in manifests.items() + if plugin_name in manifest["depends"] + ] + + +def _expand_plugins(manifests, plugin_names): + expanded = [] + + def add(name): + if name not in expanded: + expanded.append(name) + + for name in plugin_names: + add(name) + if manifests[name]["kind"] == "shared": + for dependent in _dependent_plugins(manifests, name): + add(dependent) + + return expanded + + +def _manifest_roots(manifests): + roots = [] + for name, manifest in manifests.items(): + roots.append((f"plugins/{name}", name)) + for field in ("source_roots", "operator_roots"): + for root in manifest[field]: + roots.append((_normalize_path(root), name)) + for header in manifest["device_headers"].values(): + header = _normalize_path(header) + roots.append((header, name)) + roots.append((f"src/{header}", name)) + + return roots + + +def _plugins_for_path(manifests, path): + path = _normalize_path(path) + + if path in CORE_PATHS or any(path.startswith(prefix) for prefix in CORE_PREFIXES): + return _device_plugins(manifests), True + + matches = [ + (len(root), name) + for root, name in _manifest_roots(manifests) + if _matches(path, root) + ] + + if not matches: + return _device_plugins(manifests), True + + longest = max(length for length, _ in matches) + plugin_names = [name for length, name in matches if length == longest] + + return _expand_plugins(manifests, plugin_names), False + + +def _append_unique(values, new_values): + for value in new_values: + if value not in values: + values.append(value) + + +def build_test_matrix(plugin_root, changed_paths): + manifests = infini_ops_plugin_registry.load_plugin_manifests(plugin_root) + plugins = [] + requires_full_matrix = False + + for path in changed_paths: + path_plugins, path_requires_full_matrix = _plugins_for_path(manifests, path) + _append_unique(plugins, path_plugins) + requires_full_matrix = requires_full_matrix or path_requires_full_matrix + + devices = [] + test_devices = [] + for plugin in plugins: + manifest = manifests[plugin] + _append_unique(devices, manifest["devices"]) + _append_unique(test_devices, manifest["test_devices"].values()) + + ci_platforms = [device for device in devices if device != "cpu"] + + return { + "plugins": plugins, + "devices": devices, + "test_devices": test_devices, + "ci_platforms": ci_platforms, + "requires_full_matrix": requires_full_matrix, + } + + +def _read_paths(args): + if args.paths: + return args.paths + + return [line.strip() for line in sys.stdin if line.strip()] + + +def main(argv=None): + parser = argparse.ArgumentParser( + description="Map changed paths to `infini::ops` plugin test devices." + ) + parser.add_argument( + "--plugin-root", + default="plugins", + help="Directory containing built-in `plugin.json` manifests.", + ) + parser.add_argument("paths", nargs="*", help="Changed paths to classify.") + args = parser.parse_args(argv) + + matrix = build_test_matrix(args.plugin_root, _read_paths(args)) + print(json.dumps(matrix, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tests/test_plugin_test_matrix.py b/tests/test_plugin_test_matrix.py new file mode 100644 index 000000000..d0ce683ed --- /dev/null +++ b/tests/test_plugin_test_matrix.py @@ -0,0 +1,122 @@ +import importlib.util +import json +import pathlib +import subprocess +import sys + + +def _load_matrix_module(): + path = ( + pathlib.Path(__file__).resolve().parents[1] + / "scripts" + / "infini_ops_plugin_test_matrix.py" + ) + spec = importlib.util.spec_from_file_location("infini_ops_plugin_test_matrix", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + sys.modules[spec.name] = module + spec.loader.exec_module(module) + + return module + + +def _plugin_root(): + return pathlib.Path(__file__).resolve().parents[1] / "plugins" + + +def test_device_source_change_maps_to_one_platform(): + module = _load_matrix_module() + + matrix = module.build_test_matrix( + _plugin_root(), ["src/native/cuda/nvidia/ops/add/add.cu"] + ) + + assert matrix["plugins"] == ["nvidia"] + assert matrix["devices"] == ["nvidia"] + assert matrix["test_devices"] == ["cuda"] + assert matrix["ci_platforms"] == ["nvidia"] + assert matrix["requires_full_matrix"] is False + + +def test_shared_plugin_change_maps_to_dependents(): + module = _load_matrix_module() + + matrix = module.build_test_matrix( + _plugin_root(), ["src/native/cuda/ops/gemm/gemm.cc"] + ) + + assert matrix["plugins"] == [ + "cuda-common", + "hygon", + "iluvatar", + "metax", + "moore", + "nvidia", + ] + assert matrix["devices"] == ["hygon", "iluvatar", "metax", "moore", "nvidia"] + assert matrix["test_devices"] == ["cuda", "musa"] + assert matrix["ci_platforms"] == ["hygon", "iluvatar", "metax", "moore", "nvidia"] + assert matrix["requires_full_matrix"] is False + + +def test_core_codegen_change_requests_full_matrix(): + module = _load_matrix_module() + + matrix = module.build_test_matrix( + _plugin_root(), ["scripts/generate_wrappers.py"] + ) + + assert matrix["devices"] == [ + "ascend", + "cambricon", + "cpu", + "hygon", + "iluvatar", + "metax", + "moore", + "nvidia", + ] + assert matrix["ci_platforms"] == [ + "ascend", + "cambricon", + "hygon", + "iluvatar", + "metax", + "moore", + "nvidia", + ] + assert matrix["requires_full_matrix"] is True + + +def test_plugin_manifest_change_maps_to_that_plugin(): + module = _load_matrix_module() + + matrix = module.build_test_matrix(_plugin_root(), ["plugins/cambricon/plugin.json"]) + + assert matrix["plugins"] == ["cambricon"] + assert matrix["devices"] == ["cambricon"] + assert matrix["test_devices"] == ["mlu"] + assert matrix["ci_platforms"] == ["cambricon"] + + +def test_cli_outputs_json_matrix(tmp_path): + repo = pathlib.Path(__file__).resolve().parents[1] + script = repo / "scripts" / "infini_ops_plugin_test_matrix.py" + + result = subprocess.run( + [ + sys.executable, + str(script), + "--plugin-root", + str(_plugin_root()), + "plugins/cpu/plugin.cmake", + ], + check=True, + stdout=subprocess.PIPE, + text=True, + ) + + matrix = json.loads(result.stdout) + assert matrix["plugins"] == ["cpu"] + assert matrix["devices"] == ["cpu"] + assert matrix["ci_platforms"] == []