Split UnaryOps.CAST into CAST and BITCAST #4487

wpmed92 · 2024-05-09T10:49:44Z

This PR introduces UnaryOps.BITCAST, and it is used instead of UnaryOps.CAST and bitcast=true boolean flag combination.

a = Tensor([1,2,3], dtype=dtypes.float).bitcast(dtypes.int).realize()

before

  0 ━┳ STORE MemBuffer(idx=0, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)))
  1  ┗━┳ CAST (dtypes.int, True)
  2    ┗━━ LOAD MemBuffer(idx=1, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)))
((LazyOp(op=BufferOps.STORE, src=(LazyOp(op=UnaryOps.CAST, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=1, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)))),), arg=(dtypes.int, True)),), arg=MemBuffer(idx=0, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)))),), [Opt(op=OptOps.LOCAL, axis=0, amt=3)])

after

 0 ━┳ STORE MemBuffer(idx=0, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)))
 1  ┗━┳ BITCAST dtypes.int
 2    ┗━━ LOAD MemBuffer(idx=1, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)))
((LazyOp(op=BufferOps.STORE, src=(LazyOp(op=UnaryOps.BITCAST, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=1, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)))),), arg=dtypes.int),), arg=MemBuffer(idx=0, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)))),), [Opt(op=OptOps.LOCAL, axis=0, amt=3)])

geohot · 2024-05-09T21:01:19Z

Good to see you back!

geohot · 2024-05-09T22:02:18Z

arg of cast should be changed to only be the dtype, we don't need the True/False anymore since CAST/BITCAST are different ops

tinygrad/engine/schedule.py

wpmed92 · 2024-05-12T12:38:29Z

@Qazalin @geohot

cast arg is now just the dtype (it's not a tuple anymore)
no bitcast on ImageDType
regenerated dataset so that sops.gz has the new UnaryOps.CAST and UnaryOps.BITCAST in the ASTs (for this I also had to update the UNROLL action in search, because I got a wrong action for UNROLL, axis=4 for the new dataset)

tinygrad/codegen/kernel.py

tinygrad/lazy.py

tinygrad/codegen/kernel.py

Qazalin · 2024-05-13T07:29:03Z

nice - new graph:

I'm not sure about the wrong UNROLL action though - What was the error?

wpmed92 · 2024-05-13T07:34:23Z

I was getting this for PYTHONPATH="." GPU=1 python3 extra/optimization/get_action_space.py:

github-actions · 2024-05-13T15:04:01Z

Changes

Name                              Lines    Diff    Tokens/Line    Diff
------------------------------  -------  ------  -------------  ------
tinygrad/codegen/kernel.py          460      +0           18.4    -0.0
tinygrad/codegen/linearizer.py      337      +0           19.3    +0.0
tinygrad/engine/schedule.py         243      +0           14.0    -0.0
tinygrad/lazy.py                    160      +1           18.9    -0.2
tinygrad/ops.py                     108      +1           17.4    +0.2


total lines changes: +2

Qazalin · 2024-05-13T19:03:05Z

did you try bumping the DB version for the actions issue? you might just need to revalidate the sqlite db instead of changing action space script.

wpmed92 · 2024-05-13T21:03:12Z

Still wrong after db bump, only seems to be solved by bumping axis in unroll.

chenyuxyz · 2024-05-13T21:19:25Z

the error is saying the handcoded opt had used a action that's not in the search action lists. can you print the kernel that caused this?

This reverts commit dee6911.

wpmed92 · 2024-05-15T09:04:43Z

@chenyuxyz

(LazyOp(op=BufferOps.STORE, src=(LazyOp(op=ReduceOps.SUM, src=(LazyOp(op=BinaryOps.MUL, src=(LazyOp(op=UnaryOps.CAST, src=(LazyOp(op=BinaryOps.MUL, src=(LazyOp(op=BinaryOps.MUL, src=(LazyOp(op=BinaryOps.MUL, src=(LazyOp(op=BinaryOps.MUL, src=(LazyOp(op=BinaryOps.CMPEQ, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=1, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(2, 3, 4, 5, 6, 1, 2, 5, 6, 5, 3, 4), strides=(0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)))), LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=2, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(2, 3, 4, 5, 6, 1, 2, 5, 6, 5, 3, 4), strides=(1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),))))), arg=None), LazyOp(op=BinaryOps.CMPEQ, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=3, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(2, 3, 4, 5, 6, 1, 2, 5, 6, 5, 3, 4), strides=(0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)))), LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=4, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(2, 3, 4, 5, 6, 1, 2, 5, 6, 5, 3, 4), strides=(0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),))))), arg=None)), arg=None), LazyOp(op=BinaryOps.CMPEQ, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=5, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(2, 3, 4, 5, 6, 1, 2, 5, 6, 5, 3, 4), strides=(0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0), offset=0, mask=None, contiguous=False),)))), LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=6, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(2, 3, 4, 5, 6, 1, 2, 5, 6, 5, 3, 4), strides=(0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),))))), arg=None)), arg=None), LazyOp(op=BinaryOps.CMPEQ, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=3, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(2, 3, 4, 5, 6, 1, 2, 5, 6, 5, 3, 4), strides=(0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0), offset=0, mask=None, contiguous=False),)))), LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=7, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(2, 3, 4, 5, 6, 1, 2, 5, 6, 5, 3, 4), strides=(5, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),))))), arg=None)), arg=None), LazyOp(op=BinaryOps.CMPEQ, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=8, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(2, 3, 4, 5, 6, 1, 2, 5, 6, 5, 3, 4), strides=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0), offset=0, mask=None, contiguous=False),)))), LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=9, dtype=dtypes.int, st=ShapeTracker(views=(View(shape=(2, 3, 4, 5, 6, 1, 2, 5, 6, 5, 3, 4), strides=(0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),))))), arg=None)), arg=None),), arg=dtypes.float), LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=10, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(2, 3, 4, 5, 6, 1, 2, 5, 6, 5, 3, 4), strides=(1440, 480, 120, 24, 4, 0, 0, 0, 0, 0, 0, 1), offset=0, mask=None, contiguous=False),))))), arg=None),), arg=(0, 1, 2, 3, 4)),), arg=MemBuffer(idx=0, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(1, 1, 1, 1, 1, 1, 2, 5, 6, 5, 3, 4), strides=(0, 0, 0, 0, 0, 0, 1800, 360, 60, 12, 4, 1), offset=0, mask=None, contiguous=True),)))),)

chenyuxyz · 2024-05-15T15:43:26Z

it triggered a new tc, might be related to #4427. @flammit fyi

chenyuxyz · 2024-05-15T15:44:23Z

congrats! george@tinygrad.org to claim

* Separate cast and bitcast * Fix lint * No more arg[0] * Revert "No more arg[0]" This reverts commit dee6911. * CAST/BITCAST arg is the dtype only, no more tuple * No image bitcast, regenerate dataset * Small fixes

wpmed92 changed the title ~~Separate cast and bitcast~~ Separate UnaryOps.CAST into CAST and BITCAST May 9, 2024

wpmed92 changed the title ~~Separate UnaryOps.CAST into CAST and BITCAST~~ Split UnaryOps.CAST into CAST and BITCAST May 9, 2024

geohot added the bounty locked Bounty is locked to someone label May 9, 2024

wpmed92 force-pushed the bitcast-unary-op branch 3 times, most recently from 99157dd to fbb09ae Compare May 11, 2024 07:35