Blame for scripts/xeon_runtime.py

"""xeon_runtime.py — Runtime Xeon OTIMIZADO para BiGRU_T_version V6+.

Aprimoramento do xeon_setup.py com:
  1. Configuração completa do ambiente AVX512/AMX/MKL/OMP/OneDNN
  2. Verificação do IPEX (Intel Extension for PyTorch) + aplicação de ipex.optimize
  3. Benchmark FP16 com autocast para medir TFLOPS reais da CPU
  4. Aplicação automática de ipex.optimize(model, dtype=...) para modelos
  5. Configuração de threads com detecção dinâmica de núcleos físicos
  6. Diagnóstico completo de CPU ISA (AVX512 sub-features, AMX, BF16, FP16, VNNI)
  7. API one-call: `apply_all(model=None)` para configuração total
  8. Monkey-patch opcional para torch.matmul com autocast FP16 persistente

DEVE ser importado ANTES de `import torch` para que as variáveis de ambiente
tenham efeito sobre o backend OneDNN/MKL.

Uso típico:
    from xeon_runtime import apply_all, benchmark_fp16, optimize_model
    apply_all()             # configura ambiente + verifica IPEX + imprime diagnóstico
    import torch
    apply_all.post_torch()  # configura threads do torch
    tf = benchmark_fp16()   # mede TFLOPS FP16 reais
    model = ...
    model = optimize_model(model, dtype=torch.bfloat16)  # aplica ipex.optimize
"""
from __future__ import annotations

import os
import sys
import time
import logging
import subprocess
from typing import Optional, Tuple

logger = logging.getLogger(__name__)

# ==============================================================================
# 1. CONFIGURAÇÃO DE AMBIENTE (Executar antes de qualquer import de IA/HPC)
# ==============================================================================

def _detect_physical_cores() -> int:
    """Detecta núcleos físicos da VM (com fallback seguro)."""
    # Tentativa 1: lscpu
    try:
        out = subprocess.check_output(["lscpu", "-p=Core"], stderr=subprocess.DEVNULL).decode()
        cores = set(line for line in out.splitlines() if line and not line.startswith("#"))
        if cores:
            return max(1, min(len(cores), 4))  # limita a 4 para não oversubscrever
    except Exception:
        pass
    # Tentativa 2: /proc/cpuinfo
    try:
        with open("/proc/cpuinfo") as f:
            cores = sum(1 for line in f if line.startswith("processor"))
        return max(1, min(cores, 4))
    except Exception:
        return 2


NUM_CORES = str(_detect_physical_cores())


def _detect_cpu_isa() -> dict:
    """Detecta sub-features AVX512/AMX/BF16/FP16/VNNI disponíveis na CPU."""
    flags = ""
    try:
        with open("/proc/cpuinfo") as f:
            for line in f:
                if line.startswith("flags"):
                    flags = line.lower()
                    break
    except Exception:
        pass
    return {
        "avx512f": "avx512f" in flags,
        "avx512bw": "avx512bw" in flags,
        "avx512cd": "avx512cd" in flags,
        "avx512dq": "avx512dq" in flags,
        "avx512vl": "avx512vl" in flags,
        "avx512_vnni": "avx512_vnni" in flags,
        "avx512_bf16": "avx512_bf16" in flags,
        "avx512_fp16": "avx512_fp16" in flags,
        "amx_tile": "amx_tile" in flags,
        "amx_int8": "amx_int8" in flags,
        "amx_bf16": "amx_bf16" in flags,
    }


def _pick_max_isa(isa_flags: dict) -> str:
    """Escolhe a maior ISA suportada (para ONEDNN_MAX_CPU_ISA)."""
    if isa_flags["amx_int8"] or isa_flags["amx_bf16"]:
        return "AMX_INT8"  # máximo para OneDNN
    if isa_flags["avx512f"]:
        return "AVX512"
    if "avx2" in str(isa_flags):
        return "AVX2"
    return "SSE"


