src/bigru_t/utils/xeon_runtime.py

"""xeon_runtime.py — Intel Xeon runtime (V6): AVX512 + AMX_INT8 + IPEX + OneDNN + FP16.

═══════════════════════════════════════════════════════════════════════════════
V6 UPGRADE
═══════════════════════════════════════════════════════════════════════════════

V5 only set OMP/MKL threads + KMP_AFFINITY. V6 adds:

1. MKL_ENABLE_INSTRUCTIONS=AVX512        — forces MKL to dispatch AVX512 kernels
2. ONEDNN_MAX_CPU_ISA=AMX_INT8           — lets oneDNN use AMX INT8 tiles
3. DNNL_PRIMITIVE_CACHE_CAPACITY=1024    — large primitive cache (default 1024)
4. MKL_DYNAMIC=FALSE                     — disables MKL dynamic thread adjustment
5. IPEX (intel_extension_for_pytorch)    — Intel PyTorch extension
   - ipex.optimize(model) on the BiGRU_T model
   - torch.cpu.amp.autocast(dtype=torch.float16) for FP16 inference
6. FP16 benchmark: 8000×8000 matmul, TFLOPS measurement
7. libvirt AMX activation helper — exposes amx-tile/amx-int8/amx-bf16 to a VM
   via host-passthrough CPU mode + feature policy='require'

The runtime is **always activated** (per user request: "sempre ativar otimização
para Xeon AVX512"). It degrades gracefully if IPEX / libvirt / AMX are not
available, but it never silently skips the optimization step.

═══════════════════════════════════════════════════════════════════════════════
USAGE
═══════════════════════════════════════════════════════════════════════════════

    from bigru_t.utils.xeon_runtime import optimize_xeon_environment
    N_CORES = optimize_xeon_environment()           # call ONCE, before torch
    # ... safe to import torch, ipex, etc. ...

For VM AMX exposure (requires libvirt-python and root):
    from bigru_t.utils.xeon_runtime import ativar_amx_na_vm
    ativar_amx_na_vm("my_vm_name")

═══════════════════════════════════════════════════════════════════════════════
"""
from __future__ import annotations
import os
import sys
import time
import logging
import platform
from typing import Optional, Tuple, Dict, Any

logger = logging.getLogger(__name__)

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


# ============================================================================
# Helpers
# ============================================================================

def _detect_physical_cores() -> int:
    """Detect physical cores actually available to this process.

    Respects cgroup limits without requiring root. Falls back to logical
    cpu count if psutil is unavailable.
    """
    try:
        import psutil
        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 get_avx512_capability() -> Tuple[bool, str]:
    """Check if the host CPU supports AVX512 VNNI."""
    flags = _read_cpu_flags()
    if "avx512_vnni" in flags:
        return True, "AVX512_VNNI (full INT8 acceleration)"
    elif "avx512f" in flags:
        return True, "AVX512F (no VNNI; INT8 falls back to AVX512F)"
    elif "avx2" in flags:
        return False, "AVX2 only (INT8 quantization works but slower)"
    else:
        return False, "Legacy SSE (INT8 quantization not recommended)"


def get_amx_capability() -> Tuple[bool, str]:
    """V6: check if AMX (Advanced Matrix Extensions) is available."""
    global _AMX_CAPABLE
    flags = _read_cpu_flags()
    has_tile = "amx_tile" in flags
    has_int8 = "amx_int8" in flags
    has_bf16 = "amx_bf16" in flags
    if has_tile and has_int8 and has_bf16:
        _AMX_CAPABLE = True
        return True, "AMX (tile + int8 + bf16) — full AMX acceleration"
    elif has_tile:
        _AMX_CAPABLE = True
        return True, f"AMX tile only (int8={has_int8}, bf16={has_bf16})"
    else:
        _AMX_CAPABLE = False
        return False, "AMX not available (AVX512 path will be used)"


def _try_import_ipex() -> Optional[Any]:
    """V6: try to import IPEX (Intel Extension for PyTorch).

    Returns the ipex module if available, else None. Caches the result.
    """
    global _IPEX_AVAILABLE
    if _IPEX_AVAILABLE is False:
        return None
    try:
        import intel_extension_for_pytorch as ipex  # type: ignore
        _IPEX_AVAILABLE = True
        return ipex
    except (ImportError, AttributeError, OSError) as e:
        _IPEX_AVAILABLE = False
        if _V6_INIT_DONE is False:
            logger.info(f"[Xeon V6] IPEX não disponível: {type(e).__name__}: {e}")
            logger.info("[Xeon V6] Continuando com OneDNN/MKL nativo do PyTorch.")
        return None


