vars_from_ast -> LazyOp.vars (#2965)

This commit is contained in:
chenyu
2024-01-01 18:12:38 -05:00
committed by GitHub
parent 980f421442
commit 58d3d5030b
8 changed files with 17 additions and 20 deletions

View File

@@ -1,7 +1,7 @@
from typing import List
from extra.models.resnet import ResNet50
from tinygrad.tensor import Tensor
from tinygrad.ops import LoadOps, vars_from_ast
from tinygrad.ops import LoadOps
from tinygrad.device import Device, Compiled
from tinygrad.codegen.linearizer import Linearizer
from tinygrad.features.search import time_linearizer, beam_search, bufs_from_lin
@@ -57,7 +57,7 @@ if __name__ == "__main__":
choices = []
for lin in lins:
tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10)
gflops = sym_infer(lin.info.flops, {k:k.min for k in vars_from_ast(lin.ast)})*1e-9/tm
gflops = sym_infer(lin.info.flops, {k:k.min for k in lin.ast.vars()})*1e-9/tm
choices.append((tm, gflops, lin.linearize()))
# print all kernels