"""trn2 e4m3 range handling (deepinfer/neuron_ops/fp8.py, issue #24), on the host."""

import torch

from deepinfer.neuron_ops import fp8

FP8 = torch.float8_e4m3fn


def test_halve_matches_torch_cast():
    b = torch.arange(256, dtype=torch.int32).to(torch.uint8)
    ref = (b.view(FP8).float() / 2).to(FP8).view(torch.uint8)
    got = fp8.halve_e4m3fn_bytes(b)
    nan = (b & 0x7F) == 0x7F
    assert torch.equal(got[~nan], ref[~nan])
    assert bool(((got[nan] & 0x7F) == 0x7F).all())


def test_halved_values_fit_hardware_e4m3():
    b = torch.arange(256, dtype=torch.int32).to(torch.uint8)
    got = fp8.halve_e4m3fn_bytes(b)
    finite = (b & 0x7F) != 0x7F
    assert bool((((got[finite] >> 3) & 15) < 15).all())  # exponent field 15 is inf/NaN on trn2


def test_rescale_keeps_dequantized_values():
    torch.manual_seed(0)
    w = torch.nn.Parameter((torch.randn(64, 128) * 100).clamp(-448, 448).to(FP8), requires_grad=False)
    w.sf = torch.randint(110, 130, (64, 4), dtype=torch.uint8)  # ue8m0 per 32 columns
    before = w.float() * torch.exp2(w.sf.float() - 127).repeat_interleave(32, 1)
    lossy = ((w.data.view(torch.uint8) >> 3) & 15) <= 1  # halving rounds only below the smallest normal
    model = torch.nn.Module()
    model.w = w
    if fp8.FP8_MAX > 240:
        return
    assert fp8.rescale_fp8_weights(model) == 1
    after = w.float() * torch.exp2(w.sf.float() - 127).repeat_interleave(32, 1)
    assert torch.equal(after[~lossy], before[~lossy])
    ulp = 2.0**-9 * torch.exp2(w.sf.float() - 127).repeat_interleave(32, 1)  # one smallest-subnormal step, after rescale
    assert bool(((after - before).abs()[lossy] <= ulp[lossy]).all())
    assert float(w.float().abs().max()) <= 224.0