File size: 1,189 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 | """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)
|