wehub-resource-sync
244 lines
8.6 KiB
Python
244 lines
8.6 KiB
Python
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import re
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import TypedDict
|
|
|
|
# --- Type Aliases ---
|
|
# "{data_type[::phi::dtype::bfloat16]; data_layout[STRIDED]; place[Place(gpu:0)]; library_type[PLAIN]}"
|
|
KernelConfig = dict[str, str]
|
|
# "fused_transpose": [
|
|
# "{data_type[uint8_t]; data_layout[ONEDNN]; place[Place(cpu)]; library_type[MKLDNN]}",
|
|
# "{data_type[int8_t]; data_layout[ONEDNN]; place[Place(cpu)]; library_type[MKLDNN]}",
|
|
# "{data_type[::phi::dtype::bfloat16]; data_layout[ONEDNN]; place[Place(cpu)]; library_type[MKLDNN]}",
|
|
# "{data_type[float]; data_layout[ONEDNN]; place[Place(cpu)]; library_type[MKLDNN]}"
|
|
# ]
|
|
KernelManifest = dict[str, list[KernelConfig]]
|
|
|
|
|
|
class GpuKernelChangeSummary(TypedDict):
|
|
# New kernel names that contain "gpu".
|
|
new_kernels_with_gpu: list[str]
|
|
# Existing kernels that have new GPU configurations added.
|
|
kernels_with_new_gpu_support: list[str]
|
|
# Existing GPU kernels that ignore data_type changes.
|
|
gpu_kernels_with_new_data_types: dict[str, list[str]]
|
|
|
|
|
|
class KernelManifestComparer:
|
|
def __init__(self, baseline_manifest_data: dict[str, list[str]]) -> None:
|
|
self.baseline_manifest: KernelManifest = self._process_raw_manifest(
|
|
baseline_manifest_data
|
|
)
|
|
self._preprocessed_baseline = {
|
|
kernel_name: {
|
|
"configs_set": {tuple(sorted(cfg.items())) for cfg in configs},
|
|
"gpu_signatures": {
|
|
self._create_config_signature_without_dtype(cfg)
|
|
for cfg in configs
|
|
if "gpu" in cfg.get("place", "")
|
|
},
|
|
}
|
|
for kernel_name, configs in self.baseline_manifest.items()
|
|
}
|
|
|
|
def _parse_config_string(self, config_str: str) -> KernelConfig:
|
|
pattern = r"([\w:]+)\[([^\]]+)\]"
|
|
matches = re.findall(pattern, config_str)
|
|
return dict(matches)
|
|
|
|
def _create_config_signature_without_dtype(
|
|
self, config: KernelConfig
|
|
) -> tuple[tuple[str, str], ...]:
|
|
signature_items = {
|
|
key: value for key, value in config.items() if key != "data_type"
|
|
}
|
|
return tuple(sorted(signature_items.items()))
|
|
|
|
def _process_raw_manifest(
|
|
self, raw_manifest: dict[str, list[str]]
|
|
) -> KernelManifest:
|
|
processed_manifest: KernelManifest = {}
|
|
for kernel_name, config_strings in raw_manifest.items():
|
|
processed_manifest[kernel_name] = [
|
|
self._parse_config_string(cs) for cs in config_strings
|
|
]
|
|
return processed_manifest
|
|
|
|
def compare(
|
|
self, target_manifest_data: dict[str, list[str]]
|
|
) -> GpuKernelChangeSummary:
|
|
target_manifest = self._process_raw_manifest(target_manifest_data)
|
|
|
|
new_kernels_with_gpu: set[str] = set()
|
|
kernels_with_new_gpu_support: set[str] = set()
|
|
gpu_kernels_with_new_data_types: dict[str, set[str]] = {}
|
|
|
|
baseline_kernel_names = set(self._preprocessed_baseline.keys())
|
|
target_kernel_names = set(target_manifest.keys())
|
|
|
|
# 1. Find newly added kernels that are GPU-related.
|
|
added_kernel_names = target_kernel_names - baseline_kernel_names
|
|
for kernel_name in added_kernel_names:
|
|
if any(
|
|
"gpu" in cfg.get("place", "")
|
|
for cfg in target_manifest[kernel_name]
|
|
):
|
|
new_kernels_with_gpu.add(kernel_name)
|
|
|
|
# 2. Find changes within existing kernels.
|
|
common_kernel_names = baseline_kernel_names.intersection(
|
|
target_kernel_names
|
|
)
|
|
for kernel_name in common_kernel_names:
|
|
baseline_data = self._preprocessed_baseline[kernel_name]
|
|
target_configs = target_manifest[kernel_name]
|
|
|
|
target_configs_set = {
|
|
tuple(sorted(cfg.items())) for cfg in target_configs
|
|
}
|
|
|
|
if baseline_data["configs_set"] == target_configs_set:
|
|
continue
|
|
|
|
added_config_tuples = (
|
|
target_configs_set - baseline_data["configs_set"]
|
|
)
|
|
added_configs = [dict(t) for t in added_config_tuples]
|
|
|
|
baseline_gpu_signatures = baseline_data["gpu_signatures"]
|
|
has_baseline_gpu_support = bool(baseline_gpu_signatures)
|
|
|
|
for added_cfg in added_configs:
|
|
if "gpu" in added_cfg.get("place", ""):
|
|
kernels_with_new_gpu_support.add(kernel_name)
|
|
added_signature = (
|
|
self._create_config_signature_without_dtype(added_cfg)
|
|
)
|
|
if (
|
|
has_baseline_gpu_support
|
|
and added_signature in baseline_gpu_signatures
|
|
):
|
|
new_data_type = added_cfg.get("data_type", "N/A")
|
|
gpu_kernels_with_new_data_types.setdefault(
|
|
kernel_name, set()
|
|
).add(new_data_type)
|
|
|
|
return {
|
|
"new_kernels_with_gpu": sorted(new_kernels_with_gpu),
|
|
"kernels_with_new_gpu_support": sorted(
|
|
kernels_with_new_gpu_support
|
|
),
|
|
"gpu_kernels_with_new_data_types": {
|
|
k: sorted(v) for k, v in gpu_kernels_with_new_data_types.items()
|
|
},
|
|
}
|
|
|
|
|
|
def cli():
|
|
parser = argparse.ArgumentParser(
|
|
description="Compare GPU kernel configurations between two manifests."
|
|
)
|
|
parser.add_argument(
|
|
"baseline_manifest",
|
|
type=Path,
|
|
help="Path to the baseline manifest JSON file.",
|
|
)
|
|
parser.add_argument(
|
|
"target_manifest",
|
|
type=Path,
|
|
help="Path to the target manifest JSON file.",
|
|
)
|
|
parser.add_argument(
|
|
"--ignore-data-type-changes",
|
|
action="store_true",
|
|
help="If set, ignore data_type changes in GPU kernel comparisons.",
|
|
)
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def main():
|
|
args = cli()
|
|
if not args.baseline_manifest.exists():
|
|
raise ValueError(
|
|
f"Baseline manifest file not found: {args.baseline_manifest.resolve()}"
|
|
)
|
|
if not args.target_manifest.exists():
|
|
raise ValueError(
|
|
f"Target manifest file not found: {args.target_manifest.resolve()}"
|
|
)
|
|
|
|
with args.baseline_manifest.open("r", encoding="utf-8") as f:
|
|
baseline_data = json.load(f)
|
|
|
|
with args.target_manifest.open("r", encoding="utf-8") as f:
|
|
target_data = json.load(f)
|
|
|
|
comparer = KernelManifestComparer(baseline_data)
|
|
summary = comparer.compare(target_data)
|
|
|
|
has_reportable_changes = False
|
|
if args.ignore_data_type_changes:
|
|
all_gpu_related_changes = set(summary["new_kernels_with_gpu"]) | set(
|
|
summary["kernels_with_new_gpu_support"]
|
|
)
|
|
kernels_with_only_dtype_changes = set(
|
|
summary["gpu_kernels_with_new_data_types"].keys()
|
|
)
|
|
unignored_changes = (
|
|
all_gpu_related_changes - kernels_with_only_dtype_changes
|
|
)
|
|
if unignored_changes:
|
|
has_reportable_changes = True
|
|
else:
|
|
if (
|
|
summary["new_kernels_with_gpu"]
|
|
or summary["kernels_with_new_gpu_support"]
|
|
):
|
|
has_reportable_changes = True
|
|
|
|
if summary["new_kernels_with_gpu"]:
|
|
print("New GPU Kernels Added:")
|
|
for kernel in summary["new_kernels_with_gpu"]:
|
|
print(f" - {kernel}")
|
|
|
|
if summary["kernels_with_new_gpu_support"]:
|
|
print("\nKernels with New GPU Support:")
|
|
for kernel in summary["kernels_with_new_gpu_support"]:
|
|
print(f" - {kernel}")
|
|
|
|
if summary["gpu_kernels_with_new_data_types"]:
|
|
print("\nGPU Kernels with new Data Type:")
|
|
for kernel, data_types in summary[
|
|
"gpu_kernels_with_new_data_types"
|
|
].items():
|
|
data_types_str = ", ".join(data_types)
|
|
print(f" - {kernel}: New data types - {data_types_str}")
|
|
|
|
if has_reportable_changes:
|
|
sys.exit(1)
|
|
|
|
sys.exit(0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|