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

msprobe/pytorch/api_accuracy_checker/run_ut/multi_run_ut.py

@@ -41,6 +41,7 @@ from msprobe.core.common.utils import CompareException
 
 def split_json_file(input_file, num_splits, filter_api):
     forward_data, backward_data, real_data_path = parse_json_info_forward_backward(input_file)
+    input_dir = os.path.dirname(os.path.abspath(input_file))
     if filter_api:
         forward_data = preprocess_forward_content(forward_data)
     for data_name in list(forward_data.keys()):
@@ -71,7 +72,7 @@ def split_json_file(input_file, num_splits, filter_api):
                 **backward_data
             }
         }
-        split_filename = f"temp_part{i}.json"
+        split_filename = os.path.join(input_dir, f"temp_part{i}.json")
         save_json(split_filename, temp_data)
         split_files.append(split_filename)
 

msprobe/pytorch/api_accuracy_checker/run_ut/run_overflow_check.py

@@ -23,12 +23,14 @@ try:
     import torch_npu
 except ImportError:
     is_gpu = True
+    current_device = "cuda"
 else:
     is_gpu = False
+    current_device = "npu"
 import torch
 from tqdm import tqdm
 from msprobe.pytorch.api_accuracy_checker.run_ut.run_ut import generate_device_params, get_api_info
-from msprobe.pytorch.api_accuracy_checker.run_ut.run_ut_utils import exec_api, is_unsupported_api
+from msprobe.pytorch.api_accuracy_checker.run_ut.run_ut_utils import exec_api, is_unsupported_api, ExecParams
 from msprobe.core.common.file_utils import check_link, FileChecker
 from msprobe.pytorch.api_accuracy_checker.common.utils import extract_basic_api_segments
 from msprobe.core.common.const import FileCheckConst, Const
@@ -61,19 +63,33 @@ def check_tensor_overflow(x):
         return False
 
 
-def check_data_overflow(x):
-    if isinstance(x, (tuple, list)) and x:
-        for _, item in enumerate(x):
-            if check_data_overflow(item):
-                return True
-        return False
+def check_data_overflow(x, device):
+    if isinstance(x, (tuple, list)):
+        if not x:
+            return False
+        return any(check_data_overflow(item, device) for item in x)
     else:
-        return check_tensor_overflow(x)
+        if device == Const.CPU_LOWERCASE:
+            return check_tensor_overflow(x)
+        else:
+            return torch_npu.npu.utils.npu_check_overflow(x)
+
+
+def is_bool_output(x):
+    if isinstance(x, (tuple, list)):
+        if not x:
+            return False
+        return any(is_bool_output(item) for item in x)
+    else:
+        return isinstance(x, bool)
 
 
 def run_overflow_check(forward_file):
     logger.info("start UT test")
     forward_content, _, real_data_path = parse_json_info_forward_backward(forward_file)
+    if real_data_path:
+        dump_path = os.path.dirname(forward_file)
+        real_data_path = os.path.join(dump_path, Const.DUMP_TENSOR_DATA)
     for api_full_name, api_info_dict in tqdm(forward_content.items()):
         if is_unsupported_api(api_full_name, is_overflow_check=True):
             continue
@@ -87,6 +103,9 @@ def run_overflow_check(forward_file):
             elif "expected scalar type Long" in str(err):
                 logger.warning(f"API {api_name} not support int32 tensor in CPU, please add {api_name} to CONVERT_API "
                                "'int32_to_int64' list in accuracy_tools/msprobe/core/common/const.py file.")
+            elif "could not create a primitive descriptor for a matmul primitive" in str(err):
+                logger.warning(f"API {api_name} not support matmul primitive in CPU due to pytorch bug, "
+                               "so it will be skipped.")
             else:
                 logger.error(f"Run {api_full_name} UT Error: %s" % str(err))
 
