forked from ccf-ai-infra/GPUCodeForces
294 lines
9.5 KiB
Python
294 lines
9.5 KiB
Python
#!/usr/bin/env python3
|
||
"""Generate the implemented-operator checklist in README.md.
|
||
|
||
The checklist is driven by two sources:
|
||
- scripts/operator_targets.txt: known operators to track, one per line.
|
||
- S1 codes/: submitted implementations containing *cuda*.py files.
|
||
|
||
New implementations are marked as checked automatically when the discovered
|
||
operator name matches a target name. Operators discovered from submissions but
|
||
missing from the target file are included as checked rows as well.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import re
|
||
from pathlib import Path
|
||
from urllib.parse import quote
|
||
|
||
|
||
ROOT = Path(__file__).resolve().parents[1]
|
||
README = ROOT / "README.md"
|
||
S1_DIR = ROOT / "S1 codes"
|
||
TARGETS = ROOT / "scripts" / "operator_targets.txt"
|
||
|
||
START = "<!-- OPERATOR_CHECKLIST_START -->"
|
||
END = "<!-- OPERATOR_CHECKLIST_END -->"
|
||
|
||
|
||
def clean_operator(value: str | None) -> str | None:
|
||
if not value:
|
||
return None
|
||
|
||
text = value.strip()
|
||
text = re.sub(r"\s*\((?:Fused\s+)?CUDA\s+Kernel\)\s*$", "", text)
|
||
text = re.sub(r"\s+with custom CUDA kernel\s*$", "", text)
|
||
text = re.sub(r"（一次核内.*$", "", text)
|
||
text = re.sub(r"，.*$", "", text)
|
||
text = re.sub(r"。.*$", "", text)
|
||
text = text.strip(" :：.。")
|
||
return text or None
|
||
|
||
|
||
def normalize_key(value: str) -> str:
|
||
return re.sub(r"\s+", "", value).casefold()
|
||
|
||
|
||
def read_text(path: Path) -> str:
|
||
return path.read_text(encoding="utf-8", errors="ignore")
|
||
|
||
|
||
def operator_from_torch(torch_path: Path) -> str | None:
|
||
if not torch_path.exists():
|
||
return None
|
||
|
||
text = read_text(torch_path)
|
||
if re.search(r"torch\.relu\(x \* self\.scale \+ self\.bias\)", text):
|
||
return "Affine+ReLU"
|
||
if re.search(
|
||
r"torch\.maximum\(x,\s*torch\.zeros_like\(x\)\).*torch\.log1p\(torch\.exp\(-ax\)\)",
|
||
text,
|
||
re.S,
|
||
):
|
||
return "BCEWithLogitsLoss"
|
||
|
||
match = re.search(r"return\s+F\.([A-Za-z0-9_]+)\(", text)
|
||
if match:
|
||
return match.group(1)
|
||
|
||
match = re.search(r"return\s+torch\.([A-Za-z0-9_]+)\(", text)
|
||
if match:
|
||
return match.group(1)
|
||
|
||
return None
|
||
|
||
|
||
def operator_from_prompt(prompt_path: Path) -> str | None:
|
||
if not prompt_path.exists():
|
||
return None
|
||
|
||
skip_prefixes = (
|
||
"You write custom CUDA kernels",
|
||
"你需要为下面给定的架构",
|
||
"你可以自由选择",
|
||
"给定架构",
|
||
"torchcode.py",
|
||
"cudacode.py",
|
||
"run_code.py",
|
||
)
|
||
|
||
for raw_line in prompt_path.read_text(encoding="utf-8", errors="ignore").splitlines():
|
||
line = re.sub(r"^\s*[#>*\-`]+\s*", "", raw_line).strip()
|
||
if not line:
|
||
continue
|
||
if line.startswith(skip_prefixes):
|
||
continue
|
||
|
||
match = re.match(r"^Operator\s*[:：]\s*(.+)$", line)
|
||
if match:
|
||
return clean_operator(match.group(1))
|
||
|
||
match = re.match(r"^Objective\s*[:：].*?\bfor\s+(.+?)\s+to\s+achieve", line)
|
||
if match:
|
||
return clean_operator(match.group(1))
|
||
|
||
match = re.match(r"^Implement\s+(?:a\s+|an\s+)?(.+?)(?:\s+on\b|\s+for\b|\s*[:：]|$)", line)
|
||
if match:
|
||
return clean_operator(match.group(1))
|
||
|
||
match = re.search(r"“([^”]+)”", line)
|
||
if match:
|
||
return clean_operator(match.group(1))
|
||
|
||
match = re.search(r"算子\s*[:：]\s*([^。；;，,]+)", line)
|
||
if match:
|
||
return clean_operator(match.group(1))
|
||
|
||
match = re.match(r"^(.+?)融合\s*[:：]", line)
|
||
if match:
|
||
return clean_operator(match.group(1) + "融合")
|
||
|
||
if len(line) <= 80:
|
||
return clean_operator(line)
|
||
|
||
return None
|
||
|
||
|
||
def discover_operator(directory: Path) -> str | None:
|
||
cuda_files = sorted(directory.glob("*cuda*.py"), key=lambda p: p.name.casefold())
|
||
if not cuda_files:
|
||
return None
|
||
|
||
cuda_base = cuda_files[0].stem
|
||
if re.match(r"^(example_)?cudacode$", cuda_base, re.I):
|
||
torch_files = sorted(directory.glob("*torch*.py"), key=lambda p: p.name.casefold())
|
||
specific_torch = [
|
||
path for path in torch_files if not re.match(r"^(example_)?torchcode$", path.stem, re.I)
|
||
]
|
||
if specific_torch:
|
||
return re.sub(r"_?torch(code)?$", "", specific_torch[0].stem, flags=re.I).strip()
|
||
|
||
return (
|
||
operator_from_prompt(directory / "prompt.txt")
|
||
or operator_from_torch(directory / "torchcode.py")
|
||
or cuda_base
|
||
)
|
||
|
||
return re.sub(r"_?cudacode$|_?cuda$", "", cuda_base, flags=re.I).strip()
|
||
|
||
|
||
def discover_implemented() -> dict[str, dict[str, object]]:
|
||
implemented: dict[str, dict[str, object]] = {}
|
||
if not S1_DIR.exists():
|
||
return implemented
|
||
|
||
for directory in sorted((path for path in S1_DIR.iterdir() if path.is_dir()), key=lambda p: p.name.casefold()):
|
||
operator = clean_operator(discover_operator(directory))
|
||
if not operator:
|
||
continue
|
||
|
||
key = normalize_key(operator)
|
||
row = implemented.setdefault(key, {"name": operator, "directories": []})
|
||
directories = row["directories"]
|
||
assert isinstance(directories, list)
|
||
directories.append(directory.name)
|
||
|
||
return implemented
|
||
|
||
|
||
def read_targets() -> list[str]:
|
||
if not TARGETS.exists():
|
||
return []
|
||
|
||
targets: list[str] = []
|
||
seen: set[str] = set()
|
||
for raw_line in TARGETS.read_text(encoding="utf-8").splitlines():
|
||
line = raw_line.strip()
|
||
if not line or line.startswith("#"):
|
||
continue
|
||
key = normalize_key(line)
|
||
if key in seen:
|
||
continue
|
||
seen.add(key)
|
||
targets.append(line)
|
||
return targets
|
||
|
||
|
||
def write_targets(targets: list[str]) -> None:
|
||
body = [
|
||
"# Operator targets tracked in README.md",
|
||
"# Add unimplemented operators here, one per line.",
|
||
"# scripts/update_operator_checklist.py checks them when matching implementations appear under S1 codes/.",
|
||
"",
|
||
*targets,
|
||
"",
|
||
]
|
||
TARGETS.write_text("\n".join(body), encoding="utf-8", newline="\n")
|
||
|
||
|
||
def link_for_directory(directory: str) -> str:
|
||
return f"[{directory}](S1%20codes/{quote(directory)})"
|
||
|
||
|
||
def build_section(targets: list[str], implemented: dict[str, dict[str, object]]) -> str:
|
||
target_by_key = {normalize_key(target): target for target in targets}
|
||
all_keys = set(target_by_key) | set(implemented)
|
||
|
||
checked = len([key for key in all_keys if key in implemented])
|
||
unchecked = len(all_keys) - checked
|
||
|
||
rows: list[tuple[bool, str, list[str]]] = []
|
||
for key in sorted(all_keys, key=lambda item: (item not in implemented, (target_by_key.get(item) or implemented[item]["name"]).casefold())):
|
||
item = implemented.get(key)
|
||
name = target_by_key.get(key) or str(item["name"])
|
||
directories = list(item["directories"]) if item else []
|
||
rows.append((item is not None, name, directories))
|
||
|
||
lines = [
|
||
START,
|
||
"## ✅ 算子实现状态",
|
||
"",
|
||
"该清单由 `scripts/update_operator_checklist.py` 根据 `scripts/operator_targets.txt` 和 `S1 codes/` 自动生成；新增待实现算子请写入目标清单，新增实现目录后运行 `python scripts/update_operator_checklist.py --sync-targets` 即可自动勾选。",
|
||
"",
|
||
f"- 已实现：{checked}",
|
||
f"- 未实现：{unchecked}",
|
||
f"- 跟踪总数：{len(all_keys)}",
|
||
"",
|
||
"<details>",
|
||
"<summary>展开查看算子实现状态</summary>",
|
||
"",
|
||
"| 状态 | 算子 | 实现目录 |",
|
||
"| --- | --- | --- |",
|
||
]
|
||
|
||
for done, name, directories in rows:
|
||
status = "[x]" if done else "[ ]"
|
||
links = "<br>".join(link_for_directory(directory) for directory in directories) if directories else "-"
|
||
lines.append(f"| {status} | {name.replace('|', r'\|')} | {links} |")
|
||
|
||
lines.extend(["", "</details>", END])
|
||
return "\n".join(lines)
|
||
|
||
|
||
def update_readme(section: str) -> None:
|
||
content = README.read_text(encoding="utf-8")
|
||
old_static_pattern = r"(?s)\n?<!-- IMPLEMENTED_OPERATORS_START -->.*?<!-- IMPLEMENTED_OPERATORS_END -->\n?"
|
||
content = re.sub(old_static_pattern, "\n", content)
|
||
|
||
pattern = rf"(?s){re.escape(START)}.*?{re.escape(END)}"
|
||
if re.search(pattern, content):
|
||
updated = re.sub(pattern, section, content)
|
||
else:
|
||
heading = re.compile(r"(?m)^## 📥")
|
||
if heading.search(content):
|
||
updated = heading.sub(section + "\n\n---\n\n## 📥", content, count=1)
|
||
else:
|
||
updated = content.rstrip() + "\n\n" + section + "\n"
|
||
|
||
updated = re.sub(r"(?m)(?:^---\n\s*){2,}(?=<!-- OPERATOR_CHECKLIST_START -->)", "---\n\n", updated)
|
||
README.write_text(updated.rstrip() + "\n", encoding="utf-8", newline="\n")
|
||
|
||
|
||
def main() -> int:
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument(
|
||
"--sync-targets",
|
||
action="store_true",
|
||
help="Append discovered implemented operators to scripts/operator_targets.txt.",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
implemented = discover_implemented()
|
||
targets = read_targets()
|
||
|
||
if args.sync_targets:
|
||
by_key = {normalize_key(target): target for target in targets}
|
||
for key, item in implemented.items():
|
||
by_key.setdefault(key, str(item["name"]))
|
||
targets = sorted(by_key.values(), key=str.casefold)
|
||
write_targets(targets)
|
||
|
||
section = build_section(targets, implemented)
|
||
update_readme(section)
|
||
|
||
print(
|
||
f"Updated README operator checklist: "
|
||
f"{len(implemented)} implemented operators, {len(targets)} tracked targets."
|
||
)
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|