| """im2col_tile on a bf16 raw-pixel crop (0..255 integers, exact in bf16) vs host_im2col; fp32 path unchanged.""" |
| 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() |
| try: |
| torch.manual_seed(0) |
| L1 = ttnn.L1_MEMORY_CONFIG |
| u = torch.randint(0, 256, (1, 3, 256, 192)).float() |
| xt = ttnn.from_torch(u.to(torch.bfloat16), dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=L1) |
| o = ttnn.to_torch(tk.im2col_tile(xt)).float() |
| print("[i2c] bf16 exact", torch.equal(o, host_im2col(u)), float((o - host_im2col(u)).abs().max()), flush=True) |
| x = torch.randn(1, 3, 256, 192) |
| xf = ttnn.from_torch(x, dtype=ttnn.float32, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=L1) |
| print("[i2c] fp32 exact", torch.equal(ttnn.to_torch(tk.im2col_tile(xf)), host_im2col(x)), flush=True) |
| if "time" in sys.argv: |
| print(f"[i2c] bf16 {time_op(dev, lambda: tk.im2col_tile(xt)):.2f} us fp32 {time_op(dev, lambda: tk.im2col_tile(xf)):.2f} us", flush=True) |
| finally: |
| ttnn.close_device(dev) |
|
|
