nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

remove iter from uopgraph (#6110)

Browse Source 
* remove iter from uopgraph

* linearize returns uops

* fix tests

* linearize in linearize

* tests fix

* touchup

* test failures
This commit is contained in:
George Hotz
2024-08-16 15:58:29 -07:00
committed by GitHub
parent 28c75bf2a6
commit 74ee9febec
15 changed files with 104 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/kernel.py
+8 -7
View File
@@ -4,7 +4,7 @@ from dataclasses import dataclass, replace
from collections import defaultdict
from typing import Literal, Optional, List, Tuple, Union, cast, Dict, Final, DefaultDict
from tinygrad.ops import BinaryOps, ReduceOps, UNSAFE_PAD_OPS, KernelInfo, BUFFER_UOPS, UOp, UOps, verify_ast
from tinygrad.ops import BinaryOps, ReduceOps, UNSAFE_PAD_OPS, KernelInfo, BUFFER_UOPS, UOp, UOps, verify_ast, print_uops
from tinygrad.device import Device
from tinygrad.renderer import Renderer, TensorCore, Program
from tinygrad.dtype import DType, ImageDType, PtrDType
@@ -738,15 +738,16 @@ class Kernel:
verify_ast(modified_ast)
# generate the UOpGraph
self.uops:UOpGraph = UOpGraph(ast_to_uop(modified_ast, self.opts), self.opts)
if DEBUG >= 5: self.uops.print()
if getenv("GRAPHUOPS"): self.uops.graph()
self.uops:List[UOp] = UOpGraph(ast_to_uop(modified_ast, self.opts), self.opts).linearize(self.opts.extra_matcher)
if DEBUG >= 5: print_uops(self.uops)
if getenv("GRAPHUOPS"):
from tinygrad.engine.graph import graph_uops
graph_uops(self.uops)
return self
def to_program(self, name_override:Optional[str]=None) -> Program:
self.linearize()
self.uops.linearize(self.opts.extra_matcher)
src = self.opts.render(name:=to_function_name(ansiname:=(name_override if name_override is not None else self.name)), self.uops.uops)
src = self.opts.render(name:=to_function_name(ansiname:=(name_override if name_override is not None else self.name)), self.uops)
if getenv("RUN_PROCESS_REPLAY"):
table_name = f"process_replay_{getenv('GITHUB_RUN_ID', 'HEAD')}_{getenv('GITHUB_RUN_ATTEMPT')}"
@@ -757,5 +758,5 @@ class Kernel:
mem_bytes = sum(max(cast(DType, x.src[0].dtype).itemsize * x.src[-1].arg.real_size() for x in group)
for _, group in itertools.groupby([x for x in self.ast.parents if x.op in BUFFER_UOPS and x.src[0].op is UOps.DEFINE_GLOBAL],
key=lambda x: (x.op, x.src[0].arg)))
return Program(ansiname, src, self.opts.device, self.uops.uops, mem_estimate=mem_bytes,
return Program(ansiname, src, self.opts.device, self.uops, mem_estimate=mem_bytes,
global_size=[1,1,1] if self.opts.has_local else None, local_size=[1,1,1] if self.opts.has_local else None)
tinygrad/codegen/uopgraph.py
+16 -35
View File
@@ -1,5 +1,5 @@
from __future__ import annotations
from typing import Iterator, Optional, Tuple, Dict, List, Set, Union, cast, TYPE_CHECKING, Any, DefaultDict, Callable
from typing import Optional, Tuple, Dict, List, Set, Union, cast, TYPE_CHECKING, Any, DefaultDict, Callable
import functools, itertools, heapq, math, operator
from collections import defaultdict
from tinygrad.dtype import dtypes, PtrDType, ImageDType, DType
@@ -521,34 +521,14 @@ class UOpGraph:
def __init__(self, sink:Union[UOp, List[UOp]], opts:Optional[Renderer]=None):
self.sink: UOp = sink if isinstance(sink, UOp) else UOp(UOps.SINK, None, tuple(sink))
assert self.sink.op is UOps.SINK, f"sink isn't sink, it's {self.sink.op}"
# used by linearizer
self._uops: Optional[List[UOp]] = None
self.opts = opts
self.folder = constant_folder + transcendental_folding({} if TRANSCENDENTAL >= 2 or opts is None else opts.code_for_op.keys())
def __reduce__(self): return self.__class__, (self.sink, self.opts)
def __iter__(self) -> Iterator[UOp]: return iter(self.uops)
def __getitem__(self, index) -> UOp: return self.uops[index]
@property
def uops(self) -> List[UOp]:
if self._uops is None: self.linearize()
return cast(List[UOp], self._uops)
def graph(self):
from tinygrad.engine.graph import graph_uops
graph_uops(self.uops)
def print(self): print_uops(self.uops)
cnt = 0
def linearize(self, extra_pm:Optional[PatternMatcher]=None, skip_check=False) -> UOpGraph:
def linearize(self, extra_pm:Optional[PatternMatcher]=None, skip_check=False) -> List[UOp]:
global acc_number
acc_number = 0
# NOTE: relinearizering should be okay
#assert self._uops is None, "already linearized"
# do graph rewrite
sink = graph_rewrite(self.sink, self.folder)
@@ -598,15 +578,15 @@ class UOpGraph:
if in_degree[u] == 0: push(u)
scope_end: Dict[UOp, UOp] = {}
self._uops = []
_uops: List[UOp] = []
while queue:
p,x = heapq.heappop(queue)
if DEBUG >= 7: print(f"{p:5d}",x)
if x in scope_children: scope_end[x] = x
if x.op is UOps.DEFINE_ACC:
idx = min([self._uops.index(l) for l in x.src if l.op is UOps.RANGE])
self._uops.insert(idx, x)
else: self._uops.append(x)
idx = min([_uops.index(l) for l in x.src if l.op is UOps.RANGE])
_uops.insert(idx, x)
else: _uops.append(x)
for u, ss in scope_children.items():
if x in ss:
ss.remove(x)
@@ -616,24 +596,25 @@ class UOpGraph:
if in_degree[u] == 0: push(u)
# end scopes in toposort order
for u, x in scope_end.items(): self._uops.insert(self._uops.index(x)+1, UOp(END_FOR_UOP[u.op][1], None, (u,)))
for u, x in scope_end.items(): _uops.insert(_uops.index(x)+1, UOp(END_FOR_UOP[u.op][1], None, (u,)))
# sanity checks (NOTE: these can cause things to be skipped in BEAM)
if not skip_check:
bad_ops = dedup([x.op for x in self._uops if x.op in {UOps.EXPAND, UOps.CONTRACT, UOps.REDUCE}])
bad_ops = dedup([x.op for x in _uops if x.op in {UOps.EXPAND, UOps.CONTRACT, UOps.REDUCE}])
try:
type_verify(self.uops)
assert self._uops[-1].op is UOps.SINK, f"didn't end with SINK, ended with {self._uops[-1]}"
type_verify(_uops)
assert _uops[-1].op is UOps.SINK, f"didn't end with SINK, ended with {_uops[-1]}"
assert len(bad_ops) == 0, f"bad UOps left in list: {bad_ops}"
# TODO: this should be enabled, and the valid clause should be removed
# NOTE: multiple identical stores to DEFINE_LOCAL is okay
assert len(all_stores := [x.src[0:2]+x.src[3:] for x in self._uops if x.op is UOps.STORE and x.src[0].op is not UOps.DEFINE_LOCAL]) \
assert len(all_stores := [x.src[0:2]+x.src[3:] for x in _uops if x.op is UOps.STORE and x.src[0].op is not UOps.DEFINE_LOCAL]) \
== len(dedup(all_stores)), "repeated stores in uops"
except AssertionError as e:
self.print()
if not CI: self.graph()
print_uops(_uops)
if not CI:
from tinygrad.engine.graph import graph_uops
graph_uops(_uops)
raise e
# strip the SINK
self._uops = self._uops[:-1]
return self
return _uops[:-1]