diff --git a/test/unit/test_function.py b/test/unit/test_function.py index 866c3bd9c9..69a1c221fd 100644 --- a/test/unit/test_function.py +++ b/test/unit/test_function.py @@ -557,9 +557,9 @@ class TestFunctionTuple(unittest.TestCase): def f(a:Tensor): return Tensor.custom_kernel(Tensor.empty(*a.shape, dtype=a.dtype, device=a.device), a, fxn=inplace_add)[0] with self.assertRaisesRegex(RuntimeError, "implicit buffer"): f(Tensor([1., 2., 3., 4.]).contiguous().realize()) - def test_vector_load_is_program_input(self): + def test_shrink_load_is_program_input(self): out, inp = UOp.param(0, dtypes.float, (1,)), UOp.param(1, dtypes.float, (8,)) - values = UOp(Ops.VLOAD, dtypes.float, (inp[0],), arg=8) + values = UOp(Ops.SHRINK, src=(inp, UOp.const(dtypes.weakint, 0), UOp.const(dtypes.weakint, 8))).load() sink = out[0].store(values.index(0)).sink(arg=KernelInfo(name="vector_load")) info = ProgramInfo.from_sink(sink) self.assertEqual(info.outs, (0,)) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index b10d267e71..1fcc3edb01 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -120,8 +120,17 @@ pm_expand_broadcast = pm_wmma_add+PatternMatcher([ (UPat(Ops.WMMA, name="b"), broadcast_and_devec_wmma), ]) -def do_devectorize(b:UOp): - if b.shape == () or b.tag == "vectorized": return None +@functools.cache +def _uses_shrink_memory(x:UOp) -> bool: + if x.op in (Ops.LOAD, Ops.STORE): return x.src[0].op is Ops.SHRINK + if x.op is Ops.AFTER: return _uses_shrink_memory(x.src[0]) + if x.op in GroupOp.Elementwise or x.op in (Ops.STACK, Ops.RESHAPE, Ops.PERMUTE): + return any(_uses_shrink_memory(y) for y in x.src) + return False + +def do_devectorize(ctx:Renderer|tuple[set[UOp], Renderer], b:UOp): + preserved = ctx[0] if isinstance(ctx, tuple) else set() + if b.shape == () or b in preserved or _uses_shrink_memory(b): return None # broadcasting needs to be already unpacked, Invalid matches any dtype and shape if not all(x.shape == b.shape or x.base.arg is Invalid for x in b.src): return None src = [] @@ -187,6 +196,7 @@ def fix_group_for_reduce(x:UOp): @dataclass class ReduceContext: acc_num: int = 0 + renderer: Renderer|None = None def merge_reduce_ends(sink:UOp): # merge ENDs that share the same range and nesting context (only those created by reduce_to_acc) @@ -220,7 +230,8 @@ def reduce_ranges_to_acc(ctx:ReduceContext, r:UOp): acc_out = acc_initted.store(acc_initted.alu(r.arg[0], inp)).end(*r.src[1:]).rtag("mergeable") return acc.after(acc_out) -def expand_horizontal_reduce(r:UOp): +def expand_horizontal_reduce(ctx:ReduceContext, r:UOp): + if ctx.renderer is not None and ctx.renderer.has_native_reduce(r): return None inp = r.src[0] vals = [inp.index(*idx) for idx in itertools.product(*[range(inp.max_shape[a]) for a in range(r.arg[1])])] return functools.reduce(lambda x,y: x.alu(r.arg[0], y), vals) @@ -317,7 +328,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: sink = graph_rewrite(sink, expander2, ctx=build_range_map(sink), name="expander") # remove reduce - sink = graph_rewrite(sink, mop_cleanup+pm_reduce_local, ctx=ReduceContext(), name="remove reduces") + sink = graph_rewrite(sink, mop_cleanup+pm_reduce_local, ctx=ReduceContext(renderer=ren), name="remove reduces") # add locals sink = graph_rewrite(sink, pm_add_local_buffers, ctx=itertools.count(0), name="add local buffers") @@ -330,7 +341,9 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: sink = graph_rewrite(sink, symbolic_simple+pm_expand_broadcast+pm_add_loads, name="*** expand broadcast / add loads") # devectorize - sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=ren, name="devectorize2") + sink = graph_rewrite(sink, symbolic_simple, name="pre-devectorize symbolic") + native_reduce_uops = {u for r in sink.toposort() if ren.has_native_reduce(r) for u in r.src[0].backward_slice_with_self} + sink = graph_rewrite(sink, devectorizer2, ctx=(native_reduce_uops, ren), name="devectorize2") # simplify indexing sink = graph_rewrite(sink, indexing_simplify, name="simplify load/store indexing") @@ -340,7 +353,9 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # do memory coalescing (late) sink = memory_coalescing(sink, ren) - sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True) + sink = graph_rewrite(sink, symbolic_simple, name="pre-image symbolic", bottom_up=True) + native_reduce_uops = {u for r in sink.toposort() if ren.has_native_reduce(r) for u in r.src[0].backward_slice_with_self} + sink = graph_rewrite(sink, ew_devectorizer+pm_simplify_add_image, name="add images", ctx=(native_reduce_uops, ren), bottom_up=True) # extra symbolic before decomp. crashes without this? sink = graph_rewrite(sink, sym, name="extra symbolic") diff --git a/tinygrad/codegen/late/coalesce.py b/tinygrad/codegen/late/coalesce.py index 4cc27a3cc0..cd82421288 100644 --- a/tinygrad/codegen/late/coalesce.py +++ b/tinygrad/codegen/late/coalesce.py @@ -106,7 +106,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp: # TODO: this should handle images too, it's just memory coalescing if u.op in {Ops.LOAD, Ops.STORE}: assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalescing does not support gated loads/stores" - if u.tag == "vectorized": continue + if u.src[0].op is Ops.SHRINK: continue assert u.src[0].op is Ops.INDEX, f"memory coalescing should be on INDEX, not {u.src[0].op}" buf, idx_u = u.src[0].src if buf.addrspace == AddrSpace.REG: continue diff --git a/tinygrad/llm/cpu.py b/tinygrad/llm/cpu.py index 51b6e519b3..b09955476f 100644 --- a/tinygrad/llm/cpu.py +++ b/tinygrad/llm/cpu.py @@ -5,7 +5,7 @@ from tinygrad import Tensor, nn, UOp, getenv, dtypes from tinygrad.dtype import DType, AddrSpace from tinygrad.llm.gguf import _GGML_QUANT from tinygrad.renderer import Estimates -from tinygrad.uop.ops import Ops, GroupOp, KernelInfo, AxisType +from tinygrad.uop.ops import Ops, KernelInfo, AxisType if TYPE_CHECKING: from tinygrad.llm.model import ExpertWeights, FFNBlock, Linear @@ -17,14 +17,33 @@ def _concrete_int(value:int|UOp) -> int: def _dot_byte_parts(a:tuple[UOp, ...], b:tuple[UOp, ...]) -> UOp: av, bv = UOp.stack(*a), UOp.stack(*b) - return UOp(Ops.DOT, dtypes.int32, (av, bv), arg=4) + return _dot_byte_vectors(av, bv) + +def _dot_byte_vectors(a:UOp, b:UOp, scale:UOp|None=None) -> UOp: + assert a.shape == b.shape and len(a.shape) == 1 and a.shape[0] % 4 == 0 + product = a.cast(dtypes.int32) * b.cast(dtypes.int32) + if scale is not None: product = product * scale + return product.reshape(a.shape[0]//4, 4)._rop(Ops.ADD, (1,)) + +def _unpack_nibbles(packed:UOp, values:tuple[int, ...]) -> UOp: + assert packed.shape == (16,) and len(values) == 16 + def lookup(index:UOp) -> UOp: + value = UOp.stack(*(UOp.const(dtypes.int8, values[-1]) for _ in range(packed.shape[0]))) + for i in range(14, -1, -1): value = index.eq(i).where(UOp.const(dtypes.int8, values[i]), value) + return value + lowv, highv = lookup(packed & 15), lookup((packed >> 4) & 15) + return UOp.stack(lowv, highv).reshape(32) def _dot_bytes(a:tuple[UOp, ...], b:tuple[UOp, ...]) -> UOp: parts = _dot_byte_parts(a, b) return sum((parts.index(i) for i in range(len(a)//4)), UOp.const(dtypes.int32, 0)) def _contiguous_vector_load(ptr:UOp, lanes:int, dtype:DType|None=None) -> UOp: - return UOp(Ops.VLOAD, dtype or ptr.dtype, (ptr,), arg=lanes) + assert ptr.op is Ops.INDEX + buf, coords = ptr.src[0], ptr.src[1:] + index = sum((coord * math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(dtypes.weakint, 0)) + address = UOp(Ops.SHRINK, src=(buf.flatten(), index, UOp.const(dtypes.weakint, lanes))) + return address.load(dtype=dtype or ptr.dtype) def _contiguous_vector_ptr(buf:UOp, index:UOp, lanes:int) -> UOp: return UOp(Ops.SHRINK, src=(buf, index, UOp.const(dtypes.weakint, lanes))) @@ -32,14 +51,14 @@ def _contiguous_vector_ptr(buf:UOp, index:UOp, lanes:int) -> UOp: def _dot_bytes_ptr(a:UOp, b:UOp) -> UOp: av = _contiguous_vector_load(a, 32, dtypes.int8) bv = _contiguous_vector_load(b, 32, dtypes.int8) - return UOp(Ops.DOT, dtypes.int32, (av, bv), arg=4) + return _dot_byte_vectors(av, bv) def _dot_nibbles_ptr(packed:UOp, x:UOp, values:tuple[int, ...]) -> UOp: assert len(values) == 16 pv = _contiguous_vector_load(packed, 16) - qvalues = UOp(Ops.UNPACK_LUT, dtypes.int8, (pv,), arg=values) + qvalues = _unpack_nibbles(pv, values) xv = _contiguous_vector_load(x, 32, dtypes.int8) - return UOp(Ops.DOT, dtypes.int32, (qvalues, xv), arg=4) + return _dot_byte_vectors(qvalues, xv) def _dot_nibbles_pair_ptr(packed:UOp, x:UOp, values:tuple[int, ...]) -> tuple[UOp, UOp]: assert len(values) == 16 @@ -56,7 +75,7 @@ def _dot_q6_ptr(block:UOp, x:UOp, subgroup:int) -> UOp: hi = ((_contiguous_vector_load(hi_ptr, 32) >> high_shift) & 3) << 4 qvalues = ((lo | hi) - 32).bitcast(dtypes.int8) xv = _contiguous_vector_load(x, 32, dtypes.int8) - return UOp(Ops.DOT, dtypes.int32, (qvalues, xv), arg=4) + return _dot_byte_vectors(qvalues, xv) def _rms_f16_product_ptr(x:UOp, norm:UOp, weight:UOp, scale:UOp) -> UOp: return _contiguous_vector_load(x, 8).cast(dtypes.float32) * _contiguous_vector_load(norm, 8).cast(dtypes.float32) * \ @@ -250,23 +269,18 @@ def _load_f16(raw:UOp, offset:UOp) -> UOp: return bits.bitcast(dtypes.float16).cast(dtypes.float32) def _load_f16x8_ptr(raw:UOp) -> UOp: - return UOp(Ops.VLOAD, dtypes.float16, (raw,), arg=8).cast(dtypes.float32).rtag("vectorized") + return _contiguous_vector_load(raw, 8, dtypes.float16).cast(dtypes.float32) def _vector_reg(reg:UOp, *deps:UOp) -> UOp: reg = reg.after(*deps) return UOp(Ops.SHRINK, src=(reg, UOp.const(dtypes.weakint, 0), UOp.const(dtypes.weakint, reg.max_numel()))) def _vector_acc_init(reg:UOp, *deps:UOp) -> UOp: - return reg.after(_vector_reg(reg, *deps).store(reg.const_like(0), tag="vectorized")) - -def _vectorized_tree(x:UOp) -> UOp: - if x.max_numel() == 1: return x - src = tuple(_vectorized_tree(y) if y.max_numel() > 1 else y for y in x.src) - return x.replace(src=src, tag="vectorized") if x.op in GroupOp.Elementwise else x.replace(src=src) + return reg.after(_vector_reg(reg, *deps).store(reg.const_like(0))) def _vector_acc_update(reg:UOp, value:UOp, *deps:UOp) -> UOp: - previous = _vector_reg(reg, *deps).load(tag="vectorized") - return _vector_reg(reg, *deps).store((previous + _vectorized_tree(value)).rtag("vectorized"), tag="vectorized") + previous = _vector_reg(reg, *deps).load() + return _vector_reg(reg, *deps).store(previous + value) def _finite_exp2(x:UOp) -> UOp: # The causal-convolution activation is finite. Bounding the exponent preserves sigmoid saturation while avoiding @@ -512,7 +526,7 @@ def _cpu_expert_weighted_grouped_uop(out:UOp, raw:UOp, probs:UOp, head:UOp, next qvalues, scales = [], [] for subgroup in range(8): packed = _contiguous_vector_load(raw[base + 8 + subgroup * 16], 16) - qvalues.append(UOp(Ops.UNPACK_LUT, dtypes.int8, (packed,), arg=values)) + qvalues.append(_unpack_nibbles(packed, values)) low_byte = raw[base + 4 + subgroup // 2].load() low = (low_byte >> (4 * (subgroup % 2))) & 15 high = (high_word >> (2 * subgroup)) & 3 @@ -526,10 +540,10 @@ def _cpu_expert_weighted_grouped_uop(out:UOp, raw:UOp, probs:UOp, head:UOp, next route = matched_routes[route_idx].load().cast(dtypes.weakint) input_idx = route // routes_per_input block_acc = UOp.placeholder((8,), dtypes.int32, slot=4, addrspace=AddrSpace.REG) - stage = _vector_reg(block_acc, route_loop).store(block_acc.const_like(0), tag="vectorized") + stage = _vector_reg(block_acc, route_loop).store(block_acc.const_like(0)) for subgroup in range(8): xv = _contiguous_vector_load(xq[route, block, subgroup * 32], 32, dtypes.int8) - stage = _vector_acc_update(block_acc, UOp(Ops.DOT, dtypes.int32, (qvalues[subgroup], xv), arg=4) * scales[subgroup], stage) + stage = _vector_acc_update(block_acc, _dot_byte_vectors(qvalues[subgroup], xv, scales[subgroup]), stage) dot = sum((block_acc.after(stage).index(i) for i in range(8)), UOp.const(dtypes.int32, 0)).cast(dtypes.float32) contribution = dot * xd[route, block] * block_scale * matched_probs[route_idx].load() updated = totals[output_lane, input_idx].store(totals.after(route_loop)[output_lane, input_idx].load() + contribution) @@ -573,7 +587,7 @@ def _cpu_expert_silu_uop(out:UOp, raw0:UOp, raw1:UOp, sel:UOp, xq:UOp, xd:UOp, l vectorized = ggml_type in (14, 23) or (ggml_type == 21 and repacked) acc0 = UOp.placeholder((8 if vectorized else 1,), dtypes.float32, slot=0, addrspace=AddrSpace.REG) acc1 = UOp.placeholder((8 if vectorized else 1,), dtypes.float32, slot=1, addrspace=AddrSpace.REG) - initialized = UOp.group(*((_vector_reg(acc, job).store(acc.const_like(0), tag="vectorized") if vectorized else + initialized = UOp.group(*((_vector_reg(acc, job).store(acc.const_like(0)) if vectorized else acc.after(job).store(acc.const_like(0))) for acc in (acc0, acc1))) if ggml_type == 23: block = UOp.range(in_features // 256, 110) @@ -684,9 +698,9 @@ def _cpu_expert_silu_grouped_uop(out:UOp, raw0:UOp, raw1:UOp, head:UOp, next_rou qvalues:list[UOp] = [] scales:list[UOp] = [] for subgroup in range(0, 8, 2): - packed = _contiguous_vector_load(raw[row_base + meta_size + block * 128 + subgroup * 16], 32) - unpacked = UOp(Ops.UNPACK_LUT, dtypes.int8, (packed,), arg=values) - qvalues.extend(UOp.stack(*(unpacked.index(i) for i in range(off, off + 32))) for off in (0, 32)) + for offset in (0, 16): + packed = _contiguous_vector_load(raw[row_base + meta_size + block * 128 + subgroup * 16 + offset], 16) + qvalues.append(_unpack_nibbles(packed, values)) for subgroup in range(8): scale_byte = raw[meta + 2 + subgroup // 2].load() scales.append(1 + 2 * ((scale_byte >> (4 * (subgroup % 2))) & 15).cast(dtypes.int32)) @@ -698,11 +712,11 @@ def _cpu_expert_silu_grouped_uop(out:UOp, raw0:UOp, raw1:UOp, head:UOp, next_rou route = matched[route_idx].load().cast(dtypes.weakint) xidx = (route.cast(dtypes.uint32) // routes_per_input).cast(dtypes.weakint) block_acc = UOp.placeholder((8,), dtypes.int32, slot=6 + projection, addrspace=AddrSpace.REG) - stage = _vector_reg(block_acc, route_loop).store(block_acc.const_like(0), tag="vectorized") + stage = _vector_reg(block_acc, route_loop).store(block_acc.const_like(0)) for subgroup in range(8): xv = _contiguous_vector_load(xq[xidx, block, subgroup * 32], 32, dtypes.int8) - parts = UOp(Ops.DOT, dtypes.int32, (qvalues[subgroup], xv), arg=4) - stage = _vector_acc_update(block_acc, parts * scales[subgroup], stage) + parts = _dot_byte_vectors(qvalues[subgroup], xv, scales[subgroup]) + stage = _vector_acc_update(block_acc, parts, stage) block_sum = block_acc.after(stage) value = sum((block_sum.index(i) for i in range(8)), UOp.const(dtypes.int32, 0)).cast(dtypes.float32) * \ xd[xidx, block] * _load_f16(raw, meta) @@ -771,7 +785,7 @@ def _moe_stage1_uop(rhidden:UOp, shidden:UOp, rgate:UOp, rup:UOp, sgate:UOp, sup expert = sel[route].load().cast(dtypes.weakint) racc0 = UOp.placeholder((8 if expert_repacked else 1,), dtypes.float32, slot=0, addrspace=AddrSpace.REG) racc1 = UOp.placeholder((8 if expert_repacked else 1,), dtypes.float32, slot=1, addrspace=AddrSpace.REG) - rinit = UOp.group(*((_vector_reg(acc, routed_job).store(acc.const_like(0), tag="vectorized") if expert_repacked else + rinit = UOp.group(*((_vector_reg(acc, routed_job).store(acc.const_like(0)) if expert_repacked else acc.after(routed_job).store(acc.const_like(0))) for acc in (racc0, racc1))) if expert_repacked: rblock = UOp.range(dim // 256, 100) @@ -802,7 +816,7 @@ def _moe_stage1_uop(rhidden:UOp, shidden:UOp, rgate:UOp, rup:UOp, sgate:UOp, sup shared_output, groups = shared_begin + shared_job, dim // 32 sacc0 = UOp.placeholder((8,), dtypes.float32, slot=2, addrspace=AddrSpace.REG) sacc1 = UOp.placeholder((8,), dtypes.float32, slot=3, addrspace=AddrSpace.REG) - sinit = UOp.group(*(_vector_reg(acc, shared_job).store(acc.const_like(0), tag="vectorized") for acc in (sacc0, sacc1))) + sinit = UOp.group(*(_vector_reg(acc, shared_job).store(acc.const_like(0)) for acc in (sacc0, sacc1))) if shared_repacked and groups % 8 == 0: sblock = UOp.range(groups // 8, 101) svalues = [] @@ -1264,7 +1278,7 @@ def _attention_prefill_online_uop(out:UOp, q:UOp, cache:UOp, start_pos:UOp) -> U row_sums = tuple(UOp.placeholder((1,), dtypes.float32, slot=2*token_tile+i, addrspace=AddrSpace.REG) for i in range(token_tile)) clear = UOp.range(chunks, 90) cleared = UOp.group(*(_contiguous_vector_ptr(numerator.after(job), clear * 8, 8).store( - UOp.stack(*(UOp.const(dtypes.float32, 0) for _ in range(8))), tag="vectorized") for numerator in numerators)).end(clear) + UOp.stack(*(UOp.const(dtypes.float32, 0) for _ in range(8)))) for numerator in numerators)).end(clear) initialized = UOp.group(cleared, *(row_max.after(job).store(-math.inf) for row_max in row_maxes), *(row_sum.after(job).store(0.0) for row_sum in row_sums)) @@ -1299,8 +1313,7 @@ def _attention_prefill_online_uop(out:UOp, q:UOp, cache:UOp, start_pos:UOp) -> U nbase = vchunk * 8 previous = _contiguous_vector_load(numerator.after(initialized, position)[nbase], 8) updated = previous * old_scale + values * weight - numerator_updates.append(_contiguous_vector_ptr(numerator.after(qk_done), nbase, 8).store( - _vectorized_tree(updated), tag="vectorized")) + numerator_updates.append(_contiguous_vector_ptr(numerator.after(qk_done), nbase, 8).store(updated)) values_done = UOp.group(*numerator_updates).end(vchunk) state_updates = [row_max[0].store(next_max) for row_max,next_max in zip(row_maxes, next_maxes)] state_updates += [row_sum[0].store(row_sum.after(initialized, position)[0].load() * old_scale + weight) @@ -1312,7 +1325,7 @@ def _attention_prefill_online_uop(out:UOp, q:UOp, cache:UOp, start_pos:UOp) -> U for token_offset,(numerator,row_sum) in enumerate(zip(numerators, row_sums)): query = bh * tokens + token_base + token_offset value = _contiguous_vector_load(numerator.after(positions_done)[output * 8], 8) / row_sum.after(positions_done)[0].load() - stores.append(_contiguous_vector_ptr(outf, query * dim + output * 8, 8).store(_vectorized_tree(value), tag="vectorized")) + stores.append(_contiguous_vector_ptr(outf, query * dim + output * 8, 8).store(value)) return UOp.group(*stores).end(output, job, core).sink( arg=KernelInfo(name=f"attention_prefill_online_uop_{batch}_{heads}_{tokens}_{kv_heads}_{dim}_{cache_len}", opts_to_apply=())).rtag("cpu_parallel") @@ -1364,8 +1377,7 @@ def _gated_delta_prefill_uop(core:UOp, next_state:UOp, q:UOp, k:UOp, v:UOp, beta current = UOp.placeholder((dim * dim,), dtypes.float32, slot=0, addrspace=AddrSpace.REG) init_chunk = UOp.range(dim * dim // 8, 90) initial_values = _contiguous_vector_load(statef[bh * dim * dim + init_chunk * 8], 8).cast(dtypes.float32) - initialized_state = _contiguous_vector_ptr(current, init_chunk * 8, 8).store( - _vectorized_tree(initial_values), tag="vectorized").end(init_chunk) + initialized_state = _contiguous_vector_ptr(current, init_chunk * 8, 8).store(initial_values).end(init_chunk) token = UOp.range(tokens, 91) token_base = (bh * tokens + token) * dim @@ -1381,8 +1393,8 @@ def _gated_delta_prefill_uop(core:UOp, next_state:UOp, q:UOp, k:UOp, v:UOp, beta row = UOp.range(dim, 93) state_k = UOp.placeholder((8,), dtypes.float32, slot=3, addrspace=AddrSpace.REG) state_q = UOp.placeholder((8,), dtypes.float32, slot=4, addrspace=AddrSpace.REG) - dot_init = UOp.group(_vector_reg(state_k, row).store(state_k.const_like(0), tag="vectorized"), - _vector_reg(state_q, row).store(state_q.const_like(0), tag="vectorized")) + dot_init = UOp.group(_vector_reg(state_k, row).store(state_k.const_like(0)), + _vector_reg(state_q, row).store(state_q.const_like(0))) col = UOp.range(chunks, 94) state_vec = _contiguous_vector_load(current.after(initialized_state, token)[row * dim + col * 8], 8) q_vec = _contiguous_vector_load(qf[token_base + col * 8], 8) @@ -1399,8 +1411,7 @@ def _gated_delta_prefill_uop(core:UOp, next_state:UOp, q:UOp, k:UOp, v:UOp, beta current_values = _contiguous_vector_load(current.after(initialized_state, token)[update_base], 8) key_values = _contiguous_vector_load(kf[token_base + update * 8], 8) next_values = current_values * av + delta * key_values - updated_state = _contiguous_vector_ptr(current.after(dots), update_base, 8).store( - _vectorized_tree(next_values), tag="vectorized").end(update) + updated_state = _contiguous_vector_ptr(current.after(dots), update_base, 8).store(next_values).end(update) rows_done = UOp.group(saved_core, updated_state).end(row) norm_acc = UOp.placeholder((8,), dtypes.float32, slot=5, addrspace=AddrSpace.REG) @@ -1416,8 +1427,7 @@ def _gated_delta_prefill_uop(core:UOp, next_state:UOp, q:UOp, k:UOp, v:UOp, beta save_chunk = UOp.range(dim * dim // 8, 98) saved_values = _contiguous_vector_load(current.after(token_done)[save_chunk * 8], 8).cast(next_state.dtype) - saved_state = _contiguous_vector_ptr(nextf, bh * dim * dim + save_chunk * 8, 8).store( - _vectorized_tree(saved_values), tag="vectorized").end(save_chunk) + saved_state = _contiguous_vector_ptr(nextf, bh * dim * dim + save_chunk * 8, 8).store(saved_values).end(save_chunk) return saved_state.end(bh).sink( arg=KernelInfo(name=f"gated_delta_prefill_uop_{batch}_{heads}_{tokens}_{dim}_{state.dtype.name}", opts_to_apply=())).rtag("cpu_parallel") @@ -1446,8 +1456,8 @@ def _gated_delta_uop(core:UOp|None, next_state:UOp, q:UOp, k:UOp, v:UOp, beta:UO row = UOp.range(dim, 91) state_k = UOp.placeholder((8,), dtypes.float32, slot=2, addrspace=AddrSpace.REG) state_q = UOp.placeholder((8,), dtypes.float32, slot=3, addrspace=AddrSpace.REG) - initialized = UOp.group(_vector_reg(state_k, row).store(state_k.const_like(0), tag="vectorized"), - _vector_reg(state_q, row).store(state_q.const_like(0), tag="vectorized")) + initialized = UOp.group(_vector_reg(state_k, row).store(state_k.const_like(0)), + _vector_reg(state_q, row).store(state_q.const_like(0))) col = UOp.range(chunks, 92) state_vec = _contiguous_vector_load(statef[bh * dim * dim + row * dim + col * 8], 8).cast(dtypes.float32) q_vec = _contiguous_vector_load(qf[bh * dim + col * 8], 8) @@ -1464,8 +1474,8 @@ def _gated_delta_uop(core:UOp|None, next_state:UOp, q:UOp, k:UOp, v:UOp, beta:UO state_base = bh * dim * dim + row * dim + update * 8 state_values = _contiguous_vector_load(statef[state_base], 8).cast(dtypes.float32) key_values = _contiguous_vector_load(kf[bh * dim + update * 8], 8) - next_values = _vectorized_tree(state_values * av + delta * key_values).cast(next_state.dtype) - update_state = _contiguous_vector_ptr(nextf.after(dots), state_base, 8).store(next_values, tag="vectorized").end(update) + next_values = (state_values * av + delta * key_values).cast(next_state.dtype) + update_state = _contiguous_vector_ptr(nextf.after(dots), state_base, 8).store(next_values).end(update) rows_done = UOp.group(saved_core, update_state).end(row) if normalize: @@ -1498,14 +1508,14 @@ def _gated_delta_uop(core:UOp|None, next_state:UOp, q:UOp, k:UOp, v:UOp, beta:UO norm_values = _contiguous_vector_load(norm_weight[quant_group * 32 + offset], 8).cast(dtypes.float32) gate_values = _contiguous_vector_load(gatef[base + offset], 8).cast(dtypes.float32) gate_sigmoid = UOp.stack(*(gate_values.index(i).sigmoid() for i in range(8))) - value_vecs.append(_vectorized_tree(source_values * scale * norm_values * gate_values * gate_sigmoid)) + value_vecs.append(source_values * scale * norm_values * gate_values * gate_sigmoid) values = tuple(value_vecs[i // 8].index(i % 8) for i in range(32)) amax = functools.reduce(lambda a,b:a.maximum(b), (value.abs() for value in values)) d = (amax / 127).maximum(1e-8) quant_stores = [] for chunk_idx,value_vec in enumerate(value_vecs): quant_values = UOp.stack(*((value_vec.index(i) / d).round().maximum(-127).minimum(127).cast(dtypes.int8) for i in range(8))) - quant_stores.append(_contiguous_vector_ptr(quantf, base + chunk_idx * 8, 8).store(quant_values, tag="vectorized")) + quant_stores.append(_contiguous_vector_ptr(quantf, base + chunk_idx * 8, 8).store(quant_values)) stores = UOp.group(scalef[bh * (dim // 32) + quant_group].store(d), *quant_stores).end(quant_group) name = f"gated_delta_q8_uop_{batch}_{heads}_{dim}_{state.dtype.name}" return stores.end(bh).sink(arg=KernelInfo(name=name, opts_to_apply=())).rtag("cpu_parallel") @@ -1611,8 +1621,8 @@ def _gdn_qkv_uop(q:UOp, k:UOp, v:UOp, conv:UOp, k_heads:int, v_heads:int, dim:in qsum = UOp.placeholder((8,), dtypes.float32, slot=0, addrspace=AddrSpace.REG) ksum = UOp.placeholder((8,), dtypes.float32, slot=1, addrspace=AddrSpace.REG) - initialized = UOp.group(_vector_reg(qsum, job).store(qsum.const_like(0), tag="vectorized"), - _vector_reg(ksum, job).store(ksum.const_like(0), tag="vectorized")) + initialized = UOp.group(_vector_reg(qsum, job).store(qsum.const_like(0)), + _vector_reg(ksum, job).store(ksum.const_like(0))) chunk = UOp.range(dim // 8, 100) qvalue = _contiguous_vector_load(conv[qbase + chunk * 8], 8).cast(dtypes.float32) kvalue = _contiguous_vector_load(conv[kbase + chunk * 8], 8).cast(dtypes.float32) @@ -1627,7 +1637,7 @@ def _gdn_qkv_uop(q:UOp, k:UOp, v:UOp, conv:UOp, k_heads:int, v_heads:int, dim:in kvalues = _contiguous_vector_load(conv[kbase + out_chunk * 8], 8).cast(dtypes.float32) * kscale vvalues = _contiguous_vector_load(conv[vbase + out_chunk * 8], 8).cast(dtypes.float32) stores = UOp.group(*( - target[batch_idx, head, token, out_chunk * 8 + lane].store(_vectorized_tree(values).index(lane)) + target[batch_idx, head, token, out_chunk * 8 + lane].store(values.index(lane)) for target,values in ((q, qvalues), (k, kvalues), (v, vvalues)) for lane in range(8))).end(out_chunk) return stores.end(job, core).sink( arg=KernelInfo(name=f"gdn_qkv_uop_{batch}_{tokens}_{k_heads}_{v_heads}_{dim}", opts_to_apply=())).rtag("cpu_parallel") @@ -1744,7 +1754,7 @@ def _f16_matvec_uop(out:UOp, x:UOp, weight:UOp) -> UOp: token_base = token_block * token_tile xf, wf, outf = x.flatten(), weight.flatten(), out.flatten() accs = tuple(UOp.placeholder((8,), dtypes.float32, slot=i, addrspace=AddrSpace.REG) for i in range(token_tile)) - accs = tuple(acc.after(_vector_reg(acc, output_job, token_block).store(acc.const_like(0), tag="vectorized")) for acc in accs) + accs = tuple(acc.after(_vector_reg(acc, output_job, token_block).store(acc.const_like(0))) for acc in accs) chunk = UOp.range(in_features // 8, 92) weights = _load_f16x8_ptr(wf[output * in_features + chunk * 8]) updates = [] @@ -1752,9 +1762,8 @@ def _f16_matvec_uop(out:UOp, x:UOp, weight:UOp) -> UOp: xbase = (token_base + token) * in_features + chunk * 8 values = _load_f16x8_ptr(xf[xbase]) if x.dtype == dtypes.float16 else \ UOp.stack(*(xf[xbase + lane].load() for lane in range(8))) - previous = _vector_reg(acc, chunk).load(tag="vectorized") - updates.append(_vector_reg(acc, chunk).store( - (previous + (values * weights).rtag("vectorized")).rtag("vectorized"), tag="vectorized")) + previous = _vector_reg(acc, chunk).load() + updates.append(_vector_reg(acc, chunk).store(previous + values * weights)) done = UOp.group(*updates).end(chunk) stores = [outf[(token_base + token) * out_features + output].store( sum((acc.after(done).index(lane) for lane in range(8)), UOp.const(dtypes.float32, 0)).cast(out.dtype)) diff --git a/tinygrad/renderer/__init__.py b/tinygrad/renderer/__init__.py index 8da54ef911..06ffb7ff73 100644 --- a/tinygrad/renderer/__init__.py +++ b/tinygrad/renderer/__init__.py @@ -30,12 +30,12 @@ class Estimates: if u.op in {Ops.INDEX, Ops.SHRINK}: excluded = excluded.union(set(UOp.sink(*u.src[1:]).toposort(lambda x: x.op is not Ops.END))) for u in uops: - if u.op in {Ops.LOAD, Ops.VLOAD, Ops.STORE}: + if u.op in {Ops.LOAD, Ops.STORE}: buf = u while len(buf.src) and buf.op is not Ops.PARAM: buf = buf.src[0] if buf.op is Ops.PARAM: # u.src[0] is INDEX, cap at buffer size for re-reads (e.g. matmul) - elements = u.max_numel() if u.op is Ops.VLOAD else u.src[0].max_numel() + elements = u.src[0].max_numel() accessed = mem.get((buf, u.op), 0) + elements * u.src[0].dtype.scalar().itemsize * mults mem[(buf, u.op)] = smin(accessed, buf.max_numel() * buf.dtype.scalar().itemsize) if u.op is Ops.RANGE: @@ -47,7 +47,7 @@ class Estimates: elif u.op is Ops.END: mults = mult_stack.pop(-1) elif u.op is Ops.SPECIAL: mults *= cast(sint, u.src[0].ssimplify()) # NOTE: we don't push to the mult_stack here, you can't end these elif u.op is Ops.PARAM and u.arg.addrspace == AddrSpace.ALU and u.expr == 'core_id': mults *= int(u.vmax) + 1 - elif u.op in {Ops.LOAD, Ops.VLOAD} and u.src[0].addrspace != AddrSpace.REG: + elif u.op is Ops.LOAD and u.src[0].addrspace != AddrSpace.REG: lds += u.max_numel() * u.dtype.scalar().itemsize * mults elif u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG: lds += u.max_numel() * u.src[1].dtype.scalar().itemsize * mults @@ -62,6 +62,7 @@ class Renderer: suffix: str = "" # TODO: make this generic with a list of supported types supports_float4: bool = True + def has_native_reduce(self, x:UOp) -> bool: return False has_local: bool = True has_threads: bool = False has_shared: bool = True diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 5c41a7a785..d3b20eda8c 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -22,7 +22,7 @@ base_rewrite = PatternMatcher([ # casting (UPat(Ops.CAST, name="x"), lambda ctx,x: f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_type(x)})" \ - if x.max_numel() > 1 and (x.addrspace is AddrSpace.REG or x.tag == "vectorized") else None), + if x.max_numel() > 1 and x.addrspace is AddrSpace.REG else None), (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx[x.src[0]])})"), (UPat(Ops.BITCAST, name="x"), lambda ctx,x: ctx[x.src[0]] if x.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None), (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"__builtin_bit_cast({ctx.render_type(x)}, ({ctx.render_type(x.src[0])})({ctx[x.src[0]]}))"), @@ -131,6 +131,7 @@ class CStyleLanguage(Renderer): float4: str|None = None float4_style: tuple[str, str] = ('(', ')') gep_arr_threshold: int = 4 + def should_inline_where(self, x:UOp) -> bool: return False type_map: dict[DType, str] = {} infinity: str = "INFINITY" nan: str = "NAN" @@ -239,7 +240,7 @@ class CStyleLanguage(Renderer): if (u.op is not Ops.CAST or u.max_numel() == 1) and (u.op in {Ops.CONST, Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \ (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG and child_count[u] == 1) or \ (u.op is Ops.CAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \ - (u.op in {Ops.STACK, Ops.DOT, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and + (u.op in {Ops.STACK, Ops.REDUCE, *GroupOp.ALU, Ops.CAST, Ops.BITCAST} and (u.op is not Ops.WHERE or self.should_inline_where(u)) and child_count[u] == 1 and not getenv("EXPAND_SSA"))): r[u] = l else: @@ -260,6 +261,23 @@ class ClangRenderer(CStyleLanguage): gep_arr_threshold = 0 has_local = False has_threads = bool(getenv("THREADS", 1)) + def should_inline_where(self, x:UOp) -> bool: return self._render_lut_lookup(self, x) is not None + @staticmethod + def _byte_dot_reduce(x:UOp) -> tuple[UOp, UOp, UOp|None]|None: + if x.op is not Ops.REDUCE or x.arg != (Ops.ADD, 1) or len(x.src) != 1: return None + permute = x.src[0] + if permute.op is not Ops.PERMUTE or permute.arg != (1, 0) or permute.src[0].op is not Ops.RESHAPE: return None + product, scale = permute.src[0].src[0], None + if product.op is Ops.MUL: + for base, candidate_scale in (product.src, product.src[::-1]): + if base.op is Ops.MUL and all(s.op is Ops.CAST and s.dtype is dtypes.int32 and s.src[0].dtype is dtypes.int8 for s in base.src): + product, scale = base, candidate_scale + break + if product.op is not Ops.MUL or not all(s.op is Ops.CAST and s.dtype is dtypes.int32 and s.src[0].dtype is dtypes.int8 for s in product.src): + return None + a, b = (s.src[0] for s in product.src) + return (a, b, scale) if a.shape == b.shape and a.shape in ((16,), (32,)) else None + def has_native_reduce(self, x:UOp) -> bool: return self._byte_dot_reduce(x) is not None global_max = (NUM_CPU_THREADS.value, 0, 0) infinity = "__builtin_inff()" nan = '__builtin_nanf("")' @@ -283,6 +301,32 @@ class ClangRenderer(CStyleLanguage): f"__builtin_ia32_pmaddubsw{suffix}(__builtin_elementwise_abs({ctx[a]}), __builtin_ia32_psignb{suffix}({ctx[b]}, {ctx[a]}))), " \ f"(short __attribute__((ext_vector_type({lanes})))){{{','.join([multiplier] * lanes)}}})" @staticmethod + def _uniform_vector_const(x:UOp): + return x.src[0].arg if x.op is Ops.STACK and x.src and all(v.op is Ops.CONST and v.arg == x.src[0].arg for v in x.src) else None + @staticmethod + def _render_lut_lookup(ctx, x:UOp): + values:dict[int, int] = {} + index:UOp|None = None + node = x + while node.op is Ops.WHERE and node.src[0].op in (Ops.CMPEQ, Ops.CMPNE): + cond, true_value, false_value = node.src + left, right = cond.src + key = ClangRenderer._uniform_vector_const(right) + if key is None: left, right, key = right, left, ClangRenderer._uniform_vector_const(left) + if not isinstance(key, int) or not 0 <= key < 16 or (index is not None and left is not index): return None + index = left + selected, node = (true_value, false_value) if cond.op is Ops.CMPEQ else (false_value, true_value) + value = ClangRenderer._uniform_vector_const(selected) + if not isinstance(value, int): return None + values[key] = value + default = ClangRenderer._uniform_vector_const(node) + if index is None or not isinstance(default, int) or index.shape not in ((16,), (32,)): return None + lut_values = tuple(values.get(i, default) for i in range(16)) + lanes, suffix = index.shape[0], "128" if index.shape == (16,) else "256" + charv = f"signed char __attribute__((ext_vector_type({lanes})))" + lut = f"({charv}){{{','.join(map(str, lut_values * (lanes // 16)))}}}" + return f"__builtin_ia32_pshufb{suffix}({lut}, __builtin_bit_cast({charv}, {ctx[index]}))" + @staticmethod def _render_contiguous_stack(ctx, x:UOp): if not x.src or not all(v.op is Ops.INDEX and len(v.src) == 2 and v.src[0] is x.src[0] and v.src[1].op is Ops.CONST for v in x.src): return None @@ -290,19 +334,25 @@ class ClangRenderer(CStyleLanguage): if indices != tuple(range(indices[0], indices[0] + len(indices))): return None return f"__builtin_shufflevector({ctx[x.src[0]]}, {ctx[x.src[0]]}, {','.join(map(str, indices))})" @staticmethod - def _render_unpack_lut(ctx, x:UOp, packed:UOp): - if packed.shape not in ((16,), (32,)): return None - lanes, suffix = packed.shape[0], "128" if packed.shape[0] == 16 else "256" - charv = f"signed char __attribute__((ext_vector_type({lanes})))" - pv = f"__builtin_bit_cast({charv}, {ctx[packed]})" - lut_values = x.arg if lanes == 16 else x.arg * 2 - lut = f"({charv}){{{','.join(map(str, lut_values))}}}" - mask = f"({charv}){{{','.join(['15'] * lanes)}}}" - lo = f"__builtin_ia32_pshufb{suffix}({lut}, {pv} & {mask})" - hi = f"__builtin_ia32_pshufb{suffix}({lut}, ({pv} >> 4) & {mask})" - if lanes == 16: indices = tuple(range(32)) - else: indices = (*range(16), *range(32, 48), *range(16, 32), *range(48, 64)) - return f"__builtin_shufflevector({lo}, {hi}, {','.join(map(str, indices))})" + def _render_concat_stack(ctx, x:UOp): + if len(x.src) != 2 or x.src[0].shape != x.src[1].shape or len(x.src[0].shape) != 1: return None + lanes = x.src[0].shape[0] + return f"__builtin_shufflevector({ctx[x.src[0]]}, {ctx[x.src[1]]}, {','.join(map(str, range(lanes*2)))})" + @staticmethod + def _render_vector_permute(ctx, x:UOp): + old_shape, order = x.src[0].shape, x.arg + new_shape = tuple(old_shape[i] for i in order) + indices = [] + for flat in range(x.max_numel()): + coord, rem = [], flat + for size in reversed(new_shape): + coord.append(rem % size) + rem //= size + new_coord = tuple(reversed(coord)) + old_coord = tuple(new_coord[order.index(i)] for i in range(len(order))) + old_flat = sum(c * math.prod(old_shape[i+1:]) for i,c in enumerate(old_coord)) + indices.append(old_flat) + return f"__builtin_shufflevector({ctx[x.src[0]]}, {ctx[x.src[0]]}, {','.join(map(str, indices))})" @staticmethod def _render_vector_load(ctx, x:UOp, address:UOp): vec_type = ctx._render_dtype(x.dtype, x.max_numel(), AddrSpace.REG) @@ -312,23 +362,25 @@ class ClangRenderer(CStyleLanguage): @staticmethod def _render_vector_store(ctx, address:UOp, value:UOp): vec_type = ctx._render_dtype(value.dtype, value.max_numel(), AddrSpace.REG) - # Like VLOAD, explicit vector stores can target unaligned packed data or register-backed arrays. + # Explicit vector stores can target unaligned packed data or register-backed arrays. return f"do {{ {vec_type} _v = {ctx[value]}; __builtin_memcpy({ctx[address]}, &_v, sizeof(_v)); }} while (0);" string_rewrite = PatternMatcher([ + (UPat(Ops.WHERE, name="x"), lambda ctx,x: ClangRenderer._render_lut_lookup(ctx, x)), + (UPat(Ops.MUL, dtypes.int32, src=(UPat(Ops.REDUCE, name="dot"), UPat.var("scale"))), + lambda ctx,dot,scale: ClangRenderer._render_byte_dot(ctx, pair[0], pair[1], scale) + if (pair:=ClangRenderer._byte_dot_reduce(dot)) is not None and pair[2] is None and scale.vmin >= -32768 and scale.vmax <= 32767 else None), + (UPat(Ops.REDUCE, name="x"), lambda ctx,x: ClangRenderer._render_byte_dot(ctx, *pair) + if (pair:=ClangRenderer._byte_dot_reduce(x)) is not None else None), + (UPat(Ops.CAST, name="x"), lambda ctx,x: f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_type(x)})" + if x.max_numel() > 1 else None), + (UPat(Ops.RESHAPE, name="x"), lambda ctx,x: ctx[x.src[0]]), + (UPat(Ops.PERMUTE, name="x"), lambda ctx,x: ClangRenderer._render_vector_permute(ctx, x)), + (UPat(Ops.STACK, name="x"), lambda ctx,x: ClangRenderer._render_concat_stack(ctx, x)), (UPat(Ops.STACK, name="x"), lambda ctx,x: ClangRenderer._render_contiguous_stack(ctx, x)), - (UPat(Ops.MUL, dtypes.int32, - src=(UPat(Ops.DOT, dtypes.int32, src=(UPat.var("a", dtypes.int8), UPat.var("b", dtypes.int8)), arg=4), - UPat.var("scale", dtypes.int32))), - lambda ctx,a,b,scale: ClangRenderer._render_byte_dot(ctx, a, b, scale) - if scale.vmin >= -32768 and scale.vmax <= 32767 else None), - (UPat(Ops.DOT, dtypes.int32, src=(UPat.var("a", dtypes.int8), UPat.var("b", dtypes.int8)), arg=4), - lambda ctx,a,b: ClangRenderer._render_byte_dot(ctx, a, b)), - (UPat(Ops.UNPACK_LUT, dtypes.int8, src=(UPat.var("packed", dtypes.uint8),), name="x"), - lambda ctx,x,packed: ClangRenderer._render_unpack_lut(ctx, x, packed)), (UPat(Ops.STORE, src=(UPat(Ops.SHRINK, name="address"), UPat.var("value")), name="x"), lambda ctx,x,address,value: ClangRenderer._render_vector_store(ctx, address, value) - if x.tag == "vectorized" and value.max_numel() > 1 else None), - (UPat(Ops.VLOAD, src=(UPat.var("address"),), name="x"), + if value.max_numel() > 1 else None), + (UPat(Ops.LOAD, src=(UPat(Ops.SHRINK, name="address"),), name="x"), lambda ctx,x,address: ClangRenderer._render_vector_load(ctx, x, address)), ]) + base_rewrite diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index c778d9be6f..e39fc8b876 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -52,16 +52,12 @@ class Ops(FastEnum): INDEX = auto(); SHRINK = auto() # load/store before math - LOAD = auto(); VLOAD = auto(); STORE = auto() + LOAD = auto(); STORE = auto() # ** 4 -- math ** # tensor core math op, not elementwise WMMA = auto() - # grouped dot product. arg is the number of input lanes reduced into each output lane - DOT = auto() - # unpack low/high nibbles through a byte lookup table. arg is the tuple of table values - UNPACK_LUT = auto() # UnaryOps CAST = auto(); BITCAST = auto(); EXP2 = auto(); LOG2 = auto(); SIN = auto() diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 59fcb14cf9..9a68080601 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -150,8 +150,6 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None: case Ops.WMMA: # WMMA output dtype is the accumulator dtype (src[2]) return src[2].dtype - case Ops.DOT | Ops.UNPACK_LUT | Ops.VLOAD: - return None case Ops.GETTUPLE: # GETTUPLE extracts from a TUPLE (possibly through a FUNCTION) in_tuple = src[0].src[0] if src[0].op is Ops.FUNCTION else src[0] @@ -396,17 +394,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass): case Ops.WMMA: wmma_b = _broadcast_shape(self.src[0].shape[:-1], self.src[1].shape[:-1], self.src[2].shape[:-1]) return wmma_b + (self.src[2].shape[-1],) - case Ops.DOT: - assert len(self.src) == 2 and self.src[0].shape == self.src[1].shape and len(self.src[0].shape) and \ - isinstance(self.arg, int) and self.arg > 0 and self.src[0].shape[-1] % self.arg == 0 - return self.src[0].shape[:-1] + (self.src[0].shape[-1] // self.arg,) - case Ops.UNPACK_LUT: - assert len(self.src) == 1 and isinstance(self.arg, tuple) - return self.src[0].shape[:-1] + (self.src[0].shape[-1] * 2,) - case Ops.VLOAD: - assert len(self.src) == 1 and self.src[0].shape == () and isinstance(self.arg, int) and self.arg > 0 - return (self.arg,) - # passthrough ops case Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.LOAD | \ Ops.COPY | Ops.ALLREDUCE | Ops.STORE | Ops.END: @@ -872,12 +859,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass): if self.op is Ops.PARAM: return self.arg.addrspace if self.op is Ops.BUFFER: return self.arg.addrspace if self.op in {Ops.SPECIAL, Ops.RANGE}: return AddrSpace.ALU - if self.op in {Ops.LOAD, Ops.VLOAD}: return AddrSpace.ALU # LOAD brings things into the ALU + if self.op is Ops.LOAD: return AddrSpace.ALU # LOAD brings things into the ALU if self.op in {Ops.CUSTOM, Ops.CUSTOMI}: return AddrSpace.ALU if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.STORE, Ops.MSTACK, Ops.MSELECT, Ops.END, Ops.UNSHARD}: return self.src[0].addrspace if self.op in GroupOp.Movement: return self.src[0].addrspace - if self.op in {Ops.STACK, Ops.WMMA, Ops.DOT, Ops.UNPACK_LUT, Ops.GROUP} or self.op in GroupOp.Elementwise: + if self.op in {Ops.STACK, Ops.WMMA, Ops.GROUP} or self.op in GroupOp.Elementwise: ad = [x.addrspace for x in self.src if x.addrspace is not None] if not len(ad) or not all_same(ad): return None return ad[0] @@ -1245,7 +1232,7 @@ class ProgramInfo: for u in sink.toposort(): if u.op is Ops.PARAM and u.addrspace == AddrSpace.ALU: _vars.append(u) if u.op is Ops.PARAM and u.addrspace != AddrSpace.ALU: _globals.append(u.arg.slot) - if u.op in (Ops.STORE, Ops.LOAD, Ops.VLOAD): + if u.op in (Ops.STORE, Ops.LOAD): if (idx:=u.src[0]).op in (Ops.INDEX, Ops.SHRINK) or (u.src[0].op is Ops.CAST and (idx:=u.src[0].src[0]).op is Ops.INDEX): if (buf:=idx.src[0].buf_uop).op is Ops.PARAM: (outs if u.op is Ops.STORE else ins).append(buf.arg.slot) if u.op is Ops.SPECIAL: diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 566e05803c..da083f33e8 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -120,22 +120,11 @@ spec_shared = PatternMatcher([ lambda uidx,gate,alt,load: validate_index(uidx, gate) if matches_dtype(alt, load.dtype) else False), (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().store(UPat()), validate_index), (UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted().store(UPat(), UPat.var("gate", dtype=dtypes.bool)), validate_index), - # VLOAD reinterprets a scalar address as a contiguous vector of its output dtype - (UPat(Ops.VLOAD, src=(UPat((Ops.INDEX, Ops.SHRINK), name="uidx").or_casted(),), name="x"), - lambda x,uidx: isinstance(x.arg, int) and x.arg > 0 and validate_index(uidx)), # STORE in tensor graph: store a value into a target (UPat(Ops.STORE, dtypes.void, (UPat(name="x"), UPat())), lambda x: True), # WMMA has a (UPat(Ops.WMMA, src=(UPat(), UPat(), UPat()), name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) == 5), - # DOT reduces equal vector inputs in fixed-size groups - (UPat(Ops.DOT, src=(UPat(), UPat()), name="x"), lambda x: - isinstance(x.arg, int) and x.arg > 0 and x.src[0].shape == x.src[1].shape and len(x.src[0].shape) > 0 and - isinstance(x.src[0].shape[-1], int) and x.src[0].shape[-1] % x.arg == 0), - # UNPACK_LUT maps the low nibbles followed by the high nibbles through an immutable byte table - (UPat(Ops.UNPACK_LUT, dtypes.int8, src=(UPat(dtype=dtypes.uint8),), name="x"), lambda x: - len(x.src[0].shape) > 0 and isinstance(x.src[0].shape[-1], int) and isinstance(x.arg, tuple) and len(x.arg) == 16 and - all(isinstance(v, int) and -128 <= v <= 127 for v in x.arg)), ]) def is_device(d): return isinstance(d, str) or (isinstance(d, tuple) and all(isinstance(s, str) for s in d)) @@ -217,6 +206,11 @@ spec_program = PatternMatcher([ # allow special SHRINK (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.BUFFER, Ops.AFTER)), UPat(), UPat(Ops.CONST))), lambda: True), + # register-vector views can remain when the renderer consumes their shape (for example a native horizontal reduction) + (UPat((Ops.RESHAPE, Ops.PERMUTE), name="x"), lambda x: x.addrspace is AddrSpace.ALU), + (UPat(Ops.REDUCE, src=(UPat(),), name="x"), + lambda x: isinstance(x.arg, tuple) and len(x.arg) == 2 and x.arg[0] in GroupOp.Reduce and isinstance(x.arg[1], int)), + # movement ops are not allowed in programs (UPat(GroupOp.Movement), lambda: False),