mirrors/tinygrad
2
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-24 02:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity

dedup COPY UOp [pr] (#8506)

Browse source
This commit is contained in:
qazal 2025-01-06 10:37:20 +02:00 • committed by GitHub
commit eb7df92136
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 7 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

12
test/test_schedule.py
View file

@ -230,16 +230,16 @@ class TestSchedule(unittest.TestCase):
# a and b share the same underlying device memory
self.assertIs(a.lazydata.realized, b.lazydata.realized)
# EMPTY and COPY are assigned to unique device Buffers
def test_no_dedup_copy(self):
def test_copy_dedups(self):
src = Tensor.ones(4).contiguous().realize()
a = src.clone()
b = src.clone()
sched = check_schedule([a, b], 2, filter_sink=False)
sched = check_schedule([a, b], 1, filter_sink=False)
run_schedule(sched)
# a and b are assigned to different device Buffers
self.assertIsNot(a.lazydata.realized, b.lazydata.realized)
# a and b are assigned to the same device Buffer
self.assertIs(a.lazydata.realized, b.lazydata.realized)
# EMPTY is assigned to a unique device Buffer
def test_no_dedup_empty(self):
a = Tensor.empty((4,))

1
tinygrad/ops.py
View file

@ -444,6 +444,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
if op is Ops.BIND:
assert isinstance(arg, UOp) and arg.op is Ops.BIND and shape == (), f"trying to create BIND with {arg=} {shape=}"
return UOp(Ops.VIEW, dtype, (UOp(Ops.DEVICE, arg=device), arg), ShapeTracker.from_shape(()))
if op is Ops.COPY: return UOp(op, dtype, src, arg)
# otherwise it's a contiguous st
return UOp(Ops.VIEW, dtype, (UOp.new_buffer(device, (st:=ShapeTracker.from_shape(shape)).size, dtype), UOp(op, dtype, src, arg)), st)
def copy_to_device(self, device:str, force=False, clone:bool=False) -> UOp: