| 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) |
|
|
