diff --git a/tinygrad/codegen/decomp/dtype.py b/tinygrad/codegen/decomp/dtype.py index 75d0becfd5..84caaddac1 100644 --- a/tinygrad/codegen/decomp/dtype.py +++ b/tinygrad/codegen/decomp/dtype.py @@ -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) diff --git a/tinygrad/codegen/decomp/op.py b/tinygrad/codegen/decomp/op.py index e34cbf4f86..8848ccf8ff 100644 --- a/tinygrad/codegen/decomp/op.py +++ b/tinygrad/codegen/decomp/op.py @@ -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) diff --git a/tinygrad/codegen/decomp/transcendental.py b/tinygrad/codegen/decomp/transcendental.py index 9a15729192..10ee275c1c 100644 --- a/tinygrad/codegen/decomp/transcendental.py +++ b/tinygrad/codegen/decomp/transcendental.py @@ -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, diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 72f04255be..f840ae0098 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -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")> 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) diff --git a/tinygrad/nn/onnx.py b/tinygrad/nn/onnx.py index 6e53f88c4c..2ba85adede 100644 --- a/tinygrad/nn/onnx.py +++ b/tinygrad/nn/onnx.py @@ -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):