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