[SOT] Mark some APIs can be directly run in simulation mode (#70293)

SigureMo · web-flow · commit 93d8812435ac · 2024-12-19T16:27:06.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py b/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
@@ -36,6 +36,7 @@
     is_break_graph_api,
     is_break_graph_tensor_methods,
     is_builtin_fn,
+    is_directly_run_api,
     is_not_supported_paddle_layer,
     is_paddle_api,
     magic_method_builtin_dispatch,
@@ -699,6 +700,22 @@ def call_function(self, /, *args, **kwargs):
                 )
                 return handler(*args, **kwargs)
 
+        # If API can be directly called in simulation mode (e.g. user defined native code
+        # without graph affect), we can directly call it.
+        if is_directly_run_api(self.value):
+            from ..function_graph import convert_to_py_value
+
+            res = self.value(
+                *convert_to_py_value(args),
+                **convert_to_py_value(kwargs),
+            )
+
+            return VariableFactory.from_value(
+                res,
+                self.graph,
+                DummyTracker([self, *list(args), *list(kwargs.values())]),
+            )
+
         # Try to inline call the magic function
         magic_methods = magic_method_builtin_dispatch(self.value)
         for magic_method in magic_methods:
diff --git a/python/paddle/jit/sot/utils/__init__.py b/python/paddle/jit/sot/utils/__init__.py
@@ -49,6 +49,7 @@
 from .paddle_api_config import (  # noqa: F401
     get_tensor_methods,
     is_break_graph_tensor_methods,
+    is_directly_run_api,
     is_inplace_api,
     is_not_supported_paddle_layer,
 )
diff --git a/python/paddle/jit/sot/utils/paddle_api_config.py b/python/paddle/jit/sot/utils/paddle_api_config.py
@@ -129,3 +129,28 @@ def is_break_graph_tensor_methods(method_name):
 
 def add_break_graph_apis(apis: list):
     break_graph_set.update(apis)
+
+
+def is_directly_run_api(api):
+    from .utils import hashable
+
+    if not hashable(api):
+        return False
+    NATIVE_CODE_PURE_FUNCTIONS = {
+        paddle.base.libpaddle.is_compiled_with_avx,
+        paddle.base.libpaddle.is_compiled_with_cuda,
+        paddle.base.libpaddle.is_compiled_with_cudnn_frontend,
+        paddle.base.libpaddle.is_compiled_with_rocm,
+        paddle.base.libpaddle.is_compiled_with_custom_device,
+        paddle.base.libpaddle.is_compiled_with_ipu,
+        paddle.base.libpaddle.is_compiled_with_xpu,
+        paddle.base.libpaddle.is_compiled_with_mkldnn,
+        paddle.base.libpaddle.is_compiled_with_nccl,
+        paddle.base.libpaddle.is_compiled_with_mpi,
+        paddle.base.libpaddle.is_compiled_with_mpi_aware,
+        paddle.base.libpaddle.is_compiled_with_cinn,
+        paddle.base.libpaddle.is_compiled_with_distribute,
+        paddle.base.libpaddle.is_compiled_with_brpc,
+        paddle.base.libpaddle.is_compiled_with_dist,
+    }
+    return api in NATIVE_CODE_PURE_FUNCTIONS
diff --git a/test/sot/test_builtin_dispatch.py b/test/sot/test_builtin_dispatch.py
@@ -388,5 +388,46 @@ def test_builtin_type_conversion_breakgraph(self):
         )
 
 
+@check_no_breakgraph
+def test_native_code_function():
+    res1 = paddle.base.libpaddle.is_compiled_with_avx()
+    res2 = paddle.base.libpaddle.is_compiled_with_cuda()
+    res3 = paddle.base.libpaddle.is_compiled_with_cudnn_frontend()
+    res4 = paddle.base.libpaddle.is_compiled_with_rocm()
+    res5 = paddle.base.libpaddle.is_compiled_with_custom_device("npu")
+    res6 = paddle.base.libpaddle.is_compiled_with_ipu()
+    res7 = paddle.base.libpaddle.is_compiled_with_xpu()
+    res8 = paddle.base.libpaddle.is_compiled_with_mkldnn()
+    res9 = paddle.base.libpaddle.is_compiled_with_nccl()
+    res10 = paddle.base.libpaddle.is_compiled_with_mpi()
+    res11 = paddle.base.libpaddle.is_compiled_with_mpi_aware()
+    res12 = paddle.base.libpaddle.is_compiled_with_cinn()
+    res13 = paddle.base.libpaddle.is_compiled_with_distribute()
+    res14 = paddle.base.libpaddle.is_compiled_with_brpc()
+    res15 = paddle.base.libpaddle.is_compiled_with_dist()
+    return (
+        res1,
+        res2,
+        res3,
+        res4,
+        res5,
+        res6,
+        res7,
+        res8,
+        res9,
+        res10,
+        res11,
+        res12,
+        res13,
+        res14,
+        res15,
+    )
+
+
+class TestNativeCodeFunction(TestCaseBase):
+    def test_native_code_function(self):
+        self.assert_results(test_native_code_function)
+
+
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -49,6 +49,7 @@`
`49`	`49`	`from .paddle_api_config import ( # noqa: F401`
`50`	`50`	`get_tensor_methods,`
`51`	`51`	`is_break_graph_tensor_methods,`
	`52`	`+ is_directly_run_api,`
`52`	`53`	`is_inplace_api,`
`53`	`54`	`is_not_supported_paddle_layer,`
`54`	`55`	`)`