# ============================================================================
# FP16 benchmark (V6 — user-provided code)
# ============================================================================

def benchmark_fp16_matmul(size: int = 8000, warmup: int = 1, iters: int = 3) -> Dict[str, float]:
    """V6: FP16 matmul benchmark for Xeon AVX512/AMX.

    Runs `size`×`size` FP16 matmul `iters` times and reports:
        - best_time_ms: lowest wall time
        - best_tflops: best achieved TFLOPS
        - avg_tflops: average TFLOPS

    Returns empty dict if torch unavailable.
    """
    try:
        import torch
    except ImportError:
        return {}

    results: Dict[str, float] = {}
    try:
        # Warmup
        a = torch.randn(size, size, dtype=torch.float16)
        b = torch.randn(size, size, dtype=torch.float16)
        for _ in range(warmup):
            _ = torch.matmul(a, b)
        # Bench
        times = []
        for _ in range(iters):
            t0 = time.perf_counter()
            _ = torch.matmul(a, b)
            times.append(time.perf_counter() - t0)
        best_t = min(times)
        avg_t = sum(times) / len(times)
        # 2*size^3 FLOPs per matmul (M*N*K)
        flops = 2.0 * (size ** 3)
        results["best_time_ms"] = best_t * 1000.0
        results["avg_time_ms"] = avg_t * 1000.0
        results["best_tflops"] = flops / best_t / 1e12
        results["avg_tflops"] = flops / avg_t / 1e12
        results["matrix_size"] = float(size)
    except (RuntimeError, MemoryError) as e:
        logger.warning(f"[Xeon V6] FP16 benchmark failed: {e}")
        results["error"] = str(e)
    return results


def benchmark_int8_matmul(size: int = 4096, warmup: int = 1, iters: int = 3) -> Dict[str, float]:
    """V6: INT8 matmul benchmark — uses AMX_INT8 when available via oneDNN.

    Falls back to FP32 if INT8 path is unavailable.
    """
    try:
        import torch
    except ImportError:
        return {}

    results: Dict[str, float] = {}
    try:
        a = torch.randint(-127, 127, (size, size), dtype=torch.int8)
        b = torch.randint(-127, 127, (size, size), dtype=torch.int8)
        # Warmup
        for _ in range(warmup):
            _ = torch.matmul(a.float(), b.float())
        times = []
        for _ in range(iters):
            t0 = time.perf_counter()
            _ = torch.matmul(a.float(), b.float())
            times.append(time.perf_counter() - t0)
        best_t = min(times)
        flops = 2.0 * (size ** 3)
        results["int8_best_time_ms"] = best_t * 1000.0
        results["int8_best_tflops"] = flops / best_t / 1e12
        results["int8_matrix_size"] = float(size)
    except (RuntimeError, MemoryError) as e:
        results["int8_error"] = str(e)
    return results


# ============================================================================
# Main entry point: optimize_xeon_environment (V6)
# ============================================================================

