Function apply_opts

tinygrad/codegen/opt/postrange.py:335–352

(ast:UOp, ren:Renderer, beam:int=0)

  return [Buffer(dname, x.ptrdtype.size, x.dtype.base) for x in glbls]

def apply_opts(ast:UOp, ren:Renderer, beam:int=0) -> UOp:
  if ast.tag is not None: return ast
  k = Scheduler(ast, ren)
  k.convert_loop_to_global()
  if ast.arg is not None and ast.arg.opts_to_apply is not None:
    for opt in ast.arg.opts_to_apply: k.apply_opt(opt)
  elif beam >= 1:
    from tinygrad.codegen.opt.search import beam_search
    rawbufs = bufs_from_ast(ast, ren.target.device)
    # beam search may open devices
    with Context(ALLOW_DEVICE_USAGE=1):
      k = beam_search(k, rawbufs, beam, bool(getenv("BEAM_ESTIMATE", 1)))
  elif not NOOPT and (ast.arg is None or ast.arg.applied_opts == ()):
    from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
    # NOTE: hand_coded_optimizations doesn't support multiblock opts yet
    if not any(u.op is Ops.STAGE for u in ast.backward_slice):
      k = hand_coded_optimizations(k)
  return k.get_optimized_ast(name_override=ast.arg.name if ast.arg is not None and ast.arg.name != "test" else None)