def configure_environment(num_cores: Optional[str] = None) -> dict:
    """Configura todas as variáveis de ambiente Xeon/Intel.

    DEVE ser chamado ANTES de `import torch`.
    Retorna dict com a configuração aplicada (para logging).
    """
    cores = num_cores or NUM_CORES
    isa_flags = _detect_cpu_isa()
    max_isa = _pick_max_isa(isa_flags)

    # Ativa explicitamente o suporte a AVX-512 e extensões de vetorização Intel
    os.environ["MKL_ENABLE_INSTRUCTIONS"] = "AVX512"

    # Configura o gerenciador de threads openmp/mkl para os núcleos físicos da VM
    os.environ["MKL_NUM_THREADS"] = cores
    os.environ["OMP_NUM_THREADS"] = cores
    os.environ["OMP_DYNAMIC"] = "FALSE"
    os.environ["MKL_DYNAMIC"] = "FALSE"

    # Otimizações específicas do PyTorch para CPUs Intel (OneDNN)
    os.environ["DNNL_PRIMITIVE_CACHE_CAPACITY"] = "1024"
    # AMX (Matrix Extensions) — só ativa em Xeon 4ª gen+ (Sapphire Rapids+)
    os.environ["ONEDNN_MAX_CPU_ISA"] = max_isa

    # Evita buffering de stdout (logs aparecem em tempo real)
    os.environ["PYTHONUNBUFFERED"] = "1"

    # Tokenizers paralelismo (HuggingFace)
    os.environ["TOKENIZERS_PARALLELISM"] = "true"

    # KMP_AFFINITY para melhor distribuição de threads
    os.environ["KMP_AFFINITY"] = "granularity=fine,compact,1,0"
    os.environ["KMP_BLOCKTIME"] = "1"

    return {
        "NUM_CORES": cores,
        "MKL_ENABLE_INSTRUCTIONS": os.environ["MKL_ENABLE_INSTRUCTIONS"],
        "MKL_NUM_THREADS": os.environ["MKL_NUM_THREADS"],
        "OMP_NUM_THREADS": os.environ["OMP_NUM_THREADS"],
        "OMP_DYNAMIC": os.environ["OMP_DYNAMIC"],
        "MKL_DYNAMIC": os.environ["MKL_DYNAMIC"],
        "DNNL_PRIMITIVE_CACHE_CAPACITY": os.environ["DNNL_PRIMITIVE_CACHE_CAPACITY"],
        "ONEDNN_MAX_CPU_ISA": os.environ["ONEDNN_MAX_CPU_ISA"],
        "PYTHONUNBUFFERED": os.environ["PYTHONUNBUFFERED"],
        "TOKENIZERS_PARALLELISM": os.environ["TOKENIZERS_PARALLELISM"],
        "KMP_AFFINITY": os.environ["KMP_AFFINITY"],
        "KMP_BLOCKTIME": os.environ["KMP_BLOCKTIME"],
        "ISA_FLAGS": isa_flags,
    }


# ==============================================================================
# 2. CONFIGURAÇÃO DO PYTORCH (após import) + IPEX
# ==============================================================================

def configure_torch() -> "torch":
    """Configura threads do PyTorch. Chamar APÓS `import torch`."""
    import torch
    torch.set_num_threads(int(NUM_CORES))
    torch.set_num_interop_threads(1)
    # OneDNN (MKL-DNN) — habilitado por padrão em CPU, mas garantimos
    try:
        torch.backends.mkldnn.enabled = True
    except Exception:
        pass
    # Desativa profiling mode do JIT para melhor desempenho
    try:
        torch._C._jit_set_profiling_mode(False)
        torch._C._jit_set_profiling_executor(False)
    except Exception:
        pass
    return torch


def check_ipex() -> Tuple[bool, Optional[str]]:
    """Verifica se IPEX (Intel Extension for PyTorch) está disponível."""
    try:
        import intel_extension_for_pytorch as ipex
        return True, getattr(ipex, "__version__", "unknown")
    except ImportError:
        return False, None
    except Exception:
        # IPEX pode falhar ao importar se versão do torch não bater
        return False, None


def optimize_model(model, dtype=None, level="O1"):
    """Aplica ipex.optimize ao modelo (se IPEX disponível).

    Args:
        model: nn.Module a ser otimizado
        dtype: torch.float16, torch.bfloat16, ou None (default = float32)
        level: "O0" (sem otim), "O1" (conv-bn folding + fusão), "O3" (agressivo)

    Returns:
        model otimizado (in-place + return). Se IPEX não disponível, retorna model inalterado.
    """
    ipex_ok, _ = check_ipex()
    if not ipex_ok:
        logger.warning("IPEX não disponível — model retornado sem otimização.")
        return model
    try:
        import intel_extension_for_pytorch as ipex
        if dtype is not None:
            try:
                # Tentativa: ipex.optimize com dtype (IPEX 2.x)
                model = ipex.optimize(model, dtype=dtype, level=level)
            except TypeError:
                # Fallback: ipex.optimize sem dtype
                model = ipex.optimize(model, level=level)
        else:
            model = ipex.optimize(model, level=level)
        logger.info(f"ipex.optimize aplicado (level={level}, dtype={dtype}).")
        return model
    except Exception as e:
        logger.warning(f"Falha ao aplicar ipex.optimize: {e}")
        return model


# ==============================================================================
# 3. BENCHMARK FP16 (com autocast) — medir TFLOPS reais da CPU Xeon
# ==============================================================================