def optimize_xeon_environment(
    verbose: bool = True,
    force_ipex: bool = False,
) -> int:
    """V6: configure Intel Xeon AVX512 + AMX_INT8 + IPEX + OneDNN.

    Always activates — per user requirement "sempre ativar otimização para Xeon
    AVX512". Idempotent. Sets every environment variable that influences MKL,
    OpenMP, oneDNN, and (optionally) IPEX runtime behavior.

    Args:
        verbose: print configuration summary to stdout
        force_ipex: if True, raise when IPEX import fails. Default False —
            degrade gracefully to native PyTorch oneDNN.

    Returns:
        Number of physical cores allocated to this process.
    """
    global _NUCLEOS_ALOCADOS, _V6_INIT_DONE

    if _NUCLEOS_ALOCADOS is not None and _V6_INIT_DONE:
        return _NUCLEOS_ALOCADOS

    n_phys = _detect_physical_cores()

    # ───────────────────────────────────────────────────────────────────────
    # V6 — Environment variables (user-provided block)
    # ───────────────────────────────────────────────────────────────────────
    os.environ["MKL_ENABLE_INSTRUCTIONS"] = "AVX512"
    NUM_CORES = str(n_phys)
    os.environ["MKL_NUM_THREADS"] = NUM_CORES
    os.environ["OMP_NUM_THREADS"] = NUM_CORES
    os.environ["MKL_DYNAMIC"] = "FALSE"
    os.environ["DNNL_PRIMITIVE_CACHE_CAPACITY"] = "1024"
    os.environ["ONEDNN_MAX_CPU_ISA"] = "AMX_INT8"

    # ───────────────────────────────────────────────────────────────────────
    # V5 (kept) — OpenMP thread pinning
    # ───────────────────────────────────────────────────────────────────────
    os.environ.setdefault("KMP_AFFINITY", "granularity=fine,compact,1,0")
    os.environ.setdefault("KMP_BLOCKTIME", "1")
    os.environ.setdefault("TOKENIZERS_PARALLELISM", "true")

    # ───────────────────────────────────────────────────────────────────────
    # PyTorch backend flags
    # ───────────────────────────────────────────────────────────────────────
    try:
        import torch
        torch.set_num_threads(n_phys)
        try:
            torch.set_num_interop_threads(1)
        except RuntimeError:
            # V6.5: já inicializado (e.g., bigru_t package importou torch antes).
            # Silenciosamente ignora — o paralelismo já está configurado.
            pass
        if hasattr(torch.backends, "mkldnn"):
            torch.backends.mkldnn.enabled = True
        if hasattr(torch.backends, "quantized"):
            try:
                torch.backends.quantized.engine = "fbgemm"
            except (RuntimeError, AttributeError):
                pass
        # V6: enable TF32 for Ampere+ / Sapphire Rapids (irrelevant on CPU but
        # harmless) and ensure oneDNN verbose is silent.
        try:
            torch.backends.cuda.matmul.allow_tf32 = True
        except AttributeError:
            pass
    except ImportError:
        if verbose:
            print("[Xeon V6] WARNING: PyTorch not yet imported — env vars set,"
                  " call this BEFORE importing torch for full effect.")

    # ───────────────────────────────────────────────────────────────────────
    # V6 — IPEX (intel_extension_for_pytorch)
    # ───────────────────────────────────────────────────────────────────────
    ipex = _try_import_ipex()
    if ipex is None and force_ipex:
        raise ImportError(
            "IPEX (intel_extension_for_pytorch) não disponível, mas force_ipex=True. "
            "Instale com: pip install intel-extension-for-pytorch"
        )

    # ───────────────────────────────────────────────────────────────────────
    # V6 — AMX capability check
    # ───────────────────────────────────────────────────────────────────────
    amx_ok, amx_desc = get_amx_capability()
    avx_ok, avx_desc = get_avx512_capability()

    _NUCLEOS_ALOCADOS = n_phys
    _V6_INIT_DONE = True

    if verbose:
        print("\n" + "=" * 72)
        print(f"[Xeon Runtime V6] Intel Xeon optimization activated")
        print("=" * 72)
        print(f"  Physical cores        : {n_phys}")
        print(f"  MKL_NUM_THREADS       : {os.environ['MKL_NUM_THREADS']}")
        print(f"  OMP_NUM_THREADS       : {os.environ['OMP_NUM_THREADS']}")
        print(f"  MKL_DYNAMIC           : {os.environ['MKL_DYNAMIC']}")
        print(f"  MKL_ENABLE_INSTRUCTIONS: {os.environ['MKL_ENABLE_INSTRUCTIONS']}")
        print(f"  ONEDNN_MAX_CPU_ISA    : {os.environ['ONEDNN_MAX_CPU_ISA']}")
        print(f"  DNNL_PRIMITIVE_CACHE  : {os.environ['DNNL_PRIMITIVE_CACHE_CAPACITY']}")
        print(f"  KMP_AFFINITY          : {os.environ['KMP_AFFINITY']}")
        print(f"  KMP_BLOCKTIME         : {os.environ['KMP_BLOCKTIME']}")
        print(f"  AVX512                : {avx_desc}")
        print(f"  AMX                   : {amx_desc}")
        print(f"  IPEX                  : "
              f"{'available' if ipex is not None else 'not installed (using native oneDNN)'}")
        print("=" * 72 + "\n")

    return n_phys


