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