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

opt

Browse Source
This commit is contained in:
geohot
2025-03-22 18:52:56 +08:00
parent 1b4e9f5e91
commit e1d2bec4a4
2 changed files with 15 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extra/replay_pkl.py
+14
View File
@@ -116,6 +116,11 @@ if __name__ == "__main__":
k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
k.apply_opt(Opt(OptOps.UPCAST, 1, 24))
k.apply_opt(Opt(OptOps.UPCAST, 0, 16))
elif knum == 9:
# same as 6
k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
k.apply_opt(Opt(OptOps.UPCAST, 1, 24))
k.apply_opt(Opt(OptOps.UPCAST, 0, 16))
elif knum in [7,11]:
k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
k.apply_opt(Opt(OptOps.UPCAST, 1, 144))
@@ -129,6 +134,15 @@ if __name__ == "__main__":
k.apply_opt(Opt(OptOps.UPCAST, 1, 64))
#k.apply_opt(Opt(OptOps.UPCAST, 0, 2))
pass
elif knum == 1:
k.apply_opt(Opt(OptOps.UPCAST, 2, 32))
k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
k.apply_opt(Opt(OptOps.UPCAST, 0, 2))
elif knum == 26:
#k.apply_opt(Opt(OptOps.UNROLL, 1, 0))
#k.apply_opt(Opt(OptOps.UNROLL, 0, 0))
#k.apply_opt(Opt(OptOps.UPCAST, 2, 128))
pass
#elif knum == 18:
# k.apply_opt(Opt(OptOps.UNROLL, 0, 4))
# k.apply_opt(Opt(OptOps.UPCAST, 1, 192))
tinygrad/codegen/kernel.py
+1 -1
View File
@@ -666,7 +666,7 @@ class Kernel:
print(self.name)
if getenv("RAWAST"): print(self.ast)
for i,(buf,st) in enumerate([(buf,st) for buf,st in zip(self.bufs, self.sts) if buf.op not in {Ops.CONST, Ops.VALID}]):
print(f"{i:2d}: {str(st.shape):25s} {str(buf.src[0].dtype).replace('dtypes.',''):20s} {str(st.real_strides()):20s}",
print(f"{i:2d}: {str(st.shape):25s} {str(buf.src[0].dtype).replace('dtypes.',''):20s} {str(st.real_strides()):30s}",
str(st) if DEBUG >= 4 else "")
print(self.applied_opts)
if DEBUG >= 5: print(modified_ast)