test_batch_norm.py

from tileops.utils import get_backend_name, is_available

DEVICE = get_backend_name()
"""Tests for BatchNormFwdOp and BatchNormBwdOp.

Correctness is validated against torch.nn.functional.batch_norm and the
analytical gradient via torch.autograd.

Run:
    conda run -n tileops python -m pytest tests/ops/test_batch_norm.py -vvs
"""


import pytest
import torch

from tests.test_base import FixtureBase, TestBase
from tileops.ops.norm.batch_norm import BatchNormBwdOp, BatchNormFwdOp
from workloads.batch_norm import (
    BatchNormBwdTest as _BatchNormBwdTestWorkload,
)
from workloads.batch_norm import (
    BatchNormFwdTest as _BatchNormFwdTestWorkload,
)


def _ref_fwd(x, weight, bias, running_mean, running_var, training, momentum=0.1, eps=1e-5):
    """Reference: torch.nn.functional.batch_norm (float32 upcast)."""
    x32 = x.float()
    rm = running_mean.clone()
    rv = running_var.clone()
    y32 = torch.nn.functional.batch_norm(
        x32, rm, rv, weight.float(), bias.float(),
        training=training, momentum=momentum, eps=eps)
    return y32.to(x.dtype), rm, rv


class BatchNormBwdTest(_BatchNormBwdTestWorkload, TestBase):
    def ref_program(self, grad_out, x, weight, mean, rstd):
        """Reference via torch.autograd on a float32 graph."""
        x32 = x.float().requires_grad_(True)
        w32 = weight.float().requires_grad_(True)
        b32 = torch.zeros(self.C, device=x.device, dtype=torch.float32, requires_grad=True)
        rm = torch.zeros(self.C, device=x.device, dtype=torch.float32)
        rv = torch.ones(self.C, device=x.device, dtype=torch.float32)
        y32 = torch.nn.functional.batch_norm(
            x32, rm, rv, w32, b32, training=True, momentum=0.1, eps=1e-5)
        y32.backward(grad_out.float())
        return x32.grad.to(x.dtype), w32.grad, b32.grad

class BatchNormFwdTest(_BatchNormFwdTestWorkload, TestBase):
    def ref_program(self, x, weight, bias, running_mean, running_var):
        y, rm, rv = _ref_fwd(x, weight, bias, running_mean, running_var,
                             training=self.training)
        return (y,)


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

class BatchNormFwdFixture(FixtureBase):
    """(N, C, *spatial, dtype, training)"""
    PARAMS = [
        ("N, C, spatial, dtype, training", [
            # BatchNorm1d – (N, C)
            pytest.param(32, 64, (), torch.float16, True, marks=pytest.mark.smoke),
            pytest.param(32, 64, (), torch.bfloat16, True, marks=pytest.mark.smoke),
            pytest.param(32, 64, (), torch.float16, False, marks=pytest.mark.full),
            pytest.param(32, 256, (), torch.bfloat16, True, marks=pytest.mark.full),
            # BatchNorm1d – (N, C, L)
            pytest.param(16, 64, (512,), torch.float16, True, marks=pytest.mark.full),
            # Non-persistent path (L > 8192): smallest representative case L=16384.
            pytest.param(4, 64, (64, 64), torch.float16, True, marks=pytest.mark.full),
            # BatchNorm2d – (N, C, H, W)
            pytest.param(8, 64, (1024, 1024), torch.float16, True, marks=pytest.mark.full),
            pytest.param(8, 64, (2048, 2048), torch.float16, False, marks=pytest.mark.full),
            pytest.param(4, 128, (32, 32), torch.bfloat16, True, marks=pytest.mark.full),
            # Non-aligned spatial: H*W=900, exercises partial-tile path
            pytest.param(8, 64, (30, 30), torch.float16, True, marks=pytest.mark.full),
            pytest.param(8, 64, (30, 30), torch.bfloat16, True, marks=pytest.mark.full),
        ]),
    ]


class BatchNormBwdFixture(FixtureBase):
    """(N, C, *spatial, dtype)"""
    PARAMS = [
        ("N, C, spatial, dtype", [
            pytest.param(32, 64, (), torch.float16, marks=pytest.mark.smoke),
            pytest.param(32, 64, (), torch.bfloat16, marks=pytest.mark.smoke),
            pytest.param(8, 64, (32, 32), torch.float16, marks=pytest.mark.full),
            pytest.param(4, 128, (32, 32), torch.bfloat16, marks=pytest.mark.full),
            # Non-persistent backward path (L=16384 > 8192).
            pytest.param(4, 64, (64, 64), torch.float16, marks=pytest.mark.full),
            # Non-aligned spatial: H*W=900, exercises partial-tile path
            pytest.param(8, 64, (30, 30), torch.float16, marks=pytest.mark.full),
            pytest.param(8, 64, (30, 30), torch.bfloat16, marks=pytest.mark.full),
        ]),
    ]


