diff --git a/test/backend/test_dtype.py b/test/backend/test_dtype.py index 7ce56ad3cd..27de82721c 100644 --- a/test/backend/test_dtype.py +++ b/test/backend/test_dtype.py @@ -328,7 +328,10 @@ class TestUint16DType(TestDType): class TestInt32DType(TestDType): DTYPE = dtypes.int32 class TestUint32DType(TestDType): DTYPE = dtypes.uint32 -class TestInt64DType(TestDType): DTYPE = dtypes.int64 +class TestInt64DType(TestDType): + DTYPE = dtypes.int64 + def test_int64_to_uint32_to_int64(self): + _test_op(lambda: Tensor([0x12345678ABCDEF01], dtype=dtypes.int64).cast(dtypes.uint32).cast(dtypes.int64), dtypes.int64, [2882400001]) @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs") class TestEmulatedInt64DType(TestInt64DType): diff --git a/tinygrad/codegen/late/regalloc.py b/tinygrad/codegen/late/regalloc.py index b43c934c84..611f2780c4 100644 --- a/tinygrad/codegen/late/regalloc.py +++ b/tinygrad/codegen/late/regalloc.py @@ -4,7 +4,7 @@ from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat from tinygrad.renderer.isa import ISARenderer, Register, greg from tinygrad.dtype import dtypes -PSEUDO_OPS = {Ops.CONST, Ops.CAST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK} +PSEUDO_OPS = {Ops.CONST, Ops.CAST, Ops.BITCAST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK} class LinearScanRegallocContext: # returns the uop that defines the virtual register diff --git a/tinygrad/renderer/isa/__init__.py b/tinygrad/renderer/isa/__init__.py index fb74f1569d..a9d1b38eea 100644 --- a/tinygrad/renderer/isa/__init__.py +++ b/tinygrad/renderer/isa/__init__.py @@ -27,7 +27,7 @@ class IselContext: return Register(f"v{next(self.reg_n)}", 0, _cons=cons if isinstance(cons, tuple) else (cons,)) def greg(u:UOp): - if u.op in {Ops.NOOP, Ops.AFTER} and u.src: return greg(u.src[0]) + if u.op in {Ops.NOOP, Ops.AFTER, Ops.BITCAST} and u.src: return greg(u.src[0]) if isinstance(u.tag, tuple): return u.tag[0] return u.tag diff --git a/tinygrad/renderer/isa/x86.py b/tinygrad/renderer/isa/x86.py index f359ae253f..256e9aacc3 100644 --- a/tinygrad/renderer/isa/x86.py +++ b/tinygrad/renderer/isa/x86.py @@ -182,15 +182,10 @@ def flag_gate(m:UOp) -> UOp|None: # legalize the new style graph for isel. NOTE: this runs after the spec is verified, some of these rewrites violate it pre_isel_matcher = PatternMatcher([ - # noop casts: zero extending scalar 32bit int, same-width signed/unsigned, narrowing scalar int - (UPat.var("y", dtypes.uint32).cast(dtypes.int64s, name="x"), lambda y,x: x.replace(op=Ops.NOOP, arg=None) if y.max_numel() == 1 else None), + # widening a scalar uint32 is free, the 32bit write that produced it already zeroed the upper half + (UPat.var("y", dtypes.uint32).cast(dtypes.int64s, name="x"), lambda y,x: x.replace(op=Ops.BITCAST) if y.max_numel() == 1 else None), (UPat.var("y", dtypes.ints+(dtypes.bool,)).cast(dtypes.ints, name="x"), - lambda y,x: x.replace(op=Ops.NOOP, arg=None) if x.dtype.itemsize == y.dtype.itemsize and y.max_numel() == 1 else None), - (UPat.var("y", dtypes.ints).cast(dtypes.ints, name="x"), - lambda y,x: x.replace(op=Ops.NOOP, arg=None) if x.dtype.itemsize < y.dtype.itemsize and y.max_numel() == 1 else None), - # bitcasts between scalar floats and ints are real, rest are noops - (UPat.var("y").bitcast().named("x"), lambda y,x: None if y.dtype in dtypes.floats and x.dtype in dtypes.ints or \ - y.dtype in dtypes.ints and x.dtype in dtypes.floats else x.replace(op=Ops.NOOP, arg=None)), + lambda y,x: x.replace(op=Ops.BITCAST) if x.dtype.itemsize == y.dtype.itemsize else None), # gated load/store become a conditional move on the address, the load/store are unconditional (UPat((Ops.INDEX, Ops.SHRINK), name="addr").load(UPat.var("alt"), UPat.var("gate"), name="x"), gated_load), (UPat((Ops.INDEX, Ops.SHRINK), name="addr").store(UPat.var("val"), UPat.var("gate")), gated_store), @@ -503,11 +498,12 @@ isel_matcher = PatternMatcher([ (UPat.var("y", dtypes.float64).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSD2SS, src=(y, y))), (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SS, src=(def_reg(x.dtype), y))), (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SD, src=(def_reg(x.dtype), y))), - (UPat(dtype=dtypes.uints+(dtypes.bool,)).cast(dtypes.ints, name="x"), lambda x: + (UPat(dtype=(dtypes.uint8, dtypes.uint16, dtypes.bool)).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOVZX) if x.max_numel() == 1 and x.src[0].dtype.itemsize < x.dtype.itemsize else None), (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.MOVSXD) if x.max_numel() == 1 else None), (UPat(dtype=dtypes.sints).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOVSX) if x.max_numel() == 1 and x.src[0].dtype.itemsize < x.dtype.itemsize else None), + (UPat(dtype=dtypes.ints).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOV) if x.max_numel() == 1 else None), (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBW)), (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBD)), (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBQ)),