diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 38bc2b9e4f..c994145682 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -87,14 +87,14 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) - sink = graph_rewrite(sink, pm_lower_index_dtype+load_store_indexing, ctx=ren.device, name="lower all index dtypes") sink = graph_rewrite(sink, symbolic, name="post index symbolic") + # optional pre matcher + if ren.pre_matcher is not None: sink = graph_rewrite(sink, ren.pre_matcher, name="pre_matcher") + # decompositions supported_ops = tuple(ren.code_for_op.keys()) pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, TRANSCENDENTAL>=2) sink = graph_rewrite(sink, pm_decomp, ctx=ren.device, name="decompositions") - # optional pre matcher - if ren.pre_matcher is not None: sink = graph_rewrite(sink, ren.pre_matcher, ctx=ren, name="pre_matcher") - # final rules for the renderer (without sym) extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([]) pm_final_rewrite = pm_decomp+pm_render+extra_matcher+pm_split_ends @@ -128,7 +128,7 @@ def line_rewrite(lst:list[UOp], pm:PatternMatcher, ctx=None) -> list[UOp]: def do_linearize(prg:UOp, sink:UOp) -> UOp: lst = line_rewrite(linearize(sink), pm_linearize_cleanups) - #if SPEC: type_verify(lst, program_spec) + if SPEC: type_verify(lst, program_spec) return prg.replace(src=prg.src + (UOp(Ops.LINEAR, src=tuple(lst)),)) def do_render(ctx:Renderer, prg:UOp, lin:UOp) -> UOp: diff --git a/tinygrad/dtype.py b/tinygrad/dtype.py index 75142c5409..5e7c2fe8ec 100644 --- a/tinygrad/dtype.py +++ b/tinygrad/dtype.py @@ -124,8 +124,6 @@ class dtypes: @staticmethod def is_bool(x: DType) -> bool: return x.scalar() == dtypes.bool @staticmethod - def is_mask(x: DType) -> bool: return x.scalar() in dtypes.masks - @staticmethod def from_py(x) -> DType: if x.__class__ is float: return dtypes.default_float if x.__class__ is int: return dtypes.default_int @@ -162,11 +160,6 @@ class dtypes: def fields() -> dict[str, DType]: return DTYPES_DICT void: Final[DType] = DType.new(-1, 0, "void", None) index: Final[DType] = DType.new(-1,100, "index", None) - # mask dtypes are used in x86/arm64 backends - mask8: Final[DType] = DType.new(-1, 1, "mask8", None) - mask16: Final[DType] = DType.new(-1, 2, "mask16", None) - mask32: Final[DType] = DType.new(-1, 4, "mask32", None) - mask64: Final[DType] = DType.new(-1, 8, "mask64", None) bool: Final[DType] = DType.new(0, 1, "bool", '?') int8: Final[DType] = DType.new(1, 1, "signed char", 'b') uint8: Final[DType] = DType.new(2, 1, "unsigned char", 'B') @@ -200,7 +193,6 @@ class dtypes: fp8s = (fp8e4m3, fp8e5m2) floats = fp8s + (float16, bfloat16, float32, float64) - masks = (mask8, mask16, mask32, mask64) int8s = (uint8, int8) int16s = (uint16, int16) int32s = (uint32, int32) @@ -235,10 +227,8 @@ def least_upper_dtype(*ds:DType) -> DType: if not (images:=[d for d in ds if isinstance(d, ImageDType)]) else images[0] def least_upper_float(dt:DType) -> DType: return dt if dtypes.is_float(dt) else least_upper_dtype(dt, dtypes.default_float) -DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void", "index", "mask"))} -INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, - **{v.name:k for k,v in dtypes.__dict__.items() if isinstance(v, DType) and k.startswith("mask")}, - "void": "void", "index":"index"} +DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void", "index"))} +INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, "void": "void", "index":"index"} @functools.cache def can_lossless_cast(dt0:DType, dt1:DType) -> bool: diff --git a/tinygrad/mixin/math.py b/tinygrad/mixin/math.py index 91a031a408..ef30d883d7 100644 --- a/tinygrad/mixin/math.py +++ b/tinygrad/mixin/math.py @@ -31,7 +31,7 @@ class MathMixin: if (dtype := getattr(self, "dtype")) is not None: if isinstance(dtype, tuple): dtype = dtype[0] - if not (dtypes.is_bool(dtype) or dtypes.is_int(dtype) or dtypes.is_mask(dtype)): + if not (dtypes.is_bool(dtype) or dtypes.is_int(dtype)): raise RuntimeError(f"{dtype} is not supported") def add(self, x: Self | ConstType, reverse: bool = False): diff --git a/tinygrad/renderer/isa.py b/tinygrad/renderer/isa.py index e6d7800f84..4c2f66c271 100644 --- a/tinygrad/renderer/isa.py +++ b/tinygrad/renderer/isa.py @@ -33,7 +33,8 @@ class IselContext: self.stack_size += amt return ret - def vreg(self, cons:tuple[Register, ...]|Register): return Register(f"v{next(self.reg_n)}", 0, cons=cons if isinstance(cons, tuple) else (cons,)) + def vreg(self, cons:tuple[Register, ...]|Register|None=None): + return Register(f"v{next(self.reg_n)}", 0, cons=cons if isinstance(cons, tuple) else (cons,) if cons is not None else ()) isel_fixup = PatternMatcher([ # NOOP / AFTER have the same register as first src diff --git a/tinygrad/renderer/x86.py b/tinygrad/renderer/x86.py index 795f0fe5b9..dda9d13e96 100644 --- a/tinygrad/renderer/x86.py +++ b/tinygrad/renderer/x86.py @@ -3,52 +3,18 @@ from typing import cast from tinygrad.dtype import dtypes, PtrDType, DType, truncate from tinygrad.uop import Ops, X86Ops, GroupOp, X86GroupOp from tinygrad.uop.ops import UOp, UPat, PatternMatcher -from tinygrad.renderer import Renderer from tinygrad.uop.spec import x86_spec from tinygrad.renderer.isa import Register, ISARenderer, IselContext from tinygrad.codegen.late.regalloc import assign -# ***** X86 legalization matchers ***** +# ***** X86 legalization ***** -def to_mask(dt:DType): return {1:dtypes.mask8, 2:dtypes.mask16, 4:dtypes.mask32, 8:dtypes.mask64}[dt.scalar().itemsize].vec(dt.count) -def to_int(dt:DType): return {1:dtypes.int8, 2:dtypes.int16, 4:dtypes.int32, 8:dtypes.int64}[dt.scalar().itemsize].vec(dt.count) -# on x86/arm64 certain comparisons create masks instead of booleans -mask_matcher = PatternMatcher([ - # bool CMPNE is XOR, bool CMPEQ is XOR+XOR, bool CMPLT is XOR+AND, NOTE: cmp of masks is not valid for floats (true mask == nan) - (UPat.var('x', (dtypes.bool,)+dtypes.masks).ne(UPat.var('y')), lambda x,y: x^y), - (UPat.var('x', (dtypes.bool,)+dtypes.masks).alu(Ops.CMPEQ, UPat.var('y')), lambda x,y: (x^y)^True), - (UPat.var('x', (dtypes.bool,)+dtypes.masks) 1 else None), - # convert bools to masks in bitwise source - (UPat(GroupOp.Comparison | {Ops.AND, Ops.OR, Ops.XOR}, src=(UPat.var("a", dtypes.bool), UPat.var("b", dtypes.masks)), name="x"), - lambda a,b,x: x.replace(dtype=(dt:=to_mask(b.dtype)), src=(a.cast(to_int(dt)).mul(-1).bitcast(dt), b))), - (UPat(GroupOp.Comparison | {Ops.AND, Ops.OR, Ops.XOR}, src=(UPat.var("a", dtypes.masks), UPat.var("b", dtypes.bool)), name="x"), - lambda a,b,x: x.replace(dtype=(dt:=to_mask(a.dtype)), src=(a, b.cast(to_int(dt)).mul(-1).bitcast(dt)))), - # convert bool to mask in float/packed where - (UPat.var("m", dtypes.bool).where(UPat.var("a"), UPat.var("b")), - lambda m,a,b: m.cast(to_int(a.dtype)).mul(-1).bitcast(to_mask(a.dtype)).where(a, b) if dtypes.is_float(a.dtype) or a.dtype.count > 1 else None), - # convert mask to bool in scalar int where - (UPat.var("m", (dtypes.mask32, dtypes.mask64)).where(UPat.var("a", dtypes.ints), UPat.var("b")), - lambda m,a,b: m.bitcast(to_int(m.dtype)).cast(dtypes.bool).where(a, b) if a.dtype.count == 1 else None), - # cast mask to correct size in where - (UPat.var("m", dtypes.masks).where(UPat.var("a"), UPat.var("b")), lambda m,a,b: m.cast(to_mask(a.dtype)).where(a, b)), - # cast from mask is 1 if True, 0 if False - (UPat.var("y", dtypes.masks).cast(dtypes.ints, name="x"), lambda y,x: y.bitcast(x.dtype).mul(-1)), - (UPat.var("y", dtypes.masks).cast(dtypes.floats, name="x"), lambda y,x: y.where(x.const_like(1), x.const_like(0))), - # convert bool vectorize to mask if src is mask - (UPat(Ops.VECTORIZE, dtypes.bool, (UPat.var("y", dtypes.masks),), allow_any_len=True, name="x"), - lambda y,x: x.replace(dtype=y.dtype.vec(len(x.src)))), - # mask is converted to bool in store - (UPat.var("a").store(UPat.var("b", dtypes.masks), allow_any_len=True), - lambda a,b: a.store(b.bitcast(to_int(b.dtype)).mul(-1).cast(dtypes.int8).bitcast(dtypes.bool.vec(b.dtype.count)))), - # mask is converted to bool in index - (UPat.var("buf").index(UPat.var("idx"), UPat.var("m", dtypes.masks)), lambda buf,idx,m: buf.index(idx, m.bitcast(to_int(m.dtype)).ne(0), ptr=True)), -]) - -base_extra_matcher = PatternMatcher([ +extra_matcher = PatternMatcher([ + # bool CMPNE is XOR, bool CMPEQ is XOR+XOR, bool CMPLT is XOR+AND + # TODO: how does this work for vector dtypes? + (UPat.var('x', dtypes.bool).ne(UPat.var('y')), lambda x,y: x^y), + (UPat.var('x', dtypes.bool).alu(Ops.CMPEQ, UPat.var('y')), lambda x,y: (x^y)^True), + (UPat.var('x', dtypes.bool) !(y==x) (UPat(Ops.CMPNE, src=(UPat.var("y", dtypes.ints), UPat.var("x")), name="cmp"), lambda y,x,cmp: UOp(Ops.CMPEQ, cmp.dtype, (y,x))^True if y.dtype.count > 1 else None), -]) - -# TODO: this should be removed, vectors > max len shouldn't happen -powers_of_two = {2**i:i for i in range(64)} -def split_vectorized_alu(ctx:Renderer, alu:UOp): - dt = max([alu.src[-1].dtype, alu.dtype], key=lambda x: x.itemsize) - if dt.itemsize <= ctx.max_vec_sz and dt.count in powers_of_two: return None - szs, src, offset = [4,2,1], [], 0 - while offset < dt.count: - for sz in szs: - if sz*dt.scalar().itemsize > ctx.max_vec_sz or offset+sz > dt.count: continue - src.append(UOp(alu.op, alu.dtype.scalar().vec(sz), tuple(s.gep(tuple(range(offset, offset+sz))) for s in alu.src))) - offset += sz - break - return UOp(Ops.CAT, alu.dtype, tuple(src)) - -# TODO: handle tails, define reg probably shouldn't have a vector dtype -def split_vectorized_acc(ctx:Renderer, acc:UOp, c:UOp): - if acc.dtype.itemsize <= ctx.max_vec_sz and acc.dtype.count in powers_of_two: return None - l = next(x for x in [4,2,1] if acc.dtype.count % x == 0 and acc.dtype.base.scalar().vec(x).itemsize <= ctx.max_vec_sz) - new_acc = acc.replace(dtype=acc.dtype.base.scalar().vec(l).ptr(acc.dtype.count // l, cast(PtrDType, acc.dtype).addrspace)) - return UOp(Ops.PTRCAT, acc.dtype, tuple([new_acc.index(UOp.const(dtypes.int, i)) for i in range(0, acc.dtype.count, l)])) - -# patterns that change size (bool to mask, intermediate casts) need to run before vector splitting -# patterns that cast cmp/where to different dtypes (float16 where is casted to float32) need to run before mask patterns -# the mask matcher goes after cause splitting can result in a scalar tail and scalar int cmp is a bool not mask -# we want gep pushing but not through alus -from tinygrad.codegen.late.devectorizer import no_vectorized_alu, load_store_folding -from tinygrad.uop.symbolic import gep_pushing -x86_pre_matcher = PatternMatcher(gep_pushing.patterns[:-1]) + load_store_folding + x86_matcher + PatternMatcher([ - # TODO: try not to devectorize this - (UPat(dtype=dtypes.int64s).cast(dtypes.floats, name="alu"), no_vectorized_alu), - (UPat(dtype=dtypes.floats).cast(dtypes.int64s, name="alu"), no_vectorized_alu), - # TODO: use shuffle for these casts instead of devectorizing - (UPat(dtype=dtypes.int32s+(dtypes.mask32,)).cast(dtypes.int8s+dtypes.int16s+(dtypes.mask8,dtypes.mask16), name="alu"), no_vectorized_alu), - (UPat(dtype=dtypes.int16s+(dtypes.mask16,)).cast(dtypes.int8s+(dtypes.mask8,), name="alu"), no_vectorized_alu), - (UPat(Ops.SHR, dtypes.int64, name="alu"), no_vectorized_alu), - (UPat(Ops.MUL, dtypes.int64s, name="alu"), no_vectorized_alu), - (UPat(Ops.IDIV, name="alu"), no_vectorized_alu), - (UPat((*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.ASSIGN), name="alu"), split_vectorized_alu), - (UPat(Ops.DEFINE_REG, name="acc").index(UPat.cvar("c")), split_vectorized_acc), - # no narrowing int casts, shuffle instead, NOTE: this needs to be after split_vectorized_alu - (UPat.var("y", dtypes.int64s+(dtypes.mask64,)).cast(dtypes.int32s+(dtypes.mask32,), name="x"), lambda y,x: UOp(Ops.VECTORIZE, x.dtype, - tuple(y.bitcast(x.dtype.scalar().vec(x.dtype.count*2)).gep(i*2) for i in range(2))) if y.dtype.count > 1 else None), -]) + mask_matcher - -x86_extra_matcher = base_extra_matcher + PatternMatcher([ # noop of a noop is removed (UPat(Ops.NOOP, src=(UPat(Ops.NOOP),), name="x"), lambda x: x.replace(src=x.src[0].src)), # cast to < scalar int is a noop (UPat.var("y", dtypes.ints).cast(dtypes.ints, name="x"), lambda y,x: x.replace(op=Ops.NOOP) if x.dtype.itemsize < y.dtype.itemsize and y.dtype.count == 1 else None), - # if gate in scalar int cmove is not a comparison need to add one to set the flag - (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.ints), UPat.var("b")), - lambda m,a,b: m.ne(0).where(a,b) if m.op not in GroupOp.Comparison and a.dtype.count == 1 else None), + # float where expects a mask TODO: handle float64 cmp to float32 where + (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats), UPat.var("b")), + lambda m,a,b: m.cast(a.dtype).ne(0).where(a, b) if m.src[0].dtype not in dtypes.floats else None), # TODO: do we want this? Kinda not needed if DEVECTORIZE=0. If yes make it general (UPat(Ops.VECTORIZE, dtypes.float16, name="x"), lambda x: x.replace(dtype=dtypes.float32.vec(x.dtype.count), src=tuple(s.src[0] for s in x.src)).cast(x.dtype) if all(s.op is Ops.CAST for s in x.src) else None), + # moving elements of a single register to another without shuffling is a noop + (UPat(Ops.VECTORIZE, src=(UPat.var("y"),), allow_any_len=True, name="x"), + lambda y,x: UOp(Ops.NOOP, x.dtype, y.src) if all(s.op is Ops.GEP and s.src == y.src and s.arg[0] == i for i,s in enumerate(x.src)) else None), ]) # ***** X86 instruction selection pre matcher ***** @@ -179,7 +93,7 @@ pre_isel_matcher = PatternMatcher([ (UPat(Ops.STORE, src=(UPat.var("buf"), UPat.var("a")), name="x"), lambda buf,a,x: x.replace(src=(buf, UOp.const(dtypes.int32, 0), a))), # after extracting displacement cast idx to 64bit if it can be negative #(UPat.var("base").index(UPat.var("idx", dtypes.int32)), lambda base,idx: base.index(idx.cast(dtypes.int64), ptr=True) if idx.vmin < 0 else None), - # NOTE: shared with x86_extra_matcher + # TODO: remove this once we allow all flag producing ops in cmove # if gate in scalar int cmove is not a comparison need to add one to set the flag (UPat.var("m", dtypes.bool).where(UPat.var("a"), UPat.var("b")), lambda m,a,b: m.ne(0).where(a,b) if m.op not in GroupOp.Comparison and a.dtype.count == 1 else None), @@ -209,8 +123,11 @@ def to_imm(c:UOp) -> UOp|None: if c.dtype is dtypes.uint64: return imm(dtypes.uint32, c.arg) if not c.overflows(dtypes.uint32) else None if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, c.arg) return None +def cmp(x:UOp): + if x.src[0].dtype is dtypes.float32: return UOp(X86Ops.VUCOMISS, src=x.src) + if x.src[0].dtype is dtypes.float64: return UOp(X86Ops.VUCOMISD, src=x.src) + return UOp(X86Ops.CMP, src=x.src) if (i:=to_imm(x.src[1])) is None else UOp(X86Ops.CMPi, src=(x.src[0], i)) def disp(c:UOp) -> UOp: return imm(dtypes.int32 if c.overflows(dtypes.int8) else dtypes.int8, c.arg) -def cmp(x:UOp): return UOp(X86Ops.CMP, src=x.src) if (i:=to_imm(x.src[1])) is None else UOp(X86Ops.CMPi, src=(x.src[0], i)) def def_reg(dt:DType, reg:Register|None=None): return UOp(X86Ops.DEFINE_REG, dt, arg=reg) # vshufps takes 2 registers, it gets its lower 64 bits from the first register and its upper 64 bits from the second @@ -271,7 +188,7 @@ def abi(ctx:IselContext, x:UOp): if sys.platform == "win32": return x.replace(op=X86Ops.DEFINE_REG, arg=ctx.vreg(((RCX, RDX, GPR[8], GPR[9])[x.arg],))) if x.arg < 4 else _stack_arg((x.arg-3)*8+32) return x.replace(op=X86Ops.DEFINE_REG, arg=ctx.vreg(((RDI, RSI, RDX, RCX, GPR[8], GPR[9])[x.arg],))) if x.arg < 6 else _stack_arg((x.arg-5)*8) -dts = dtypes.ints + dtypes.masks + (dtypes.bool, dtypes.float16, dtypes.float32, dtypes.float64) +dts = dtypes.ints + (dtypes.bool, dtypes.float16, dtypes.float32, dtypes.float64) dt_16bit = tuple(dt.vec(l) for dt in dts for l in [2,1] if dt.vec(l).itemsize == 2 and dt.vec(l) not in dtypes.int16s) dt_32bit = tuple(dt.vec(l) for dt in dts for l in [4,2,1] if dt.vec(l).itemsize == 4 and dt.vec(l) not in dtypes.int32s) dt_64bit = tuple(dt.vec(l) for dt in dts for l in [8,4,2,1] if dt.vec(l).itemsize == 8 and dt.vec(l) not in dtypes.int64s) @@ -279,9 +196,9 @@ dt_128bit = tuple(dt.vec(l) for dt in dts for l in [16,8,4,2,1] if dt.vec(l).ite isel_matcher = PatternMatcher([ # **** Op rewrites **** - # TODO: add callee saved registers on windows to RET - # RET, add stack pointer to it. Also add add frame pointer, this makes it so the prologue and epilogue are automatically setup by the register allocator - (UPat(Ops.SINK, name="x"), lambda x: x.replace(op=X86Ops.RET, src=x.src + (UOp(X86Ops.DEFINE_REG, dtypes.uint64, arg=RSP),) + (UOp(X86Ops.DEFINE_REG, dtypes.uint64, arg=RBP),))), + # add callee saved registers to the RET, these will be scheduled at the top of the kernel and will be saved/restored if they are used in regalloc + # so regalloc builds the prologue/epilogue naturally + (UPat(Ops.SINK, name="x"), lambda x: x.replace(op=X86Ops.RET, src=x.src + tuple(def_reg(dtypes.uint64, r) for r in [RSP, RBP]))), # TODO: RANGE and END is tricky. Both linearizer and regalloc need them so they stay as Ops. This gets into a broader issue with tinygrad # not being able to represent control flow properly. For now they are rewritten after regalloc # HACK: annoying hack so const doesn't get rewritten because linearizer needs it @@ -291,7 +208,7 @@ isel_matcher = PatternMatcher([ # HACK: the register that holds the DEFINE_VAR is unknown until after linearizing, we add vreg to it that can't be allocated to any register # after linearizing we know the position of DEFINE_VAR in the function args and rewrite the vreg to the real reg # the right fix for this is to add the function arg position to DEFINE_VAR like DEFINE_GLOBAL - #(UPat(Ops.DEFINE_VAR, name="x")), + #(UPat(Ops.DEFINE_VAR, name="x"), lambda ctx,x: x.replace(arg=ctx.vreg())), # these are treated the same for now (UPat((Ops.DEFINE_REG, Ops.DEFINE_LOCAL), name="x"), lambda ctx,x: x.replace(op=X86Ops.LEA, src=(UOp(X86Ops.DEFINE_REG, x.dtype, arg=RSP), UOp(Ops.NOOP), imm(dtypes.int32, ctx.inc_stack(x.dtype.nbytes()))), arg=None)), # noqa: E501 @@ -306,49 +223,33 @@ isel_matcher = PatternMatcher([ (UPat(Ops.INDEX, src=(UPat.var("base"), UPat.var("idx") + UPat.cvar("dis")), name="x"), lambda base,idx,dis,x: x.replace(op=X86Ops.LEA, src=(base, idx.cast(dtypes.int64) if idx.vmin < 0 else idx, disp(dis.const_like(dis.arg * base.dtype.itemsize))))), (UPat(Ops.INDEX, src=(UPat.var("base"), UPat.cvar("dis")), name="x"), lambda base,dis,x: x.replace(op=X86Ops.LEA, src=(base, UOp(Ops.NOOP), disp(dis.const_like(dis.arg * base.dtype.itemsize))))), (UPat(Ops.INDEX, src=(UPat.var("base"), UPat.var("idx")), name="x"), lambda base,idx,x: x.replace(op=X86Ops.LEA, src=(base, idx.cast(dtypes.int64) if idx.vmin < 0 else idx, imm(dtypes.int8, 0)))), - # conditional moves that use flags (implicitly) - (UPat(Ops.CMPLT, dtypes.bool, (UPat(dtype=dtypes.uints), UPat()), name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: UOp(X86Ops.CMOVB, a.dtype, src=(b, a, cmp(m)))), # noqa: E501 - (UPat(Ops.CMPLT, dtypes.bool, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: UOp(X86Ops.CMOVL, a.dtype, src=(b, a, cmp(m)))), # noqa: E501 - (UPat(Ops.CMPEQ, dtypes.bool, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: UOp(X86Ops.CMOVE, a.dtype, src=(b, a, cmp(m)))), # noqa: E501 - (UPat(Ops.CMPNE, dtypes.bool, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: UOp(X86Ops.CMOVNE, a.dtype, src=(b, a, cmp(m)))), # noqa: E501 - # jumps + # jumps, use flags (UPat(Ops.IF, src=(UPat(Ops.CMPLT, dtypes.bool, (UPat(dtype=dtypes.uints), UPat()), name="y"),), name="x"), lambda y,x: UOp(X86Ops.JB, x.dtype, (cmp(y),))), # noqa: E501 (UPat(Ops.IF, src=(UPat(Ops.CMPLT, name="y"),)), lambda y: UOp(X86Ops.JL, src=(cmp(y),))), (UPat(Ops.IF, src=(UPat(Ops.CMPEQ, name="y"),)), lambda y: UOp(X86Ops.JE, src=(cmp(y),))), (UPat(Ops.IF, src=(UPat(Ops.CMPNE, name="y"),)), lambda y: UOp(X86Ops.JNE, src=(cmp(y),))), + # TODO: now how do you handle int cmp to float where? + # TODO: how do I deal with bitwise? + # answer: deal with them the same way, if int cmp or bitwise (bool) cast to int of float size, mul -1 and bitcast + # if float cmp and int where use ucomiss all otehr cases just use the vcmpss, convert to bool with bitcast -> and 1 -> noop bool + # conditional moves that use masks NOTE: these currently assume a mask producing cmp exists + (UPat(name="m").where(UPat.var("a", dtypes.ints), UPat.var("b")).named("x"), lambda m,a,b,x: x.replace(op=X86Ops.VPBLENDVB, src=(b, a, m.replace(dtype=m.src[0].dtype))) if x.dtype.count > 1 else None), + (UPat(name="m").where(UPat.var("a", dtypes.float32), UPat.var("b")).named("x"), lambda m,a,b,x: x.replace(op=X86Ops.VBLENDVPS, src=(b, a, m.replace(dtype=m.src[0].dtype)))), + (UPat(name="m").where(UPat.var("a", dtypes.float64), UPat.var("b")).named("x"), lambda m,a,b,x: x.replace(op=X86Ops.VBLENDVPD, src=(b, a, m.replace(dtype=m.src[0].dtype)))), + # in this case we have a mask producing comparison whose user expects a bool, so we convert to bool + (UPat(GroupOp.Comparison, dtypes.bool, (UPat(dtype=dtypes.float32), UPat()), name="x"), lambda x: x.replace(dtype=x.src[0].dtype).bitcast(dtypes.int32).bitwise_and(1).f(Ops.NOOP, dtype=dtypes.bool)), # noqa: E501 + (UPat(GroupOp.Comparison, dtypes.bool, (UPat(dtype=dtypes.float64), UPat()), name="x"), lambda x: x.replace(dtype=x.src[0].dtype).bitcast(dtypes.int64).bitwise_and(1).f(Ops.NOOP, dtype=dtypes.bool)), # noqa: E501 + # conditional moves that use flags + (UPat(Ops.CMPLT, src=(UPat(dtype=dtypes.sints), UPat()), name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: UOp(X86Ops.CMOVL, a.dtype, src=(b, a, cmp(m)))), # noqa: E501 + (UPat(Ops.CMPLT, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: UOp(X86Ops.CMOVB, a.dtype, src=(b, a, cmp(m)))), # noqa: E501 + (UPat(Ops.CMPEQ, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: UOp(X86Ops.CMOVE, a.dtype, src=(b, a, cmp(m)))), # noqa: E501 + (UPat(Ops.CMPNE, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: UOp(X86Ops.CMOVNE, a.dtype, src=(b, a, cmp(m)))), # noqa: E501 # comparisons whose user doesn't use the flag, move flag result to register (UPat(Ops.CMPLT, dtypes.bool, (UPat(dtype=dtypes.uints), UPat()), name="x"), lambda x: UOp(X86Ops.SETB, x.dtype, (cmp(x),))), (UPat(Ops.CMPLT, dtypes.bool, name="x"), lambda x: UOp(X86Ops.SETL, x.dtype, (cmp(x),))), (UPat(Ops.CMPEQ, dtypes.bool, name="x"), lambda x: UOp(X86Ops.SETE, x.dtype, (cmp(x),))), (UPat(Ops.CMPNE, dtypes.bool, name="x"), lambda x: UOp(X86Ops.SETNE, x.dtype, (cmp(x),))), - # float unary - (UPat.var("y", dtypes.float32).sqrt().named("x"), lambda y,x: UOp(X86Ops.VSQRTSS, x.dtype, (y, y)) if x.dtype.count == 1 else x.replace(op=X86Ops.VSQRTPS)), # noqa: E501 - (UPat.var("y", dtypes.float64).sqrt().named("x"), lambda y,x: UOp(X86Ops.VSQRTSD, x.dtype, (y, y)) if x.dtype.count == 1 else x.replace(op=X86Ops.VSQRTPD)), # noqa: E501 - (UPat.var("y", dtypes.float32).trunc().named("x"), lambda y,x: UOp(X86Ops.VROUNDSS, x.dtype, (y, y, imm(dtypes.uint8, 3))) if x.dtype.count == 1 else None), # noqa: E501 - (UPat.var("y", dtypes.float64).trunc().named("x"), lambda y,x: UOp(X86Ops.VROUNDSD, x.dtype, (y, y, imm(dtypes.uint8, 3))) if x.dtype.count == 1 else None), # noqa: E501 - (UPat.var("y", dtypes.float32).trunc().named("x"), lambda y,x: UOp(X86Ops.VROUNDPS, x.dtype, (y, imm(dtypes.uint8, 3)))), - (UPat.var("y", dtypes.float64).trunc().named("x"), lambda y,x: UOp(X86Ops.VROUNDPD, x.dtype, (y, imm(dtypes.uint8, 3)))), - # broadcasts TODO: not quite right, what about load fusion? Also, bitcast should be x86op and reg is xmm? - (UPat.var("y", dtypes.int8s+(dtypes.bool,)).broadcast(name="x"), lambda y,x: UOp(X86Ops.VPBROADCASTB, x.dtype, (y.bitcast(dtypes.float32),))), # noqa: E501 - (UPat.var("y", dtypes.int16s).broadcast(name="x"), lambda y,x: UOp(X86Ops.VPBROADCASTW, x.dtype, (y.bitcast(dtypes.float32),))), - (UPat.var("y", dtypes.int32s).broadcast(name="x"), lambda y,x: UOp(X86Ops.VPBROADCASTD, x.dtype, (y.bitcast(dtypes.float32),))), - (UPat.var("y", dtypes.int64s).broadcast(name="x"), lambda y,x: UOp(X86Ops.VPBROADCASTQ, x.dtype, (y.bitcast(dtypes.float64),))), - (UPat.var("y", dtypes.float32).broadcast(name="x"), lambda y,x: UOp(X86Ops.VBROADCASTSS, x.dtype, (y,))), - # shufles - (UPat.var("y", dtypes.int8s).bitcast(dtypes.mask8).named("x"), lambda y,x: UOp(X86Ops.VPINSRB, x.dtype, (def_reg(x.dtype), y, imm(dtypes.uint8, 0)))), - (UPat.var("y", dtypes.int16s).bitcast((dtypes.float16, dtypes.mask16)).named("x"), lambda y,x: UOp(X86Ops.VPINSRW, x.dtype, (def_reg(x.dtype), y, imm(dtypes.uint8, 0)))), # noqa: E501 - (UPat(Ops.VECTORIZE, dtypes.ints+(dtypes.bool,), name="x"), vpins), - (UPat(Ops.VECTORIZE, (dtypes.float32, dtypes.mask32), name="x"), vshufps), - (UPat(Ops.VECTORIZE, (dtypes.float32, dtypes.mask32), name="x"), vinsertps), - (UPat.var("y", dtypes.float32).gep(name="x"), lambda y,x: UOp(X86Ops.VINSERTPS, x.dtype, (y, y, imm(dtypes.uint8, x.arg[0] << 6)))), - # extract - (UPat.var("y", dtypes.mask8).bitcast(dtypes.int8s).named("x"), lambda y,x: UOp(X86Ops.VPEXTRB, x.dtype, (y, imm(dtypes.uint8, 0)))), - (UPat.var("y", (dtypes.float16, dtypes.mask16)).bitcast(dtypes.int16s).named("x"), lambda y,x: UOp(X86Ops.VPEXTRW, x.dtype, (y, imm(dtypes.uint8, 0)))), # noqa: E501 - (UPat.var("y", dtypes.int8s).gep(name="x"), lambda y,x: UOp(X86Ops.VPEXTRB, x.dtype, (y, imm(dtypes.uint8, x.arg[0])))), - (UPat.var("y", dtypes.int16s).gep(name="x"), lambda y,x: UOp(X86Ops.VPEXTRW, x.dtype, (y, imm(dtypes.uint8, x.arg[0])))), - (UPat.var("y", dtypes.int32s).gep(name="x"), lambda y,x: UOp(X86Ops.VPEXTRD, x.dtype, (y, imm(dtypes.uint8, x.arg[0])))), - (UPat.var("y", dtypes.int64s).gep(name="x"), lambda y,x: UOp(X86Ops.VPEXTRQ, x.dtype, (y, imm(dtypes.uint8, x.arg[0])))), - # comparisons that produce masks + # comparisons that produce masks (these aren't bool dtype) (UPat(Ops.CMPLT, src=(UPat(dtype=dtypes.float32), UPat()), name="x"), lambda x: x.replace(op=X86Ops.VCMPSS if x.dtype.count == 1 else X86Ops.VCMPPS, src=x.src + (imm(dtypes.uint8, 1),))), # noqa: E501 (UPat(Ops.CMPLT, src=(UPat(dtype=dtypes.float64), UPat()), name="x"), lambda x: x.replace(op=X86Ops.VCMPSD if x.dtype.count == 1 else X86Ops.VCMPPD, src=x.src + (imm(dtypes.uint8, 1),))), # noqa: E501 (UPat(Ops.CMPNE, src=(UPat(dtype=dtypes.float32), UPat()), name="x"), lambda x: x.replace(op=X86Ops.VCMPSS if x.dtype.count == 1 else X86Ops.VCMPPS, src=x.src + (imm(dtypes.uint8, 4),))), # noqa: E501 @@ -363,17 +264,38 @@ isel_matcher = PatternMatcher([ (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int16s), UPat.var("b")), name="x"), lambda a,b,x: x.replace(op=X86Ops.VPCMPGTW, src=(b, a))), (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int32s), UPat.var("b")), name="x"), lambda a,b,x: x.replace(op=X86Ops.VPCMPGTD, src=(b, a))), (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int64s), UPat.var("b")), name="x"), lambda a,b,x: x.replace(op=X86Ops.VPCMPGTQ, src=(b, a))), - # conditional moves that use masks - (UPat(name="m").where(UPat.var("a", dtypes.ints), UPat.var("b")).named("x"), lambda m,a,b,x: x.replace(op=X86Ops.VPBLENDVB, src=(b, a, m))), - (UPat(name="m").where(UPat.var("a", dtypes.float32), UPat.var("b")).named("x"), lambda m,a,b,x: x.replace(op=X86Ops.VBLENDVPS, src=(b, a, m))), - (UPat(name="m").where(UPat.var("a", dtypes.float64), UPat.var("b")).named("x"), lambda m,a,b,x: x.replace(op=X86Ops.VBLENDVPD, src=(b, a, m))), + # float unary + (UPat.var("y", dtypes.float32).sqrt().named("x"), lambda y,x: UOp(X86Ops.VSQRTSS, x.dtype, (y, y)) if x.dtype.count == 1 else x.replace(op=X86Ops.VSQRTPS)), # noqa: E501 + (UPat.var("y", dtypes.float64).sqrt().named("x"), lambda y,x: UOp(X86Ops.VSQRTSD, x.dtype, (y, y)) if x.dtype.count == 1 else x.replace(op=X86Ops.VSQRTPD)), # noqa: E501 + (UPat.var("y", dtypes.float32).trunc().named("x"), lambda y,x: UOp(X86Ops.VROUNDSS, x.dtype, (y, y, imm(dtypes.uint8, 3))) if x.dtype.count == 1 else None), # noqa: E501 + (UPat.var("y", dtypes.float64).trunc().named("x"), lambda y,x: UOp(X86Ops.VROUNDSD, x.dtype, (y, y, imm(dtypes.uint8, 3))) if x.dtype.count == 1 else None), # noqa: E501 + (UPat.var("y", dtypes.float32).trunc().named("x"), lambda y,x: UOp(X86Ops.VROUNDPS, x.dtype, (y, imm(dtypes.uint8, 3)))), + (UPat.var("y", dtypes.float64).trunc().named("x"), lambda y,x: UOp(X86Ops.VROUNDPD, x.dtype, (y, imm(dtypes.uint8, 3)))), + # broadcasts TODO: not quite right, what about load fusion? Also, bitcast should be x86op and reg is xmm? + (UPat.var("y", dtypes.int8s+(dtypes.bool,)).broadcast(name="x"), lambda y,x: UOp(X86Ops.VPBROADCASTB, x.dtype, (y.bitcast(dtypes.float32),))), # noqa: E501 + (UPat.var("y", dtypes.int16s).broadcast(name="x"), lambda y,x: UOp(X86Ops.VPBROADCASTW, x.dtype, (y.bitcast(dtypes.float32),))), + (UPat.var("y", dtypes.int32s).broadcast(name="x"), lambda y,x: UOp(X86Ops.VPBROADCASTD, x.dtype, (y.bitcast(dtypes.float32),))), + (UPat.var("y", dtypes.int64s).broadcast(name="x"), lambda y,x: UOp(X86Ops.VPBROADCASTQ, x.dtype, (y.bitcast(dtypes.float64),))), + (UPat.var("y", dtypes.float32).broadcast(name="x"), lambda y,x: UOp(X86Ops.VBROADCASTSS, x.dtype, (y,))), + # shufles + (UPat.var("y", dtypes.int16s).bitcast(dtypes.float16).named("x"), lambda y,x: UOp(X86Ops.VPINSRW, x.dtype, (def_reg(x.dtype), y, imm(dtypes.uint8, 0)))), # noqa: E501 + (UPat(Ops.VECTORIZE, dtypes.ints+(dtypes.bool,), name="x"), vpins), + (UPat(Ops.VECTORIZE, dtypes.float32, name="x"), vshufps), + (UPat(Ops.VECTORIZE, dtypes.float32, name="x"), vinsertps), + (UPat.var("y", dtypes.float32).gep(name="x"), lambda y,x: UOp(X86Ops.VINSERTPS, x.dtype, (y, y, imm(dtypes.uint8, x.arg[0] << 6)))), + # extract + (UPat.var("y", dtypes.float16).bitcast(dtypes.int16s).named("x"), lambda y,x: UOp(X86Ops.VPEXTRW, x.dtype, (y, imm(dtypes.uint8, 0)))), + (UPat.var("y", dtypes.int8s).gep(name="x"), lambda y,x: UOp(X86Ops.VPEXTRB, x.dtype, (y, imm(dtypes.uint8, x.arg[0])))), + (UPat.var("y", dtypes.int16s).gep(name="x"), lambda y,x: UOp(X86Ops.VPEXTRW, x.dtype, (y, imm(dtypes.uint8, x.arg[0])))), + (UPat.var("y", dtypes.int32s).gep(name="x"), lambda y,x: UOp(X86Ops.VPEXTRD, x.dtype, (y, imm(dtypes.uint8, x.arg[0])))), + (UPat.var("y", dtypes.int64s).gep(name="x"), lambda y,x: UOp(X86Ops.VPEXTRQ, x.dtype, (y, imm(dtypes.uint8, x.arg[0])))), # fused multiply add (UPat(Ops.MULACC, dtypes.float32, name="x"), lambda x: x.replace(op=X86Ops.VFMADD213SS if x.dtype.count == 1 else X86Ops.VFMADD213PS)), (UPat(Ops.MULACC, dtypes.float64, name="x"), lambda x: x.replace(op=X86Ops.VFMADD213SD if x.dtype.count == 1 else X86Ops.VFMADD213PD)), # packed bitwise - ((UPat() & UPat()).named("x"), lambda x: x.replace(op=X86Ops.VPAND) if x.dtype.count > 1 or x.dtype in dtypes.masks else None), - ((UPat() | UPat()).named("x"), lambda x: x.replace(op=X86Ops.VPOR) if x.dtype.count > 1 or x.dtype in dtypes.masks else None), - ((UPat() ^ UPat()).named("x"), lambda x: x.replace(op=X86Ops.VPXOR) if x.dtype.count > 1 or x.dtype in dtypes.masks else None), + ((UPat() & UPat()).named("x"), lambda x: x.replace(op=X86Ops.VPAND) if x.dtype.count > 1 else None), + ((UPat() | UPat()).named("x"), lambda x: x.replace(op=X86Ops.VPOR) if x.dtype.count > 1 else None), + ((UPat() ^ UPat()).named("x"), lambda x: x.replace(op=X86Ops.VPXOR) if x.dtype.count > 1 else None), # packed int binary ((UPat(dtype=dtypes.int32s) << UPat()).named("x"), lambda x: x.replace(op=X86Ops.VPSLLVD) if x.dtype.count > 1 else None), ((UPat(dtype=dtypes.int64s) << UPat()).named("x"), lambda x: x.replace(op=X86Ops.VPSLLVQ) if x.dtype.count > 1 else None), @@ -442,10 +364,10 @@ isel_matcher = PatternMatcher([ (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.replace(op=X86Ops.MOVSXD)), (UPat(dtype=dtypes.sints).cast(dtypes.ints, name="x"), lambda x: x.replace(op=X86Ops.MOVSX)), # bitcasts - (UPat(dtype=dtypes.int32s).bitcast((dtypes.float32, dtypes.mask32)).named("x"), lambda x: x.replace(op=X86Ops.VMOVD)), - (UPat(dtype=dtypes.int64s).bitcast((dtypes.float64, dtypes.mask64)).named("x"), lambda x: x.replace(op=X86Ops.VMOVQ)), - (UPat(dtype=(dtypes.float32, dtypes.mask32)).bitcast(dtypes.int32s).named("x"), lambda x: x.replace(op=X86Ops.VMOVDm)), - (UPat(dtype=(dtypes.float64, dtypes.mask64)).bitcast(dtypes.int64s).named("x"), lambda x: x.replace(op=X86Ops.VMOVQm)), + (UPat(dtype=dtypes.int32s).bitcast(dtypes.float32).named("x"), lambda x: x.replace(op=X86Ops.VMOVD)), + (UPat(dtype=dtypes.int64s).bitcast(dtypes.float64).named("x"), lambda x: x.replace(op=X86Ops.VMOVQ)), + (UPat(dtype=dtypes.float32).bitcast(dtypes.int32s).named("x"), lambda x: x.replace(op=X86Ops.VMOVDm)), + (UPat(dtype=dtypes.float64).bitcast(dtypes.int64s).named("x"), lambda x: x.replace(op=X86Ops.VMOVQm)), # TODO: fuse stores, very few cases -- store cmp becomes setcc, store gep int becomes vpextr, store bitcast to int becomes vmovd/q # assign, load, store # NOTE: assign here violates the spec, it only happens in register allocation when a reg to reg move needs to be inserted @@ -464,12 +386,12 @@ isel_matcher = PatternMatcher([ (UPat(Ops.STORE, src=(UPat(), UPat(), UPat(dtype=dt_16bit)), name="x"), lambda ctx,x: x.replace(op=X86Ops.VPEXTRW, src=fuse_index(ctx, x) + (x.src[-1], imm(dtypes.uint8, 0)))), # noqa: E501 (UPat(Ops.STORE, src=(UPat(), UPat(), UPat(dtype=dtypes.ints+(dtypes.bool,),)), name="x"), lambda ctx,x: x.replace(op=X86Ops.MOVm, src=fuse_index(ctx, x) + (x.src[-1],))), # noqa: E501 # **** X86Op rewrites **** - # allocate virtual register to X86Op, ones with specific constraints have already been allocated - (UPat(X86GroupOp.All, name="x"), lambda ctx,x: x.replace(arg=ctx.vreg(XMM if x.dtype in dtypes.floats+dtypes.masks or x.dtype.count > 1 else WGPR)) if x.arg is None and x.dtype != dtypes.void else None), # noqa: E501 # fuse loads into X86Ops that allow it, if beneficial (UPat(X86GroupOp.ReadMem1st, src=(UPat(Ops.LOAD),), allow_any_len=True, name="x"), lambda ctx,x: fuse_load(ctx, x, 0)), (UPat(X86GroupOp.ReadMem2nd, src=(UPat(), UPat(Ops.LOAD)), allow_any_len=True, name="x"), lambda ctx,x: fuse_load(ctx, x, 1)), (UPat(X86GroupOp.ReadMem3rd, src=(UPat(), UPat(), UPat(Ops.LOAD)), name="x"), lambda ctx,x: fuse_load(ctx, x, 2)), + # allocate virtual register to X86Op, ones with specific constraints have already been allocated + (UPat(X86GroupOp.All, name="x"), lambda ctx,x: x.replace(arg=ctx.vreg(XMM if x.dtype in dtypes.floats or x.dtype.count > 1 else WGPR)) if x.arg is None and x.dtype != dtypes.void else None), # noqa: E501 ]) # ***** post register allocation ***** @@ -667,6 +589,8 @@ encodings = PatternMatcher([ (UPat(X86Ops.VPSUBB, name="x"), lambda x: encode(x, 0xF8, pp=1, sel=1)), (UPat(X86Ops.VPSUBW, name="x"), lambda x: encode(x, 0xF9, pp=1, sel=1)), (UPat(X86Ops.VPSUBD, name="x"), lambda x: encode(x, 0xFA, pp=1, sel=1)), (UPat(X86Ops.VPSUBQ, name="x"), lambda x: encode(x, 0xFB, pp=1, sel=1)), (UPat(X86Ops.VPSRAVD, name="x"), lambda x: encode(x, 0x46, pp=1, sel=2)), + # float cmp + (UPat(X86Ops.VUCOMISS, name="x"), lambda x: encode(x, 0x2E, pp=0, sel=1)), (UPat(X86Ops.VUCOMISD, name="x"), lambda x: encode(x, 0x2E, pp=1, sel=1)), # scalar / packed float binary (UPat(X86Ops.VADDSS, name="x"), lambda x: encode(x, 0x58, pp=2, sel=1)), (UPat(X86Ops.VADDPS, name="x"), lambda x: encode(x, 0x58, pp=0, sel=1)), (UPat(X86Ops.VADDSD, name="x"), lambda x: encode(x, 0x58, pp=3, sel=1)), (UPat(X86Ops.VADDPD, name="x"), lambda x: encode(x, 0x58, pp=1, sel=1)), @@ -709,8 +633,7 @@ class X86Renderer(ISARenderer): max_vec_sz = 16 has_local = False global_max = None - pre_matcher = x86_pre_matcher - extra_matcher = x86_extra_matcher + extra_matcher = extra_matcher pre_isel_matcher = pre_isel_matcher isel_matcher = isel_matcher post_regalloc_matcher = post_regalloc_matcher diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index f25215d0d2..a4dda058c7 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -161,6 +161,7 @@ class X86Ops(FastEnum): # bitcasts VMOVD = auto(); VMOVQ = auto(); VMOVDm = auto(); VMOVQm = auto() # noqa: E702 # comparisons + VUCOMISS = auto(); VUCOMISD = auto() # noqa: E702 VCMPSS = auto(); VCMPSD = auto(); VCMPPS = auto(); VCMPPD = auto() # noqa: E702 VPCMPGTB = auto(); VPCMPGTW = auto(); VPCMPGTD = auto(); VPCMPGTQ = auto() # noqa: E702 VPCMPEQB = auto(); VPCMPEQW = auto(); VPCMPEQD = auto(); VPCMPEQQ = auto() # noqa: E702 @@ -231,7 +232,7 @@ class X86GroupOp: X86Ops.VPMULLW, X86Ops.VPMULLD, X86Ops.VROUNDSS, X86Ops.VROUNDSD, X86Ops.VSQRTSS, X86Ops.VSQRTSD, X86Ops.VSHUFPS, X86Ops.VINSERTPS, X86Ops.VPINSRB, X86Ops.VPINSRW, X86Ops.VPINSRD, X86Ops.VPINSRQ, X86Ops.VPAND, X86Ops.VPOR, X86Ops.VPXOR, X86Ops.VPSLLVD, X86Ops.VPSLLVQ, X86Ops.VPSRLVD, X86Ops.VPSRLVQ, X86Ops.VPSRAVD, X86Ops.VCVTSI2SS, X86Ops.VCVTSI2SD, X86Ops.VCVTSS2SD, X86Ops.VCVTSD2SS, - X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB} + X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB, X86Ops.VUCOMISS, X86Ops.VUCOMISD} # X86Ops whose third src can read from memory NOTE: these are TwoAddress1st so the third src is actually the second ReadMem3rd = {X86Ops.VFMADD213SS, X86Ops.VFMADD213SD, X86Ops.VFMADD213PS, X86Ops.VFMADD213PD} @@ -248,6 +249,6 @@ class X86GroupOp: # X86Ops that write flags or can modify flags to undefined values WriteFlags = {X86Ops.CMP, X86Ops.CMPi, X86Ops.ADD, X86Ops.ADDi, X86Ops.SUB, X86Ops.SUBi, X86Ops.IMUL, X86Ops.IMULi, X86Ops.IDIV, X86Ops.DIV, X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR, X86Ops.SARi, X86Ops.AND, X86Ops.ANDi, X86Ops.XOR, X86Ops.XORi, - X86Ops.OR, X86Ops.ORi} + X86Ops.OR, X86Ops.ORi, X86Ops.VUCOMISS, X86Ops.VUCOMISD} All = set(X86Ops) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 1dc672d100..6e64ec8bf4 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -18,6 +18,9 @@ from tinygrad.uop.validate import validate_index shared_spec = PatternMatcher([ (UPat(Ops.SINK, dtypes.void), lambda: True), # NOTE: for testing, we let sinks be anything + # NOOP + (UPat(Ops.NOOP), lambda: True), + # CONST/DEFINE_VAR are everywhere (UPat(Ops.CONST, src=(), name="x"), lambda x: type(x.arg) is type(dtypes.as_const(x.arg, x.dtype))), (UPat(Ops.DEFINE_VAR, name="x"), lambda x: isinstance(x.arg[1], int) and isinstance(x.arg[2], int)), @@ -279,6 +282,8 @@ x86_spec = PatternMatcher([ (UPat((Ops.NOOP, Ops.GROUP, Ops.AFTER, Ops.BARRIER)), lambda: True), (UPat(GroupOp.All), lambda: False), (UPat(X86GroupOp.All), lambda: True), + # vblends take mask which is float or int dtype + # cmove take flag producing instruction not just CMP ]) # ***** uop helpers *****