github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 23:48:01 -05:00
Code Issues Packages Projects Releases Wiki Activity

add some types to the scheduler [run_process_replay] (#5941)

Browse Source 
* add some types to the scheduler [run_process_replay]

* set -> dedup
This commit is contained in:
George Hotz
2024-08-06 12:23:54 -07:00
committed by GitHub
parent e077bc7baf
commit 75154d7ae2
2 changed files with 7 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
tinygrad/engine/schedule.py
View File

@@ -125,7 +125,7 @@ def _recurse_reduceops(buf:LazyBuffer, st:ShapeTracker, realizes:Dict[LazyBuffer
return (buf, st)
return top_reduce
def _lower_lazybuffer(outs:List[LazyBuffer], realizes:Dict[LazyBuffer, None]):
def _lower_lazybuffer(outs:List[LazyBuffer], realizes:Dict[LazyBuffer, None]) -> Tuple[LazyOp, List[LazyBuffer], Dict[Variable, int], List[Metadata]]:
"""describe the computation for a LazyBuffer with LazyOp + inputs + var_vals"""
if (out:=outs[0]).op is MetaOps.COPY and getenv("USE_COPY_KERNEL") and out.device.split(":")[0] == out.srcs[0].device.split(":")[0]:
rd = LazyOp(BufferOps.LOAD, (), MemBuffer(1, dtypes.uint8, st:=ShapeTracker.from_shape((out.arg,))))
@@ -164,7 +164,7 @@ def _lower_lazybuffer(outs:List[LazyBuffer], realizes:Dict[LazyBuffer, None]):
def _recurse_lb(buf:LazyBuffer, realizes:Dict[LazyBuffer, None], allbufs:Dict[LazyBuffer, None], simple_pads:Dict[LazyBuffer, None],\
children:DefaultDict[LazyBuffer, Dict[LazyBuffer, None]], assign_targets:Dict[LazyBuffer, LazyBuffer],\
double_reduces:Dict[LazyBuffer, None], scheduled=False):
double_reduces:Dict[LazyBuffer, None], scheduled=False) -> None:
"""recursively search the entire graph for all LazyBuffers, insert realizes after expands"""
if buf in allbufs or buf.base.realized is not None: return
if GRAPH: log_lazybuffer(buf, scheduled)
@@ -236,7 +236,10 @@ def _get_isolated_children(r:LazyBuffer, reduce_for_op:Dict[LazyBuffer, LazyBuff
return merge_dicts([group, {} if any(tr in group for tr in descendants) else descendants])
SCHEDULES: List = []
def _graph_schedule(outs:List[LazyBuffer], seen:Set[LazyBuffer]):
def _graph_schedule(outs:List[LazyBuffer], seen:Set[LazyBuffer]) -> \
Tuple[DefaultDict[LazyBuffer, List[LazyBuffer]], # this is the graph
DefaultDict[LazyBuffer, int], # this is the in-degree of the graph
Dict[LazyBuffer, Tuple[List[LazyBuffer], LazyOp, List[LazyBuffer], Dict[Variable, int], List[Metadata]]]]: # this is ???
"""create a graph for realizing the outputs"""
# start by just realizing the buffers passed in
realizes: Dict[LazyBuffer, None] = {x.base:None for x in outs if x.base.realized is None}

2
tinygrad/renderer/cstyle.py
View File

@@ -201,7 +201,7 @@ class ClangRenderer(CStyleLanguage):
code_for_op = {**CStyleLanguage().code_for_op, BinaryOps.MAX: lambda a,b,dtype: f"(({a}>{b})?{a}:{b})"}
def render_kernel(self, function_name, kernel, bufs, uops, prefix=None) -> str:
prefix = [_make_clang_dtype(self, dtype) for dtype in set(uop.dtype for uop in uops if uop.dtype is not None and uop.dtype.count>1)]
prefix = [_make_clang_dtype(self, dtype) for dtype in dedup(uop.dtype for uop in uops if uop.dtype is not None and uop.dtype.count>1)]
return super().render_kernel(function_name, kernel, bufs, uops, prefix)
class OpenCLRenderer(CStyleLanguage):