geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

remove e( function just alu( [run_process_replay] (#6589)

Browse Source 
* remove e( function just alu( [run_process_replay]

* missed two
This commit is contained in:
George Hotz
2024-09-19 10:24:02 +08:00
committed by GitHub
parent fa0f678d5a
commit e015b41ce9
4 changed files with 16 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/abstractions2.py
+1 -1
View File
@@ -89,7 +89,7 @@ del a.srcs
del b.srcs
# describe the computation
out = a.e(BinaryOps.ADD, b)
out = a.alu(BinaryOps.ADD, b)
# schedule the computation as a list of kernels
sched = create_schedule([out])
test/test_schedule.py
+4 -4
View File
@@ -1269,16 +1269,16 @@ class TestSchedule(unittest.TestCase):
@unittest.skipIf(Device.DEFAULT not in view_supported_devices, "subbuffer not supported")
def test_bitcast_subbufer(self):
x = cast(LazyBuffer, Tensor.empty(1, dtype=dtypes.float32).realize().lazydata)
a = x.e(UnaryOps.EXP2).cast(dtypes.int32, True, allow_buffer_view=True)
a = x.alu(UnaryOps.EXP2).cast(dtypes.int32, True, allow_buffer_view=True)
b = x.cast(dtypes.int32, True, allow_buffer_view=True)
b = a.e(BinaryOps.ADD, b)
b = a.alu(BinaryOps.ADD, b)
check_schedule(b, 2) # this should fuse when it makes sense
def test_bitcast_disable_subbufer(self):
x = cast(LazyBuffer, Tensor.empty(1, dtype=dtypes.float32).realize().lazydata)
a = x.e(UnaryOps.EXP2).cast(dtypes.int32, True, allow_buffer_view=False)
a = x.alu(UnaryOps.EXP2).cast(dtypes.int32, True, allow_buffer_view=False)
b = x.cast(dtypes.int32, True, allow_buffer_view=False)
b = a.e(BinaryOps.ADD, b)
b = a.alu(BinaryOps.ADD, b)
check_schedule(b, 1)
def test_reduceop_reshape_dont_push(self):
tinygrad/lazy.py
+3 -4
View File
@@ -86,7 +86,7 @@ class LazyBuffer(MathTrait):
def contiguous(self, allow_buffer_view=True):
if not self.st.contiguous or self.size != self.base.size or self.is_unrealized_const():
ret = self.e(MetaOps.VIEW) if allow_buffer_view and self.can_view() else self.e(MetaOps.CONTIGUOUS)
ret = self.alu(MetaOps.VIEW) if allow_buffer_view and self.can_view() else self.alu(MetaOps.CONTIGUOUS)
if (sti := self.st.invert(self.base.shape)) is not None: self.base.contiguous_child = ref(ret), sti
return ret
self.base.forced_realize = True
@@ -134,8 +134,7 @@ class LazyBuffer(MathTrait):
# copy the base and apply the shapetracker on the new device
return self.base._copy(device)._view(self.st)
def alu(self, op, *in_srcs): return self.e(op, *in_srcs)
def e(self, op:Union[MetaOps, UnaryOps, BinaryOps, TernaryOps], *in_srcs:LazyBuffer, arg:Optional[Any]=None) -> LazyBuffer:
def alu(self, op:Union[MetaOps, UnaryOps, BinaryOps, TernaryOps], *in_srcs:LazyBuffer) -> LazyBuffer:
srcs: List[LazyBuffer] = []
for s in (self,)+in_srcs:
if s == s.base and s.base.contiguous_child and (root:=s.base.contiguous_child[0]()) is not None:
@@ -162,7 +161,7 @@ class LazyBuffer(MathTrait):
if y.is_unrealized_unmasked_const() and (val := y.base.arg) in (1, 0): return x if val == 1 else x.const_like(0)
if op is BinaryOps.IDIV and y.is_unrealized_unmasked_const() and y.base.arg == 1: return x
return create_lazybuffer(self.device, ShapeTracker.from_shape(self.shape), out_dtype, op, arg, tuple(srcs))
return create_lazybuffer(self.device, ShapeTracker.from_shape(self.shape), out_dtype, op, None, tuple(srcs))
# *** reduce ops ***
tinygrad/multi.py
+8 -8
View File
@@ -1,5 +1,5 @@
from __future__ import annotations
from typing import Optional, Union, Any, Tuple, List, Dict
from typing import Optional, Union, Tuple, List, Dict
import functools, itertools, operator
from tinygrad.helpers import all_same, all_int, dedup, prod, DEBUG, RING, getenv
from tinygrad.dtype import DType
@@ -18,7 +18,7 @@ def all_reduce(op: ReduceOps, lbs: List[LazyBuffer]) -> List[LazyBuffer]:
use_ring = (RING >= 2 or (n_lbs > 2 and dim > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1))
if DEBUG >= 2: print(f"{'RING ALLREDUCE' if use_ring else 'NAIVE ALLREDUCE'} {n_lbs}x{dim} | {lbs[0].dtype}")
if not use_ring:
return [functools.reduce(lambda x,y: x.e(bop, y), [x.copy_to_device(lb.device) for x in lbs]) for lb in lbs]
return [functools.reduce(lambda x,y: x.alu(bop, y), [x.copy_to_device(lb.device) for x in lbs]) for lb in lbs]
factor = max(f for f in [32, 16, 8, 4, 2, 1] if dim % f == 0)
base, left = (dim // factor) // n_lbs, (dim // factor) % n_lbs
c_lens = [(base + 1) * factor if i < left else base * factor for i in range(n_lbs)]
@@ -30,7 +30,7 @@ def all_reduce(op: ReduceOps, lbs: List[LazyBuffer]) -> List[LazyBuffer]:
for step in range(n_lbs - 1):
for i in range(len(chunks)):
s, r = (i+step)%n_lbs, (i+step+1)%n_lbs
chunked[r][i] = chunked[r][i].e(bop, chunked[s][i].copy_to_device(chunked[r][i].device, force=True))
chunked[r][i] = chunked[r][i].alu(bop, chunked[s][i].copy_to_device(chunked[r][i].device, force=True))
# Allgather step
for step in range(n_lbs - 1):
@@ -40,7 +40,8 @@ def all_reduce(op: ReduceOps, lbs: List[LazyBuffer]) -> List[LazyBuffer]:
# Assemble chunks back
pads = [((s,dim-e),) for s,e in chunks]
return [functools.reduce(lambda x,y: x.e(BinaryOps.ADD, y), [c.pad(pads[i]) for i,c in enumerate(lb_c)]).reshape(lbs[0].shape) for lb_c in chunked]
return [functools.reduce(lambda x,y: x.alu(BinaryOps.ADD, y),
[c.pad(pads[i]) for i,c in enumerate(lb_c)]).reshape(lbs[0].shape) for lb_c in chunked]
def to_sharded(lbs:List[LazyBuffer], axis:int, bounds: Tuple[Tuple[int, int], ...]) -> List[LazyBuffer]:
if DEBUG >= 3 and lbs[0].shape[axis] % len(lbs) != 0: print(f"multi axis uneven: {lbs[0].shape=} {axis=} {len(lbs)=}, bounds={bounds}")
@@ -85,7 +86,7 @@ class MultiLazyBuffer(MathTrait):
if not real: continue
pad_arg = tuple((0,0) if a != self.axis else (start, self.bounds[-1][1]-end) for a in range(len(lb.shape)))
llbs.append(lb.copy_to_device(device).pad(pad_arg))
return functools.reduce(lambda x,y: x.e(BinaryOps.ADD, y), llbs)
return functools.reduce(lambda x,y: x.alu(BinaryOps.ADD, y), llbs)
# passthroughs
def is_realized(self) -> bool: return all(lb.base.realized is not None for lb, r in zip(self.lbs, self.real) if r is True)
@@ -96,8 +97,7 @@ class MultiLazyBuffer(MathTrait):
def contiguous(self): return MultiLazyBuffer([x.contiguous() for x in self.lbs], self.axis, self.real)
# elementwise is simple
def alu(self, op, *in_srcs): return self.e(op, *in_srcs)
def e(self, op:Union[MetaOps, UnaryOps, BinaryOps, TernaryOps], *in_srcs:MultiLazyBuffer, arg:Optional[Any]=None) -> MultiLazyBuffer:
def alu(self, op:Union[MetaOps, UnaryOps, BinaryOps, TernaryOps], *in_srcs:MultiLazyBuffer) -> MultiLazyBuffer:
msrcs = (self,)+in_srcs
assert all(isinstance(x, MultiLazyBuffer) for x in msrcs), f"all buffers must be MultiLazyBuffer {msrcs}"
assert all_same([x.device for x in msrcs]), f"all buffers must have the same device {[x.device for x in msrcs]}"
@@ -112,7 +112,7 @@ class MultiLazyBuffer(MathTrait):
if (mlb.axis == axis and (mlb.axis is None or mlb.bounds == bounds)) or not_all_real: srcs.append(mlb.lbs)
elif mlb.axis is None and axis is not None: srcs.append(to_sharded(mlb.lbs, axis, bounds))
else: srcs.append(to_sharded([mlb.copy_to_device(lb.device) for lb in mlb.lbs], axis, bounds))
new_real_lbs:Dict[int,LazyBuffer] = {i:lsrcs[0].e(op, *lsrcs[1:], arg=arg) for i,(lsrcs,r) in enumerate(zip(zip(*srcs), new_real)) if r}
new_real_lbs:Dict[int,LazyBuffer] = {i:lsrcs[0].alu(op, *lsrcs[1:]) for i,(lsrcs,r) in enumerate(zip(zip(*srcs), new_real)) if r}
# NOTE: const dtype should match real
real_dtype = next(iter(new_real_lbs.values())).dtype
return MultiLazyBuffer([new_real_lbs.get(i, lsrcs[0].const_like(0).cast(real_dtype)) for i,lsrcs in enumerate(zip(*srcs))], axis, new_real)