tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 10:16:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

5 and 8

Browse Source
This commit is contained in:
geohot
2025-03-22 19:14:04 +08:00
parent fd73ec2b1b
commit 2634975d5a
2 changed files with 12 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/onnx.py
+3
View File
@@ -744,6 +744,9 @@ def get_onnx_ops():
return y, scale, zero_point
def DequantizeLinear(x:Tensor, x_scale:Tensor, x_zero_point:Tensor|int=0, axis:int=1, block_size:int=0):
if getenv("NHWC") and len(x.shape) == 4 and x.shape[1:] == (1,3,3):
# 3x3 depthwise
x = x.permute(2,3,0,1).contiguous().permute(2,3,0,1)
WEIGHT_SHIFT = 4
if getenv("NHWC") and len(x.shape) == 4 and x.shape[2:] == (1,1) and x.shape[1]%WEIGHT_SHIFT == 0:
if x.shape[0]%32 == 0:
extra/replay_pkl.py
+9
View File
@@ -112,6 +112,12 @@ if __name__ == "__main__":
k.apply_opt(Opt(OptOps.UNROLL, 0, 8))
k.apply_opt(Opt(OptOps.UPCAST, 1, 96))
k.apply_opt(Opt(OptOps.UPCAST, 0, 4))
elif knum == 5:
k.apply_opt(Opt(OptOps.UPCAST, 2, 96))
# this breaks something
#k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
elif knum == 8:
k.apply_opt(Opt(OptOps.UPCAST, 2, 144))
elif knum == 6:
k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
k.apply_opt(Opt(OptOps.UPCAST, 1, 24))
@@ -168,6 +174,9 @@ if __name__ == "__main__":
full_shape = k.full_shape
out_shape = k.sts[0].shape
out_strides = k.sts[0].real_strides()
# there's some bug here with this
#if len(out_strides) == 5 and full_shape[-2:] == (3,3):
#if knum not in [2,30]: k.apply_opt(Opt(OptOps.UPCAST, 2, 0))
if len(out_strides) == 3:
if full_shape[1] == 192 and full_shape[0]%2 == 0:
k.apply_opt(Opt(OptOps.UNROLL, 0, 4))