import unittest

import torch
import torch.nn.functional as F

from prain import gemv

GPU = torch.cuda.is_available() and gemv.fast_device(0)


def rel(a, b):
    return ((a.float() - b.float()).norm() / b.float().norm()).item()


@unittest.skipUnless(GPU, "needs a gfx10 GPU")
class TestGemv(unittest.TestCase):
    def test_up_to_four_rows_match_a_float32_linear(self):
        torch.manual_seed(0)
        for n, k in ((1000, 2816), (2816, 2112), (130, 70)):
            w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") * 0.02
            for rows in (1, 2, 3, 4):
                x = torch.randn(rows, k, dtype=torch.bfloat16, device="cuda")
                y = gemv.linear(x, w)
                self.assertEqual((y.shape, y.dtype), ((rows, n), torch.bfloat16))
                self.assertLess(rel(y, F.linear(x.float(), w.float())), 1e-2, (n, k, rows))
        x = torch.randn(2, 1, 70, dtype=torch.float16, device="cuda")  # leading dims, fp16
        w = torch.randn(130, 70, dtype=torch.float16, device="cuda")
        self.assertEqual(gemv.linear(x, w).shape, (2, 1, 130))
        self.assertLess(rel(gemv.linear(x, w), F.linear(x.float(), w.float())), 1e-2)

    def test_more_rows_bias_or_strided_weights_use_f_linear(self):
        w = torch.randn(64, 96, dtype=torch.bfloat16, device="cuda")
        bias = torch.randn(64, dtype=torch.bfloat16, device="cuda")
        x = torch.randn(5, 96, dtype=torch.bfloat16, device="cuda")
        self.assertTrue(torch.equal(gemv.linear(x, w), F.linear(x, w)))
        self.assertTrue(torch.equal(gemv.linear(x[:1], w, bias), F.linear(x[:1], w, bias)))
        strided = torch.randn(96, 64, dtype=torch.bfloat16, device="cuda").t()
        self.assertTrue(torch.equal(gemv.linear(x[:1], strided), F.linear(x[:1], strided)))

    def test_wrapped_dispatch_sends_only_decode_shapes_to_the_kernel(self):
        calls = []

        def fallback(layer, x, weight, bias=None):
            calls.append(x.shape[0])
            return F.linear(x, weight, bias)

        gemm = gemv.wrap(lambda *args, **kwargs: fallback)("auto")
        w = torch.randn(64, 96, dtype=torch.bfloat16, device="cuda")
        gemm(None, torch.randn(1, 96, dtype=torch.bfloat16, device="cuda"), w)
        gemm(None, torch.randn(9, 96, dtype=torch.bfloat16, device="cuda"), w)
        self.assertEqual(calls, [9])


if __name__ == "__main__":
    unittest.main()