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）。

設計原則（避免誤擋）：
- 只有在 evaluator 報告中出現「明確的」不支援訊號時才擋。
- 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 名稱」（抓不到也沒關係，會退回泛用訊息）。
_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"
    ("hw_not_support", r"(?:HW[_ ]NOT[_ ]SUPPORT|Hardware\s*not\s*support|HardwareNotSupport)\s*[:\-]?\s*([A-Za-z0-9_]+)?"),
    # "UNIMPLEMENTED_FEATURE: ..." / "UnimplementedFeature: undefined CPU op [Softmax]"
    ("unimplemented", r"(?:UNIMPLEMENTED[_ ]FEATURE|UnimplementedFeature)\s*[:\-]?\s*(.*)"),
    # 泛用「HW not support」報告欄位（toolchain 報告字串）
    ("hw_not_support_col", r"HW\s+not\s+support\b"),
)

# 從一段訊息中撈出被中括號 / 引號包住的 op 名稱，例如 "undefined CPU op [Softmax]"。
_OP_IN_BRACKETS = re.compile(r"[\[\(]\s*([A-Za-z][A-Za-z0-9_]*)\s*[\]\)]")
# 常見 ONNX op 名稱（PascalCase / 首字母大寫），用來在自由文字裡挑一個像 op 的字。
_OP_TOKEN = re.compile(r"\b([A-Z][A-Za-z0-9]{2,})\b")


def _extract_op_name(fragment: Optional[str]) -> Optional[str]:
    """盡力從訊息片段中抽出 operator 名稱。抓不到就回 None。"""
    if not fragment:
        return None
    fragment = fragment.strip()
    if not fragment:
        return None

    m = _OP_IN_BRACKETS.search(fragment)
    if m:
        return m.group(1)

    # 直接就是一個 op token（例如 "Softmax"）
    if re.fullmatch(r"[A-Za-z][A-Za-z0-9_]*", fragment):
        return fragment

    m = _OP_TOKEN.search(fragment)
    if m:
        return m.group(1)

    return None


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

    - 沒有偵測到任何明確不支援訊號 → 回空 list（呼叫端應放行）。
    - 偵測到訊號但抽不出 op 名稱 → 回含一個占位符 ``"<unknown>"`` 的 list，
      讓呼叫端仍能 fail-fast（但訊息較泛用）。
    """
    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:
        # 有明確不支援訊號、但抽不到 op 名稱：仍回報（用占位符）。
        return ["<unknown>"]
    return []


def build_error_message(platform: str, ops: Iterable[str]) -> str:
    """產生使用者看得懂的錯誤訊息。"""
    op_list = [o for o in ops if o and o != "<unknown>"]
    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