from common import *
dev = open_dev()
torch.manual_seed(0)
L1 = ttnn.L1_MEMORY_CONFIG
x_t = torch.randn(1, 1, 192, 1280); w_t = torch.randn(1, 48, 1280, 96) / 1280 ** .5
x = ttnn.from_torch(x_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=L1)
for wd in [ttnn.bfloat16, ttnn.bfloat8_b]:
    w = ttnn.from_torch(w_t, dtype=wd, layout=ttnn.TILE_LAYOUT, device=dev)
    try:
        o = ttnn.matmul(x, w, memory_config=L1)
        print("bcast-in0 matmul shape", o.shape, pcc(ttnn.to_torch(o), x_t @ w_t))
        print("time", time_op(dev, lambda: ttnn.matmul(x, w, memory_config=L1)))
    except Exception as e:
        print("bcast-in0 FAIL", str(e)[:300])
    # x repeated?
    xr = ttnn.repeat(x, ttnn.Shape([1, 48, 1, 1]))
    try:
        o = ttnn.matmul(xr, w, memory_config=L1)
        print("batched matmul", pcc(ttnn.to_torch(o), x_t @ w_t), time_op(dev, lambda: ttnn.matmul(xr, w, memory_config=L1)))
    except Exception as e:
        print("batched FAIL", str(e)[:300])
    print("repeat time", time_op(dev, lambda: ttnn.repeat(x, ttnn.Shape([1, 48, 1, 1]))))
ttnn.close_device(dev)