From e015b41ce9dcfa841e3736ab7755bb07334dc2a0 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Thu, 19 Sep 2024 10:24:02 +0800 Subject: [PATCH] remove e( function just alu( [run_process_replay] (#6589) * remove e( function just alu( [run_process_replay] * missed two --- docs/abstractions2.py | 2 +- test/test_schedule.py | 8 ++++---- tinygrad/lazy.py | 7 +++---- tinygrad/multi.py | 16 ++++++++-------- 4 files changed, 16 insertions(+), 17 deletions(-) diff --git a/docs/abstractions2.py b/docs/abstractions2.py index 528f44a8ab..fe9e140a44 100644 --- a/docs/abstractions2.py +++ b/docs/abstractions2.py @@ -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]) diff --git a/test/test_schedule.py b/test/test_schedule.py index 91acece3a7..c0d10f5e83 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -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): diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py index a9c8ea715d..47bc6803dc 100644 --- a/tinygrad/lazy.py +++ b/tinygrad/lazy.py @@ -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 *** diff --git a/tinygrad/multi.py b/tinygrad/multi.py index 0c1217bc53..5ccd096e2b 100644 --- a/tinygrad/multi.py +++ b/tinygrad/multi.py @@ -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)