[Utils] Allow classmethod and staticmethod in TVMDerivedObject (#14249)

Instance methods that exist in the user-defined class but not in the
TVM base are forward using `__getattr__`.  However, this is only
applied for attribute look of instances, and doesn't apply for
attribute lookup on the class object itself, such as when calling a
classmethod or staticmethod.

This commit exposes class methods and static methods in the wrapper
class, if they are defined in the user-defined subclass.

Author: Lunderberg

python/tvm/meta_schedule/utils.py

@@ -128,6 +128,9 @@ def __setattr__(self, name, value):
     TVMDerivedObject.__name__ = cls.__name__
     TVMDerivedObject.__doc__ = cls.__doc__
     TVMDerivedObject.__module__ = cls.__module__
+    for key, value in cls.__dict__.items():
+        if isinstance(value, (classmethod, staticmethod)):
+            setattr(TVMDerivedObject, key, value)
     return TVMDerivedObject
 
 

tests/python/unittest/test_meta_schedule_post_order_apply.py

@@ -404,5 +404,26 @@ def _get_sch(filter_fn):
     assert len(schs) == 8
 
 
+def test_meta_schedule_derived_object():
+    @derived_object
+    class RemoveBlock(PyScheduleRule):
+        @classmethod
+        def class_construct(cls):
+            return cls()
+
+        @staticmethod
+        def static_construct():
+            return RemoveBlock()
+
+    inst_by_init = RemoveBlock()
+    assert isinstance(inst_by_init, RemoveBlock)
+
+    inst_by_classmethod = RemoveBlock.class_construct()
+    assert isinstance(inst_by_classmethod, RemoveBlock)
+
+    inst_by_staticmethod = RemoveBlock.static_construct()
+    assert isinstance(inst_by_staticmethod, RemoveBlock)
+
+
 if __name__ == "__main__":
     tvm.testing.main()