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

use .val to access the value of Ops.CONST (#17347)

Browse Source
This commit is contained in:
George Hotz
2026-07-31 19:51:42 -07:00
committed by GitHub
parent 099d69ff7d
commit 9082ecef5d
23 changed files with 97 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/codegen/decomp/dtype.py
+1 -1
View File
@@ -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
tinygrad/codegen/decomp/op.py
+7 -7
View File
@@ -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),
(UPat.var("x", dtypes.sints)*-1 < UPat.cvar("c"), lambda x,c:-c<x),
((UPat.cvar("c1")<UPat.var("x", dtypes.sints)) & (UPat.var("x", dtypes.sints)<UPat.cvar("c2")),
lambda x,c1,c2: x.eq(c1+1) if c1.arg+1==c2.arg-1 else None), # (c-1)<x & x<(c+1) -> x==c
lambda x,c1,c2: x.eq(c1+1) if c1.val+1==c2.val-1 else None), # (c-1)<x & x<(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<<n.arg), c))]
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<<n.val), c))]
# some backends emit FDIV for RECIP, in that case: a*(1/b) -> a/b
if Ops.FDIV in ops:
pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))]
tinygrad/codegen/decomp/transcendental.py
+2 -2
View File
@@ -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"""
tinygrad/codegen/late/coalesce.py
+5 -5
View File
@@ -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)
tinygrad/codegen/opt/heuristic.py
+2 -2
View File
@@ -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)
tinygrad/codegen/simplify.py
+1 -1
View File
@@ -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})
tinygrad/engine/realize.py
+1 -1
View File
@@ -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
tinygrad/mixin/elementwise.py
+1 -1
View File
@@ -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)
tinygrad/renderer/cstyle.py
+15 -15
View File
@@ -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;")
tinygrad/renderer/isa/x86.py
+13 -13
View File
@@ -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:])
tinygrad/renderer/llvmir.py
+2 -2
View File
@@ -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:
tinygrad/renderer/nir.py
+2 -2
View File
@@ -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)<UPat.var('y'), lambda x,y: (x^True)&y),
# load/store 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"))),
tinygrad/renderer/ptx.py
+3 -3
View File
@@ -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:
tinygrad/renderer/wgsl.py
+3 -3
View File
@@ -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<u32>({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<u32>({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{'<workgroup>' 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),),)),
tinygrad/runtime/ops_dsp.py
+1 -1
View File
@@ -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):
tinygrad/runtime/ops_python.py
+1 -1
View File
@@ -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:
tinygrad/runtime/support/hcq2.py
+1 -1
View File
@@ -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 ()))
tinygrad/schedule/multi.py
+1 -1
View File
@@ -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
tinygrad/schedule/rangeify.py
+2 -2
View File
@@ -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
tinygrad/uop/divandmod.py
+5 -5
View File
@@ -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)
tinygrad/uop/ops.py
+15 -11
View File
@@ -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
tinygrad/uop/symbolic.py
+12 -12
View File
@@ -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)<UPat.cvar("c"), lambda x,c: lt_folding(x, c.arg) if 0 < c.arg else None),
(UPat.var("x", dtypes.weakint)<UPat.cvar("c"), lambda x,c: lt_folding(x, c.val) if 0 < c.val else None),
(UPat.var("x", dtypes.weakint)*-1 < UPat.var("y")*-1, lambda x,y: y<x),
# canonicalize a simplex with positive coefficients > 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) **
tinygrad/uop/upat.py
+1 -1
View File
@@ -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]: