shapetracker from newgpu (#456)

* shapetracker from newgpu

* touchup ops

* test

* testst

* thneed deletes unused inputs

* test

* bugfix
This commit is contained in:
George Hotz
2023-01-09 12:40:01 -08:00
committed by GitHub
parent 538b1d7f5b
commit 4885fce56e
9 changed files with 104 additions and 37 deletions

View File

@@ -187,9 +187,8 @@ class LLVMBuffer(ExplicitExecAST):
k = ASTKernel(ast)
# cached kernel
key = str(ast) # TODO: does this uniquely determine the AST? No! The shapetracker can change. Do this better.
if key in LLVMBuffer.func_cache:
LLVMBuffer.func_cache[key](*[x._buf for x in k.bufs])
if k.key in LLVMBuffer.func_cache:
LLVMBuffer.func_cache[k.key](*[x._buf for x in k.bufs])
return k.ret
# cache miss, we have to process the kernel
@@ -362,5 +361,5 @@ class LLVMBuffer(ExplicitExecAST):
loop_entry[-1].branch(loop_exit[-1]._block)
loop_exit[0].ret_void()
LLVMBuffer.func_cache[key] = LLVM().exec(module, k.bufs, k.info.flops, sum(len(x._buf) for x in k.bufs))
LLVMBuffer.func_cache[k.key] = LLVM().exec(module, k.bufs, k.info.flops, sum(len(x._buf) for x in k.bufs))
return k.ret