nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

fast 66

Browse Source
This commit is contained in:
geohot
2025-03-27 16:47:58 +08:00
parent a8bd26d9bc
commit d32ad080c3
3 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/onnx.py
+1 -1
View File
@@ -812,7 +812,7 @@ def get_onnx_ops():
else:
# DSP swizzle memory
x = x.reshape(x.shape[0], x.shape[1]//WEIGHT_SHIFT, WEIGHT_SHIFT).permute(1,0,2).contiguous().permute(1,0,2).reshape(x.shape)
if False and getenv("NHWC") and x.shape == (1000, 1280):
if getenv("NHWC") and x.shape == (1000, 1280):
x = x.reshape(-1, 320, 4)
order = (1,0,2)
x = x.permute(*order).contiguous().permute(*argsort(order))
extra/replay_pkl.py
+4 -1
View File
@@ -32,7 +32,8 @@ if __name__ == "__main__":
k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
elif knum == 66:
k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
k.apply_opt(Opt(OptOps.UPCAST, 0, 8))
k.apply_opt(Opt(OptOps.PADTO, 0, 128))
k.apply_opt(Opt(OptOps.UPCAST, 0, 128))
elif k.full_shape[-3:] == (32,3,3):
# NOTE: there's an issue with the 7 -> 8 upcast and the valid removal
if knum not in [52,56,60] and k.full_shape[-4]%4 != 0: k.apply_opt(Opt(OptOps.PADTO, len(k.full_shape)-4, 4))
@@ -55,6 +56,8 @@ if __name__ == "__main__":
k.apply_opt(Opt(OptOps.UNROLL, 0, 8))
k.apply_opt(Opt(OptOps.UPCAST, 2, 32))
if k.full_shape[1]%4 == 0: k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
elif len(k.full_shape) == 2:
if k.full_shape[0]%128 == 0: k.apply_opt(Opt(OptOps.UPCAST, 0, 128))
elif len(k.full_shape) == 1:
for sz in [128,64,32]:
if k.full_shape[0]%sz == 0:
tinygrad/codegen/kernel.py
+1 -1
View File
@@ -387,7 +387,7 @@ class Kernel:
self.group_for_reduces += 1
elif opt.op is OptOps.UNROLL: # purple
check(axis < self.first_upcast, "can't upcasted already upcasted")
check(amt <= 32, "don't unroll more than 32")
#check(amt <= 32, "don't unroll more than 32")
# TODO: fix upcast_count to put purples before yellows. broken because of METAL tensor cores
#upcast_count = sum(x == y for x,y in zip(self.full_shape[-self.upcasted:], self.output_shape[-self.upcasted:])) if self.upcasted else 0
#self.shift_to(axis, amt, insert_before=None if upcast_count == 0 else self.shape_len-upcast_count)