mindstudio-probe 1.0.3__py3-none-any.whl → 1.1.0__py3-none-any.whl

msprobe/mindspore/dump/hook_cell/api_registry.py

@@ -1,104 +1,128 @@
-# 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.
-# ============================================================================
-
-import mindspore as ms
-from msprobe.mindspore.dump.hook_cell.wrap_functional import get_functional_ops, setup_hooks, \
-    HOOKFunctionalOP, HOOKMintOP, HOOKMintNNFunctionalOP
-from msprobe.mindspore.dump.hook_cell.wrap_tensor import get_tensor_ops, wrap_tensor_ops_and_bind, HOOKTensor
-from msprobe.core.common.utils import Const
-
-
-class ApiRegistry:
-    def __init__(self):
-        self.tensor_ori_attr = {}
-        self.functional_ori_attr = {}
-        self.mint_ops_ori_attr = {}
-        self.mint_func_ops_ori_attr = {}
-        self.norm_inner_ops_ori_attr = {}
-
-        self.tensor_hook_attr = {}
-        self.functional_hook_attr = {}
-        self.mint_ops_hook_attr = {}
-        self.mint_func_ops_hook_attr = {}
-        self.norm_inner_ops_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)
-                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 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(ms.ops, self.norm_inner_ops_hook_attr)
-
-    def norm_inner_op_set_ori_func(self):
-        self.set_api_attr(ms.ops, self.norm_inner_ops_ori_attr)
-
-    def api_set_hook_func(self):
-        self.set_api_attr(ms.Tensor, self.tensor_hook_attr)
-        self.set_api_attr(ms.ops, self.functional_hook_attr)
-        self.set_api_attr(ms.mint, self.mint_ops_hook_attr)
-        self.set_api_attr(ms.mint.nn.functional, self.mint_func_ops_hook_attr)
-
-    def api_set_ori_func(self):
-        self.set_api_attr(ms.Tensor, self.tensor_ori_attr)
-        self.set_api_attr(ms.ops, self.functional_ori_attr)
-        self.set_api_attr(ms.mint, self.mint_ops_ori_attr)
-        self.set_api_attr(ms.mint.nn.functional, self.mint_func_ops_ori_attr)
-
-    def initialize_hook(self, hook):
-        self.store_ori_attr(ms.Tensor, get_tensor_ops(), self.tensor_ori_attr)
-        wrap_tensor_ops_and_bind(hook)
-        for attr_name in dir(HOOKTensor):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                self.tensor_hook_attr[attr_name[Const.ATTR_NAME_PREFIX_LEN:]] = getattr(HOOKTensor, attr_name)
-
-        functional_ops, mint_ops, mint_func_ops = get_functional_ops()
-        self.store_ori_attr(ms.ops, self.norm_inner_ops, self.norm_inner_ops_ori_attr)
-        self.store_ori_attr(ms.ops, functional_ops, self.functional_ori_attr)
-        self.store_ori_attr(ms.mint, mint_ops, self.mint_ops_ori_attr)
-        self.store_ori_attr(ms.mint.nn.functional, mint_func_ops, self.mint_func_ops_ori_attr)
-        setup_hooks(hook)
-        for attr_name in dir(HOOKFunctionalOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                self.functional_hook_attr[attr_name[Const.ATTR_NAME_PREFIX_LEN:]] = getattr(HOOKFunctionalOP, attr_name)
-                if attr_name[Const.ATTR_NAME_PREFIX_LEN:] in self.norm_inner_ops:
-                    self.norm_inner_ops_hook_attr[attr_name[Const.ATTR_NAME_PREFIX_LEN:]] = getattr(HOOKFunctionalOP, attr_name)
-        for attr_name in dir(HOOKMintOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                self.mint_ops_hook_attr[attr_name[Const.ATTR_NAME_PREFIX_LEN:]] = getattr(HOOKMintOP, attr_name)
-        for attr_name in dir(HOOKMintNNFunctionalOP):
-            if attr_name.startswith(Const.ATTR_NAME_PREFIX):
-                self.mint_func_ops_hook_attr[attr_name[Const.ATTR_NAME_PREFIX_LEN:]] = getattr(HOOKMintNNFunctionalOP, attr_name)
-
-
-api_register = ApiRegistry()
+# 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.
+# ============================================================================
+
+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,
+                                                       get_wrap_api_list, setup_hooks)
+from msprobe.core.common.utils import Const
+
+
+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.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.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)
+                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 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):
+        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):
+        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):
+        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)
+        setup_hooks(hook)
+        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/hook_cell.py

