From 487a5b005c4eeba1025a1e8d17fe831288db39b8 Mon Sep 17 00:00:00 2001 From: TecJesh Date: Wed, 12 Aug 2026 05:05:48 +0000 Subject: [PATCH] [Format](fix) Add pre-commit code style check --- .github/PULL_REQUEST_TEMPLATE.md | 4 +- .pre-commit-config.yaml | 27 ++ .../agent/opencode_adapter.py | 86 +++- src/TA_main2main_workflow/flow.py | 55 +-- src/TA_main2main_workflow/main.py | 46 ++- src/TA_main2main_workflow/pipeline/build.py | 63 +-- src/TA_main2main_workflow/pipeline/commit.py | 7 +- .../pipeline/finalize.py | 64 +-- src/TA_main2main_workflow/pipeline/fix.py | 23 +- .../pipeline/ir_patch.py | 382 ++++++++++++------ src/TA_main2main_workflow/pipeline/merge.py | 24 +- src/TA_main2main_workflow/pipeline/plan.py | 70 +++- src/TA_main2main_workflow/pipeline/push_pr.py | 147 +++++-- src/TA_main2main_workflow/pipeline/resolve.py | 2 +- src/TA_main2main_workflow/pipeline/test.py | 103 +++-- .../01-merge-upstream-conflict-resolution.md | 4 +- ...vm-version-adaptation-and-compile-fixes.md | 2 +- ...3-unit-test-failure-diagnosis-and-fixes.md | 2 +- ...ir-compatibility-and-backend-adaptation.md | 4 +- .../ir_compatibility_patch_example.patch | 84 ++-- src/TA_main2main_workflow/utils/config.py | 30 +- src/TA_main2main_workflow/utils/git.py | 24 +- src/TA_main2main_workflow/utils/logging.py | 19 +- src/TA_main2main_workflow/utils/submodule.py | 6 +- 24 files changed, 835 insertions(+), 443 deletions(-) create mode 100644 .pre-commit-config.yaml diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index 6d2714b..6ed4f5a 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -2,11 +2,11 @@ -- [ ] run `pre-commit` to ensure code style +- [ ] run `pre-commit run --all-files` to ensure code style - [ ] unit tests added or updated (if applicable) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..8d50658 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,27 @@ +default_stages: [pre-commit, pre-push, manual] +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + - id: check-ast + - id: check-added-large-files + - id: check-merge-conflict + - id: debug-statements + - id: detect-private-key + + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.9.1 + hooks: + - id: ruff + args: ["--fix", "--exit-non-zero-on-fix"] + - id: ruff-format + + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.15.0 + hooks: + - id: mypy + args: ["--ignore-missing-imports", "src/"] + pass_filenames: false diff --git a/src/TA_main2main_workflow/agent/opencode_adapter.py b/src/TA_main2main_workflow/agent/opencode_adapter.py index c1d9689..c5172f4 100644 --- a/src/TA_main2main_workflow/agent/opencode_adapter.py +++ b/src/TA_main2main_workflow/agent/opencode_adapter.py @@ -40,6 +40,7 @@ # ── backend detection ──────────────────────────────────────────────────────── + def _detect_backend() -> str: """Detect which AI backend to use. Checks AI_BACKEND env var first, then falls back to whatever is available on PATH.""" @@ -72,8 +73,10 @@ def _detect_backend() -> str: # ── prompt loader ──────────────────────────────────────────────────────────── + def _build_prompt(inputs: dict[str, Any]) -> str: from collections import defaultdict + template = _PROMPT_PATH.read_text(encoding="utf-8") ctx = defaultdict(str, {k: str(v) for k, v in inputs.items()}) return template.format_map(ctx) @@ -81,6 +84,7 @@ def _build_prompt(inputs: dict[str, Any]) -> str: # ── result model ───────────────────────────────────────────────────────────── + class AIResult(BaseModel): modified_files: list[str] = Field(default_factory=list) is_noop: bool = Field(default=False) @@ -92,6 +96,7 @@ class AIResult(BaseModel): # ── main entry point ───────────────────────────────────────────────────────── + def run_opencode_adapter(inputs: dict[str, Any]) -> AIResult: """Run the AI adapter for conflict resolution or test fixing. @@ -125,9 +130,12 @@ def run_opencode_adapter(inputs: dict[str, Any]) -> AIResult: if result.modified_files: print(f" Modified: {', '.join(result.modified_files)}", flush=True) if result.resolved_conflicts: - print(f" Resolved conflicts: {', '.join(result.resolved_conflicts)}", flush=True) + print( + f" Resolved conflicts: {', '.join(result.resolved_conflicts)}", + flush=True, + ) if result.is_noop: - print(f" (no changes needed)", flush=True) + print(" (no changes needed)", flush=True) return result @@ -136,6 +144,7 @@ def run_opencode_adapter(inputs: dict[str, Any]) -> AIResult: # opencode backend (JSONL streaming) # ═══════════════════════════════════════════════════════════════════════════════ + def _run_opencode(inputs: dict[str, Any]) -> AIResult: """opencode backend: JSONL streaming with stale-timeout retry.""" base_prompt = _build_prompt(inputs) @@ -170,12 +179,17 @@ def _run_opencode(inputs: dict[str, Any]) -> AIResult: if reason == "stale_timeout" and attempt < _MAX_STALE_RETRIES: retry = attempt + 1 - print(f"\n[opencode] retrying after stale timeout ({retry}/{_MAX_STALE_RETRIES})", flush=True) + print( + f"\n[opencode] retrying after stale timeout ({retry}/{_MAX_STALE_RETRIES})", + flush=True, + ) prompt = _build_opencode_continue(base_prompt, inputs, retry) continue if stderr_path and stderr_path.exists(): - stderr_content = stderr_path.read_text(encoding="utf-8", errors="replace")[-2000:] + stderr_content = stderr_path.read_text(encoding="utf-8", errors="replace")[ + -2000: + ] if stderr_content: print(f"\n[opencode] stderr tail:\n{stderr_content}", flush=True) break @@ -186,8 +200,10 @@ def _run_opencode(inputs: dict[str, Any]) -> AIResult: return result -def _build_opencode_continue(base_prompt: str, inputs: dict[str, Any], retry: int) -> str: - return f"""Continue the task for step {inputs.get('step_id', '')}. +def _build_opencode_continue( + base_prompt: str, inputs: dict[str, Any], retry: int +) -> str: + return f"""Continue the task for step {inputs.get("step_id", "")}. The previous opencode run produced no output for {_STALE_SECONDS} seconds and was terminated. This is continuation retry {retry}/{_MAX_STALE_RETRIES}. @@ -212,7 +228,9 @@ def _print_prompt(prompt: str, attempt: int) -> None: print(f"{'━' * 60}", flush=True) if len(prompt) > 8000: print(prompt[:4000]) - print(f"\n... [{len(prompt) - 8000} chars truncated, see log for full prompt] ...\n") + print( + f"\n... [{len(prompt) - 8000} chars truncated, see log for full prompt] ...\n" + ) print(prompt[-4000:]) else: print(prompt) @@ -249,8 +267,10 @@ def _run_opencode_once( stderr_fh = stderr_path.open("a", encoding="utf-8") if stderr_path else None proc = subprocess.Popen( [ - "opencode", "run", - "--format", "json", + "opencode", + "run", + "--format", + "json", "--dangerously-skip-permissions", prompt, ], @@ -287,12 +307,18 @@ def _stdout_reader(): except queue.Empty: now = time.monotonic() if now > deadline: - print(f"\n[opencode] TOTAL TIMEOUT ({_TIMEOUT_MINUTES}min), killing process", flush=True) + print( + f"\n[opencode] TOTAL TIMEOUT ({_TIMEOUT_MINUTES}min), killing process", + flush=True, + ) proc.kill() stop_reason = "total_timeout" break if now - last_output_time > _STALE_SECONDS: - print(f"\n[opencode] STALE TIMEOUT ({_STALE_SECONDS}s no output), killing process", flush=True) + print( + f"\n[opencode] STALE TIMEOUT ({_STALE_SECONDS}s no output), killing process", + flush=True, + ) proc.kill() stop_reason = "stale_timeout" break @@ -363,8 +389,15 @@ def _print_opencode_event(line: str, state: _EventState) -> None: elif status == "completed": output = st.get("output", "") if output: - display = output if len(output) <= 2000 else output[:2000] + "\n... [truncated]" - print(f"\n {'─' * 56}\n [AI output]\n {display}\n {'─' * 56}", flush=True) + display = ( + output + if len(output) <= 2000 + else output[:2000] + "\n... [truncated]" + ) + print( + f"\n {'─' * 56}\n [AI output]\n {display}\n {'─' * 56}", + flush=True, + ) def _log_opencode_event(line: str, state: _EventState, fh: Any) -> None: @@ -402,6 +435,7 @@ def _log_opencode_event(line: str, state: _EventState, fh: Any) -> None: # in the meantime), this implementation streams stdout line-by-line in real # time. A heartbeat "." is printed every 15s of silence to show liveness. + def _run_claude(inputs: dict[str, Any]) -> AIResult: """Run Claude Code with real-time streaming output. @@ -431,8 +465,11 @@ def _run_claude(inputs: dict[str, Any]) -> AIResult: if log_path: _log_prompt(prompt, 0, log_path) - print(f"\n > [claude] Starting Claude Code (timeout={_TIMEOUT_MINUTES}min)...", flush=True) - print(f" (streaming output in real time — '.' = still thinking)", flush=True) + print( + f"\n > [claude] Starting Claude Code (timeout={_TIMEOUT_MINUTES}min)...", + flush=True, + ) + print(" (streaming output in real time — '.' = still thinking)", flush=True) # ── Launch claude ────────────────────────────────────────────────────── stderr_fh = stderr_path.open("a", encoding="utf-8") if stderr_path else None @@ -475,12 +512,18 @@ def _write_stdin(): if line is None: now = time.monotonic() if now > deadline: - print(f"\n [claude] TOTAL TIMEOUT ({_TIMEOUT_MINUTES}min), killing process", flush=True) + print( + f"\n [claude] TOTAL TIMEOUT ({_TIMEOUT_MINUTES}min), killing process", + flush=True, + ) proc.kill() stop_reason = "total_timeout" break if now - last_output_time > _STALE_SECONDS: - print(f"\n [claude] STALE TIMEOUT ({_STALE_SECONDS}s no output), killing process", flush=True) + print( + f"\n [claude] STALE TIMEOUT ({_STALE_SECONDS}s no output), killing process", + flush=True, + ) proc.kill() stop_reason = "stale_timeout" break @@ -520,7 +563,9 @@ def _write_stdin(): if proc.returncode != 0: print(f"\n [claude] exited with code {proc.returncode}", flush=True) if stderr_path and stderr_path.exists(): - stderr_tail = stderr_path.read_text(encoding="utf-8", errors="replace")[-2000:] + stderr_tail = stderr_path.read_text(encoding="utf-8", errors="replace")[ + -2000: + ] if stderr_tail: print(f" [claude] stderr tail:\n{stderr_tail}", flush=True) else: @@ -556,7 +601,10 @@ def _read_line_with_timeout(stream: Any, timeout: float) -> str | None: # shared result builder # ═══════════════════════════════════════════════════════════════════════════════ -def _build_result(step_dir: Path | None, ascend_path: str, output_text: str) -> AIResult: + +def _build_result( + step_dir: Path | None, ascend_path: str, output_text: str +) -> AIResult: """Build AIResult from AI output: extract summary, detect modified files.""" summary = "" if step_dir: diff --git a/src/TA_main2main_workflow/flow.py b/src/TA_main2main_workflow/flow.py index ed9c83a..58e2b18 100644 --- a/src/TA_main2main_workflow/flow.py +++ b/src/TA_main2main_workflow/flow.py @@ -23,8 +23,6 @@ from TA_main2main_workflow.utils import ( UpgradeCompleted, UpgradeFailed, - HasNewCommits, - HasNoNewCommits, WORKSPACE_DIR, ) from TA_main2main_workflow.pipeline.prepare import prepare @@ -36,7 +34,10 @@ from TA_main2main_workflow.pipeline.test import test_and_fix_loop from TA_main2main_workflow.pipeline.commit import commit_step from TA_main2main_workflow.pipeline.finalize import finalize -from TA_main2main_workflow.pipeline.ir_patch import build_baseline_llvm, per_step_ir_patch +from TA_main2main_workflow.pipeline.ir_patch import ( + build_baseline_llvm, + per_step_ir_patch, +) log = get_logger(__name__) @@ -97,9 +98,7 @@ def run(self) -> str: step_id = step["id"] ctx = ctx.copy_with(retry_count=0) - log.header( - f"Step {ctx.current_step + 1}/{ctx.total_steps}: {step_id}" - ) + log.header(f"Step {ctx.current_step + 1}/{ctx.total_steps}: {step_id}") log.key_value("commits in step", str(step["commit_count"])) log.key_value("end commit", step["end_commit"][:12]) reason = step.get("reason", "line_budget") @@ -129,13 +128,15 @@ def run(self) -> str: log.status(True, "Conflicts resolved") # ── Step C: Build/Test — IR patch or standard ─────── - llvm_hash_changed = (reason == "llvm_version") + llvm_hash_changed = reason == "llvm_version" if not llvm_hash_changed: llvm_hash_changed = llvm_hash_changed_after_merge(ctx) if llvm_hash_changed: if reason != "llvm_version": - log.info(f"[{step_id}] LLVM hash changed during merge " - f"(post-merge detection) — routing to IR patch pipeline") + log.info( + f"[{step_id}] LLVM hash changed during merge " + f"(post-merge detection) — routing to IR patch pipeline" + ) log.section(f"LLVM Version Change in {step_id} — IR Patch Pipeline") with timed("ir-patch"): ctx = per_step_ir_patch(ctx, self.config, step) @@ -176,25 +177,26 @@ def run(self) -> str: ctx.step_pr_descriptions.append(desc) # Record per-step detail for sync report - ctx.step_details.append({ - "step_id": step_id, - "step_index": ctx.current_step + 1, - "commits": step["commit_count"], - "end_commit": step["end_commit"][:12], - "source_lines": step.get("source_changed_lines", 0), - "conflict_files": len(ctx.conflict_files), - "build_fixes": ctx.build_fix_count, - "test_fixes": ctx.test_fix_count, - "retries": ctx.retry_count, - "reason": reason, - }) + ctx.step_details.append( + { + "step_id": step_id, + "step_index": ctx.current_step + 1, + "commits": step["commit_count"], + "end_commit": step["end_commit"][:12], + "source_lines": step.get("source_changed_lines", 0), + "conflict_files": len(ctx.conflict_files), + "build_fixes": ctx.build_fix_count, + "test_fixes": ctx.test_fix_count, + "retries": ctx.retry_count, + "reason": reason, + } + ) # Advance to next step ctx = ctx.copy_with(current_step=ctx.current_step + 1) log.status( True, - f"Step {step_id} completed " - f"({ctx.current_step}/{ctx.total_steps})", + f"Step {step_id} completed ({ctx.current_step}/{ctx.total_steps})", ) # ── Phase 4: Finalize ─────────────────────────────────────── @@ -207,8 +209,11 @@ def run(self) -> str: self._push_pr(ctx) ctx.summary_rows.append( - ("Single-Step Sync", "PASS", - f"{ctx.total_steps} step(s), branch: {ctx.work_branch}") + ( + "Single-Step Sync", + "PASS", + f"{ctx.total_steps} step(s), branch: {ctx.work_branch}", + ) ) log.table(ctx.summary_rows) log.elapsed(total_elapsed()) diff --git a/src/TA_main2main_workflow/main.py b/src/TA_main2main_workflow/main.py index 80ca79a..8714a06 100644 --- a/src/TA_main2main_workflow/main.py +++ b/src/TA_main2main_workflow/main.py @@ -32,9 +32,7 @@ """ import argparse -import os import sys -from pathlib import Path from TA_main2main_workflow.flow import TA_Main2MainFlow from TA_main2main_workflow.utils import UpgradeFailed @@ -49,40 +47,47 @@ def kickoff(): description="Triton-Ascend Main2Main Upstream Sync (Single-Step Mode)" ) parser.add_argument( - "--triton-ascend-path", default=None, - help="Local path to the triton-ascend repository (default: TRITON_ASCEND_PATH env)" + "--triton-ascend-path", + default=None, + help="Local path to the triton-ascend repository (default: TRITON_ASCEND_PATH env)", ) parser.add_argument( - "--triton-path", default=None, - help="Local path to the upstream triton repository (default: TRITON_PATH env)" + "--triton-path", + default=None, + help="Local path to the upstream triton repository (default: TRITON_PATH env)", ) parser.add_argument( - "--target-commit", default=None, - help="Upstream triton commit SHA to merge (default: upstream HEAD)" + "--target-commit", + default=None, + help="Upstream triton commit SHA to merge (default: upstream HEAD)", ) parser.add_argument( - "--llvm-prefix", default=None, - help="LLVM install prefix path for building" + "--llvm-prefix", default=None, help="LLVM install prefix path for building" ) parser.add_argument( - "--conda-env", default=None, - help="Conda environment name (default: ta-upgrade)" + "--conda-env", default=None, help="Conda environment name (default: ta-upgrade)" ) parser.add_argument( - "--build-procs", type=int, default=None, - help="Parallel build workers (default: 32)" + "--build-procs", + type=int, + default=None, + help="Parallel build workers (default: 32)", ) parser.add_argument( - "--test-procs", type=int, default=None, - help="Parallel pytest workers (default: 8)" + "--test-procs", + type=int, + default=None, + help="Parallel pytest workers (default: 8)", ) parser.add_argument( - "--extra-test-dirs", default=None, - help="Extra test directories (comma/space separated, appended to default pytest ut)" + "--extra-test-dirs", + default=None, + help="Extra test directories (comma/space separated, appended to default pytest ut)", ) parser.add_argument( - "--test-command", default=None, - help="Additional custom test command (runs after pytest ut)" + "--test-command", + default=None, + help="Additional custom test command (runs after pytest ut)", ) args = parser.parse_args() @@ -117,6 +122,7 @@ def kickoff(): except Exception as exc: log.error(f"WORKFLOW CRASHED: {exc}") import traceback + traceback.print_exc() sys.exit(1) diff --git a/src/TA_main2main_workflow/pipeline/build.py b/src/TA_main2main_workflow/pipeline/build.py index 0914918..fba71b5 100644 --- a/src/TA_main2main_workflow/pipeline/build.py +++ b/src/TA_main2main_workflow/pipeline/build.py @@ -24,11 +24,16 @@ from TA_main2main_workflow.utils.tracker import timed from TA_main2main_workflow.utils.git import run_git, stream_cmd from TA_main2main_workflow.utils import ( - BUILD_RESULT_FILE, STEPS_DIR, WORKSPACE_DIR, + BUILD_RESULT_FILE, + STEPS_DIR, + WORKSPACE_DIR, ) from TA_main2main_workflow.pipeline.fix import ai_fix from TA_main2main_workflow.pipeline.pre_ci import cleanup_temp_files -from TA_main2main_workflow.utils.submodule import commit_submodule, submodule_has_changes +from TA_main2main_workflow.utils.submodule import ( + commit_submodule, + submodule_has_changes, +) log = get_logger(__name__) @@ -112,10 +117,12 @@ def llvm_setup(ctx: WorkflowContext, config: TAConfig) -> WorkflowContext: def build_llvm(ctx: WorkflowContext, num_procs: int = 32) -> WorkflowContext: """cmake + ninja for LLVM. Pure build — no retry logic.""" - llvm_project = Path(os.path.expanduser( - os.getenv("LLVM_PROJECT_PATH", "~/llvm-project"))) - llvm_install = Path(os.path.expanduser( - os.getenv("LLVM_INSTALL_PREFIX_SYNC", "~/llvm-install-sync"))) + llvm_project = Path( + os.path.expanduser(os.getenv("LLVM_PROJECT_PATH", "~/llvm-project")) + ) + llvm_install = Path( + os.path.expanduser(os.getenv("LLVM_INSTALL_PREFIX_SYNC", "~/llvm-install-sync")) + ) ascend_path = Path(ctx.triton_ascend_path) required_hash = ( (ascend_path / "cmake" / "llvm-hash.txt").read_text(encoding="utf-8").strip() @@ -128,6 +135,7 @@ def build_llvm(ctx: WorkflowContext, num_procs: int = 32) -> WorkflowContext: build_dir = WORKSPACE_DIR / "llvm-build" if build_dir.exists(): import shutil + shutil.rmtree(build_dir) build_dir.mkdir(parents=True, exist_ok=True) @@ -137,8 +145,10 @@ def build_llvm(ctx: WorkflowContext, num_procs: int = 32) -> WorkflowContext: # ── cmake configure ── cmake_cmd = [ - "cmake", str(llvm_project / "llvm"), - "-G", "Ninja", + "cmake", + str(llvm_project / "llvm"), + "-G", + "Ninja", "-DCMAKE_BUILD_TYPE=Release", "-DLLVM_ENABLE_ASSERTIONS=ON", "-DLLVM_ENABLE_PROJECTS=mlir;llvm;lld", @@ -150,33 +160,34 @@ def build_llvm(ctx: WorkflowContext, num_procs: int = 32) -> WorkflowContext: with open(llvm_build_log, "w", encoding="utf-8") as fh: fh.write(f"=== cmake ===\n{' '.join(cmake_cmd)}\n\n") fh.flush() - rc = stream_cmd(cmake_cmd, build_dir, fh, timeout=300, - label="Configuring LLVM with cmake") + rc = stream_cmd( + cmake_cmd, build_dir, fh, timeout=300, label="Configuring LLVM with cmake" + ) if rc != 0: log.error(f"LLVM cmake FAILED — see {llvm_build_log}") - return ctx.copy_with( - build_passed=False, fix_errors=[str(llvm_build_log)] - ) + return ctx.copy_with(build_passed=False, fix_errors=[str(llvm_build_log)]) log.status(True, "cmake configure OK") # ── ninja build + install ── log.info(f"ninja -j{num_procs} install (this may take a while)...") with open(llvm_build_log, "a", encoding="utf-8") as fh: - fh.write(f"\n=== ninja install ===\n") + fh.write("\n=== ninja install ===\n") fh.flush() rc = stream_cmd( ["ninja", "-j", str(num_procs), "install"], - build_dir, fh, timeout=7200, label="ninja install", + build_dir, + fh, + timeout=7200, + label="ninja install", ) if rc != 0: log.error(f"LLVM ninja FAILED — see {llvm_build_log}") - return ctx.copy_with( - build_passed=False, fix_errors=[str(llvm_build_log)] - ) + return ctx.copy_with(build_passed=False, fix_errors=[str(llvm_build_log)]) log.status(True, "LLVM ninja install OK") # Copy FileCheck import shutil + fc = build_dir / "bin" / "FileCheck" if fc.exists(): shutil.copy2(fc, llvm_install / "bin" / "FileCheck") @@ -203,8 +214,10 @@ def build_llvm(ctx: WorkflowContext, num_procs: int = 32) -> WorkflowContext: def build_triton( - ctx: WorkflowContext, config: TAConfig, - clean: bool = False, python_exe: str = "", + ctx: WorkflowContext, + config: TAConfig, + clean: bool = False, + python_exe: str = "", ) -> WorkflowContext: """Build triton-ascend. Pure build — no retry logic.""" ascend_path = Path(ctx.triton_ascend_path) @@ -257,9 +270,7 @@ def build_triton( result = { "all_passed": passed, - "steps": [ - {"step": "setup_py_install", "passed": passed, "exit_code": rc} - ], + "steps": [{"step": "setup_py_install", "passed": passed, "exit_code": rc}], } (step_dir / BUILD_RESULT_FILE).write_text( json.dumps(result, indent=2, ensure_ascii=False) + "\n", encoding="utf-8" @@ -267,9 +278,7 @@ def build_triton( if not passed: log.error(f"Build FAILED — see {build_log}") - return ctx.copy_with( - build_passed=False, fix_errors=[str(build_log)] - ) + return ctx.copy_with(build_passed=False, fix_errors=[str(build_log)]) log.status(True, "Build passed") return ctx.copy_with(build_passed=True) @@ -318,7 +327,7 @@ def commit_fixes(ctx: WorkflowContext, config: TAConfig) -> None: # ── 5. Stage and commit ──────────────────────────────────────────── log.section("Commit AI Fixes") try: - staged_files = run_git(ascend_path, "diff", "--name-only", "HEAD").strip() + _ = run_git(ascend_path, "diff", "--name-only", "HEAD").strip() run_git(ascend_path, "add", "-A") staged = run_git(ascend_path, "diff", "--cached", "--name-only").strip() if staged: diff --git a/src/TA_main2main_workflow/pipeline/commit.py b/src/TA_main2main_workflow/pipeline/commit.py index 4c31c51..c4808c7 100644 --- a/src/TA_main2main_workflow/pipeline/commit.py +++ b/src/TA_main2main_workflow/pipeline/commit.py @@ -35,10 +35,7 @@ def commit_step(ctx: WorkflowContext, config: TAConfig) -> WorkflowContext: # ── 1. Submodule first ──────────────────────────────────────────── if submodule_has_changes(ascend_path): target_short = ctx.target_commit[:12] if ctx.target_commit else "HEAD" - commit_submodule( - ascend_path, - f"[Sync](fix) AI fix for {target_short}\n" - ) + commit_submodule(ascend_path, f"[Sync](fix) AI fix for {target_short}\n") # ── 2. Clean temp files ─────────────────────────────────────────── cleanup_temp_files(ascend_path) @@ -50,7 +47,7 @@ def commit_step(ctx: WorkflowContext, config: TAConfig) -> WorkflowContext: return ctx # Print staged files for visibility - staged_files = [l[3:] for l in staged.splitlines() if l.strip()] + staged_files = [line[3:] for line in staged.splitlines() if line.strip()] log.info(f"Files staged ({len(staged_files)}):") for f in staged_files[:30]: log.info(f" {f}") diff --git a/src/TA_main2main_workflow/pipeline/finalize.py b/src/TA_main2main_workflow/pipeline/finalize.py index 09eb2e8..f1ac728 100644 --- a/src/TA_main2main_workflow/pipeline/finalize.py +++ b/src/TA_main2main_workflow/pipeline/finalize.py @@ -76,7 +76,9 @@ def finalize(ctx: WorkflowContext, config: TAConfig | None = None) -> WorkflowCo def _generate_pr_description( - ctx: WorkflowContext, config: TAConfig, summary_path: Path, + ctx: WorkflowContext, + config: TAConfig, + summary_path: Path, ) -> None: """Invoke AI (report mode) to write the PR description.""" step = ctx.steps[-1] if ctx.steps else {"id": "step-0"} @@ -87,30 +89,32 @@ def _generate_pr_description( context = _build_report_context(ctx) log.section("AI PR Description (report mode)") - result = run_opencode_adapter({ - "step_id": "finalize-report", - "previous_step_id": "", - "previous_step_summary_path": "", - "is_last_step": "true", - "step_dir": str(step_dir), - "fix_dir": str(step_dir), - "conflict_dir": str(WORKSPACE_DIR / "conflicts"), - "ascend_path": str(Path(ctx.triton_ascend_path)), - "triton_path": ctx.triton_ascend_path, - "reference_dir": _REF, - "mode": "report", - "error_logs": json.dumps(context, ensure_ascii=False, default=str), - "target_commit": ctx.target_commit, - "step_index": f"{ctx.total_steps}/{ctx.total_steps}", - "upstream_commits_count": str(ctx.upstream_commits_count), - "total_steps": str(ctx.total_steps), - "conflict_files_resolved": str(ctx.conflict_files_resolved), - "build_fix_count": str(ctx.build_fix_count), - "test_fix_count": str(ctx.test_fix_count), - "final_status": ctx.final_status or "Success", - "ascend_npu_ir_fix": "false", - "ascend_npu_ir_compat_ref": "", - }) + result = run_opencode_adapter( + { + "step_id": "finalize-report", + "previous_step_id": "", + "previous_step_summary_path": "", + "is_last_step": "true", + "step_dir": str(step_dir), + "fix_dir": str(step_dir), + "conflict_dir": str(WORKSPACE_DIR / "conflicts"), + "ascend_path": str(Path(ctx.triton_ascend_path)), + "triton_path": ctx.triton_ascend_path, + "reference_dir": _REF, + "mode": "report", + "error_logs": json.dumps(context, ensure_ascii=False, default=str), + "target_commit": ctx.target_commit, + "step_index": f"{ctx.total_steps}/{ctx.total_steps}", + "upstream_commits_count": str(ctx.upstream_commits_count), + "total_steps": str(ctx.total_steps), + "conflict_files_resolved": str(ctx.conflict_files_resolved), + "build_fix_count": str(ctx.build_fix_count), + "test_fix_count": str(ctx.test_fix_count), + "final_status": ctx.final_status or "Success", + "ascend_npu_ir_fix": "false", + "ascend_npu_ir_compat_ref": "", + } + ) # AI writes to step_dir/step_summary.md; copy to final location ai_summary = step_dir / "step_summary.md" @@ -181,7 +185,7 @@ def _write_summary_fallback(ctx: WorkflowContext, summary_path: Path) -> None: f"- **Target**: `{ctx.target_commit[:12]}`", f"- **Steps**: {ctx.total_steps}", f"- **Upstream commits**: {ctx.upstream_commits_count}", - f"- **Status**: Success", + "- **Status**: Success", f"- **Date**: {time.strftime('%Y-%m-%d %H:%M:%S')}", f"- **Work branch**: `{ctx.work_branch}`", ] @@ -208,16 +212,16 @@ def _write_sync_report(ctx: WorkflowContext) -> None: report_path = WORKSPACE_DIR / "SYNC_REPORT.md" try: parts = [ - f"# Triton-Ascend 上游同步报告\n", - f"## 基本信息\n", + "# Triton-Ascend 上游同步报告\n", + "## 基本信息\n", f"- 目标提交: `{ctx.target_commit[:12]}`", f"- 步骤数: {ctx.total_steps}", f"- 上游提交数: {ctx.upstream_commits_count}", f"- 工作分支: `{ctx.work_branch}`", - f"- 状态: 成功", + "- 状态: 成功", ] if ctx.step_details: - parts.append(f"\n## 步骤详情\n") + parts.append("\n## 步骤详情\n") for d in ctx.step_details: parts.append( f"### {d['step_id']}\n" diff --git a/src/TA_main2main_workflow/pipeline/fix.py b/src/TA_main2main_workflow/pipeline/fix.py index dd39df8..007b156 100644 --- a/src/TA_main2main_workflow/pipeline/fix.py +++ b/src/TA_main2main_workflow/pipeline/fix.py @@ -16,8 +16,9 @@ _REF = str(Path(__file__).parent.parent / "reference") -def ai_fix(ctx: WorkflowContext, config: TAConfig, attempt: int = 1, - mode: str = "fix") -> WorkflowContext: +def ai_fix( + ctx: WorkflowContext, config: TAConfig, attempt: int = 1, mode: str = "fix" +) -> WorkflowContext: """Invoke AI to fix build or test failures. Args: @@ -50,10 +51,12 @@ def ai_fix(ctx: WorkflowContext, config: TAConfig, attempt: int = 1, prev_step_id = prev["id"] prev_summary = WORKSPACE_DIR / STEPS_DIR / prev_step_id / "step_summary.md" prev_summary_path = str(prev_summary) if prev_summary.exists() else "" - is_last_step = (ctx.current_step >= ctx.total_steps - 1) + is_last_step = ctx.current_step >= ctx.total_steps - 1 ascend_npu_ir_fix = _detect_ascend_npu_ir_errors(ascend_path, step_id) ascend_npu_ir_compat_ref = str( - Path(__file__).parent.parent / "reference" / "AscendNPU-IR_LLVM_VERSION_COMPAT.md" + Path(__file__).parent.parent + / "reference" + / "AscendNPU-IR_LLVM_VERSION_COMPAT.md" ) conflict_dir = str(WORKSPACE_DIR / "conflicts") @@ -84,9 +87,7 @@ def ai_fix(ctx: WorkflowContext, config: TAConfig, attempt: int = 1, ) # ── Fix validation gate ──────────────────────────────────────── - is_valid, reason = validate_fix( - ascend_path, pre_files, result.modified_files - ) + is_valid, reason = validate_fix(ascend_path, pre_files, result.modified_files) if not is_valid: log.warning(f"Fix validation FAILED: {reason}") log.warning("Reverting invalid changes...") @@ -169,8 +170,12 @@ def _detect_ascend_npu_ir_errors(ascend_path: Path, step_id: str) -> bool: try: content = build_log.read_text(encoding="utf-8", errors="replace").lower() indicators = [ - "AscendNPU-IR".lower(), "ascendnpu-ir", - "llvm::", "mlir::", "fatal error", "undefined reference", + "AscendNPU-IR".lower(), + "ascendnpu-ir", + "llvm::", + "mlir::", + "fatal error", + "undefined reference", ] return any(ind in content for ind in indicators) except Exception: diff --git a/src/TA_main2main_workflow/pipeline/ir_patch.py b/src/TA_main2main_workflow/pipeline/ir_patch.py index 8b48dc2..d359c49 100644 --- a/src/TA_main2main_workflow/pipeline/ir_patch.py +++ b/src/TA_main2main_workflow/pipeline/ir_patch.py @@ -28,14 +28,20 @@ from TA_main2main_workflow.utils.git import run_git, run_git_no_check, stream_cmd from TA_main2main_workflow.pipeline.build import build_and_fix_loop from TA_main2main_workflow.pipeline.test import ( - run_tests, detect_oom_in_tests, rerun_tests_reduced_concurrency, - test_and_fix_loop, _run_pretest_and_fix, + run_tests, + detect_oom_in_tests, + rerun_tests_reduced_concurrency, + test_and_fix_loop, + _run_pretest_and_fix, ) -from TA_main2main_workflow.pipeline.fix import ai_fix from TA_main2main_workflow.utils import ( - WORKSPACE_DIR, STEPS_DIR, BUILD_RESULT_FILE, TEST_RESULT_FILE, - IR_ANALYSIS_DIR, IR_OPS_REPORT_FILE, IR_CHANGES_REPORT_FILE, - IR_DIAGNOSIS_FILE, IR_MAX_ITERATIONS, LLVM_CHANGE_ANALYSIS_DIR, + WORKSPACE_DIR, + STEPS_DIR, + TEST_RESULT_FILE, + IR_ANALYSIS_DIR, + IR_OPS_REPORT_FILE, + IR_CHANGES_REPORT_FILE, + IR_DIAGNOSIS_FILE, _ASCEND_BASELINE_LLVM_HASH, ) @@ -82,7 +88,9 @@ def build_baseline_llvm(ctx: WorkflowContext, config: TAConfig) -> WorkflowConte # Allow skipping baseline LLVM build (LLVM already built for current TA) if config.skip_baseline_llvm: - log.status(True, "SKIP_BASELINE_LLVM set — assuming baseline LLVM already built") + log.status( + True, "SKIP_BASELINE_LLVM set — assuming baseline LLVM already built" + ) return ctx.copy_with(build_passed=True) # Ensure llvm-project exists @@ -133,8 +141,9 @@ def build_baseline_llvm(ctx: WorkflowContext, config: TAConfig) -> WorkflowConte # ═══════════════════════════════════════════════════════════════════════════ -def per_step_ir_patch(ctx: WorkflowContext, config: TAConfig, - step: dict) -> WorkflowContext: +def per_step_ir_patch( + ctx: WorkflowContext, config: TAConfig, step: dict +) -> WorkflowContext: """Full per-step IR patch pipeline for LLVM version changes. Strategy: apply existing patch first → build → test → AI supplement. @@ -180,7 +189,10 @@ def per_step_ir_patch(ctx: WorkflowContext, config: TAConfig, def _do_apply_existing_patch( - ctx: WorkflowContext, config: TAConfig, step: dict, target_llvm_hash: str, + ctx: WorkflowContext, + config: TAConfig, + step: dict, + target_llvm_hash: str, ) -> bool: """Apply existing Ascend LLVM patch to llvm-project, with AI fix retry. @@ -189,7 +201,7 @@ def _do_apply_existing_patch( ascend_path = Path(ctx.triton_ascend_path) llvm_project = config.llvm_project llvm_install = config.llvm_install - step_id = step["id"] + _ = step["id"] patch_dir = ascend_path / "third_party/ascend/patch" patch_files = sorted(patch_dir.glob("*.patch")) if patch_dir.exists() else [] @@ -233,7 +245,10 @@ def _do_apply_existing_patch( log.info(f"Applying existing patch Successfully (attempt {attempt})") # ── Build LLVM ── if config.skip_llvm_rebuild: - log.status(True, f"SKIP_LLVM_REBUILD set — assuming LLVM already built (attempt {attempt})") + log.status( + True, + f"SKIP_LLVM_REBUILD set — assuming LLVM already built (attempt {attempt})", + ) return True try: _do_llvm_build(llvm_project, llvm_install, target_llvm_hash) @@ -257,7 +272,9 @@ def _do_apply_existing_patch( def _do_ta_build_with_fix( - ctx: WorkflowContext, config: TAConfig, step: dict, + ctx: WorkflowContext, + config: TAConfig, + step: dict, ) -> WorkflowContext: """Build Triton-Ascend with AI fix loop for compile errors.""" return build_and_fix_loop(ctx, config) @@ -269,7 +286,10 @@ def _do_ta_build_with_fix( def _do_test_and_fix_with_ir_retry( - ctx: WorkflowContext, config: TAConfig, step: dict, target_llvm_hash: str, + ctx: WorkflowContext, + config: TAConfig, + step: dict, + target_llvm_hash: str, ) -> WorkflowContext: """Test with IR supplement loop. @@ -277,7 +297,7 @@ def _do_test_and_fix_with_ir_retry( rebuild LLVM → rebuild TA → retest. Code issues → AI fix loop. Max 3 IR supplement iterations. """ - step_id = step["id"] + _ = step["id"] ascend_path = Path(ctx.triton_ascend_path) llvm_project = config.llvm_project llvm_install = config.llvm_install @@ -314,17 +334,25 @@ def _do_test_and_fix_with_ir_retry( # ── Classify failures: IR vs code ── is_ir_issue = _classify_test_failures(ctx, config, step) if is_ir_issue: - log.info(f"IR issues detected — supplementing patch (iter {ir_iter + 1}/{ir_max})") - _ir_supplement_patch(ctx, config, step, target_llvm_hash, - supplement_iter=ir_iter + 1) + log.info( + f"IR issues detected — supplementing patch (iter {ir_iter + 1}/{ir_max})" + ) + _ir_supplement_patch( + ctx, config, step, target_llvm_hash, supplement_iter=ir_iter + 1 + ) # Rebuild LLVM with updated patch if config.skip_llvm_rebuild: - log.status(True, "SKIP_LLVM_REBUILD set — skipping LLVM rebuild after supplement") + log.status( + True, + "SKIP_LLVM_REBUILD set — skipping LLVM rebuild after supplement", + ) else: try: # Re-apply updated patch to clean workspace before building _clean_checkout_apply_patch( - llvm_project, ascend_path, target_llvm_hash, + llvm_project, + ascend_path, + target_llvm_hash, reason=f"IR supplement iter {ir_iter + 1}", ) _do_llvm_build(llvm_project, llvm_install, target_llvm_hash) @@ -350,7 +378,10 @@ def _do_test_and_fix_with_ir_retry( def _per_step_ir_patch_fallback( - ctx: WorkflowContext, config: TAConfig, step: dict, target_llvm_hash: str, + ctx: WorkflowContext, + config: TAConfig, + step: dict, + target_llvm_hash: str, ) -> WorkflowContext: """Fallback: Full OP analysis pipeline when existing-patch-first fails. @@ -393,7 +424,9 @@ def _per_step_ir_patch_fallback( else: try: _clean_checkout_apply_patch( - llvm_project, ascend_path, target_llvm_hash, + llvm_project, + ascend_path, + target_llvm_hash, reason="fallback IR pipeline", ) _do_llvm_build(llvm_project, llvm_install, target_llvm_hash) @@ -416,9 +449,15 @@ def _per_step_ir_patch_fallback( # ═══════════════════════════════════════════════════════════════════════════ -def _ir_ai_base(ctx: WorkflowContext, config: TAConfig, step_id: str, - ir_dir: Path, mode: str, error_logs: str = "[]", - extra: dict | None = None) -> dict: +def _ir_ai_base( + ctx: WorkflowContext, + config: TAConfig, + step_id: str, + ir_dir: Path, + mode: str, + error_logs: str = "[]", + extra: dict | None = None, +) -> dict: """Build the common AI context dict for IR patch calls.""" ascend_path = Path(ctx.triton_ascend_path) base = { @@ -450,53 +489,71 @@ def _run_ir_op_analysis(ctx: WorkflowContext, config: TAConfig) -> dict: ir_dir = WORKSPACE_DIR / STEPS_DIR / step_id / IR_ANALYSIS_DIR ir_dir.mkdir(parents=True, exist_ok=True) - result = run_opencode_adapter(_ir_ai_base(ctx, config, step_id, ir_dir, - "ir_op_analysis")) + result = run_opencode_adapter( + _ir_ai_base(ctx, config, step_id, ir_dir, "ir_op_analysis") + ) ops_file = ir_dir / IR_OPS_REPORT_FILE - if hasattr(result, 'step_summary') and result.step_summary: + if hasattr(result, "step_summary") and result.step_summary: try: ops_data = json.loads(result.step_summary) - ops_file.write_text(json.dumps(ops_data, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") + ops_file.write_text( + json.dumps(ops_data, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) return ops_data except json.JSONDecodeError: pass return {} -def _run_ir_change_analysis(ctx: WorkflowContext, config: TAConfig, - target_llvm_hash: str) -> dict: +def _run_ir_change_analysis( + ctx: WorkflowContext, config: TAConfig, target_llvm_hash: str +) -> dict: """AI compares OP .td definitions between baseline and target LLVM.""" step = ctx.steps[ctx.current_step] if ctx.steps else {"id": "step-0"} step_id = step["id"] ir_dir = WORKSPACE_DIR / STEPS_DIR / step_id / IR_ANALYSIS_DIR ir_dir.mkdir(parents=True, exist_ok=True) - result = run_opencode_adapter(_ir_ai_base(ctx, config, step_id, ir_dir, - "ir_change_analysis", - error_logs=json.dumps({ - "baseline_llvm_hash": _ASCEND_BASELINE_LLVM_HASH, - "target_llvm_hash": target_llvm_hash, - }, ensure_ascii=False), - extra={ - "baseline_llvm_hash": _ASCEND_BASELINE_LLVM_HASH, - "target_llvm_hash": target_llvm_hash, - }, - )) + result = run_opencode_adapter( + _ir_ai_base( + ctx, + config, + step_id, + ir_dir, + "ir_change_analysis", + error_logs=json.dumps( + { + "baseline_llvm_hash": _ASCEND_BASELINE_LLVM_HASH, + "target_llvm_hash": target_llvm_hash, + }, + ensure_ascii=False, + ), + extra={ + "baseline_llvm_hash": _ASCEND_BASELINE_LLVM_HASH, + "target_llvm_hash": target_llvm_hash, + }, + ) + ) changes_file = ir_dir / IR_CHANGES_REPORT_FILE - if hasattr(result, 'step_summary') and result.step_summary: + if hasattr(result, "step_summary") and result.step_summary: try: changes_data = json.loads(result.step_summary) - changes_file.write_text(json.dumps(changes_data, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") + changes_file.write_text( + json.dumps(changes_data, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) return changes_data except json.JSONDecodeError: pass return {} -def _run_ir_generate_patches(ctx: WorkflowContext, config: TAConfig, - step: dict, target_llvm_hash: str = "") -> None: +def _run_ir_generate_patches( + ctx: WorkflowContext, config: TAConfig, step: dict, target_llvm_hash: str = "" +) -> None: """AI modifies the existing patch file in-place for new LLVM version.""" step_id = step["id"] ir_dir = WORKSPACE_DIR / STEPS_DIR / step_id / IR_ANALYSIS_DIR @@ -507,18 +564,28 @@ def _run_ir_generate_patches(ctx: WorkflowContext, config: TAConfig, patch_files = sorted(patch_dir.glob("*.patch")) if patch_dir.exists() else [] ascend_patch_file = str(patch_files[0]) if patch_files else "" - run_opencode_adapter(_ir_ai_base(ctx, config, step_id, ir_dir, - "ir_patch_gen", - error_logs=json.dumps({ - "ops_report": ctx.ir_ops_report, - "changes_report": ctx.ir_changes_report, - }, ensure_ascii=False, default=str), - extra={ - "baseline_llvm_hash": _ASCEND_BASELINE_LLVM_HASH, - "target_llvm_hash": target_llvm_hash, - "ascend_patch_file": ascend_patch_file, - }, - )) + run_opencode_adapter( + _ir_ai_base( + ctx, + config, + step_id, + ir_dir, + "ir_patch_gen", + error_logs=json.dumps( + { + "ops_report": ctx.ir_ops_report, + "changes_report": ctx.ir_changes_report, + }, + ensure_ascii=False, + default=str, + ), + extra={ + "baseline_llvm_hash": _ASCEND_BASELINE_LLVM_HASH, + "target_llvm_hash": target_llvm_hash, + "ascend_patch_file": ascend_patch_file, + }, + ) + ) # ═══════════════════════════════════════════════════════════════════════════ @@ -527,7 +594,10 @@ def _run_ir_generate_patches(ctx: WorkflowContext, config: TAConfig, def _ai_adjust_patch_for_failure( - ctx: WorkflowContext, config: TAConfig, step: dict, error_info: str, + ctx: WorkflowContext, + config: TAConfig, + step: dict, + error_info: str, ) -> None: """AI adjusts the LLVM patch after build failure.""" step_id = step["id"] @@ -540,22 +610,30 @@ def _ai_adjust_patch_for_failure( ascend_patch_file = str(patch_files[0]) if patch_files else "" target_llvm_hash = _get_current_llvm_hash(ascend_path) - run_opencode_adapter(_ir_ai_base(ctx, config, step_id, ir_dir, - "ir_patch_fix", - error_logs=json.dumps({"error": error_info[:5000]}, ensure_ascii=False), - extra={ - "baseline_llvm_hash": _ASCEND_BASELINE_LLVM_HASH, - "target_llvm_hash": target_llvm_hash, - "ascend_patch_file": ascend_patch_file, - "adjust_mode": "patch_apply_failure", - "patch_error_type": "apply_or_build", - "patch_error_msg": error_info[:2000], - }, - )) + run_opencode_adapter( + _ir_ai_base( + ctx, + config, + step_id, + ir_dir, + "ir_patch_fix", + error_logs=json.dumps({"error": error_info[:5000]}, ensure_ascii=False), + extra={ + "baseline_llvm_hash": _ASCEND_BASELINE_LLVM_HASH, + "target_llvm_hash": target_llvm_hash, + "ascend_patch_file": ascend_patch_file, + "adjust_mode": "patch_apply_failure", + "patch_error_type": "apply_or_build", + "patch_error_msg": error_info[:2000], + }, + ) + ) def _build_focused_change_report( - ctx: WorkflowContext, config: TAConfig, step: dict, + ctx: WorkflowContext, + config: TAConfig, + step: dict, target_llvm_hash: str, ) -> Path | None: """Analyze OP definition diffs for affected OPs identified by ir_diagnose. @@ -585,11 +663,13 @@ def _build_focused_change_report( if f.get("classification") == "ir_compatibility": op_name = f.get("affected_op", "").strip() if op_name and op_name not in {o["name"] for o in affected_ops}: - affected_ops.append({ - "name": op_name, - "error_summary": f.get("error_summary", ""), - "rationale": f.get("rationale", ""), - }) + affected_ops.append( + { + "name": op_name, + "error_summary": f.get("error_summary", ""), + "rationale": f.get("rationale", ""), + } + ) if not affected_ops: log.info("No ir_compatibility OPs in diagnosis — nothing to analyze") return None @@ -617,7 +697,9 @@ def _build_focused_change_report( # Only diff each .td file once (multiple OPs in same file) if td_relative not in seen_td_files: seen_td_files.add(td_relative) - diff = _diff_td_file(llvm_project, baseline_hash, target_llvm_hash, td_relative) + diff = _diff_td_file( + llvm_project, baseline_hash, target_llvm_hash, td_relative + ) entry["td_diff"] = diff[:8000] if diff else "(no diff)" if diff and len(diff) > 8000: entry["td_diff_truncated"] = True @@ -673,8 +755,8 @@ def _diagnosis_candidates(step_id: str, ir_dir: Path) -> list[Path]: """Candidate paths for IR diagnosis, in priority order.""" step_dir = WORKSPACE_DIR / STEPS_DIR / step_id return [ - step_dir / "ir_diagnosis.json", # where AI writes per prompt - ir_dir / IR_DIAGNOSIS_FILE, # where classify writes parsed + step_dir / "ir_diagnosis.json", # where AI writes per prompt + ir_dir / IR_DIAGNOSIS_FILE, # where classify writes parsed ] @@ -689,7 +771,9 @@ def _find_td_file(llvm_project: Path, target_hash: str, op_name: str) -> str | N result = subprocess.run( ["git", "grep", "-l", f"def {short_name}", target_hash, "--", "*.td"], cwd=str(llvm_project), - capture_output=True, text=True, timeout=30, + capture_output=True, + text=True, + timeout=30, ) except (subprocess.TimeoutExpired, OSError): return None @@ -699,14 +783,17 @@ def _find_td_file(llvm_project: Path, target_hash: str, op_name: str) -> str | N return result.stdout.strip().split("\n")[0] -def _diff_td_file(llvm_project: Path, baseline_hash: str, target_hash: str, - td_relative: str) -> str: +def _diff_td_file( + llvm_project: Path, baseline_hash: str, target_hash: str, td_relative: str +) -> str: """Get the diff of a .td file between baseline and target LLVM.""" try: result = subprocess.run( ["git", "diff", f"{baseline_hash}..{target_hash}", "--", td_relative], cwd=str(llvm_project), - capture_output=True, text=True, timeout=60, + capture_output=True, + text=True, + timeout=60, ) except (subprocess.TimeoutExpired, OSError): return "" @@ -714,7 +801,10 @@ def _diff_td_file(llvm_project: Path, baseline_hash: str, target_hash: str, def _ir_supplement_patch( - ctx: WorkflowContext, config: TAConfig, step: dict, target_llvm_hash: str, + ctx: WorkflowContext, + config: TAConfig, + step: dict, + target_llvm_hash: str, supplement_iter: int = 1, ) -> None: """AI supplements the existing IR patch with missing OP IR changes. @@ -752,7 +842,10 @@ def _ir_supplement_patch( # Extract affected OPs from diagnosis, diff their .td definitions # between baseline and target LLVM, write to focused_changes.json focused_report_path = _build_focused_change_report( - ctx, config, step, target_llvm_hash, + ctx, + config, + step, + target_llvm_hash, ) if focused_report_path: error_log_paths.append(str(focused_report_path)) @@ -769,7 +862,9 @@ def _ir_supplement_patch( except Exception: pass - log.key_value("Existing patch", str(ascend_patch_file) if ascend_patch_file else "(none)") + log.key_value( + "Existing patch", str(ascend_patch_file) if ascend_patch_file else "(none)" + ) log.key_value("Target LLVM", target_llvm_hash[:12]) log.key_value("Test error logs", str(len(error_log_paths))) if focused_report_path: @@ -777,24 +872,32 @@ def _ir_supplement_patch( if diagnosis_path: log.key_value("Diagnosis", str(diagnosis_path)) - run_opencode_adapter(_ir_ai_base(ctx, config, step_id, ir_dir, - "ir_generate_patch", - error_logs=json.dumps(error_log_paths, ensure_ascii=False), - extra={ - "previous_step_id": "ir-diagnose", - "previous_step_summary_path": str(diagnosis_path or ""), - "focused_changes_path": str(focused_report_path or ""), - "target_llvm_hash": target_llvm_hash, - "baseline_llvm_hash": _ASCEND_BASELINE_LLVM_HASH, - "ascend_patch_file": ascend_patch_file, - "patch_content_snippet": patch_content_snippet, - "adjust_mode": "supplement", - "supplement_iteration": str(supplement_iter), - "ascend_npu_ir_compat_ref": str( - Path(__file__).parent.parent / "reference" - / "AscendNPU-IR_LLVM_VERSION_COMPAT.md"), - }, - )) + run_opencode_adapter( + _ir_ai_base( + ctx, + config, + step_id, + ir_dir, + "ir_generate_patch", + error_logs=json.dumps(error_log_paths, ensure_ascii=False), + extra={ + "previous_step_id": "ir-diagnose", + "previous_step_summary_path": str(diagnosis_path or ""), + "focused_changes_path": str(focused_report_path or ""), + "target_llvm_hash": target_llvm_hash, + "baseline_llvm_hash": _ASCEND_BASELINE_LLVM_HASH, + "ascend_patch_file": ascend_patch_file, + "patch_content_snippet": patch_content_snippet, + "adjust_mode": "supplement", + "supplement_iteration": str(supplement_iter), + "ascend_npu_ir_compat_ref": str( + Path(__file__).parent.parent + / "reference" + / "AscendNPU-IR_LLVM_VERSION_COMPAT.md" + ), + }, + ) + ) def _collect_test_error_logs() -> list[str]: @@ -824,7 +927,9 @@ def _collect_test_error_logs() -> list[str]: def _classify_test_failures( - ctx: WorkflowContext, config: TAConfig, step: dict, + ctx: WorkflowContext, + config: TAConfig, + step: dict, ) -> bool: """AI classifies test failures: IR compatibility vs code issues. @@ -852,10 +957,16 @@ def _classify_test_failures( log.info(f" ... and {len(error_log_paths) - 5} more") try: - result = run_opencode_adapter(_ir_ai_base(ctx, config, step_id, ir_dir, - "ir_diagnose", - error_logs=json.dumps(error_log_paths, ensure_ascii=False), - )) + result = run_opencode_adapter( + _ir_ai_base( + ctx, + config, + step_id, + ir_dir, + "ir_diagnose", + error_logs=json.dumps(error_log_paths, ensure_ascii=False), + ) + ) except Exception as e: log.error(f"IR diagnosis failed: {e}") return True # Default to IR issue on failure @@ -866,7 +977,10 @@ def _classify_test_failures( try: diagnosis_data = json.loads(summary) except json.JSONDecodeError: - diagnosis_data = {"summary": summary, "has_ir_issues": "ir_issue" in summary.lower()} + diagnosis_data = { + "summary": summary, + "has_ir_issues": "ir_issue" in summary.lower(), + } diagnosis_path.write_text( json.dumps(diagnosis_data, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", @@ -878,7 +992,9 @@ def _classify_test_failures( def _do_ai_fix_loop( - ctx: WorkflowContext, config: TAConfig, step: dict, + ctx: WorkflowContext, + config: TAConfig, + step: dict, ) -> WorkflowContext: """Standard AI fix loop for code issues (not IR-related).""" return test_and_fix_loop(ctx, config) @@ -929,8 +1045,10 @@ def _ensure_commit_available(llvm_project: Path, commit_hash: str) -> None: if result.returncode == 0: return if attempt < max_attempts: - log.info(f"Commit {commit_hash[:12]} not found locally — fetching " - f"(attempt {attempt}/{max_attempts})...") + log.info( + f"Commit {commit_hash[:12]} not found locally — fetching " + f"(attempt {attempt}/{max_attempts})..." + ) try: run_git(llvm_project, "fetch", "origin", commit_hash) except Exception: @@ -941,7 +1059,10 @@ def _ensure_commit_available(llvm_project: Path, commit_hash: str) -> None: def _clean_checkout_apply_patch( - llvm_project: Path, ascend_path: Path, target_llvm_hash: str, reason: str = "", + llvm_project: Path, + ascend_path: Path, + target_llvm_hash: str, + reason: str = "", ) -> bool: """Clean workspace, checkout target hash, apply the current Ascend patch. @@ -974,9 +1095,7 @@ def _get_current_llvm_hash(ascend_path: Path) -> str: return "" - -def _do_llvm_build(llvm_project: Path, llvm_install: Path, - required_hash: str) -> str: +def _do_llvm_build(llvm_project: Path, llvm_install: Path, required_hash: str) -> str: """Build and install LLVM from current working tree state. Cleans build directory, runs cmake + ninja install, copies FileCheck. @@ -1000,8 +1119,10 @@ def _do_llvm_build(llvm_project: Path, llvm_install: Path, # ── cmake configure ────────────────────────────────────────────── cmake_cmd = [ - "cmake", str(llvm_project / "llvm"), - "-G", "Ninja", + "cmake", + str(llvm_project / "llvm"), + "-G", + "Ninja", "-DCMAKE_BUILD_TYPE=Release", "-DLLVM_ENABLE_ASSERTIONS=ON", "-DLLVM_ENABLE_PROJECTS=mlir;llvm;lld", @@ -1014,8 +1135,9 @@ def _do_llvm_build(llvm_project: Path, llvm_install: Path, with open(llvm_build_log, "w", encoding="utf-8") as fh: fh.write(f"=== cmake ===\n{' '.join(cmake_cmd)}\n\n") fh.flush() - rc = stream_cmd(cmake_cmd, build_dir, fh, timeout=300, - label="Configuring LLVM with cmake") + rc = stream_cmd( + cmake_cmd, build_dir, fh, timeout=300, label="Configuring LLVM with cmake" + ) if rc != 0: raise RuntimeError( f"LLVM cmake configure failed (exit {rc}). See: {llvm_build_log}" @@ -1025,10 +1147,11 @@ def _do_llvm_build(llvm_project: Path, llvm_install: Path, # ── ninja build + install ─────────────────────────────────────── log.info("ninja install (this may take a while)...") with open(llvm_build_log, "a", encoding="utf-8") as fh: - fh.write(f"\n=== ninja install ===\n") + fh.write("\n=== ninja install ===\n") fh.flush() - rc = stream_cmd(["ninja", "install"], build_dir, fh, timeout=7200, - label="ninja install") + rc = stream_cmd( + ["ninja", "install"], build_dir, fh, timeout=7200, label="ninja install" + ) if rc != 0: raise RuntimeError( f"LLVM ninja build failed (exit {rc}). See: {llvm_build_log}" @@ -1060,9 +1183,12 @@ def _detect_ascend_npu_ir_errors(ctx: WorkflowContext) -> bool: try: content = build_log.read_text(encoding="utf-8", errors="replace").lower() indicators = [ - "AscendNPU-IR".lower(), "ascendnpu-ir", - "llvm::", "mlir::", - "fatal error", "undefined reference", + "AscendNPU-IR".lower(), + "ascendnpu-ir", + "llvm::", + "mlir::", + "fatal error", + "undefined reference", ] return any(ind in content for ind in indicators) except Exception: diff --git a/src/TA_main2main_workflow/pipeline/merge.py b/src/TA_main2main_workflow/pipeline/merge.py index 71f658e..dd1c423 100644 --- a/src/TA_main2main_workflow/pipeline/merge.py +++ b/src/TA_main2main_workflow/pipeline/merge.py @@ -15,7 +15,9 @@ from TA_main2main_workflow.utils.logging import get_logger from TA_main2main_workflow.utils.git import run_git, run_git_no_check from TA_main2main_workflow.utils import ( - WORKSPACE_DIR, STEPS_DIR, get_base_branch_ref, + WORKSPACE_DIR, + STEPS_DIR, + get_base_branch_ref, ) log = get_logger(__name__) @@ -59,7 +61,9 @@ def merge_upstream_commit(ctx: WorkflowContext, config: TAConfig) -> WorkflowCon # Ensure we're on the work branch current_branch = run_git(ascend_path, "branch", "--show-current").strip() if ctx.work_branch and current_branch != ctx.work_branch: - log.warning(f"Expected '{ctx.work_branch}' but on '{current_branch}' — switching") + log.warning( + f"Expected '{ctx.work_branch}' but on '{current_branch}' — switching" + ) run_git(ascend_path, "checkout", ctx.work_branch) # ── Do the merge ──────────────────────────────────────────────── @@ -68,11 +72,11 @@ def merge_upstream_commit(ctx: WorkflowContext, config: TAConfig) -> WorkflowCon ascend_path, "merge", "--no-ff", "--no-edit", step["end_commit"] ) - conflict_files = run_git( + conflict_raw = run_git( ascend_path, "diff", "--name-only", "--diff-filter=U" ).strip() - conflict_files = ( - [f for f in conflict_files.splitlines() if f] if conflict_files else [] + conflict_files: list[str] = ( + [f for f in conflict_raw.splitlines() if f] if conflict_raw else [] ) has_conflicts = len(conflict_files) > 0 @@ -89,8 +93,10 @@ def merge_upstream_commit(ctx: WorkflowContext, config: TAConfig) -> WorkflowCon if has_conflicts: log.conflict_list(conflict_files) elif merge_proc.returncode != 0: - log.warning(f"Merge exited with code {merge_proc.returncode} " - f"but no conflict markers found — continuing") + log.warning( + f"Merge exited with code {merge_proc.returncode} " + f"but no conflict markers found — continuing" + ) else: log.key_value("merge exit code", str(merge_proc.returncode)) log.key_value("conflict files", "0") @@ -135,7 +141,9 @@ def _create_work_branch(repo: Path, config: TAConfig) -> None: try: run_git(repo, "fetch", config.work_branch_base) except Exception: - log.warning(f"Could not fetch remote '{config.work_branch_base}' — using origin") + log.warning( + f"Could not fetch remote '{config.work_branch_base}' — using origin" + ) # ── 4. Checkout base ref and create work branch ──────────────── try: diff --git a/src/TA_main2main_workflow/pipeline/plan.py b/src/TA_main2main_workflow/pipeline/plan.py index f37bed6..7bb2fe9 100644 --- a/src/TA_main2main_workflow/pipeline/plan.py +++ b/src/TA_main2main_workflow/pipeline/plan.py @@ -13,7 +13,11 @@ from TA_main2main_workflow.utils.config import TAConfig from TA_main2main_workflow.utils.context import WorkflowContext from TA_main2main_workflow.utils import ( - WORKSPACE_DIR, STEPS_FILE, STEPS_DIR, LLVM_HASH_FILE, SOURCE_DIRS, + WORKSPACE_DIR, + STEPS_FILE, + STEPS_DIR, + LLVM_HASH_FILE, + SOURCE_DIRS, ) from TA_main2main_workflow.utils.git import run_git from TA_main2main_workflow.utils.logging import get_logger @@ -54,7 +58,9 @@ def plan_steps(ctx: WorkflowContext, config: TAConfig) -> WorkflowContext: log.info(f"[plan] Line budget: {line_budget} (no commit-count limit)") if config.progressive_merge: - lines_per_commit, llvm_commits, source_touching = _scan_commits(triton_path, commits) + lines_per_commit, llvm_commits, source_touching = _scan_commits( + triton_path, commits + ) steps = _plan_steps_inner( commits, lines_per_commit, base, line_budget, llvm_commits ) @@ -67,7 +73,9 @@ def plan_steps(ctx: WorkflowContext, config: TAConfig) -> WorkflowContext: reason_tag = " [LLVM VERSION]" elif s.get("reason") == "oversized": reason_tag = " [OVERSIZED]" - budget_label = "OVERSIZED" if s["source_changed_lines"] > line_budget else "OK" + budget_label = ( + "OVERSIZED" if s["source_changed_lines"] > line_budget else "OK" + ) log.info( f" {s['id']}: {s['commit_count']} commits, " f"{s['source_changed_lines']} lines ({budget_label})" @@ -130,8 +138,7 @@ def llvm_hash_changed_after_merge(ctx: WorkflowContext) -> bool: if ctx.step_start_ascend_head: try: old_content = run_git( - ascend_path, "show", - f"{ctx.step_start_ascend_head}:{LLVM_HASH_FILE}" + ascend_path, "show", f"{ctx.step_start_ascend_head}:{LLVM_HASH_FILE}" ).strip() return old_content != current_hash except Exception: @@ -158,8 +165,14 @@ def _source_lines_for_commit(repo: Path, sha: str) -> int: for d in SOURCE_DIRS: try: output = run_git( - repo, "diff-tree", "--no-commit-id", "-r", "--numstat", - sha, "--", f":(top){d}", + repo, + "diff-tree", + "--no-commit-id", + "-r", + "--numstat", + sha, + "--", + f":(top){d}", ) except Exception: continue @@ -177,8 +190,7 @@ def _source_lines_for_commit(repo: Path, sha: str) -> int: def _commit_changed_llvm_hash(repo: Path, sha: str) -> bool: try: - output = run_git( - repo, "diff-tree", "--no-commit-id", "--name-only", "-r", sha) + output = run_git(repo, "diff-tree", "--no-commit-id", "--name-only", "-r", sha) return LLVM_HASH_FILE in output except Exception: return False @@ -197,16 +209,22 @@ def _scan_commits( source_touching += 1 if _commit_changed_llvm_hash(repo, c["sha"]): llvm_commits.add(c["sha"]) - log.info(f"[plan] LLVM version change detected: {c['sha'][:8]} {c['subject'][:80]}") + log.info( + f"[plan] LLVM version change detected: {c['sha'][:8]} {c['subject'][:80]}" + ) if (i + 1) % 50 == 0: log.info(f"[plan] ... scanned {i + 1}/{len(commits)} commits") # Print zero-lines summary (matching pre-refactor) if source_touching < len(commits): - log.info(f"[plan] {len(commits) - source_touching} commits touch zero " - f"source lines — included in steps with 0 line contribution") + log.info( + f"[plan] {len(commits) - source_touching} commits touch zero " + f"source lines — included in steps with 0 line contribution" + ) if llvm_commits: - log.info(f"[plan] {len(llvm_commits)} commit(s) changed LLVM hash " - f"— each will be a solo merge step") + log.info( + f"[plan] {len(llvm_commits)} commit(s) changed LLVM hash " + f"— each will be a solo merge step" + ) return lines_per_commit, llvm_commits, source_touching @@ -257,7 +275,12 @@ def _plan_steps_inner( step_commits, step_lines = [], 0 steps.append( _make_step( - len(steps) + 1, [commit], start, lines, budget, reason="llvm_version" + len(steps) + 1, + [commit], + start, + lines, + budget, + reason="llvm_version", ) ) start = steps[-1]["end_commit"] @@ -272,7 +295,9 @@ def _plan_steps_inner( start = steps[-1]["end_commit"] step_commits, step_lines = [], 0 steps.append( - _make_step(len(steps) + 1, [commit], start, lines, budget, reason="oversized") + _make_step( + len(steps) + 1, [commit], start, lines, budget, reason="oversized" + ) ) start = steps[-1]["end_commit"] continue @@ -301,14 +326,19 @@ def _enrich_steps(repo: Path, steps: list[dict[str, Any]]) -> None: # Build pathspec args for SOURCE_DIRS filtering pathspecs = [f":(top){d}" for d in SOURCE_DIRS] step["upstream_patch"] = run_git( - repo, "diff", + repo, + "diff", f"{step['start_commit']}..{step['end_commit']}", - "--", *pathspecs, + "--", + *pathspecs, ) step["changed_files"] = run_git( - repo, "diff", "--name-only", + repo, + "diff", + "--name-only", f"{step['start_commit']}..{step['end_commit']}", - "--", *pathspecs, + "--", + *pathspecs, ) step["files_changed"] = sorted( f for f in step["changed_files"].strip().splitlines() if f diff --git a/src/TA_main2main_workflow/pipeline/push_pr.py b/src/TA_main2main_workflow/pipeline/push_pr.py index e535baa..cb38cea 100644 --- a/src/TA_main2main_workflow/pipeline/push_pr.py +++ b/src/TA_main2main_workflow/pipeline/push_pr.py @@ -17,7 +17,7 @@ from TA_main2main_workflow.utils.logging import get_logger from TA_main2main_workflow.utils.git import run_git, run_git_no_check from TA_main2main_workflow.utils.submodule import push_submodule -from TA_main2main_workflow.utils import WORKSPACE_DIR, FINAL_TARGET_PATCH_FILE, FINAL_SUMMARY_FILE +from TA_main2main_workflow.utils import WORKSPACE_DIR, FINAL_SUMMARY_FILE log = get_logger(__name__) @@ -71,8 +71,13 @@ def push_and_create_pr( # ── 6. Create PR from fork → upstream ────────────────────────── pr_url = _create_pr( - ascend_path, github_repo, branch, fork_owner, token, - summary_file, target_commit, + ascend_path, + github_repo, + branch, + fork_owner, + token, + summary_file, + target_commit, ) return pr_url @@ -99,7 +104,9 @@ def _ensure_gh_auth(repo: Path) -> None: try: subprocess.run( ["gh", "auth", "status"], - check=True, capture_output=True, text=True, + check=True, + capture_output=True, + text=True, ) log.info("gh CLI already authenticated") except subprocess.CalledProcessError: @@ -107,7 +114,9 @@ def _ensure_gh_auth(repo: Path) -> None: try: subprocess.run( ["gh", "auth", "setup-git"], - check=True, capture_output=True, text=True, + check=True, + capture_output=True, + text=True, ) except Exception: pass @@ -121,7 +130,9 @@ def _ensure_gh_auth(repo: Path) -> None: try: subprocess.run( ["gh", "auth", "login", "--with-token", "--hostname", "github.com"], - input=token.encode(), capture_output=True, timeout=30, + input=token.encode(), + capture_output=True, + timeout=30, ) log.info("gh auth login OK") except Exception as e: @@ -131,7 +142,9 @@ def _ensure_gh_auth(repo: Path) -> None: try: result = subprocess.run( ["gh", "auth", "status", "--hostname", "github.com"], - capture_output=True, text=True, timeout=30, + capture_output=True, + text=True, + timeout=30, ) log.info(f"gh auth status: {result.stdout.strip()}") except Exception: @@ -141,7 +154,9 @@ def _ensure_gh_auth(repo: Path) -> None: try: subprocess.run( ["gh", "auth", "setup-git", "--hostname", "github.com"], - capture_output=True, text=True, timeout=30, + capture_output=True, + text=True, + timeout=30, ) log.info("gh auth setup-git OK") except Exception: @@ -193,7 +208,9 @@ def _push_to_fork(repo: Path, branch: str, fork_owner: str, token: str) -> None: log.info(f"Pushing to fork {fork_owner}/triton-ascend via proxy...") log.debug(f"fork remote: {fork_remote}") - log.debug(f"fork URL (masked): https://x-access-token:***@gh-proxy.test.osinfra.cn/https://github.com/{fork_owner}/triton-ascend.git") + log.debug( + f"fork URL (masked): https://x-access-token:***@gh-proxy.test.osinfra.cn/https://github.com/{fork_owner}/triton-ascend.git" + ) last_error = "" for attempt in range(1, _MAX_PUSH_RETRIES + 1): @@ -204,10 +221,18 @@ def _push_to_fork(repo: Path, branch: str, fork_owner: str, token: str) -> None: run_git(repo, "remote", "add", fork_remote, fork_url) push_result = subprocess.run( - ["git", - "-c", "http.https://github.com/.extraheader=", - "push", "--force-with-lease", fork_remote, branch], - cwd=str(repo), capture_output=True, text=True, + [ + "git", + "-c", + "http.https://github.com/.extraheader=", + "push", + "--force-with-lease", + fork_remote, + branch, + ], + cwd=str(repo), + capture_output=True, + text=True, ) # Clean up temp remote run_git_no_check(repo, "remote", "remove", fork_remote) @@ -235,8 +260,13 @@ def _push_to_fork(repo: Path, branch: str, fork_owner: str, token: str) -> None: def _create_pr( - repo: Path, github_repo: str, branch: str, fork_owner: str, token: str, - summary_file: Path, target_commit: str, + repo: Path, + github_repo: str, + branch: str, + fork_owner: str, + token: str, + summary_file: Path, + target_commit: str, ) -> str: """Create PR via gh CLI (with fork-aware origin swap). @@ -256,8 +286,7 @@ def _create_pr( saved_origin = run_git(repo, "config", "--get", "remote.origin.url").strip() if token and fork_owner: pr_origin = ( - f"https://x-access-token:{token}@" - f"github.com/{fork_owner}/triton-ascend.git" + f"https://x-access-token:{token}@github.com/{fork_owner}/triton-ascend.git" ) else: pr_origin = saved_origin @@ -272,7 +301,9 @@ def _create_pr( return pr_url except Exception as e: last_error = str(e) - log.warning(f"PR create attempt {attempt}/{_MAX_PR_RETRIES} FAILED: {last_error}") + log.warning( + f"PR create attempt {attempt}/{_MAX_PR_RETRIES} FAILED: {last_error}" + ) if attempt < _MAX_PR_RETRIES: time.sleep(_RETRY_DELAY_BASE * attempt) finally: @@ -293,7 +324,9 @@ def _create_pr( try: return _create_pr_via_api(github_repo, head, title, pr_body, base_branch, token) except Exception as e: - raise RuntimeError(f"PR creation failed after all attempts: {last_error}; API fallback: {e}") + raise RuntimeError( + f"PR creation failed after all attempts: {last_error}; API fallback: {e}" + ) def _build_pr_title(target_commit: str = "") -> str: @@ -308,14 +341,19 @@ def _build_pr_title(target_commit: str = "") -> str: author = os.getenv("PR_AUTHOR", "Sync").strip() pr_type = os.getenv("PR_TYPE", "feat").strip() if target_commit: - return f"[{author}]({pr_type}) Merge upstream triton commits {target_commit[:8]}" + return ( + f"[{author}]({pr_type}) Merge upstream triton commits {target_commit[:8]}" + ) ts = datetime.now().strftime("%Y%m%d-%H%M%S") return f"[{author}]({pr_type}) Merge upstream triton commits {ts}" def _create_pr_via_gh( - github_repo: str, title: str, body: str, - head_ref: str, base_branch: str, + github_repo: str, + title: str, + body: str, + head_ref: str, + base_branch: str, ) -> str: """Create a GitHub PR via the gh CLI. @@ -324,26 +362,36 @@ def _create_pr_via_gh( """ gh_token = os.environ.get("GH_TOKEN") or "" cmd = [ - "gh", "pr", "create", - "--title", title, - "--body", body, - "--head", head_ref, - "--base", base_branch, - "--repo", github_repo, + "gh", + "pr", + "create", + "--title", + title, + "--body", + body, + "--head", + head_ref, + "--base", + base_branch, + "--repo", + github_repo, ] log.info(f"Running: GH_HOST=github.com {' '.join(cmd)}") result = subprocess.run( cmd, - capture_output=True, text=True, timeout=60, - env={**os.environ, - "GITHUB_TOKEN": gh_token, - "GH_TOKEN": gh_token, - "GH_HOST": "github.com"}, + capture_output=True, + text=True, + timeout=60, + env={ + **os.environ, + "GITHUB_TOKEN": gh_token, + "GH_TOKEN": gh_token, + "GH_HOST": "github.com", + }, ) if result.returncode != 0: raise RuntimeError( - f"gh pr create failed (exit {result.returncode}): " - f"{result.stderr.strip()}" + f"gh pr create failed (exit {result.returncode}): {result.stderr.strip()}" ) pr_url = result.stdout.strip() if not pr_url: @@ -352,8 +400,12 @@ def _create_pr_via_gh( def _create_pr_via_api( - github_repo: str, head: str, title: str, body: str, - base: str, token: str, + github_repo: str, + head: str, + title: str, + body: str, + base: str, + token: str, ) -> str: """Create a GitHub PR via the REST API (fallback). @@ -372,12 +424,21 @@ def _create_pr_via_api( url = f"https://api.github.com/repos/{github_repo}/pulls" cmd = [ - "curl", "-s", "-X", "POST", url, - "-H", f"Authorization: Bearer {token}", - "-H", "Accept: application/vnd.github+json", - "-H", "X-GitHub-Api-Version: 2022-11-28", - "-H", "Content-Type: application/json", - "-d", json.dumps(data), + "curl", + "-s", + "-X", + "POST", + url, + "-H", + f"Authorization: Bearer {token}", + "-H", + "Accept: application/vnd.github+json", + "-H", + "X-GitHub-Api-Version: 2022-11-28", + "-H", + "Content-Type: application/json", + "-d", + json.dumps(data), ] result = subprocess.run(cmd, capture_output=True, text=True, timeout=30) if result.returncode != 0: diff --git a/src/TA_main2main_workflow/pipeline/resolve.py b/src/TA_main2main_workflow/pipeline/resolve.py index 2bfea41..a33ec9e 100644 --- a/src/TA_main2main_workflow/pipeline/resolve.py +++ b/src/TA_main2main_workflow/pipeline/resolve.py @@ -38,7 +38,7 @@ def resolve_conflicts(ctx: WorkflowContext, config: TAConfig) -> WorkflowContext prev_summary_path = str( WORKSPACE_DIR / STEPS_DIR / prev_step_id / "step_summary.md" ) - is_last_step = (ctx.current_step >= ctx.total_steps - 1) + is_last_step = ctx.current_step >= ctx.total_steps - 1 log.header("AI Conflict Resolution") for attempt in range(1, config.max_retries + 1): diff --git a/src/TA_main2main_workflow/pipeline/test.py b/src/TA_main2main_workflow/pipeline/test.py index 8b72683..b2b091a 100644 --- a/src/TA_main2main_workflow/pipeline/test.py +++ b/src/TA_main2main_workflow/pipeline/test.py @@ -17,7 +17,7 @@ from TA_main2main_workflow.utils.context import WorkflowContext from TA_main2main_workflow.utils.logging import get_logger from TA_main2main_workflow.utils.tracker import timed -from TA_main2main_workflow.utils import TEST_RESULT_FILE, WORKSPACE_DIR, STEPS_DIR +from TA_main2main_workflow.utils import WORKSPACE_DIR from TA_main2main_workflow.pipeline.build import build_triton, commit_fixes from TA_main2main_workflow.pipeline.fix import ai_fix @@ -57,7 +57,8 @@ def test_and_fix_loop(ctx: WorkflowContext, config: TAConfig) -> WorkflowContext oom_ctx = rerun_tests_reduced_concurrency(ascend_path, config) if oom_ctx is not None and oom_ctx.test_passed: return ctx.copy_with( - test_passed=True, pytest_passed=True, + test_passed=True, + pytest_passed=True, test_fix_count=ctx.test_fix_count, ) @@ -80,7 +81,8 @@ def test_and_fix_loop(ctx: WorkflowContext, config: TAConfig) -> WorkflowContext if attempt > 0: commit_fixes(ctx, config) return ctx.copy_with( - test_passed=True, pytest_passed=True, + test_passed=True, + pytest_passed=True, test_fix_count=ctx.test_fix_count + (1 if attempt > 0 else 0), ) @@ -96,7 +98,9 @@ def test_and_fix_loop(ctx: WorkflowContext, config: TAConfig) -> WorkflowContext def _run_pretest_and_fix( - ctx: WorkflowContext, config: TAConfig, ascend_path: Path, + ctx: WorkflowContext, + config: TAConfig, + ascend_path: Path, ) -> WorkflowContext: """Run a single-file pre-test with its own fix loop. @@ -127,8 +131,9 @@ def _run_pretest_and_fix( continue with timed("pretest"): - ctx = _run_pytest(ctx, config, [_PRETEST_FILE], - test_procs=1, label="pretest") + ctx = _run_pytest( + ctx, config, [_PRETEST_FILE], test_procs=1, label="pretest" + ) if ctx.test_passed: if pretest_attempt > 0: commit_fixes(ctx, config) @@ -141,8 +146,9 @@ def _run_pretest_and_fix( return ctx.copy_with(test_passed=False) -def run_tests(ctx: WorkflowContext, config: TAConfig, - python_exe: str = "", test_procs: int = 0) -> WorkflowContext: +def run_tests( + ctx: WorkflowContext, config: TAConfig, python_exe: str = "", test_procs: int = 0 +) -> WorkflowContext: """Execute tests sequentially. 1. Default pytest UT (primary test dir) — always runs first @@ -155,7 +161,9 @@ def run_tests(ctx: WorkflowContext, config: TAConfig, test_dirs = list(config.test_dirs) if not test_dirs and not config.test_command: log.warning("No test directories or test command configured") - return ctx.copy_with(test_passed=True, test_log_dir=str(WORKSPACE_DIR / "test-logs")) + return ctx.copy_with( + test_passed=True, test_log_dir=str(WORKSPACE_DIR / "test-logs") + ) primary = test_dirs[0] if test_dirs else None extras = test_dirs[1:] if len(test_dirs) > 1 else [] @@ -166,16 +174,28 @@ def run_tests(ctx: WorkflowContext, config: TAConfig, # ── Step 1: Default pytest UT ───────────────────────────────────── if primary: log.section("Default pytest UT") - ctx = _run_pytest(ctx, config, [primary], python_exe=python_exe, - test_procs=test_procs, label="primary") + ctx = _run_pytest( + ctx, + config, + [primary], + python_exe=python_exe, + test_procs=test_procs, + label="primary", + ) all_passed = all_passed and ctx.test_passed all_errors.extend(ctx.fix_errors) # ── Step 2: Extra test dirs — one by one ────────────────────────── for i, extra_dir in enumerate(extras): log.section(f"Extra Tests ({i + 1}/{len(extras)}): {extra_dir}") - extra_ctx = _run_pytest(ctx, config, [extra_dir], python_exe=python_exe, - test_procs=test_procs, label=f"extra-{i + 1}") + extra_ctx = _run_pytest( + ctx, + config, + [extra_dir], + python_exe=python_exe, + test_procs=test_procs, + label=f"extra-{i + 1}", + ) if not extra_ctx.test_passed: all_passed = False all_errors.extend(extra_ctx.fix_errors) @@ -228,7 +248,8 @@ def _run_custom_test(ctx: WorkflowContext, config: TAConfig) -> WorkflowContext: proc = subprocess.Popen( ["bash", "-c", cmd], cwd=str(ascend_path), - stdout=fh, stderr=subprocess.STDOUT, + stdout=fh, + stderr=subprocess.STDOUT, ) try: rc = proc.wait(timeout=7200) @@ -244,6 +265,7 @@ def _run_custom_test(ctx: WorkflowContext, config: TAConfig) -> WorkflowContext: pf = pe = tp = 0 junit_xml: Path | None = None import re as _re + m = _re.search(r"--junitxml[= ](\S+)", cmd) if m: junit_xml = Path(m.group(1)) @@ -253,7 +275,9 @@ def _run_custom_test(ctx: WorkflowContext, config: TAConfig) -> WorkflowContext: try: tree = ET.parse(str(junit_xml)) root = tree.getroot() - suites = [root] if root.tag != "testsuites" else root.findall("testsuite") + suites = ( + [root] if root.tag != "testsuites" else root.findall("testsuite") + ) for s in suites: tp += int(s.get("tests", 0)) pf += int(s.get("failures", 0)) @@ -286,7 +310,9 @@ def _run_custom_test(ctx: WorkflowContext, config: TAConfig) -> WorkflowContext: fix_errors=[str(junit_xml or output_log), str(result_file)], test_log_dir=str(test_log_dir), ) - log.status(True, f"All tests passed ({tp} passed)" if tp else f"Tests passed (exit 0)") + log.status( + True, f"All tests passed ({tp} passed)" if tp else "Tests passed (exit 0)" + ) return ctx.copy_with(test_passed=True, test_log_dir=str(test_log_dir)) @@ -295,10 +321,14 @@ def _run_custom_test(ctx: WorkflowContext, config: TAConfig) -> WorkflowContext: # --------------------------------------------------------------------------- -def _run_pytest(ctx: WorkflowContext, config: TAConfig, - test_dirs: list[str], - python_exe: str = "", test_procs: int = 0, - label: str = "pytest") -> WorkflowContext: +def _run_pytest( + ctx: WorkflowContext, + config: TAConfig, + test_dirs: list[str], + python_exe: str = "", + test_procs: int = 0, + label: str = "pytest", +) -> WorkflowContext: """Execute pytest for the given *test_dirs* in a single invocation. Each call with a unique *label* writes to a separate JUnit XML file @@ -327,14 +357,14 @@ def _run_pytest(ctx: WorkflowContext, config: TAConfig, junit_xml = test_log_dir / f"pytest-junit-{label}.xml" pytest_bin = shutil.which("pytest") - cmd = ( - [pytest_bin] if pytest_bin - else [python_exe, "-m", "pytest"] - ) + cmd = [pytest_bin] if pytest_bin else [python_exe, "-m", "pytest"] cmd += [str(p) for p in test_paths] cmd += ["-n", str(procs), f"--junitxml={junit_xml}"] - log.key_value(f"[{label}] test dirs", ", ".join(str(p.relative_to(ascend_path)) for p in test_paths)) + log.key_value( + f"[{label}] test dirs", + ", ".join(str(p.relative_to(ascend_path)) for p in test_paths), + ) log.info(f"[{label}] cmd: {' '.join(cmd)}") _start = time.time() try: @@ -396,10 +426,17 @@ def detect_oom_in_tests(ctx: WorkflowContext) -> bool: return False oom_keywords = [ - "OutOfMemoryError", "out of memory", "MemoryError", - "Cannot allocate memory", "OOM", "Killed", - "Exit code 137", "exit code 137", - "CUDA error", "cuMemAlloc", "NPU error", + "OutOfMemoryError", + "out of memory", + "MemoryError", + "Cannot allocate memory", + "OOM", + "Killed", + "Exit code 137", + "exit code 137", + "CUDA error", + "cuMemAlloc", + "NPU error", ] # Scan ALL JUnit XML files (each test suite writes its own) for junit_xml in sorted(test_log_dir.glob("pytest-junit-*.xml")): @@ -440,7 +477,9 @@ def rerun_tests_reduced_concurrency( procs = max(1, original_procs // (2 ** (r + 1))) if procs >= original_procs: break - log.info(f"Rerunning tests with {procs} workers (attempt {r + 1}/{max_reruns})...") + log.info( + f"Rerunning tests with {procs} workers (attempt {r + 1}/{max_reruns})..." + ) ascend_path_str = str(ascend_path) ctx = WorkflowContext(triton_ascend_path=ascend_path_str) @@ -450,7 +489,9 @@ def rerun_tests_reduced_concurrency( return ctx # Stop if OOM is gone — remaining failures are code issues if not detect_oom_in_tests(ctx): - log.info("OOM resolved — remaining failures are not memory-related, stopping rerun") + log.info( + "OOM resolved — remaining failures are not memory-related, stopping rerun" + ) return ctx log.error(f"Tests still failing after {max_reruns} concurrency reductions") diff --git a/src/TA_main2main_workflow/reference/01-merge-upstream-conflict-resolution.md b/src/TA_main2main_workflow/reference/01-merge-upstream-conflict-resolution.md index 28de386..0a865cb 100644 --- a/src/TA_main2main_workflow/reference/01-merge-upstream-conflict-resolution.md +++ b/src/TA_main2main_workflow/reference/01-merge-upstream-conflict-resolution.md @@ -91,10 +91,10 @@ def program_id(axis: int, builder: ir.builder) -> tl.tensor: # 新代码(3.5.x) class TritonSemantic(Generic[TensorTy]): builder: ir.builder - + def __init__(self, builder): self.builder = builder - + def program_id(self, axis: int) -> TensorTy: return self.tensor(self.builder.create_get_program_id(axis), tl.int32) ``` diff --git a/src/TA_main2main_workflow/reference/02-llvm-version-adaptation-and-compile-fixes.md b/src/TA_main2main_workflow/reference/02-llvm-version-adaptation-and-compile-fixes.md index 6e6753a..8cc027c 100644 --- a/src/TA_main2main_workflow/reference/02-llvm-version-adaptation-and-compile-fixes.md +++ b/src/TA_main2main_workflow/reference/02-llvm-version-adaptation-and-compile-fixes.md @@ -94,7 +94,7 @@ FILES=$(grep -rl "\.get\|\.get\|\.is" \ third_party/ascend/lib/ --include="*.cpp" --include="*.h") for f in $FILES; do - # 替换 .get() → cast(...) + # 替换 .get() → cast(...) # 替换 .get() → cast(...) # 替换 .is() → isa(...) echo "Processing: $f" diff --git a/src/TA_main2main_workflow/reference/03-unit-test-failure-diagnosis-and-fixes.md b/src/TA_main2main_workflow/reference/03-unit-test-failure-diagnosis-and-fixes.md index 672fb3b..a17b5ba 100644 --- a/src/TA_main2main_workflow/reference/03-unit-test-failure-diagnosis-and-fixes.md +++ b/src/TA_main2main_workflow/reference/03-unit-test-failure-diagnosis-and-fixes.md @@ -310,7 +310,7 @@ if __name__ == "__main__": def test_vector_add(): # ...测试逻辑... torch.testing.assert_close(output, expected) - + test_vector_add() ``` diff --git a/src/TA_main2main_workflow/reference/04-ir-compatibility-and-backend-adaptation.md b/src/TA_main2main_workflow/reference/04-ir-compatibility-and-backend-adaptation.md index 197c1b7..39ade92 100644 --- a/src/TA_main2main_workflow/reference/04-ir-compatibility-and-backend-adaptation.md +++ b/src/TA_main2main_workflow/reference/04-ir-compatibility-and-backend-adaptation.md @@ -113,7 +113,7 @@ class NPUOptions: ascend.indirect_store %ptr, %indices, %data : , tensor<8x32xi64>, tensor<8x32xf32> # 新 IR (3.5.x): -%result = ascend.unstructured_load %ptr, %indices : , tensor<8x32xi64> +%result = ascend.unstructured_load %ptr, %indices : , tensor<8x32xi64> unstructured_dims = [0, 1] -> tensor<8x32xf32> ascend.unstructured_store %ptr, %indices, %data : , tensor<8x32xi64>, tensor<8x32xf32> unstructured_dims = [0, 1] @@ -251,7 +251,7 @@ name, mix_mode = name.rsplit("_", 1) @dataclass class NPUOptions: ir_override: Optional[str] = None # 覆盖 IR 文件的路径前缀 - + # 使用方式: # 如果 ir_override = "/path/to/kernel" # 则会查找: diff --git a/src/TA_main2main_workflow/reference/ir_compatibility_patch_example.patch b/src/TA_main2main_workflow/reference/ir_compatibility_patch_example.patch index b5dd2e3..4ea3075 100644 --- a/src/TA_main2main_workflow/reference/ir_compatibility_patch_example.patch +++ b/src/TA_main2main_workflow/reference/ir_compatibility_patch_example.patch @@ -21,7 +21,7 @@ index e735651d5366..26632e26cbbf 100644 @@ -61,6 +61,12 @@ LogicalResult foldToBufferToTensorPair(RewriterBase &rewriter, ToBufferOp toBuffer, const BufferizationOptions &options); - + +/// Try to fold to_memref(to_tensor(x)). If x's type and the result type of the +/// to_memref op are different, a memref.cast is needed. +LogicalResult foldToMemrefToTensorPair(RewriterBase &rewriter, @@ -38,7 +38,7 @@ index 6724d4c48310..33cfd50c6849 100644 @@ -483,13 +483,85 @@ def Bufferization_ToTensorOp : Bufferization_Op<"to_tensor", [ } }]; - + + let hasCustomAssemblyFormat = 1; + + let hasCanonicalizer = 1; @@ -117,20 +117,20 @@ index 6724d4c48310..33cfd50c6849 100644 - `:` type($buffer) `to` type($result) + $tensor (`read_only` $read_only^)? attr-dict `:` type($memref) }]; - + - let hasCanonicalizer = 1; let hasFolder = 1; + let hasCanonicalizer = 1; } - - + + diff --git a/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td b/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td index a441fd82546e..543ab84ef896 100644 --- a/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td +++ b/mlir/include/mlir/Dialect/ControlFlow/IR/ControlFlowOps.td @@ -227,12 +227,7 @@ def CondBranchOp }]; - + let hasCanonicalizer = 1; - let assemblyFormat = [{ - $condition (`weights` `(` $branch_weights^ `)` )? `,` @@ -140,7 +140,7 @@ index a441fd82546e..543ab84ef896 100644 - }]; + let hasCustomAssemblyFormat = 1; } - + //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td index 06ce4f16c867..a17cec470904 100644 @@ -149,13 +149,13 @@ index 06ce4f16c867..a17cec470904 100644 @@ -119,9 +119,7 @@ def CallOp : Func_Op<"call", } }]; - + - let assemblyFormat = [{ - $callee `(` $operands `)` attr-dict `:` functional-type($operands, results) - }]; + let hasCustomAssemblyFormat = 1; } - + //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td index 398388bd720b..877ebc89e9f0 100644 @@ -164,14 +164,14 @@ index 398388bd720b..877ebc89e9f0 100644 @@ -557,12 +557,6 @@ def LLVM_AssumeOp dag args = (ins I1:$cond); let arguments = !con(args, baseArgs); - + - let assemblyFormat = [{ - $cond - ( custom($op_bundle_operands, type($op_bundle_operands), - $op_bundle_tags)^ )? - `:` type($cond) attr-dict - }]; - + let builders = [ OpBuilder<(ins "Value":$cond)>, diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td @@ -221,9 +221,9 @@ index 9753dca67c23..48bf6d8d4893 100644 LLVM_ScalarOrVectorOf, LLVM_ScalarOrVectorOf>; @@ -2394,16 +2385,7 @@ def LLVM_InlineAsmOp : LLVM_Op<"inline_asm", [DeclareOpInterfaceMethods:$res); - + - let assemblyFormat = [{ - (`has_side_effects` $has_side_effects^)? - (`is_align_stack` $is_align_stack^)? @@ -235,7 +235,7 @@ index 9753dca67c23..48bf6d8d4893 100644 - operands `:` functional-type(operands, results) - }]; + let hasCustomAssemblyFormat = 1; - + let extraClassDeclaration = [{ static StringRef getElementTypeAttrName() { diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp @@ -245,7 +245,7 @@ index 7cfd6d3a98df..e78cb55f8bae 100644 @@ -1561,6 +1561,83 @@ LogicalResult arith::ScalingExtFOp::verify() { // TruncIOp //===----------------------------------------------------------------------===// - + +ParseResult TruncIOp::parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result) { + ::mlir::OpAsmParser::UnresolvedOperand inRawOperand{}; + ::llvm::ArrayRef<::mlir::OpAsmParser::UnresolvedOperand> inOperands(&inRawOperand, 1); ::llvm::SMLoc inOperandsLoc; @@ -332,7 +332,7 @@ index e0cf353da207..6bd0b89e5517 100644 +++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp @@ -654,7 +654,9 @@ bool AnalysisState::canOmitTensorCopy(OpOperand &opOperand) const { } - + bool AnalysisState::isInPlace(OpOperand &opOperand) const { - // ToBufferOps are always in-place. + // ToMemrefOps and ToBufferOps are always in-place. @@ -340,11 +340,11 @@ index e0cf353da207..6bd0b89e5517 100644 + return true; if (isa(opOperand.getOwner())) return true; - + @@ -690,6 +692,16 @@ static void ensureToBufferOpIsValid(Value tensor, Type memrefType) { #endif } - + +// bufferization.to_memref is not allowed to change the rank. +static void ensureToMemrefOpIsValid(Value tensor, Type memrefType) { +#ifndef NDEBUG @@ -361,7 +361,7 @@ index e0cf353da207..6bd0b89e5517 100644 @@ -702,12 +714,13 @@ FailureOr bufferization::getBuffer(RewriterBase &rewriter, Value value, if (auto toTensorOp = value.getDefiningOp()) return toTensorOp.getBuffer(); - + - // Insert to_buffer op. + // Insert to_memref or to_buffer op. OpBuilder::InsertionGuard g(rewriter); @@ -383,12 +383,12 @@ index 56ff2121e462..5d1d7a165c3d 100644 #include "mlir/IR/Matchers.h" +#include "mlir/Interfaces/InferTypeOpInterface.h" #include - + using namespace mlir; @@ -127,6 +128,53 @@ LogicalResult mlir::bufferization::foldToBufferToTensorPair( return success(); } - + +/// Try to fold to_memref(to_tensor(x)). If x's type and the result type of the +/// to_memref op are different, a memref.cast is needed. +LogicalResult mlir::bufferization::foldToMemrefToTensorPair( @@ -442,7 +442,7 @@ index 56ff2121e462..5d1d7a165c3d 100644 @@ -750,6 +798,101 @@ bool ToTensorOp::isWritable(Value value, const AnalysisState &state) { return getWritable(); } - + + +ParseResult ToTensorOp::parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result) { + ::mlir::OpAsmParser::UnresolvedOperand bufferRawOperand{}; @@ -553,11 +553,11 @@ index 56ff2121e462..5d1d7a165c3d 100644 + return toMemref.getTensor(); return {}; } - + @@ -1198,6 +1347,107 @@ void bufferization::populateDeallocOpCanonicalizationPatterns( RemoveAllocDeallocPairWhenNoOtherUsers>(context); } - + +//===----------------------------------------------------------------------===// +// ToMemrefOp +//===----------------------------------------------------------------------===// @@ -669,7 +669,7 @@ index 582593adfa5c..a81a84bae601 100644 @@ -437,6 +437,114 @@ struct CondBranchTruthPropagation : public OpRewritePattern { }; } // namespace - + +ParseResult CondBranchOp::parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result) { + ::mlir::OpAsmParser::UnresolvedOperand conditionRawOperand{}; + ::llvm::ArrayRef<::mlir::OpAsmParser::UnresolvedOperand> conditionOperands(&conditionRawOperand, 1); ::llvm::SMLoc conditionOperandsLoc; @@ -788,7 +788,7 @@ index 3c09a2124bd7..a1bf2bdf5a09 100644 @@ -59,6 +59,72 @@ Operation *FuncDialect::materializeConstant(OpBuilder &builder, Attribute value, // CallOp //===----------------------------------------------------------------------===// - + +ParseResult CallOp::parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result) { + ::mlir::FlatSymbolRefAttr calleeAttr; + ::llvm::SmallVector<::mlir::OpAsmParser::UnresolvedOperand, 4> operandsOperands; @@ -865,7 +865,7 @@ index 5d08cccb4faa..219e64280985 100644 @@ -834,6 +834,172 @@ LogicalResult LLVM::GEPOp::verify() { // LoadOp //===----------------------------------------------------------------------===// - + +ParseResult LoadOp::parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result) { + ::mlir::OpAsmParser::UnresolvedOperand addrRawOperand{}; + ::llvm::ArrayRef<::mlir::OpAsmParser::UnresolvedOperand> addrOperands(&addrRawOperand, 1); ::llvm::SMLoc addrOperandsLoc; @@ -1204,13 +1204,13 @@ index 5d08cccb4faa..219e64280985 100644 + _odsPrinter << ' '; + _odsPrinter << getAddr().getType(); +} - + void StoreOp::getEffects( SmallVectorImpl> @@ -4209,6 +4540,234 @@ LogicalResult InlineAsmOp::verify() { return success(); } - + +ParseResult InlineAsmOp::parse(::mlir::OpAsmParser &parser, ::mlir::OperationState &result) { + ::mlir::LLVM::AsmDialectAttr asm_dialectAttr; + ::mlir::ArrayAttr operand_attrsAttr; @@ -1448,7 +1448,7 @@ index 59013a23b3e3..b44c025c9c33 100644 +++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp @@ -3898,33 +3898,14 @@ static FailureOr parseIndexingMapsAttr(OpAsmParser &parser) { } - + ParseResult MatmulOp::parse(OpAsmParser &parser, OperationState &result) { - FailureOr indexingMapsAttr = parseIndexingMapsAttr(parser); - if (failed(indexingMapsAttr)) @@ -1484,7 +1484,7 @@ index 59013a23b3e3..b44c025c9c33 100644 + // LinalgNamedStructuredOps.yamlgen.cpp.inc + "linalg.memoized_indexing_maps"}); } - + /// Verify the user defined indexing maps. @@ -4678,11 +4659,9 @@ void BatchMatmulOp::print(OpAsmPrinter &p) { SmallVector indexingMaps = llvm::map_to_vector<3>( @@ -1492,7 +1492,7 @@ index 59013a23b3e3..b44c025c9c33 100644 [](AffineMap map) -> Attribute { return AffineMapAttr::get(map); }); - if (!llvm::equal(getIndexingMaps(), indexingMaps)) - p << " indexing_maps = " << llvm::interleaved_array(getIndexingMaps()); - + std::array elidedAttrs = { - "operandSegmentSizes", "linalg.memoized_indexing_maps", "indexing_maps"}; + "operandSegmentSizes", "linalg.memoized_indexing_maps"}; @@ -1505,18 +1505,18 @@ index a9da6c2c8320..4ebc46deae3d 100644 +++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp @@ -486,9 +486,6 @@ static void printInitializationList(OpAsmPrinter &p, } - + void ForOp::print(OpAsmPrinter &p) { - if (getUnsignedCmp()) - p << " unsigned"; - p << " " << getInductionVar() << " = " << getLowerBound() << " to " << getUpperBound() << " step " << getStep(); - + @@ -512,10 +509,6 @@ ParseResult ForOp::parse(OpAsmParser &parser, OperationState &result) { OpAsmParser::Argument inductionVariable; OpAsmParser::UnresolvedOperand lb, ub, step; - + - if (succeeded(parser.parseOptionalKeyword("unsigned"))) - result.addAttribute(getUnsignedCmpAttrName(result.name), - builder.getUnitAttr()); @@ -1531,7 +1531,7 @@ index a9da6c2c8320..4ebc46deae3d 100644 - [](OpBuilder &, Location, Value, ValueRange) {}, getUnsignedCmp()); + [](OpBuilder &, Location, Value, ValueRange) {}, false); newLoop->setAttrs(getPrunedAttributeList(getOperation(), {})); - + // Generate the new yield values and append them to the scf.yield operation. @@ -854,7 +847,7 @@ mlir::scf::replaceAndCastForOpIterArg(RewriterBase &rewriter, scf::ForOp forOp, scf::ForOp newForOp = scf::ForOp::create( @@ -1550,22 +1550,22 @@ index a9da6c2c8320..4ebc46deae3d 100644 + /*bodyBuilder=*/nullptr, false); newForOp->setAttrs(forOp->getAttrs()); Block &newBlock = newForOp.getRegion().front(); - + @@ -1175,7 +1168,7 @@ Speculation::Speculatability ForOp::getSpeculatability() { - + std::optional ForOp::getStaticTripCount() { return constantTripCount(getLowerBound(), getUpperBound(), getStep(), - /*isSigned=*/!getUnsignedCmp(), computeUbMinusLb); + /*isSigned=*/!false, computeUbMinusLb); } - + //===----------------------------------------------------------------------===// diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp index 3d19c5ad8fbc..8b8d6d4b95dc 100644 --- a/mlir/lib/IR/AsmPrinter.cpp +++ b/mlir/lib/IR/AsmPrinter.cpp @@ -3733,8 +3733,14 @@ void OperationPrinter::printGenericOp(Operation *op, bool printOpName) { - + // Print the properties. if (Attribute prop = op->getPropertiesAsAttribute()) { + auto dict = cast(prop); @@ -1579,7 +1579,7 @@ index 3d19c5ad8fbc..8b8d6d4b95dc 100644 + Impl::printAttribute(DictionaryAttr::get(op->getContext(), filtered)); os << '>'; } - + diff --git a/mlir/lib/Interfaces/FunctionImplementation.cpp b/mlir/lib/Interfaces/FunctionImplementation.cpp index 90f32896e818..855410ae2580 100644 --- a/mlir/lib/Interfaces/FunctionImplementation.cpp @@ -1592,4 +1592,4 @@ index 90f32896e818..855410ae2580 100644 + p, op, {visibilityAttrName, typeAttrName, argAttrsName, resAttrsName, "no_inline"}); // Print the body if this is not an external function. Region &body = op->getRegion(0); - if (!body.empty()) { \ No newline at end of file + if (!body.empty()) { diff --git a/src/TA_main2main_workflow/utils/config.py b/src/TA_main2main_workflow/utils/config.py index 695d71b..f95475e 100644 --- a/src/TA_main2main_workflow/utils/config.py +++ b/src/TA_main2main_workflow/utils/config.py @@ -11,7 +11,7 @@ import os from dataclasses import dataclass, field from pathlib import Path -from typing import Literal +from typing import Literal, cast AIBackendChoice = Literal["opencode", "claude", "auto"] @@ -50,7 +50,9 @@ class TAConfig: skip_build: bool = False skip_e2e_test: bool = False skip_llvm_rebuild: bool = False # skip LLVM rebuild when version changes (IR patch) - skip_baseline_llvm: bool = False # skip initial baseline LLVM build at workflow start + skip_baseline_llvm: bool = ( + False # skip initial baseline LLVM build at workflow start + ) skip_ir_patch: bool = False # skip entire IR patch phase (SKIP_IR_PATCH) # ── Git / Branch ────────────────────────────────────────────────────── @@ -92,8 +94,9 @@ def from_env(cls) -> TAConfig: "TRITON_UPSTREAM_URL", "https://github.com/triton-lang/triton.git" ), target_commit=os.getenv("TRITON_TARGET_COMMIT", ""), - ai_backend=_env_choice( - "AI_BACKEND", ["opencode", "claude", "auto"], "auto" + ai_backend=cast( + AIBackendChoice, + _env_choice("AI_BACKEND", ["opencode", "claude", "auto"], "auto"), ), ai_timeout_minutes=_env_int("TA_AI_TIMEOUT_MINUTES", 30), ai_stale_seconds=_env_int("TA_AI_STALE_SECONDS", 1200), @@ -119,7 +122,9 @@ def from_env(cls) -> TAConfig: push_to_github=_env_bool("PUSH_TO_GITHUB", False), github_repo=os.getenv("GITHUB_REPO", "triton-lang/triton-ascend"), llvm_project_path=os.getenv("LLVM_PROJECT_PATH", "~/llvm-project"), - llvm_install_prefix_sync=os.getenv("LLVM_INSTALL_PREFIX_SYNC", "~/llvm-install-sync"), + llvm_install_prefix_sync=os.getenv( + "LLVM_INSTALL_PREFIX_SYNC", "~/llvm-install-sync" + ), conda_env=os.getenv("CONDA_ENV", "ta-upgrade"), test_dir=os.getenv("TA_TEST_DIR", "third_party/ascend/unittest/pytest_ut"), test_dirs=_resolve_test_dirs(), @@ -133,14 +138,19 @@ def from_env(cls) -> TAConfig: def llvm_project(self) -> Path: if self.llvm_project_path: return Path(os.path.expanduser(self.llvm_project_path)) - return Path(os.path.expanduser(os.getenv("LLVM_PROJECT_PATH", "~/llvm-project"))) + return Path( + os.path.expanduser(os.getenv("LLVM_PROJECT_PATH", "~/llvm-project")) + ) @property def llvm_install(self) -> Path: if self.llvm_install_prefix_sync: return Path(os.path.expanduser(self.llvm_install_prefix_sync)) - return Path(os.path.expanduser( - os.getenv("LLVM_INSTALL_PREFIX_SYNC", "~/llvm-install-sync"))) + return Path( + os.path.expanduser( + os.getenv("LLVM_INSTALL_PREFIX_SYNC", "~/llvm-install-sync") + ) + ) def _env_bool(name: str, default: bool) -> bool: @@ -186,7 +196,9 @@ def _resolve_test_dirs(primary: str = "", extra: str = "") -> list[str]: When called from main.py with CLI args, *primary* and *extra* override the env vars. """ - _primary = primary or os.getenv("TA_TEST_DIR", "third_party/ascend/unittest/pytest_ut") + _primary = primary or os.getenv( + "TA_TEST_DIR", "third_party/ascend/unittest/pytest_ut" + ) dirs = [_primary] if _primary else [] extra_raw = extra or os.getenv("TA_EXTRA_TEST_DIRS", "") diff --git a/src/TA_main2main_workflow/utils/git.py b/src/TA_main2main_workflow/utils/git.py index a3ee663..e4e1e14 100644 --- a/src/TA_main2main_workflow/utils/git.py +++ b/src/TA_main2main_workflow/utils/git.py @@ -29,7 +29,10 @@ def run_git(repo: Path | str, *args: str) -> str: for attempt in range(1, MAX_RETRIES + 1 if is_retryable else 2): result = subprocess.run( - cmd, capture_output=True, text=True, timeout=600, + cmd, + capture_output=True, + text=True, + timeout=600, ) if result.returncode == 0: return result.stdout @@ -49,8 +52,14 @@ def run_git_no_check(repo: Path | str, *args: str) -> subprocess.CompletedProces return subprocess.run(cmd, capture_output=True, text=True, timeout=600) -def stream_cmd(cmd: list[str], cwd: Path, log_fh, timeout: int, - label: str = "", env: dict | None = None) -> int: +def stream_cmd( + cmd: list[str], + cwd: Path, + log_fh, + timeout: int, + label: str = "", + env: dict | None = None, +) -> int: """Stream subprocess output line-by-line to console and log file. Each output line is: @@ -65,12 +74,17 @@ def stream_cmd(cmd: list[str], cwd: Path, log_fh, timeout: int, """ import os as _os import sys + proc_env = _os.environ.copy() if env: proc_env.update(env) proc = subprocess.Popen( - cmd, cwd=str(cwd), env=proc_env, - stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, + cmd, + cwd=str(cwd), + env=proc_env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, ) assert proc.stdout is not None last_line = "" diff --git a/src/TA_main2main_workflow/utils/logging.py b/src/TA_main2main_workflow/utils/logging.py index 8c9a904..42c2a62 100644 --- a/src/TA_main2main_workflow/utils/logging.py +++ b/src/TA_main2main_workflow/utils/logging.py @@ -30,7 +30,7 @@ # ═══════════════════════════════════════════════════════════════════════════ -class TALogger(logging.getLoggerClass()): +class TALogger(logging.Logger): """Logger with extra formatting methods for workflow output.""" def header(self, title: str) -> None: @@ -53,11 +53,10 @@ def status(self, ok: bool, msg: str) -> None: icon = "✔" if ok else "✘" self.info(f" {icon} {msg}") - def warn(self, msg: str, *args, **kwargs) -> None: - # Override to use consistent prefix + def warn(self, msg: object, *args: Any, **kwargs: Any) -> None: # type: ignore[override] super().warning(f" ⚠ {msg}", *args, **kwargs) - def error(self, msg: str, *args, **kwargs) -> None: + def error(self, msg: object, *args: Any, **kwargs: Any) -> None: # type: ignore[override] super().error(f" ✘ {msg}", *args, **kwargs) def key_value(self, key: str, value: Any) -> None: @@ -78,18 +77,20 @@ def conflict_list(self, files: list[str]) -> None: def ai_call(self, backend: str, mode: str, attempt: int, max_attempts: int) -> None: ts = datetime.now().strftime("%H:%M:%S") - self.info(f"\n ╭─ AI Call ─────────────────────────────────────────────") + self.info("\n ╭─ AI Call ─────────────────────────────────────────────") self.info(f" │ Backend: {backend}") self.info(f" │ Mode: {mode}") self.info(f" │ Attempt: {attempt}/{max_attempts}") self.info(f" │ Time: {ts}") - self.info(f" ╰──────────────────────────────────────────────────────") + self.info(" ╰──────────────────────────────────────────────────────") def ai_result( - self, ok: bool, modified_files: list[str] = (), summary: str = "" + self, ok: bool, modified_files: list[str] | None = None, summary: str = "" ) -> None: + if modified_files is None: + modified_files = [] icon = "✔" if ok else "✘" - self.info(f"\n ╭─ AI Result ───────────────────────────────────────────") + self.info("\n ╭─ AI Result ───────────────────────────────────────────") self.info(f" │ Status: {icon} {'Success' if ok else 'Failed'}") if modified_files: self.info(f" │ Modified files ({len(modified_files)}):") @@ -98,7 +99,7 @@ def ai_result( if summary: preview = summary[:500] + "..." if len(summary) > 500 else summary self.info(f" │ Summary: {preview}") - self.info(f" ╰──────────────────────────────────────────────────────") + self.info(" ╰──────────────────────────────────────────────────────") def table(self, rows: list[tuple[str, str, str]]) -> None: ts = datetime.now().strftime("%H:%M:%S") diff --git a/src/TA_main2main_workflow/utils/submodule.py b/src/TA_main2main_workflow/utils/submodule.py index 9aa0f83..b28a75f 100644 --- a/src/TA_main2main_workflow/utils/submodule.py +++ b/src/TA_main2main_workflow/utils/submodule.py @@ -7,7 +7,6 @@ from __future__ import annotations import os -import subprocess from pathlib import Path from TA_main2main_workflow.utils.git import run_git, run_git_no_check @@ -46,7 +45,7 @@ def commit_submodule(repo: Path, commit_msg: str) -> bool: if not submodule_has_changes(repo): return False - log.info(f"Committing AscendNPU-IR submodule changes...") + log.info("Committing AscendNPU-IR submodule changes...") try: run_git(sp, "add", "-A") run_git(sp, "commit", "-s", "-m", commit_msg) @@ -58,8 +57,7 @@ def commit_submodule(repo: Path, commit_msg: str) -> bool: return False -def push_submodule(repo: Path, branch: str | None = None, - force: bool = True) -> bool: +def push_submodule(repo: Path, branch: str | None = None, force: bool = True) -> bool: """Push the AscendNPU-IR submodule to its dedicated remote. Creates/updates a branch at current HEAD and pushes using