fix bitcast

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()

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