From afb0463f103d8fb6406145bc2944e1dbd18a010c Mon Sep 17 00:00:00 2001 From: George Hotz Date: Thu, 25 Jun 2026 10:11:13 -0700 Subject: [PATCH] fix bitcast --- test/null/test_bitcast_spec.py | 21 +++++++++++++++------ tinygrad/uop/ops.py | 9 ++++++--- 2 files changed, 21 insertions(+), 9 deletions(-) diff --git a/test/null/test_bitcast_spec.py b/test/null/test_bitcast_spec.py index 3a281a8431..f08cedd90f 100644 --- a/test/null/test_bitcast_spec.py +++ b/test/null/test_bitcast_spec.py @@ -1,36 +1,45 @@ import unittest from tinygrad import UOp, dtypes +from tinygrad.uop.ops import shape_to_shape_arg, ParamArg, Ops, AddrSpace + +def placeholder(shape, dtype, slot): + return UOp(Ops.PARAM, dtype, (shape_to_shape_arg(shape),), arg=ParamArg(slot, AddrSpace.GLOBAL)) class TestBitcastSpec(unittest.TestCase): def test_bitcast_no_shape_change(self): - pl = UOp.placeholder((10,10), dtypes.int, 0) + pl = placeholder((10,10), dtypes.int, 0) out = pl.bitcast(dtypes.float) self.assertEqual(out.shape, (10,10)) def test_bitcast_increase_shape(self): - pl = UOp.placeholder((10,10), dtypes.int, 0) + pl = placeholder((10,10), dtypes.int, 0) out = pl.bitcast(dtypes.short) self.assertEqual(out.shape, (10,20)) def test_bitcast_decrease_shape(self): - pl = UOp.placeholder((10,10), dtypes.int, 0) + pl = placeholder((10,10), dtypes.int, 0) out = pl.bitcast(dtypes.long) self.assertEqual(out.shape, (10,5)) def test_bitcast_remove_ones(self): - pl = UOp.placeholder((10,2), dtypes.int, 0) + pl = placeholder((10,2), dtypes.int, 0) out = pl.bitcast(dtypes.long) self.assertEqual(out.shape, (10,)) def test_bitcast_remove_ones_full(self): - pl = UOp.placeholder((2,), dtypes.int, 0) + pl = placeholder((2,), dtypes.int, 0) out = pl.bitcast(dtypes.long) self.assertEqual(out.shape, ()) def test_bitcast_add_ones_full(self): - pl = UOp.placeholder((), dtypes.long, 0) + pl = placeholder((), dtypes.long, 0) out = pl.bitcast(dtypes.int) self.assertEqual(out.shape, (2,)) + def test_bitcast_add_ones_full_uchar(self): + pl = placeholder((), dtypes.long, 0) + out = pl.bitcast(dtypes.uchar) + self.assertEqual(out.shape, (8,)) + if __name__ == '__main__': unittest.main() diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index f7e7cd94e1..6200765879 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -312,10 +312,13 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # TODO: disallow shape changing bitcast case Ops.BITCAST: - ps = self.src[0]._shape - if ps is None: return None + ps = self.src[0].shape if (output_sz:=self.dtype.itemsize) != (input_sz:=self.src[0].dtype.itemsize): - return ps[:-1]+(ssimplify((ps[-1]*input_sz) // output_sz),) if len(ps) > 0 else ps + if ps == (): + if output_sz > input_sz: raise RuntimeError(f"shape () must be an expanding bitcast {output_sz=} {input_sz=}") + ps = (1,) + ps = ps[:-1]+(ssimplify((ps[-1]*input_sz) // output_sz),) + return ps[:-1] if ps[-1] == 1 else ps # remove trailing one return ps # MULTI marker has no shape