From 8428c694b3b33b5a984cae8d414a6a5d69a5a302 Mon Sep 17 00:00:00 2001 From: Christopher Milan Date: Wed, 17 Sep 2025 14:59:20 -0700 Subject: [PATCH] diskcache --- tinygrad/helpers.py | 4 +-- tinygrad/runtime/support/lvp.py | 50 +++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) create mode 100644 tinygrad/runtime/support/lvp.py diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index c3cc6cbfdf..f14d6b9445 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -247,9 +247,9 @@ def diskcache_clear(): drop_tables = cur.execute("SELECT 'DROP TABLE IF EXISTS ' || quote(name) || ';' FROM sqlite_master WHERE type = 'table';").fetchall() cur.executescript("\n".join([s[0] for s in drop_tables] + ["VACUUM;"])) -def diskcache_get(table:str, key:dict|str|int) -> Any: +def diskcache_get(table:str, key:dict|str|bytes|int) -> Any: if CACHELEVEL < 1: return None - if isinstance(key, (str,int)): key = {"key": key} + if isinstance(key, (str,bytes,int)): key = {"key": key} cur = db_connection().cursor() try: res = cur.execute(f"SELECT val FROM '{table}_{VERSION}' WHERE {' AND '.join([f'{x}=?' for x in key.keys()])}", tuple(key.values())) diff --git a/tinygrad/runtime/support/lvp.py b/tinygrad/runtime/support/lvp.py new file mode 100644 index 0000000000..1713592b32 --- /dev/null +++ b/tinygrad/runtime/support/lvp.py @@ -0,0 +1,50 @@ +from tinygrad.device import Compiler +from tinygrad.helpers import cpu_objdump +from tinygrad.runtime.autogen.nir import nir_shader_compiler_options +from tinygrad.runtime.support.compiler_cpu import cerr, expect +import tinygrad.runtime.autogen.llvm as llvm +import base64, ctypes, gzip +import tinygrad.runtime.autogen.lvp as lvp + +class LVPCompiler(Compiler): + nir_options = ctypes.pointer(nir_shader_compiler_options.from_buffer_copy( + gzip.decompress(base64.b64decode("H4sIAAAAAAAAA2NgZGRkYGAAkYxgCsQFsxigwgwQBoxmhCqFq2WEKwIrAEGIkQxoAEMALwCqVsCiGUwLMHA0QPn29nBJ0syHAgCuUNSvAAEAAA==")))) + def __init__(self, cache_key="lvp"): + # FIXME: this is wrong if mesa is compiled using ORCJIT + self.ctx = lvp.lp_context_ref(ctypes.cast(llvm.LLVMContextCreate(), ctypes.POINTER(lvp.struct_LLVMOpaqueContext)), True) + super().__init__(f"compile_{cache_key}") + + def compile(self, src) -> bytes: + # import os + # input(f"pid: {os.getpid()}") + blobreader = lvp.struct_blob_reader() + lvp.blob_reader_init(blobreader, src, len(src)) + shader = lvp.nir_deserialize(None, ctypes.cast(self.nir_options, ctypes.POINTER(lvp.nir_shader_compiler_options)), blobreader) + + gallivm = lvp.gallivm_create(None, self.ctx, None) + module = ctypes.cast(gallivm.contents.module, ctypes.POINTER(llvm.struct_LLVMOpaqueModule)) + + params = lvp.struct_lp_build_tgsi_params(lvp.struct_lp_type(floating=True, sign=True, width=32, length=4), + resources_type=lvp.lp_build_jit_resources_type(gallivm), mask=ctypes.pointer(lvp.struct_lp_build_mask_context())) + + ctx = ctypes.cast(gallivm.contents.context, ctypes.POINTER(llvm.struct_LLVMOpaqueContext)) + builder = ctypes.cast(gallivm.contents.builder, ctypes.POINTER(llvm.struct_LLVMOpaqueBuilder)) + pt = llvm.LLVMPointerType(ctypes.cast(params.resources_type, ctypes.POINTER(llvm.struct_LLVMOpaqueType)), 0) + fn = llvm.LLVMAddFunction(module, b"aaa", llvm.LLVMFunctionType(llvm.LLVMVoidTypeInContext(ctx), pt, 1, 0)) + llvm.LLVMPositionBuilderAtEnd(builder, llvm.LLVMAppendBasicBlockInContext(ctx, fn, b"entry")) + + params.consts_ptr = lvp.lp_build_struct_get_ptr2(gallivm, params.resources_type, + ctypes.cast(llvm.LLVMGetParam(fn, 0), ctypes.POINTER(lvp.struct_LLVMOpaqueValue)), lvp.LP_JIT_RES_CONSTANTS, b"constants") + lvp.lp_build_mask_begin(params.mask, gallivm, params.type, lvp.lp_build_one(gallivm, params.type)) + lvp.lp_build_mask_end(params.mask) + + lvp.lp_build_nir_soa(gallivm, shader, params, None) + llvm.LLVMBuildRetVoid(builder) + lvp.gallivm_verify_function(gallivm, ctypes.cast(fn, ctypes.POINTER(lvp.struct_LLVMOpaqueValue))) + lvp.gallivm_compile_module(gallivm) + t = llvm.LLVMGetExecutionEngineTargetMachine(ctypes.cast(gallivm.contents.engine, ctypes.POINTER(llvm.struct_LLVMOpaqueExecutionEngine))) + obj_buf = expect(llvm.LLVMTargetMachineEmitToMemoryBuffer(t, module, llvm.LLVMObjectFile, e:=cerr(), + ctypes.pointer(b:=llvm.LLVMMemoryBufferRef())), e, b) + return ctypes.string_at(llvm.LLVMGetBufferStart(obj_buf), llvm.LLVMGetBufferSize(obj_buf)) + + def disassemble(self, lib:bytes): cpu_objdump(lib)