Add caching for class vars

Summary: Let's leverage the slower fallback method for class vars too!

Reviewed By: alexmalyshev

Differential Revision: D109748680

fbshipit-source-id: 3ba2cb4f03652ad8055c517db750d968016a5b99

Author: DinoV

cinderx/Jit/inline_cache.cpp

@@ -110,6 +110,25 @@ bool loadMethodValueIsUnbound(uintptr_t bits) {
   return (bits & kLoadMethodUnboundTag) != 0;
 }
 
+// For a class-method descriptor `descr` found on a type, returns the callable
+// to cache (which lookup() binds to the receiver's type), or nullptr if it
+// isn't safe to cache.
+BorrowedRef<> classMethodCacheableCallable(BorrowedRef<> descr) {
+  if (Py_TYPE(descr) == &PyClassMethodDescr_Type) {
+    // A C-level classmethod descriptor (e.g. dict.fromkeys). It is immutable
+    // and directly callable with the bound type as its first argument, so cache
+    // the descriptor itself.
+    return descr;
+  }
+  // A Python-level classmethod. Only cache when it wraps a plain function;
+  // other callables may run arbitrary user code when bound.
+  BorrowedRef<> callable = Ci_PyClassMethod_GetFunc(descr);
+  if (Py_TYPE(callable) == &PyFunction_Type) {
+    return callable;
+  }
+  return nullptr;
+}
+
 // Sentinel PyTypeObject that must never escape into user code.
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wmissing-field-initializers"
@@ -1445,6 +1464,12 @@ LoadMethodResult LoadMethodCache::lookup(
         // Bound method (common case): the low bit is clear, so the value is an
         // untagged PyObject* which is a method-like object.
         return {Py_NewRef(reinterpret_cast<PyObject*>(value)), Py_NewRef(obj)};
+      } else if (entry.is_class_method) {
+        // Class method: the (tagged) value is the underlying callable. Bind it
+        // to the receiver's type rather than the receiver itself.
+        return {
+            Py_NewRef(loadMethodValuePtr(value)),
+            Py_NewRef(reinterpret_cast<PyObject*>(tp.get()))};
       } else if (value != kLoadMethodGetAttrSentinel) {
         // A tagged pointer (value > the tag bit): a staticmethod or class
         // variable. Untag it and return it as a plain attribute without binding
@@ -1510,6 +1535,7 @@ LoadMethodResult __attribute__((noinline)) LoadMethodCache::lookupSlowPath(
   PyObject* attr;
   bool is_method = false;
   bool is_static_method = false;
+  bool is_class_method = false;
 
   // A type with a __getattr__ hook is cacheable as long as the hook wraps the
   // generic getattr we replicate below: the type-dict lookup is authoritative,
@@ -1547,6 +1573,16 @@ LoadMethodResult __attribute__((noinline)) LoadMethodCache::lookupSlowPath(
       // instance attribute, so defer caching until after the instance dict
       // check below.
       is_static_method = true;
+    } else if (
+        Py_TYPE(descr) == &PyClassMethod_Type ||
+        Py_TYPE(descr) == &PyClassMethodDescr_Type) {
+      // A class method (Python-level classmethod or C-level
+      // classmethod_descriptor). Both are non-data descriptors that can be
+      // shadowed by an instance attribute, so defer caching until after the
+      // instance dict check below. `f` is set so we can still dispatch through
+      // the descriptor if it turns out not to be cacheable.
+      is_class_method = true;
+      f = descr->ob_type->tp_descr_get;
     } else {
       f = descr->ob_type->tp_descr_get;
       if (f != nullptr && PyDescr_IsData(descr)) {
@@ -1596,6 +1632,27 @@ LoadMethodResult __attribute__((noinline)) LoadMethodCache::lookupSlowPath(
     return {Py_None, Py_NewRef(callable)};
   }
 
+  if (is_class_method) {
+    // The class method was found in the type dict and is not shadowed by an
+    // instance attribute. If it is safe to cache, cache the underlying callable
+    // and return it bound to the type (class methods bind to the type, not the
+    // instance). Otherwise fall through to the generic descriptor dispatch.
+    BorrowedRef<> callable = classMethodCacheableCallable(descr);
+    if (callable != nullptr) {
+      fill(
+          tp,
+          callable,
+          name,
+          /*has_getattr_hook=*/false,
+          /*is_bound_method=*/false,
+          /*is_class_method=*/true);
+      LoadMethodResult result = {
+          Py_NewRef(callable), Py_NewRef(reinterpret_cast<PyObject*>(tp))};
+      Py_DECREF(descr);
+      return result;
+    }
+  }
+
   if (f != nullptr) {
     maybeCollectCacheStats(
         cache_stats_, tp, name, CacheMissReason::kUncategorized);
@@ -1635,7 +1692,8 @@ void LoadMethodCache::fill(
     BorrowedRef<> value,
     BorrowedRef<> name,
     bool has_getattr_hook,
-    bool is_bound_method) {
+    bool is_bound_method,
+    bool is_class_method) {
   if (!Ci_Type_HasValidVersionTag(type)) {
     // The type must have a valid version tag in order for us to be able to
     // invalidate the cache when the type is modified. See the comment at
@@ -1660,6 +1718,7 @@ void LoadMethodCache::fill(
       entry.value = tagLoadMethodValue(value, is_bound_method);
       entry.keys_version = keys_version;
       entry.has_getattr_hook = has_getattr_hook;
+      entry.is_class_method = is_class_method;
       return;
     }
   }

cinderx/Jit/inline_cache.h

@@ -308,7 +308,7 @@ class LoadMethodCache {
     //     and must be resolved via __getattr__ / __getattribute__ dispatch.
     //   * value > 1: a tagged PyObject* for a staticmethod descriptor or class
     //     variable; untag it and return it as a plain attribute (no self
-    //     binding).
+    //     binding) -- unless is_class_method is set (see below).
     // Use the tag helpers in inline_cache.cpp to read it.
     uintptr_t value{0};
     uint32_t keys_version;
@@ -319,6 +319,12 @@ class LoadMethodCache {
     // value != nullptr.
     bool has_getattr_hook{false};
 
+    // Set when the entry caches a class method (a Python-level classmethod or a
+    // C-level classmethod_descriptor). The (tagged, unbound) value is the
+    // underlying callable, which lookup() binds to the receiver's type rather
+    // than to the receiver itself. Only meaningful for unbound entries.
+    bool is_class_method{false};
+
     bool isValidKeysVersion(BorrowedRef<> obj);
   };
   static_assert(sizeof(Entry) == 24, "Entry must be small");
@@ -341,7 +347,8 @@ class LoadMethodCache {
       BorrowedRef<> value,
       BorrowedRef<> name,
       bool has_getattr_hook,
-      bool is_bound_method = true);
+      bool is_bound_method = true,
+      bool is_class_method = false);
 
   std::array<Entry, 4> entries_;
   std::unique_ptr<CacheStats> cache_stats_;

cinderx/PythonLib/test_cinderx/test_jit_attr_cache.py

@@ -43,6 +43,13 @@ def call_static_with_arg(obj):
     return obj.compute(10)
 
 
+@cinder_support.failUnlessJITCompiled
+def call_fromkeys(obj):
+    # `fromkeys` resolves to a classmethod_descriptor (a C-level classmethod)
+    # when accessed via an instance.
+    return obj.fromkeys([1, 2])
+
+
 class LoadMethodCacheTests(unittest.TestCase):
     def test_type_modified(self):
         class Oracle:
@@ -404,6 +411,161 @@ class Oracle:
         self.assertEqual(get_meaning_of_life(obj), 0)
 
 
+class LoadMethodClassMethodTests(unittest.TestCase):
+    """LoadMethodCache should cache class methods found on the type. Both a
+    Python-level classmethod and a C-level classmethod_descriptor are unwrapped
+    to their underlying callable, cached, and bound to the receiver's type (not
+    the receiver itself) when called.
+    """
+
+    def test_class_method_on_type(self):
+        class Oracle:
+            @classmethod
+            def meaning_of_life(cls):
+                return 42
+
+        obj = Oracle()
+        # Uncached, then cached.
+        self.assertEqual(get_meaning_of_life(obj), 42)
+        self.assertEqual(get_meaning_of_life(obj), 42)
+        for _ in range(100):
+            self.assertEqual(get_meaning_of_life(obj), 42)
+
+    def test_class_method_binds_to_type(self):
+        class Oracle:
+            @classmethod
+            def meaning_of_life(cls):
+                return cls.__name__
+
+        obj = Oracle()
+        # cls is bound to the receiver's type, not the receiver.
+        self.assertEqual(get_meaning_of_life(obj), "Oracle")
+        self.assertEqual(get_meaning_of_life(obj), "Oracle")
+
+    def test_class_method_binds_to_most_derived_type(self):
+        class Base:
+            @classmethod
+            def meaning_of_life(cls):
+                return cls.__name__
+
+        class Derived(Base):
+            pass
+
+        # An inherited classmethod binds to the most-derived type of the
+        # receiver. The two receiver types get independent cache entries.
+        self.assertEqual(get_meaning_of_life(Base()), "Base")
+        self.assertEqual(get_meaning_of_life(Derived()), "Derived")
+        self.assertEqual(get_meaning_of_life(Base()), "Base")
+        self.assertEqual(get_meaning_of_life(Derived()), "Derived")
+
+    def test_class_method_passes_args_after_cls(self):
+        class Oracle:
+            @classmethod
+            def compute(cls, x):
+                return x * 2
+
+        obj = Oracle()
+        # The explicit argument follows the implicit cls.
+        self.assertEqual(call_static_with_arg(obj), 20)
+        self.assertEqual(call_static_with_arg(obj), 20)
+
+    def test_class_method_type_modified(self):
+        class Oracle:
+            @classmethod
+            def meaning_of_life(cls):
+                return 42
+
+        obj = Oracle()
+        self.assertEqual(get_meaning_of_life(obj), 42)
+        self.assertEqual(get_meaning_of_life(obj), 42)
+
+        # Replace with a different classmethod; the cache must be invalidated.
+        # pyrefly: ignore [bad-assignment]
+        Oracle.meaning_of_life = classmethod(lambda cls: 0)
+        self.assertEqual(get_meaning_of_life(obj), 0)
+
+    def test_class_method_base_modified(self):
+        class Base:
+            @classmethod
+            def meaning_of_life(cls):
+                return 42
+
+        class Derived(Base):
+            pass
+
+        obj = Derived()
+        self.assertEqual(get_meaning_of_life(obj), 42)
+        self.assertEqual(get_meaning_of_life(obj), 42)
+
+        # Mutating the base should propagate to Derived and invalidate the cache.
+        # pyrefly: ignore [bad-assignment]
+        Base.meaning_of_life = classmethod(lambda cls: 0)
+        self.assertEqual(get_meaning_of_life(obj), 0)
+
+    def test_class_method_shadowed_by_instance(self):
+        class Oracle:
+            @classmethod
+            def meaning_of_life(cls):
+                return 42
+
+        obj = Oracle()
+        # Cache the classmethod first.
+        self.assertEqual(get_meaning_of_life(obj), 42)
+        self.assertEqual(get_meaning_of_life(obj), 42)
+
+        # A classmethod is a non-data descriptor, so an instance attribute
+        # shadows it.
+        # pyrefly: ignore [missing-attribute]
+        obj.meaning_of_life = nothing
+        self.assertEqual(get_meaning_of_life(obj), 0)
+
+    def test_class_method_with_getattr_defined(self):
+        """A classmethod present on the type is returned directly, not routed
+        through __getattr__, even when the type defines __getattr__."""
+
+        class Oracle:
+            @classmethod
+            def meaning_of_life(cls):
+                return 42
+
+            def __getattr__(self, name):
+                raise AttributeError(name)
+
+        obj = Oracle()
+        self.assertEqual(get_meaning_of_life(obj), 42)
+        self.assertEqual(get_meaning_of_life(obj), 42)
+        for _ in range(100):
+            self.assertEqual(get_meaning_of_life(obj), 42)
+
+    def test_class_method_wrapping_non_function(self):
+        """A classmethod wrapping a non-function callable is not cached, but is
+        still dispatched correctly through the descriptor."""
+
+        class Callable:
+            def __call__(self, cls):
+                return cls.__name__
+
+        class Oracle:
+            meaning_of_life = classmethod(Callable())
+
+        obj = Oracle()
+        self.assertEqual(get_meaning_of_life(obj), "Oracle")
+        self.assertEqual(get_meaning_of_life(obj), "Oracle")
+        for _ in range(100):
+            self.assertEqual(get_meaning_of_life(obj), "Oracle")
+
+    def test_classmethod_descriptor_on_builtin(self):
+        """A C-level classmethod_descriptor (e.g. dict.fromkeys) accessed via an
+        instance is cached and bound to the type."""
+
+        obj = {}
+        expected = {1: None, 2: None}
+        self.assertEqual(call_fromkeys(obj), expected)
+        self.assertEqual(call_fromkeys(obj), expected)
+        for _ in range(100):
+            self.assertEqual(call_fromkeys(obj), expected)
+
+
 class LoadMethodGetAttrTests(unittest.TestCase):
     """LoadMethodCache should support types that define __getattr__.