BiGRU_T_version / scripts /xeon_runtime.py

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

e719654 verified 3 months ago

17.4 kB

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

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