@@ -98,17 +117,26 @@ def run_torch_api(api_full_name, api_info_dict, real_data_path):
     if not need_grad:
         logger.warning("%s function with out=... arguments don't support automatic differentiation, skip backward." 
                        % api_full_name)
+    device_info_kwargs = kwargs.get(Const.DEVICE)
+    if device_info_kwargs and device_info_kwargs.get(Const.VALUE):
+        kwargs[Const.DEVICE] = current_device
     npu_args, npu_kwargs = generate_device_params(args, kwargs, False, api_name)
-    if kwargs.get("device"):
-        del kwargs["device"]
-    out = exec_api(api_type, api_name, Const.CPU_LOWERCASE, args, kwargs)
-    npu_out = exec_api(api_type, api_name, Const.NPU_LOWERCASE, npu_args, npu_kwargs)
+    if kwargs.get(Const.DEVICE):
+        del kwargs[Const.DEVICE]
+    cpu_exec_params = ExecParams(api_type, api_name, Const.CPU_LOWERCASE, args, kwargs, False, None)
+    device_exec_params = ExecParams(api_type, api_name, Const.NPU_LOWERCASE, npu_args, npu_kwargs, False, None)
+    out = exec_api(cpu_exec_params)
+    npu_out = exec_api(device_exec_params)
     if out is None and npu_out is None:
         logger.warning("The %s overflow is a normal overflow, out and npu_out is None." % api_full_name)
         return
+    if is_bool_output(out) or is_bool_output(npu_out):
+        logger.warning("The output of %s is bool type.This dtype not support overflow, so it will be skipped."
+                       % api_full_name)
+        return
 
-    cpu_overflow = check_data_overflow(out)
-    npu_overflow = torch_npu.npu.utils.npu_check_overflow(npu_out)
+    cpu_overflow = check_data_overflow(out, Const.CPU_LOWERCASE)
+    npu_overflow = check_data_overflow(npu_out, Const.NPU_LOWERCASE)
     if cpu_overflow == npu_overflow:
         logger.warning("The %s overflow is a normal overflow." % api_full_name)
     else:

msprobe/pytorch/api_accuracy_checker/run_ut/run_ut.py

@@ -31,6 +31,7 @@ except ImportError:
 else:
     is_gpu = False
     current_device = "npu"
+
 import torch
 from tqdm import tqdm
 
@@ -48,10 +49,12 @@ from msprobe.core.common.file_utils import FileChecker, change_mode, \
 from msprobe.pytorch.common.log import logger
 from msprobe.pytorch.pt_config import parse_json_config
 from msprobe.core.common.const import Const, FileCheckConst, CompareConst
-from msprobe.core.common.utils import safe_get_value
+from msprobe.core.common.utils import safe_get_value, CompareException
+from msprobe.pytorch.common.utils import seed_all
 from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.attl import ATTL, ATTLConfig, move2device_exec
 from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.device_dispatch import ConsumerDispatcher
-from msprobe.pytorch.api_accuracy_checker.run_ut.run_ut_utils import generate_cpu_params, generate_device_params
+from msprobe.pytorch.api_accuracy_checker.run_ut.run_ut_utils import generate_cpu_params, generate_device_params, \
+    ExecParams
 
 
 current_time = time.strftime("%Y%m%d%H%M%S")
@@ -61,6 +64,7 @@ DETAILS_FILE_NAME = "accuracy_checking_details_" + current_time + ".csv"
 
 
 not_backward_list = ['repeat_interleave']