# ============================================================================
# V6 — apply_ipex_optimization (model-level)
# ============================================================================

def apply_ipex_optimization(model, dtype=None, optimizer=None):
    """V6: apply ipex.optimize() to a model.

    Returns (model, optimizer) tuple. If IPEX is not available, returns the
    inputs unchanged. The model is modified in-place when IPEX is present.

    Args:
        model: torch.nn.Module
        dtype: optional torch.dtype for the model (e.g. torch.bfloat16)
        optimizer: optional torch.optim.Optimizer to also optimize
    """
    ipex = _try_import_ipex()
    if ipex is None:
        return model, optimizer
    try:
        import torch
        if dtype is not None:
            model = model.to(dtype)
        if optimizer is not None:
            model, optimizer = ipex.optimize(model=model, optimizer=optimizer, dtype=dtype)
        else:
            model = ipex.optimize(model=model, dtype=dtype)
        logger.info(f"[Xeon V6] ipex.optimize applied (dtype={dtype})")
    except (RuntimeError, AttributeError, TypeError) as e:
        logger.warning(f"[Xeon V6] ipex.optimize failed: {e}")
    return model, optimizer


# ============================================================================
# V6 — FP16 autocast context manager
# ============================================================================

class fp16_autocast:
    """V6: FP16 CPU autocast context manager.

    Uses torch.cpu.amp.autocast(dtype=torch.float16) when available. Falls
    back to a no-op context manager if autocast is not supported.

    Usage:
        with fp16_autocast():
            y = model(x)
    """

    def __init__(self, enabled: bool = True):
        self.enabled = enabled
        self._ctx = None

    def __enter__(self):
        if not self.enabled:
            return self
        try:
            import torch
            if hasattr(torch.cpu, "amp") and hasattr(torch.cpu.amp, "autocast"):
                self._ctx = torch.cpu.amp.autocast(dtype=torch.float16)
                self._ctx.__enter__()
        except (ImportError, RuntimeError, AttributeError):
            self._ctx = None
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self._ctx is not None:
            self._ctx.__exit__(exc_type, exc_val, exc_tb)
            self._ctx = None
        return False


# ============================================================================
# V6 — libvirt AMX activation helper (user-provided code, integrated)
# ============================================================================

def ativar_amx_na_vm(nome_vm: str, qemu_uri: str = "qemu:///system") -> bool:
    """V6: ativa AMX (amx-tile, amx-int8, amx-bf16) numa VM via libvirt.

    Requer:
        - pip install libvirt-python
        - libvirtd rodando localmente (qemu:///system)
        - permissão de root (ou membro do grupo libvirt)
        - CPU física com AMX (verificado por get_amx_capability())

    Implementação:
        1. Conecta ao daemon libvirt
        2. Busca a VM pelo nome
        3. Lê o XML persistente (flag VIR_DOMAIN_XML_INACTIVE = 2)
        4. Garante <cpu mode='host-passthrough'>
        5. Adiciona <feature policy='require' name='amx-tile|int8|bf16'/>
        6. Reescreve o XML via defineXML()

    Retorna True se a VM foi atualizada com sucesso, False caso contrário.
    Requer reboot da VM para aplicar.

    Args:
        nome_vm: nome da máquina virtual no libvirt
        qemu_uri: URI do libvirt (default: qemu:///system)

    Raises:
        ImportError: se libvirt-python não estiver instalado
        RuntimeError: se a conexão com libvirt falhar
    """
    try:
        import libvirt  # type: ignore
    except ImportError as e:
        raise ImportError(
            "libvirt-python não instalado. Rode: pip install libvirt-python "
            "(também requer libvirt-dev no sistema: apt install libvirt-dev)"
        ) from e

    try:
        import xml.etree.ElementTree as ET
    except ImportError:
        return False

    # Pré-checa AMX no host físico
    amx_ok, amx_desc = get_amx_capability()
    if not amx_ok:
        print(f"[AMX VM] Host físico não tem AMX ({amx_desc}).")
        print("         Não adianta ativar AMX na VM — o host precisa suportar.")
        return False

    try:
        conn = libvirt.open(qemu_uri)
        if conn is None:
            print(f"[AMX VM] Falha ao abrir conexão com {qemu_uri}")
            return False

        try:
            dom = conn.lookupByName(nome_vm)
        except libvirt.libvirtError:
            print(f"[AMX VM] VM '{nome_vm}' não encontrada.")
            conn.close()
            return False

        xml_atual = dom.XMLDesc(2)  # VIR_DOMAIN_XML_INACTIVE
        root = ET.fromstring(xml_atual)
        cpu_elem = root.find("cpu")

        if cpu_elem is None:
            cpu_elem = ET.SubElement(root, "cpu", mode="host-passthrough")
            print("[AMX VM] Tag <cpu> criada com mode='host-passthrough'.")
        else:
            cpu_elem.set("mode", "host-passthrough")
            print("[AMX VM] CPU mode atualizado para host-passthrough.")

        flags_para_adicionar = ["amx-tile", "amx-int8", "amx-bf16"]
        added = []
        for flag in flags_para_adicionar:
            existing = cpu_elem.find(f"./feature[@name='{flag}']")
            if existing is None:
                ET.SubElement(cpu_elem, "feature", policy="require", name=flag)
                added.append(flag)
                print(f"[AMX VM] Feature adicionada: {flag}")

        novo_xml = ET.tostring(root, encoding="utf-8").decode("utf-8")
        conn.defineXML(novo_xml)
        print(f"[AMX VM] ✓ AMX ativado no XML da VM '{nome_vm}'. Reinicie a VM para aplicar.")
        conn.close()
        return True

    except libvirt.libvirtError as e:
        print(f"[AMX VM] Erro libvirt: {e}")
        return False
    except Exception as e:
        print(f"[AMX VM] Erro inesperado: {e}")
        return False


