geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

delete explicit casts [pr] (#17255)

Browse Source 
handled by broadcast
This commit is contained in:
chenyu
2026-07-28 13:57:13 -04:00
committed by GitHub
parent a9ad08064e
commit c9e11544df
6 changed files with 13 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/decomp/dtype.py
+4 -4
View File
@@ -30,7 +30,7 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
sign = lo if x.dtype is dtypes.bool else x
return lo, (sign < sign.const_like(0)).where(lo.const_like(-1), lo.const_like(0))
case Ops.CAST if dt in (dtypes.long, dtypes.ulong):
return (lo:=uops[0].cast(l2i_dt[dt])), (uops[0] / 2**32).cast(l2i_dt[dt]) - ((uops[0] < 0) & lo.ne(0)).cast(l2i_dt[dt])
return (lo:=uops[0].cast(l2i_dt[dt])), (uops[0] / 2**32).cast(l2i_dt[dt]) - ((uops[0] < 0) & lo.ne(0))
case Ops.CAST if dt in dtypes.floats:
small = (a1.eq(0) & (a0 >= 0)) | (a1.eq(-1) & (a0 < 0))
return small.where(a0.cast(dt), ((a1.cast(dtypes.float32) * (2**32)) + a0.bitcast(dtypes.uint).cast(dtypes.float32)).cast(dt))
@@ -45,8 +45,8 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
lo, hi = ((a0u >> n) | ((a1u << 1) << (31 - n))).bitcast(dt), a1 >> (b0 & 31)
fill = a1 >> 31 if dt == dtypes.int else zero # vacated high word: sign bits when signed, else 0
return (b0 >= 32).where(hi, lo), (b0 >= 32).where(fill, hi)
case Ops.ADD: return (low:=a0+b0), a1 + b1 + (low.bitcast(dtypes.uint) < a0.bitcast(dtypes.uint)).cast(dt)
case Ops.SUB: return a0 - b0, a1 - b1 - (a0.bitcast(dtypes.uint) < b0.bitcast(dtypes.uint)).cast(dt)
case Ops.ADD: return (low:=a0+b0), a1 + b1 + (low.bitcast(dtypes.uint) < a0.bitcast(dtypes.uint))
case Ops.SUB: return a0 - b0, a1 - b1 - (a0.bitcast(dtypes.uint) < b0.bitcast(dtypes.uint))
case Ops.MUL:
(a00, a01), (b00, b01) = unpack32(a0), unpack32(b0)
mid = l2i(Ops.ADD, dt, shl(a00*b01, 16).bitcast(dt), shr(a00*b01, 16).bitcast(dt), shl(a01*b00, 16).bitcast(dt), shr(a01*b00, 16).bitcast(dt))
@@ -85,7 +85,7 @@ def split_l2i(op: Ops, dt: DType, *uops:UOp):
# ***** floats *****
f2f_dt = { f:getattr(dtypes, f"uint{f.bitsize}") for f in dtypes.floats }
def rne(v: UOp, s) -> UOp: return shr(v, s) + ((shr(v, s - 1) & 1) & ((v & ((1 << (s - 1)) - 1)).ne(0).cast(v.dtype) | (shr(v, s) & 1)))
def rne(v: UOp, s) -> UOp: return shr(v, s) + ((shr(v, s - 1) & 1) & ((v & ((1 << (s - 1)) - 1)).ne(0) | (shr(v, s) & 1)))
def f2f(v, fr:DType, to:DType, sat=True):
fs, fb, (fe, fm), ts, tb, (te, tm) = fr.bitsize, exponent_bias(fr), dtypes.finfo(fr), to.bitsize, exponent_bias(to), dtypes.finfo(to)
tinygrad/codegen/decomp/op.py
+1 -1
View File
@@ -63,7 +63,7 @@ def threefry2x32(x: UOp, key: UOp):
def floordiv_to_idiv(a:UOp, b:UOp) -> UOp:
if (a.vmin >= 0 and b.vmin > 0) or (a.vmax <= 0 and b.vmax < 0): return a.alu(Ops.CDIV, b)
return a.alu(Ops.CDIV, b) - (a.alu(Ops.CMOD, b).ne(0) & (a<0).ne(b<0)).cast(a.dtype)
return a.alu(Ops.CDIV, b) - (a.alu(Ops.CMOD, b).ne(0) & (a<0).ne(b<0))
def floormod_to_mod(a:UOp, b:UOp) -> UOp:
if (a.vmin >= 0 and b.vmin > 0) or (a.vmax <= 0 and b.vmax < 0): return a.alu(Ops.CMOD, b)
tinygrad/codegen/decomp/transcendental.py
+1 -1
View File
@@ -200,7 +200,7 @@ def xexp2(d:UOp) -> UOp:
x = _lazy_map_numbers(d, d.const_like(0.0), d.const_like(0.0), d.const_like(0.0), d)
q = rintk(x)
# s = d - round(d)
s = x - q.cast(x.dtype)
s = x - q
# a polynomial approximation with 13 non-zero terms in the range of [−(log 2)/2,(log 2)/2].
if d.dtype == dtypes.float64:
u = polyN(s, [0.4434359082926529454e-9, 0.7073164598085707425e-8, 0.1017819260921760451e-6, 0.1321543872511327615e-5, 0.1525273353517584730e-4,
tinygrad/codegen/simplify.py
+5 -5
View File
@@ -84,9 +84,9 @@ def reduce_unparented(red:UOp) -> UOp|None:
if len(reduce_unparented) == 0: return None
ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0]
if red.arg[0] is Ops.ADD:
for r in reduce_unparented: ret = ret * r.src[0].cast(ret.dtype)
for r in reduce_unparented: ret = ret * r.src[0]
if red.arg[0] is Ops.MUL:
for r in reduce_unparented: ret = ret ** r.src[0].cast(ret.dtype)
for r in reduce_unparented: ret = ret ** r.src[0]
return ret
pm_reduce_unparented = PatternMatcher([
@@ -96,7 +96,7 @@ pm_reduce_unparented = PatternMatcher([
pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
# lift x+y out of reduce on lt
((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c-y)) if no_range(y) and no_range(c) else None),
# lift x*y out of reduce
((UPat.var("x")*UPat.var("y")) < UPat.var("c"),
lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and dtypes.is_int(y.dtype) and y.vmin > 0 else None),
@@ -107,13 +107,13 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
((UPat.var("r")<UPat.var("lower")).logical_not()&(UPat(Ops.RANGE, name="r")<UPat.var("upper"))).where(UPat.var("val"), 0),
).reduce(UPat.var("r"), arg=Ops.ADD), lambda r,val,lower=None,upper=None:
((upper.minimum(r.src[0]) if upper is not None else r.src[0]) -
(lower.maximum(0) if lower is not None else r.const_like(0))).maximum(0).minimum(r.src[0]).cast(val.dtype) * val if no_range(val) else None),
(lower.maximum(0) if lower is not None else r.const_like(0))).maximum(0).minimum(r.src[0]) * val if no_range(val) else None),
# REDUCE on ADD
((UPat.var("x")+UPat.var("y")).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
lambda x,y,r: x.reduce(*r.src[1:], arg=Ops.ADD) + y.reduce(*r.src[1:],arg=Ops.ADD)),
# AND on WHERE
((UPat(Ops.PARAM, name="x") & UPat.var("y")).where(UPat.var("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)),
lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x),
# MUL casted bool
((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast()), lambda x,gate: gate.where(x, 0)),
])+symbolic
tinygrad/mixin/rand.py
+1 -1
View File
@@ -21,7 +21,7 @@ class RandMixin(OpMixin):
for i in range(0, num, dtypes.uint32.max):
chunk_num = min(num - i, dtypes.uint32.max)
c_low = low + (i & 0xffffffff)
c_high = high + (i >> 32) + (c_low < low).cast(dtypes.uint32)
c_high = high + (i >> 32) + (c_low < low)
new_key = cls._threefry_random_bits(key, c_low, c_high)
counts0 = cls.arange(ceildiv(chunk_num, 2), dtype=dtypes.uint32)
counts1 = counts0 + ceildiv(chunk_num, 2)
tinygrad/nn/onnx.py
+1 -1
View File
@@ -1093,7 +1093,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
def RMSNormalization(X:Tensor, scale:Tensor, axis:int=-1, epsilon:float=1e-5, stash_type:int=1):
assert stash_type == 1, "only float32 is supported"
norm = X.cast(dtypes.float).square().mean(axis=tuple(range(axis + X.ndim if axis < 0 else axis, X.ndim)), keepdim=True).add(epsilon).rsqrt()
return X.cast(X.dtype) * norm * scale
return X * norm * scale
def RotaryEmbedding(X:Tensor, cos_cache:Tensor, sin_cache:Tensor, position_ids:Tensor|None=None, interleaved:int=0, num_heads:int|None=None,
rotary_embedding_dim:int=0):