diskcache

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()))

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)