mindstudio-probe 1.2.2__py3-none-any.whl → 8.1.0__py3-none-any.whl

msprobe/mindspore/dump/hook_cell/primitive_hooks.py

@@ -21,6 +21,7 @@ from mindspore.common.tensor import Tensor
 from msprobe.core.common.utils import Const, DumpException
 from msprobe.core.data_dump.data_processor.base import (ModuleBackwardInputs, ModuleBackwardOutputs,
                                                         ModuleForwardInputsOutputs)
+from msprobe.core.hook_manager import BaseHookManager
 from msprobe.mindspore.common.log import logger
 
 
@@ -58,7 +59,7 @@ class PrimitiveHookService:
             def backward_hook(grad):
                 captured_grads.extend(grad)
                 backward_primitive_name = f"{updated_primitive_name}{Const.SEP}{Const.BACKWARD}"
-
+                self.service_instance.inner_switch = True
                 try:
                     if hook_type == Const.INPUT:
                         self.service_instance.data_collector.update_api_or_module_name(backward_primitive_name)
@@ -77,6 +78,7 @@ class PrimitiveHookService:
                     logger.error(f"This is a primitive op {hook_type}_backward dump error: {exception}, "
                                  f"updated_primitive_name: {updated_primitive_name}")
                     raise DumpException(DumpException.BACKWARD_DATA_COLLECTION_ERROR) from exception
+                self.service_instance.inner_switch = False
 
             return backward_hook
 
@@ -137,6 +139,7 @@ class PrimitiveHookService:
 
         def pre_forward_hook(primitive_name, primitive_instance, args, kwargs):
             module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=None)
