diff --git a/test/test_rangeify.py b/test/test_rangeify.py index a7b04338a6..0e4aa40249 100644 --- a/test/test_rangeify.py +++ b/test/test_rangeify.py @@ -9,6 +9,13 @@ class TestRangeify(unittest.TestCase): C = Tensor.empty(N, N) (A@B@C).realize() + def test_double_gemm_exp(self): + N = 1024 + A = Tensor.empty(N, N) + B = Tensor.empty(N, N) + C = Tensor.empty(N, N) + (((A@B).exp()@C).exp()).realize() + def test_double_gemm_relu(self): N = 1024 A = Tensor.empty(N, N) diff --git a/test/test_tiny.py b/test/test_tiny.py index bcf69162da..905996bf9e 100644 --- a/test/test_tiny.py +++ b/test/test_tiny.py @@ -91,12 +91,12 @@ class TestTiny(unittest.TestCase): @unittest.skipIf(IMAGE>0 or (CI and Device.DEFAULT == "DSP"), "failing because of make things that can't be images not images") def test_mnist(self): layers = [ - nn.Conv2d(1, 32, 5), Tensor.relu, Tensor.contiguous, - nn.Conv2d(32, 32, 5), Tensor.relu, Tensor.contiguous, - nn.BatchNorm(32), Tensor.max_pool2d, Tensor.contiguous, - nn.Conv2d(32, 64, 3), Tensor.relu, Tensor.contiguous, - nn.Conv2d(64, 64, 3), Tensor.relu, Tensor.contiguous, - nn.BatchNorm(64), Tensor.max_pool2d, Tensor.contiguous, + nn.Conv2d(1, 32, 5), Tensor.relu, #Tensor.contiguous, + nn.Conv2d(32, 32, 5), Tensor.relu, #Tensor.contiguous, + nn.BatchNorm(32), Tensor.max_pool2d, #Tensor.contiguous, + nn.Conv2d(32, 64, 3), Tensor.relu, #Tensor.contiguous, + nn.Conv2d(64, 64, 3), Tensor.relu, #Tensor.contiguous, + nn.BatchNorm(64), Tensor.max_pool2d, #Tensor.contiguous, lambda x: x.flatten(1), nn.Linear(576, 10)] # replace random weights with ones diff --git a/tinygrad/apps/llm.py b/tinygrad/apps/llm.py index b717d5c9fb..40a5f2e4f7 100644 --- a/tinygrad/apps/llm.py +++ b/tinygrad/apps/llm.py @@ -116,7 +116,7 @@ class TransformerBlock: def _feed_forward(self, h: Tensor) -> Tensor: h_norm = self.ffn_norm(h) - gated = (self.ffn_gate(h_norm).silu() * self.ffn_up(h_norm)).contiguous(arg=(1,)) + gated = (self.ffn_gate(h_norm).silu() * self.ffn_up(h_norm)) #.contiguous(arg=(1,)) return h + self.ffn_down(gated) def __call__(self, x: Tensor, start_pos: int|UOp): diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 45c98da724..e2aa7dac6a 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -15,8 +15,10 @@ class RangeifyContext: #indexed_child: dict[UOp, list[UOp]] = field(default_factory=dict) #ended_child: dict[UOp, int] = field(default_factory=dict) +reduce_map = {} def map_contiguous(ctx:RangeifyContext, x:UOp, idx:UOp|None=None): if x.tag == 1: return None + if x in reduce_map: return reduce_map[x] ranges = [] new_ranges = [] passthrough_idx = [] @@ -45,7 +47,9 @@ def map_contiguous(ctx:RangeifyContext, x:UOp, idx:UOp|None=None): new_ranges.extend(sub.values()) ret = x.src[0].index(*ranges).contiguous(*new_ranges, arg=x.shape, tag=1) - return ret.index(*passthrough_idx) if idx is not None else ret + ret = ret.index(*passthrough_idx) if idx is not None else ret + reduce_map[x] = ret + return ret def map_reshape(x:UOp, r:UOp): acc = 1 @@ -81,7 +85,6 @@ def map_pad(x:UOp, r:UOp): # PAD is with 0 return bigwhere.simplify().where(UOp(Ops.INDEX, r.dtype, src=(r.src[0],)+tuple(ret)), UOp.const(r.dtype, 0)) -reduce_map = {} def map_reduce(ctx:RangeifyContext, idx:UOp, red:UOp): # TODO: this should be in the cache if red in reduce_map: return reduce_map[red] @@ -153,6 +156,26 @@ def might_end_axes(x:UOp): if r.op is Ops.ADD and r.src[0].op is Ops.RANGE and r.src[1].op is Ops.RANGE: end_axes.append(i) return x.replace(src=(x.src[0].contiguous(arg=tuple(end_axes)),)+x.src[1:]) if end_axes else None +def map_expand(r:UOp, x:UOp): + new_rngs = [] + ending_ranges = [] + axis_to_range = [] + for a,x,y in zip(x.src[1:], r.src[0].shape, r.shape): + axis_to_range.append([u.arg[0] for u in a.toposort() if u.op is Ops.RANGE]) + if resolve(x!=y, False): + ending_ranges.extend(axis_to_range[-1]) + new_rngs.append(a.const_like(0)) + else: new_rngs.append(a) + earliest_ending_axis = min(ending_ranges) + to_end_axis = [] + for i,a in enumerate(axis_to_range): + if any(x > earliest_ending_axis for x in a): + to_end_axis.append(i) + ret = r.src[0] + if to_end_axis and any(x.op is Ops.REDUCE_AXIS for x in ret.toposort()): + ret = ret.contiguous(arg=tuple(to_end_axis)) + return ret.index(*new_rngs) + pm_rangeify = PatternMatcher([ # if there are new ended children, tag the SINK (UPat(Ops.INDEX, src=(UPat(Ops.CHILD, src=(UPat(name="c"), ), name="x"),), allow_any_len=True, name="idx"), index_child), @@ -168,13 +191,15 @@ pm_rangeify = PatternMatcher([ lambda r,x: r.src[0].index(*[a+ss if resolve(ss != 0) else a for a,(ss,_) in zip(x.src[1:], r.arg)])), (UPat(Ops.INDEX, src=(UPat(Ops.FLIP, name="r"),), allow_any_len=True, name="x"), lambda r,x: r.src[0].index(*[((s-1)-a) if f else a for a,s,f in zip(x.src[1:], r.shape, r.arg)])), - (UPat(Ops.INDEX, src=(UPat(Ops.EXPAND, name="r"),), allow_any_len=True, name="x"), - lambda r,x: r.src[0].index(*[a.const_like(0) if resolve(x!=y, False) else a for a,x,y in zip(x.src[1:], r.src[0].shape, r.shape)])), + #(UPat(Ops.INDEX, src=(UPat(Ops.EXPAND, name="r"),), allow_any_len=True, name="x"), + # lambda r,x: r.src[0].index(*[a.const_like(0) if resolve(x!=y, False) else a for a,x,y in zip(x.src[1:], r.src[0].shape, r.shape)])), + # lambda r,x: r.src[0].index(*[a.const_like(0) if resolve(x!=y, False) else a for a,x,y in zip(x.src[1:], r.src[0].shape, r.shape)])), + (UPat(Ops.INDEX, src=(UPat(Ops.EXPAND, name="r"),), allow_any_len=True, name="x"), map_expand), (UPat(Ops.INDEX, src=(UPat(Ops.RESHAPE, name="r"),), allow_any_len=True, name="x"), map_reshape), (UPat(Ops.INDEX, src=(UPat(Ops.PAD, name="r"),), allow_any_len=True, name="x"), map_pad), # might end axes - (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.REDUCE_AXIS})),), allow_any_len=True, name="x"), might_end_axes), + #(UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.REDUCE_AXIS})),), allow_any_len=True, name="x"), might_end_axes), # move MAP through elementwise ALU / reduce. these are the items with cost (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.STORE})),), allow_any_len=True, name="x"),