mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-28 00:53:45 +08:00
firestar5683
66 lines
2.7 KiB
Python
66 lines
2.7 KiB
Python
"""Shared compilation and input allocation for model and warp artifacts."""
|
|
import io, pathlib, pickle, shutil, struct, time, tempfile
|
|
import numpy as np
|
|
from typing import Callable
|
|
from tinygrad import Tensor, Device, Context
|
|
from tinygrad.device import Buffer
|
|
from tinygrad.uop.ops import PatternMatcher, UPat, Ops, graph_rewrite
|
|
from tinygrad.nn.state import get_parameters
|
|
|
|
def allocate_inputs(input_specs, initialize=None):
|
|
"""Initialize inputs before copying to devices."""
|
|
arrays = {name: np.zeros(shape, dtype=dtype) for name, (shape, dtype, _) in input_specs.items()}
|
|
if initialize is not None: initialize(arrays)
|
|
return {name: Tensor(arrays[name], device=device).realize() for name, (_, _, device) in input_specs.items()}
|
|
|
|
def dump_pickle(obj, path):
|
|
with tempfile.TemporaryFile(dir=".") as buffers:
|
|
opcodes, offset = io.BytesIO(), 0
|
|
|
|
def persistent_id(b):
|
|
nonlocal offset
|
|
if not isinstance(b, Buffer) or b.device != Device.DEFAULT: return None
|
|
offset += buffers.write(bytes((-offset) % 256))
|
|
offset = (off:=offset) + buffers.write(b.as_memoryview())
|
|
return b.nbytes, off
|
|
|
|
p = pickle.Pickler(opcodes)
|
|
p.persistent_id = persistent_id
|
|
p.dump(obj)
|
|
with open(path, "wb") as f:
|
|
f.write(struct.pack('<q', len(opcodes_b:=opcodes.getvalue())))
|
|
f.write(opcodes_b)
|
|
buffers.seek(0)
|
|
shutil.copyfileobj(buffers, f)
|
|
|
|
def load_pickle(path, *, out_of_band=False):
|
|
with open(path, "rb") as f:
|
|
if not out_of_band: return pickle.load(f)
|
|
opcodes = f.read(struct.unpack('<q', f.read(8))[0])
|
|
if f.tell() == pathlib.Path(path).stat().st_size:
|
|
return pickle.loads(opcodes)
|
|
buffers = Tensor(pathlib.Path(path))[f.tell():].to(Device.DEFAULT).realize().uop.buffer
|
|
|
|
def persistent_load(pid): return buffers.view(*pid)
|
|
|
|
u = pickle.Unpickler(io.BytesIO(opcodes))
|
|
u.persistent_load = persistent_load
|
|
return u.load()
|
|
|
|
def benchmark(fxn:Callable, cb=None, **kwargs):
|
|
# Select benchmark options when called, without probing devices on import.
|
|
with Context(OPENPILOT_HACKS=1, **{'AMD': {'TC_OPT': 2, 'TC_MIN_GLOBALS': 32}}.get(Device.DEFAULT, {})):
|
|
Device.default.synchronize()
|
|
start = time.perf_counter()
|
|
if (output := fxn(**kwargs)) is not None: output.realize()
|
|
Device.default.synchronize()
|
|
end = time.perf_counter()
|
|
if cb: cb(end-start)
|
|
return [t.numpy().copy() for t in get_parameters(kwargs.get('output_buffers', output))]
|
|
|
|
pm_retargetable = PatternMatcher([
|
|
(UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat()), name="p"), lambda p: p.replace(src=p.src[:-1]) if p.arg.target.device == "CPU" else None)
|
|
])
|
|
|
|
def make_retargetable(jit): jit.captured._linear = graph_rewrite(jit.captured._linear, pm_retargetable, walk=True, enter_calls=True)
|