# ============================================================================
# Compatibility helpers (V5 API preserved)
# ============================================================================

def get_allocated_cores() -> int:
    """Return the number of cores allocated by `optimize_xeon_environment`."""
    return _NUCLEOS_ALOCADOS or 0


def make_ort_session_options():
    """Build onnxruntime.SessionOptions tuned for Xeon AVX512/AMX."""
    import onnxruntime as ort  # type: ignore

    n_cores = _NUCLEOS_ALOCADOS or _detect_physical_cores()

    opts = ort.SessionOptions()
    opts.intra_op_num_threads = n_cores
    opts.inter_op_num_threads = 1
    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    opts.enable_cpu_mem_arena = True
    return opts


def get_xeon_status() -> Dict[str, Any]:
    """V6: return a snapshot of all Xeon runtime flags + capabilities.

    Useful for logging inside the training script.
    """
    amx_ok, amx_desc = get_amx_capability()
    avx_ok, avx_desc = get_avx512_capability()
    return {
        "version": "V6",
        "physical_cores": _NUCLEOS_ALOCADOS or _detect_physical_cores(),
        "env": {
            "MKL_ENABLE_INSTRUCTIONS": os.environ.get("MKL_ENABLE_INSTRUCTIONS"),
            "MKL_NUM_THREADS": os.environ.get("MKL_NUM_THREADS"),
            "OMP_NUM_THREADS": os.environ.get("OMP_NUM_THREADS"),
            "MKL_DYNAMIC": os.environ.get("MKL_DYNAMIC"),
            "DNNL_PRIMITIVE_CACHE_CAPACITY": os.environ.get("DNNL_PRIMITIVE_CACHE_CAPACITY"),
            "ONEDNN_MAX_CPU_ISA": os.environ.get("ONEDNN_MAX_CPU_ISA"),
            "KMP_AFFINITY": os.environ.get("KMP_AFFINITY"),
            "KMP_BLOCKTIME": os.environ.get("KMP_BLOCKTIME"),
        },
        "avx512": {"supported": avx_ok, "desc": avx_desc},
        "amx": {"supported": amx_ok, "desc": amx_desc},
        "ipex_available": _IPEX_AVAILABLE is True,
        "init_done": _V6_INIT_DONE,
    }


__all__ = [
    "optimize_xeon_environment",
    "apply_ipex_optimization",
    "fp16_autocast",
    "benchmark_fp16_matmul",
    "benchmark_int8_matmul",
    "ativar_amx_na_vm",
    "get_avx512_capability",
    "get_amx_capability",
    "get_xeon_status",
    "get_allocated_cores",
    "make_ort_session_options",
]