sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

bitcast to mixin [PR] (#16924)

Browse Source
This commit is contained in:
chenyu
2026-07-07 20:03:29 -04:00
committed by GitHub
parent 9f388d42b7
commit c43a3fdebb
7 changed files with 38 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/backend/test_dtype.py
+1 -1
View File
@@ -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])
test/null/test_tensor_uop_mixin.py
+4
View File
@@ -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):
test/unit/test_disk_tensor.py
+1 -1
View File
@@ -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()
tinygrad/mixin/dtype.py
+14 -1
View File
@@ -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:
"""
tinygrad/schedule/rangeify.py
+15 -1
View File
@@ -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
tinygrad/tensor.py
-27
View File
@@ -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")
tinygrad/uop/ops.py
+3 -3
View File
@@ -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