V6+: add xeon_runtime.py with IPEX/AMX/AVX512 + FP16/BF16 benchmark

scripts/xeon_runtime.py

@@ -0,0 +1,472 @@
+"""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}")