"""[add + LN + s2i + qkv] vs [addln + qkv] in a trace (op-to-op effects included).  usage: addln_chain.py [eth|worker]"""
import sys
from pathlib import Path
from common import *
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from hamer import tt_kernels as tk
dev = open_dev()
torch.manual_seed(0)
L1 = ttnn.L1_MEMORY_CONFIG
res = ttnn.create_sharded_memory_config((192, 1280), core_grid=ttnn.CoreGrid(y=6, x=10), strategy=ttnn.ShardStrategy.BLOCK,
                                        orientation=ttnn.ShardOrientation.ROW_MAJOR)
x = ttnn.from_torch(torch.randn(1, 1, 192, 1280), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res)
y = ttnn.from_torch(torch.randn(1, 1, 192, 1280), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev, memory_config=res)
gt = 1 + 0.3 * torch.randn(1280); bt = 0.2 * torch.randn(1280)
C = tk.addln_constants(gt, bt, dev)
ga = ttnn.from_torch(gt.to(torch.bfloat16).reshape(1, 1, -1, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev)
be = ttnn.from_torch(bt.to(torch.bfloat16).reshape(1, 1, -1, 32), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev)
lnpc = ttnn.LayerNormShardedMultiCoreProgramConfig(compute_with_storage_grid_size=(10, 6), subblock_w=4, block_h=1, block_w=4, inplace=False)
w = ttnn.from_torch(torch.randn(1280, 4608) / 36, dtype=ttnn.bfloat8_b, layout=ttnn.TILE_LAYOUT, device=dev)
b = ttnn.from_torch(torch.randn(1, 1, 1, 4608), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev)
G = dev.compute_with_storage_grid_size()
pc = ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig(compute_with_storage_grid_size=ttnn.CoreCoord(G.x, G.y), in0_block_w=10,
      out_subblock_h=2, out_subblock_w=2, per_core_M=6, per_core_N=2, fuse_batch=True, fused_activation=None, mcast_in0=True)
ck = ttnn.init_device_compute_kernel_config(dev.arch(), math_fidelity=ttnn.MathFidelity.HiFi2, fp32_dest_acc_en=True, packer_l1_acc=True)
eps = 1e-6
def lib():
    s = ttnn.add(x, y, memory_config=res)
    h = ttnn.layer_norm(s, weight=ga, bias=be, epsilon=eps, memory_config=res, program_config=lnpc)
    hi = ttnn.sharded_to_interleaved(h, L1); ttnn.deallocate(h)
    o = ttnn.linear(hi, w, bias=b, program_config=pc, memory_config=L1, compute_kernel_config=ck); ttnn.deallocate(hi)
    return [s, o]
def cus():
    s, h = tk.add_layernorm(x, y, C, eps)
    o = ttnn.linear(h, w, bias=b, program_config=pc, memory_config=L1, compute_kernel_config=ck); ttnn.deallocate(h)
    return [s, o]
def cus_only():
    return list(tk.add_layernorm(x, y, C, eps))
def mm_only():
    hi = ttnn.sharded_to_interleaved(x, L1)
    return [hi]
hpre = ttnn.sharded_to_interleaved(x, L1)
def qkv_only():
    return [ttnn.linear(hpre, w, bias=b, program_config=pc, memory_config=L1, compute_kernel_config=ck)]
def addln_then_qkv_pre():
    s, h = tk.add_layernorm(x, y, C, eps)
    o = ttnn.linear(hpre, w, bias=b, program_config=pc, memory_config=L1, compute_kernel_config=ck)
    return [s, h, o]
def ln_then_qkv_pre():
    s = ttnn.add(x, y, memory_config=res)
    h = ttnn.layer_norm(s, weight=ga, bias=be, epsilon=eps, memory_config=res, program_config=lnpc)
    o = ttnn.linear(hpre, w, bias=b, program_config=pc, memory_config=L1, compute_kernel_config=ck)
    return [s, h, o]
for name, fn in [("qkv only", qkv_only), ("addln; qkv(pre)", addln_then_qkv_pre), ("add+ln; qkv(pre)", ln_then_qkv_pre), ("lib add+ln+s2i+qkv", lib), ("addln+qkv", cus), ("addln", cus_only), ("lib", lib), ("addln+qkv", cus)]:
    print(f"{name:22s} {time_op(dev, fn, reps=32, iters=20):7.2f} us", flush=True)
ttnn.close_device(dev)