Change JIT reference tags to 2 bit scheme

Author: reidenong

Include/internal/pycore_optimizer.h

@@ -210,9 +210,12 @@ static inline uint16_t uop_get_error_target(const _PyUOpInstruction *inst)
 
 
 #define REF_IS_BORROWED 1
-#define REF_IS_INVALID  2
-#define REF_IS_UNIQUE   4
-#define REF_TAG_BITS    7
+#define REF_IS_UNIQUE   2
+#define REF_IS_INVALID  3
+#define REF_TAG_BITS    3
+
+#define REF_GET_TAG(x)   ((uintptr_t)(x) & (REF_TAG_BITS))
+#define REF_CLEAR_TAG(x) ((uintptr_t)(x) & (~REF_TAG_BITS))
 
 #define JIT_BITS_TO_PTR_MASKED(REF) ((JitOptSymbol *)(((REF).bits) & (~REF_TAG_BITS)))
 
@@ -234,38 +237,32 @@ PyJitRef_Wrap(JitOptSymbol *sym)
 static inline JitOptRef
 PyJitRef_WrapInvalid(void *ptr)
 {
-    return (JitOptRef){.bits=(uintptr_t)ptr | REF_IS_INVALID};
+    return (JitOptRef){.bits = REF_CLEAR_TAG((uintptr_t)ptr) | REF_IS_INVALID};
 }
 
 static inline bool
 PyJitRef_IsInvalid(JitOptRef ref)
 {
-    return (ref.bits & REF_IS_INVALID) == REF_IS_INVALID;
+    return REF_GET_TAG(ref.bits) == REF_IS_INVALID;
 }
 
 static inline JitOptRef
 PyJitRef_MakeUnique(JitOptRef ref)
 {
-    assert((ref.bits & REF_IS_UNIQUE) == 0);
-    return (JitOptRef){ ref.bits | REF_IS_UNIQUE };
-}
-
-static inline JitOptRef
-PyJitRef_RemoveUnique(JitOptRef ref)
-{
-    return (JitOptRef){ ref.bits & (~REF_IS_UNIQUE) };
+    return (JitOptRef){ REF_CLEAR_TAG(ref.bits) | REF_IS_UNIQUE };
 }
 
 static inline bool
 PyJitRef_IsUnique(JitOptRef ref)
 {
-    return (ref.bits & REF_IS_UNIQUE) == REF_IS_UNIQUE;
+    return REF_GET_TAG(ref.bits) == REF_IS_UNIQUE;
 }
 
 static inline JitOptRef
 PyJitRef_StripBorrowInfo(JitOptRef ref)
 {
-    return (JitOptRef){ .bits = ref.bits & ~(REF_IS_BORROWED | REF_IS_INVALID) };
+    if (PyJitRef_IsUnique(ref)) return ref;
+    return (JitOptRef){ .bits = REF_CLEAR_TAG(ref.bits) };
 }
 
 static inline JitOptRef
@@ -277,7 +274,7 @@ PyJitRef_StripReferenceInfo(JitOptRef ref)
 static inline JitOptRef
 PyJitRef_Borrow(JitOptRef ref)
 {
-    return (JitOptRef){ .bits = ref.bits | REF_IS_BORROWED };
+    return (JitOptRef){ .bits = REF_CLEAR_TAG(ref.bits) | REF_IS_BORROWED };
 }
 
 static const JitOptRef PyJitRef_NULL = {.bits = REF_IS_BORROWED};
@@ -291,7 +288,7 @@ PyJitRef_IsNull(JitOptRef ref)
 static inline int
 PyJitRef_IsBorrowed(JitOptRef ref)
 {
-    return (ref.bits & REF_IS_BORROWED) == REF_IS_BORROWED;
+    return REF_GET_TAG(ref.bits) == REF_IS_BORROWED;
 }
 
 extern bool _Py_uop_sym_is_null(JitOptRef sym);

Python/optimizer_bytecodes.c

@@ -89,24 +89,29 @@ dummy_func(void) {
         if (sym_is_null(value)) {
             ctx->done = true;
         }
-        value = PyJitRef_RemoveUnique(value);
+        if (PyJitRef_IsUnique(value)) {
+            value = PyJitRef_StripReferenceInfo(value);
+        }
     }
 
     op(_LOAD_FAST, (-- value)) {
         value = GETLOCAL(oparg);
-        value = PyJitRef_RemoveUnique(value);
+        if (PyJitRef_IsUnique(value)) {
+            value = PyJitRef_StripReferenceInfo(value);
+        }
     }
 
     op(_LOAD_FAST_BORROW, (-- value)) {
         value = PyJitRef_Borrow(GETLOCAL(oparg));
-        value = PyJitRef_RemoveUnique(value);
     }
 
     op(_LOAD_FAST_AND_CLEAR, (-- value)) {
         value = GETLOCAL(oparg);
         JitOptRef temp = sym_new_null(ctx);
         GETLOCAL(oparg) = temp;
-        value = PyJitRef_RemoveUnique(value);
+        if (PyJitRef_IsUnique(value)) {
+            value = PyJitRef_StripReferenceInfo(value);
+        }
     }
 
     op(_STORE_ATTR_INSTANCE_VALUE, (offset/1, value, owner -- o)) {
@@ -700,7 +705,7 @@ dummy_func(void) {
 
     op(_COPY, (bottom, unused[oparg-1] -- bottom, unused[oparg-1], top)) {
         assert(oparg > 0);
-        bottom = PyJitRef_RemoveUnique(bottom);
+        bottom = PyJitRef_IsUnique(bottom) ? PyJitRef_StripReferenceInfo(bottom) : bottom;
         top = bottom;
     }