File size: 1,087 Bytes
58ba6df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 | import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from common import *
from hamer import tt_kernels as tk
from hamer.fused import host_im2col
dev = open_dev()
torch.manual_seed(0)
for mem in [ttnn.L1_MEMORY_CONFIG]:
x = torch.randn(1, 3, 256, 192)
xt = ttnn.from_torch(x, dtype=ttnn.float32, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=mem)
o = ttnn.to_torch(tk.im2col_tile(xt))
ref = host_im2col(x)
print("exact", torch.equal(o, ref), float((o - ref).abs().max()))
x2 = torch.randn(1, 3, 256, 192)
h = ttnn.from_torch(x2, dtype=ttnn.float32, layout=ttnn.ROW_MAJOR_LAYOUT)
ttnn.copy_host_to_device_tensor(h, xt)
print("exact after copy", torch.equal(ttnn.to_torch(tk.im2col_tile(xt)), host_im2col(x2)))
print("us", time_op(dev, lambda: tk.im2col_tile(xt)))
print("tilize us", time_op(dev, lambda: ttnn.tilize_with_zero_padding(ttnn.from_torch(host_im2col(x), dtype=ttnn.float32, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev), use_multicore=True)) if False else "")
ttnn.close_device(dev)
|