IQ.Pilot Release Commit @ bec7652
This commit is contained in:
IQ.Lvbs history cleanup
@@ -0,0 +1,91 @@
|
||||
import unittest, numpy as np, os
|
||||
from tinygrad import Tensor, Device, TinyJit
|
||||
from tinygrad.helpers import Timing, getenv
|
||||
import multiprocessing.shared_memory as shared_memory
|
||||
|
||||
N = getenv("NSZ", 256)
|
||||
class TestCopySpeed(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls): Device[Device.DEFAULT].synchronize()
|
||||
|
||||
def testCopySHMtoDefault(self):
|
||||
s = shared_memory.SharedMemory(name="test_X", create=True, size=N*N*4)
|
||||
s.close()
|
||||
if os.path.exists("/dev/shm"):
|
||||
t = Tensor.empty(N, N, device="disk:/dev/shm/test_X").realize()
|
||||
else:
|
||||
t = Tensor.empty(N, N, device="disk:shm:test_X").realize()
|
||||
for _ in range(3):
|
||||
with Timing("sync: ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
|
||||
with Timing("queue: "):
|
||||
t.to(Device.DEFAULT).realize()
|
||||
Device[Device.DEFAULT].synchronize()
|
||||
s.unlink()
|
||||
|
||||
def testCopyCPUtoDefault(self):
|
||||
t = Tensor.ones(N, N, device="CPU").contiguous().realize()
|
||||
print(f"buffer: {t.nbytes()*1e-9:.2f} GB")
|
||||
for _ in range(3):
|
||||
with Timing("sync: ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
|
||||
with Timing("queue: "):
|
||||
t.to(Device.DEFAULT).realize()
|
||||
Device[Device.DEFAULT].synchronize()
|
||||
|
||||
def testCopyCPUtoDefaultFresh(self):
|
||||
print("fresh copy")
|
||||
for _ in range(3):
|
||||
t = Tensor.ones(N, N, device="CPU").contiguous().realize()
|
||||
with Timing("sync: ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"): # noqa: F821
|
||||
with Timing("queue: "):
|
||||
t.to(Device.DEFAULT).realize()
|
||||
Device[Device.DEFAULT].synchronize()
|
||||
del t
|
||||
|
||||
def testCopyDefaulttoCPU(self):
|
||||
t = Tensor.ones(N, N).contiguous().realize()
|
||||
print(f"buffer: {t.nbytes()*1e-9:.2f} GB")
|
||||
for _ in range(3):
|
||||
with Timing("sync: ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
|
||||
t.to('CPU').realize()
|
||||
|
||||
def testCopyDefaulttoCPUJit(self):
|
||||
if Device.DEFAULT == "CPU": return unittest.skip("CPU to CPU copy is a no-op")
|
||||
|
||||
@TinyJit
|
||||
def _do_copy(t): return t.to('CPU').realize()
|
||||
|
||||
t = Tensor.randn(N, N).contiguous().realize()
|
||||
Device[Device.DEFAULT].synchronize()
|
||||
for _ in range(5):
|
||||
with Timing(f"copy {Device.DEFAULT} -> CPU {t.nbytes()/(1024**2)}M: ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
|
||||
x = _do_copy(t)
|
||||
Device[Device.DEFAULT].synchronize()
|
||||
np.testing.assert_equal(t.numpy(), x.numpy())
|
||||
|
||||
def testCopyCPUtoDefaultJit(self):
|
||||
if Device.DEFAULT == "CPU": return unittest.skip("CPU to CPU copy is a no-op")
|
||||
|
||||
@TinyJit
|
||||
def _do_copy(x): return x.to(Device.DEFAULT).realize()
|
||||
|
||||
for _ in range(5):
|
||||
t = Tensor.randn(N, N, device="CPU").contiguous().realize()
|
||||
Device["CPU"].synchronize()
|
||||
with Timing(f"copy CPU -> {Device.DEFAULT} {t.nbytes()/(1024**2)}M: ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
|
||||
x = _do_copy(t)
|
||||
Device[Device.DEFAULT].synchronize()
|
||||
np.testing.assert_equal(t.numpy(), x.numpy())
|
||||
|
||||
@unittest.skipIf(Device.DEFAULT != "CL" or Device[Device.DEFAULT].count() != 6, "only test this on CL, with 6 gpus")
|
||||
def testCopyCPUto6GPUs(self):
|
||||
t = Tensor.ones(N, N, device="CPU").contiguous().realize()
|
||||
print(f"buffer: {t.nbytes()*1e-9:.2f} GB")
|
||||
for _ in range(3):
|
||||
with Timing("sync: ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s ({t.nbytes()*6/ns:.2f} GB/s total)"):
|
||||
with Timing("queue: "):
|
||||
for g in range(6):
|
||||
t.to(f"CL:{g}").realize()
|
||||
Device["CL"].synchronize()
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,37 @@
|
||||
import unittest
|
||||
from tinygrad import Device
|
||||
from tinygrad.helpers import Timing, Profiling
|
||||
|
||||
class TestDeviceSpeed(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.dev = Device[Device.DEFAULT]
|
||||
cls.empty = Device[Device.DEFAULT].renderer.render([])
|
||||
|
||||
def test_empty_compile(self):
|
||||
with Timing("compiler "):
|
||||
self.dev.compiler.compile(self.empty)
|
||||
|
||||
def test_empty_compile_twice(self):
|
||||
self.dev.compiler.compile(self.empty)
|
||||
with Timing("compiler "):
|
||||
self.dev.compiler.compile(self.empty)
|
||||
|
||||
def test_launch_speed(self):
|
||||
prg_bin = self.dev.compiler.compile(self.empty)
|
||||
prg = self.dev.runtime("test", prg_bin)
|
||||
for _ in range(10): prg() # ignore first launches
|
||||
with Timing("launch 1000x "):
|
||||
for _ in range(1000): prg()
|
||||
with Timing("launch 1000x with wait "):
|
||||
for _ in range(1000): prg(wait=True)
|
||||
|
||||
def test_profile_launch_speed(self):
|
||||
prg_bin = self.dev.compiler.compile(self.empty)
|
||||
prg = self.dev.runtime("test", prg_bin)
|
||||
for _ in range(10): prg() # ignore first launches
|
||||
with Profiling():
|
||||
for _ in range(1000): prg()
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,54 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor, Device, dtypes
|
||||
from tinygrad.helpers import DEV
|
||||
# similar to test/external/external_test_gpu_ast.py, but universal
|
||||
|
||||
@unittest.skipIf(Device.DEFAULT in {"CUDA", "NV"} and DEV.interface.startswith("MOCK"), "slow on ocelot")
|
||||
class TestSpecific(unittest.TestCase):
|
||||
# from openpilot
|
||||
|
||||
# 1x1 6 <- 24
|
||||
def test_1x1_6_24(self):
|
||||
x = Tensor.randn(1, 24*4, 32, 64)
|
||||
w = Tensor.randn(6*4, 24*4, 1, 1)
|
||||
x.conv2d(w).permute(0,2,3,1).reshape(32, 384, 4).contiguous().realize()
|
||||
|
||||
def test_vec_mul(self):
|
||||
# this forces it to be an image...
|
||||
x = Tensor.ones(1, 512, 4).contiguous().reshape(1, 2048)
|
||||
w = Tensor.randn(2048, 512)
|
||||
(x @ w).reshape(1, 128, 4).contiguous().realize()
|
||||
|
||||
@unittest.skipUnless(dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need float16 support")
|
||||
def test_big_vec_mul(self):
|
||||
# from LLaMA
|
||||
# 0 buffer<4096, dtypes.float> [View((1024, 1, 1, 4), (4, 0, 0, 1), 0, None)]
|
||||
# 1 buffer<4096, dtypes.float> [View((1024, 1024, 4, 4), (0, 4, 1, 0), 0, None)]
|
||||
# 2 buffer<16777216, dtypes.half> [View((1024, 1024, 4, 4), (16384, 4, 1, 4096), 0, None)]
|
||||
x = Tensor.randn(4096).realize()
|
||||
w = Tensor.randn(4096, 4096, dtype=dtypes.float16).realize()
|
||||
(x @ w.T).realize()
|
||||
|
||||
# from https://dl.acm.org/doi/pdf/10.1145/3495243.3517020
|
||||
|
||||
# ~260 GFLOPS on Adreno 640, should be 260*(720/890)*(596/710) = 176.5 on downclocked 630
|
||||
# we get 170
|
||||
def test_1x1_28_28(self):
|
||||
x = Tensor.randn(1, 256, 28, 28)
|
||||
w = Tensor.randn(256, 256, 1, 1)
|
||||
x.conv2d(w).permute(0,2,3,1).reshape(28, 28*256//4, 4).contiguous().realize()
|
||||
|
||||
# 132 GFLOPS on Adreno 640, should be 132*(720/890)*(596/710) = 90 on downclocked 630
|
||||
# gets 54 with broken opt, 74 without opt, and 146 if we pad and opt 3!
|
||||
def test_3x3_28_28_stride_2(self):
|
||||
x = Tensor.randn(1, 288, 36, 36)
|
||||
w = Tensor.randn(384, 288, 3, 3)
|
||||
x.conv2d(w, stride=2).permute(0,2,3,1).reshape(17, 17*384//4, 4).contiguous().realize()
|
||||
|
||||
def test_3x3_28_28_stride_2_padded(self):
|
||||
x = Tensor.randn(1, 288, 36, 36)
|
||||
w = Tensor.randn(384, 288, 3, 3)
|
||||
x.conv2d(w, stride=2, padding=1).permute(0,2,3,1).reshape(18, 18*384//4, 4).contiguous().realize()
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,288 @@
|
||||
import os
|
||||
os.environ["NVIDIA_TF32_OVERRIDE"] = "0"
|
||||
os.environ["MKL_NUM_THREADS"] = "1"
|
||||
os.environ["NUMEXPR_NUM_THREADS"] = "1"
|
||||
os.environ["OMP_NUM_THREADS"] = "1"
|
||||
os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
|
||||
import unittest
|
||||
import torch
|
||||
torch.set_num_threads(1)
|
||||
import time
|
||||
import numpy as np
|
||||
import sys
|
||||
np.set_printoptions(linewidth=160)
|
||||
from tinygrad import Tensor, Device, GlobalCounters, TinyJit
|
||||
from tinygrad.nn import Conv2d
|
||||
from tinygrad.helpers import colorize_float, getenv, DEV
|
||||
|
||||
IN_CHANS = [int(x) for x in getenv("IN_CHANS", "4,16,64").split(",")]
|
||||
|
||||
torch_dt = torch.float16 if getenv("HALF", 0) else torch.float32
|
||||
torch_device = torch.device('mps' if getenv("MPS", 0) else ('cuda' if getenv("TORCHCUDA", 0) else 'cpu'))
|
||||
if str(torch_device) == "mps":
|
||||
import torch.mps
|
||||
def sync(): torch.mps.synchronize()
|
||||
elif str(torch_device) == "cuda":
|
||||
import torch.cuda
|
||||
def sync(): torch.cuda.synchronize()
|
||||
else:
|
||||
def sync(): pass
|
||||
|
||||
save_ops, save_mem = 0, 0
|
||||
CNT = getenv("CNT", 8)
|
||||
def helper_test_speed(f1, *args):
|
||||
global save_ops, save_mem
|
||||
ets = []
|
||||
ret = None
|
||||
cache_defeat = np.zeros((2048,2048))
|
||||
for i in range(CNT):
|
||||
del ret
|
||||
|
||||
# operation cache defeats
|
||||
args = [(x+1).realize() if isinstance(x, Tensor) else (None if x is None else (x+1)) for x in args]
|
||||
args = [(x-1).realize() if isinstance(x, Tensor) else (None if x is None else (x-1)) for x in args]
|
||||
|
||||
# force syncing
|
||||
[x.numpy() if isinstance(x, Tensor) or str(torch_device) == "cpu" else x.cpu().numpy() for x in args if x is not None]
|
||||
|
||||
# clear 32MB global memory cache (CPU and global memory only)
|
||||
cache_defeat += 1
|
||||
|
||||
# manual pre sync
|
||||
if isinstance(args[0], Tensor):
|
||||
local_device = Device[args[0].device]
|
||||
local_device.synchronize()
|
||||
else: sync()
|
||||
|
||||
GlobalCounters.global_ops = 0
|
||||
GlobalCounters.global_mem = 0
|
||||
st = time.perf_counter()
|
||||
ret = f1(*args)
|
||||
if isinstance(ret, Tensor): local_device.synchronize()
|
||||
else: sync()
|
||||
et = (time.perf_counter() - st) * 1000
|
||||
if i >= 1: ets.append(et)
|
||||
if GlobalCounters.global_ops:
|
||||
save_ops, save_mem = GlobalCounters.global_ops, GlobalCounters.global_mem
|
||||
return ret.numpy() if isinstance(ret, Tensor) else ret.cpu().numpy(), np.min(ets)
|
||||
|
||||
def helper_test_generic_square(name, N, f1, f2, onearg=False):
|
||||
torch.manual_seed(0)
|
||||
torch_a = (torch.rand(N, N, dtype=torch_dt) - 0.5).to(torch_device)
|
||||
torch_b = (torch.rand(N, N, dtype=torch_dt) - 0.5).to(torch_device) if not onearg else None
|
||||
|
||||
tiny_a = Tensor(torch_a.cpu().numpy())
|
||||
tiny_b = Tensor(torch_b.cpu().numpy()) if not onearg else None
|
||||
|
||||
helper_test_generic(f"{name:30s} {N:5d}x{N:5d}", f1, (torch_a, torch_b), TinyJit(f2), (tiny_a, tiny_b))
|
||||
|
||||
def helper_test_matvec(name, N, M):
|
||||
torch.manual_seed(0)
|
||||
torch_a = (torch.rand(N, dtype=torch_dt) - 0.5).to(torch_device)
|
||||
torch_b = (torch.rand(N, M, dtype=torch_dt) - 0.5).to(torch_device)
|
||||
|
||||
tiny_a = Tensor(torch_a.cpu().numpy())
|
||||
tiny_b = Tensor(torch_b.cpu().numpy())
|
||||
|
||||
helper_test_generic(f"{name:30s} {N:5d}x{M:5d}", lambda a,b: a@b, (torch_a, torch_b), TinyJit(lambda a,b:a@b), (tiny_a, tiny_b))
|
||||
|
||||
prefix = None
|
||||
def helper_test_generic(name, f1, f1_args, f2, f2_args):
|
||||
global prefix
|
||||
with torch.no_grad():
|
||||
val_torch, et_torch = helper_test_speed(f1, *f1_args)
|
||||
val_tinygrad, et_tinygrad = helper_test_speed(f2, *f2_args)
|
||||
|
||||
desc = "faster" if et_torch > et_tinygrad else "slower"
|
||||
flops = save_ops*1e-6
|
||||
mem = save_mem*1e-6
|
||||
print(("\r" if sys.stdout.isatty() else "")+f"{name:42s} {et_torch:7.2f} ms ({flops/et_torch:9.2f} GFLOPS {mem/et_torch:7.2f} GB/s) in torch, {et_tinygrad:7.2f} ms ({flops/et_tinygrad:9.2f} GFLOPS {mem/et_tinygrad:7.2f} GB/s) in tinygrad, {colorize_float(et_tinygrad/et_torch)} {desc} {flops:10.2f} MOPS {mem:8.2f} MB") # noqa: E501
|
||||
atol, rtol = (1e-2, 1e-2) if torch_dt == torch.float16 else (1e-3, 1e-3)
|
||||
np.testing.assert_allclose(val_tinygrad, val_torch, atol=atol, rtol=rtol)
|
||||
|
||||
def helper_test_conv(bs, in_chans, out_chans, kernel_size, img_size_y, img_size_x):
|
||||
torch.manual_seed(0)
|
||||
torch_dat = torch.rand(bs, in_chans, img_size_y, img_size_x, dtype=torch_dt).to(torch_device)
|
||||
torch_conv = torch.nn.Conv2d(in_chans, out_chans, kernel_size, bias=None, dtype=torch_dt).to(torch_device)
|
||||
|
||||
tiny_dat = Tensor(torch_dat.cpu().numpy())
|
||||
tiny_conv = Conv2d(in_chans, out_chans, kernel_size, bias=None)
|
||||
tiny_conv.weight = Tensor(torch_conv.weight.detach().cpu().numpy())
|
||||
|
||||
def f1(torch_dat): return torch_conv(torch_dat)
|
||||
def f2(tiny_dat): return tiny_conv(tiny_dat).realize()
|
||||
helper_test_generic(f"conv bs:{bs:3d} chans:{in_chans:3d} -> {out_chans:3d} k:{kernel_size}", f1, (torch_dat,), TinyJit(f2), (tiny_dat,))
|
||||
|
||||
@unittest.skipIf(getenv("BIG") == 0, "no big tests")
|
||||
@unittest.skipIf(DEV.interface.startswith("MOCK"), "no MOCKGPUs")
|
||||
class TestBigSpeed(unittest.TestCase):
|
||||
def test_add(self):
|
||||
def f(a, b): return a+b
|
||||
helper_test_generic_square('add', 8192, f, f)
|
||||
def test_exp(self):
|
||||
def f(a, b): return a.exp()
|
||||
helper_test_generic_square('exp', 8192, f, f, onearg=True)
|
||||
def test_gemm_2048(self):
|
||||
def f(a, b): return a @ b
|
||||
helper_test_generic_square('gemm', 2048, f, f)
|
||||
def test_gemm_4096(self):
|
||||
def f(a, b): return a @ b
|
||||
helper_test_generic_square('gemm', 4096, f, f)
|
||||
def test_large_conv_1x1(self): helper_test_conv(bs=32, in_chans=128, out_chans=128, kernel_size=1, img_size_y=128, img_size_x=128)
|
||||
def test_large_conv_3x3(self): helper_test_conv(bs=4, in_chans=128, out_chans=128, kernel_size=3, img_size_y=130, img_size_x=130)
|
||||
def test_large_conv_5x5(self): helper_test_conv(bs=4, in_chans=128, out_chans=128, kernel_size=5, img_size_y=132, img_size_x=132)
|
||||
def test_matvec_4096_16384(self): helper_test_matvec('matvec_4096_16384', 4096, 16384)
|
||||
def test_matvec_16384_4096(self): helper_test_matvec('matvec_16384_4096', 16384, 4096)
|
||||
|
||||
@unittest.skipIf(getenv("BIG") == 1, "only big tests")
|
||||
@unittest.skipIf(DEV.interface.startswith("MOCK"), "no MOCKGPUs")
|
||||
class TestSpeed(unittest.TestCase):
|
||||
def test_sub(self):
|
||||
def f(a, b): return a-b
|
||||
helper_test_generic_square('sub', 4096, f, f)
|
||||
|
||||
def test_pow(self):
|
||||
def f(a, b): return a.pow(b)
|
||||
helper_test_generic_square('pow', 2048, f, f)
|
||||
|
||||
def test_sum(self):
|
||||
def f(a, b): return a.sum()
|
||||
helper_test_generic_square('sum', 2048, f, f, onearg=True)
|
||||
helper_test_generic_square('sum', 4096, f, f, onearg=True)
|
||||
|
||||
def test_partial_sum(self):
|
||||
R = 256
|
||||
def f(a, b): return a.reshape(int(4096//R), int(4096*R)).sum(axis=1)
|
||||
helper_test_generic_square('partial_sum', 4096, f, f, onearg=True)
|
||||
|
||||
@unittest.skip("not really used in models")
|
||||
def test_cumsum(self):
|
||||
def f0(a, b): return a.cumsum(axis=0)
|
||||
def f1(a, b): return a.cumsum(axis=1)
|
||||
helper_test_generic_square('cumsum_0', 256, f0, f0, onearg=True)
|
||||
helper_test_generic_square('cumsum_1', 256, f1, f1, onearg=True)
|
||||
|
||||
def test_cat(self):
|
||||
helper_test_generic_square('cat_0', 2048, lambda x,y: torch.cat((x,y),dim=0), lambda x,y: x.cat(y,dim=0))
|
||||
helper_test_generic_square('cat_1', 2048, lambda x,y: torch.cat((x,y),dim=1), lambda x,y: x.cat(y,dim=1))
|
||||
|
||||
def test_array_packing(self):
|
||||
N = 2048
|
||||
def f(a, b): return a.reshape(N, N // 32, 32).permute(1,0,2).contiguous()
|
||||
helper_test_generic_square('array_packing', N, f, f, onearg=True)
|
||||
|
||||
def test_permute(self):
|
||||
for N in [1024, 4096]:
|
||||
# this is a 64MB tensor, M1 L1 cache is 128kB
|
||||
# to fit easily in L1, rotations should be 128x128 chunks.
|
||||
def f(a, b): return a.permute(1,0).contiguous()
|
||||
helper_test_generic_square('permute', N, f, f, onearg=True)
|
||||
|
||||
def test_double_permute(self):
|
||||
N = 64
|
||||
torch.manual_seed(0)
|
||||
torch_a = (torch.rand(N, N, N, N, dtype=torch_dt) - 0.5).to(torch_device)
|
||||
tiny_a = Tensor(torch_a.cpu().numpy())
|
||||
def f(a): return a.permute(1,0,3,2).contiguous()
|
||||
helper_test_generic(f"double_permute {tiny_a.shape}", f, (torch_a,), TinyJit(lambda a: f(a).realize()), (tiny_a,))
|
||||
|
||||
def test_neg(self):
|
||||
def f(a, b): return -a
|
||||
helper_test_generic_square('neg', 4096, f, f, onearg=True)
|
||||
|
||||
def test_exp(self):
|
||||
def f(a, b): return a.exp()
|
||||
helper_test_generic_square('exp', 2048, f, f, onearg=True)
|
||||
|
||||
def test_sqrt(self):
|
||||
def f(a, b): return a.sqrt()
|
||||
helper_test_generic_square('sqrt', 2048, f, f, onearg=True)
|
||||
|
||||
def test_relu(self):
|
||||
def f(a, b): return a.relu()
|
||||
helper_test_generic_square('relu', 4096, f, f, onearg=True)
|
||||
|
||||
def test_max(self):
|
||||
def f(a, b): return a.max()
|
||||
helper_test_generic_square('max', 4096, f, f, onearg=True)
|
||||
|
||||
def test_mul_sum(self):
|
||||
def f(a, b): return (a*b).sum()
|
||||
helper_test_generic_square('mul_sum', 4096, f, f)
|
||||
|
||||
def test_add_a(self):
|
||||
def f(a, b): return a + b
|
||||
helper_test_generic_square('add', 1, f, f)
|
||||
|
||||
def test_add_big(self):
|
||||
for N in [1024, 4096]:
|
||||
def f(a, b): return a + b
|
||||
helper_test_generic_square('add', N, f, f)
|
||||
|
||||
def test_add_constant(self):
|
||||
def f(a, b): return a+2.0
|
||||
helper_test_generic_square('add_constant', 4096, f, f, onearg=True)
|
||||
|
||||
def test_add_sq(self):
|
||||
def f(a, b): return a*a + b*b
|
||||
helper_test_generic_square('add_sq', 4096, f, f)
|
||||
|
||||
def test_gemm(self):
|
||||
def f(a, b): return a @ b
|
||||
helper_test_generic_square('gemm', 1024, f, f)
|
||||
|
||||
def test_gemm_small(self):
|
||||
def f(a, b): return a @ b
|
||||
helper_test_generic_square('gemm', 256, f, f)
|
||||
|
||||
def test_gemm_unrolled(self):
|
||||
N = 512
|
||||
def f1(a, b): return a@b.T
|
||||
def f2(a, b): return (a.reshape(N, 1, N).expand(N, N, N) * b.reshape(1, N, N).expand(N, N, N)).sum(axis=2)
|
||||
helper_test_generic_square('gemm_unrolled', N, f1, f2)
|
||||
|
||||
def test_gemm_unrolled_permute_l(self):
|
||||
N = 512
|
||||
def f1(a, b): return a.T@b.T
|
||||
def f2(a, b): return (a.permute(1,0).reshape(N, 1, N).expand(N, N, N) * b.reshape(1, N, N).expand(N, N, N)).sum(axis=2)
|
||||
helper_test_generic_square('gemm_unrolled_permute_l', N, f1, f2)
|
||||
|
||||
def test_gemm_unrolled_permute_r(self):
|
||||
N = 512
|
||||
def f1(a, b): return a@b
|
||||
def f2(a, b): return (a.reshape(N, 1, N).expand(N, N, N) * b.permute(1,0).reshape(1, N, N).expand(N, N, N)).sum(axis=2)
|
||||
helper_test_generic_square('gemm_unrolled_permute_r', N, f1, f2)
|
||||
|
||||
def test_gemm_unrolled_permute_lr(self):
|
||||
N = 512
|
||||
def f1(a, b): return a.T@b
|
||||
def f2(a, b): return (a.permute(1,0).reshape(N, 1, N).expand(N, N, N) * b.permute(1,0).reshape(1, N, N).expand(N, N, N)).sum(axis=2)
|
||||
helper_test_generic_square('gemm_unrolled_permute_lr', N, f1, f2)
|
||||
|
||||
def test_matvec_1024_1024(self): helper_test_matvec('matvec_1024_1024', 1024, 1024)
|
||||
def test_matvec_1024_4096(self): helper_test_matvec('matvec_1024_4096', 1024, 4096)
|
||||
def test_matvec_4096_1024(self): helper_test_matvec('matvec_4096_1024', 4096, 1024)
|
||||
def test_matvec_4096_4096(self): helper_test_matvec('matvec_4096_4096', 4096, 4096)
|
||||
|
||||
def test_openpilot_conv2d(self):
|
||||
bs, in_chans, out_chans = 1,12,32
|
||||
torch.manual_seed(0)
|
||||
torch_dat = torch.rand(bs, 64, 128, 12, dtype=torch_dt).to(torch_device)
|
||||
torch_conv = torch.nn.Conv2d(in_chans, out_chans, 3, bias=None, padding=1, dtype=torch_dt).to(torch_device)
|
||||
|
||||
tiny_dat = Tensor(torch_dat.cpu().numpy())
|
||||
tiny_conv = Conv2d(in_chans, out_chans, 3, bias=None, padding=1)
|
||||
tiny_conv.weight = Tensor(torch_conv.weight.detach().cpu().numpy())
|
||||
|
||||
def f1(torch_dat): return torch_conv(torch_dat.permute(0,3,1,2))
|
||||
def f2(tiny_dat): return tiny_conv(tiny_dat.permute(0,3,1,2)).realize()
|
||||
helper_test_generic(f"conv bs:{bs:3d} chans:{in_chans:3d} -> {out_chans:3d} k:3", f1, (torch_dat,), TinyJit(f2), (tiny_dat,))
|
||||
|
||||
def test_conv2d(self):
|
||||
for bs in [32]:
|
||||
for in_chans in IN_CHANS:
|
||||
for out_chans in [32]:
|
||||
helper_test_conv(bs, in_chans, out_chans, 3, 34, 34)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user