File size: 1,121 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 | 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])
# x repeated?
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)
|