From 8db92bd060ae8a813baef59c7e3057d2f2f07d67 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Sun, 8 Oct 2023 05:57:41 -0700 Subject: [PATCH] fix tvm gemm example --- extra/gemm/tvm_gemm.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/extra/gemm/tvm_gemm.py b/extra/gemm/tvm_gemm.py index 126a609f7b..11318a348e 100644 --- a/extra/gemm/tvm_gemm.py +++ b/extra/gemm/tvm_gemm.py @@ -31,9 +31,6 @@ except ImportError: import os from tinygrad.tensor import Tensor -# disable optimizations -os.environ["NOOPT"] = "1" - # define the compute A = Tensor.rand(M, K, device="clang") B = Tensor.rand(K, N, device="clang") @@ -42,7 +39,7 @@ C = (A.reshape(M, 1, K) * B.permute(1,0).reshape(1, N, K)).sum(axis=2) sched = C.lazydata.schedule() from tinygrad.codegen.linearizer import Linearizer from tinygrad.codegen.kernel import LinearizerOptions -lin = Linearizer(sched[-1][0], LinearizerOptions(has_local=False, supports_float4=False)) +lin = Linearizer(sched[-1].ast, LinearizerOptions(has_local=False, supports_float4=False)) #lin.hand_coded_optimizations() lin.linearize() from tinygrad.runtime.ops_clang import renderer