[dynamo][guards-cpp-refactor] GetAttrGuardAccessor (pytorch#119833)

Pull Request resolved: pytorch#119833
Approved by: https://github.com/jansel
ghstack dependencies: pytorch#119822, pytorch#119827

Author: anijain2305

test/dynamo/test_guard_manager.py

@@ -7,6 +7,7 @@
 from torch._C._dynamo import guards
 
 RootGuardManager = guards.RootGuardManager
+GetAttrGuardAccessor = guards.GetAttrGuardAccessor
 
 
 def id_type(x):
@@ -137,6 +138,44 @@ def test_guard_manager_leaf_guard(self):
         self.assertFalse(guard_manager.check(4))
         self.assertFalse(guard_manager.check("foo"))
 
+    def test_attr_guard_manager(self):
+        class Foo:
+            def __init__(self, x, y):
+                self.x = x
+                self.y = y
+
+        foo = Foo(1, 2)
+        guard_manager = RootGuardManager()
+        guard_manager.add_type_match_guard(id_type(foo), ["type(x) == Foo"])
+        guard_manager.getattr_manager("x", 1).add_lambda_guard(
+            functools.partial(equals_match, expected=foo.x),
+            equals_match_verbose_code_parts(foo.x),
+        )
+        guard_manager.getattr_manager("y", 2).add_lambda_guard(
+            functools.partial(equals_match, expected=foo.y),
+            equals_match_verbose_code_parts(foo.y),
+        )
+        self.assertEqual(len(guard_manager.get_leaf_guards()), 1)
+        # 2 child managers, one for x and one for y
+        self.assertEqual(len(guard_manager.get_accessors()), 2)
+        self.assertTrue(
+            isinstance(guard_manager.get_accessors()[0], GetAttrGuardAccessor)
+        )
+        self.assertTrue(
+            isinstance(guard_manager.get_accessors()[1], GetAttrGuardAccessor)
+        )
+        # Check leaf guards on child managers
+        self.assertEqual(
+            len(guard_manager.getattr_manager("x", None).get_leaf_guards()), 1
+        )
+        self.assertEqual(
+            len(guard_manager.getattr_manager("y", None).get_leaf_guards()), 1
+        )
+
+        self.assertTrue(guard_manager.check(foo))
+        self.assertFalse(guard_manager.check(Foo(3, 4)))
+        self.assertFalse(guard_manager.check("foo"))
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests

torch/csrc/dynamo/guards.cpp

@@ -1353,6 +1353,60 @@ std::unique_ptr<GuardManager> make_guard_manager(
   return std::make_unique<GuardManager>(root);
 }
 
+/**
+ * Represents __getattr__ acccessor.
+ */
+class GetAttrGuardAccessor : public GuardAccessor {
+ public:
+  GetAttrGuardAccessor(
+      RootGuardManager* root,
+      py::str name,
+      py::handle example_value)
+      : GuardAccessor(root, name, example_value), _attr_name(name.ptr()) {}
+
+  // NB: Intentional duplication between check_nopybind and
+  // check_verbose_nopybind.
+  bool check_nopybind(PyObject* obj) override { // borrowed ref
+    PyObject* x = PyObject_GetAttr(obj, _attr_name); // new ref
+    if (x == nullptr) {
+      // Attribute absent, clear the exception and return false.
+      PyErr_Clear();
+      return false;
+    }
+    bool result = _guard_manager->check_nopybind(x);
+    Py_DECREF(x);
+    return result;
+  }
+
+  GuardDebugInfo check_verbose_nopybind(
+      PyObject* obj) override { // borrowed ref
+    PyObject* x = PyObject_GetAttr(obj, _attr_name); // new ref
+    if (x == nullptr) {
+      // Attribute absent, clear the exception and return false.
+      PyErr_Clear();
+      return GuardDebugInfo(
+          false,
+          std::string("get attr failed for attr name ") +
+              py::str(_attr_name).cast<std::string>(),
+          0);
+    }
+    GuardDebugInfo result = _guard_manager->check_verbose_nopybind(x);
+    Py_DECREF(x);
+    return result;
+  }
+
+  std::string repr() const override {
+    // Helpful when priting GuardManager tree structure.
+    return "GetAttrGuardAccessor(" + py::str(_attr_name).cast<std::string>() +
+        ")";
+  }
+
+ private:
+  // no need of py::object here because the attr_name is already passed on to
+  // the base class as accessor_key which is a py::object.
+  PyObject* _attr_name;
+};
+
 } // namespace
 
 static void* _torchinductor_pyobject_tensor_data_ptr(PyObject* obj) {
@@ -1458,6 +1512,10 @@ PyObject* torch_c_dynamo_guards_init() {
   py::class_<GuardAccessor, std::unique_ptr<GuardAccessor>>(
       py_m, "GuardAccessor")
       .def("repr", &GuardAccessor::repr);
+  py::class_<
+      GetAttrGuardAccessor,
+      GuardAccessor,
+      std::unique_ptr<GetAttrGuardAccessor>>(py_m, "GetAttrGuardAccessor");
 
   // Guard Manager - No constructor in python, python should use
   // RootGuardManager.
@@ -1510,7 +1568,13 @@ PyObject* torch_c_dynamo_guards_init() {
              py::object verbose_code_parts) -> void {
             self.add_leaf_guard(
                 std::make_shared<EQUALS_MATCH>(value, verbose_code_parts));
-          });
+          })
+      // return by reference because C++ GuardManager has the ownership of
+      // accessors and guard managers
+      .def(
+          "getattr_manager",
+          &GuardManager::get_child_manager<GetAttrGuardAccessor>,
+          py::return_value_policy::reference);
 
   // Root Guard Manager
   py::class_<RootGuardManager, GuardManager, std::unique_ptr<RootGuardManager>>(