Merge pull request #491 from cangtianhuang/develop

Fix `API Tracer`

Author: cangtianhuang

tools/api_tracer/api_alias.py

@@ -1,6 +1,9 @@
+import os
 import re
 from pathlib import Path
 
+import yaml
+
 INPUT_DIR = Path("tools/api_tracer/trace_output_tmp")
 
 
@@ -12,33 +15,55 @@ def parse_api(api):
     return api
 
 
-def process_file(input_path):
+def process_file(input_path, target_apis):
     input_name = input_path.name
     output_name = input_name[4:]
     output_path = input_path.parent / output_name
 
+    output_excluded_name = output_name.replace(".txt", "_excluded.txt")
+    output_excluded_path = input_path.parent / output_excluded_name
+
     apis = set()
     with input_path.open("r") as f:
         apis = set([line.strip() for line in f if line.strip()])
     print(f"Read {len(apis)} apis from {input_path}", flush=True)
 
     alias_apis = set()
+    excluded_apis = set()
     for api in apis:
+        alias_api = api
+        if alias_api not in target_apis:
+            excluded_apis.add(api)
+            continue
         alias_apis.add(parse_api(api))
 
     with output_path.open("w") as f:
         f.writelines(f"{line}\n" for line in sorted(alias_apis))
     print(f"Write {len(alias_apis)} alias apis to {output_path}", flush=True)
 
+    with output_excluded_path.open("w") as f:
+        f.writelines(f"{line}\n" for line in sorted(excluded_apis))
+    print(f"Write {len(excluded_apis)} excluded apis to {output_excluded_path}", flush=True)
+
 
 def main():
-    input_files = list(INPUT_DIR.glob("api_apis*.txt"))
+    yaml_path = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)),
+        "api_list",
+        "torch_api_list.yaml",
+    )
+    target_apis = []
+    with open(yaml_path, "r", encoding="utf-8") as f:
+        target_apis = yaml.safe_load(f)
+    print(f"Loaded {len(target_apis)} target APIs.")
+
+    input_files = list(INPUT_DIR.glob("api_apis.txt"))
     if not input_files:
         print(f"No input files found in {INPUT_DIR}", flush=True)
         return
 
     for input_file in sorted(input_files):
-        process_file(input_file)
+        process_file(input_file, target_apis)
 
 
 if __name__ == "__main__":

tools/api_tracer/framework_dialect.py

@@ -123,7 +123,7 @@ def install(self):
                         types.BuiltinMethodType,
                     ),
                 ):
-                    wrapped_func = self._create_wrapper(
+                    wrapper = self._create_wrapper(
                         api_name, original_api, self.serializer, self.level
                     )
                 elif isinstance(original_api, (classmethod, staticmethod)):
@@ -154,6 +154,8 @@ def install(self):
                     setattr(parent_obj, func_name, wrapper)
                     self._original_apis[api_name] = original_api
                     patched_count += 1
+                else:
+                    skipped_count += 1
             except (TypeError, AttributeError) as e:
                 error_msg = str(e).lower()
                 if (
@@ -453,7 +455,7 @@ class PyTorchDialect(FrameworkDialect):
         "torch.fx.experimental.unification.multipledispatch.dispatcher.str_signature",
         "torch.nn.functional.handle_torch_function",
         "torch.nn.functional.has_torch_function_unary",
-        "torch.optim.Optimizer.profile_hook_step",
+        "torch.optim.Optimizer.profile_hook_step",  # it will be overridden by subclass of Optimizer
     }
 
     def get_framework_name(self) -> str:
@@ -534,6 +536,11 @@ def discover_apis(self) -> List[str]:
                             if cls_member_name in self.IGNORE_ATTRIBUTES:
                                 continue
                             full_cls_name = f"{full_name}.{cls_member_name}"
+                            if full_cls_name in self.IGNORE_CLASSES_OR_METHODS:
+                                continue
+                            # it will be overridden by subclass of Optimizer
+                            if cls_member_name == "profile_hook_step":
+                                continue
                             if isinstance(
                                 cls_member,
                                 (