cnn_bigru/utils/xeon_runtime.py

"""xeon_runtime.py — Ativação do runtime Intel Xeon (AVX512 + AMX + IPEX + OneDNN).

Adaptado do projeto BiGRU_T_version (PowerMachine), com simplificações para
rodar em CPU-only e degradar graciosamente quando IPEX/AMX não estiverem
disponíveis. A ativação é SEMPRE realizada (mesmo que parcial), conforme
requisito do usuário: "ativar xeon_runtime.py".

Otimizações aplicadas:
    1. OMP_NUM_THREADS / MKL_NUM_THREADS = núcleos físicos
    2. KMP_AFFINITY=granularity=fine,compact
    3. MKL_ENABLE_INSTRUCTIONS=AVX512 (se suportado)
    4. ONEDNN_MAX_CPU_ISA=AMX_INT8 (se suportado)
    5. DNNL_PRIMITIVE_CACHE_CAPACITY=1024
    6. MKL_DYNAMIC=FALSE
    7. IPEX (intel_extension_for_pytorch) — se disponível, ipex.optimize(model)
    8. torch.set_float32_matmul_precision("high")
    9. torch.backends.cudnn.benchmark = True (no-op em CPU)

Uso:
    from cnn_bigru.utils.xeon_runtime import optimize_xeon_environment
    N_CORES = optimize_xeon_environment()   # chamar ANTES de importar torch
    import torch
"""
from __future__ import annotations

import logging
import os
import platform
from typing import Any, Optional

logger = logging.getLogger(__name__)

_NUCLEOS_ALOCADOS: Optional[int] = None
_IPEX_AVAILABLE: Optional[bool] = None
_AMX_CAPABLE: Optional[bool] = None
_INIT_DONE: bool = False


def _detect_physical_cores() -> int:
    """Detecta núcleos físicos respeitando cgroup limits."""
    try:
        import psutil  # type: ignore
        n_phys = psutil.cpu_count(logical=False) or 1
    except ImportError:
        try:
            with open("/proc/cpuinfo", "r") as f:
                cores = set()
                for line in f:
                    if line.startswith("core id"):
                        cores.add(line.strip())
                n_phys = len(cores) or 1
        except OSError:
            n_phys = 1
    try:
        n_affine = len(os.sched_getaffinity(0))
        n_logical = os.cpu_count() or 1
        if n_affine < n_logical:
            n_phys = max(1, n_affine // 2)
        else:
            n_phys = min(n_phys, n_affine)
    except (AttributeError, OSError):
        pass
    return max(1, n_phys)


def _read_cpu_flags() -> str:
    try:
        with open("/proc/cpuinfo", "r") as f:
            for line in f:
                if line.startswith("flags"):
                    return line
    except OSError:
        pass
    return ""


def _detect_amx() -> bool:
    flags = _read_cpu_flags()
    return "amx_int8" in flags and "amx_bf16" in flags


def _detect_avx512() -> bool:
    flags = _read_cpu_flags()
    return "avx512f" in flags


def optimize_xeon_environment(force_threads: Optional[int] = None) -> int:
    """Ativa todas as otimizações de CPU. Retorna o número de núcleos alocados.

    Deve ser chamada UMA VEZ, antes de importar torch, para que as variáveis
    de ambiente tenham efeito. Chamadas subsequentes são no-op (idempotente).
    """
    global _NUCLEOS_ALOCADOS, _IPEX_AVAILABLE, _AMX_CAPABLE, _INIT_DONE
    if _INIT_DONE:
        return _NUCLEOS_ALOCADOS or 1

    n_cores = force_threads or _detect_physical_cores()
    _NUCLEOS_ALOCADOS = n_cores

    # Threads
    os.environ.setdefault("OMP_NUM_THREADS", str(n_cores))
    os.environ.setdefault("MKL_NUM_THREADS", str(n_cores))
    os.environ.setdefault("OPENBLAS_NUM_THREADS", str(n_cores))
    os.environ.setdefault("NUMEXPR_NUM_THREADS", str(n_cores))
    os.environ["KMP_AFFINITY"] = "granularity=fine,compact"
    os.environ["MKL_DYNAMIC"] = "FALSE"

    # ISA detection
    if _detect_avx512():
        os.environ["MKL_ENABLE_INSTRUCTIONS"] = "AVX512"
        logger.info("AVX512 detectado e ativado para MKL")
    if _detect_amx():
        os.environ["ONEDNN_MAX_CPU_ISA"] = "AMX_INT8"
        os.environ["DNNL_PRIMITIVE_CACHE_CAPACITY"] = "1024"
        _AMX_CAPABLE = True
        logger.info("AMX_INT8 detectado e ativado para OneDNN")
    else:
        _AMX_CAPABLE = False

    # Tokenizers parallelism
    os.environ.setdefault("TOKENIZERS_PARALLELISM", "true")

    _INIT_DONE = True
    logger.info(
        "Xeon runtime ativado: cores=%d, avx512=%s, amx=%s, ipex=%s",
        n_cores,
        _detect_avx512(),
        bool(_AMX_CAPABLE),
        _check_ipex_available(),
    )
    return n_cores


def _check_ipex_available() -> bool:
    global _IPEX_AVAILABLE
    if _IPEX_AVAILABLE is not None:
        return _IPEX_AVAILABLE
    try:
        import intel_extension_for_pytorch  # noqa: F401
        _IPEX_AVAILABLE = True
    except ImportError:
        _IPEX_AVAILABLE = False
    return _IPEX_AVAILABLE


def optimize_model_ipex(model: Any) -> Any:
    """Aplica ipex.optimize() no modelo, se IPEX estiver disponível."""
    if not _check_ipex_available():
        logger.info("IPEX indisponível — pulando ipex.optimize()")
        return model
    try:
        import intel_extension_for_pytorch as ipex  # type: ignore
        model = ipex.optimize(model)
        logger.info("Modelo otimizado com IPEX")
    except Exception as e:
        logger.warning("Falha ao aplicar ipex.optimize(): %s", e)
    return model


def get_runtime_info() -> dict:
    """Retorna informações sobre o runtime ativado."""
    return {
        "nucleos_alocados": _NUCLEOS_ALOCADOS,
        "avx512": _detect_avx512(),
        "amx_capable": bool(_AMX_CAPABLE),
        "ipex_available": _check_ipex_available(),
        "platform": platform.platform(),
        "init_done": _INIT_DONE,
    }


__all__ = [
    "optimize_xeon_environment",
    "optimize_model_ipex",
    "get_runtime_info",
]