"""Pinned precision policy and backend identity for request-owned state."""
from contextlib import contextmanager
import torch

ARITHMETIC_VERSION='current_recipe_b200_alignment_dense_attention_output_attention_8192_v8'


def execution_identity(device):
    device=torch.device(device)
    capability=tuple(torch.cuda.get_device_capability(device)) if device.type=='cuda' else None
    backend='native_b200_mma' if capability and capability==(10,0) else 'canonical_b200_mma'
    matmul=torch.backends.cuda.matmul
    return (ARITHMETIC_VERSION,str(torch.__version__),torch.version.cuda,str(device),capability,backend,
            bool(matmul.allow_tf32),bool(matmul.allow_bf16_reduced_precision_reduction),
            bool(matmul.allow_fp16_reduced_precision_reduction),torch.get_float32_matmul_precision())


@contextmanager
def release_precision():
    # The public facade holds its process-wide execution lock around this scope.
    # Both switches are restored even if evaluation raises.
    matmul=torch.backends.cuda.matmul
    tf32=matmul.allow_tf32;bf16=matmul.allow_bf16_reduced_precision_reduction
    try:
        matmul.allow_tf32=False
        matmul.allow_bf16_reduced_precision_reduction=False
        yield
    finally:
        matmul.allow_bf16_reduced_precision_reduction=bf16
        matmul.allow_tf32=tf32