File size: 3,526 Bytes
58ba6df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | 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)
|