Add function for dumping paddle thirdparty apis

README.md

@@ -7,6 +7,7 @@
 
 ## 最近更新
 
+-   支持添加Paddle自定义算子
 -   支持单模型运行并dump相关数据
 -   提供离线对齐工具
 

docs/Tutorial.md

@@ -436,6 +436,14 @@ assign_weight(layer, module)
 
 设置环境变量可以打开该功能： `export PADIFF_API_CHECK=ON`
 
+### 添加自定义算子
+
+当`PADIFF_API_CHECK`开启时，可以添加PD_BUILD_OP方式注册的自定义算子，该功能通过`PADDLE_THIRDPARTY_API`环境变量开启
+
+设置时，需将自定义算子的module、api名写全，多个算子之间用逗号分隔，如：
+
+`export PADDLE_THIRDPARTY_API=paddle3d.ops.iou3d_nms,paddle3d.ops.hard_voxelize,paddle_xpu_nn.xpu_rms_norm`
+
 ### 略过 wrap_layer
 
 `export PADIFF_SIKP_WRAP_LAYER=TRUE`

padiff/__init__.py

@@ -228,12 +228,33 @@ def create_module(self, spec):
         return None
 
 
+def add_thirdparty_apis(thirdparty_apis):
+    # Exp. 1:
+    #   thirdparty_apis = "paddle3d.ops.iou3d_nms,paddle3d.ops.hard_voxelize,paddle_xpu_nn.xpu_rms_norm"
+    #   json.THIRD_PARTY = {"paddle3d.ops": {"iou3d_nms", "hard_voxelize"}, "paddle_xpu_nn": {"xpu_rms_norm"}}
+    thirdparty_apis = thirdparty_api.replace(" ", "").split(",")
+    json.THIRD_PARTY = {}
+    for fullname in thirdparty_apis:
+        module = fullname.rpartition(".")[0]
+        api = fullname.rpartition(".")[2]
+        if not module in json.THIRD_PARTY:
+            json.THIRD_PARTY[module] = {api}
+        else:
+            json.THIRD_PARTY[module].add(api)
+
+    self.paddle_apis.update(self.THIRD_PARTY)
+
+
 if os.getenv("PADIFF_API_CHECK") == "ON":
     for name in jsons.TORCH_PATH:
         if name in sys.modules.keys():
             module = sys.modules[name]
             wrap_api_method(module)
 
+    thirdparty_apis = os.getenv("PADDLE_THIRDPARTY_API")
+    if thirdparty_apis is not None:
+        add_thirdparty_apis(thirdparty_apis)
+
     for name in jsons.PADDLE_PATH:
         if name in sys.modules.keys():
             module = sys.modules[name]