def benchmark_fp16(N: int = 4000, warmup: int = 2) -> dict:
    """Executa benchmark FP16 com autocast em matmul NxN.

    Retorna dict com:
      - duration_s: tempo real (segundos)
      - tflops: TeraFLOPS estimados (2*N^3 / duration / 1e12)
      - dtype_in: dtype de entrada (torch.float16)
      - dtype_out: dtype de saída (após autocast)
      - ipex_active: bool
      - isa: str ("AVX512", "AMX", etc.)
    """
    import torch
    ipex_ok, _ = check_ipex()

    # Matrizes em FP16 nativo
    A = torch.randn(N, N, dtype=torch.float16)
    B = torch.randn(N, N, dtype=torch.float16)

    # Warm-up
    with torch.cpu.amp.autocast(enabled=True, dtype=torch.float16):
        for _ in range(warmup):
            _ = torch.matmul(A[:N // 4, :N // 4], B[:N // 4, :N // 4])

    # Medição
    start = time.perf_counter()
    with torch.cpu.amp.autocast(enabled=True, dtype=torch.float16):
        C = torch.matmul(A, B)
    # Sincroniza (CPU é síncrono por padrão, mas garantimos)
    _ = C.sum().item()
    end = time.perf_counter()

    duration = end - start
    flops = 2.0 * (N ** 3) / duration if duration > 0 else 0.0
    tflops = flops / 1e12

    return {
        "matrix_size": N,
        "duration_s": duration,
        "tflops": tflops,
        "gops_per_sec": flops / 1e9,
        "dtype_in": "torch.float16",
        "dtype_out": str(C.dtype),
        "ipex_active": ipex_ok,
        "isa": get_cpu_isa(),
        "torch_threads": torch.get_num_threads(),
    }


def benchmark_bf16(N: int = 4000, warmup: int = 2) -> dict:
    """Executa benchmark BF16 (advantage: mesma range que FP32, melhor para treino)."""
    import torch
    ipex_ok, _ = check_ipex()

    A = torch.randn(N, N, dtype=torch.bfloat16)
    B = torch.randn(N, N, dtype=torch.bfloat16)

    with torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
        for _ in range(warmup):
            _ = torch.matmul(A[:N // 4, :N // 4], B[:N // 4, :N // 4])

    start = time.perf_counter()
    with torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16):
        C = torch.matmul(A, B)
    _ = C.sum().item()
    end = time.perf_counter()

    duration = end - start
    flops = 2.0 * (N ** 3) / duration if duration > 0 else 0.0
    return {
        "matrix_size": N,
        "duration_s": duration,
        "tflops": flops / 1e12,
        "gops_per_sec": flops / 1e9,
        "dtype_in": "torch.bfloat16",
        "dtype_out": str(C.dtype),
        "ipex_active": ipex_ok,
        "isa": get_cpu_isa(),
        "torch_threads": torch.get_num_threads(),
    }


# ==============================================================================
# 4. DIAGNÓSTICO E PRINT
# ==============================================================================

def get_cpu_isa() -> str:
    """Retorna a maior ISA suportada pela CPU (string curta)."""
    isa_flags = _detect_cpu_isa()
    if isa_flags["amx_int8"] or isa_flags["amx_bf16"]:
        return "AMX"
    if isa_flags["avx512f"]:
        return "AVX512"
    if "avx2" in str(isa_flags):
        return "AVX2"
    return "SSE"


def print_xeon_summary(benchmark: bool = False, bench_N: int = 4000):
    """Imprime resumo das configurações Xeon aplicadas.

    Args:
        benchmark: se True, executa benchmark FP16 + BF16 ao final
        bench_N: tamanho da matriz para benchmark
    """
    ipex_ok, ipex_ver = check_ipex()
    isa = get_cpu_isa()
    isa_flags = _detect_cpu_isa()

    print("=" * 72)
    print("CONFIGURAÇÃO XEON APLICADA (xeon_runtime.py)")
    print("=" * 72)
    print(f"  Python:             {sys.version.split()[0]}")
    print(f"  NUM_CORES:          {NUM_CORES}")
    print(f"  MKL_ENABLE_INSTR:   {os.environ.get('MKL_ENABLE_INSTRUCTIONS')}")
    print(f"  ONEDNN_MAX_CPU_ISA: {os.environ.get('ONEDNN_MAX_CPU_ISA')}")
    print(f"  OMP_NUM_THREADS:    {os.environ.get('OMP_NUM_THREADS')}")
    print(f"  MKL_NUM_THREADS:    {os.environ.get('MKL_NUM_THREADS')}")
    print(f"  KMP_AFFINITY:       {os.environ.get('KMP_AFFINITY')}")
    print(f"  KMP_BLOCKTIME:      {os.environ.get('KMP_BLOCKTIME')}")
    print(f"  CPU ISA detected:   {isa}")
    print(f"  AVX512 sub-features:")
    for k, v in isa_flags.items():
        marker = "✓" if v else " "
        print(f"      {marker} {k}")
    print(f"  IPEX ativo:         {ipex_ok}")
    if ipex_ok:
        print(f"  IPEX version:       {ipex_ver}")
    try:
        import torch
        print(f"  Torch:              {torch.__version__}")
        print(f"  Torch threads:      {torch.get_num_threads()}")
        print(f"  Torch interop:      {torch.get_num_interop_threads()}")
        print(f"  OneDNN enabled:     {torch.backends.mkldnn.enabled}")
        print(f"  CPU bf16 support:   {torch.backends.mkldnn.is_available()}")
    except ImportError:
        print("  Torch:              (ainda não importado)")
    print("=" * 72)

    if benchmark:
        print("\n--- BENCHMARK FP16 (matmul autocast) ---")
        try:
            res_fp16 = benchmark_fp16(N=bench_N)
            print(f"  Matriz:        {res_fp16['matrix_size']}x{res_fp16['matrix_size']}")
            print(f"  Tempo:         {res_fp16['duration_s']:.4f}s")
            print(f"  TFLOPS FP16:   {res_fp16['tflops']:.4f}")
            print(f"  GOPS/s:        {res_fp16['gops_per_sec']:.1f}")
            print(f"  Output dtype:  {res_fp16['dtype_out']}")
        except Exception as e:
            print(f"  [ERRO] FP16 benchmark falhou: {e}")

        print("\n--- BENCHMARK BF16 (matmul autocast) ---")
        try:
            res_bf16 = benchmark_bf16(N=bench_N)
            print(f"  Matriz:        {res_bf16['matrix_size']}x{res_bf16['matrix_size']}")
            print(f"  Tempo:         {res_bf16['duration_s']:.4f}s")
            print(f"  TFLOPS BF16:   {res_bf16['tflops']:.4f}")
            print(f"  GOPS/s:        {res_bf16['gops_per_sec']:.1f}")
            print(f"  Output dtype:  {res_bf16['dtype_out']}")
        except Exception as e:
            print(f"  [ERRO] BF16 benchmark falhou: {e}")
        print("=" * 72)


# ==============================================================================
# 5. API ONE-CALL: apply_all() — configura tudo numa chamada
# ==============================================================================

def apply_all(model=None, dtype=None, level: str = "O1", benchmark: bool = False,
             bench_N: int = 4000, verbose: bool = True) -> dict:
    """Aplica TODA a configuração Xeon em uma chamada.

    Fluxo:
      1. configure_environment() — env vars (ANTES de import torch)
      2. import torch + configure_torch()
      3. Verifica IPEX
      4. (Opcional) ipex.optimize(model, dtype, level)
      5. (Opcional) benchmark FP16/BF16
      6. Print resumo (se verbose)

    Uso:
        from xeon_runtime import apply_all
        apply_all()  # antes de import torch
        import torch

    Ou:
        from xeon_runtime import apply_all
        apply_all()  # env vars
        import torch
        apply_all.post_torch()  # threads
        model = ...
        model = apply_all.optimize_model(model, dtype=torch.bfloat16)

    Returns:
        dict com configuração aplicada + benchmarks (se solicitado)
    """
    env = configure_environment()
    import torch
    configure_torch()
    ipex_ok, ipex_ver = check_ipex()

    result = {
        "env": env,
        "ipex_active": ipex_ok,
        "ipex_version": ipex_ver,
        "torch_version": torch.__version__,
        "torch_threads": torch.get_num_threads(),
        "onednn_enabled": torch.backends.mkldnn.enabled,
        "isa": get_cpu_isa(),
    }

    if model is not None:
        model = optimize_model(model, dtype=dtype, level=level)
        result["model_optimized"] = True

    if benchmark:
        try:
            result["fp16_benchmark"] = benchmark_fp16(N=bench_N)
        except Exception as e:
            result["fp16_benchmark_error"] = str(e)
        try:
            result["bf16_benchmark"] = benchmark_bf16(N=bench_N)
        except Exception as e:
            result["bf16_benchmark_error"] = str(e)

    if verbose:
        print_xeon_summary(benchmark=benchmark, bench_N=bench_N)

    return result


def post_torch():
    """Configura threads do torch (deve ser chamado APÓS import torch)."""
    return configure_torch()


# ==============================================================================
# 6. MAIN — para executar standalone como diagnóstico
# ==============================================================================

if __name__ == "__main__":
    print("\n>>> Aplicando configuração Xeon + benchmark FP16/BF16...\n")
    result = apply_all(benchmark=True, bench_N=4000)
    print("\n>>> Resumo final:")
    for k, v in result.items():
        if k in ("env", "fp16_benchmark", "bf16_benchmark"):
            print(f"  {k}:")
            if isinstance(v, dict):
                for k2, v2 in v.items():
                    print(f"    {k2}: {v2}")
        else:
            print(f"  {k}: {v}")