diff --git a/tinygrad/callify.py b/tinygrad/callify.py index 5ff272a4ec..b7bb8d4972 100644 --- a/tinygrad/callify.py +++ b/tinygrad/callify.py @@ -57,7 +57,7 @@ def _make_buffer_view(src:UOp) -> UOp|None: if (offset := src.contiguous_view_offset()) is None: return None buf = src.base if buf.op is Ops.SLICE: - byte_offset = buf.src[1].arg * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize + byte_offset = buf.src[1].val * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize buf = buf.src[0] if byte_offset % buf.dtype.itemsize != 0: return None offset = byte_offset // buf.dtype.itemsize diff --git a/tinygrad/codegen/simplify.py b/tinygrad/codegen/simplify.py index 608a656d25..4aadeb6d6f 100644 --- a/tinygrad/codegen/simplify.py +++ b/tinygrad/codegen/simplify.py @@ -47,7 +47,7 @@ def mark_gated(ctx, idx): guards = {r:c for v in cond.split_uop(Ops.AND) if v.op is Ops.CMPLT and (r:=v.src[0]).op is Ops.RANGE and (c:=v.src[1]).op is Ops.CONST} else: x, guards = idx, {} # ensure that we choose max(c_i) for all i where r < c_i - ctx |= {r:c for r,c in guards.items() if (r not in ctx or ctx[r].arg < c.arg)} + ctx |= {r:c for r,c in guards.items() if (r not in ctx or ctx[r].val < c.val)} # but if a range is ever ungated, we cannot shrink it ctx |= {r:r.src[0] for r in x.ranges if r not in guards} diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index a5f8155162..91b9e331a6 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -28,7 +28,7 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N ast, arg_uops = call.src[0], get_call_arg_uops(call) if ast.op is Ops.PROGRAM: return ast.arg.name if ast.op is Ops.SLICE: - offset = ast.src[1].arg * arg_uops[1].dtype.itemsize + offset = ast.src[1].val * arg_uops[1].dtype.itemsize return colored(f"view {_uop_sz_to_str(arg_uops[0]):>10} @ {offset:<10d}", "yellow") if ast.op is Ops.COPY: return colored(f"copy {_uop_sz_to_str(arg_uops[0]):>10}, {_dev_str(bufs[0]):>7s} <- {_dev_str(bufs[1]):7s}", "yellow") if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return colored(f"enc/dec {_uop_sz_to_str(arg_uops[0])}", "yellow") @@ -156,7 +156,7 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d def exec_view(ctx:ExecContext, call:UOp, ast:UOp) -> float|None: resolved = resolve_params(call, ctx.input_uops) bufs = [cast(Buffer, b.buffer) for b in resolved] - bv = bufs[1].view(resolved[0].max_numel(), ast.dtype, ast.src[1].arg*bufs[1].dtype.itemsize) + bv = bufs[1].view(resolved[0].max_numel(), ast.dtype, ast.src[1].val*bufs[1].dtype.itemsize) with track_stats(ctx, call, bv.device, [bv, bufs[1]], ctx.var_vals): buffers[resolved[0]] = bv return None diff --git a/tinygrad/function.py b/tinygrad/function.py index e63adaad70..8b540e3381 100644 --- a/tinygrad/function.py +++ b/tinygrad/function.py @@ -22,7 +22,7 @@ def invalid_outputs(uret:UOp) -> set[UOp]: # invalids() returns fresh write-only scratch: a clone storing CONST(Invalid) # don't capture it as an input; only skip fresh buffers, not realized ones return {u.src[0].buf_uop for u in uret.backward_slice_with_self - if u.op is Ops.STORE and u.src[1].base.op is Ops.CONST and u.src[1].base.arg is Invalid + if u.op is Ops.STORE and u.src[1].base.op is Ops.CONST and u.src[1].base.val is Invalid and not u.src[0].buf_uop.is_realized} ReturnType = TypeVar('ReturnType') diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 6dbf090bb4..e25ab0c605 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -320,7 +320,7 @@ class OpenCLRenderer(CStyleLanguage): (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_{ctx.render_dtype(x.dtype)}(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"), # bfloat16 constants need to be rendered as their bit pattern since bf16 is stored as ushort (UPat(Ops.CONST, dtypes.bfloat16, name="x"), - lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.arg)))[0] >> 16)}u"), + lambda ctx,x: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(x.val)))[0] >> 16)}u"), # load/store image (OpenCL) (UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), lambda ctx,buf,idx_y,idx_x: f"IMAGE<{ctx[buf]}, {ctx[idx_y]}, {ctx[idx_x]}>"), (UPat(Ops.LOAD, dtype=dtypes.float, src=(UPat.var('buf').index(UPat.var('idx_y'), UPat.var('idx_x')), UPat.var("var"), UPat.var("gate"))), diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py index 4ea41f643c..1346a02678 100644 --- a/tinygrad/renderer/isa/x86.py +++ b/tinygrad/renderer/isa/x86.py @@ -221,15 +221,15 @@ reg_strs = {"rax": {4:"eax", 2:"ax", 1:"al"}, "rcx": {4:"ecx", 2:"cx", 1:"cl"}, # ***** X86 instruction selection ***** def base(x:UOp, i:int) -> UOp: return s.src[0] if (s:=x.src[i]).op is Ops.INDEX else s -def lane(x:UOp, i:int) -> int: return s.src[1].arg if (s:=x.src[i]).op is Ops.INDEX else 0 +def lane(x:UOp, i:int) -> int: return s.src[1].val if (s:=x.src[i]).op is Ops.INDEX else 0 def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt] def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, dt, arg=X86Ops.DEFINE, tag=None if reg is None else (reg,)) def imm(dt:DType, v:int) -> UOp: return UOp.const(truncate[dt](v), dt).rtag() def to_imm(c:UOp) -> UOp|None: if c.op is not Ops.CONST: return None - if c.dtype is dtypes.int64: return imm(dtypes.int32, c.arg) if not c.overflows(dtypes.int32) else None - if c.dtype is dtypes.uint64: return imm(dtypes.uint32, c.arg) if not c.overflows(dtypes.uint32) else None - if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, c.arg) + if c.dtype is dtypes.int64: return imm(dtypes.int32, c.val) if not c.overflows(dtypes.int32) else None + if c.dtype is dtypes.uint64: return imm(dtypes.uint32, c.val) if not c.overflows(dtypes.uint32) else None + if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, c.val) return None def cmp(x:UOp) -> UOp: if x.src[0].dtype is dtypes.float32: return x.ins(X86Ops.VUCOMISS, dtype=dtypes.void) @@ -421,15 +421,15 @@ isel_matcher = PatternMatcher([ (UPat(Ops.STACK, dtypes.ints+(dtypes.bool,), name="x"), vpins), # INDEX on a vector register value extracts a single element (UPat.var("y", dtypes.int8s+(dtypes.bool,)).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPEXTRB, src=(y, imm(dtypes.uint8, c.arg))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPEXTRB, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None), (UPat.var("y", dtypes.int16s).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, c.arg))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None), (UPat.var("y", dtypes.int32s).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPEXTRD, src=(y, imm(dtypes.uint8, c.arg))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPEXTRD, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None), (UPat.var("y", dtypes.int64s).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPEXTRQ, src=(y, imm(dtypes.uint8, c.arg))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPEXTRQ, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None), (UPat.var("y", dtypes.floats).index(UPat.cvar("c"), name="x"), - lambda y,c,x: x.ins(X86Ops.VPSRLDQ, src=(y, imm(dtypes.uint8, c.arg * x.dtype.itemsize))) if _is_vec_xmm(y) else None), + lambda y,c,x: x.ins(X86Ops.VPSRLDQ, src=(y, imm(dtypes.uint8, c.val * x.dtype.itemsize))) if _is_vec_xmm(y) else None), # packed bitwise ((UPat() & UPat()).named("x"), lambda x: x.ins(X86Ops.VPAND) if x.max_numel() > 1 else None), ((UPat() | UPat()).named("x"), lambda x: x.ins(X86Ops.VPOR) if x.max_numel() > 1 else None), @@ -614,7 +614,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) -> rm = cast(Register, greg(rm_uop)).index idx = cast(Register, greg(idx_uop)).index if idx_uop is not None and greg(idx_uop) is not None else 4 # for a memory operand the rm size is the element size from the address, otherwise it's the size of the value in the register - rm_sz = sz_uop.arg if sz_uop is not None else rm_uop.dtype.itemsize + rm_sz = sz_uop.val if sz_uop is not None else rm_uop.dtype.itemsize reg_sz = reg_uop.dtype.itemsize if reg_uop is not None else 0 sz = reg_sz or rm_sz @@ -650,7 +650,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) -> assert disp_uop.op is Ops.CONST, "displacement must be a constant" assert disp_uop.dtype in (dtypes.int8, dtypes.int32), "displacement can only be 1 or 4 byte signed int" # rbp/r13 always require a displacement - if disp_uop.arg != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10 + if disp_uop.val != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10 else: mod = 0b00 else: mod = 0b11 # x 0b0 and idx 0b100 means rsp which means no index exists @@ -701,7 +701,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) -> encodings = { # moves X86Ops.MOVABS: lambda x: - bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | greg(x).index >> 3, 0xB8 + (greg(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].arg), + bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | greg(x).index >> 3, 0xB8 + (greg(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].val), X86Ops.MOV: lambda x: encode(x, 0x8B), X86Ops.MOVi: lambda x: encode(x, 0xC7, reg=0), X86Ops.MOVm: lambda x: encode(x, 0x89), X86Ops.LEA: lambda x: encode(x, 0x8D), X86Ops.VMOVSS: lambda x: encode(x, 0x10, pp=2, sel=1), X86Ops.VMOVSSm: lambda x: encode(x, 0x11, pp=2, sel=1), @@ -724,8 +724,8 @@ encodings = { X86Ops.VCVTPS2PD: lambda x: encode(x, 0x5A, pp=0, sel=1), X86Ops.VCVTPD2PS: lambda x: encode(x, 0x5A, pp=1, sel=1), X86Ops.VCVTTPS2DQ: lambda x: encode(x, 0x5B, pp=2, sel=1), X86Ops.VCVTTPD2DQ: lambda x: encode(x, 0xE6, pp=1, sel=1), # the int src is the 2nd src (the rm field), if it was folded into a memory operand its width is the element size of the address - X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=(x.src[4].arg if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8), - X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=(x.src[4].arg if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8), + X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=(x.src[4].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8), + X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=(x.src[4].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8), X86Ops.VCVTTSS2SI: lambda x: encode(x, 0x2C, pp=2, sel=1, we=x.dtype.itemsize == 8), X86Ops.VCVTTSD2SI: lambda x: encode(x, 0x2C, pp=3, sel=1, we=x.dtype.itemsize == 8), # int division diff --git a/tinygrad/renderer/nir.py b/tinygrad/renderer/nir.py index a531007ff7..b0d3a38f2c 100644 --- a/tinygrad/renderer/nir.py +++ b/tinygrad/renderer/nir.py @@ -185,7 +185,7 @@ class NIRRenderer(Renderer): def render(self, uops:list[UOp]): self.prerender(uops) - for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]: self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].arg + for u in [u for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]: self.b.shader.contents.info.workgroup_size[int(u.arg[-1])] = u.src[0].val self.r: dict[UOp, Any] = {} self.param_idx = 0 ranges: list[mesa.nir_def|None] = [] @@ -194,7 +194,7 @@ class NIRRenderer(Renderer): if u.op in {Ops.NOOP, Ops.GROUP} or (u.op is Ops.STACK and len(u.src) == 0): pass elif u.op in {Ops.INDEX, Ops.SHRINK}: # INDEX on a register value picks the element, memory INDEX is handled in the LOAD/STORE patterns - if u.src[0].op not in {Ops.PARAM, Ops.BUFFER, Ops.AFTER}: self.r[u] = nchannel(self.b, self.r[u.src[0]], u.src[1].arg) + if u.src[0].op not in {Ops.PARAM, Ops.BUFFER, Ops.AFTER}: self.r[u] = nchannel(self.b, self.r[u.src[0]], u.src[1].val) elif u.op is Ops.AFTER: self.r[u] = self.r[u.src[0]] elif u.op == Ops.SINK: diff --git a/tinygrad/runtime/graph/metal.py b/tinygrad/runtime/graph/metal.py index 9094f36692..409cfb973d 100644 --- a/tinygrad/runtime/graph/metal.py +++ b/tinygrad/runtime/graph/metal.py @@ -113,5 +113,5 @@ class MetalGraph(GraphRunner): @staticmethod def supports_uop(batch_devs, new_call:UOp) -> bool: # Metal ICB replay encodes offsets as uint32; reject if any Metal buffer offset exceeds 32-bit range. - if any(b.op is Ops.SLICE and b.src[1].arg * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False + if any(b.op is Ops.SLICE and b.src[1].val * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False return GraphRunner.supports_uop(batch_devs, new_call) diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index 498703bc3a..0f36f516de 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -334,7 +334,7 @@ pm_replace_params = PatternMatcher([ def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp: base = bv.src[0].after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()) itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize - return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(bv.src[1].arg * itemsize, dtypes.uint64) + return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(bv.src[1].val * itemsize, dtypes.uint64) pm_early_simplify = PatternMatcher([ (UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat(Ops.SLICE, name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice), diff --git a/tinygrad/schedule/__init__.py b/tinygrad/schedule/__init__.py index 0845e2e354..99d228945f 100644 --- a/tinygrad/schedule/__init__.py +++ b/tinygrad/schedule/__init__.py @@ -186,7 +186,7 @@ def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]: if b.op is Ops.BIND: nm = b.src[0].expr if nm not in used_vars: continue - val = b.src[1].arg + val = b.src[1].val if var_vals.get(nm, val) != val: raise RuntimeError(f"bind mismatch on {nm}, {var_vals[nm]} != {val}") var_vals[nm] = val diff --git a/tinygrad/uop/divandmod.py b/tinygrad/uop/divandmod.py index bb082bc3d3..03dbd1e138 100644 --- a/tinygrad/uop/divandmod.py +++ b/tinygrad/uop/divandmod.py @@ -87,7 +87,7 @@ def fold_divmod_general(d: UOp) -> UOp|None: if (q:=u.divide_exact(y)) is not None: quo.append(q) 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)) + quo.append(u.divides(c)*(c//y.val) if d.op is Ops.FLOORDIV else u.const_like(0)) else: rem.append(u) if not quo: return None @@ -101,8 +101,8 @@ div_and_mod_symbolic = PatternMatcher([ ((UPat.var("x")//UPat.cvar("c") + UPat.cvar("a"))//UPat.cvar("d"), lambda x,c,a,d: (x+a*c)//(c*d) if d.vmin>0 else None), # (x+c)//d -> (x+c%d)//d + c//d ; (x+c)%d -> (x+c%d)%d (split the multiple of d out of the const, holds for any d!=0) (UPat((Ops.FLOORDIV, Ops.FLOORMOD), src=(UPat.var("x", dtypes.weakint)+UPat.cvar("c"), UPat.cvar("d")), name="n"), - lambda n,x,c,d: None if d.arg==0 or c.arg%d.arg==c.arg else - (x+c.arg%d.arg)//d + c.arg//d.arg if n.op is Ops.FLOORDIV else (x+c.arg%d.arg)%d), + lambda n,x,c,d: None if d.val==0 or c.val%d.val==c.val else + (x+c.val%d.val)//d + c.val//d.val if n.op is Ops.FLOORDIV else (x+c.val%d.val)%d), # ** 2. Slow Rules ** (UPat((Ops.FLOORDIV, Ops.FLOORMOD), dtypes.weakint, name="d"), lambda d: fold_divmod_general(d)), diff --git a/tinygrad/uop/movement.py b/tinygrad/uop/movement.py index 686f7a7ae7..295efced07 100644 --- a/tinygrad/uop/movement.py +++ b/tinygrad/uop/movement.py @@ -13,10 +13,10 @@ mop_cleanup = PatternMatcher([ (UPat(Ops.PERMUTE, name="x"), lambda x: x.src[0] if list(x.arg) == list(range(len(x.arg))) else None), # STACK on INDEX CONST (UPat(Ops.STACK, src=UPat(Ops.INDEX, src=(UPat.var("src"), UPat(Ops.CONST))), name="stk"), - lambda src,stk: src if stk.shape == src.shape and list(range(len(stk.src))) == [x.src[1].arg for x in stk.src] else None), + lambda src,stk: src if stk.shape == src.shape and list(range(len(stk.src))) == [x.src[1].val for x in stk.src] else None), # const INDEX into STACK is src (UPat(Ops.INDEX, src=(UPat(Ops.STACK, name="a"), UPat.cvar("i")), name="idx", allow_any_len=True), - lambda a,i,idx: a.src[i.arg] if len(idx.src) <= 2 else a.src[i.arg].index(*idx.src[2:])), + lambda a,i,idx: a.src[i.val] if len(idx.src) <= 2 else a.src[i.val].index(*idx.src[2:])), # INDEX on INDEX is INDEX (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, name="idx1", allow_any_len=True),), allow_any_len=True, name="idx2"), lambda idx1,idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:]) if all(x.shape == () for x in idx1.src[1:]+idx2.src[1:]) else None), diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 17899ce79f..027508e2be 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -567,7 +567,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None])) def index(self, *srcs:UOp|int|None, **kwargs): new_srcs: list[UOp] = [UOp.const(x) if isinstance(x, int) else x for x in srcs if x is not None] - if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].arg] + if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].val] return UOp(Ops.INDEX, src=(self,)+tuple(new_srcs), **kwargs) def __getitem__(self, idx): # buffers index into INDEX UOps (scalar lookup); everything else uses the shared mixin view path @@ -931,7 +931,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op is Ops.SLICE: if (cret:=buffers.get(self)) is not None: return cret buf = self.src[0].buffer - offset = self.src[1].arg + offset = self.src[1].val if isinstance(buf, MultiBuffer): mbuf = MultiBuffer.__new__(MultiBuffer) mbuf.bufs = [b.view(self.arg, self.dtype, offset * self.src[0].dtype.itemsize) for b in buf.bufs] @@ -1787,7 +1787,7 @@ def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None: def commit_weak(s:UOp, dt:DType) -> UOp: # a bare weak CONST commits directly (its number must fit), a weak non-const src takes the demand cast - return UOp.const(s.arg, dt) if s.op is Ops.CONST else s.cast(dt) + return UOp.const(s.val, dt) if s.op is Ops.CONST else s.cast(dt) def commit_weak_srcs(u:UOp) -> UOp|None: if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None diff --git a/tinygrad/uop/render.py b/tinygrad/uop/render.py index 97ebe79db4..c8467ae654 100644 --- a/tinygrad/uop/render.py +++ b/tinygrad/uop/render.py @@ -18,7 +18,7 @@ def pretty_print(x:UOp, cache=None, d=0)->str: def print_uops(uops:list[UOp]): uops_index = {u:i for i,u in enumerate(uops)} for i,u in enumerate(uops): - formatted_srcs = [(uops_index[x] if x.op is not Ops.CONST else f"{x.arg}") if x in uops else "--" for x in u.src] + formatted_srcs = [(uops_index[x] if x.op is not Ops.CONST else f"{x.val}") if x in uops else "--" for x in u.src] print(f"{i:4d} {str(u.op):20s}: {multirange_str(u.ranges, color=True, pad=10)} {str(u.dtype):40s} " f"{str(formatted_srcs):32s} {u.arg}") # for debug @@ -36,7 +36,7 @@ renderer = PatternMatcher([ (UPat((Ops.SPECIAL), name="x"), lambda x: x.arg), (UPat(Ops.RANGE, dtypes.void, name="x"), lambda x: f"loop{x.arg[0]}"), (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"), - (UPat(Ops.CONST, name="x"), lambda x: str(x.arg)), + (UPat(Ops.CONST, name="x"), lambda x: str(x.val)), (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({str(x.dtype)[7:]})({ctx[x.src[0]]})"), (UPat(Ops.BIND, name="x"), lambda ctx,x: ctx[x.src[0]]), (UPat(Ops.NEG, name="x"), lambda ctx,x: f"(-{ctx[x.src[0]]})"), @@ -79,9 +79,9 @@ def render_marg(ctx,x:UOp): sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY, Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH} pm_pyrender_extra = PatternMatcher([ - (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.arg}, {x.dtype})"), + (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.val}, {x.dtype})"), (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})"), - (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].arg}, {repr(x.arg)}, dtype={x.dtype})"), + (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].val}, {repr(x.arg)}, dtype={x.dtype})"), (UPat(Ops.BUFFER, src=(UPat(),), name="x"), lambda x: f"UOp.new_buffer({repr(x.arg.device)}, {x.max_numel()}, {x.dtype}, {x.arg.slot})" if isinstance(x.arg, ParamArg) and x.addrspace is AddrSpace.GLOBAL else None), @@ -90,7 +90,7 @@ pm_pyrender_extra = PatternMatcher([ (UPat(Ops.REDUCE, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}._rop({r.arg[0]}, {tuple(range(r.arg[1]))})" if r.arg[1] else None), # NOTE: range has srcs sometimes after control flow (UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c: - "UOp.range("+', '.join([str(c.arg)] + [repr(y) for y in x.arg])+ + "UOp.range("+', '.join([str(c.val)] + [repr(y) for y in x.arg])+ (f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else '')+")"), # TODO: index shouldn't mismatch dtype (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x: diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index b61b971f14..4f07447a60 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -10,7 +10,7 @@ from tinygrad.helpers import DEBUG, Context, SPEC, Metadata, panic, CHECK_OOB, a def validate_index(uidx:UOp, gate:UOp|None=None): if len(uidx.src) != 2: return True # skip for non final index. TODO: check more complex index with shape buf,idx = uidx.src - if idx.op is Ops.CONST and idx.arg is Invalid: return True + if idx.op is Ops.CONST and idx.val is Invalid: return True if gate is None: gate = UOp.const(True) # TODO: check for overflow if not CHECK_OOB or is_image_shape(buf._shape): return True @@ -54,7 +54,7 @@ spec_shared = PatternMatcher([ (UPat(Ops.NOOP), lambda: True), # CONST is everywhere - (UPat(Ops.CONST, src=(), name="x"), lambda x: type(x.arg) is type(x.dtype.const(x.arg))), + (UPat(Ops.CONST, src=(), name="x"), lambda x: type(x.val) is type(x.dtype.const(x.val))), # STACK is everywhere too (UPat(Ops.STACK, dtype=dtypes.void, src=()), lambda: True), diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 3e21f6eb0b..4e9a5e1c66 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -23,11 +23,11 @@ def fold_bitcast(root:UOp, c:UOp) -> UOp|None: if (from_fmt:=c.dtype.fmt) is None or (to_fmt:=root.dtype.fmt) is None: return None if c.dtype.itemsize != root.dtype.itemsize: return None def convert(v:ConstType) -> ConstType: return struct.unpack(to_fmt, struct.pack(from_fmt, v))[0] - return root.const_like(convert(c.arg)) + return root.const_like(convert(c.val)) def const_arg(u:UOp) -> ConstType|tuple[ConstType, ...]|None: - if u.op is Ops.CONST: return u.arg - if u.op is Ops.STACK and all(s.op is Ops.CONST for s in u.src): return tuple(s.arg for s in u.src) + if u.op is Ops.CONST: return u.val + if u.op is Ops.STACK and all(s.op is Ops.CONST for s in u.src): return tuple(s.val for s in u.src) return None def fold_const_alu(a:UOp) -> UOp|None: @@ -39,8 +39,8 @@ def _quotient_base(q:UOp, base:UOp, div:int) -> UOp|None: # moves consts freely: the quotient may be merged ((x//c + a)//div -> (x + a*c)//(c*div) for div>0) and shifted ((y + k*D)//D == y//D + k) (q, s), (num, a) = q.pop_const(), base.pop_const() if q.op is not Ops.FLOORDIV or q.src[1].op is not Ops.CONST: return None - if div > 0 and num.op is Ops.FLOORDIV and num.src[1].op is Ops.CONST and q.src[1].arg == (c:=num.src[1].arg)*div: num, a, D = num.src[0], a*c, c*div - elif q.src[1].arg == div: D = div + if div > 0 and num.op is Ops.FLOORDIV and num.src[1].op is Ops.CONST and q.src[1].val == (c:=num.src[1].val)*div: num, a, D = num.src[0], a*c, c*div + elif q.src[1].val == div: D = div else: return None (x, xa), (p, pa) = num.pop_const(), q.src[0].pop_const() if p is not x or (t:=xa + a - pa) % D: return None @@ -54,13 +54,13 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None: for i,u in enumerate(terms): mod, mul = u.pop_const(Ops.MUL) if mod.op is not Ops.FLOORMOD or mod.src[1].op is not Ops.CONST: continue - base, div = mod.src[0], mod.src[1].arg + base, div = mod.src[0], mod.src[1].val for j,v in enumerate(terms): q, scale = v.pop_const(Ops.MUL) if i == j or scale != div*mul: continue rest = [t for k,t in enumerate(terms) if k not in (i,j)] if (b:=_quotient_base(q, base, div)) is not None: return (b*mul).usum(*rest) - if q.op is Ops.FLOORMOD and q.src[1].op is Ops.CONST and (d:=q.src[1].arg) > 0 and (b:=_quotient_base(q.src[0], base, div)) is not None: + if q.op is Ops.FLOORMOD and q.src[1].op is Ops.CONST and (d:=q.src[1].val) > 0 and (b:=_quotient_base(q.src[0], base, div)) is not None: return ((b % (div*d))*mul).usum(*rest) return None @@ -119,7 +119,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ (UPat.var("x", dtype=dtypes.bool).where(UPat.const(False, dtypes.bool), UPat.const(True, dtypes.bool)), lambda x: x.logical_not()), # CAST(bool -> int) != const — CAST(True)=1, CAST(False)=0, so fold based on const value (UPat.var("x", dtype=dtypes.bool).cast(dtypes.ints+(dtypes.weakint,)) != UPat.cvar("c"), - lambda x,c: x if c.arg == 0 else x.logical_not() if c.arg == 1 else x.const_like(True)), + lambda x,c: x if c.val == 0 else x.logical_not() if c.val == 1 else x.const_like(True)), (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)).trunc(), lambda x: x), # ** zero folding ** (UPat.var("x") < UPat.var("x"), lambda x: x.const_like(False, dtypes.bool)), # x < x -> False @@ -129,9 +129,9 @@ symbolic_simple = pm_data_invalid + PatternMatcher([ # (x&mask)>>k -> x>>k when mask only clears bits below k # TODO: combine this with "# rules for threefry" below ((UPat.var("x") & UPat.cvar("mask")) >> UPat.cvar("k"), - lambda x,mask,k: x >> k.arg if mask.arg | ((1 << k.arg) - 1) == -1 else None), + lambda x,mask,k: x >> k.val if mask.val | ((1 << k.val) - 1) == -1 else None), ((UPat.var("x") & UPat.cvar("mask")) // UPat.cvar("c"), - lambda x,mask,c: x // c.arg if c.arg > 0 and c.arg & (c.arg-1) == 0 and mask.arg | (c.arg-1) == -1 else None), + lambda x,mask,c: x // c.val if c.val > 0 and c.val & (c.val-1) == 0 and mask.val | (c.val-1) == -1 else None), (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) != UPat.var("x"), lambda x: x.const_like(False, dtypes.bool)), # x != x -> False (only ints) # ** constant folding ** @@ -199,7 +199,7 @@ def canonicalize_simplex(X:UOp) -> UOp|None: changed, ret = False, [] for u in X.split_uop(Ops.ADD): # assumed the const is the last src of MUL - if u.op is Ops.MUL and u.src[1].op is Ops.CONST and u.src[1].arg > 0: + if u.op is Ops.MUL and u.src[1].op is Ops.CONST and u.src[1].val > 0: changed = True u = u.src[0] if not (u.op in GroupOp.Irreducible and u.vmin >= 0): return None @@ -265,10 +265,10 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ # ** lt ** # c0*x sign(c0)*x < ceil(c1/abs(c0)) ((UPat.cvar("c0")*UPat.var("x", dtype=dtypes.weakint)) 0 else -x)<-(-c1.arg//abs(c0.arg)) if abs(c0.arg) > 1 else None), + lambda x,c0,c1: (x if c0.val > 0 else -x)<-(-c1.val//abs(c0.val)) if abs(c0.val) > 1 else None), # x//d x0, and -> c*d 0 else (x>c.arg*d.arg) if d.arg < 0 else None), + lambda x,d,c: (x 0 else (x>c.val*d.val) if d.val < 0 else None), # ** move add/mul consts to end (NOTE: this is still happening before constant folding) ** ((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), ((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), @@ -304,12 +304,12 @@ def parse_valid(v:UOp) -> tuple[UOp, bool, int]|None: # if it's X <= c, returns X, True, c # if it's X >= c, returns X, False, c - if v.op is Ops.CMPNE and v.src[1].op is Ops.CONST and v.src[1].arg == 1 and (s0:=v.src[0]).op is Ops.CMPLT and dtypes.is_int(s0.src[0].dtype): + if v.op is Ops.CMPNE and v.src[1].op is Ops.CONST and v.src[1].val == 1 and (s0:=v.src[0]).op is Ops.CMPLT and dtypes.is_int(s0.src[0].dtype): # (X < c).ne(True) -> X >= c return s0.src[0], False, int(s0.src[1].vmin) if v.op is Ops.CMPLT and dtypes.is_int(v.src[0].dtype): # c < X -> X >= c+1 (a const on the left is a lower bound on the right) - if v.src[0].op is Ops.CONST: return v.src[1], False, int(v.src[0].arg)+1 + if v.src[0].op is Ops.CONST: return v.src[1], False, int(v.src[0].val)+1 # X < c -> X <= c-1 return v.src[0], True, int((v.src[1]).vmax)-1 return None diff --git a/tinygrad/uop/validate.py b/tinygrad/uop/validate.py index 7c50b76ea4..d87fce5a31 100644 --- a/tinygrad/uop/validate.py +++ b/tinygrad/uop/validate.py @@ -44,8 +44,8 @@ z3_renderer = PatternMatcher([ (UPat((Ops.LOAD, Ops.INDEX), dtypes.bool), lambda ctx: (z3.Bool(f"load{len(ctx[1])}", ctx=ctx[0]), None)), # constants (UPat(Ops.CONST, arg=Invalid), lambda ctx: (z3.Int("Invalid", ctx=ctx[0]), None)), - (UPat(Ops.CONST, dtypes.ints+(dtypes.weakint,), name="x"), lambda x,ctx: (z3.IntVal(x.arg, ctx=ctx[0]), None)), - (UPat(Ops.CONST, dtypes.bool, name="x"), lambda x,ctx: (z3.BoolVal(x.arg, ctx=ctx[0]), None)), + (UPat(Ops.CONST, dtypes.ints+(dtypes.weakint,), name="x"), lambda x,ctx: (z3.IntVal(x.val, ctx=ctx[0]), None)), + (UPat(Ops.CONST, dtypes.bool, name="x"), lambda x,ctx: (z3.BoolVal(x.val, ctx=ctx[0]), None)), # casts from floats create new variables (UPat(Ops.CAST, dtypes.ints+(dtypes.weakint,), src=(UPat(dtype=dtypes.floats),), name="x"), lambda x,ctx: create_bounded(f"cast{len(ctx[1])}", x.dtype.min, x.dtype.max, ctx[0])), diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index d9c32c0a31..8574bb1a71 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -129,7 +129,7 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]: with soft_err(): if u.op in GroupOp.Movement and u.marg: argst = (mask_to_str if u.op in {Ops.SHRINK, Ops.PAD} else shape_to_str)(u.marg) if u.op is Ops.BINARY: argst = f"<{len(u.arg)} bytes>" - if u.op is Ops.CONST and dtypes.is_float(u.dtype): argst = f"{u.arg:g}" + if u.op is Ops.CONST and dtypes.is_float(u.dtype): argst = f"{u.val:g}" wrap_len = 200 if u.op is Ops.SOURCE else 80 label = f"{str(u.op).split('.')[1]}{(chr(10)+word_wrap(argst.replace(':', ''), wrap=wrap_len)) if u.arg is not None else ''}" if u.dtype != dtypes.void: label += f"\n{u.dtype}" @@ -138,7 +138,7 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]: # walk through excluded movement ops to find the underlying CONST cx = x while cx.op in GroupOp.Movement and len(cx.src) >= 1 and cx.src[0] in excluded: cx = cx.src[0] - arg = f"{cx.arg:g}" if cx.op is Ops.CONST and dtypes.is_float(cx.dtype) else cx.render() if cx.op is Ops.STACK else f"{cx.arg}" + arg = f"{cx.val:g}" if cx.op is Ops.CONST and dtypes.is_float(cx.dtype) else cx.render() if cx.op is Ops.STACK else f"{cx.arg}" label += f"\n{cx.op.name}{idx} {arg}" + (f" {cx.src[0].op}" if len(cx.src) else "") try: if len(rngs:=u.ranges):