Files
paddlepaddle--paddle/python/paddle/jit/sot/utils/paddle_api_config.py
T
2026-07-13 12:40:42 +08:00

169 lines
5.2 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import paddle
def is_inplace_api(func):
inplace_apis = {paddle.static.setitem}
return func in inplace_apis
def get_tensor_methods():
return [
member_name
for member_name, member in inspect.getmembers(paddle.pir.Value)
if inspect.isfunction(member) or inspect.ismethoddescriptor(member)
]
def get_paddle_api():
modules = [
paddle,
paddle.nn.functional,
paddle.nn.quant,
paddle.incubate.nn.functional,
paddle.linalg,
paddle.signal,
paddle.fft,
paddle.vision.ops,
paddle.metric,
paddle.geometric,
]
distributed_apis = [
paddle.distributed.all_reduce,
paddle.distributed.shard_tensor,
paddle.distributed.reshard,
paddle.distributed.all_gather,
paddle.distributed.alltoall,
paddle.distributed.barrier,
paddle.distributed.recv,
paddle.distributed.send,
paddle.distributed.broadcast,
paddle.distributed.unshard_dtensor,
paddle.distributed.auto_parallel.api.dtensor_to_local,
paddle.distributed.auto_parallel.api.dtensor_from_local,
paddle.distributed.auto_parallel.api.moe_global_mesh_tensor,
paddle.distributed.auto_parallel.api.moe_sub_mesh_tensors,
]
special_paddle_apis = [
paddle.tensor.fill_constant,
paddle.tensor.top_p_sampling,
]
non_operator_related_apis = [
paddle.in_dynamic_mode,
paddle.save,
paddle.load,
paddle.get_cuda_rng_state,
paddle.set_rng_state,
paddle.set_cuda_rng_state,
paddle.get_rng_state,
paddle.set_default_dtype,
paddle.check_shape,
paddle.summary,
paddle.finfo,
paddle.iinfo,
paddle.enable_static,
paddle.disable_static,
paddle.is_grad_enabled,
]
# TODO: users should not call static_apis, but we need to use, so add static_apis here temporary
static_apis = [paddle.static.setitem, paddle.static.accuracy]
paddle_api_list = []
for module in modules:
for fn_name in getattr(module, "__all__", []):
fn = getattr(module, fn_name)
if inspect.isfunction(fn):
paddle_api_list.append(fn)
return list(
set(special_paddle_apis)
| set(distributed_apis)
| set(static_apis)
| set(paddle_api_list) - set(non_operator_related_apis)
)
paddle_api_list = get_paddle_api()
# TODO(Aurelius84): It seems that we use it to judge 'in_paddle_module()'.
# Bug what does 'is_paddle_module' really means? Is all paddle.xx sub module
# considered as paddle module
paddle_api_module_prefix = {
"paddle.nn.functional",
}
break_graph_functions = set()
break_graph_layer_classes = set()
break_graph_tensor_method = {
'register_hook',
'numpy',
'clear_gradient',
'tolist',
'item',
# TODO: Browse all possible functions and make prior judgments.
}
not_supported_paddle_layer = {paddle.nn.RNN}
def is_not_supported_paddle_layer(layer_class):
return layer_class in not_supported_paddle_layer
def is_break_graph_tensor_methods(method_name):
return method_name in break_graph_tensor_method
def add_break_graph_function(fn):
break_graph_functions.add(fn)
def add_break_graph_layer_class(layer_class: type[paddle.nn.Layer]):
break_graph_layer_classes.add(layer_class)
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_onednn,
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,
paddle.base.libpaddle.is_compiled_with_flagcx,
}
if hasattr(paddle.base.libpaddle, "get_device_properties"):
NATIVE_CODE_PURE_FUNCTIONS.add(
paddle.base.libpaddle.get_device_properties
)
return api in NATIVE_CODE_PURE_FUNCTIONS