github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-24 14:28:09 -05:00
Code Issues Packages Projects Releases Wiki Activity

remove uop symobilc rendering [run_process_replay] (#6108)

Browse Source
This commit is contained in:
George Hotz
2024-08-16 09:02:15 -07:00
committed by GitHub
parent e8ae9af962
commit f82ecd8802
1 changed files with 2 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
tinygrad/codegen/lowerer.py
View File

@@ -7,7 +7,7 @@ from tinygrad.dtype import dtypes, PtrDType, ImageDType, DType
from tinygrad.ops import BufferOps, LazyOp, ReduceOps, UnaryOps, MetaOps, KernelInfo, MemBuffer, BinaryOps
from tinygrad.codegen.uops import UOp, UOps
from tinygrad.renderer import Renderer
from tinygrad.helpers import getenv, all_int, get_contraction, prod, partition, flatten
from tinygrad.helpers import all_int, get_contraction, prod, partition, flatten
# TODO: this needs to be replaced, there shouldn't be variables in the shapetracker, only ints and UOps
from tinygrad.shape.symbolic import Variable, NumNode, SumNode, MulNode, DivNode, ModNode, LtNode, AndNode
@@ -33,7 +33,7 @@ def _uop_view(view:View, idxs:List[UOp], vexpr:UOp) -> Tuple[UOp, UOp]:
return iexpr, vexpr
# TODO: change this once UOps is ready to replace symbolic
def st_to_uops_graph(st:ShapeTracker, idxs:List[UOp], dtype:DType) -> Tuple[UOp, UOp]:
def st_to_uops(st:ShapeTracker, idxs:List[UOp], dtype:DType) -> Tuple[UOp, UOp]:
idx, valid = _uop_view(st.views[-1], idxs, UOp.const(dtypes.bool, True))
for view in reversed(st.views[0:-1]):
view = view.minify()
@@ -47,38 +47,6 @@ def st_to_uops_graph(st:ShapeTracker, idxs:List[UOp], dtype:DType) -> Tuple[UOp,
idx = UOp(UOps.VECTORIZE, dtypes.int.vec(3), ((idx // 4) % dtype.shape[1], (idx // (4 * dtype.shape[1])), idx % 4))
return idx, valid
# TODO: this is the old one, delete when ready
def st_to_uops_symbolic(st:ShapeTracker, idxs:List[UOp], dtype:DType) -> Tuple[UOp, UOp]:
fake_idxs = [Variable(f"__idx{i}", 0, s-1) for i,s in enumerate(st.shape)]
idx, valid = st.expr_idxs(fake_idxs)
ctx = dict(zip(fake_idxs, idxs))
uvalid = valid.render(render_ops, ctx)
if isinstance(dtype, ImageDType):
image_idxs = (idx // 4) % dtype.shape[1], (idx // (4 * dtype.shape[1])), idx % 4
uidx = UOp(UOps.VECTORIZE, dtypes.int.vec(3), tuple(x.render(render_ops, ctx) for x in image_idxs))
else:
uidx = idx.render(render_ops, ctx)
if uvalid.op is UOps.CONST: uvalid = UOp.const(dtypes.bool, uvalid.arg)
assert uvalid.dtype == dtypes.bool
return uidx, uvalid
def st_to_uops(st:ShapeTracker, idxs:List[UOp], dtype:DType) -> Tuple[UOp, UOp]:
if getenv("SYMBOLIC_DIFF"):
symbolic_idx, symbolic_valid = st_to_uops_symbolic(st, idxs, dtype)
graph_idx, graph_valid = st_to_uops_graph(st, idxs, dtype)
import ocdiff
from tinygrad.codegen.uopgraph import UOpGraph
from tinygrad.renderer.cstyle import OpenCLRenderer
def render(s1, s2):
glbl = UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), arg="idxs")
st = tuple(UOp(UOps.STORE, None, (glbl, UOp.const(dtypes.int, i), s)) for i,s in enumerate([s1,s2]))
return OpenCLRenderer().render("indexing", UOpGraph(UOp(UOps.SINK, None, st)).linearize(skip_check=True).uops)
cmp_symbolic, cmp_graph = render(symbolic_idx, symbolic_valid), render(graph_idx, graph_valid)
if cmp_symbolic != cmp_graph: print(ocdiff.console_diff(f"SYMBOLIC {len(cmp_symbolic)}\n"+cmp_symbolic, f"GRAPH {len(cmp_graph)}\n"+cmp_graph))
return st_to_uops_graph(st, idxs, dtype) if getenv("UOP_IS_SYMBOLIC", 1) else st_to_uops_symbolic(st, idxs, dtype)
def _limit_dims(dims:Tuple[sint, ...], max_sizes:Tuple[int, ...]):
# TODO: symbolic shape
if not all_int(dims): return dims