diff --git a/test/torch_npu_schema.json b/test/torch_npu_schema.json index 603f723f8a5a2c64bb4e0415b906e3a8d3184cc6..f11c5872dc401f63927154f0c5541dee97724e90 100644 --- a/test/torch_npu_schema.json +++ b/test/torch_npu_schema.json @@ -1331,6 +1331,15 @@ "torch_npu.npu.synchronize": { "signature": "(device=None)" }, + "torch_npu.npu.obfuscation_initialize": { + "signature": "(hidden_size, tp_rank, cmd, *, data_type=None, model_obf_seed_id=0, data_obf_seed_id=0, thread_num=4, obf_coefficient=1.0)" + }, + "torch_npu.npu.obfuscation_finalize": { + "signature": "(fd_to_close)" + }, + "torch_npu.npu.obfuscation_calculate": { + "signature": "(fd, x, param, *, obf_coefficient=1.0)" + }, "torch_npu.npu.utilization": { "signature": "(device=None)" }, @@ -1730,6 +1739,15 @@ "torch_npu.npu.utils.synchronize": { "signature": "(device=None)" }, + "torch_npu.npu.utils.obfuscation_initialize": { + "signature": "(hidden_size, tp_rank, cmd, *, data_type=None, model_obf_seed_id=0, data_obf_seed_id=0, thread_num=4, obf_coefficient=1.0)" + }, + "torch_npu.npu.utils.obfuscation_finalize": { + "signature": "(fd_to_close)" + }, + "torch_npu.npu.utils.obfuscation_calculate": { + "signature": "(fd, x, param, *, obf_coefficient=1.0)" + }, "torch_npu.npu.utils.utilization": { "signature": "(device=None)" }, diff --git a/torch_npu/npu/__init__.py b/torch_npu/npu/__init__.py index 957ad7e7c8b16674d68693825b5846f7e318138e..6d1d51b5b6a1daf9bfd28f577dabea9860d0d6fa 100644 --- a/torch_npu/npu/__init__.py +++ b/torch_npu/npu/__init__.py @@ -120,7 +120,10 @@ __all__ = [ "get_device_limit", "set_stream_limit", "reset_stream_limit", - "get_stream_limit" + "get_stream_limit", + "obfuscation_initialize", + "obfuscation_finalize", + "obfuscation_calculate" ] from typing import Tuple, Union @@ -135,7 +138,8 @@ from torch_npu.utils import _should_print_warning import torch_npu from torch_npu.utils._error_code import ErrCode, pta_error, prof_error -from .utils import (synchronize, device_count, can_device_access_peer, set_device, current_device, get_device_name, +from .utils import (obfuscation_initialize, obfuscation_calculate, obfuscation_finalize, + synchronize, device_count, can_device_access_peer, set_device, current_device, get_device_name, get_device_properties, get_device_capability, _get_device_index, device, device_of, StreamContext, stream, set_stream, current_stream, default_stream, set_sync_debug_mode, get_sync_debug_mode, init_dump, current_blas_handle, is_bf16_supported, diff --git a/torch_npu/npu/utils.py b/torch_npu/npu/utils.py index 1fa5d3848ff9eb92a5cfded983541cf845a9af6d..588ccfff866bfd2b4b2f0608f95c96fe5feb9f77 100644 --- a/torch_npu/npu/utils.py +++ b/torch_npu/npu/utils.py @@ -14,7 +14,8 @@ from torch_npu.utils._error_code import ErrCode, pta_error, _except_handler from torch_npu.npu._backends import get_soc_version -__all__ = ["synchronize", "device_count", "can_device_access_peer", "set_device", "current_device", "get_device_name", +__all__ = ["obfuscation_initialize", "obfuscation_finalize", "obfuscation_calculate", + "synchronize", "device_count", "can_device_access_peer", "set_device", "current_device", "get_device_name", "get_device_properties", "mem_get_info", "get_device_capability", "utilization", "device", "device_of", "StreamContext", "stream", "set_stream", "current_stream", "default_stream", "set_sync_debug_mode", "get_sync_debug_mode", "init_dump", "set_dump", "finalize_dump", "is_support_inf_nan", "is_bf16_supported", @@ -22,6 +23,19 @@ __all__ = ["synchronize", "device_count", "can_device_access_peer", "set_device" "check_uce_in_memory", "stress_detect", "get_cann_version"] +def obfuscation_initialize(hidden_size, tp_rank, cmd, *, data_type=None, model_obf_seed_id=0, data_obf_seed_id=0, thread_num=4, obf_coefficient=1.0): + return torch_npu.obfuscation_initialize(hidden_size, tp_rank, cmd, data_type=data_type, model_obf_seed_id=model_obf_seed_id, + data_obf_seed_id=data_obf_seed_id, thread_num=thread_num, obf_coefficient=obf_coefficient) + + +def obfuscation_finalize(fd_to_close): + return torch_npu.obfuscation_finalize(fd_to_close) + + +def obfuscation_calculate(fd, x, param, *, obf_coefficient=1.0): + return torch_npu.obfuscation_calculate(fd, x, param, obf_coefficient=obf_coefficient) + + def get_cann_version(module="CANN"): r""" Args: