tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 21:46:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

move max/min and add test

Browse Source
This commit is contained in:
ttomsa
2026-02-27 22:52:50 +00:00
parent ce2c690721
commit eb4ad1ebf0
2 changed files with 21 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/test_isel.py
+12
View File
@@ -18,6 +18,18 @@ class TestIselX86(unittest.TestCase):
# both comparisons become the same instruction
self.assertTrue(n.src[0].src[2] == n.src[1].src[2] and n.src[0].src[2].arg is X86Ops.CMP)
def test_vmax(self):
a = UOp.variable("a", 0, 0, dtypes.float32)
b = UOp.variable("b", 0, 0, dtypes.float32)
n = self.isel_rewrite((a < b).where(b, a))
self.assertTrue(n.arg is X86Ops.VMAXSS)
def test_vmin(self):
a = UOp.variable("a", 0, 0, dtypes.float32)
b = UOp.variable("b", 0, 0, dtypes.float32)
n = self.isel_rewrite((a < b).where(a, b))
self.assertTrue(n.arg is X86Ops.VMINSS)
def test_vpbroadcast(self):
a = UOp.variable("a", 0, 0, dtypes.int32)
n = self.isel_rewrite(a.broadcast(4))
tinygrad/renderer/isa/x86.py
+9 -9
View File
@@ -344,6 +344,15 @@ isel_matcher = PatternMatcher([
(UPat.cvar("x", dtypes.ints+(dtypes.bool,)), lambda x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, x.arg),))),
(UPat.cvar("x", dtypes.floats), lambda x:
UOp.const(dt:=to_int(x.dtype), struct.unpack(dt.fmt, struct.pack(x.dtype.fmt, x.arg))[0]).bitcast(x.dtype)),
# TODO: these should use a.maximum(b) / a.minimum(b)
((UPat.var("a") < UPat.var("b")).where(UPat.var("b", dtypes.float32), UPat.var("a")), lambda a,b:
a.ins(X86Ops.VMAXSS if a.dtype.count == 1 else X86Ops.VMAXPS, src=(a, b))),
((UPat.var("a") < UPat.var("b")).where(UPat.var("b", dtypes.float64), UPat.var("a")), lambda a,b:
a.ins(X86Ops.VMAXSD if a.dtype.count == 1 else X86Ops.VMAXPD, src=(a, b))),
((UPat.var("a") < UPat.var("b")).where(UPat.var("a", dtypes.float32), UPat.var("b")), lambda a,b:
a.ins(X86Ops.VMINSS if a.dtype.count == 1 else X86Ops.VMINPS, src=(a, b))),
((UPat.var("a") < UPat.var("b")).where(UPat.var("a", dtypes.float64), UPat.var("b")), lambda a,b:
a.ins(X86Ops.VMINSD if a.dtype.count == 1 else X86Ops.VMINPD, src=(a, b))),
# conditional moves that use masks NOTE: these currently assume a mask producing cmp exists
(UPat.var("m").where(UPat.var("a", dtypes.ints), UPat.var("b")), lambda m,a,b:
a.ins(X86Ops.VPBLENDVB, src=(b, a, m.replace(dtype=m.src[0].dtype))) if a.dtype.count > 1 else None),
@@ -456,15 +465,6 @@ isel_matcher = PatternMatcher([
(UPat(Ops.SUB, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VSUBSD if x.dtype.count == 1 else X86Ops.VSUBPD)),
(UPat(Ops.FDIV, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VDIVSS if x.dtype.count == 1 else X86Ops.VDIVPS)),
(UPat(Ops.FDIV, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VDIVSD if x.dtype.count == 1 else X86Ops.VDIVPD)),
# TODO: these should use a.maximum(b) / a.minimum(b)
((UPat.var("a", dtypes.float32) < UPat.var("b")).where(UPat.var("b"), UPat.var("a")), lambda a,b:
a.ins(X86Ops.VMAXSS if a.dtype.count == 1 else X86Ops.VMAXPS, src=(a, b))),
((UPat.var("a", dtypes.float64) < UPat.var("b")).where(UPat.var("b"), UPat.var("a")), lambda a,b:
a.ins(X86Ops.VMAXSD if a.dtype.count == 1 else X86Ops.VMAXPD, src=(a, b))),
((UPat.var("a", dtypes.float32) < UPat.var("b")).where(UPat.var("a"), UPat.var("b")), lambda a,b:
a.ins(X86Ops.VMINSS if a.dtype.count == 1 else X86Ops.VMINPS, src=(a, b))),
((UPat.var("a", dtypes.float64) < UPat.var("b")).where(UPat.var("a"), UPat.var("b")), lambda a,b:
a.ins(X86Ops.VMINSD if a.dtype.count == 1 else X86Ops.VMINPD, src=(a, b))),
# casts
(UPat(dtype=dtypes.int32).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTDQ2PS) if x.dtype.count > 1 else None),
(UPat(dtype=dtypes.int32).cast(dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VCVTDQ2PD) if x.dtype.count > 1 else None),