@@ -1,53 +1,55 @@
-# 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.
-# ============================================================================
-from collections import defaultdict
-
-from mindspore import nn
-from msprobe.core.common.const import Const
-
-
-class HOOKCell(nn.Cell):
-    cell_count = defaultdict(int)
-    g_stop_hook = False
-
-    def __init__(self, build_hook) -> None:
-        super(HOOKCell, self).__init__()
-        self.changed_status = False
-        self.input_kwargs = {}
-        self.prefix = ""
-        if not HOOKCell.g_stop_hook:
-            HOOKCell.g_stop_hook = True
-            self.changed_status = True
-            if hasattr(self, "prefix_op_name_"):
-                self.prefix = self.prefix_op_name_
-
-            HOOKCell.cell_count[self.prefix] += 1
-            self.prefix = self.prefix + str(HOOKCell.cell_count[self.prefix] - 1) + Const.SEP
-            forward_hook, backward_hook = build_hook(self.prefix)
-            self.register_forward_hook(forward_hook)
-            self.register_backward_hook(backward_hook)
-
-    # 重载call，加全局标志。
-    def __call__(self, *args, **kwargs):
-        try:
-            self.input_kwargs = kwargs
-            out = super(HOOKCell, self).__call__(*args, **kwargs)
-        except Exception as e:
-            raise e
-        finally:
-            if self.changed_status:
-                self.changed_status = False
-                HOOKCell.g_stop_hook = False
-        return out
+# 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.
+# ============================================================================
+
+from collections import defaultdict
+
+from mindspore import nn
+
+from msprobe.core.common.const import Const
+
+
+class HOOKCell(nn.Cell):
+    cell_count = defaultdict(int)
+    g_stop_hook = False
+
+    def __init__(self, build_hook) -> None:
+        super(HOOKCell, self).__init__()
+        self.changed_status = False
+        self.input_kwargs = {}
+        self.prefix = ""
+        if not HOOKCell.g_stop_hook:
+            HOOKCell.g_stop_hook = True
+            self.changed_status = True
+            if hasattr(self, "prefix_api_name"):
+                self.prefix = self.prefix_api_name
+
+            HOOKCell.cell_count[self.prefix] += 1
+            self.prefix = self.prefix + str(HOOKCell.cell_count[self.prefix] - 1) + Const.SEP
+            forward_hook, backward_hook = build_hook(self.prefix)
+            self.register_forward_hook(forward_hook)
+            self.register_backward_hook(backward_hook)
+
+    # 重载call，加全局标志。
+    def __call__(self, *args, **kwargs):
+        try:
+            self.input_kwargs = kwargs
+            out = super(HOOKCell, self).__call__(*args, **kwargs)
+        except Exception as e:
+            raise e
+        finally:
+            if self.changed_status:
+                self.changed_status = False
+                HOOKCell.g_stop_hook = False
+        return out

msprobe/mindspore/dump/hook_cell/primitive_hooks.py

