tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 14:56:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

insert Ops.FUSE for arange (#10140)

Browse Source 
* insert Ops.FUSE for arange

* reshape does not collapse

* do not fuse reshapes

* add children

* fixups

* work

* add Ops.WHERE support to z3

* fix fuse for cast

* diff

* ugh

* don't need this anymore

* contiguous

* add always_contiguous

* there too
This commit is contained in:
qazal
2025-05-05 08:32:12 +03:00
committed by GitHub
parent cf626e23cb
commit 62e86bc5ec
2 changed files with 48 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_schedule.py
+12 -11
View File
@@ -1688,9 +1688,9 @@ class TestIndexing(unittest.TestCase):
Tensor.manual_seed(0)
a = Tensor.arange(4,)
b = Tensor.randn(4, 4).realize()
out = a+b
out = (a+b).sum()
self.check_schedule(out, 1)
np.testing.assert_allclose(out.numpy(), np.arange(4)+b.numpy())
np.testing.assert_allclose(out.numpy(), (np.arange(4)+b.numpy()).sum(), atol=1e-5)
def test_argmin(self):
Tensor.manual_seed(0)
@@ -1709,25 +1709,25 @@ class TestIndexing(unittest.TestCase):
def test_arange_transposed(self):
Tensor.manual_seed(0)
x = Tensor.randint(4, 1).realize()
a = (Tensor.arange(4,)*x).T
a = ((Tensor.arange(4,)*x).T).sum()
self.check_schedule(a, 1)
np.testing.assert_equal(a.numpy(), (np.arange(4)*x.numpy()).T)
np.testing.assert_equal(a.numpy(), (np.arange(4)*x.numpy()).T.sum())
def test_arange_transposed_descendants(self):
Tensor.manual_seed(0)
x = Tensor.randint(4, 1).realize()
a = (Tensor.arange(4,)*x).T
b = Tensor.randint(4, 4).realize()
out = a+b
out = (a+b).sum()
self.check_schedule(out, 1)
np.testing.assert_equal(out.numpy(), (np.arange(4)*x.numpy()).T+b.numpy())
np.testing.assert_equal(out.numpy(), ((np.arange(4)*x.numpy()).T+b.numpy()).sum())
def test_arange_index(self):
Tensor.manual_seed(0)
x = Tensor.randn(5, 2).realize()
a = Tensor.arange(10)
out = (x + a[2]).sum()
self.check_schedule(out, 2)
self.check_schedule(out, 1)
np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6)
@unittest.skip("TOOD: FUSE_ARANGE overrules Tensor.arange().contiguous()")
@@ -1744,7 +1744,7 @@ class TestIndexing(unittest.TestCase):
x = Tensor.randn(5, 2).realize()
a = Tensor.arange(10)+1
out = (x + a[2]).sum()
self.check_schedule(out, 2)
self.check_schedule(out, 1)
np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6)
@unittest.skip("TOOD: FUSE_ARANGE overrules Tensor.arange().contiguous()")
@@ -1791,15 +1791,16 @@ class TestIndexing(unittest.TestCase):
@unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from ext device")
def test_arange_shrink_copy(self):
a = Tensor.arange(12).reshape(4, 3).shrink(((1, 2), (1, 3))).to("CPU")
sched = self.check_schedule(a, 1)
sched = self.check_schedule(a, 2) # NOTE: there is a contiguous between REDUCE_AXIS and COPY
self.assertIs(sched[-1].ast.op, Ops.COPY)
np.testing.assert_equal(a.numpy(), [[4, 5]])
@unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from ext device")
def test_arange_expand_copy(self):
a = Tensor.arange(4).reshape(2, 2, 1).expand(2, 2, 2).contiguous().to("CPU")
sched = self.check_schedule(a, 1)
self.assertIs(sched[1].ast.op, Ops.COPY)
sched = self.check_schedule(a, 2) # NOTE: there is a contiguous between REDUCE_AXIS and COPY
self.assertIs(sched[2].ast.op, Ops.COPY)
self.assertIs(sched[1].ast.src[0].src[2].op, Ops.LOAD)
self.assertIs(sched[0].ast.src[0].src[2].op, Ops.ADD)
np.testing.assert_equal(a.numpy(), [[[0, 0], [1, 1]], [[2, 2], [3, 3]]])
tinygrad/engine/grouper.py
+36 -8
View File
@@ -4,7 +4,7 @@ from tinygrad.ops import UOp, Ops, GroupOp, PatternMatcher, UPat, graph_rewrite,
from tinygrad.ops import can_pad, sint, track_rewrites, _substitute
from tinygrad.codegen.lowerer import get_contraction_with_reduce
from tinygrad.codegen.symbolic import symbolic_simple
from tinygrad.helpers import Metadata, all_int, all_same, colored, prod, dedup, unwrap, flatten, getenv, pluralize, ContextVar, Context, diskcache_put
from tinygrad.helpers import Metadata, all_int, all_same, colored, prod, dedup, unwrap, getenv, pluralize, ContextVar, Context, diskcache_put
from tinygrad.helpers import FUSE_CONV_BW, FUSE_ARANGE, DEBUG, DONT_REALIZE_EXPAND, DONT_GROUP_REDUCES, SPLIT_REDUCEOP, CAPTURE_PROCESS_REPLAY
from tinygrad.dtype import ImageDType, dtypes
from tinygrad.shape.shapetracker import ShapeTracker
@@ -223,11 +223,6 @@ def group_realizes(sink:UOp) -> dict[UOp, None]:
group = {tr: None}
realizes[tr] = None
reduce_for_op.update((tr, r) for tr in group)
if FUSE_ARANGE and r.arg[0] is Ops.ADD and r.src[0].base.op is Ops.CONST and (getenv("FUSE_ARANGE_UINT", 1) or not dtypes.is_unsigned(r.dtype)):
# maybe fuse arange with its children
# TODO: FUSE_ARANGE_UINT is for not fusing rand. should not be here.
if len(flatten(children[tr] for tr in group)) != 0:
for tr in group: del realizes[tr]
# fuse double reduces with no other child
for reduceop in double_reduces:
top_reduce = reduceop.src[0].base
@@ -427,7 +422,7 @@ pm_fuse = PatternMatcher([
lambda r: r.replace(src=(r.src[0].fuse(),), arg=r.arg+(True,)) if len(r.arg) == 2 else None),
# FUSE elementwise. TODO: check for PAD
(UPat(Ops.VIEW, src=(UPat(GroupOp.ALU, name="alu"),), name="view").fuse(),
(UPat(Ops.VIEW, src=(UPat({*GroupOp.ALU, Ops.CAST}, name="alu"),), name="view").fuse(),
lambda alu, view: alu.replace(src=tuple(x.view(view.arg).fuse() for x in alu.src))),
# push FUSE through to srcs
@@ -444,6 +439,39 @@ def do_fusion(x:UOp):
return graph_rewrite(x.substitute(found_contiguous), pm_fuse, name="local fusion").substitute({v:k for k,v in found_contiguous.items()})
do_fuse = PatternMatcher([(UPat(Ops.FUSE, name="x"), do_fusion),])
def fuse_arange(root:UOp):
# skip if root is arange
if not FUSE_ARANGE or (dtypes.is_unsigned(root.dtype) and not getenv("FUSE_ARANGE_UINT",1)) or root.src[0].base.op is Ops.CONST: return None
# gather all local aranges (including any fused ones)
local_arange: list[UOp] = []
def gate_reduce(u):
if u.op is Ops.REDUCE_AXIS and u.src[0].base.op is Ops.CONST: local_arange.append(u)
return u.op not in {*ALWAYS_CONTIGUOUS, Ops.REDUCE_AXIS} or u is root
toposort = root.toposort(gate=gate_reduce)
if not local_arange: return None
# fuse the nearest expand child of arange
local_children: dict[UOp, list[UOp]] = {}
for u in toposort:
for s in u.src: local_children.setdefault(s, []).append(u)
fuse_rep: dict[UOp, UOp] = {}
for r in local_arange:
# skip if already fused
if len(r.arg) > 2: continue
q = list(local_children[r])
while q:
u = q.pop()
if not (curr_children:=local_children.get(u, [])): continue
for child in curr_children:
other_paths = {s for s in child.toposort() if s.op in {Ops.REDUCE_AXIS, Ops.BUFFER} and s not in {root, r}}
fuse_rep[child] = child.replace(src=tuple(s.fuse() if s is u else s for s in child.src))
if other_paths: break
else: q.extend(curr_children)
return root.substitute(fuse_rep, name="fuse_arange") if fuse_rep else None
insert_fuse = PatternMatcher([
(UPat(Ops.REDUCE_AXIS, name="root"), fuse_arange),
])
PROCESS_REPLAY_CAPTURE:dict[str, bytes] = {}
if CAPTURE_PROCESS_REPLAY:
import atexit
@@ -458,7 +486,7 @@ def get_name(becomes_map:dict[UOp, UOp]) -> str:
@track_rewrites(name_fxn=get_name)
def get_becomes_map(big_sink:UOp) -> dict[UOp, UOp]:
# merge_views + simplify
tensor_map = graph_rewrite_map(big_sink, do_fuse+merge_views+sym+replace_contiguous, ctx={})
tensor_map = graph_rewrite_map(big_sink, insert_fuse+do_fuse+merge_views+sym+replace_contiguous, ctx={})
# display the cleaned up tensor graph
if getenv("VIZ"): graph_rewrite(tensor_map[big_sink], PatternMatcher([]), name="View Tensor Graph")