mindstudio-probe 1.2.1__py3-none-any.whl → 1.3.0__py3-none-any.whl

msprobe/mindspore/dump/hook_cell/api_registry.py

@@ -1,205 +0,0 @@
-# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
-# 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.
-
-from mindspore import Tensor, ops, mint
-from mindspore.mint.nn import functional
-from mindspore.common._stub_tensor import StubTensor
-from mindspore.communication import comm_func
-
-from msprobe.mindspore.dump.hook_cell.wrap_api import (HOOKTensor, HOOKStubTensor, HOOKFunctionalOP,
-                                                       HOOKMintOP, HOOKMintNNFunctionalOP, HOOKDistributedOP,
-                                                       HOOKTorchOP, HOOKTorchTensor, HOOKTorchFunctionalOP,
-                                                       HOOKTorchDistributedOP, HOOKTorchNpuOP,
-                                                       get_wrap_api_list, get_wrap_torch_api_list, setup_hooks)
-from msprobe.core.common.utils import Const
-from msprobe.mindspore.common.utils import is_mindtorch
-
-if is_mindtorch():
-    import torch
-    import torch_npu
-
-
-def stub_method(method):
-    def wrapped_method(*args, **kwargs):
-        return method(*args, **kwargs)
-    return wrapped_method
-
-
-class ApiRegistry:
-    def __init__(self):
-        self.tensor_ori_attr = {}
-        self.stub_tensor_ori_attr = {}
-        self.functional_ori_attr = {}
-        self.mint_ops_ori_attr = {}
-        self.mint_func_ops_ori_attr = {}
-        self.distributed_ori_attr = {}
-        self.norm_inner_ops_ori_attr = {}
-
-        self.torch_ori_attr = {}
-        self.torch_tensor_ori_attr = {}
-        self.torch_functional_ori_attr = {}
-        self.torch_distributed_ori_attr = {}
-        self.torch_npu_ori_attr = {}
-
-        self.tensor_hook_attr = {}
-        self.stub_tensor_hook_attr = {}
-        self.functional_hook_attr = {}
-        self.mint_ops_hook_attr = {}
-        self.mint_func_ops_hook_attr = {}
-        self.distibuted_hook_attr = {}
-        self.norm_inner_ops_hook_attr = {}
-
-        self.torch_hook_attr = {}
-        self.torch_tensor_hook_attr = {}
-        self.torch_functional_hook_attr = {}
-        self.torch_distributed_hook_attr = {}
-        self.torch_npu_hook_attr = {}
-
-        self.norm_inner_ops = ["norm", "square", "sqrt", "is_complex"]
-
-    @staticmethod
-    def store_ori_attr(ori_api_group, api_list, api_ori_attr):
-        for api in api_list:
-            if Const.SEP in api:
-                sub_module_name, sub_op = api.rsplit(Const.SEP, 1)
-                sub_module = getattr(ori_api_group, sub_module_name)
-                ori_api_func = getattr(sub_module, sub_op)
-            else:
-                ori_api_func = getattr(ori_api_group, api)
-            if ori_api_group == StubTensor:
-                api_ori_attr[api] = stub_method(ori_api_func)
-                continue
-            api_ori_attr[api] = ori_api_func
-
-    @staticmethod
-    def set_api_attr(api_group, attr_dict):
-        for api, api_attr in attr_dict.items():
-            if Const.SEP in api:
-                sub_module_name, sub_op = api.rsplit(Const.SEP, 1)
-                sub_module = getattr(api_group, sub_module_name, None)
-                if sub_module is not None:
-                    setattr(sub_module, sub_op, api_attr)
-            else:
-                setattr(api_group, api, api_attr)
-
-    def norm_inner_op_set_hook_func(self):
-        self.set_api_attr(ops, self.norm_inner_ops_hook_attr)
-
-    def norm_inner_op_set_ori_func(self):
-        self.set_api_attr(ops, self.norm_inner_ops_ori_attr)
-
-    def api_set_hook_func(self):
-        if is_mindtorch():
-            self.set_api_attr(torch, self.torch_hook_attr)
-            self.set_api_attr(torch.Tensor, self.torch_tensor_hook_attr)
-            self.set_api_attr(torch.nn.functional, self.torch_functional_hook_attr)
-            self.set_api_attr(torch.distributed, self.torch_distributed_hook_attr)
-            self.set_api_attr(torch_npu, self.torch_npu_hook_attr)
-        else:
-            self.set_api_attr(Tensor, self.tensor_hook_attr)
-            self.set_api_attr(StubTensor, self.stub_tensor_hook_attr)
-            self.set_api_attr(ops, self.functional_hook_attr)
-            self.set_api_attr(mint, self.mint_ops_hook_attr)
-            self.set_api_attr(functional, self.mint_func_ops_hook_attr)
-            self.set_api_attr(comm_func, self.distibuted_hook_attr)
-
-    def api_set_ori_func(self):
-        if is_mindtorch():
-            self.set_api_attr(torch, self.torch_ori_attr)
-            self.set_api_attr(torch.Tensor, self.torch_tensor_ori_attr)
-            self.set_api_attr(torch.nn.functional, self.torch_functional_ori_attr)
-            self.set_api_attr(torch.distributed, self.torch_distributed_ori_attr)
-            self.set_api_attr(torch_npu, self.torch_npu_ori_attr)
-        else:
-            self.set_api_attr(Tensor, self.tensor_ori_attr)
-            self.set_api_attr(StubTensor, self.stub_tensor_ori_attr)
-            self.set_api_attr(ops, self.functional_ori_attr)
-            self.set_api_attr(mint, self.mint_ops_ori_attr)
-            self.set_api_attr(functional, self.mint_func_ops_ori_attr)
-            self.set_api_attr(comm_func, self.distributed_ori_attr)
-
-    def initialize_hook(self, hook):
-        setup_hooks(hook)
-        if is_mindtorch():
-            wrap_torch_api_name = get_wrap_torch_api_list()
-            self.store_ori_attr(torch,
-                                wrap_torch_api_name.torch_api_names, self.torch_ori_attr)
-            self.store_ori_attr(torch.Tensor,
-                                wrap_torch_api_name.tensor_api_names, self.torch_tensor_ori_attr)
-            self.store_ori_attr(torch.nn.functional,
-                                wrap_torch_api_name.functional_api_names, self.torch_functional_ori_attr)
-            self.store_ori_attr(torch.distributed,
-                                wrap_torch_api_name.distributed_api_names, self.torch_distributed_ori_attr)
-            self.store_ori_attr(torch_npu,
-                                wrap_torch_api_name.npu_api_names, self.torch_npu_ori_attr)
-            for attr_name in dir(HOOKTorchOP):
-                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
-                    self.torch_hook_attr[api_name] = getattr(HOOKTorchOP, attr_name)
-            for attr_name in dir(HOOKTorchTensor):
-                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
-                    self.torch_tensor_hook_attr[api_name] = getattr(HOOKTorchTensor, attr_name)
-            for attr_name in dir(HOOKTorchFunctionalOP):
-                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
-                    self.torch_functional_hook_attr[api_name] = getattr(HOOKTorchFunctionalOP, attr_name)
-            for attr_name in dir(HOOKTorchDistributedOP):
-                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
-                    self.torch_distributed_hook_attr[api_name] = getattr(HOOKTorchDistributedOP, attr_name)
-            for attr_name in dir(HOOKTorchNpuOP):
-                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                    api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
-                    self.torch_npu_hook_attr[api_name] = getattr(HOOKTorchNpuOP, attr_name)
-            return
-
-        wrap_api_name = get_wrap_api_list()
-        self.store_ori_attr(Tensor, wrap_api_name.tensor_api_names, self.tensor_ori_attr)
-        self.store_ori_attr(StubTensor, wrap_api_name.stub_tensor_api_names, self.stub_tensor_ori_attr)
-        self.store_ori_attr(ops, wrap_api_name.ops_api_names, self.functional_ori_attr)
-        self.store_ori_attr(mint, wrap_api_name.mint_api_names, self.mint_ops_ori_attr)
-        self.store_ori_attr(functional, wrap_api_name.mint_nn_func_api_names, self.mint_func_ops_ori_attr)
-        self.store_ori_attr(comm_func, wrap_api_name.distributed_api_names, self.distributed_ori_attr)
-        self.store_ori_attr(ops, self.norm_inner_ops, self.norm_inner_ops_ori_attr)
-        for attr_name in dir(HOOKTensor):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
-                self.tensor_hook_attr[api_name] = getattr(HOOKTensor, attr_name)
-        for attr_name in dir(HOOKStubTensor):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
-                self.stub_tensor_hook_attr[api_name] = getattr(HOOKStubTensor, attr_name)
-        for attr_name in dir(HOOKFunctionalOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
-                self.functional_hook_attr[api_name] = getattr(HOOKFunctionalOP, attr_name)
-                if api_name in self.norm_inner_ops:
-                    self.norm_inner_ops_hook_attr[api_name] = getattr(HOOKFunctionalOP, attr_name)
-        for attr_name in dir(HOOKMintOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
-                self.mint_ops_hook_attr[api_name] = getattr(HOOKMintOP, attr_name)
-        for attr_name in dir(HOOKMintNNFunctionalOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
-                self.mint_func_ops_hook_attr[api_name] = getattr(HOOKMintNNFunctionalOP, attr_name)
-        for attr_name in dir(HOOKDistributedOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                api_name = attr_name[Const.ATTR_NAME_PREFIX_LEN:]
-                self.distibuted_hook_attr[api_name] = getattr(HOOKDistributedOP, attr_name)
-
-
-api_register = ApiRegistry()

msprobe/mindspore/dump/hook_cell/wrap_api.py

@@ -1,212 +0,0 @@
-# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
-# 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 os
-
-from mindspore import Tensor, mint, ops
-from mindspore.common._stub_tensor import StubTensor
-from mindspore.communication import comm_func
-from mindspore.mint.nn import functional
-
-from msprobe.core.common.const import Const
-from msprobe.core.common.file_utils import load_yaml
-from msprobe.mindspore.common.const import Const as MsConst
-from msprobe.mindspore.common.utils import is_mindtorch
-from msprobe.mindspore.dump.hook_cell.hook_cell import HOOKCell
-
-if is_mindtorch():
-    import torch
-    import torch_npu
-
-cur_path = os.path.dirname(os.path.realpath(__file__))
-yaml_path = os.path.join(cur_path, MsConst.SUPPORTED_API_LIST_FILE)
-torch_yaml_path = os.path.join(cur_path, "../../../pytorch/hook_module", MsConst.SUPPORTED_API_LIST_FILE)
-
-
-class HOOKTensor(object):
-    pass
-
-
-class HOOKStubTensor(object):
-    pass
-
-
-class HOOKFunctionalOP(object):
-    pass
-
-
-class HOOKMintOP(object):
-    pass
-
-
-class HOOKMintNNFunctionalOP(object):
-    pass
-
-
-class HOOKDistributedOP(object):
-    pass
-
-
-class HOOKTorchOP(object):
-    pass
-
-
-class HOOKTorchTensor(object):
-    pass
-
-
-class HOOKTorchFunctionalOP(object):
-    pass
-
-
-class HOOKTorchDistributedOP(object):
-    pass
-
-
-class HOOKTorchNpuOP(object):
-    pass
-
-
-class ApiTemplate(HOOKCell):
-    def __init__(self, api_name, api_dict, prefix, hook):
-        self.api_name = api_name
-        self.api_func = api_dict[api_name]
-        self.prefix_api_name = prefix + str(api_name.split(Const.SEP)[-1]) + Const.SEP
-        super().__init__(hook)
-
-    @staticmethod
-    def async_to_sync(output):
-        # Fake handle, used to return after the CommHandle executes the wait method
-        fake_handle = type("FakeHandle", (), {"wait": lambda self: None})()
-        if isinstance(output, tuple) and len(output) == 2 and hasattr(output[1], "wait"):
-            output[1].wait()
-            output = (output[0], fake_handle)
-        elif hasattr(output, "wait"):
-            output.wait()
-            output = fake_handle
-        return output
-
-    def construct(self, *args, **kwargs):
-        if self.api_name.startswith(MsConst.DROPOUT_API_NAME_PREFIX):
-            return args[0] if args else kwargs.get(Const.INPUT)
-        
-        output = self.api_func(*args, **kwargs)
-
-        if self.prefix_api_name.startswith(MsConst.DISTRIBUTED_DATA_PREFIX):
-            if kwargs.get("async_op") or self.api_name in ["isend", "irecv"]:
-                output = self.async_to_sync(output)
-        return output
-
-    def forward(self, *args, **kwargs):
-        if self.api_name.startswith(MsConst.DROPOUT_API_NAME_PREFIX):
-            return args[0] if args else kwargs.get(Const.INPUT)
-        return self.api_func(*args, **kwargs)
-
-
-class WrapApiName:
-    def __init__(self, tensor_api_names, stub_tensor_api_names, ops_api_names, mint_api_names, mint_nn_func_api_names,
-                 distributed_api_names):
-        self.tensor_api_names = tensor_api_names
-        self.stub_tensor_api_names = stub_tensor_api_names
-        self.ops_api_names = ops_api_names
-        self.mint_api_names = mint_api_names
-        self.mint_nn_func_api_names = mint_nn_func_api_names
-        self.distributed_api_names = distributed_api_names
-
-
-class WrapTorchApiName:
-    def __init__(self, torch_api_names, tensor_api_names, functional_api_names, distributed_api_names, npu_api_names):
-        self.torch_api_names = torch_api_names
-        self.tensor_api_names = tensor_api_names
-        self.functional_api_names = functional_api_names
-        self.distributed_api_names = distributed_api_names
-        self.npu_api_names = npu_api_names
-
-
-def get_wrap_api_list():
-    api_list = load_yaml(yaml_path)
-    tensor_api = api_list.get(MsConst.SUPPORTED_TENSOR_LIST_KEY)
-    ops_api = api_list.get(MsConst.SUPPORTED_OPS_LIST_KEY)
-    mint_api = api_list.get(MsConst.SUPPORTED_MINT_LIST_KEY)
-    mint_nn_func_api = api_list.get(MsConst.SUPPORTED__MINT_NN_FUNC_LIST_KEY)
-    distributed_api = api_list.get(MsConst.SUPPORTED_COMM_LIST_KEY)
-    wrap_api_name = WrapApiName(set(tensor_api) & set(dir(Tensor)),
-                                set(tensor_api) & set(dir(StubTensor)),
-                                set(ops_api) & set(dir(ops)),
-                                set(mint_api) & set(dir(mint)),
-                                set(mint_nn_func_api) & set(dir(functional)),
-                                set(distributed_api) & set(dir(comm_func)))
-    return wrap_api_name
-
-
-def get_wrap_torch_api_list():
-    api_list = load_yaml(torch_yaml_path)
-    torch_api = api_list.get("torch")
-    tensor_api = api_list.get("tensor")
-    functional_api = api_list.get("functional")
-    distributed_api = api_list.get("distributed")
-    npu_api = api_list.get("torch_npu")
-    wrap_api_name = WrapTorchApiName(set(torch_api) & set(dir(torch)),
-                                     set(tensor_api) & set(dir(torch.Tensor)),
-                                     set(functional_api) & set(dir(torch.nn.functional)),
-                                     set(distributed_api) & set(dir(torch.distributed)),
-                                     set(npu_api) & set(dir(torch_npu)))
-    return wrap_api_name
-
-
-def wrap_api_func(api_name, api_dict, prefix, hook):
-    def api_function(*args, **kwargs):
-        return ApiTemplate(api_name, api_dict, prefix, hook)(*args, **kwargs)
-    return api_function
-
-
-def wrap_api_func_and_bind(api_list, api_dict, prefix, hook, hook_class):
-    for api_name in api_list:
-        if callable(api_dict[api_name]):
-            setattr(hook_class, Const.ATTR_NAME_PREFIX + api_name, wrap_api_func(api_name, api_dict, prefix, hook))
-
-
-def setup_hooks(hook):
-    if is_mindtorch():
-        torch_wrap_api_name = get_wrap_torch_api_list()
-        wrap_api_func_and_bind(torch_wrap_api_name.torch_api_names,
-                               {f: getattr(torch, f) for f in dir(torch)},
-                               MsConst.TORCH_DATA_PREFIX, hook, HOOKTorchOP)
-        wrap_api_func_and_bind(torch_wrap_api_name.tensor_api_names,
-                               {f: getattr(torch.Tensor, f) for f in dir(torch.Tensor)},
-                               MsConst.TENSOR_DATA_PREFIX, hook, HOOKTorchTensor)
-        wrap_api_func_and_bind(torch_wrap_api_name.functional_api_names,
-                               {f: getattr(torch.nn.functional, f) for f in dir(torch.nn.functional)},
-                               MsConst.OPS_DATA_PREFIX, hook, HOOKTorchFunctionalOP)
-        wrap_api_func_and_bind(torch_wrap_api_name.distributed_api_names,
-                               {f: getattr(torch.distributed, f) for f in dir(torch.distributed)},
-                               MsConst.DISTRIBUTED_DATA_PREFIX, hook, HOOKTorchDistributedOP)
-        wrap_api_func_and_bind(torch_wrap_api_name.npu_api_names, {f: getattr(torch_npu, f) for f in dir(torch_npu)},
-                               MsConst.TORCH_NPU_DATA_PREFIX, hook, HOOKTorchNpuOP)
-        return
-
-    wrap_api_name = get_wrap_api_list()
-    wrap_api_func_and_bind(wrap_api_name.tensor_api_names, {f: getattr(Tensor, f) for f in dir(Tensor)},
-                           MsConst.TENSOR_DATA_PREFIX, hook, HOOKTensor)
-    wrap_api_func_and_bind(wrap_api_name.stub_tensor_api_names, {f: getattr(StubTensor, f) for f in dir(StubTensor)},
-                           MsConst.STUB_TENSOR_DATA_PREFIX, hook, HOOKStubTensor)
-    wrap_api_func_and_bind(wrap_api_name.ops_api_names, {f: getattr(ops, f) for f in dir(ops)},
-                           MsConst.OPS_DATA_PREFIX, hook, HOOKFunctionalOP)
-    wrap_api_func_and_bind(wrap_api_name.mint_api_names, {f: getattr(mint, f) for f in dir(mint)},
-                           MsConst.MINT_DATA_PREFIX, hook, HOOKMintOP)
-    wrap_api_func_and_bind(wrap_api_name.mint_nn_func_api_names, {f: getattr(functional, f) for f in dir(functional)},
-                           MsConst.MINT_NN_FUNC_DATA_PREFIX, hook, HOOKMintNNFunctionalOP)
-    wrap_api_func_and_bind(wrap_api_name.distributed_api_names, {f: getattr(comm_func, f) for f in dir(comm_func)},
-                           MsConst.DISTRIBUTED_DATA_PREFIX, hook, HOOKDistributedOP)

msprobe/pytorch/hook_module/api_registry.py

@@ -1,166 +0,0 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
-# 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 torch
-import torch.distributed as dist
-
-from msprobe.pytorch.hook_module import wrap_torch, wrap_functional, wrap_tensor, wrap_vf, wrap_distributed, wrap_aten
-from msprobe.pytorch.hook_module.wrap_aten import get_aten_ops
-from msprobe.pytorch.hook_module.wrap_distributed import get_distributed_ops
-from msprobe.pytorch.hook_module.wrap_functional import get_functional_ops
-from msprobe.pytorch.hook_module.wrap_tensor import get_tensor_ops
-from msprobe.pytorch.hook_module.wrap_torch import get_torch_ops
-from msprobe.pytorch.hook_module.wrap_vf import get_vf_ops
-from msprobe.pytorch.common.utils import torch_without_guard_version, npu_distributed_api, is_gpu
-from msprobe.core.common.const import Const
-
-torch_version_above_2 = torch.__version__.split('+')[0] > '2.0'
-
-if not is_gpu:
-    import torch_npu
-    from . import wrap_npu_custom
-    from .wrap_npu_custom import get_npu_ops
-
-
-class ApiRegistry:
-    def __init__(self):
-        self.tensor_ori_attr = {}
-        self.torch_ori_attr = {}
-        self.functional_ori_attr = {}
-        self.distributed_ori_attr = {}
-        self.npu_distributed_ori_attr = {}
-        self.vf_ori_attr = {}
-        self.aten_ori_attr = {}
-        self.torch_npu_ori_attr = {}
-
-        self.tensor_hook_attr = {}
-        self.torch_hook_attr = {}
-        self.functional_hook_attr = {}
-        self.distributed_hook_attr = {}
-        self.npu_distributed_hook_attr = {}
-        self.vf_hook_attr = {}
-        self.aten_hook_attr = {}
-        self.torch_npu_hook_attr = {}
-
-    @staticmethod
-    def store_ori_attr(ori_api_group, api_list, api_ori_attr):
-        for api in api_list:
-            if '.' in api:
-                sub_module_name, sub_op = api.rsplit('.', 1)
-                sub_module = getattr(ori_api_group, sub_module_name)
-                api_ori_attr[api] = getattr(sub_module, sub_op)
-            else:
-                api_ori_attr[api] = getattr(ori_api_group, api)
-
-    @staticmethod
-    def set_api_attr(api_group, attr_dict):
-        for api, api_attr in attr_dict.items():
-            if '.' in api:
-                sub_module_name, sub_op = api.rsplit('.', 1)
-                sub_module = getattr(api_group, sub_module_name, None)
-                if sub_module is not None:
-                    setattr(sub_module, sub_op, api_attr)
-            else:
-                setattr(api_group, api, api_attr)
-
-    def api_modularity(self):
-        self.set_api_attr(torch.Tensor, self.tensor_hook_attr)
-        self.set_api_attr(torch, self.torch_hook_attr)
-        self.set_api_attr(torch.nn.functional, self.functional_hook_attr)
-        self.set_api_attr(dist, self.distributed_hook_attr)
-        self.set_api_attr(dist.distributed_c10d, self.distributed_hook_attr)
-        if not is_gpu and not torch_without_guard_version:
-            self.set_api_attr(torch_npu.distributed, self.npu_distributed_hook_attr)
-            self.set_api_attr(torch_npu.distributed.distributed_c10d, self.npu_distributed_hook_attr)
-        if torch_version_above_2:
-            self.set_api_attr(torch.ops.aten, self.aten_hook_attr)
-        self.set_api_attr(torch._VF, self.vf_hook_attr)
-        if not is_gpu:
-            self.set_api_attr(torch_npu, self.torch_npu_hook_attr)
-
-    def api_originality(self):
-        self.set_api_attr(torch.Tensor, self.tensor_ori_attr)
-        self.set_api_attr(torch, self.torch_ori_attr)
-        self.set_api_attr(torch.nn.functional, self.functional_ori_attr)
-        self.set_api_attr(dist, self.distributed_ori_attr)
-        self.set_api_attr(dist.distributed_c10d, self.distributed_ori_attr)
-        if not is_gpu and not torch_without_guard_version:
-            self.set_api_attr(torch_npu.distributed, self.npu_distributed_ori_attr)
-            self.set_api_attr(torch_npu.distributed.distributed_c10d, self.npu_distributed_ori_attr)
-        if torch_version_above_2:
-            self.set_api_attr(torch.ops.aten, self.aten_ori_attr)
-        self.set_api_attr(torch._VF, self.vf_ori_attr)
-        if not is_gpu:
-            self.set_api_attr(torch_npu, self.torch_npu_ori_attr)
-
-    def initialize_hook(self, hook, online_run_ut=False):
-        """
-        initialize_hook
-        Args:
-            hook (_type_): initialize_hook
-            online_run_ut (bool): default False, whether online run_ut or not.
-                If online_run_ut is True, the hook will not wrap the aten ops.
-        """
-        self.store_ori_attr(torch.Tensor, get_tensor_ops(), self.tensor_ori_attr)
-        wrap_tensor.wrap_tensor_ops_and_bind(hook)
-        for attr_name in dir(wrap_tensor.HOOKTensor):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                self.tensor_hook_attr[attr_name[5:]] = getattr(wrap_tensor.HOOKTensor, attr_name)
-
-        self.store_ori_attr(torch, get_torch_ops(), self.torch_ori_attr)
-        wrap_torch.wrap_torch_ops_and_bind(hook)
-        for attr_name in dir(wrap_torch.HOOKTorchOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                self.torch_hook_attr[attr_name[5:]] = getattr(wrap_torch.HOOKTorchOP, attr_name)
-
-        self.store_ori_attr(torch.nn.functional, get_functional_ops(), self.functional_ori_attr)
-        wrap_functional.wrap_functional_ops_and_bind(hook)
-        for attr_name in dir(wrap_functional.HOOKFunctionalOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                self.functional_hook_attr[attr_name[5:]] = getattr(wrap_functional.HOOKFunctionalOP, attr_name)
-
-        self.store_ori_attr(dist, get_distributed_ops(), self.distributed_ori_attr)
-        wrap_distributed.wrap_distributed_ops_and_bind(hook)
-        if not is_gpu and not torch_without_guard_version:
-            self.store_ori_attr(torch_npu.distributed, npu_distributed_api, self.npu_distributed_ori_attr)
-        for attr_name in dir(wrap_distributed.HOOKDistributedOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                self.distributed_hook_attr[attr_name[5:]] = getattr(wrap_distributed.HOOKDistributedOP, attr_name)
-                if not is_gpu and not torch_without_guard_version and attr_name[5:] in npu_distributed_api:
-                    self.npu_distributed_hook_attr[attr_name[5:]] = getattr(wrap_distributed.HOOKDistributedOP,
-                                                                            attr_name)
-
-        if torch_version_above_2 and not online_run_ut:
-            self.store_ori_attr(torch.ops.aten, get_aten_ops(), self.aten_ori_attr)
-            wrap_aten.wrap_aten_ops_and_bind(hook)
-            for attr_name in dir(wrap_aten.HOOKAtenOP):
-                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                    self.aten_hook_attr[attr_name[5:]] = getattr(wrap_aten.HOOKAtenOP, attr_name)
-
-        self.store_ori_attr(torch._VF, get_vf_ops(), self.vf_ori_attr)
-        wrap_vf.wrap_vf_ops_and_bind(hook)
-        for attr_name in dir(wrap_vf.HOOKVfOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                self.vf_hook_attr[attr_name[5:]] = getattr(wrap_vf.HOOKVfOP, attr_name)
-
-        if not is_gpu:
-            self.store_ori_attr(torch_npu, get_npu_ops(), self.torch_npu_ori_attr)
-            wrap_npu_custom.wrap_npu_ops_and_bind(hook)
-            for attr_name in dir(wrap_npu_custom.HOOKNpuOP):
-                if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                    self.torch_npu_hook_attr[attr_name[5:]] = getattr(wrap_npu_custom.HOOKNpuOP, attr_name)
-
-
-api_register = ApiRegistry()

msprobe/pytorch/hook_module/wrap_distributed.py

@@ -1,75 +0,0 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
-# 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 os
-from functools import wraps
-import torch.distributed as dist
-
-from msprobe.pytorch.hook_module.hook_module import HOOKModule
-from msprobe.pytorch.common.utils import torch_device_guard
-from msprobe.core.common.const import Const
-from msprobe.core.common.file_utils import load_yaml
-
-
-cur_path = os.path.dirname(os.path.realpath(__file__))
-yaml_path = os.path.join(cur_path, "support_wrap_ops.yaml")
-
-
-distributed_func = {}
-for f in dir(dist):
-    distributed_func[f] = getattr(dist, f)
-
-
-def get_distributed_ops():
-    _all_distributed_ops = dir(dist)
-    yaml_data = load_yaml(yaml_path)
-    wrap_distributed_ops = yaml_data.get('distributed')
-    return set(wrap_distributed_ops) & set(_all_distributed_ops)
-
-
-class HOOKDistributedOP(object):
-    pass
-
-
-class DistributedOPTemplate(HOOKModule):
-    def __init__(self, op_name, build_hook):
-        self.op_name_ = op_name
-        self.prefix_op_name_ = "Distributed" + Const.SEP + str(op_name) + Const.SEP
-        super().__init__(build_hook)
-        if not self.stop_hook:
-            self.op_is_distributed = True
-
-    @torch_device_guard
-    def forward(self, *args, **kwargs):
-        handle = distributed_func.get(self.op_name_)(*args, **kwargs)
-        if kwargs.get("async_op") or self.op_name_ in ["isend", "irecv"]:
-            if handle and hasattr(handle, 'wait'):
-                handle.wait()
-        return handle
-
-
-def wrap_distributed_op(op_name, hook):
-    @wraps(DistributedOPTemplate)
-    def distributed_op_template(*args, **kwargs):
-        return DistributedOPTemplate(op_name, hook)(*args, **kwargs)
-
-    distributed_op_template.__name__ = op_name
-    return distributed_op_template
-
-
-def wrap_distributed_ops_and_bind(hook):
-    _distributed_ops = get_distributed_ops()
-    for op_name in _distributed_ops:
-        setattr(HOOKDistributedOP, "wrap_" + str(op_name), wrap_distributed_op(op_name, hook))