@@ -0,0 +1,206 @@
+# 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.
+# ============================================================================
+
+import os
+
+import mindspore as ms
+from mindspore.common.tensor import Tensor
+from mindspore import ops
+
+from msprobe.mindspore.common.log import logger
+from msprobe.core.common.utils import Const, DumpException
+from msprobe.core.data_dump.data_processor.base import ModuleBackwardInputsOutputs, ModuleForwardInputsOutputs, \
+    ModuleBackwardInputs, ModuleBackwardOutputs
+
+
+class PrimitiveHookService:
+    def __init__(self, service_instance):
+        self.primitive_counters = {}
+        self.service_instance = service_instance
+
+    def wrap_primitive(self, origin_func, primitive_name):
+        """
+        包装原始的 primitive 函数，添加输入和输出的 hook 以捕获前向和反向数据。
+
+        Args:
+            origin_func (callable): 原始 的 primitive 函数。
+            primitive_name (str): 原始的 primitive 名称。
+
+        Returns:
+            callable: 包装后的 primitive 函数。
+        """
+        def create_backward_hook(captured_grads, num_tensors, updated_primitive_name, hook_type):
+            """
+            创建反向 hook 函数，用于捕获梯度。
+
+            Args:
+                captured_grads (list): 用于保存捕获的梯度。
+                num_tensors (int): 张量数量。
+                updated_primitive_name (str): 更新后的 primitive 名称。
+                hook_type (str): hook 类型 (输入/输出)。
+
+            Returns:
+                callable: 反向 hook 函数。
+            """
+            def backward_hook(grad):
+
+                captured_grads.append(grad)
+                backward_primitive_name = f"{updated_primitive_name}{Const.SEP}{Const.BACKWARD}"
+
+                try:
+                    if len(captured_grads) == num_tensors and hook_type == Const.INPUT:
+                        self.service_instance.data_collector.update_api_or_module_name(backward_primitive_name)
+                        new_module_input_output = ModuleBackwardOutputs(grad_output=tuple(captured_grads))
+                        self.service_instance.data_collector.backward_output_data_collect(
+                            backward_primitive_name, self, os.getpid(), new_module_input_output
+                        )
+                        captured_grads.clear()
+                    elif len(captured_grads) == num_tensors and hook_type == Const.OUTPUT:
+                        self.service_instance.data_collector.update_api_or_module_name(backward_primitive_name)
+                        new_module_input_output = ModuleBackwardInputs(grad_input=tuple(captured_grads))
+                        self.service_instance.data_collector.backward_input_data_collect(
+                            backward_primitive_name, self, os.getpid(), new_module_input_output
+                        )
+                        captured_grads.clear()
+
+                except Exception as exception:
+                    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
+
+            return backward_hook
+
+        def hook_primitive_inputs(args, captured_grads_input, updated_primitive_name):
+            """
+            针对前向输入添加 hook。
+
+            Args:
+                args (tuple): primitive 输入参数。
+                captured_grads_input (list): 捕获的输入梯度。
+                updated_primitive_name (str): 更新后的 primitive 名称。
+
+            Returns:
+                list: 添加了 hook 的输入。
+            """
+            hooked_inputs = []
+            num_tensors = sum(isinstance(arg, Tensor) for arg in args)
+            input_backward_hook = create_backward_hook(captured_grads_input, num_tensors, updated_primitive_name,
+                                                       Const.INPUT)
+            for arg in args:
+                if isinstance(arg, Tensor):
+                    arg_hooked = ops.HookBackward(input_backward_hook)(arg)
+                    hooked_inputs.append(arg_hooked)
+                else:
+                    hooked_inputs.append(arg)
+            return hooked_inputs
+
+        def hook_primitive_outputs(out, captured_grads_output, updated_primitive_name):
+            """
+            针对前向输出添加 hook。
+
+            Args:
+                out (Tensor/tuple): primitive 输出。
+                captured_grads_output (list): 捕获的输出梯度。
+                updated_primitive_name (str): 更新后的 primitive 名称。
+
+            Returns:
+                Tensor/tuple: 添加了 hook 的输出。
+            """
+            if isinstance(out, tuple):
+                num_output_tensors = sum(isinstance(tensor, Tensor) for tensor in out)
+            else:
+                num_output_tensors = 1
+            output_backward_hook = create_backward_hook(captured_grads_output, num_output_tensors,
+                                                        updated_primitive_name, Const.OUTPUT)
+
+            if isinstance(out, Tensor):
+                return ops.HookBackward(output_backward_hook)(out)
+            elif isinstance(out, tuple):
+                hooked_outputs = []
+                for tensor in out:
+                    if isinstance(tensor, Tensor):
+                        hooked_outputs.append(ops.HookBackward(output_backward_hook)(tensor))
+                    else:
+                        hooked_outputs.append(tensor)
+                return tuple(hooked_outputs)
+            return out
+
+        def wrapped_primitive_call(instance_self, *args, **kwargs):
+            """
+            包装后的 primitive 调用函数，添加输入和输出的 hook。
+
+            Args:
+                instance_self (object): primitive 的实例。
+                *args: primitive 输入参数。
+                **kwargs: primitive 关键字参数。
+
+            Returns:
+                Tensor/tuple: primitive 的返回值。
+            """
+            self.update_primitive_counters(primitive_name)
+            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:
+                return origin_func(*args, **kwargs)
+
+            captured_grads_input, captured_grads_output = [], []
+
+            try:
+                hooked_inputs = hook_primitive_inputs(args, captured_grads_input, updated_primitive_name)
+            except Exception as exception:
+                logger.error(f"This is a primitive op dump error during input hooking: {exception}, "
+                             f"primitive_name: {primitive_name}")
+                raise DumpException(DumpException.INPUT_HOOK_ERROR) from exception
+
+            try:
+                out = origin_func(*hooked_inputs, **kwargs)
+            except Exception as exception:
+                logger.error(f"This is a primitive op dump error during function call: {exception}, "
+                             f"primitive_name: {primitive_name}")
+                raise DumpException(DumpException.FUNCTION_CALL_ERROR) from exception
+
+            forward_primitive_name = f"{updated_primitive_name}{Const.SEP}{Const.FORWARD}"
+            self.service_instance.data_collector.update_api_or_module_name(forward_primitive_name)
+            if self.service_instance.data_collector:
+                module_input_output = ModuleForwardInputsOutputs(args=hooked_inputs, kwargs=kwargs, output=out)
+                try:
+                    self.service_instance.data_collector.forward_data_collect(forward_primitive_name, instance_self,
+                                                             os.getpid(), module_input_output)
+                except Exception as exception:
+                    logger.error(f"This is a primitive op dump error during forward data collection: {exception}, "
+                                 f"primitive_name: {primitive_name}")
+                    raise DumpException(DumpException.FORWARD_DATA_COLLECTION_ERROR) from exception
+
+                if self.service_instance.data_collector.if_return_forward_new_output():
+                    out = self.service_instance.data_collector.get_forward_new_output()
+
+            try:
+                out = hook_primitive_outputs(out, captured_grads_output, updated_primitive_name)
+            except Exception as exception:
+                logger.error(f"This is a primitive op dump error during output hooking: {exception}, "
+                             f"primitive_name: {primitive_name}")
+                raise DumpException(DumpException.OUTPUT_HOOK_ERROR) from exception
+
+            return out
+
+        return wrapped_primitive_call
+
+    def update_primitive_counters(self, primitive_name):
+        if primitive_name not in self.primitive_counters:
+            self.primitive_counters[primitive_name] = 0
+        else:
+            self.primitive_counters[primitive_name] += 1
+