# ---------------------------------------------------------------------------
# Test helpers
# ---------------------------------------------------------------------------


# ---------------------------------------------------------------------------
# Test functions
# ---------------------------------------------------------------------------

@BatchNormFwdFixture
def test_batch_norm_fwd(N, C, spatial, dtype, training):
    test = BatchNormFwdTest(N, C, spatial, dtype, training)
    x, weight, bias, running_mean, running_var = test.gen_inputs()

    # Clone before op call so reference sees the same initial state.
    running_mean_ref = running_mean.clone()
    running_var_ref = running_var.clone()

    op = BatchNormFwdOp(N, C, tuple(spatial), dtype=dtype, training=training)
    # Manifest input order: (x, running_mean, running_var, weight, bias).
    y = op(x, running_mean, running_var, weight, bias)

    ref_y, ref_rm, ref_rv = _ref_fwd(x, weight, bias, running_mean_ref, running_var_ref,
                                      training=training)

    # float16 accumulates more error; use loose tolerances.
    atol, rtol = (1e-2, 1e-2) if dtype == torch.float16 else (2e-2, 2e-2)
    max_err = (y.float() - ref_y.float()).abs().max()
    assert torch.allclose(y.float(), ref_y.float(), atol=atol, rtol=rtol), \
        f"fwd mismatch (training={training}): max_err={max_err:.4e}"

    if training:
        rm_err = (running_mean.float() - ref_rm.float()).abs().max()
        assert torch.allclose(running_mean.float(), ref_rm.float(), atol=atol, rtol=rtol), \
            f"running_mean mismatch: max_err={rm_err:.4e}"
        rv_err = (running_var.float() - ref_rv.float()).abs().max()
        assert torch.allclose(running_var.float(), ref_rv.float(), atol=atol, rtol=rtol), \
            f"running_var mismatch: max_err={rv_err:.4e}"

    print(f"test_batch_norm_fwd passed [training={training}]: max_err={max_err:.4e}")


@BatchNormBwdFixture
def test_batch_norm_bwd(N, C, spatial, dtype):
    test = BatchNormBwdTest(N, C, spatial, dtype)
    grad_out, x, weight, mean, rstd = test.gen_inputs()

    op = BatchNormBwdOp(N, C, *spatial, dtype=dtype)
    grad_x, grad_weight, grad_bias = op(grad_out, x, weight, mean, rstd)

    ref_gx, ref_gw, ref_gb = test.ref_program(grad_out, x, weight, mean, rstd)

    atol, rtol = (1e-2, 1e-2) if dtype == torch.float16 else (2e-2, 2e-2)

    for name, got, ref in [
        ("grad_x",      grad_x.float(),    ref_gx.float()),
        ("grad_weight", grad_weight.float(), ref_gw.float()),
        ("grad_bias",   grad_bias.float(),   ref_gb.float()),
    ]:
        max_err = (got - ref).abs().max()
        assert torch.allclose(got, ref, atol=atol, rtol=rtol), \
            f"bwd {name} mismatch: max_err={max_err:.4e}"
    print("test_batch_norm_bwd passed: grad_x/weight/bias all match")


@pytest.mark.smoke
def test_batch_norm_fwd_returns_single_tensor() -> None:
    """BatchNormFwdOp forward must produce one tensor — manifest declares
    a single output. ``training`` is bound at ctor; the runtime kwarg is
    no longer accepted."""
    if not is_available():
        pytest.skip(f"{DEVICE.upper()} required for forward call")

    N, C, H, W = 4, 8, 4, 4
    op = BatchNormFwdOp(N, C, (H, W), dtype=torch.float16, training=False)
    x = torch.randn(N, C, H, W, device=DEVICE, dtype=torch.float16)
    weight = torch.randn(C, device=DEVICE, dtype=torch.float32)
    bias = torch.randn(C, device=DEVICE, dtype=torch.float32)
    rm = torch.zeros(C, device=DEVICE, dtype=torch.float32)
    rv = torch.ones(C, device=DEVICE, dtype=torch.float32)

    y = op(x, rm, rv, weight, bias)
    assert isinstance(y, torch.Tensor)
    assert y.shape == x.shape


if __name__ == "__main__":
    pytest.main([__file__, "-vvs"])