+            self.service_instance.inner_switch = True
             try:
                 self.service_instance.data_collector.forward_input_data_collect(
                     primitive_name,
@@ -148,9 +151,11 @@ class PrimitiveHookService:
                 logger.error(f"This is a primitive op dump error during forward input data collection: {exception}, "
                              f"primitive_name: {primitive_name}")
                 raise DumpException(DumpException.FORWARD_DATA_COLLECTION_ERROR) from exception
+            self.service_instance.inner_switch = False
 
         def post_forward_hook(primitive_name, primitive_instance, args, kwargs, output):
             module_input_output = ModuleForwardInputsOutputs(args=args, kwargs=kwargs, output=output)
+            self.service_instance.inner_switch = True
             try:
                 self.service_instance.data_collector.forward_output_data_collect(
                     primitive_name,
@@ -162,6 +167,7 @@ class PrimitiveHookService:
                 logger.error(f"This is a primitive op dump error during forward output data collection: {exception}, "
                              f"primitive_name: {primitive_name}")
                 raise DumpException(DumpException.FORWARD_DATA_COLLECTION_ERROR) from exception
+            self.service_instance.inner_switch = False
 
         def wrapped_primitive_call(instance_self, *args, **kwargs):
             """
@@ -179,7 +185,7 @@ class PrimitiveHookService:
             current_count = self.primitive_counters.get(primitive_name, 0)
             updated_primitive_name = f"{Const.PRIMITIVE_PREFIX}{Const.SEP}{primitive_name}{Const.SEP}{current_count}"
 
-            if not self.service_instance.primitive_switch:
+            if not self.service_instance.primitive_switch or BaseHookManager.inner_switch:
                 return origin_func(*args, **kwargs)
 
             captured_grads_input, captured_grads_output = [], []

msprobe/mindspore/dump/hook_cell/support_wrap_ops.yaml

@@ -564,15 +564,15 @@ tensor:
   - all
   - amax
   - amin
+  - angle
   - any
   - arccos
   - arccosh
-  - argmax
-  - angle
   - arcsin
   - arcsinh
   - arctan
   - arctanh
+  - argmax
   - argmin
   - argsort
   - asin
@@ -582,19 +582,23 @@ tensor:
   - atanh
   - baddbmm
   - bernoulli
+  - bfloat16
   - bincount
   - bitwise_and
   - bitwise_or
   - bitwise_xor
   - bmm
   - bool
+  - bool astype
   - broadcast_to
+  - byte
   - ceil
-  - cholesky_solve
   - cholesky
+  - cholesky_solve
   - clamp
   - clip
   - conj
+  - copy
   - copysign
   - cos
   - cosh
@@ -606,11 +610,13 @@ tensor:
   - deg2rad
   - diag
   - diagflat
+  - diagonal
   - diff
   - digamma
   - div
   - div_
   - divide
+  - double
   - equal
   - erf
   - erfc
@@ -618,13 +624,16 @@ tensor:
   - exp
   - expand_as
   - expm1
+  - flatten
   - flip
   - fliplr
   - flipud
+  - float
   - float_power
   - floor
   - fmod
   - frac
+  - from_numpy
   - gather_elements
   - ge
   - geqrf
@@ -648,12 +657,12 @@ tensor:
   - inner
   - int
   - inverse
+  - is_complex
+  - is_signed
   - isclose
   - isfinite
   - isinf
   - isnan
-  - is_complex
-  - is_signed
   - isneginf
   - isposinf
   - isreal
@@ -704,28 +713,27 @@ tensor:
   - new_ones
   - new_zeros
   - nextafter
-  - norm
   - nonzero
+  - norm
   - not_equal
   - ormqr
   - permute
   - pow
   - prod
   - qr
+  - rad2deg
   - ravel
   - real
   - reciprocal
   - remainder
   - renorm
-  - rad2deg
-  - tile
   - repeat_interleave
   - reshape
   - reshape
-  - round
+  - resize
   - rot90
+  - round
   - rsqrt
-  - sum_to_size
   - scatter
   - sgn
   - short
@@ -745,7 +753,8 @@ tensor:
   - sub
   - sub_
   - subtract
-  - subtract
+  - sum
+  - sum_to_size
   - svd
   - swapaxes
   - swapdims
@@ -753,13 +762,13 @@ tensor:
   - take
   - tan
   - tanh
-  - trace
-  - swapaxes
+  - tensor_split
   - tile
+  - to
   - topk
-  - tril
-  - tensor_split
+  - trace
   - transpose
+  - tril
   - true_divide
   - trunc
   - unbind
@@ -769,17 +778,6 @@ tensor:
   - view
   - where
   - xlogy
-  - from_numpy
-  - std
-  - take
-  - var
-  - all
-  - any
-  - copy
-  - diagonal
-  - flatten
-  - resize
-  - sum
 
 mint.ops:
   - abs
@@ -1027,3 +1025,21 @@ communication.comm_func:
   - recv
   - isend
   - irecv
+
+mint.distributed:
+  - send
+  - recv
+  - broadcast
+  - all_reduce
+  - reduce
+  - all_gather
+  - gather
+  - isend
+  - irecv
+  - scatter
+  - reduce_scatter
+  - all_to_all_single
+  - all_to_all
+  - all_gather_into_tensor
+  - reduce_scatter_tensor
+  - batch_isend_irecv

msprobe/mindspore/dump/jit_dump.py

@@ -13,9 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
 from collections import defaultdict
+import os
+import types
 
+import mindspore
+from mindspore import nn
 from mindspore._c_expression import PyNativeExecutor_
 try:
     from mindspore.common.api import _MindsporeFunctionExecutor
@@ -24,30 +27,31 @@ except ImportError:
 
 from msprobe.core.common.log import logger
 from msprobe.core.common.const import Const
+from msprobe.core.common.runtime import Runtime
 from msprobe.core.data_dump.data_processor.base import ModuleForwardInputsOutputs, ModuleBackwardInputsOutputs
-from msprobe.mindspore.dump.hook_cell.api_registry import api_register
+from msprobe.mindspore.common.const import Const as MsConst
+from msprobe.mindspore.dump.hook_cell.api_register import get_api_register
+
+
+_api_register = get_api_register()
 
 
 def dump_jit(name, in_feat, out_feat, is_forward):
     pid = os.getpid()
-    ori_args = str(name)
-    index = ori_args.find("<")
-    if index != 0 and index != -1:
-        result = ori_args[0:index]
-    elif name is not None and "<" not in str(name):
-        result = str(name)
-    else:
-        result = "JitFunction"
+    name = name if name else "JitFunction"
     if JitDump.need_dump():
         if is_forward:
-            JitDump.jit_count[result] += 1
-            name_template = (Const.JIT + Const.SEP + result + Const.SEP +
-                             str(JitDump.jit_count[result]) + Const.SEP + Const.FORWARD)
+            if name in JitDump.jit_count:
+                JitDump.jit_count[name] += 1
+            else:
+                JitDump.jit_count[name] = 0
+            name_template = (Const.JIT + Const.SEP + name + Const.SEP +
+                             str(JitDump.jit_count[name]) + Const.SEP + Const.FORWARD)
             JitDump.data_collector.update_api_or_module_name(name_template)
             module_input_output = ModuleForwardInputsOutputs(args=in_feat, kwargs={}, output=out_feat)
             JitDump.data_collector.forward_data_collect(name_template, None, pid, module_input_output)
         else:
-            name_template = Const.JIT + Const.SEP + result + Const.SEP + str(JitDump.jit_count[result]) + Const.SEP + \
+            name_template = Const.JIT + Const.SEP + name + Const.SEP + str(JitDump.jit_count[name]) + Const.SEP + \
                             Const.BACKWARD
             JitDump.data_collector.update_api_or_module_name(name_template)
             module_input_output = ModuleBackwardInputsOutputs(grad_input=in_feat, grad_output=out_feat)
@@ -57,7 +61,7 @@ def dump_jit(name, in_feat, out_feat, is_forward):
 class JitDump(_MindsporeFunctionExecutor):
     dump_config = None
     jit_enable = False
-    jit_dump_switch = True
+    jit_dump_switch = False
     jit_count = defaultdict(int)
 
     def __init__(self, *args, **kwargs):
@@ -68,19 +72,17 @@ class JitDump(_MindsporeFunctionExecutor):
         self._executor = PyNativeExecutor_.get_instance()
 
     def __call__(self, *args, **kwargs):
-        if JitDump.jit_dump_switch:
-            api_register.api_set_ori_func()
+        _api_register.restore_all_api()
         out = super().__call__(*args, **kwargs)
-        if JitDump.jit_dump_switch and len(args) > 0:
-            if self.name and self.name != "construct":
+        if JitDump.jit_dump_switch and len(args) > 0 and self.name:
+            if self.name != "construct":
                 dump_jit(self.name, args, out, True)
-            else:
-                dump_jit(args[0], args, out, True)
+            elif Runtime.run_mode != MsConst.PYNATIVE_GRAPH_MODE and isinstance(args[0], nn.Cell):
+                dump_jit(args[0].__class__.__name__, args, out, True)
             JitDump.jit_enable = True
         elif len(args) == 0:
             logger.warning(f"The jit function {self.name} has no input arguments, nothing will be dumped.")
-        if JitDump.jit_dump_switch:
-            api_register.api_set_hook_func()
+        _api_register.register_all_api()
         return out
 
     @classmethod
@@ -101,9 +103,15 @@ class JitDump(_MindsporeFunctionExecutor):
 
     def grad(self, obj, grad, weights, grad_position, *args, **kwargs):
         if JitDump.jit_dump_switch and JitDump.jit_enable:
-            api_register.api_set_ori_func()
-        output = self._executor.grad(grad, obj, weights, grad_position, *args, *(kwargs.values()))
+            _api_register.restore_all_api()
+        if mindspore.__version__ >= "2.5":
+            output = self._executor.grad(grad, obj, weights, grad_position, False, *args, *(kwargs.values()))
+        else:
+            output = self._executor.grad(grad, obj, weights, grad_position, *args, *(kwargs.values()))
         if JitDump.jit_dump_switch and JitDump.jit_enable:
-            dump_jit(obj, args, None, False)
-            api_register.api_set_hook_func()
+            if isinstance(obj, types.FunctionType):
+                dump_jit(obj.__name__, args, None, False)
+            elif Runtime.run_mode != MsConst.PYNATIVE_GRAPH_MODE and isinstance(obj, nn.Cell):
+                dump_jit(obj.__class__.__name__, args, None, False)
+            _api_register.register_all_api()
         return output

msprobe/mindspore/dump/kernel_kbyk_dump.py

@@ -39,9 +39,12 @@ class KernelKbykDump:
         common_set["input_output"] = 0
         common_set["kernels"] = []
         common_set["support_device"] = [0, 1, 2, 3, 4, 5, 6, 7]
-        e2e_set = dict()
-        e2e_set["enable"] = True
-        e2e_set["trans_flag"] = True
+        e2e_set = {
+                    "enable": not config.async_dump,
+                    "trans_flag": True,
+                    "stat_calc_mode": config.stat_cal_mode,
+                    "device_stat_precision_mode": config.device_stat_precision_mode,
+                  }
 
         if config.list:
             common_set["dump_mode"] = 1

msprobe/mindspore/dym_loader/hook_dynamic_loader.cpp

@@ -0,0 +1,110 @@
+/**
+ * Copyright 2024 Huawei Technologies Co., Ltd
+ *
+ * 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.
+ */
+
+#include "hook_dynamic_loader.h"
+#include <sys/stat.h>
+#include <cstdlib>
+#include <cstring>
+#include <pybind11/embed.h>
+#include "utils/log_adapter.h"
+
+namespace py = pybind11;
+
+HookDynamicLoader &HookDynamicLoader::GetInstance() 
+{
+    static HookDynamicLoader instance;
+    return instance;
+}
+
+bool HookDynamicLoader::LoadFunction(void *handle, const std::string &functionName) {
+    void *func = dlsym(handle, functionName.c_str());
+    if (!func) {
+        MS_LOG(WARNING) << "Could not load function: " << functionName << ", error: " << dlerror();
+        return false;
+    }
+    funcMap_[functionName] = func;
+    return true;
+}
+
+bool HookDynamicLoader::LoadLibrary() 
+{
+    std::string msprobePath = "";
+    // 获取gil锁
+    py::gil_scoped_acquire acquire;
+    try {
+        py::module msprobeMod = py::module::import("msprobe.lib._msprobe_c");
+        if (!py::hasattr(msprobeMod, "__file__")) {
+            MS_LOG(WARNING) << "Adump mod not found";
+            return false;
+    }
+    msprobePath = msprobeMod.attr("__file__").cast<std::string>();
+    } catch (const std::exception& e) {
+    MS_LOG(WARNING) << "Adump mod path unable to get: " << e.what();
+    return false;
+    }
+    std::lock_guard<std::mutex> lock(mutex_);
+    if (handle_) {
+        MS_LOG(WARNING) << "Hook library already loaded!";
+        return false;
+    }
+    if (msprobePath == "") {
+        MS_LOG(WARNING) << "Adump path not loaded";
+        return false;
+    }
+    handle_ = dlopen(msprobePath.c_str(), RTLD_LAZY | RTLD_LOCAL);
+    if (!handle_) {
+        MS_LOG(WARNING) << "Failed to load Hook library: " << dlerror();
+        return false;
+    }
+
+    for (const auto &functionName : functionList_) {
+        if (!LoadFunction(handle_, functionName)) {
+            MS_LOG(WARNING) << "Failed to load adump function";
+            dlclose(handle_);
+            handle_ = nullptr;
+            return false;
+        }
+    }
+
+    MS_LOG(INFO) << "Hook library loaded successfully.";
+    return true;
+}
+
+bool HookDynamicLoader::UnloadLibrary() 
+{
+    std::lock_guard<std::mutex> lock(mutex_);
+    if (!handle_) {
+        MS_LOG(WARNING) << "Hook library hasn't been loaded.";
+        return false;
+    }
+
+    dlclose(handle_);
+    handle_ = nullptr;
+    funcMap_.clear();
+    MS_LOG(INFO) << "Library unloaded successfully.";
+    return true;
+}
+
+void *HookDynamicLoader::GetHooker(const std::string &funcName) 
+{
+    std::lock_guard<std::mutex> lock(mutex_);
+    auto iter = funcMap_.find(funcName);
+    if (iter == funcMap_.end()) {
+      MS_LOG(WARNING) << "Function not found: " << funcName;
+      return nullptr;
+    }
+    return iter->second;
+}

msprobe/mindspore/dym_loader/hook_dynamic_loader.h

@@ -27,27 +27,26 @@ constexpr auto kHookBegin = "MS_DbgOnStepBegin";
 constexpr auto kHookEnd = "MS_DbgOnStepEnd";
 
 class HookDynamicLoader {
- public:
-  static HookDynamicLoader &GetInstance();
+public:
+    static HookDynamicLoader &GetInstance();
 
-  HookDynamicLoader(const HookDynamicLoader &) = delete;
-  HookDynamicLoader &operator=(const HookDynamicLoader &) = delete;
+    HookDynamicLoader(const HookDynamicLoader &) = delete;
+    HookDynamicLoader &operator=(const HookDynamicLoader &) = delete;
 
-  bool LoadLibrary();
-  bool UnloadLibrary();
-  void *GetHooker(const std::string &funcName);
+    bool LoadLibrary();
+    bool UnloadLibrary();
+    void *GetHooker(const std::string &funcName);
 
- private:
-  // Helper functions
-  bool loadFunction(void *handle, const std::string &functionName);
-  bool validateLibraryPath(const std::string &libPath);
+private:
+    // Helper functions
+    bool LoadFunction(void *handle, const std::string &functionName);
 
-  HookDynamicLoader() = default;
+    HookDynamicLoader() = default;
 
-  void *handle_ = nullptr;
-  std::vector<std::string> functionList_ = {kHookBegin, kHookEnd};
-  std::map<std::string, void *> funcMap_;
-  std::mutex mutex_;
+    void *handle_ = nullptr;
+    std::vector<std::string> functionList_ = {kHookBegin, kHookEnd};
+    std::map<std::string, void *> funcMap_;
+    std::mutex mutex_;
 };
 
 #endif  // HOOK_DYNAMIC_LOADER_H

msprobe/mindspore/free_benchmark/api_pynative_self_check.py

@@ -19,22 +19,27 @@ import os
 import traceback
 
 import mindspore as ms
+
 from msprobe.core.common.const import Const
 from msprobe.core.common.exceptions import DistributedNotInitializedError
 from msprobe.core.common.file_utils import check_path_length, load_yaml
+from msprobe.core.common.runtime import Runtime
+from msprobe.core.hook_manager import HookSet
 from msprobe.mindspore.common.const import Const as MsConst
 from msprobe.mindspore.common.const import FreeBenchmarkConst
 from msprobe.mindspore.common.log import logger
 from msprobe.mindspore.common.utils import get_rank_if_initialized
 from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
-from msprobe.mindspore.dump.hook_cell.api_registry import api_register
+from msprobe.mindspore.dump.hook_cell.api_register import get_api_register
 from msprobe.mindspore.dump.hook_cell.hook_cell import HOOKCell
 from msprobe.mindspore.free_benchmark.common.config import Config
 from msprobe.mindspore.free_benchmark.common.handler_params import HandlerParams
 from msprobe.mindspore.free_benchmark.common.utils import Tools
 from msprobe.mindspore.free_benchmark.handler.handler_factory import HandlerFactory
 from msprobe.mindspore.free_benchmark.perturbation.perturbation_factory import PerturbationFactory
-from msprobe.mindspore.runtime import Runtime
+
+
+_api_register = get_api_register()
 
 
 class ApiPyNativeSelfCheck:
@@ -60,8 +65,8 @@ class ApiPyNativeSelfCheck:
         self.store_original_func()
 
     def handle(self):
-        api_register.initialize_hook(self.build_hook)
-        api_register.api_set_hook_func()
+        _api_register.initialize_hook(self.build_hook)
+        _api_register.register_all_api()
 
     def build_hook(self, api_name):
         def pre_hook(cell, input_data):
@@ -71,7 +76,7 @@ class ApiPyNativeSelfCheck:
             ret = None
 
             if not need_wrapper_func():
-                del cell.input_kwargs
+                del cell.msprobe_input_kwargs
                 return ret
 
             api_name_with_id = api_name_with_id[:-1]
@@ -80,9 +85,9 @@ class ApiPyNativeSelfCheck:
                         api_name_with_id[api_name_with_id.find(Const.SEP) + 1:api_name_with_id.rfind(Const.SEP)])
             if api_name in self.api_list:
                 ret = check_self(api_name_with_id, output_data, self.ori_func.get(api_name),
-                                 *input_data, **cell.input_kwargs)
+                                 *input_data, **cell.msprobe_input_kwargs)
 
-            del cell.input_kwargs
+            del cell.msprobe_input_kwargs
             return ret
 
         def backward_hook(cell, grad_input, grad_output):
@@ -101,8 +106,13 @@ class ApiPyNativeSelfCheck:
 
         def pre_backward_hook(cell, grad_input):
             return None
-
-        return pre_hook, wrap_forward_hook, wrap_backward_hook, pre_backward_hook
+        
+        return HookSet(
+            forward_hook=wrap_forward_hook,
+            forward_pre_hook=pre_hook,
+            backward_hook=wrap_backward_hook,
+            backward_pre_hook=pre_backward_hook
+        )
 
     def store_original_func(self):
         for api_name in self.api_list:
@@ -166,13 +176,13 @@ def check_self(api_name_with_id, output, ori_func, *args, **kwargs):
         return ret
 
     logger.info(f"[{api_name_with_id}] is {Config.handler_type}ing.")
-    api_register.api_set_ori_func()
+    _api_register.restore_all_api()
 
     try:
         perturbation = PerturbationFactory.create(api_name_with_id)
         params.fuzzed_result = perturbation.handle(params)
         if params.fuzzed_result is False:
-            api_register.api_set_hook_func()
+            _api_register.register_all_api()
             return ret
         if Config.stage == Const.BACKWARD:
             params.original_result = Tools.get_grad(params.original_func, *params.args, **params.kwargs)
@@ -183,7 +193,7 @@ def check_self(api_name_with_id, output, ori_func, *args, **kwargs):
         logger.error(f"[{api_name_with_id}] Error: {str(e)}")
         logger.error(f"[{api_name_with_id}] Error detail: {traceback.format_exc()}")
 
-    api_register.api_set_hook_func()
+    _api_register.register_all_api()
     return ret
 
 

msprobe/mindspore/free_benchmark/common/utils.py

@@ -19,10 +19,10 @@ from typing import Any, Optional
 import mindspore as ms
 from mindspore import Tensor, ops
 
+from msprobe.core.common.runtime import Runtime
 from msprobe.mindspore.common.const import FreeBenchmarkConst
 from msprobe.mindspore.free_benchmark.common.config import Config
 from msprobe.mindspore.free_benchmark.common.handler_params import HandlerParams
-from msprobe.mindspore.runtime import Runtime
 
 
 class Tools:

msprobe/mindspore/free_benchmark/perturbation/perturbation_factory.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# Copyright (c) 2024-2025, Huawei Technologies Co., Ltd.
 # All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0  (the "License");
@@ -14,6 +14,7 @@
 # limitations under the License.
 
 from msprobe.mindspore.common.const import FreeBenchmarkConst
+from msprobe.mindspore.common.log import logger
 from msprobe.mindspore.free_benchmark.common.config import Config
 from msprobe.mindspore.free_benchmark.perturbation.add_noise import AddNoisePerturbation
 from msprobe.mindspore.free_benchmark.perturbation.bit_noise import BitNoisePerturbation
@@ -41,4 +42,5 @@ class PerturbationFactory:
         if perturbation:
             return perturbation(api_name_with_id)
         else:
-            raise Exception(f'{Config.pert_type} is a invalid perturbation type')
+            logger.error(f'{Config.pert_type} is a invalid perturbation type')
+            raise ValueError