+unsupported_backward_list = ['masked_select']
 
 
 tqdm_params = {
@@ -237,7 +241,8 @@ def run_torch_api(api_full_name, real_data_path, backward_content, api_info_dict
     in_fwd_data_list = []
     backward_message = ''
     api_type, api_name = extract_basic_api_segments(api_full_name)
-    args, kwargs, need_grad = get_api_info(api_info_dict, api_name, real_data_path)
+    args, kwargs, output_dtype = get_api_info(api_info_dict, api_name, real_data_path)
+    need_grad = check_need_grad(api_info_dict)
     in_fwd_data_list.append(args)
     in_fwd_data_list.append(kwargs)
     need_backward = api_full_name in backward_content
@@ -248,14 +253,30 @@ def run_torch_api(api_full_name, real_data_path, backward_content, api_info_dict
         need_grad = False
         logger.info("%s %s" % (api_full_name, BackwardMessage.NO_BACKWARD_RESULT_MESSAGE))
         backward_message += BackwardMessage.NO_BACKWARD_RESULT_MESSAGE
+    if api_name in unsupported_backward_list:
+        need_grad = False
+        logger.info("%s %s" % (api_full_name, BackwardMessage.UNSUPPORT_API_MESSAGE))
+        backward_message += BackwardMessage.UNSUPPORT_API_MESSAGE
     need_backward = need_backward and need_grad
-    if kwargs.get("device"):
-        del kwargs["device"]
-    cpu_args, cpu_kwargs = generate_cpu_params(args, kwargs, need_backward, api_name)
+    
+    device_info_kwargs = kwargs.get(Const.DEVICE)
+    if device_info_kwargs and device_info_kwargs.get(Const.VALUE):
+        kwargs[Const.DEVICE] = current_device
     device_args, device_kwargs = generate_device_params(args, kwargs, need_backward, api_name)
+    if kwargs.get(Const.DEVICE):
+        del kwargs[Const.DEVICE]
+    cpu_params = generate_cpu_params(args, kwargs, need_backward, api_name)
+    cpu_args, cpu_kwargs = cpu_params.cpu_args, cpu_params.cpu_kwargs
+    autocast_dtype, is_autocast = cpu_params.autocast_dtype, cpu_params.is_autocast
+    if not is_autocast and output_dtype:
+        is_autocast = autocast_dtype != output_dtype
+        autocast_dtype = output_dtype
     bench_grad_out, device_grad_out = None, None
-    out = exec_api(api_type, api_name, Const.CPU_LOWERCASE, cpu_args, cpu_kwargs)
-    device_out = exec_api(api_type, api_name, current_device, device_args, device_kwargs)
+    cpu_exec_params = ExecParams(api_type, api_name, Const.CPU_LOWERCASE, cpu_args, cpu_kwargs, False, autocast_dtype)
+    out = exec_api(cpu_exec_params)
+    device_exec_params = ExecParams(api_type, api_name, current_device, device_args, device_kwargs, is_autocast,
+                                     autocast_dtype)
+    device_out = exec_api(device_exec_params)
     current_path = os.path.dirname(os.path.realpath(__file__))
     ut_setting_path = os.path.join(current_path, "torch_ut_setting.json")
     api_setting_dict = get_json_contents(ut_setting_path)
@@ -273,7 +294,8 @@ def run_torch_api(api_full_name, real_data_path, backward_content, api_info_dict
             }
             grad = gen_args(backward_args, api_name, func_options)
             grad = safe_get_value(grad, 0, "grad")
-            bench_grad, _ = generate_cpu_params(grad, {}, False, api_name)
+            grad_params = generate_cpu_params(grad, {}, False, api_name)
+            bench_grad = grad_params.cpu_args
             bench_grad_out = run_backward(cpu_args, bench_grad, grad_index, out)
             device_grad = grad.clone().detach().to(current_device)
             device_grad_out = run_backward(device_args, device_grad, grad_index, device_out)
@@ -300,13 +322,18 @@ def run_torch_api_online(api_full_name, api_data, backward_content):
     return UtDataInfo(None, None, out, device_out, None, in_fwd_data_list, None, rank=api_data.rank)
 
 
-def get_api_info(api_info_dict, api_name, real_data_path):
-    convert_type, api_info_dict = api_info_preprocess(api_name, api_info_dict)
+def check_need_grad(api_info_dict):
     need_grad = True
-    if api_info_dict.get("input_kwargs") and "out" in api_info_dict.get("input_kwargs"):
+    if api_info_dict.get(Const.INPUT_KWARGS) and "out" in api_info_dict.get(Const.INPUT_KWARGS):
         need_grad = False
-    args, kwargs = gen_api_params(api_info_dict, api_name, need_grad, convert_type, real_data_path)
-    return args, kwargs, need_grad
+    return need_grad
+
+
+def get_api_info(api_info_dict, api_name, real_data_path):
+    convert_type, api_info_dict = api_info_preprocess(api_name, api_info_dict)
+    need_grad = check_need_grad(api_info_dict)
+    args, kwargs, output_dtype = gen_api_params(api_info_dict, api_name, need_grad, convert_type, real_data_path)
+    return args, kwargs, output_dtype
 
 
 def need_to_backward(grad_index, out):
@@ -323,15 +350,25 @@ def run_backward(args, grad, grad_index, out):
         out[grad_index].backward(grad)
     else:
         out.backward(grad)
-    args_grad = []
-    for arg in args:
-        if isinstance(arg, torch.Tensor):
-            args_grad.append(arg.grad)
-    grad_out = args_grad
+
+    grad_out = extract_tensors_grad(args)
 
     return grad_out
 
 
+def extract_tensors_grad(args, depth=0):
+    if depth > Const.MAX_DEPTH:
+        logger.error("The depth of arg_in is too large, please check the arg_in.")
+        raise CompareException(CompareException.RECURSION_LIMIT_ERROR)
+    grads = []
+    for arg in args:
+        if isinstance(arg, torch.Tensor):
+            grads.append(arg.grad)
+        elif isinstance(arg, (list, tuple)):
+            grads.extend(extract_tensors_grad(arg, depth+1))
+    return grads
+
+
 def initialize_save_error_data(error_data_path):
     create_directory(error_data_path)
     error_data_path_checker = FileChecker(error_data_path, FileCheckConst.DIR,
@@ -479,6 +516,10 @@ def run_ut_command(args):
 
     if not is_gpu:
         torch.npu.set_compile_mode(jit_compile=args.jit_compile)
+        if args.jit_compile:
+            torch.npu.config.allow_internal_format = True
+        else:
+            torch.npu.config.allow_internal_format = False
     used_device = current_device + ":" + str(args.device_id[0])
     try:
         if is_gpu:
@@ -497,6 +538,9 @@ def run_ut_command(args):
                                             ability=FileCheckConst.READ_ABLE, file_type=FileCheckConst.JSON_SUFFIX)
         checked_api_info = api_info_file_checker.common_check()
         forward_content, backward_content, real_data_path = parse_json_info_forward_backward(checked_api_info)
+        if real_data_path:
+            dump_path = os.path.dirname(checked_api_info)
+            real_data_path = os.path.join(dump_path, Const.DUMP_TENSOR_DATA)
         if args.filter_api:
             logger.info("Start filtering the api in the api_info_file.")
             forward_content = preprocess_forward_content(forward_content)
@@ -538,5 +582,6 @@ def run_ut_command(args):
 
 
 if __name__ == '__main__':
+    seed_all()
     _run_ut()
     logger.info("UT task completed.")

msprobe/pytorch/api_accuracy_checker/run_ut/run_ut_utils.py

@@ -16,6 +16,7 @@
 # limitations under the License.
 
 import os
+from collections import namedtuple
 import re
 import torch
 
@@ -23,8 +24,10 @@ try:
     import torch_npu
 except ImportError:
     current_device = "cuda"
+    from torch.cuda.amp import autocast
 else:
     current_device = "npu"
+    from torch_npu.npu.amp import autocast
 
 from msprobe.core.common.const import FileCheckConst, Const, CompareConst
 from msprobe.core.common.file_utils import FileChecker
@@ -47,11 +50,17 @@ PRECISION_MAPPING = {
 }
 
 
+CpuParams = namedtuple("CpuArgs", ["cpu_args", "cpu_kwargs", "autocast_dtype", "is_autocast"])
+ExecParams = namedtuple("ExecParams", ["api_type", "api_name", "device", "args", "kwargs", 
+                                       "is_autocast", "autocast_dtype"])
+
+
 class BackwardMessage:
     MULTIPLE_BACKWARD_MESSAGE = "Multiple backward is not supported."
     UNSUPPORT_BACKWARD_MESSAGE = "function with out=... arguments don't support automatic differentiation, " \
                                   "skip backward."
     NO_BACKWARD_RESULT_MESSAGE = "This API does not have backward input data, skip backward."
+    UNSUPPORT_API_MESSAGE = "This API does not support backward ut, skip backward."
 
 
 class UtDataInfo:
@@ -91,7 +100,15 @@ def get_validated_details_csv_path(validated_result_csv_path):
     return validated_details_csv_path
 
 
-def exec_api(api_type, api_name, device, args, kwargs):
+def exec_api(exec_params):
+    api_type = exec_params.api_type
+    api_name = exec_params.api_name
+    device = exec_params.device
+    args = exec_params.args
+    kwargs = exec_params.kwargs
+    is_autocast = exec_params.is_autocast
+    autocast_dtype = exec_params.autocast_dtype
+
     if api_type == "Functional":
         torch_api = FunctionalOPTemplate(api_name, str, False)
     if api_type == "Tensor":
@@ -102,7 +119,11 @@ def exec_api(api_type, api_name, device, args, kwargs):
         torch_api = AtenOPTemplate(api_name, None, False)
     if api_type == "NPU":
         torch_api = NpuOPTemplate(api_name, None, False, device)
-    out = torch_api.forward(*args, **kwargs)
+    if is_autocast:
+        with autocast(dtype=autocast_dtype):
+            out = torch_api.forward(*args, **kwargs)
+    else:
+        out = torch_api.forward(*args, **kwargs)
     return out
 
 
@@ -196,19 +217,28 @@ def generate_cpu_params(input_args, input_kwargs, need_backward, api_name):
         return set()
 
     raise_dtype = None
+    autocast_dtype = None
+    is_autocast = False
     need_raise_dtypes = recursive_find_dtypes(input_args)
     need_raise_dtypes.update(recursive_find_dtypes(input_kwargs, check_kwargs=True))
     if len(need_raise_dtypes) == 1:
-        raise_dtype = PRECISION_MAPPING.get(need_raise_dtypes.pop(), torch.float32)
+        origin_dtype = need_raise_dtypes.pop()
+        raise_dtype = PRECISION_MAPPING.get(origin_dtype, torch.float32)
+        autocast_dtype = origin_dtype
+        
     elif len(need_raise_dtypes) >= 2:
         raise_dtype = torch.float32
+        need_raise_dtypes.discard(torch.float32)
+        autocast_dtype = need_raise_dtypes.pop()
+        is_autocast = True
 
     raise_dtype = None if api_name in not_raise_dtype_set else raise_dtype
     is_detach = api_name not in not_detach_set
     cpu_args = recursive_arg_to_cpu(input_args, is_detach, raise_dtype=raise_dtype)
     cpu_kwargs = {key: recursive_arg_to_cpu(value, key != "out" and is_detach, raise_dtype=raise_dtype) for
                   key, value in input_kwargs.items()}
-    return cpu_args, cpu_kwargs
+    cpu_params = CpuParams(cpu_args, cpu_kwargs, autocast_dtype, is_autocast)
+    return cpu_params
 
 
 def record_skip_info(api_full_name, compare, compare_alg_results):

msprobe/pytorch/api_accuracy_checker/tensor_transport_layer/device_dispatch.py

@@ -24,7 +24,7 @@ from msprobe.core.common.const import Const, CompareConst
 from msprobe.pytorch.api_accuracy_checker.compare.api_precision_compare import online_api_precision_compare
 from msprobe.pytorch.api_accuracy_checker.compare.compare_utils import DETAIL_TEST_ROWS, thousandth_standard_api, \
     binary_standard_api, absolute_standard_api
-from msprobe.pytorch.api_accuracy_checker.run_ut.run_ut_utils import UtDataInfo, exec_api
+from msprobe.pytorch.api_accuracy_checker.run_ut.run_ut_utils import UtDataInfo, exec_api, ExecParams
 from msprobe.pytorch.common.log import logger
 from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.attl import move2target_device
 from msprobe.pytorch.api_accuracy_checker.run_ut.run_ut_utils import generate_cpu_params
@@ -92,8 +92,10 @@ def online_precision_compare(api_data, device, common_config, api_precision_csv_
 
     try:
         # NPU vs CPU
-        cpu_args, cpu_kwargs = generate_cpu_params(npu_args, npu_kwargs, False, api_name)
-        cpu_out = exec_api(api_type, api_name, Const.CPU_LOWERCASE, cpu_args, cpu_kwargs)
+        cpu_params = generate_cpu_params(npu_args, npu_kwargs, False, api_name)
+        cpu_args, cpu_kwargs = cpu_params.cpu_args, cpu_params.cpu_kwargs
+        cpu_exec_params = ExecParams(api_type, api_name, Const.CPU_LOWERCASE, cpu_args, cpu_kwargs, False, None)
+        cpu_out = exec_api(cpu_exec_params)
         npu_data_info = UtDataInfo(None, None, npu_out, cpu_out, None, [], None, rank=api_data.rank)
         npu_detail = compare.compare_output(api_full_name, npu_data_info, True)
         npu_data = pd.DataFrame(npu_detail, columns=DETAIL_TEST_ROWS[-1])

msprobe/pytorch/bench_functions/apply_adam.py

@@ -0,0 +1,215 @@
+# 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 collections import namedtuple
+import torch
+
+
+VarParams = namedtuple('VarParams', ['var', 'lr_t', 'm_t', 'beta1_broad', 'grad', 'epsilon', 'v_t'])
+
+
+def _output_m_compute(m, beta1_broad, grad):
+    """
+    _output_m_compute
+    do compute m_t = m + (beta1 - 1) * (m - grad)
+    """
+    input_dtype = m.dtype
+
+    sneg_one = torch.ones((1), dtype=input_dtype) * -1
+    sneg_one = sneg_one.to(beta1_broad.device)
+
+    # `formula; beta1 -1`
+    vsub_beta1_1 = torch.add(beta1_broad, sneg_one)
+
+    # `formula; m - grad`
+    vsub_m_grad = torch.sub(m, grad)
+
+    # `formula; (beta1 - 1) * (m - grad)`
+    vmul_m = torch.mul(vsub_beta1_1, vsub_m_grad)
+
+    # `formula; m_t = m + (beta1 - 1) * (m - grad)`
+    m_t = torch.add(m, vmul_m)
+
+    return m_t
+
+
+def _output_v_compute(v, beta2, grad):
+    """
+    _output_v_compute
+    do compute v_t = v + (1 - beta2)*(grad*grad -v)
+    """
+    input_dtype = v.dtype
+
+    sneg_one = torch.ones((1), dtype=input_dtype) * -1
+
+    # `formula; broadcast beta2 to vector`
+    beta2_tensor = torch.tensor(beta2, dtype=input_dtype)
+    beta2_broad = beta2_tensor.expand_as(v)
+
+    # `formula; beta2 - 1`
+    vsub_beta2_1 = torch.add(beta2_broad, sneg_one)
+    vsub_beta2_1 = vsub_beta2_1.to(v.device)
+
+    # `formula; grad * grad`
+    vmul_grad_grad = torch.mul(grad, grad)
+
+    # `formula; (v - grad*grad)`
+    vsub_v_grad = torch.sub(v, vmul_grad_grad)
+
+    # `formula; (beta2 -1) * (v - grad * grad)`
+    vmul_grad = torch.mul(vsub_beta2_1, vsub_v_grad)
+
+    # `formula; v_t = v + (beta2 - 1) * (v - grad * grad)`
+    v_t = torch.add(v, vmul_grad)
+
+    return v_t
+
+
+def _inner_lr_compute(lr, beta2_power, beta1_power, compute_shape_tensor):
+    """
+    _inner_lr_compute
+    `formula; lr_t = learning_rate * (sqrt(1-beta2_power)) / (1 - beta1_power)`
+    """
+
+    input_dtype = compute_shape_tensor.dtype
+
+    s_one = torch.ones((1), dtype=input_dtype)
+
+    s_neg_one = torch.ones((1), dtype=input_dtype) * -1
+
+    # `formula; (1 - beta2_power)`
+    v_neg_beta2_power = torch.mul(beta2_power, s_neg_one)
+    v_add_beta2_power = torch.add(v_neg_beta2_power, s_one)
+
+    # `formula; sqrt(1 - beta2_power)`
+    v_sqrt_beta2_power = torch.sqrt(v_add_beta2_power)
+
+    # `formula; (1 - beta1_power)`
+    v_neg_beta1_power = torch.mul(beta1_power, s_neg_one)
+    v_add_beta1_power = torch.add(v_neg_beta1_power, s_one)
+
+    # `formula; learning_rate * (sqrt(1-beta2_power)`
+    res = torch.mul(lr, v_sqrt_beta2_power)
+
+    # `formula; learning_rate*(sqrt(1-beta2_power))/(1-beta1_power)`
+    res = torch.div(res, v_add_beta1_power)
+    return res.expand_as(compute_shape_tensor)
+
+
+def _inner_eps_add_sqrt_vt_compute(epsilon, v_t):
+    """
+    (epsilon + sqrt(v_t) )
+    """
+    # `formula; sqrt(v_t)`
+    sqrt_vt = torch.sqrt(v_t)
+
+    # `formula; broadcast epsilon  to vector`
+    input_dtype = v_t.dtype
+    epsilon_tensor = torch.tensor(epsilon, dtype=input_dtype)
+    epsilon_broad = epsilon_tensor.expand_as(v_t)
+    epsilon_broad = epsilon_broad.to(sqrt_vt.device)
+
+    # `formula; epsilon + sqrt(v_t)`
+    v_add_sqrt_v = torch.add(sqrt_vt, epsilon_broad)
+
+    return v_add_sqrt_v
+
+
+def _output_var_t_compute_use_nesterov(varparams):
+    """
+    _output_var_t_compute_use_nesterov
+    `formula; var_t = var - lr_t * (m_t * beta1 + (1 - beta1) * grad) / (epsilon + sqrt(v_t))`
+    `formula; var_t = var - lr_t * (m_t * beta1 + (1 - beta1) * grad) / (epsilon + sqrt(v_t))`
+    """
+    var = varparams.var
+    lr_t = varparams.lr_t
+    m_t = varparams.m_t
+    beta1_broad = varparams.beta1_broad
+    grad = varparams.grad
+    epsilon = varparams.epsilon
+    v_t = varparams.v_t
+
+    input_dtype = var.dtype
+
+    s_one = torch.ones((1), dtype=input_dtype)
+
+    s_neg_one = torch.ones((1), dtype=input_dtype) * -1
+
+    # `formula; m_t * beta1`
+    v_muls_mt_beta1 = torch.mul(m_t, beta1_broad)
+
+    # `formula; 1 -beta1`
+    v_neg_beta1 = torch.mul(beta1_broad, s_neg_one)
+    vsub_1_beta1 = torch.add(v_neg_beta1, s_one)
+
+    # `formula; (1-beta1)* grad`
+    v_mul_grad = torch.mul(vsub_1_beta1, grad)
+
+    # `formula; (m_t*beta1 + (1 - beta1)*grad)`
+    v_div_left = torch.add(v_muls_mt_beta1, v_mul_grad)
+
+    # `formula; lr_t * (m_t*beta1 + (1 - beta1) * grad)`
+    # broadcast lr_t to vector
+
+    lrt_broad = lr_t.expand_as(var)
+    v_mul_left = torch.mul(lrt_broad, v_div_left)
+
+    # `formula; (epsilon + sqrt(v_t))`
+    v_add_sqrt_v = _inner_eps_add_sqrt_vt_compute(epsilon, v_t)
+
+    # `formula; lr_t * (m_t*beta1 + (1-beta1)*grad / (epsilon + sqrt(v_t))`
+    v_div_res = torch.div(v_mul_left, v_add_sqrt_v)
+
+    # `formula; var - lr_t * (m_t*beta1 + (1-beta1)*grad) / (epsilon + sqrt(v_t))`
+    v_t = torch.sub(var, v_div_res)
+
+    return v_t
+
+
+def _output_var_t_compute(var, lr_t, m_t, epsilon, v_t):
+    """
+    _output_var_t_compute
+    `var_t = var - lr_t * m_t / (epsilon + sqrt(v_t))`
+    """
+    # `formula; lr_t * m_t`
+    lr_t = lr_t.to(m_t.device)
+    v_mul_left = torch.mul(lr_t, m_t)
+
+    # `formula; (epsilon + sqrt(v_t))`
+    v_add_sqrt_v = _inner_eps_add_sqrt_vt_compute(epsilon, v_t)
+
+    # `formula; lr_t * m_t /(epsilon + sqrt(v_t))`
+    v_div_res = torch.div(v_mul_left, v_add_sqrt_v)
+
+    # `formula; var - lr_t * m_t / (epsilon + sqrt(v_t))`
+    v_t = torch.sub(var, v_div_res)
+
+    return v_t
+
+
+def npu_apply_adam(beta1_power, beta2_power, lr, beta1, beta2, epsilon, grad, use_locking, use_nesterov, out):
+    var, m, v = out
+    input_dtype = m.dtype
+    beta1_tensor = torch.tensor(beta1, dtype=input_dtype).to(m.device)
+    beta1_broad = beta1_tensor.expand_as(m)
+    m_t = _output_m_compute(m, beta1_broad, grad)
+    v_t = _output_v_compute(v, beta2, grad)
+    lr_t = _inner_lr_compute(lr, beta2_power, beta1_power, grad)
+    if use_nesterov:
+        var_params = VarParams(var, lr_t, m_t, beta1_broad, grad, epsilon, v_t)
+        var_t = _output_var_t_compute_use_nesterov(var_params)
+    else:
+        var_t = _output_var_t_compute(var, lr_t, m_t, epsilon, v_t)
+    return var_t, m_t, v_t

msprobe/pytorch/bench_functions/group_norm_silu.py

@@ -0,0 +1,27 @@
+# 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 torch
+
+
+def npu_group_norm_silu(x, gama, beta, group, eps):
+    if len(x.shape) != 4:
+        raise ValueError("x shape should be (N, C, H, W)")
+    res = torch.ops.aten.native_group_norm(x, gama, beta, x.shape[0], x.shape[1], x.shape[2] * x.shape[3], group, eps)
+    res = list(res)
+    if not res:
+        raise ValueError("run native_group_norm failed")
+    res[0] = torch.nn.functional.silu(res[0])
+    return res

msprobe/pytorch/bench_functions/mish.py

@@ -0,0 +1,21 @@
+# 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 torch
+
+
+def npu_mish(x):
+    mish = torch.nn.Mish()
+    return mish(x)