mirrors/tinygrad
2
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-24 02:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity

control flow

Browse source
This commit is contained in:
George Hotz 2026-05-15 19:42:29 -07:00
commit f3e7efd39e
2 changed files with 9 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

6
tinygrad/codegen/__init__.py
View file

@ -153,7 +153,11 @@ def do_render(ctx:Renderer, prg:UOp, lin:UOp) -> UOp:
return prg.replace(src=prg.src + (UOp(Ops.SOURCE, arg=src),), arg=new_arg)
def do_compile(ctx:Renderer, prg:UOp, source:UOp) -> UOp|None:
lib = ctx.compiler.compile_cached(source.arg)
try: lib = ctx.compiler.compile_cached(source.arg)
except Exception as e:
print("**** FAILED TO COMPILE ****")
print(source.arg)
raise e
return prg.replace(src=prg.src + (UOp(Ops.BINARY, arg=lib),))
pm_to_program = PatternMatcher([

5
tinygrad/codegen/minigen.py
View file

@ -4,7 +4,7 @@ from tinygrad.dtype import dtypes, Invalid
from tinygrad.helpers import SPEC
from tinygrad.renderer import Renderer
from tinygrad.codegen.late.devectorizer import pm_add_loads, reduce_to_acc, ReduceContext
from tinygrad.codegen.late.linearizer import pm_split_ends
from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow
from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_transcendental_patterns
pm_lower_weakint = PatternMatcher([
@ -52,5 +52,8 @@ def minigen_to_sink(ast:UOp, ren:Renderer, optimize:bool) -> UOp:
get_transcendental_patterns(supported_ops, force_transcendental=False)
sink = graph_rewrite(sink, pm_decomp, ctx=ren.target, name="decompose ops to renderable")
# this was the linearizer, add control flow edges where they are needed
sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True)
if SPEC: type_verify(sink, spec_program)
return sink