From c43a3fdebb2e1a1b2f291c4d5b428fd19be00445 Mon Sep 17 00:00:00 2001 From: chenyu Date: Tue, 7 Jul 2026 20:03:29 -0400 Subject: [PATCH] bitcast to mixin [PR] (#16924) --- test/backend/test_dtype.py | 2 +- test/null/test_tensor_uop_mixin.py | 4 ++++ test/unit/test_disk_tensor.py | 2 +- tinygrad/mixin/dtype.py | 15 ++++++++++++++- tinygrad/schedule/rangeify.py | 16 +++++++++++++++- tinygrad/tensor.py | 27 --------------------------- tinygrad/uop/ops.py | 6 +++--- 7 files changed, 38 insertions(+), 34 deletions(-) diff --git a/test/backend/test_dtype.py b/test/backend/test_dtype.py index cd63186864..f2bd2a1fcc 100644 --- a/test/backend/test_dtype.py +++ b/test/backend/test_dtype.py @@ -281,7 +281,7 @@ class TestBitCast(unittest.TestCase): def test_shape_change_bitcast_exceptions(self): with self.assertRaises(RuntimeError): # should fail because 3 int8 is 3 bytes but float16 is two and 3 isn't a multiple of 2 - Tensor.empty((3,), dtype=dtypes.int8).bitcast(dtypes.float16) + Tensor.empty((3,), dtype=dtypes.int8).bitcast(dtypes.float16).shape def test_bitcast_float_to_int32(self): a = Tensor([1.,2,3]) diff --git a/test/null/test_tensor_uop_mixin.py b/test/null/test_tensor_uop_mixin.py index 717d968a08..ebcc9bbf3a 100644 --- a/test/null/test_tensor_uop_mixin.py +++ b/test/null/test_tensor_uop_mixin.py @@ -222,6 +222,10 @@ class TestTensorUOpBitcast(unittest.TestCase): t = _t(4) self.assertIs(t.bitcast("uint32").uop, t.uop.bitcast("uint32")) self.assertIs(t.uop.bitcast("uint32").dtype, dtypes.uint32) + def test_bitcast_same_and_diff_size(self): + _check(self, _t(4).float(), lambda x: x.bitcast(dtypes.uint32)) # same size + _check(self, _t(4).cast(dtypes.uint8), lambda x: x.bitcast(dtypes.uint16)) # widen: uint8[4] -> uint16[2] + _check(self, _t(4).cast(dtypes.uint16), lambda x: x.bitcast(dtypes.uint8)) # narrow: uint16[4] -> uint8[8] class TestTensorUOpRand(unittest.TestCase): def test_random_bits(self): diff --git a/test/unit/test_disk_tensor.py b/test/unit/test_disk_tensor.py index 182ee5a58c..2d25decc0a 100644 --- a/test/unit/test_disk_tensor.py +++ b/test/unit/test_disk_tensor.py @@ -88,7 +88,7 @@ class TestRawDiskBuffer(unittest.TestCase): # Those two should be moved to test_dtype.py:test_shape_change_bitcast after bitcast works on non-disk with self.assertRaises(RuntimeError): # should fail because 3 int8 is 3 bytes but float16 is two and 3 isn't a multiple of 2 - Tensor.empty((3,), dtype=dtypes.int8, device=f"DISK:{tmp}").bitcast(dtypes.float16) + Tensor.empty((3,), dtype=dtypes.int8, device=f"DISK:{tmp}").bitcast(dtypes.float16).shape pathlib.Path(tmp).unlink() diff --git a/tinygrad/mixin/dtype.py b/tinygrad/mixin/dtype.py index 708e760cf4..518d9994e5 100644 --- a/tinygrad/mixin/dtype.py +++ b/tinygrad/mixin/dtype.py @@ -31,7 +31,20 @@ class DTypeMixin: """ return self if self.dtype == (dt:=to_dtype(dtype)) else self._wrap_uop(self._uop.cast(dt)) - def bitcast(self, dtype:DTypeLike) -> Self: raise NotImplementedError + def bitcast(self, dtype:DTypeLike) -> Self: + """ + Bitcasts `self` to the given `dtype`. If the itemsize differs, the last axis is rescaled. + + ```python exec="true" source="above" session="tensor" result="python" + t = Tensor([-1, 2, 3], dtype=dtypes.int32) + print(t.dtype, t.numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + t = t.bitcast(dtypes.uint32) + print(t.dtype, t.numpy()) + ``` + """ + return self if self.dtype == (dt:=to_dtype(dtype)) else self._wrap_uop(self._uop.bitcast(dt)) def element_size(self) -> int: """ diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 12b929d52e..8237c465b8 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -1,7 +1,7 @@ from dataclasses import dataclass, field, replace from typing import cast import itertools -from tinygrad.dtype import dtypes, AddrSpace, Invalid +from tinygrad.dtype import dtypes, AddrSpace, Invalid, to_dtype from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, identity_element from tinygrad.uop.symbolic import symbolic @@ -112,6 +112,18 @@ def resolve_function(c:UOp, allow_param_mismatch=True) -> UOp|None: if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}") return c.src[0].substitute(dict_map, walk=True) +# shape-changing bitcast +def expand_bitcast(bc:UOp) -> UOp|None: + x = bc.src[0] + if (ns:=bc.dtype.itemsize) == (os:=x.dtype.itemsize) or (isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS"))): return None + new_uint, tmp = to_dtype(f"uint{8*ns}"), x.bitcast(to_dtype(f"uint{8*os}")) + if ns > os: + tmp = tmp.reshape(x.shape[:-1] + (x.shape[-1]//(rate := ns//os), rate)) + parts = [tmp.shrink((None,)*(len(tmp.shape)-1) + ((i, i+1),)).cast(new_uint)<<8*i*os for i in range(rate)] + return parts[0].usum(*parts[1:]).squeeze(-1).bitcast(bc.dtype) + parts = [tmp>>8*i*ns for i in range(os//ns)] + return parts[0].stack(*parts[1:], dim=-1).flatten(-2).cast(new_uint).bitcast(bc.dtype) + earliest_rewrites = mop_cleanup+PatternMatcher([ # resolve FUNCTION calls (inline the body) (UPat(Ops.FUNCTION, name="c"), resolve_function), @@ -166,6 +178,8 @@ earliest_rewrites = mop_cleanup+PatternMatcher([ (UPat(Ops.STORE, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(name="src"))), lambda target, src: target.store(src.bitcast(target.dtype))), + (UPat(Ops.BITCAST, name="bc"), expand_bitcast), + # ** size 0 ** # reduce of size 0 is the identity element diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index d9f984abca..819a7a81e0 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -611,33 +611,6 @@ class Tensor(RandMixin): fn = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(frame_pos.src[0], *[UOp.const(dtypes.int, s) for s in shape]), arg="encdec") return Tensor(out.uop.after(fn.call(*[s.uop for s in srcs], frame_pos))) - # ***** cast ops ***** - - def bitcast(self, dtype:DTypeLike) -> Tensor: - """ - Bitcasts `self` to the given `dtype` of the same itemsize. - - ```python exec="true" source="above" session="tensor" result="python" - t = Tensor([-1, 2, 3], dtype=dtypes.int32) - print(t.dtype, t.numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - t = t.bitcast(dtypes.uint32) - print(t.dtype, t.numpy()) - ``` - """ - dt = to_dtype(dtype) - if (ns:=dt.itemsize) != (os:=self.dtype.itemsize) and (self.shape[-1]*os) % ns != 0: raise RuntimeError("unsupported size in bitcast") - if (not isinstance(self.device, str) or not self.device.startswith("DISK")) and ns != os: - new_uint, old_uint = to_dtype(f"uint{8*ns}"), to_dtype(f"uint{8*os}") - tmp = self.bitcast(old_uint) - if ns > os: - tmp = tmp.reshape(self.shape[:-1] + (self.shape[-1]//(rate := ns//os), rate)) - nones = (None,) * (tmp.ndim - 1) - return Tensor.usum(*[tmp.shrink(nones + ((i, i+1),)).cast(new_uint)<<8*i*os for i in range(rate)]).squeeze(-1).bitcast(dtype) - return Tensor.stack(*(tmp>>8*i*ns for i in range(os//ns)), dim=-1).flatten(-2).cast(new_uint).bitcast(dtype) - return self._apply_uop(UOp.bitcast, dtype=dt) if self.dtype != dt else self - P = ParamSpec("P") T = TypeVar("T") diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index af237c520a..acd8c5e579 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -289,12 +289,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass): Ops.COPY | Ops.ALLREDUCE | Ops.STORE | Ops.END: return self.src[0]._shape - # TODO: disallow shape changing bitcast case Ops.BITCAST: ps = self.src[0]._shape if ps is None: return None - 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 (output_sz:=self.dtype.itemsize) != (input_sz:=self.src[0].dtype.itemsize) and len(ps) > 0: + if isinstance(ps[-1], int) and (ps[-1]*input_sz) % output_sz: raise RuntimeError("unsupported size in bitcast") + return ps[:-1]+(ssimplify((ps[-1]*input_sz) // output_sz),) return ps # MULTI marker has no shape