constant fold classmethod and staticmethod in JIT

Author: kumaraditya303

Include/internal/pycore_opcode_metadata.h

@@ -3173,24 +3173,36 @@ def m(self):
         class E(Exception):
             def m(self):
                 return 1
+        class F:
+            @classmethod
+            def class_method(cls):
+                return 1
+            @staticmethod
+            def static_method():
+                return 1
+
         def f(n):
             x = 0
             c = C()
             d = D()
             e = E()
+            f = F()
             for _ in range(n):
                 x += C.A  # _LOAD_ATTR_CLASS
                 x += c.A  # _LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES
                 x += d.A  # _LOAD_ATTR_NONDESCRIPTOR_NO_DICT
                 x += c.m()  # _LOAD_ATTR_METHOD_WITH_VALUES
                 x += d.m()  # _LOAD_ATTR_METHOD_NO_DICT
                 x += e.m()  # _LOAD_ATTR_METHOD_LAZY_DICT
+                x += f.class_method()  # _LOAD_ATTR
+                x += f.static_method()  # _LOAD_ATTR
             return x
 
         res, ex = self._run_with_optimizer(f, TIER2_THRESHOLD)
-        self.assertEqual(res, 6 * TIER2_THRESHOLD)
+        self.assertEqual(res, 8 * TIER2_THRESHOLD)
         self.assertIsNotNone(ex)
         uops = get_opnames(ex)
+        self.assertNotIn("_LOAD_ATTR", uops)
         self.assertNotIn("_LOAD_ATTR_CLASS", uops)
         self.assertNotIn("_LOAD_ATTR_NONDESCRIPTOR_WITH_VALUES", uops)
         self.assertNotIn("_LOAD_ATTR_NONDESCRIPTOR_NO_DICT", uops)

Lib/test/test_capi/test_opt.py

@@ -2728,6 +2728,7 @@ dummy_func(
 
         macro(LOAD_ATTR) =
             _SPECIALIZE_LOAD_ATTR +
+            _RECORD_TOS_TYPE +
             unused/8 +
             _LOAD_ATTR;
 

Python/bytecodes.c

@@ -873,10 +873,42 @@ dummy_func(void) {
     }
 
     op(_LOAD_ATTR, (owner -- attr, self_or_null[oparg&1])) {
-        (void)owner;
-        attr = sym_new_not_null(ctx);
-        if (oparg & 1) {
-            self_or_null[0] = sym_new_unknown(ctx);
+        PyTypeObject *type = sym_get_probable_type(owner);
+        if (oparg & 1 && type != NULL) {
+            PyObject *name = get_co_name(ctx, oparg >> 1);
+            PyObject *descr = _PyType_Lookup(type, name);
+            bool class_method = descr && Py_IS_TYPE(descr, &PyClassMethod_Type);
+            bool static_method = descr && Py_IS_TYPE(descr, &PyStaticMethod_Type);
+            if (class_method || static_method) {
+                PyObject *callable = NULL;
+                if (class_method) {
+                    callable = _PyClassMethod_GetFunc(descr);
+                }
+                else {
+                    assert(static_method);
+                    callable = _PyStaticMethod_GetFunc(descr);
+                }
+                assert(callable);
+                bool immortal = _Py_IsImmortal(callable) || (type->tp_flags & Py_TPFLAGS_IMMUTABLETYPE);
+                ADD_OP(_GUARD_TYPE_VERSION, 0, type->tp_version_tag);
+                ADD_OP(_POP_TOP, 0, 0);
+                ADD_OP(immortal ? _LOAD_CONST_INLINE_BORROW : _LOAD_CONST_INLINE, 0, (uintptr_t)callable);
+                if (class_method) {
+                    ADD_OP(_LOAD_CONST_INLINE_BORROW, 0, (uintptr_t)type);
+                    self_or_null[0] = sym_new_const(ctx, (PyObject *)type);
+                } else if (static_method) {
+                    ADD_OP(_PUSH_NULL, 0, 0);
+                    self_or_null[0] = sym_new_null(ctx);
+                }
+                attr = sym_new_const(ctx, callable);
+                PyType_Watch(TYPE_WATCHER_ID, (PyObject *)type);
+                _Py_BloomFilter_Add(dependencies, (PyTypeObject *)type);
+            } else {
+                attr = sym_new_not_null(ctx);
+                self_or_null[0] = sym_new_unknown(ctx);
+            }
+        } else {
+            attr = sym_new_not_null(ctx);
         }
     }
 

Python/optimizer_bytecodes.c

@@ -1149,7 +1149,8 @@ def add_macro(
                             f"Recording uop {part.name} must be first in macro",
                             macro.tokens[0])
                     parts.append(uop)
-                    first = False
+                    if uop.properties.tier != 1:
+                        first = False
             case parser.CacheEffect():
                 parts.append(Skip(part.size))
             case _: