kneron_model_converter/services/backends/precheck.py

"""Platform / operator compatibility pre-check.

在轉檔流程「早期」（ONNX 階段拿到 out.onnx 之後、送 bie/nef 之前）用 IP Evaluator
的分析結果判斷模型是否含有目標 platform 不支援的 operator（例如 Softmax@520）。

背景：Kneron toolchain 的 ``ktc.ModelConfig(...).evaluate()``（ip_evaluator, mode 0）
內部會跑 compiler frontend 做 node placement 分析。當某個 op 在目標 platform 無法
落在 NPU、且沒有對應的 CPU op 定義時，compiler frontend 會在 log 內留下
``HardwareNotSupport`` / ``UnimplementedFeature`` / ``creating an EmptyNode instance for
op_type:`` 等訊號，toolchain 會把這些寫進 ``model_fx_report``。因此我們可以在編譯前
（不需要跑完 quantization + batch_compile）就偵測到不支援的 op，fail-fast 並回一個
使用者看得懂的錯誤，而不是讓 job 跑到 nef 階段撞 C++ backtrace（exit 6）。

設計原則（真 fail-open、避免誤擋）：
- 「明確訊號」= 不支援 marker **加上一個具體的 operator 名稱**。只出現 marker 字樣
  （例如報告表頭的 "HW not support" 欄位名）但抽不出 op 名稱時 → 放行 + warning，
  不擋。誤擋能轉的模型（false positive）的代價高於晚一點在 nef 階段失敗。
- evaluator 分析失敗（例如環境問題、transient error）而沒有明確不支援訊號時 → 放行，
  讓後續流程照舊跑，不因為 pre-check 有疑慮就擋掉本來能跑的模型。
"""

from __future__ import annotations

import logging
import re
from typing import Iterable, List, Optional

logger = logging.getLogger(__name__)


class UnsupportedOperatorError(Exception):
    """目標 platform 不支援模型中的某個 operator。

    這是一個「明確且可預期」的失敗，訊息會直接透過 job 失敗回報機制傳回給使用者。
    """


# Compiler frontend / ip_evaluator 在遇到不支援 op 時會在報告 / log 留下的訊號。
# 對照來源：vendor/sys_flow*/test_case.py::check_compiler_HardwareNotSupport 與
# check_knerex_error 寫入 model_fx_report 的字串。
#
# 每個 pattern 都帶一個「必須命中」的 capture group 抓 op 名稱片段。
# marker 後面抽不出具體 op 名稱 → 視為弱訊號、放行（見 find_unsupported_operators）。
#
# 注意：這裡刻意**沒有**「HW not support」欄位名的泛用 pattern —— 那是
# model_fx_report 的欄位分類名稱（vendor/sys_flow/test_case.py::model_fx_report），
# 會出現在正常報告的表頭，光是出現不代表模型有不支援 op（R3 誤擋根因，
# 見 .autoflow/05-implementation/tflite-520-rootcause-2026-07-06.md §3a/§4a）。
# 「HW not support: <Op>」帶具體 op 的強訊號仍由 hw_not_support pattern
# 涵蓋（IGNORECASE + [_ ] 分隔）。
_UNSUPPORTED_MARKERS: tuple[tuple[str, str], ...] = (
    # "creating an EmptyNode instance for op_type: Softmax"
    ("empty_node", r"creating an EmptyNode instance for op_type:\s*([A-Za-z0-9_]+)"),
    # "HW_NOT_SUPPORT: Softmax" / "HardwareNotSupport: Softmax" / "Hardware not support: Softmax"
    # op capture group 為必須：marker 後面一定要跟具體 op 名稱才算訊號。
    # marker 帶 (?:ED)? 字尾容忍 + \b 詞界：報告若寫 "Hardware not supported: ..."，
    # 字尾 "ed" 必須被 marker 吃掉、不能漏進 capture group 被當成 op 名（誤擋）。
    ("hw_not_support", r"(?:HW[_ ]NOT[_ ]SUPPORT(?:ED)?|Hardware\s*not\s*support(?:ed)?|HardwareNotSupport)\b\s*[:\-]?\s*([A-Za-z0-9_]+)"),
    # "UNIMPLEMENTED_FEATURE: ..." / "UnimplementedFeature: undefined CPU op [Softmax]"
    ("unimplemented", r"(?:UNIMPLEMENTED[_ ]FEATURE|UnimplementedFeature)\s*[:\-]?\s*(.*)"),
)

# 從一段訊息中撈出被中括號 / 括號包住的 op 名稱，例如 "undefined CPU op [Softmax]"。
_OP_IN_BRACKETS = re.compile(r"[\[\(]\s*([A-Za-z][A-Za-z0-9_]*)\s*[\]\)]")
# 報告表格常見的「非 op」值：欄位名（如 "HW not support"）後面接的可能是這些字，
# 不能把它們當成 operator 名稱（否則又變回誤擋）。
_NON_OP_TOKENS = frozenset({"none", "null", "nan", "na", "n", "true", "false", "yes", "no"})


