From f3e7efd39ef7a62e4e3462a774bbe7919fcd356e Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 15 May 2026 19:42:29 -0700 Subject: [PATCH] control flow --- tinygrad/codegen/__init__.py | 6 +++++- tinygrad/codegen/minigen.py | 5 ++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 9fa5e09bf4..f584bac65b 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -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([ diff --git a/tinygrad/codegen/minigen.py b/tinygrad/codegen/minigen.py index 2ccb77e81f..d1a4cb14d2 100644 --- a/tinygrad/codegen/minigen.py +++ b/tinygrad/codegen/minigen.py @@ -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 \ No newline at end of file