mirror of
https://jihulab.com/mr-one/openpilot.git
synced 2026-10-28 05:23:46 +08:00
1okko
8d0bf5ee0a
The current recompiled42 model catalog declares tinygrad_ref=fe5d3169b. The old vendored e837 runtime accepted the pkl but left policy outputs at zero, so sync the runtime package to the compiler version used by the downloaded models.
62 lines
3.5 KiB
Python
62 lines
3.5 KiB
Python
import inspect, math
|
|
from tinygrad.device import BufferStorage, Compiled, Allocator, ProfileGraphEntry, ProfileGraphEvent, Program, TinyELF
|
|
from tinygrad.engine.jit import MultiGraphRunner
|
|
from tinygrad.renderer import Renderer, cstyle, nir, ptx, llvmir, wgsl
|
|
from tinygrad.renderer.cstyle import CStyleLanguage
|
|
from tinygrad.uop.ops import UOp, Ops
|
|
from tinygrad.helpers import cpu_profile, getenv, dedup, NULL_ALLOW_COPYOUT, PROFILE, cpu_events, perf_counter_us
|
|
|
|
class NullRenderer(CStyleLanguage):
|
|
has_local = False
|
|
float4 = "float4"
|
|
barrier = "// BARRIER"
|
|
code_for_op = {**CStyleLanguage.code_for_op, Ops.THREEFRY: lambda a,b,dtype: f"threefry({a},{b})", Ops.MAX: lambda a,b,dtype: f"max({a},{b})"}
|
|
def asm(self, prg: UOp, lin: UOp) -> bytes:
|
|
assert self.target.arch.startswith("gfx"), "only amd supports assembly"
|
|
from tinygrad.renderer.amd.elf import assemble_linear
|
|
return assemble_linear(prg, lin, self.target.arch)
|
|
|
|
class NullProgram(Program['NullDevice']):
|
|
def __init__(self, dev:'NullDevice', obj:TinyELF): self.device, self.name, self.profile_key = dev.device, obj.name, obj.profile_key
|
|
def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
|
|
with cpu_profile(self.name, self.device, profile_key=self.profile_key): return 1e-3
|
|
|
|
class NullAllocator(Allocator['NullDevice']):
|
|
def _alloc(self, size:int, options) -> BufferStorage: return BufferStorage(None)
|
|
|
|
def _copyin(self, dest, src:memoryview): pass
|
|
def _copyout(self, dest:memoryview, src):
|
|
if not NULL_ALLOW_COPYOUT: raise RuntimeError("no copyout on NULL")
|
|
def _transfer(self, dest, src, sz:int, src_dev, dest_dev):
|
|
with cpu_profile(f"{src_dev.device} -> {dest_dev.device}", f"{src_dev.device}:SDMA:0"): pass
|
|
def _offset(self, buf, size:int, offset:int): pass
|
|
|
|
class NullGraph(MultiGraphRunner):
|
|
def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None:
|
|
if PROFILE:
|
|
st, descs = perf_counter_us(), []
|
|
event_count:dict[str, int] = {}
|
|
for (_,_,bufs,_),runtime in zip(self.calls, self.runtimes):
|
|
# description based on command, copied from HCQ graph
|
|
device = runtime.device if runtime is not None else f"{bufs[1].device}:SDMA:0"
|
|
descs.append((device, runtime.name if runtime is not None else f"{bufs[1].device} -> {bufs[0].device}",
|
|
runtime.profile_key if runtime is not None else None, count:=event_count.get(device, 0)))
|
|
event_count[device] = count+1
|
|
# pack events evenly per device
|
|
dur, sigs, ents = max(1, math.ceil((perf_counter_us()-st)/max(event_count.values()))), [], []
|
|
for i,(device,name,profile_key,count) in enumerate(descs):
|
|
sigs += [st+count*dur, st+(count+1)*dur]
|
|
ents.append(ProfileGraphEntry(device, name, 2*i, 2*i+1, profile_key))
|
|
cpu_events.append(ProfileGraphEvent(ents, [], sigs))
|
|
return 1e-1
|
|
|
|
class NullDevice(Compiled):
|
|
def synchronize(self, timeout:int|None=None): pass
|
|
|
|
def __init__(self, device:str):
|
|
assert (emu:=getenv("EMULATE", "")) == "", \
|
|
"EMULATE is deprecated, use DEV=NULL:HIP:"+{"AMD":"gfx1100", "AMD_RDNA4":"gfx1201", "AMD_CDNA4":"gfx950"}.get(emu, "<arch>")
|
|
renderers = [NullRenderer] + [r for m in [cstyle, nir, ptx, llvmir, wgsl] for r in m.__dict__.values()
|
|
if inspect.isclass(r) and issubclass(r, Renderer)]
|
|
super().__init__(device, NullAllocator(self), dedup(renderers), NullProgram, NullGraph)
|