def _extract_op_name(fragment: Optional[str]) -> Optional[str]:
    """從訊息片段中抽出 operator 名稱。抓不到（或抓到的不像 op）就回 None。

    只接受兩種**明確**形式：
    1. 中括號 / 括號包住的 token（例如 "undefined CPU op [Softmax]"）；
    2. fragment 本身就是單一 token（例如 hw_not_support capture 到的 "Softmax"）。

    刻意**不**在自由文字裡猜「首字母大寫的字」——自由文字（如
    "UnimplementedFeature: Not supported in this mode"）猜出來的字
    （"Not"、"CPU"）不是 op 或不是對的 op，會造成誤擋 / 錯誤訊息；
    抓不到就交給呼叫端 fail-open。
    """
    if not fragment:
        return None
    fragment = fragment.strip()
    if not fragment:
        return None

    candidate: Optional[str] = None

    m = _OP_IN_BRACKETS.search(fragment)
    if m:
        candidate = m.group(1)
    elif re.fullmatch(r"[A-Za-z][A-Za-z0-9_]*", fragment):
        # 直接就是一個 op token（例如 "Softmax"）
        candidate = fragment

    if candidate and candidate.lower() in _NON_OP_TOKENS:
        return None
    return candidate


def find_unsupported_operators(report: Optional[str]) -> List[str]:
    """掃描 evaluator 報告字串，回傳偵測到的「不支援 operator 名稱」清單。

    - 沒有偵測到任何明確不支援訊號 → 回空 list（呼叫端應放行）。
    - 偵測到 marker 字樣但抽不出具體 op 名稱 → **放行**（回空 list）+ warning log。
      marker 字樣可能只是報告表頭的欄位名（例如 "HW not support" 欄），
      光是出現不足以斷定模型有不支援 op；擋錯的代價（誤擋能轉的模型）
      比晚一點在 nef 階段失敗更高（真 fail-open）。
    """
    if not report:
        return []

    found: List[str] = []
    saw_marker = False

    for _kind, pattern in _UNSUPPORTED_MARKERS:
        for match in re.finditer(pattern, report, flags=re.IGNORECASE):
            saw_marker = True
            captured = match.group(1) if match.groups() else None
            op = _extract_op_name(captured)
            if op and op not in found:
                found.append(op)

    if found:
        return found
    if saw_marker:
        logger.warning(
            "Pre-check saw an unsupported-op marker but could not extract a concrete "
            "operator name; treating as inconclusive and allowing the job to proceed "
            "(fail-open)."
        )
    return []


def build_error_message(platform: str, ops: Iterable[str]) -> str:
    """產生使用者看得懂的錯誤訊息。"""
    op_list = [o for o in ops if o]
    if op_list:
        ops_text = ", ".join(op_list)
        op_clause = f"operator {ops_text}"
    else:
        op_clause = "某個 operator"
    return (
        f"platform {platform} 不支援 {op_clause}，"
        f"請改用支援的 platform（如 720/730）或將該 operator 移至 host 端後處理。"
    )


def run_precheck(
    onnx_path: str,
    *,
    model_id: int,
    version: str,
    platform: str,
    evaluator=None,
) -> Optional[str]:
    """對 ``onnx_path`` 跑 platform/operator 相容性檢查。

    偵測到「明確不支援」的 op → raise :class:`UnsupportedOperatorError`。
    其餘情況（分析成功且全支援、或分析失敗但無明確訊號）→ 放行。

    放行時回傳 ``evaluate()`` 的原始報告字串（成功分析）或 ``None``
    （evaluate() raise 但無明確不支援訊號）。回傳值讓呼叫端可以**重用**這份
    報告，避免 ``enable_precheck`` 與 ``enable_evaluate`` 同開時重複跑一次
    ``evaluate()``（compiler frontend 分析成本高、秒~數十秒級）。

    Args:
        onnx_path: 已經過 onnx2onnx_flow 的 out.onnx 路徑。
        model_id / version / platform: 目標編譯設定。
        evaluator: 可注入的 evaluator backend（測試用）。None 時使用預設 Kneron backend。

    Returns:
        evaluate() 的報告字串（放行且分析成功），或 None（放行但分析失敗無訊號）。
    """
    if evaluator is None:
        from services.backends.evaluator import get_evaluator_backend

        evaluator = get_evaluator_backend()

    report: Optional[str] = None
    try:
        report = evaluator.evaluate(
            onnx_path,
            model_id=int(model_id),
            version=str(version),
            platform=str(platform),
        )
    except Exception as exc:  # noqa: BLE001 - 需要檢查例外訊息內容
        # evaluate() 在遇到不支援 op 時可能直接 raise（例如 AssertionError）。
        # 用例外訊息掃描明確訊號；掃不到就當成「分析失敗但不確定」→ 放行（不誤擋）。
        message = str(exc)
        ops = find_unsupported_operators(message)
        if ops:
            raise UnsupportedOperatorError(build_error_message(platform, ops)) from exc
        logger.warning(
            "Compatibility pre-check inconclusive (evaluator raised, no explicit "
            "unsupported-op signal); allowing job to proceed. platform=%s error=%s",
            platform,
            message,
        )
        return None

    ops = find_unsupported_operators(report)
    if ops:
        raise UnsupportedOperatorError(build_error_message(platform, ops))

    logger.info("Compatibility pre-check passed. platform=%s", platform)
    return report