diff --git a/tinygrad/codegen/decomp/dtype.py b/tinygrad/codegen/decomp/dtype.py index df9a24e47b..d682aad8f5 100644 --- a/tinygrad/codegen/decomp/dtype.py +++ b/tinygrad/codegen/decomp/dtype.py @@ -158,7 +158,7 @@ pm_long_decomp = PatternMatcher([ (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx: x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag[0]).replace(dtype=l2i_dt[x.dtype], tag=None),), tag=None) if x.tag is not None else None), (UPat(Ops.CONST, tag={(w, dt) for w in (0, 1) for dt in l2i_dt.values()}, name='x'), lambda x: - UOp.const(truncate[x.tag[1]]((x.arg >> 32) if x.tag[0] == 1 else (x.arg & 0xFFFFFFFF)), x.tag[1])) + UOp.const(truncate[x.tag[1]]((x.val >> 32) if x.tag[0] == 1 else (x.val & 0xFFFFFFFF)), x.tag[1])) ]) # float decomposition patterns - ctx is (fr, to) tuple diff --git a/tinygrad/codegen/decomp/op.py b/tinygrad/codegen/decomp/op.py index 6efe4147d1..6b0526b498 100644 --- a/tinygrad/codegen/decomp/op.py +++ b/tinygrad/codegen/decomp/op.py @@ -77,7 +77,7 @@ def get_simplifying_rewrite_patterns(ops:tuple[Ops, ...]) -> PatternMatcher: # these are rewrites that make things simpler pat: list[tuple[UPat, Callable]] = [(UPat.var("a")//UPat.var("b"), floordiv_to_idiv)] # FLOORMOD by 2**y -> x & (2**y-1) (correct floor mod for any sign in two's complement); fires before floormod_to_mod - if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None)) + if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.val-1) if c.val in powers_of_two else None)) pat.append((UPat.var("a")%UPat.var("b"), floormod_to_mod)) # no real hardware supports THREEFRY, but NullRenderer does if Ops.THREEFRY not in ops: pat.append((UPat(Ops.THREEFRY, dtype=dtypes.uint64, src=(UPat.var("x"), UPat.var("key"))), threefry2x32)) @@ -91,19 +91,19 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa if Ops.OR in ops: pat += [(UPat.var("x", dtypes.bool).logical_not()&UPat.var("y", dtypes.bool).logical_not(), lambda x,y: (x | y).logical_not())] # rewrite MUL/CDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y) - if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)] + if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.val, 0)) else None)] if Ops.SHR in ops: # uint CDIV by 2**v -> x >> v (FLOORDIV is lowered to CDIV by the rule above before reaching here) pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.uints), UPat.cvar("c"))), - lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)] + lambda x,c: x >> v if (v:=powers_of_two.get(c.val, 0)) else None)] # signed CDIV (trunc) by 2**v -> (x + (x<0 ? c-1 : 0)) >> v pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("c"))), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(c-1, 0)) >> v - if (v:=powers_of_two.get(c.arg, 0)) else None)] + if (v:=powers_of_two.get(c.val, 0)) else None)] if not disable_fast_idiv: # fast_idiv handles non-pow2: only fire on non-negative inputs (signed magic-mul is unreliable for x<0) pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))), - lambda ctx, x, d: fast_idiv(ctx, x, d.arg) if x.vmin >= 0 or x.dtype in dtypes.uints else None)] + lambda ctx, x, d: fast_idiv(ctx, x, d.val) if x.vmin >= 0 or x.dtype in dtypes.uints else None)] # rewrite raw CMOD -> x - d*CDIV(x,d) so fast_idiv can pick up the CDIV. only on non-negative inputs; # avoids disturbing floormod_to_mod's general-path output (which uses a trunc Ops.CMOD as an implementation detail) pat += [(UPat(Ops.CMOD, src=(UPat.var("x", dtypes.ints), UPat.var("d"))), @@ -119,13 +119,13 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa (UPat.var("x", dtypes.sints)*-1 < UPat.var("y", dtypes.sints)*UPat.cvar("c"), lambda x,y,c: y*(-c) x==c + lambda x,c1,c2: x.eq(c1+1) if c1.val+1==c2.val-1 else None), # (c-1) x==c ] if Ops.CMPEQ in ops: pat += [(UPat.var('x').ne(UPat.var('y')).logical_not(), lambda x,y: x.alu(Ops.CMPEQ, y))] if Ops.MULACC in ops: pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))] # also fuse (x << n) + c → MULACC(x, 2^n, c) since MUL→SHL may run first - if Ops.SHL in ops: pat += [(UPat.var('x').alu(Ops.SHL, UPat.cvar('n'))+UPat.var('c'), lambda x,n,c: x.alu(Ops.MULACC, x.const_like(1< a/b if Ops.FDIV in ops: pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))] diff --git a/tinygrad/codegen/decomp/transcendental.py b/tinygrad/codegen/decomp/transcendental.py index 23b1723503..b2771acb37 100644 --- a/tinygrad/codegen/decomp/transcendental.py +++ b/tinygrad/codegen/decomp/transcendental.py @@ -16,8 +16,8 @@ def exponent_bias(d:DType) -> int: return (1 << (dtypes.finfo(d)[0] - 1)) - (0 i def exponent_mask(d:DType) -> int: return (1 << dtypes.finfo(d)[0]) - 1 # **** utils **** -def shr(x:UOp|int, y:UOp|int) -> UOp: return x // (2**(y.simplify().arg) if isinstance(y, UOp) else 2**y) -def shl(x:UOp|int, y:UOp|int) -> UOp: return x * (2**(y.simplify().arg) if isinstance(y, UOp) else 2**y) +def shr(x:UOp|int, y:UOp|int) -> UOp: return x // (2**(y.simplify().val) if isinstance(y, UOp) else 2**y) +def shl(x:UOp|int, y:UOp|int) -> UOp: return x * (2**(y.simplify().val) if isinstance(y, UOp) else 2**y) def rintk(d:UOp) -> UOp: """round d:float to int away from 0""" diff --git a/tinygrad/codegen/late/coalesce.py b/tinygrad/codegen/late/coalesce.py index 436032e35f..47ec328017 100644 --- a/tinygrad/codegen/late/coalesce.py +++ b/tinygrad/codegen/late/coalesce.py @@ -84,7 +84,7 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None: h, w, cidx = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].index(1).simplify().backward_slice)) buf = buf.replace(src=(shape_to_shape_arg((h, w, 4)),)) shapes[buf.arg.slot] = (h, w) - if valid.op is not Ops.CONST or valid.arg is not True: + if valid.op is not Ops.CONST or valid.val is not True: return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid), dtype=dtypes.float) else: return buf.index(cidx.src[1], cidx.src[0], dtype=dtypes.float) @@ -111,10 +111,10 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp: if buf.addrspace == AddrSpace.REG: continue idx, valid = idx_u.get_idx(), idx_u.get_valid() root_src: UOp|str - if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].arg - elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].arg - elif idx.op is Ops.CONST and idx.arg is Invalid: root_src, arg = "INVALID", 0 - elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.arg + if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: root_src, arg = idx.src[0], idx.src[1].val + elif idx.op is Ops.ADD and idx.src[0].op is Ops.CONST: root_src, arg = idx.src[1], idx.src[0].val + elif idx.op is Ops.CONST and idx.val is Invalid: root_src, arg = "INVALID", 0 + elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.val else: root_src, arg = idx, 0 memory[(u.op, buf, root_src, valid)].setdefault(arg, []).append(u) diff --git a/tinygrad/codegen/opt/heuristic.py b/tinygrad/codegen/opt/heuristic.py index 5b687ae04d..f78dcb2b78 100644 --- a/tinygrad/codegen/opt/heuristic.py +++ b/tinygrad/codegen/opt/heuristic.py @@ -126,8 +126,8 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler: if rng in idx.backward_slice: num_strides += 1 for c in idx.split_uop(Ops.ADD): if c is rng: sum_strides += 1 - if c.op is Ops.MUL and c.src[0] is rng and c.src[1].op is Ops.CONST: sum_strides += c.src[1].arg - if c.op is Ops.MUL and c.src[1] is rng and c.src[0].op is Ops.CONST: sum_strides += c.src[0].arg + if c.op is Ops.MUL and c.src[0] is rng and c.src[1].op is Ops.CONST: sum_strides += c.src[1].val + if c.op is Ops.MUL and c.src[1] is rng and c.src[0].op is Ops.CONST: sum_strides += c.src[0].val xb_choices.append((num_strides, sum_strides, axis, upcast_amount)) if xb_choices: xb_choices = sorted(xb_choices) diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 647a6c1f11..608a656d25 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -62,7 +62,7 @@ pm_simplify_ranges = PatternMatcher([ def mark_range_mod(ctx:dict[UOp, UOp|None], r:UOp, c:UOp) -> None: # ranges that aren't looped over can't be split if r not in ctx and r.arg[-1] not in {AxisType.WARP, AxisType.DEVICE} \ - and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c + and r.src[0].op is Ops.CONST and r.src[0].divides(c.val) is not None: ctx[r] = c def do_substitute(ctx:dict, x: UOp, sub_fxn:Callable[[UOp, UOp], UOp]) -> UOp|None: ret = x.substitute({k:sub_fxn(k,v) for k,v in ctx.items() if v is not None}) diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index f11c9688d4..a5f8155162 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -198,7 +198,7 @@ def exec_validate(ctx:ExecContext, call:UOp, ast:UOp) -> float|None: def exec_encdec(ctx:ExecContext, call:UOp, ast:UOp) -> float|None: bufs = [cast(Buffer, b.buffer).ensure_allocated() for b in resolve_params(call, ctx.input_uops)] - shape, pos_var = tuple(s.arg for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr + shape, pos_var = tuple(s.val for s in ast.src if s.op is Ops.CONST), ast.variables()[0].expr with track_stats(ctx, call, bufs[0].device, bufs, ctx.var_vals): bufs[0].allocator._encode_decode(bufs[0]._buf, bufs[1]._buf, bufs[2]._buf, [x._buf for x in bufs[3:]], shape, ctx.var_vals[pos_var]) return None diff --git a/tinygrad/mixin/elementwise.py b/tinygrad/mixin/elementwise.py index f73e0a0bea..c74baa8b39 100644 --- a/tinygrad/mixin/elementwise.py +++ b/tinygrad/mixin/elementwise.py @@ -24,7 +24,7 @@ class ElementwiseMixin(CreationMixin): out_dtype = least_upper_dtype(x.dtype, y.dtype) # keep weak CONST weak, might lift weakint -> weakfloat def promote(t): - if t.dtype in dtypes.weaks and t._uop.base.op is Ops.CONST: return t._wrap_uop(t._uop.const_like(t._uop.base.arg, weak_dtype(out_dtype))) + if t.dtype in dtypes.weaks and t._uop.base.op is Ops.CONST: return t._wrap_uop(t._uop.const_like(t._uop.base.val, weak_dtype(out_dtype))) return t.cast(out_dtype) return promote(x), promote(y) diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index dfb18ff382..6dbf090bb4 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -34,18 +34,18 @@ base_rewrite = PatternMatcher([ # const (UPat(Ops.CONST, arg=math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x, ctx.infinity)})"), (UPat(Ops.CONST, arg=-math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x, f'-{ctx.infinity}')})"), - (UPat(Ops.CONST, dtype=dtypes.floats, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx.nan)})" if math.isnan(x.arg) else None), - (UPat(Ops.CONST, dtype=dtypes.float, name="x"), lambda ctx,x: f"{x.arg}f"), - (UPat(Ops.CONST, dtype=dtypes.int64, name="x"), lambda ctx,x: f"{x.arg}l"), - (UPat(Ops.CONST, dtype=dtypes.uint64, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}ul"), - (UPat(Ops.CONST, dtype=dtypes.uint32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}u"), - (UPat(Ops.CONST, dtype=dtypes.bool, name="x"), lambda ctx,x: "1" if x.arg else "0"), + (UPat(Ops.CONST, dtype=dtypes.floats, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx.nan)})" if math.isnan(x.val) else None), + (UPat(Ops.CONST, dtype=dtypes.float, name="x"), lambda ctx,x: f"{x.val}f"), + (UPat(Ops.CONST, dtype=dtypes.int64, name="x"), lambda ctx,x: f"{x.val}l"), + (UPat(Ops.CONST, dtype=dtypes.uint64, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}ul"), + (UPat(Ops.CONST, dtype=dtypes.uint32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}u"), + (UPat(Ops.CONST, dtype=dtypes.bool, name="x"), lambda ctx,x: "1" if x.val else "0"), # consts are rendered to larger type and casted - (UPat(Ops.CONST, (*dtypes.fp8s, dtypes.bfloat16, dtypes.half), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.arg}f')})"), - (UPat(Ops.CONST, (dtypes.uint8, dtypes.uint16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.arg}u')})"), - (UPat(Ops.CONST, (dtypes.int8, dtypes.int16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, str(x.arg))})"), + (UPat(Ops.CONST, (*dtypes.fp8s, dtypes.bfloat16, dtypes.half), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.val}f')})"), + (UPat(Ops.CONST, (dtypes.uint8, dtypes.uint16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, f'{x.val}u')})"), + (UPat(Ops.CONST, (dtypes.int8, dtypes.int16), name="x"), lambda ctx,x: f"({ctx.render_cast(x, str(x.val))})"), # default const render - (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.arg)), + (UPat(Ops.CONST, name="x"), lambda ctx,x: str(x.val)), # SHRINK/INDEX (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var('idx')), name="x"), lambda ctx,**kwargs: ctx.render_index(**kwargs)), @@ -164,7 +164,7 @@ class CStyleLanguage(Renderer): if buf.addrspace == AddrSpace.ALU: # this is lane access in C if idx.op is not Ops.CONST: return f"({self[buf]})[{self[idx]}]" - return self[buf]+(f"[{idx.arg}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.arg]}") + return self[buf]+(f"[{idx.val}]" if buf.max_numel() > self.gep_arr_threshold else f".{'xyzwabcd'[idx.val]}") return f"({self[buf]}+{strip_parens(self[idx]) if idx.arg == Ops.ADD else self[idx]})" def render_buffer(self, x:UOp): @@ -494,10 +494,10 @@ class HIPRenderer(CStyleLanguage): (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}," f" {fp8_index(x.src[0].dtype)}, {fp8_index(x.src[0].dtype)}, 0, 0, 0, 0)" if x.arg[0][2] == 128 else None), (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}, 0, 0, 0)"), - (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.nan}, {fp8_index(x.dtype)})" if math.isnan(x.arg) else None), + (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.nan}, {fp8_index(x.dtype)})" if math.isnan(x.val) else None), (UPat(Ops.CONST, dtypes.fp8s, arg=math.inf, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.infinity}, {fp8_index(x.dtype)})"), (UPat(Ops.CONST, dtypes.fp8s, arg=-math.inf, name="x"), lambda ctx,x: f"f32_to_fp8(-{ctx.infinity}, {fp8_index(x.dtype)})"), - (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({x.arg}f, {fp8_index(x.dtype)})"), + (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({x.val}f, {fp8_index(x.dtype)})"), (UPat(Ops.CAST, dtypes.fp8s, (UPat(dtype=dtypes.float),), name="x",), lambda ctx,x: f"f32_to_fp8({ctx[x.src[0]]}, {fp8_index(x.dtype)})"), (UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",), @@ -522,7 +522,7 @@ class HIPRenderer(CStyleLanguage): lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2])) if x.src[0].max_numel() == 8 and x.src[0].dtype in dtypes.fp8_ocp else None), # bfloat16 constant casting - (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(x.arg, dtypes.float))), + (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(x.val, dtypes.float))), ]) def asm(self, prg:UOp, lin:UOp) -> bytes: @@ -538,7 +538,7 @@ class HIPRenderer(CStyleLanguage): prefix, ockl = [], [] type_map = { dtypes.bfloat16: "bf16", dtypes.float: "f32", dtypes.half: "f16", dtypes.fp8e4m3: "_fp8_fp8", dtypes.fp8e5m2: "_bf8_bf8" } used_dtypes = uops_to_dtypes(uops) - if any(u.op is Ops.CONST and not math.isfinite(u.arg) for u in uops): + if any(u.op is Ops.CONST and not math.isfinite(u.val) for u in uops): prefix += ["#define INFINITY (__builtin_inff())", "#define NAN (__builtin_nanf(\"\"))"] if any(u.op is Ops.SPECIAL for u in uops): prefix.append("typedef long unsigned int size_t;") diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py index 2436d8bcf4..4ea41f643c 100644 --- a/tinygrad/renderer/isa/x86.py +++ b/tinygrad/renderer/isa/x86.py @@ -289,8 +289,8 @@ def fold_address(x:UOp) -> tuple[UOp, UOp, UOp, UOp]: # buffers are indexed by element, everything else (the stack pointer) by byte scale = base.dtype.itemsize if base.op in {Ops.PARAM, Ops.BUFFER, Ops.AFTER} else 1 sz = imm(dtypes.uint8, base.dtype.itemsize) - if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: return (base, _cast(idx.src[0]), _disp(idx.src[1].arg * scale), sz) - if idx.op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.arg * scale), sz) + if idx.op is Ops.ADD and idx.src[1].op is Ops.CONST: return (base, _cast(idx.src[0]), _disp(idx.src[1].val * scale), sz) + if idx.op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.val * scale), sz) return (base, _cast(idx), _disp(0), sz) def abi(ctx:IselContext, x:UOp) -> UOp|None: @@ -353,7 +353,7 @@ isel_matcher = PatternMatcher([ # cast of void is a noop (UPat.var("y").cast(name="x"), lambda y,x: y if y.dtype == dtypes.void else None), # range is lowered to acc, cmp, jmp after regalloc - (UPat(Ops.RANGE, src=(UPat.cvar("c"),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(c.dtype, c.arg),) + x.src[1:])), + (UPat(Ops.RANGE, src=(UPat.cvar("c"),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(c.dtype, c.val),) + x.src[1:])), (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(tag=(ctx.vreg(WGPR),)) if not isinstance(x.tag, tuple) else None), # really all a backedge END is is an IF with a tag referencing the RANGE start label (UPat(Ops.END, src=(UPat(), UPat(), UPat(GroupOp.Comparison, name="cond")), name="x"), @@ -367,10 +367,10 @@ isel_matcher = PatternMatcher([ # function abi constraints (UPat((Ops.PARAM, Ops.SPECIAL), name="x"), abi), # constants that can't be immediates, move them to registers - (UPat.cvar("x", dtypes.int64s), lambda x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, x.arg),)) if not x.tag else None), - (UPat.cvar("x", dtypes.ints+(dtypes.bool,)), lambda x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, x.arg),)) if not x.tag else None), + (UPat.cvar("x", dtypes.int64s), lambda x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, x.val),)) if not x.tag else None), + (UPat.cvar("x", dtypes.ints+(dtypes.bool,)), lambda x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, x.val),)) if not x.tag else None), (UPat.cvar("x", dtypes.floats), lambda x: - UOp.const(struct.unpack((dt:=to_int(x.dtype)).fmt, struct.pack(x.dtype.fmt, x.arg))[0], dt).bitcast(x.dtype) if not x.tag else None), + UOp.const(struct.unpack((dt:=to_int(x.dtype)).fmt, struct.pack(x.dtype.fmt, x.val))[0], dt).bitcast(x.dtype) if not x.tag else None), # conditional moves that use masks NOTE: these currently assume a mask producing cmp exists (UPat.var("m").where(UPat.var("a", dtypes.int8s+dtypes.int16s+dtypes.int32s+(dtypes.int64,)), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.VPBLENDVB, src=(b, a, m.replace(dtype=m.src[0].dtype))) if a.max_numel() > 1 else None), @@ -453,9 +453,9 @@ isel_matcher = PatternMatcher([ # scalar int binary ((UPat(dtype=dtypes.ints).alu(Ops.CDIV, UPat())).named("x"), idiv), # scalar int binary with immediate - (UPat.var("a", dtypes.ints) << UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHLi, src=(a, imm(dtypes.uint8, c.arg)))), - (UPat.var("a", dtypes.uints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.arg)))), - (UPat.var("a", dtypes.sints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.arg)))), + (UPat.var("a", dtypes.ints) << UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHLi, src=(a, imm(dtypes.uint8, c.val)))), + (UPat.var("a", dtypes.uints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.val)))), + (UPat.var("a", dtypes.sints) >> UPat.cvar("c"), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.val)))), (UPat.var("a", dtypes.ints) + UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ADDi, src=(a, i)) if (i:=to_imm(c)) is not None else None), (UPat.var("a", dtypes.ints) * UPat.cvar("c"), lambda a,c: a.ins(X86Ops.IMULi, src=(a, i)) if (i:=to_imm(c)) is not None else None), (UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.cvar("c"), lambda a,c: a.ins(X86Ops.ANDi, src=(a, i)) if (i:=to_imm(c)) is not None else None), @@ -664,10 +664,10 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) -> # DISP byte if mod == 0b01 or mod == 0b10: assert disp_uop is not None - inst += struct.pack(unwrap(disp_uop.dtype.fmt), disp_uop.arg) + inst += struct.pack(unwrap(disp_uop.dtype.fmt), disp_uop.val) # IMM byte if imm_uop is not None: - if imm_uop.op is Ops.CONST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.arg) + if imm_uop.op is Ops.CONST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.val) elif isinstance(greg(imm_uop), Register): inst += bytes([(greg(imm_uop).index & 0b1111) << 4 | 0b0000]) return inst @@ -840,10 +840,10 @@ class X86Renderer(ISARenderer): def _format_op(x:UOp) -> str: return f" {(o[7:-1] if (o:=str(x.arg))[-1] in ('i', 'm') else o[7:]).lower():7s}" def _format_operands(x:UOp) -> str: def _format(src:tuple[UOp, ...]) -> list[str]: - return [str(s.arg) if s.op is Ops.CONST else reg_strs[o].get(s.dtype.itemsize, o) if \ + return [str(s.val) if s.op is Ops.CONST else reg_strs[o].get(s.dtype.itemsize, o) if \ (o:=str(greg(s))) in reg_strs else o for s in src if greg(s) is not None] def _mem_adress(base:UOp, idx:UOp, disp:UOp, sz:UOp) -> list[str]: - return [f"[{greg(base)}" + (f" + {greg(idx)}*{sz.arg}" if greg(idx) else "") + (f" + {disp.arg}" if disp.arg else "") + "]"] + return [f"[{greg(base)}" + (f" + {greg(idx)}*{sz.val}" if greg(idx) else "") + (f" + {disp.val}" if disp.val else "") + "]"] if len(x.src) > 4 and x.arg in X86GroupOp.WriteMem: ret = _mem_adress(*x.src[:4]) + _format(x.src[4:]) elif len(x.src) > 3 and x.arg in X86GroupOp.Rm1st: ret = _format((x,)) + _mem_adress(*x.src[:4]) + _format(x.src[4:]) diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 6ae37820d8..56b3b38027 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -67,7 +67,7 @@ base_rewrite = PatternMatcher([ f" {ctx[x]} = getelementptr inbounds {ldt(x.dtype)}, {ldt(x.dtype, ptr=True)} {ctx[x.src[0]]}, {ldt(x.src[1].dtype)} {ctx[x.src[1]]}"), # register index (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.cvar("idx")), name="x"), lambda ctx,buf,idx,x: - f" {ctx[x]} = extractelement {ldt(buf.dtype, buf.max_numel())} {ctx[buf]}, i32 {idx.arg}" if buf.addrspace == AddrSpace.ALU else None), + f" {ctx[x]} = extractelement {ldt(buf.dtype, buf.max_numel())} {ctx[buf]}, i32 {idx.val}" if buf.addrspace == AddrSpace.ALU else None), # load/store (UPat(Ops.LOAD, src=(UPat.var("idx"), UPat.var("alt"), UPat.var("mask")), name="x"), @@ -170,7 +170,7 @@ class LLVMRenderer(Renderer): kernel.append(f" {r[u]} = addrspacecast [{size} x {ldt(u.dtype)}] addrspace(3)* @{r[u][1:]} to [{size} x {ldt(u.dtype)}]*") else: kernel.append(f" {r[u]} = alloca [{size} x {ldt(u.dtype)}], align 16") - elif u.op is Ops.CONST: r[u] = lconst(u.arg, u.dtype) + elif u.op is Ops.CONST: r[u] = lconst(u.val, u.dtype) elif u.op is Ops.CAST and ldt(u.dtype) == ldt(u.src[0].dtype): r[u] = r[u.src[0]] # cast from signed to unsigned of the same size is a noop, or pointer cast else: diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index fbac922a10..a531007ff7 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -122,7 +122,7 @@ class NIRRenderer(Renderer): extra_matcher = PatternMatcher([ # handle negative unsigned CONST - (UPat.cvar("x", dtypes.uints), lambda x: UOp.const(x.dtype.max+x.arg+1, x.dtype) if x.arg < 0 else None), + (UPat.cvar("x", dtypes.uints), lambda x: UOp.const(x.dtype.max+x.val+1, x.dtype) if x.val < 0 else None), # from ptx (UPat.var('x', dtype=dtypes.bool) uint8 @@ -144,7 +144,7 @@ class NIRRenderer(Renderer): ]) def_rewrite = PatternMatcher([ - (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.arg, x.dtype)), + (UPat(Ops.CONST, name="x"), lambda ctx,x: nimm(ctx.b, x.val, x.dtype)), (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx.param(ctx.b, x, x.dtype.itemsize if x.addrspace is AddrSpace.ALU else 8)), (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: nchannel(ctx.b, {'g':ngid, 'l':nlid, 'i': nid}[x.arg[0]](ctx.b), int(x.arg[-1]))), (UPat(Ops.STORE, src=(UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"),UPat.var("off")), allow_any_len=True), UPat.var("val"))), diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index 36a46a8d1d..1b6b77859a 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -79,8 +79,8 @@ def modifier(a: DType, b: DType): return '.rzi' if dtypes.is_int(a) and dtypes.i (a.itemsize < b.itemsize or dtypes.is_int(b) or b == dtypes.bool) else '' string_rewrite = PatternMatcher([ - (UPat.cvar("x", dtypes.bool), lambda ctx, x: f"setp.ne.s16 {ctx.r[x]}, {render_val(x.arg, x.dtype)}, 0;"), - (UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.arg, x.dtype)};"), + (UPat.cvar("x", dtypes.bool), lambda ctx, x: f"setp.ne.s16 {ctx.r[x]}, {render_val(x.val, x.dtype)}, 0;"), + (UPat.cvar("x"), lambda ctx, x: f"mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {render_val(x.val, x.dtype)};"), (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"mov.u32 %{x.arg}, %{'ctaid' if x.arg[0] == 'g' else 'tid'}.{chr(120+int(x.arg[-1]))};"), (UPat(Ops.PARAM, name="x"), lambda ctx, x: f"ld.param.{ctx.types[dtypes.ulong] if x.addrspace is AddrSpace.GLOBAL else ctx.mem_types[x.dtype]} {ctx.r[x]}, [data{x.arg.slot}+0];"), @@ -203,7 +203,7 @@ class PTXRenderer(Renderer): # on REG, INDEX/SHRINK pick the register (must be CONST) and LOAD is a noop if u.op is not Ops.LOAD and u.src[1].op is not Ops.CONST: raise RuntimeError(f"PTX does not support dynamic register indexing: {u}") - r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][u.src[1].arg] + r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][u.src[1].val] continue if u.op is Ops.SPECIAL: r[u] = "%" + u.arg elif u.op is Ops.LOAD: diff --git a/tinygrad/renderer/wgsl.py b/tinygrad/renderer/wgsl.py index 71c969629a..e9e2bfdcc3 100644 --- a/tinygrad/renderer/wgsl.py +++ b/tinygrad/renderer/wgsl.py @@ -68,10 +68,10 @@ class WGSLRenderer(CStyleLanguage): string_rewrite = PatternMatcher([ (UPat(Ops.NEG, dtypes.uints, src=(UPat.var('x'))), lambda ctx,x: f"(0-{ctx[x]})"), - (UPat.cvar("x", dtype=dtypes.bool), lambda x: "true" if x.arg else "false"), + (UPat.cvar("x", dtype=dtypes.bool), lambda x: "true" if x.val else "false"), (UPat(Ops.CONST, dtype=(dtypes.uchar, dtypes.ushort, dtypes.uint32), name="x"), - lambda x: f"bitcast({x.arg})" if x.arg < 0 else f"{x.arg&0xFFFFFFFF}u"), - (UPat(Ops.CONST, dtype=dtypes.int32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.arg)}"), + lambda x: f"bitcast({x.val})" if x.val < 0 else f"{x.val&0xFFFFFFFF}u"), + (UPat(Ops.CONST, dtype=dtypes.int32, name="x"), lambda ctx,x: f"{truncate[x.dtype](x.val)}"), (UPat(Ops.BUFFER, name="x"), lambda ctx,x: f"var{'' if x.addrspace == AddrSpace.LOCAL else ''} {ctx[x]}: array<{ctx.buf_map(x)},{_packed_size(x)}>;"), (UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)), diff --git a/tinygrad/runtime/ops_dsp.py b/tinygrad/runtime/ops_dsp.py index 9419ea962e..e8ffa2ef83 100644 --- a/tinygrad/runtime/ops_dsp.py +++ b/tinygrad/runtime/ops_dsp.py @@ -13,7 +13,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat # NOTE: this just increases readability of the generated code dsp_string = PatternMatcher([ - (UPat(Ops.CONST, (dtypes.int8, dtypes.uint8), name="x"), lambda ctx,x: str(x.arg)), + (UPat(Ops.CONST, (dtypes.int8, dtypes.uint8), name="x"), lambda ctx,x: str(x.val)), ]) class DSPRenderer(ClangRenderer): diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index 5b2f14483b..5112df0040 100644 --- a/tinygrad/runtime/ops_python.py +++ b/tinygrad/runtime/ops_python.py @@ -102,7 +102,7 @@ class PythonProgram(Program['PythonDevice']): elif u.op is Ops.SPECIAL: if u.arg[0] == 'g': values[u] = [idxs[2-int(u.arg[-1])]] * warp_size elif u.arg[0] == 'l': values[u] = [x[2-int(u.arg[-1])] for x in warp] - elif u.op is Ops.CONST: values[u] = [u.arg] * warp_size + elif u.op is Ops.CONST: values[u] = [u.val] * warp_size elif u.op in {Ops.INDEX, Ops.SHRINK}: ret:list = [] if u.src[0].addrspace == AddrSpace.ALU: diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index 3952f30e9e..498703bc3a 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -55,7 +55,7 @@ def make_cmdbuf(lin, devs, buf:UOp|None=None, dep:UOp|None=None): blob, patches = bytearray(), [] for s in (s for ins in lin.src for s in ins.src): if s.op is not Ops.CONST: patches.append((len(blob), s)) - blob.extend(struct.pack(f'<{s.dtype.fmt}', s.arg if s.op is Ops.CONST else 0x0)) + blob.extend(struct.pack(f'<{s.dtype.fmt}', s.val if s.op is Ops.CONST else 0x0)) cmdbuf = buf if buf is not None else UOp.placeholder((len(blob) // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("cmdbuf") writable = cmdbuf.after(dep) if dep is not None else cmdbuf return cmdbuf.after(make_binary_patch(writable, bytes(blob)), *((make_patches(writable, patches),) if patches else ())) diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index 440b1565f9..17110647f4 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -209,7 +209,7 @@ def index_multi(root:UOp, multi:UOp): continue # strided ownership: idx ≡ rng (mod shard_sz), intra-shard position is (idx - rng) // shard_sz diff = (idxs[ax] - rng).simplify() - if (mod:=(diff % shard_sz).simplify()).op is Ops.CONST and mod.arg == 0: + if (mod:=(diff % shard_sz).simplify()).op is Ops.CONST and mod.val == 0: local = (diff // shard_sz).simplify() if local.vmin >= 0 and local.vmax < shard_sz: idxs[ax] = local diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index ff15b77b24..0e6b44b376 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -283,7 +283,7 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp): # if it makes it here, the bufferize is removed # this is the ranges replaced # NOTE: if buf src is a const, we don't replace it. if idx is Invalid (dead load), don't replace it either - replaced = {k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST and not (v.op is Ops.CONST and v.arg is Invalid)} + replaced = {k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST and not (v.op is Ops.CONST and v.val is Invalid)} return src.substitute(replaced, extra_pm=pm_gate_substitute) def remove_noop_bufferize(idx,b2): @@ -304,7 +304,7 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([ (UPat(Ops.INDEX, src=(UPat(Ops.STAGE),), allow_any_len=True, name="idx").f(Ops.NOOP).f(Ops.STAGE, allow_any_len=True, name="b2"), remove_noop_bufferize), # no buffers for const (ranges don't matter for const - it's the same value everywhere) - (UPat(Ops.CONST, name='c').f(Ops.STAGE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.arg)), + (UPat(Ops.CONST, name='c').f(Ops.STAGE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.val)), # indexing a const is a const (UPat(Ops.INDEX, src=(UPat(Ops.CONST, name="c"),),), lambda c: c), # indexing an after with all fully invalid stores is invalid diff --git a/tinygrad/uop/divandmod.py b/tinygrad/uop/divandmod.py index 8ad865baca..bb082bc3d3 100644 --- a/tinygrad/uop/divandmod.py +++ b/tinygrad/uop/divandmod.py @@ -12,7 +12,7 @@ def fold_divmod_general(d: UOp) -> UOp|None: # x//y is constant if (xdiv:=x//y).vmin == xdiv.vmax: return x - xdiv.vmin*y if d.op is Ops.FLOORMOD else xdiv.const_like(xdiv.vmin) # PARAM // c is irreducible - if x.op is Ops.PARAM and y.op is Ops.CONST and x.arg.multiple_of % y.arg == 0: return d.const_like(0) if d.op is Ops.FLOORMOD else None + if x.op is Ops.PARAM and y.op is Ops.CONST and x.arg.multiple_of % y.val == 0: return d.const_like(0) if d.op is Ops.FLOORMOD else None # split uops for the rest of the processing x_peeled, const = x.pop_const() @@ -20,7 +20,7 @@ def fold_divmod_general(d: UOp) -> UOp|None: # ** Constant Denominator Rules ** # these rules strictly require y to be a scalar constant > 0 - if y.op is Ops.CONST and (c := y.arg) > 0: + if y.op is Ops.CONST and (c := y.val) > 0: # nested_div: (x%(k*c))//c -> (x//c)%k (requires k>0); the mod case is handled by remove_nested_mod below if d.op is Ops.FLOORDIV and x.op is Ops.FLOORMOD and (k := x.src[1].divides(c)) is not None and k > 0: return x.src[0] // y % k @@ -76,7 +76,7 @@ def fold_divmod_general(d: UOp) -> UOp|None: # divide_by_gcd: x//y -> (x//gcd)//(y//gcd) gcd = UOp.gcd(*all_uops, y).simplify() - if not (gcd.op is Ops.CONST and gcd.arg==1): + if not (gcd.op is Ops.CONST and gcd.val==1): ret = unwrap(x.divide_exact(gcd)).alu(d.op, unwrap(y.divide_exact(gcd))) return ret*gcd if d.op is Ops.FLOORMOD else ret @@ -85,8 +85,8 @@ def fold_divmod_general(d: UOp) -> UOp|None: quo, rem = [], [] for u in all_uops: if (q:=u.divide_exact(y)) is not None: quo.append(q) - elif y.op is Ops.CONST and (c:=u.const_factor())%y.arg!=c: - rem.append(u.divides(c)*(c%y.arg)) + elif y.op is Ops.CONST and (c:=u.const_factor())%y.val!=c: + rem.append(u.divides(c)*(c%y.val)) quo.append(u.divides(c)*(c//y.arg) if d.op is Ops.FLOORDIV else u.const_like(0)) else: rem.append(u) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 1ad814049a..17899ce79f 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -259,6 +259,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if (self.op, self.dtype, self.src, self.arg, self.tag) == new_args: return self return UOp(*new_args) def rtag(self, tag=True): return self.replace(tag=tag) + @property + def val(self): + assert self.op is Ops.CONST, f"val is only for CONST, got {self.op}" + return self.arg @recursive_property def key(self) -> bytes: return hashlib.sha256(str((self.op, self.dtype, self.arg)).encode() + b"".join([s.key for s in self.src])).digest() @@ -521,7 +525,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): from tinygrad.uop.symbolic import symbolic with Context(TRACK_MATCH_STATS=0 if not tracked else TRACK_MATCH_STATS.value): return graph_rewrite(self, symbolic, name="simplify") - def ssimplify(self) -> UOp|ConstType: return ret.arg if (ret:=self.simplify()).op is Ops.CONST else ret + def ssimplify(self) -> UOp|ConstType: return ret.val if (ret:=self.simplify()).op is Ops.CONST else ret def _eval(self, dtype, expected_type:Type[T]) -> T: assert self.dtype in dtype, f"eval with wrong dtype {self}" vmin, vmax = (simple_self:=self.simplify())._min_max @@ -765,9 +769,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # cached property here makes external_uop_gc fail, why? @property def as_shape(self) -> tuple[sint, ...]: - if self.op is Ops.CONST: return (self.arg,) + if self.op is Ops.CONST: return (self.val,) if self.op is not Ops.STACK: return (ssimplify(self),) - return tuple(s.arg if s.op is Ops.CONST else ssimplify(s) for s in self.src) + return tuple(s.val if s.op is Ops.CONST else ssimplify(s) for s in self.src) @functools.cached_property def marg(self): @@ -895,7 +899,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): idx = self.flatten().index(UOp.range(self.numel(), 0)) out = graph_rewrite(idx, pm_mops+symbolic+pm_contiguous_view_offset, ctx=self, name="contiguous_view_offset") - return out.arg if out.op is Ops.CONST and isinstance(out.arg, int) else None + return out.val if out.op is Ops.CONST and isinstance(out.val, int) else None def has_buffer_identity(self, after_ok=False): """Check if this UOp has a concrete buffer identity in the graph (RESHAPE/UNSHARD -> BUFFER chain).""" @@ -984,7 +988,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return UOp(Ops.BIND, src=(self, uval)) def unbind(self) -> tuple[Variable, int]: assert self.op is Ops.BIND and self.src[0].op is Ops.PARAM and self.src[1].op is Ops.CONST, f"can't unbind {self}" - return self.src[0], self.src[1].arg + return self.src[0], self.src[1].val def unbind_all(self) -> tuple[UOp, dict[Variable, int]]: ret:dict[Variable, int] = {} return graph_rewrite(self, pm_unbind, ctx=ret), ret @@ -997,15 +1001,15 @@ class UOp(RandMixin, metaclass=UOpMetaClass): def const_factor(self) -> int: """largest known int that divides self""" # TODO: for negatives it's not the largest - if self.op is Ops.CONST: return self.arg + if self.op is Ops.CONST: return self.val if self.op is Ops.STACK: return math.gcd(*[x.const_factor() for x in self.src]) if self.op is Ops.ADD: return math.gcd(self.src[0].const_factor(), self.src[1].const_factor()) - if self.op is Ops.MUL: return self.src[0].arg if self.src[0].op is Ops.CONST else self.src[1].arg if self.src[1].op is Ops.CONST else 1 + if self.op is Ops.MUL: return self.src[0].val if self.src[0].op is Ops.CONST else self.src[1].val if self.src[1].op is Ops.CONST else 1 if self.op is Ops.PARAM and self.arg.multiple_of is not None: return self.arg.multiple_of return 1 def divides(self, v:int) -> UOp|None: if v==1: return self - if self.op is Ops.CONST: return self.const_like(self.arg//v) if self.arg%v == 0 else None + if self.op is Ops.CONST: return self.const_like(self.val//v) if self.val%v == 0 else None if self.op is Ops.STACK: srcs = tuple(s.divides(v) for s in self.src) return None if any(s is None for s in srcs) else UOp(Ops.STACK, src=cast(tuple[UOp, ...], srcs)) @@ -1016,7 +1020,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op is Ops.PARAM and self.arg.multiple_of is not None: return self // v if self.arg.multiple_of%v == 0 else None return None # generic None if we aren't sure def pop_const(self, op=Ops.ADD) -> tuple[UOp, PyConst]: # NOTE: assume Invalid ALU is resolved - return (self.src[0], self.src[1].arg) if self.op is op and self.src[1].op is Ops.CONST else (self, identity_element(op, self.dtype)) + return (self.src[0], self.src[1].val) if self.op is op and self.src[1].op is Ops.CONST else (self, identity_element(op, self.dtype)) @staticmethod def gcd(*uops: UOp) -> UOp: terms, factors = zip(*[(u.divides(f:=u.const_factor()),f) for u in uops]) @@ -1024,7 +1028,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return math.prod([*count.elements(), terms[0].const_like(math.gcd(*factors))]) # put the const at the top def divide_exact(self, v:UOp) -> UOp|None: if self is v: return self.const_like(1) - if v.op is Ops.CONST: return self.divides(v.arg) + if v.op is Ops.CONST: return self.divides(v.val) if self.op is Ops.ADD: return None if (s0:=self.src[0].divide_exact(v)) is None or (s1:=self.src[1].divide_exact(v)) is None else s0+s1 if self.op is Ops.MUL: (fac, const), (div_fac, div_const) = self.pop_const(Ops.MUL), v.pop_const(Ops.MUL) @@ -1082,7 +1086,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op in (Ops.RANGE, Ops.SPECIAL) and self.dtype is not dtypes.void: return 0, (self.src[0]-1).vmax if self.op is Ops.BIND: return self.src[0]._min_max # ignore the bound value if self.op is Ops.STACK: return min(x.vmin for x in self.src), max(x.vmax for x in self.src) - if self.op is Ops.CONST and self.arg is not Invalid: return self.arg, self.arg + if self.op is Ops.CONST and self.val is not Invalid: return self.val, self.val if self.op is Ops.INDEX: return self.src[0]._min_max if self.op is Ops.CAST: # a cast to unsigned keeps exact bounds when the source fits diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 0c39015d09..3e21f6eb0b 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -13,10 +13,10 @@ from tinygrad.codegen.decomp.transcendental import xpow # ******** phase 1 of symbolic used to live in ops, it's the most generic folding rules ******** def simplify_pow(x:UOp, c:UOp) -> UOp|None: - if c.arg < 0: return x.reciprocal().pow(-c) - if c.arg == 0: return x.const_like(1) - if int(c.arg-0.5)+0.5 == c.arg: return x.pow(c.const_like(c.arg-0.5)) * x.sqrt() - if int(c.arg) == c.arg: return (y := x.pow(c.const_like(c.arg//2))) * y * (x if c.arg%2 == 1 else 1) + if c.val < 0: return x.reciprocal().pow(-c) + if c.val == 0: return x.const_like(1) + if int(c.val-0.5)+0.5 == c.val: return x.pow(c.const_like(c.val-0.5)) * x.sqrt() + if int(c.val) == c.val: return (y := x.pow(c.const_like(c.val//2))) * y * (x if c.val%2 == 1 else 1) return None def fold_bitcast(root:UOp, c:UOp) -> UOp|None: @@ -110,8 +110,8 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ ((UPat.var() % UPat.var("y")).named("base") % UPat.var("y"), lambda base,y: base), # (x%y)%y = -> x%y (rewritten with base for speed) # variations of (x%c)+(x//c)*c = x (UPat(Ops.ADD, dtype=dtypes.weakint, name="x"), fold_add_divmod_recombine), - (UPat.var("x", dtype=dtypes.bool) & UPat.cvar("c"), lambda x,c: x if c.arg else c), - (UPat.var("x", dtype=dtypes.bool) | UPat.cvar("c"), lambda x,c: c if c.arg else x), + (UPat.var("x", dtype=dtypes.bool) & UPat.cvar("c"), lambda x,c: x if c.val else c), + (UPat.var("x", dtype=dtypes.bool) | UPat.cvar("c"), lambda x,c: c if c.val else x), (UPat.var("x", dtype=dtypes.bool) != UPat.const(False, dtypes.bool), lambda x: x), # x != False -> x (UPat(GroupOp.Idempotent, src=(UPat.var("x"), UPat.var("x"))), lambda x: x), (UPat.var("x", dtype=dtypes.bool).logical_not().logical_not(), lambda x: x), @@ -153,10 +153,10 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ # if x is nan or inf it should render the nan value. # NOTE: this can be wrong for loaded NaN (UPat.var("x") * 0, lambda x: x.const_like(float("nan") if x.op is Ops.CONST - and isinstance(x.arg, float) and (math.isnan(x.arg) or math.isinf(x.arg)) else 0)), + and isinstance(x.val, float) and (math.isnan(x.val) or math.isinf(x.val)) else 0)), # *** cast/bitcast *** # TODO: delete this once CONST has no dtype - (UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.arg)), + (UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.val)), (UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None), (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast), # b.cast(a).cast(b) -> b if a preserves all values in b @@ -167,7 +167,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ # ** pow ** (UPat.var("x").alu(Ops.POW, UPat.cvar("c")), simplify_pow), # positive const ** x - (UPat.cvar("c").alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.arg == 1 else (x*math.log2(c.arg)).exp2() if c.arg > 0 else None), + (UPat.cvar("c").alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.val == 1 else (x*math.log2(c.val)).exp2() if c.val > 0 else None), # rules for threefry ((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)), (((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y), @@ -178,7 +178,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ # ** simple where folding ** # a conditional with the same results either way is a noop, also fold const conditionals (UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val), - (UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.arg else c1), + (UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.val else c1), # a.where(b.where(c, d), d) -> (a & b).where(c, d) (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)), # a.where(c, b.where(c, d)) -> (a | b).where(c, d) @@ -274,7 +274,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ ((UPat.var("x") * UPat.cvar("c1")) * UPat.var("y"), lambda x,c1,y: (x*y)*c1 if y.op is not Ops.CONST else None), # *** rules from symbolic *** # generic lt folding - (UPat.var("x", dtypes.weakint) 0. NOTE: not x < 1 means x > 0 ((UPat.var("x", dtypes.weakint)<1).ne(True), lambda x: (newx<1).ne(True) if (newx:=canonicalize_simplex(x)) is not None else None), @@ -453,7 +453,7 @@ sym = symbolic+pm_simplify_valid+PatternMatcher([ (UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")*UPat.var("y")), lambda x,y,d: y*(1-d)), (UPat.var("x") * ((1+UPat.var("x")).reciprocal().named("d")+UPat.var("y")), lambda x,y,d: (1-d)+x*y), # move const multiply after REDUCE (NOTE: the mul chain can do this, but only if it's a same dtype reduce) - ((UPat.var("x")*UPat.cvar("c")).reduce(arg=Ops.ADD, name="r", allow_any_len=True), lambda x,c,r: r.replace(src=(x,)+r.src[1:])*c.arg), + ((UPat.var("x")*UPat.cvar("c")).reduce(arg=Ops.ADD, name="r", allow_any_len=True), lambda x,c,r: r.replace(src=(x,)+r.src[1:])*c.val), # reduce mul chain, move muls after the reduce (UPat(Ops.MUL).reduce(name="r", allow_any_len=True), reduce_mul_chain), # ** combine terms (opinionated) ** diff --git a/tinygrad/uop/upat.py b/tinygrad/uop/upat.py index 2eb5cf18ca..9bbcab452f 100644 --- a/tinygrad/uop/upat.py +++ b/tinygrad/uop/upat.py @@ -118,7 +118,7 @@ pm_renderer = PatternMatcher([ lambda r,x: r.replace(src=(UOp(Ops.CUSTOMI, arg="(" + ' and '.join(y.arg for y in x.src) + ")"),)+r.src[1:])), (UPat(Ops.CUSTOM, src=UPat(Ops.CUSTOMI), name="x"), lambda x: UOp(Ops.CUSTOMI, arg=x.arg.format(*[y.arg for y in x.src]))), - (UPat(Ops.INDEX, src=(UPat(Ops.CUSTOMI, name="x"), UPat(Ops.CONST, name="c")), name="g"), lambda x,c,g: x.replace(arg=x.arg+f".src[{c.arg}]")) + (UPat(Ops.INDEX, src=(UPat(Ops.CUSTOMI, name="x"), UPat(Ops.CONST, name="c")), name="g"), lambda x,c,g: x.replace(arg=x.arg+f".src[{c.val}]")) ], compiled=False) def _final_render(x:UOp, has_ctx:bool, depth=1) -> list[str]: