forked from tinygrad/tinygrad
remove Ops.VCAT and Ops.PTRCAT (glm) (#16771)
* remove Ops.VCAT and Ops.PTRCAT (glm) * fix stack, and remove useless casts
This commit is contained in:
George Hotz
+15
-15
@@ -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:
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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}
|
||||
|
||||
+2
-2
@@ -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
|
||||
|
||||
@@ -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) ****
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user