"""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}")