from common import *
dev = open_dev()
torch.manual_seed(0)
DR = ttnn.DRAM_MEMORY_CONFIG
nb = 8
def wsh(ncores, width):
    cg = ttnn.num_cores_to_corerangeset(ncores, ttnn.CoreCoord(12, 10), row_wise=True)
    return ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.L1, ttnn.ShardSpec(cg, [32, width // ncores], ttnn.ShardOrientation.ROW_MAJOR))
ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=False, packer_l1_acc=True)
for K, N in [(1024, 512), (1024, 1024)]:
    x_t = torch.randn(1, 1, 32, K); w_t = torch.randn(K, N) / K ** .5; b_t = torch.randn(1, 1, 1, N) * .1
    xi = ttnn.from_torch(x_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=DR)
    grid = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(nb - 1, 0))})
    wmc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, ttnn.ShardSpec(grid, [K, N // nb], ttnn.ShardOrientation.ROW_MAJOR))
    w = ttnn.from_torch(w_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=wmc)
    print("w roundtrip", pcc(ttnn.to_torch(w), w_t))
    b = ttnn.from_torch(b_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev)
    xs = ttnn.to_memory_config(xi, wsh(8, K))
    pc = ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig(in0_block_w=K // 32 // 8, per_core_M=1, per_core_N=N // 32 // 8, fused_activation=None)
    o = ttnn.to_torch(ttnn.linear(xs, w, program_config=pc, memory_config=wsh(8, N), compute_kernel_config=ck)).float()
    ref = x_t @ w_t
    print(K, N, "nobias full", pcc(o, ref), "row0", pcc(o[..., :1, :], ref[..., :1, :]))
    err = (o - ref).abs()[0, 0]
    print("  err by col block", [round(float(err[:, i*64:(i+1)*64].max()), 3) for i in range(N // 64)])
    o = ttnn.to_torch(ttnn.linear(xs, w, bias=b, program_config=pc, memory_config=wsh(8, N), compute_kernel_config=ck)).float()
    print(K, N, "bias full", pcc(o, ref + b_t))
ttnn.close_device(dev)
import sys
dev = open_dev()
for K, N in [(1024, 512), (1024, 1024)]:
    x_t = torch.randn(1, 1, 32, K); w_t = torch.randn(K, N) / K ** .5; b_t = torch.randn(1, 1, 1, N) * .1
    xi = ttnn.from_torch(x_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=DR)
    grid = ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(nb - 1, 0))})
    wmc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, ttnn.ShardSpec(grid, [K, N // nb], ttnn.ShardOrientation.ROW_MAJOR))
    bmc = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.WIDTH_SHARDED, ttnn.BufferType.DRAM, ttnn.ShardSpec(grid, [32, N // nb], ttnn.ShardOrientation.ROW_MAJOR))
    w = ttnn.from_torch(w_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=wmc)
    b = ttnn.from_torch(b_t, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=bmc)
    xs = ttnn.to_memory_config(xi, wsh(8, K))
    pc = ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig(in0_block_w=K // 32 // 8, per_core_M=1, per_core_N=N // 32 // 8, fused_activation=None)
    o = ttnn.to_torch(ttnn.linear(xs, w, bias=b, program_config=pc, memory_config=wsh(8, N), compute_kernel_config=ck)).float()
    print(K, N, "sharded bias [1,1,1,N] row0", pcc(o[..., :1, :], (x_t @ w_t + b_t)[..., :1, :]))
    print("   time", time_op(dev, lambda: ttnn.linear(xs, w, bias=b, program_config=pc, memory_config=wsh(8, N), compute_kernel_config=ck)))
ttnn.close_device(dev)