github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 15:08:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

rangeify: fix abstractions2.py (#12386)

Browse Source 
* rangeify: fix abstractions2.py

* tests

* lint

* only abstractions2

* base
This commit is contained in:
b1tg
2025-10-01 14:58:56 +08:00
committed by GitHub
parent fe96c8d345
commit 154d114364
2 changed files with 7 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
docs/abstractions2.py
View File

@@ -80,7 +80,9 @@ print("******** third, the UOp ***********")
from tinygrad.engine.realize import run_schedule
from tinygrad.engine.schedule import create_schedule_with_vars
from tinygrad.helpers import RANGEIFY
from tinygrad.schedule.kernelize import get_kernelize_map
from tinygrad.schedule.rangeify import get_rangeify_map
# allocate some values + load in values
a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
@@ -93,10 +95,10 @@ out = a + b
s = UOp(Ops.SINK, dtypes.void, (out,))
# group the computation into kernels
becomes_map = get_kernelize_map(s)
becomes_map = get_rangeify_map(s) if RANGEIFY else get_kernelize_map(s)
# the compute maps to an assign
assign = becomes_map[a+b]
assign = becomes_map[a+b].base
# the first source is the output buffer (data)
assert assign.src[0].op is Ops.BUFFER