github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-08 06:34:03 -05:00
Code Issues Packages Projects Releases Wiki Activity

one compiler path

Browse Source
This commit is contained in:
George Hotz
2025-12-22 19:20:04 -05:00
parent 339dadf056
commit 2aec58654a
3 changed files with 6 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
test/test_renderer_failures.py
View File

@@ -23,8 +23,9 @@ def _test_uop_result(inputs:list[Tensor], stores:list[UOp], local_size=None):
initial_value=np.zeros(sz, dtype=_to_np_dtype(dtype)).data) for u in uops if u.op is Ops.STORE]
inbufs = [x.uop.base.buffer for x in inputs]
src = Device[Device.DEFAULT].renderer.render(uops)
lib = Device[Device.DEFAULT].compiler.compile_cached(src)
ei = CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test",
src, Device.DEFAULT, uops[-1], uops=uops, local_size=local_size))
src, Device.DEFAULT, uops[-1], uops=uops, lib=lib, local_size=local_size))
ei.exec(outbufs+inbufs)
return [np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in outbufs]

7
test/test_uops.py
View File

@@ -7,7 +7,6 @@ from tinygrad.dtype import dtypes, DType, AddrSpace
from tinygrad.device import Buffer, Device
from tinygrad.uop.ops import Ops, UOp, UPat, KernelInfo, exec_alu, AxisType
from tinygrad.uop.spec import shared_spec
from tinygrad.renderer import ProgramSpec
from tinygrad.renderer.cstyle import CStyleLanguage
from tinygrad.engine.realize import CompiledRunner, get_program, get_runner
from tinygrad.engine.schedule import ExecItem
@@ -26,11 +25,7 @@ def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
return ret[:-1]
def _uops_to_prg(uops_list):
uops = full_rewrite(ast:=UOp.sink(*uops_list), ren=Device[Device.DEFAULT].renderer)
src = Device[Device.DEFAULT].renderer.render(uops)
has_local = Device[Device.DEFAULT].renderer.has_local
return CompiledRunner(ProgramSpec(uops[-1].arg.name if uops[-1].arg is not None else "test", src, Device.DEFAULT, ast, uops=uops,
global_size=[1,1,1] if has_local else None, local_size=[1,1,1] if has_local else None))
return CompiledRunner(get_program(UOp.sink(*uops_list), Device[Device.DEFAULT].renderer, Device.DEFAULT))
def uop(uops:list[UOp], uop:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
uops.append(UOp(uop, dtype, tuple(src), arg))

8
tinygrad/engine/realize.py
View File

@@ -2,7 +2,7 @@ from typing import cast, Callable
import time, pprint, random, itertools, math
from dataclasses import dataclass, replace, field
from tinygrad.helpers import all_same, colored, DEBUG, GlobalCounters, ansilen, BEAM, NOOPT, all_int, CAPTURING, Metadata, TRACEMETA, TracingKey
from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, cpu_profile, PROFILE, ProfilePointEvent, cpu_events, prod, Context
from tinygrad.helpers import DEVECTORIZE, time_to_str, VALIDATE_WITH_CPU, getenv, PROFILE, ProfilePointEvent, cpu_events, prod, Context
from tinygrad.helpers import unwrap
from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, graph_rewrite, track_rewrites, KernelInfo, pyrender
from tinygrad.device import Device, Buffer
@@ -75,11 +75,9 @@ class CompiledRunner(Runner):
def __init__(self, p:ProgramSpec, prg=None):
if DEBUG >= 3: print(p.applied_opts)
if DEBUG >= 4: print(p.src)
if p.lib is None:
with cpu_profile(TracingKey(f"compile {p.name}", (p.function_name,)), "TINY"):
p = replace(p, lib=Device[p.device].compiler.compile_cached(p.src))
assert p.lib is not None, "lib must be provided"
self.p:ProgramSpec = p
if DEBUG >= 7: Device[p.device].compiler.disassemble(unwrap(p.lib))
if DEBUG >= 7: Device[p.device].compiler.disassemble(p.lib)
self._prg = Device[p.device].runtime(p.function_name, p.lib) if prg is None else prg
super().__init__(p.name, p.device, p.estimates)