diff --git a/test/mockgpu/amd/emu.py b/test/mockgpu/amd/emu.py index 1110bf9472..35bafdc19b 100644 --- a/test/mockgpu/amd/emu.py +++ b/test/mockgpu/amd/emu.py @@ -375,7 +375,7 @@ def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32, """Conditional memory store with sub-word support. Returns list of store UOps.""" adt = dtypes.uint64 if addr_bits == 64 else dtypes.uint32 word_addr = addr >> UOp.const(adt, 2) - idx = mem.index(word_addr.cast(dtypes.int).valid(active)) + idx = mem.index(word_addr.valid(active)) if data_bits == 32: return [idx.store(active.where(_to_u32(val), idx))] # Sub-word store: read-modify-write with mask byte_pos = addr.cast(dtypes.uint32) & _c(3) @@ -388,7 +388,7 @@ def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32, is_cross = byte_pos.eq(_c(3)) cross_word0 = (idx & _c(0x00FFFFFF)) | ((val_u32 & _c(0xFF)) << _c(24)) store0 = idx.store(active.where(is_cross.where(cross_word0, new_word), idx)) - next_idx = mem.index((word_addr + UOp.const(adt, 1)).cast(dtypes.int).valid(active & is_cross)) + next_idx = mem.index((word_addr + UOp.const(adt, 1)).valid(active & is_cross)) cross_word1 = (next_idx & _c(0xFFFFFF00)) | ((val_u32 >> _c(8)) & _c(0xFF)) return [store0, next_idx.store((active & is_cross).where(cross_word1, next_idx))] @@ -398,7 +398,7 @@ def _mem_store_bytes(mem: UOp, addr: UOp, val: UOp, active: UOp, data_bits: int val_u32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val for i in range(data_bits // 8): byte_val = (val_u32 >> UOp.const(dtypes.uint32, i * 8)) & UOp.const(dtypes.uint32, 0xFF) - stores.append(mem.index((addr + UOp.const(dtypes.uint64, i)).cast(dtypes.int).valid(active)).store(byte_val.cast(dtypes.uint8))) + stores.append(mem.index((addr + UOp.const(dtypes.uint64, i)).valid(active)).store(byte_val.cast(dtypes.uint8))) return stores def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode_vars: dict | None = None, op_name: str = "") -> dict[int, UOp]: @@ -463,7 +463,7 @@ class _Ctx: """Read instruction dword from vmem at PC + dword_idx*4.""" pc = self.rpc() addr = pc if dword_idx == 0 else pc + UOp.const(dtypes.uint64, dword_idx * 4) - return self.vmem.index((addr >> UOp.const(dtypes.uint64, 2)).cast(dtypes.int), ptr=True).load() + return self.vmem.index(addr >> UOp.const(dtypes.uint64, 2), ptr=True).load() def inst_field(self, field) -> UOp: """Extract field bits from instruction encoding. Tracks field for canonical key computation.""" @@ -516,14 +516,14 @@ class _Ctx: # Dynamic register access (takes UOp index instead of int) def rsgpr_dyn(self, reg: UOp, valid: UOp | None = None) -> UOp: """Read SGPR with dynamic register index.""" - if valid is not None: return self.sgpr.index(reg.cast(dtypes.int).valid(valid), ptr=True).load() - return self.sgpr.index(reg.cast(dtypes.int), ptr=True).load() + if valid is not None: return self.sgpr.index(reg.valid(valid), ptr=True).load() + return self.sgpr.index(reg, ptr=True).load() def wsgpr_dyn(self, reg: UOp, val: UOp) -> UOp: """Write SGPR with dynamic register index. On RDNA, index 124 = NULL (writes discarded). On CDNA, index 124 = M0 (read/write).""" # RDNA: NULL (124) discards writes. CDNA: M0 (124) is writable. valid = None if self.wave_size == 64 else reg.ne(_c(124)) - return self.sgpr.index(reg.cast(dtypes.int).valid(valid) if valid is not None else reg.cast(dtypes.int), ptr=True).store(val.cast(dtypes.uint32)) + return self.sgpr.index(reg.valid(valid) if valid is not None else reg, ptr=True).store(val.cast(dtypes.uint32)) def wmask(self, reg: UOp, val: UOp) -> list[UOp]: """Write a lane mask (VCC/EXEC). Splits into lo/hi for wave64.""" @@ -656,7 +656,7 @@ class _Ctx: stores = [] for dest, val in assigns: if dest.startswith('D0'): stores.append(self.wsgpr_dyn(vdst_off, val.cast(dtypes.uint32))) - elif dest.startswith('VGPR['): stores.append(self.vgpr.index(val[0].cast(dtypes.int)).store(val[1].cast(dtypes.uint32))) + elif dest.startswith('VGPR['): stores.append(self.vgpr.index(val[0]).store(val[1].cast(dtypes.uint32))) return UOp.sink(*stores, *self.inc_pc()) def compile_vop_pcode(self, op, srcs: dict[str, UOp | int], lane: UOp, vdst_reg: UOp, exec_mask: UOp, @@ -711,11 +711,11 @@ class _Ctx: # VGPR bit-slice: (vgpr_idx, rhs_val, hi_bit, lo_bit) - hi/lo are UOp constants hi_bit, lo_bit = int(val[2].arg), int(val[3].arg) width = hi_bit - lo_bit + 1 - old = self.vgpr.index(val[0].cast(dtypes.int), ptr=True).load() + old = self.vgpr.index(val[0], ptr=True).load() new_val = _set_bits(old, _val_to_bits(val[1]), width, lo_bit).cast(dtypes.uint32) active = _lane_active(exec_mask, lane) if len(val) > 4: active = active & _to_bool(val[4]) - raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].cast(dtypes.int).valid(active)).store(new_val))) + raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].valid(active)).store(new_val))) continue if 'D0' in dest and '[laneId]' in dest: old_vcc = self.rmask(_c(VCC_LO.offset)) @@ -834,7 +834,7 @@ def _compile_smem(inst: ir3.SMEM | ir4.SMEM, ctx: _Ctx) -> UOp: nval = int(part.removeprefix('DWORD').removeprefix('X') or '1') if 'DWORD' in part else int(part[1:]) / 32 * (-1 if part[0] == 'I' else 1) ndwords = max(1, int(abs(nval))) dword_base = addr >> UOp.const(dtypes.uint64, 2) - vals = [ctx.vmem.index((dword_base + UOp.const(dtypes.uint64, i)).cast(dtypes.int)) for i in range(ndwords)] + vals = [ctx.vmem.index(dword_base + UOp.const(dtypes.uint64, i)) for i in range(ndwords)] if abs(nval) < 1: nbits = int(abs(nval) * 32) byte_off = (addr & UOp.const(dtypes.uint64, 3)).cast(dtypes.uint32) * UOp.const(dtypes.uint32, 8) @@ -1819,7 +1819,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA srcs = {'ADDR': addr_reg, 'DATA0': vdata_reg, 'VDST': vdst_reg, 'OFFSET': offset, 'EXEC': exec_mask.cast(dtypes.uint64), '_vgpr': ctx.vgpr, '_wave_size': ctx.wave_size} _, assigns = parse_pcode(pcode, srcs) - stores = [ctx.vgpr.index(val[0].cast(dtypes.int)).store(val[1].cast(dtypes.uint32)) for dest, val in assigns if dest.startswith('VGPR[')] + stores = [ctx.vgpr.index(val[0]).store(val[1].cast(dtypes.uint32)) for dest, val in assigns if dest.startswith('VGPR[')] return UOp.sink(*stores, *ctx.inc_pc()) def make_addr(lane: UOp) -> UOp: @@ -1853,7 +1853,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA def wmem(addr: UOp, val: UOp, active: UOp, data_bits: int = 32) -> UOp: if data_bits < 32: # Sub-dword LDS write: read-modify-write within the uint32 slot - word_addr = (addr >> addr_shift).cast(dtypes.int) + word_addr = addr >> addr_shift idx = mem.index(word_addr.valid(active)) byte_pos = addr.cast(dtypes.uint32) & _c(3) byte_shift = byte_pos * _c(8) @@ -1861,7 +1861,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA mask = size_mask << byte_shift new_word = (idx & (mask ^ _c(0xFFFFFFFF))) | ((val.cast(dtypes.uint32) & size_mask) << byte_shift) return idx.store(active.where(new_word, idx)) - idx = mem.index((addr >> addr_shift).cast(dtypes.int)) + idx = mem.index(addr >> addr_shift) return idx.store(active.where(val, idx.load())) def make_srcs(lane: UOp) -> dict: @@ -2011,7 +2011,7 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp: for i in range(n_dwords): word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2) val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64), ptr=True).load(), _c(0)) - lds_idx = ((lds_addr + _c(i * 4)) >> _c(2)).cast(dtypes.int) + lds_idx = (lds_addr + _c(i * 4)) >> _c(2) lds_slot = ctx.lds.index(lds_idx.valid(active)) stores.append(lds_slot.store(active.where(val, lds_slot))) elif is_store: diff --git a/test/mockgpu/amd/pcode.py b/test/mockgpu/amd/pcode.py index 05e41d5f03..c892090f8a 100644 --- a/test/mockgpu/amd/pcode.py +++ b/test/mockgpu/amd/pcode.py @@ -831,7 +831,7 @@ class Parser: def mindex(idx:UOp, ptr=False): return mem.index(idx.valid(active) if active is not None else idx, ptr=ptr) byte_mem = mem.dtype.base == dtypes.uint8 if byte_mem: - idx = addr.cast(dtypes.int) + idx = addr if dt in (dtypes.uint64, dtypes.int64, dtypes.float64): val = _u32(0).cast(dtypes.uint64) for i in range(8): val = val | (mindex(idx + _const(dtypes.int, i), ptr=True).load().cast(dtypes.uint64) << _u64(i * 8)) @@ -845,10 +845,10 @@ class Parser: val = _u32(0) for i in range(4): val = val | (mindex(idx + _const(dtypes.int, i), ptr=True).load().cast(dtypes.uint32) << _u32(i * 8)) else: - idx = (addr >> _const(addr.dtype, 2)).cast(dtypes.int) + idx = addr >> _const(addr.dtype, 2) val = mindex(idx) if dt in (dtypes.uint64, dtypes.int64, dtypes.float64): - idx2 = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int) + idx2 = (addr + _const(adt, 4)) >> _const(adt, 2) val = val.cast(dtypes.uint64) | (mindex(idx2).cast(dtypes.uint64) << _u64(32)) elif dt in (dtypes.uint8, dtypes.int8): val = (val >> ((addr & _const(adt, 3)).cast(dtypes.uint32) * _u32(8))) & _u32(0xFF) elif dt in (dtypes.uint16, dtypes.int16): diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index a60d708f39..11bca61e5f 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -790,26 +790,6 @@ class TestLoadStoreFolding(unittest.TestCase): self.assertEqual(inner_load.op, Ops.LOAD) self.assertEqual(len(inner_load.src), 2) # INDEX + alt - def test_gated_load_ptrcat_preserves_alt(self): - """Test that LOAD(PTRCAT, alt) preserves alt value after rewrite""" - from tinygrad.codegen.late.devectorizer import load_store_folding - buf1 = UOp.param(0, dtypes.float.ptr()) - buf2 = UOp.param(1, dtypes.float.ptr()) - idx = UOp.const(dtypes.int, 0) - idx1 = buf1.index(idx) - idx2 = buf2.index(idx) - ptrcat = UOp(Ops.PTRCAT, dtypes.float.ptr().vec(2), (idx1, idx2)) - alt = UOp.const(dtypes.float.vec(2), 42.0) - gated_load = ptrcat.load(alt) - self.assertEqual(len(gated_load.src), 2) # PTRCAT + alt - result = graph_rewrite(gated_load, load_store_folding, name='test') - # After rewrite, should be CAT of LOADs, each preserving alt - self.assertEqual(result.op, Ops.VCAT) - for inner_load in result.src: - self.assertEqual(inner_load.op, Ops.LOAD) - self.assertEqual(len(inner_load.src), 2) # INDEX + alt - self.assertEqual(inner_load.src[1].arg, 42.0) # alt value preserved - class TestConstBufferize(unittest.TestCase): def test_const_bufferize_with_ranges(self): """Test that CONST.BUFFERIZE with ranges is folded correctly. diff --git a/tinygrad/codegen/late/devectorizer.py b/tinygrad/codegen/late/devectorizer.py index 651c5e4312..ae5a63c5e1 100644 --- a/tinygrad/codegen/late/devectorizer.py +++ b/tinygrad/codegen/late/devectorizer.py @@ -71,22 +71,6 @@ def expand_index(ctx, buf:UOp, vec:UOp): # generate the individual indexes return UOp(Ops.STACK, buf.dtype, tuple(buf.index(vec.gep(i), ptr=True) for i in range(vec.dtype.count))) -def fold_expanded_index(midx:UOp): - buf = midx.src[0].src[0] - if not all(s.src[0] is buf for s in midx.src): return None - if not all(isinstance(s.dtype, PtrDType) for s in midx.src): return None - if len(midx.src) == 1: return midx.src[0] - return UOp(Ops.PTRCAT, buf.ptrdtype.base.ptr(size=buf.max_numel(), addrspace=buf.addrspace).vec(len(midx.src)), midx.src) - -def cat_after_store(cat:UOp, data:UOp): - # TODO: this is written in many places - offset = 0 - ret: list[UOp] = [] - for s in cat.src: - ret.append(s.store(data.gep(tuple(range(offset, offset+s.dtype.count))))) - offset += s.dtype.count - return UOp.group(*ret) - def gep_on_store(gep:UOp, st:UOp): # NOTE: we need to invert the gep here, but it may be an expanding gep # fake argsort. TODO: handle duplicates @@ -95,19 +79,33 @@ def gep_on_store(gep:UOp, st:UOp): new_arg = tuple(x[1] for x in sorted(a.items())) return gep.src[0].store(st.gep(new_arg)) +def load_stack(stack:UOp, ld:UOp): + offset, ret = 0, [] + for x in stack.src: + src = [x] + for s in ld.src[1:]: + src.append(s.gep(tuple(range(offset, offset+x.dtype.count))) if s.dtype.vcount > 1 else s) + ret.append(ld.replace(dtype=x.dtype.base, src=tuple(src))) + offset += x.dtype.count + return UOp(Ops.STACK, stack.dtype.base.vec(len(stack.src)), tuple(ret)) + +def store_stack(stack:UOp, data:UOp): + offset, ret = 0, [] + for x in stack.src: + ret.append(x.store(data.gep(tuple(range(offset, offset+x.dtype.count))))) + offset += x.dtype.count + return UOp.group(*ret) + load_store_folding = PatternMatcher([ (UPat(Ops.INDEX, src=(UPat(Ops.STACK, src=UPat(name="buf")), UPat.var("vec"))), expand_index), - (UPat(Ops.STACK, src=UPat(Ops.INDEX), name="midx"), fold_expanded_index), + # put STACK of indexes after LOAD/STORE + (UPat(Ops.LOAD, src=(UPat(Ops.STACK, src=UPat(Ops.INDEX), name="stack"),), name="ld", allow_any_len=True), load_stack), # GEP after LOAD (UPat(Ops.LOAD, src=(UPat(Ops.GEP, name="gep"),), name="ld", allow_any_len=True), lambda gep, ld: ld.replace(dtype=ld.dtype.scalar().vec(gep.dtype.count), src=(gep.src[0],)+ld.src[1:]).gep(gep.arg)), # GEP on data of STORE (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st"))), gep_on_store), - # put PTRCAT after LOAD - (UPat(Ops.LOAD, src=(UPat(Ops.PTRCAT, name="cat"),), name="ld", allow_any_len=True), - lambda cat,ld: UOp(Ops.VCAT, cat.dtype.base.vec(cat.dtype.vcount), tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))), - # put PTRCAT after STORE - (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data"))), cat_after_store), + (UPat(Ops.STORE, src=(UPat(Ops.STACK, src=UPat(Ops.INDEX), name="stack"), UPat(name="data"))), store_stack), ]) # *** correct load/store *** @@ -123,7 +121,7 @@ def split_load_store(ctx, ls:UOp, idx:UOp): lidx = buf.index((offset + i).valid(mask), ptr=True) if ls.op is Ops.STORE: ret.append(ls.replace(src=(lidx, ls.src[1].gep(i)))) else: ret.append(ls.replace(src=(lidx,)+ls.src[1:], dtype=ls.dtype.scalar())) - return UOp(Ops.VCAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret) + return UOp(Ops.STACK, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret) correct_load_store = PatternMatcher([ # split LOAD/STORE diff --git a/tinygrad/codegen/late/expander.py b/tinygrad/codegen/late/expander.py index 8308034ff3..72a1057532 100644 --- a/tinygrad/codegen/late/expander.py +++ b/tinygrad/codegen/late/expander.py @@ -51,7 +51,8 @@ def do_expand(root:UOp): new_srcs.append(src) elif src.dtype.count > 1: # put any input dtype > 1 grouped together - new_srcs.append(UOp(Ops.VCAT, src.dtype.scalar().vec(expand_sz*src.dtype.count), (src,)*expand_sz)) + new_srcs.append(UOp(Ops.STACK, src.dtype.scalar().vec(expand_sz*src.dtype.count), + tuple(src.gep(i) for _ in range(expand_sz) for i in range(src.dtype.count)))) else: # repeat the arg new_srcs.append(src.broadcast(expand_sz)) diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index 37fd0a1f59..4b91bac4c9 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -106,7 +106,7 @@ class Ops(FastEnum): REDUCE = auto(); ALLREDUCE = auto() # expander ops - UNROLL = auto(); CONTRACT = auto(); VCAT = auto(); PTRCAT = auto() + UNROLL = auto(); CONTRACT = auto() class GroupOp: Unary = {Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.SQRT, Ops.RECIPROCAL, Ops.NEG, Ops.TRUNC} diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index b4943b7141..e6a83559d0 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -227,7 +227,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): match self.op: # late ops don't have shape case Ops.DEVICE | Ops.IF | Ops.BARRIER | Ops.CUSTOM | Ops.CUSTOMI | \ - Ops.SINK | Ops.REWRITE_ERROR | Ops.PTRCAT | Ops.ENDIF | \ + Ops.SINK | Ops.REWRITE_ERROR | Ops.ENDIF | \ Ops.LINEAR | Ops.PROGRAM | Ops.SOURCE | Ops.INS | Ops.TUPLE | Ops.CALL | Ops.FUNCTION: return None @@ -272,7 +272,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): else: return (len(self.src),) + self.src[0].shape # TODO: contract and unroll should be deleted - case Ops.CONST | Ops.CONTRACT | Ops.UNROLL | Ops.VCAT: + case Ops.CONST | Ops.CONTRACT | Ops.UNROLL: return (self.dtype.count,) if self.dtype.count > 1 else () # some ops init the shape diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 2e10a41d26..45d54ca6a7 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -230,7 +230,7 @@ spec_full = PatternMatcher([ # allow any AFTER (UPat(Ops.AFTER, src=(UPat(),), allow_any_len=True), lambda: True), - # expander: unroll/contract/gep/ptrcat/cat + # expander: unroll/contract/gep/cat (UPat((Ops.UNROLL, Ops.CONTRACT), src=(UPat(),)), lambda: True), # GEP multi is supported here @@ -241,13 +241,6 @@ spec_full = PatternMatcher([ # while BIND is being casted (UPat(Ops.BIND, (dtypes.int, dtypes.weakint), (UPat(), UPat()), arg=None), lambda: True), - - # TODO: PTRCAT and VCAT need to be deleted - - # PTRCAT is like VECTORIZE, but it functions on ptrs - (UPat(Ops.PTRCAT, name="x"), lambda x: x.dtype.vcount == sum([y.dtype.base.count for y in x.src])), - # VCAT is like VECTORIZE, but the srcs can be vectors - (UPat(Ops.VCAT, name="x"), lambda x: x.dtype.vcount == sum([y.dtype.vcount for y in x.src])), ])+spec_tensor+spec_program # **** pyrender (move this) **** diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 3d533514a1..8e83ce3b77 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -229,9 +229,6 @@ gep_pushing = PatternMatcher([ (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST), name='alu').f(Ops.GEP, dtype=dtypes.weakint, name='gep'), lambda gep,alu: UOp(alu.op, alu.dtype.scalar().vec(gep.dtype.count), tuple(x.gep(gep.arg) for x in alu.src), alu.arg) \ if not isinstance(gep.dtype, PtrDType) and not isinstance(alu.dtype, PtrDType) else None), - # CAT can't be rendered. it's a VECTORIZE on vectors, we expand to a single VECTORIZEs with GEPs (TODO: move this later) - (UPat(Ops.VCAT, name="x"), lambda x: UOp(Ops.STACK, x.dtype, tuple(y.gep(i) for y in x.src for i in range(y.dtype.count))) \ - if not isinstance(x.dtype, PtrDType) else None), # VECTORIZE on same GEP (UPat(Ops.STACK, name="v", src=UPat(Ops.GEP, src=(UPat.var("x"),))), lambda v,x: x.gep(tuple(get_single_element(i.arg) for i in v.src))), # push some GEPs through WMMAs