github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

LoadOps SYNC (#3223)

Browse Source 
* LoadOps SYNC and WAIT

* no wait, only sync

* DEBUG >= 1

* track cross device
This commit is contained in:
George Hotz
2024-01-23 21:59:18 -08:00
committed by GitHub
parent 2f4b3ab1c0
commit e2e4632aea
7 changed files with 48 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
test/test_multitensor.py
View File

@@ -238,7 +238,8 @@ class TestMultiTensor(unittest.TestCase):
np.testing.assert_allclose(r.numpy(), np.ones(256)+np.ones(256), atol=1e-4, rtol=1e-5)
assert len(jf.jit_cache) > 0
@unittest.skipIf(CI and Device.DEFAULT=="METAL", "no ICB in CI, creation of graph fails")
#@unittest.skipIf(CI and Device.DEFAULT=="METAL", "no ICB in CI, creation of graph fails")
@unittest.skip("test broken")
def test_multi_device_jit_graph(self):
if Device[d0].graph is None or Device[d1].graph is None: raise unittest.SkipTest("only test graphs")

18
tinygrad/device.py
View File

@@ -57,7 +57,7 @@ def update_stats(name:str, op_estimate:sint, mem_estimate:int, var_vals: Optiona
GlobalCounters.global_ops += op_estimate
GlobalCounters.global_mem += mem_estimate
if et is not None: GlobalCounters.time_sum_s += et
if DEBUG >= 2:
if DEBUG >= 1:
ptm = (colored(f"{et*1e3:9.2f}ms", "yellow") if et > 0.01 else f"{et*1e6:9.2f}us") if et is not None else ""
print(f"{colored(f'*** {device[:7]:7s} {GlobalCounters.kernel_count:4d}', ('magenta' if num_kernels == 1 else 'CYAN') if jit else ('green' if first_run else None))} {name+' '*(38-ansilen(name))} arg {buf_count:3d} mem {GlobalCounters.mem_used/1e9:5.2f} GB " + # noqa: E501
(str() if et is None else f"tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({op_estimate/((et or 1e-20)*1e9):8.2f} GFLOPS, {mem_estimate/((et or 1e-20)*1e9):7.2f} GB/s)")) # noqa: E501
@@ -84,7 +84,7 @@ class Buffer:
if not hasattr(self, '_buf'): return # happens when __init__ has raised exception
if not self.device.startswith("DISK"): GlobalCounters.mem_used -= self.nbytes
self.allocator.free(self._buf, self.nbytes, self.options)
def __repr__(self): return f"<buf device:{self.device} size:{self.size} dtype:{self.dtype}>"
def __repr__(self): return f"<buf device:{self.device} size:{self.size} dtype:{self.dtype}" + (">" if self.options is None else f"{self.options=}>")
def as_buffer(self, allow_zero_copy=False, force_zero_copy=False) -> memoryview:
# zero copy with as_buffer (disabled by default due to use after free)
if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, 'as_buffer'): return self.allocator.as_buffer(self._buf)
@@ -104,7 +104,7 @@ class Buffer:
class BufferCopy(JITRunner):
def copy(self, dest, src): dest.copyin(src.as_buffer(allow_zero_copy=True)) # may allocate a CPU buffer depending on allow_zero_copy
def __call__(self, rawbufs:List[Buffer], var_vals:Dict[Variable, int], wait=False, jit=False):
dest, src = rawbufs
dest, src = rawbufs[0:2]
assert dest.size == src.size and dest.dtype == src.dtype, f"buffer copy mismatch, {dest.size} != {src.size}, {dest.dtype} != {src.dtype}"
st = time.perf_counter()
self.copy(dest, src)
@@ -126,14 +126,10 @@ class BufferRead(BufferCopy):
class BufferXfer(BufferCopy):
def copy(self, dest, src):
# fast path, used on HIP between GPUs
# NOTE: we have to block here so the data isn't copied too early. this is probably due to buffer reuse
if hasattr(src.d, "block") and hasattr(dest.d, "event"): src.d.block(dest.d.event())
else: dest.d.synchronize()
src.allocator.transfer(dest._buf, src._buf, dest.size*dest.dtype.itemsize)
# NOTE: we have to block here so the data is ready on dest when dest needs it
if hasattr(dest.d, "block") and hasattr(src.d, "event"): dest.d.block(src.d.event())
else: src.d.synchronize()
if hasattr(dest.allocator.device, "track_cross_buffer") and hasattr(src.allocator, "track_cross_device"):
dest.allocator.device.track_cross_buffer.append(src)
src.allocator.track_cross_device.append(dest.allocator.device)
dest.allocator.transfer(dest._buf, src._buf, dest.nbytes)
# TODO: size, dest, src are the same type. can we enforce this?
class Allocator:

1
tinygrad/jit.py
View File

@@ -45,6 +45,7 @@ def apply_graph_to_jit(jit_cache: List[JitItem], input_rawbuffers: List[Buffer],
nonlocal current_batch, current_device
assert current_device is not None
try:
if len(current_batch) <= 1: raise GraphException("only one kernel doesn't graph")
graphed_jit_cache.append(JitItem(current_device.graph(current_batch, input_rawbuffers, var_vals), cast(List[Optional[Buffer]], input_rawbuffers))) # noqa: E501
if DEBUG >= 2: print(f"\tJIT GRAPHing batch with {len(current_batch)} kernels on device {current_device}")
except GraphException as e:

23
tinygrad/lazy.py
View File

@@ -17,7 +17,7 @@ sys.setrecursionlimit(10000)
lazycache: Dict[Any, ReferenceType[LazyBuffer]] = {}
def create_lazybuffer(device:str, st:ShapeTracker, dtype:DType, op:Optional[Op]=None, arg:Any=None, srcs:Tuple[LazyBuffer, ...]=(),
base:Optional[LazyBuffer]=None, enable_cache=bool(getenv("LAZYCACHE", 1))):
if st.size == 0: op, arg, srcs, base = LoadOps.CONST, 0, (), None
if st.size == 0 and op is not LoadOps.SYNC: op, arg, srcs, base = LoadOps.CONST, 0, (), None
cache_key = (device, st, dtype, op, arg, tuple(ref(x) for x in srcs)) if base is None else (st, ref(base))
if (rret := lazycache.get(cache_key, None)): return cast(LazyBuffer, rret()) # NOTE: this should always be a live reference
@@ -35,8 +35,8 @@ class LazyBuffer:
self.op, self.arg, self.srcs = op, arg, srcs # this is a LazyOp, except the src is LazyBuffers and not LazyOps
self.realized: Optional[Buffer] = None
self.output_buffer: Optional[Buffer] = None
self.forced_realize = False
self.contiguous_child: Optional[Tuple[ReferenceType[LazyBuffer], ShapeTracker]] = None
self.forced_realize = False
else:
# properties on view
assert base.base == base, "base must be a base itself"
@@ -53,8 +53,8 @@ class LazyBuffer:
def base(self) -> LazyBuffer: return self._base if self._base is not None else self
@staticmethod
def loadop(op, shape:Tuple[sint,...], dtype:DType, device:str, arg=None, src:Optional[LazyBuffer]=None) -> LazyBuffer:
return create_lazybuffer(device, ShapeTracker.from_shape(shape), dtype, op, arg, (src,) if src is not None else (), enable_cache=False)
def loadop(op, shape:Tuple[sint,...], dtype:DType, device:str, arg=None, src:Optional[LazyBuffer]=None, enable_cache=False) -> LazyBuffer:
return create_lazybuffer(device, ShapeTracker.from_shape(shape), dtype, op, arg, (src,) if src is not None else (), enable_cache=enable_cache)
def const(self, val:Union[float, int], shape:Optional[Tuple[sint,...]]=None) -> LazyBuffer:
shape = self.shape if shape is None else shape
@@ -77,6 +77,10 @@ class LazyBuffer:
def schedule(self, seen=None): return create_schedule([self], seen)
def _copy(self, device:str) -> LazyBuffer:
sync = LazyBuffer.loadop(LoadOps.SYNC, (0,), dtypes.uint32, self.device, src=self, enable_cache=True)
return create_lazybuffer(device, ShapeTracker.from_shape(self.shape), self.dtype, LoadOps.COPY, None, (self, sync), enable_cache=False)
def copy_to_device(self, device:str) -> LazyBuffer:
# no COPY
if self.device == device: return self
@@ -90,11 +94,10 @@ class LazyBuffer:
return LazyBuffer.loadop(LoadOps.CONST, tuple(), self.dtype, device, arg=self.base.arg)._view(self.st)
# if it's a shrink, do the shrink before the copy with CONTIGUOUS
if prod(self.st.shape) < prod(self.base.st.shape):
return LazyBuffer.loadop(LoadOps.COPY, self.shape, self.dtype, device, src=self.contiguous())
if prod(self.st.shape) < prod(self.base.st.shape): return self.contiguous()._copy(device)
# copy the base and apply the shapetracker on the new device
return LazyBuffer.loadop(LoadOps.COPY, self.base.shape, self.dtype, device, src=self.base)._view(self.st)
return self.base._copy(device)._view(self.st)
def e(self, op:Union[LoadOps, UnaryOps, BinaryOps, TernaryOps], *in_srcs:LazyBuffer, arg:Optional[Any]=None) -> LazyBuffer:
srcs: List[LazyBuffer] = []
@@ -191,9 +194,9 @@ def _recursive_schedule(out:LazyBuffer, seen:Set[LazyBuffer], realizes:Set[LazyB
inputs: List[LazyBuffer] = []
var_vals: Dict[Variable, int] = out.st.var_vals.copy()
if out.op is LoadOps.COPY:
op, inputs = LazyOp(LoadOps.COPY, (), out.srcs[0].base), [out.srcs[0].base]
elif out.op is LoadOps.CUSTOM:
op, inputs = LazyOp(LoadOps.CUSTOM, (), out.arg), list(out.srcs)
op, inputs = LazyOp(LoadOps.COPY, (), out.srcs[0]), list(out.srcs)
elif out.op in {LoadOps.CUSTOM, LoadOps.SYNC}:
op, inputs = LazyOp(out.op, (), out.arg), list(out.srcs)
elif out.op is LoadOps.EMPTY:
op = LazyOp(LoadOps.EMPTY)
else:

2
tinygrad/ops.py
View File

@@ -19,7 +19,7 @@ class ReduceOps(Enum): SUM = auto(); MAX = auto() # noqa: E702
class BufferOps(Enum): LOAD = auto(); CONST = auto(); STORE = auto() # noqa: E702
# Ops below this line are not allowed in ASTs
class MovementOps(Enum): RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); STRIDE = auto(); AS_STRIDED = auto() # noqa: E702
class LoadOps(Enum): EMPTY = auto(); CONST = auto(); COPY = auto(); CONTIGUOUS = auto(); CUSTOM = auto() # noqa: E702
class LoadOps(Enum): EMPTY = auto(); CONST = auto(); COPY = auto(); CONTIGUOUS = auto(); CUSTOM = auto(); SYNC = auto() # noqa: E702
Op = Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, LoadOps, TernaryOps, BufferOps]
OpType = Union[Type[UnaryOps], Type[BinaryOps], Type[ReduceOps], Type[MovementOps], Type[LoadOps], Type[TernaryOps], Type[BufferOps]]

27
tinygrad/realize.py
View File

@@ -1,8 +1,8 @@
from typing import List, Dict, Optional, cast
from tinygrad.ops import LoadOps, ScheduleItem, BufferOps, GlobalCounters
from tinygrad.device import Device, Buffer, BufferCopy, BufferXfer, BufferRead, JITRunner, update_stats, InterpretedASTRunner
from tinygrad.device import Device, Buffer, BufferCopy, BufferXfer, BufferRead, JITRunner, update_stats, InterpretedASTRunner, Compiled
from tinygrad.graph import print_tree, realized_lazybuffer
from tinygrad.helpers import colored, getenv, GRAPH
from tinygrad.helpers import colored, getenv, GRAPH, cpu_time_execution, DEBUG
from tinygrad.shape.symbolic import Variable
# *** schedule running ***
@@ -13,6 +13,14 @@ class CustomOp(JITRunner):
super().__init__()
def __call__(self, rawbufs:List[Buffer], var_vals:Dict[Variable, int], wait=False, jit=False): self.fxn(*rawbufs)
class SyncOp(JITRunner):
def __init__(self, device):
self.device = device
super().__init__()
def __call__(self, rawbufs:List[Buffer], var_vals:Dict[Variable, int], wait=False, jit=False):
et = cpu_time_execution(Device[self.device].synchronize, enable=wait or DEBUG >= 1)
update_stats(colored("synchronize", "RED"), 0, 0, {}, et, 1, device=self.device)
def lower_schedule_item(si:ScheduleItem) -> Optional[JITRunner]:
assert all(si.out.device == x.device for x in si.inputs) or si.ast.op is LoadOps.COPY, \
f"all devices must be the same, {si.out.device} != {[x.device for x in si.inputs]} {print_tree(si.ast) or ''}"
@@ -22,6 +30,7 @@ def lower_schedule_item(si:ScheduleItem) -> Optional[JITRunner]:
if si.inputs[0].device.startswith("DISK"): return BufferRead()
return BufferCopy()
if si.ast.op is LoadOps.CUSTOM: return CustomOp(si.ast.arg)
if si.ast.op is LoadOps.SYNC: return SyncOp(si.out.device) if isinstance(Device[si.out.device], Compiled) else None
return Device[si.out.device].get_runner(si.ast)
logops = open(getenv("LOGOPS", ""), "a") if getenv("LOGOPS", "") else None
@@ -42,12 +51,14 @@ def run_schedule(schedule:List[ScheduleItem]):
break
# we don't have an output buffer, we have to create it, and create to max size if it has symbolic shape
si.out.realized = si.out.output_buffer if si.out.output_buffer is not None else \
Buffer(si.out.device, si.out.size, si.out.dtype, "PLACEHOLDER" if isinstance(prg, InterpretedASTRunner) else None)
del si.out.srcs
if si.out.size > 0:
si.out.realized = si.out.output_buffer if si.out.output_buffer is not None else \
Buffer(si.out.device, si.out.size, si.out.dtype, "PLACEHOLDER" if isinstance(prg, InterpretedASTRunner) else None)
del si.out.srcs
# run the function (put it in JIT)
assert all(x.realized is not None for x in si.inputs), f"can't run, some inputs aren't realized {[x for x in si.inputs if x.realized is None]}"
if prg: prg.exec([si.out.realized] + [cast(Buffer, x.realized) for x in si.inputs], si.var_vals)
else: update_stats(colored(f"empty {si.out.st.size:10d} {si.out.dtype}", "yellow"), 0, 0, {}, None, 1, device=si.out.device)
real_buffers = [x.realized for x in (si.out,)+si.inputs if x.size != 0]
assert all(x is not None for x in real_buffers), f"can't run, some inputs aren't realized {real_buffers}"
if prg: prg.exec(cast(List[Buffer], real_buffers), si.var_vals)
elif si.out.size > 0: update_stats(colored(f"empty {si.out.st.size:10d} {si.out.dtype}", "yellow"), 0, 0, {}, None, 1, device=si.out.device)
if GRAPH: realized_lazybuffer(si.out, GlobalCounters.kernel_count)

6
tinygrad/runtime/ops_hip.py
View File

@@ -1,6 +1,6 @@
from __future__ import annotations
import ctypes, functools, subprocess, io
from typing import Tuple, TypeVar, List
from typing import Tuple, TypeVar, List, Any
import gpuctypes.hip as hip
from tinygrad.helpers import DEBUG, getenv, init_c_var, compile_cuda_style, encode_args_cuda_style, time_execution_cuda_style
from tinygrad.helpers import from_mv, round_up, to_mv
@@ -45,9 +45,11 @@ CHUNK_SIZE, PAGE_SIZE = 256*1024*1024, 0x1000
class HIPAllocator(LRUAllocator):
def __init__(self, device:HIPDevice):
self.device = device
self.track_cross_device: List[HIPDevice] = []
super().__init__()
def free_cache(self):
self.device.synchronize()
for x in self.track_cross_device: x.synchronize()
return super().free_cache()
def _alloc(self, size:int):
check(hip.hipSetDevice(self.device.device))
@@ -102,6 +104,7 @@ class HIPDevice(Compiled):
self.arch = init_c_var(hip.hipDeviceProp_t(), lambda x: check(hip.hipGetDeviceProperties(x, self.device))).gcnArchName.decode() if not MOCKHIP else "gfx1100" # noqa: E501
self.pending_copyin: List[hip.hipDeviceptr_t] = []
self.pending_events: List[hip.hipEvent_t] = []
self.track_cross_buffer: List[Any] = []
from tinygrad.runtime.graph.hip import HIPGraph
super().__init__(device, MallocAllocator if MOCKHIP else HIPAllocator(self), LinearizerOptions("HIP"), HIPRenderer,
@@ -111,6 +114,7 @@ class HIPDevice(Compiled):
check(hip.hipDeviceSynchronize())
for opaque in self.pending_copyin: check(hip.hipFree(opaque))
for opaque in self.pending_events: check(hip.hipEventDestroy(opaque))
self.track_cross_buffer.clear()
self.pending_copyin.clear()
self.pending_events.clear()
def event(self):