github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 23:48:01 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix kernelize with VIEW children (#9961)

Browse Source
This commit is contained in:
qazal
2025-04-21 18:38:46 +03:00
committed by GitHub
parent 757533cbe6
commit 36ed3c3253
2 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
test/test_schedule.py
View File

@@ -597,6 +597,12 @@ class TestSchedule(unittest.TestCase):
d = c+2
check_schedule(d, 2)
def test_kernelize_view(self):
a = Tensor.empty(4,1)
b = a*2
c = b.kernelize()+Tensor.empty(4,4)
check_schedule(c, 2)
# unlike schedule, kernelize can be called multiple times on a Tensor
def test_double_kerenlize(self):
a = Tensor.empty(10)

2
tinygrad/spec.py
View File

@@ -35,7 +35,7 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([
# "make things that can't be images not images" can change the buffer dtype
# this is fine as long as it's a realized buffer and base dtypes match.
((isinstance(mv.dtype, ImageDType) or isinstance(x.dtype, ImageDType)) and x.dtype.base == mv.dtype.base and x.base.op is Ops.BUFFER)),
(UPat(Ops.VIEW, src=(UPat(GroupOp.All-{Ops.BUFFER, Ops.BUFFER_VIEW, Ops.CONST, Ops.DEVICE}),)), lambda: False),
(UPat(Ops.VIEW, src=(UPat(GroupOp.All-{Ops.BUFFER, Ops.BUFFER_VIEW, Ops.ASSIGN, Ops.CONST, Ops.DEVICE}),)), lambda: False),
# Tensor variable bindings
(UPat(Ops.BIND, dtypes.int, (UPat(Ops.DEFINE_VAR), UPat.cvar(dtype=dtypes.int)), arg=None), lambda: True),