IQ.Pilot Release Commit @ b6534c0

artifacts/package_sources/tinygrad/extra/gemm/simple_conv.py

@@ -0,0 +1,33 @@
+from tinygrad.helpers import getenv
+from tinygrad import dtypes, Tensor
+
+dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
+acc_dtype = dtypes.half if getenv("ACC_HALF") else dtypes.bfloat16 if getenv("ACC_BFLOAT16") else None
+
+CNT = getenv("CNT", 8)
+BS = getenv("BS", 16)
+CIN = getenv("CIN", 128)
+COUT = getenv("COUT", 128)
+HW = getenv("HW", 128)
+K = getenv("K", 3)
+PADDING = getenv("PADDING", 1)
+COMP = getenv("COMP", 0)
+ATOL = getenv("ATOL", 1e-4)
+RTOL = getenv("RTOL", 3e-2)
+
+FLOPS = BS*K*K*CIN*HW*HW*COUT*2
+def rand_input(): return Tensor.rand(BS, CIN, HW, HW, dtype=dtype_in).realize(), Tensor.rand(COUT, CIN, K, K, dtype=dtype_in).realize()
+
+if __name__ == "__main__":
+  a, b = rand_input()
+  for i in range(CNT):
+    if i > 0 and getenv("RAND", 0) != 0:
+      a, b = rand_input()
+    c = a.conv2d(b, padding=PADDING, dtype=acc_dtype).realize()
+
+  if COMP:
+    import numpy as np, time, torch
+    torch_device = "cuda:0" if torch.cuda.is_available() else ("mps" if getenv("MPS", 0) else "cpu")
+    ta, tb = torch.from_numpy(a.numpy()).to(torch_device), torch.from_numpy(b.numpy()).to(torch_device)
+    tc = torch.nn.functional.conv2d(ta, tb, padding=PADDING)
+    np.testing.assert_allclose(c.numpy(), tc.cpu(), atol=ATOL, rtol=RTOL)