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

msprobe/pytorch/api_accuracy_checker/tensor_transport_layer/device_dispatch.py

@@ -0,0 +1,204 @@
+import time
+from collections import namedtuple
+
+import pandas as pd
+import torch
+import torch.multiprocessing as mp
+
+from msprobe.core.common.const import Const
+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.common.log import logger
+from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.attl import move2target_device
+
+# NPU vs GPU api list
+CompareApi = set(absolute_standard_api) | set(binary_standard_api) | set(thousandth_standard_api)
+
+current_time = time.strftime("%Y%m%d%H%M%S")
+ONLINE_API_PRECISION_COMPARE_RESULT_FILE_NAME = "api_precision_compare_result_" + current_time + "_rank*.csv"
+ONLINE_API_PRECISION_COMPARE_DETAILS_FILE_NAME = "api_precision_compare_details_" + current_time + "_rank*.csv"
+
+OnlineApiPrecisionCompareConfig = namedtuple('OnlineApiPrecisionCompareConfig',
+                                             ['npu_data', 'gpu_data', 'rank', 'result_csv_path', 'details_csv_path'])
+# namedtuple of [instance of Comparator, func of run_touch_api_online, config of run_ut_config]
+CommonCompareConfig = namedtuple('CommonCompareConfig', ['compare', 'handle_func', 'config'])
+
+
+def run_ut_process(xpu_id, consumer_queue, common_config, api_precision_csv_file):
+    """ When consumer_queue(shared with ConsumerDispatcher) is not empty, consume api data from consumer_queue.
+    :param xpu_id: int
+    :param consumer_queue: shared queues of ConsumerDispatcher
+    :param common_config: namedtuple of CommonCompareConfig
+    :param api_precision_csv_file: list, length is 2, result file name and details file name
+    :return:
+    """
+    gpu_device = torch.device(f'cuda:{xpu_id}')
+
+    while True:
+        if consumer_queue.empty():
+            time.sleep(0.1)
+            continue
+
+        api_data = consumer_queue.get()
+        if api_data == "KILL_":
+            # current consumer finish
+            return
+
+        _, api_name, _ = api_data.name.split(Const.SEP)
+        if api_name in CompareApi:
+            # NPU vs GPU
+            online_compare(api_data, gpu_device, common_config)
+        else:
+            # NPUvsCPU vs GPUvsCPU
+            online_precision_compare(api_data, gpu_device, common_config, api_precision_csv_file)
+
+
+def online_precision_compare(api_data, device, common_config, api_precision_csv_file):
+    """online run_ut for precision_compare: NPUvsCPU vs GPUvsCPU
+    1. get NPUvsCPU compare result
+    2. get GPUvsCPU compare result
+    3. call online_api_precision_compare
+    :param api_data
+    :param device
+    :param common_config: namedtuple of CommonCompareConfig
+    :param api_precision_csv_file: [result_file_name, details_file_name]
+    """
+    compare, func, config = common_config.compare, common_config.handle_func, common_config.config
+    api_full_name = api_data.name
+    [api_type, api_name, _] = api_full_name.split(Const.SEP)
+    npu_args, npu_kwargs, npu_out = api_data.args, api_data.kwargs, api_data.result
+
+    if npu_kwargs.get("device"):
+        del npu_kwargs["device"]
+
+    try:
+        # NPU vs CPU
+        cpu_out = exec_api(api_type, api_name, npu_args, npu_kwargs)
+        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])
+
+        # GPU vs CPU
+        api_data_gpu = move2target_device(api_data, device)  # args, kwargs -> gpu, result -> npu
+        data_info = func(api_full_name, api_data_gpu, config.backward_content)
+        gpu_out = data_info.bench_output
+        gpu_data_info = UtDataInfo(None, None, gpu_out, cpu_out, None, [], None, rank=api_data.rank)
+        gpu_detail = compare.compare_output(api_full_name, gpu_data_info, True)
+        gpu_data = pd.DataFrame(gpu_detail, columns=DETAIL_TEST_ROWS[-1])
+
+        # NPUvsCPU vs GPUvsCPU
+        result_file_name, details_file_name = api_precision_csv_file
+        precision_compare_config = OnlineApiPrecisionCompareConfig(npu_data, gpu_data, api_data.rank,
+                                                                   result_file_name, details_file_name)
+        online_api_precision_compare(precision_compare_config)
+
+    except Exception as err:
+        if "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 "
+                f"'int32_to_int64' list in accuracy_tools/msprobe/core/common/const.py file.")
+        elif api_type in [Const.DISTRIBUTED]:
+            logger.info(f"{api_full_name} is not supported for run ut. SKIP.")
+        else:
+            logger.error(f"Run {api_full_name} UT Error: {str(err)}")
+
+        compare.write_summary_csv((api_full_name, "SKIP", "SKIP", [[str(err)]], api_data.rank))
+
+    finally:
+        torch.cuda.empty_cache()
+
+
+def online_compare(api_data, device, common_config):
+    """online run_ut for compare：NPU vs GPU
+    """
+    compare, func, config = common_config.compare, common_config.handle_func, common_config.config
+    api_full_name = api_data.name
+    api_data = move2target_device(api_data, device)
+    try:
+        data_info = func(api_full_name, api_data, config.backward_content)
+        is_fwd_success, is_bwd_success = compare.compare_output(api_full_name, data_info)
+        logger.info(f"running api_full_name {api_full_name} ut, "
+                    f"is_fwd_success: {is_fwd_success}, "
+                    f"is_bwd_success: {is_bwd_success}")
+    except Exception as err:
+        [api_type, api_name, _] = api_full_name.split(Const.SEP)
+        if "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 "
+                f"'int32_to_int64' list in accuracy_tools/msprobe/core/common/const.py file.")
+        elif api_type in [Const.DISTRIBUTED]:
+            logger.info(f"{api_full_name} is not supported for run ut. SKIP.")
+        else:
+            logger.error(f"Run {api_full_name} UT Error: {str(err)}")
+
+        compare.write_summary_csv((api_full_name, "SKIP", "SKIP", [[str(err)]], api_data.rank))
+
+    finally:
+        torch.cuda.empty_cache()
+
+
+class ConsumerDispatcher:
+    def __init__(self, compare, capacity=10, num_workers=8, device: str = "gpu") -> None:
+        self.num_workers = num_workers
+        self.capacity = capacity
+        self.compare = compare
+        self.queues = []
+        self.processes = []
+        self.reverse_sort = False
+        self.pool = None
+        self.device = device
+        self.data_id = 0
+        self.lock = mp.Lock()
+        self.result_queue = mp.Queue()
+        mp.set_start_method("spawn", force=True)
+
+    def start(self, handle_func, config):
+        self.queues = [mp.Queue(maxsize=self.capacity) for _ in range(self.num_workers)]
+        api_precision_csv_file = [ONLINE_API_PRECISION_COMPARE_RESULT_FILE_NAME, ONLINE_API_PRECISION_COMPARE_DETAILS_FILE_NAME]
+        common_config = CommonCompareConfig(self.compare, handle_func, config)
+        for xpu_id, q in enumerate(self.queues):
+            p = mp.Process(name="run_ut_process", target=run_ut_process,
+                           args=(xpu_id, q, common_config, api_precision_csv_file))
+
+            p.start()
+            self.processes.append(p)
+        logger.info(f"Api_precision_compare task result will be saved in {ONLINE_API_PRECISION_COMPARE_RESULT_FILE_NAME}")
+        logger.info(f"Api_precision_compare task details will be saved in {ONLINE_API_PRECISION_COMPARE_DETAILS_FILE_NAME}")
+        logger.info("Successfully start unittest process.")
+
+    def stop(self):
+        for q in self.queues:
+            while q.full():
+                time.sleep(0.1)
+            q.put("KILL_")
+
+        for p in self.processes:
+            p.join()
+        logger.info("Successfully stop unittest process.")
+        logger.info(f"Api_precision_compare task result is saved in {ONLINE_API_PRECISION_COMPARE_RESULT_FILE_NAME}")
+        logger.info(f"Api_precision_compare task details is saved in {ONLINE_API_PRECISION_COMPARE_DETAILS_FILE_NAME}")
+
+    def update_consume_queue(self, api_data):
+        while True:
+            index = self._choose_max_empty_site_strategy()
+            if index != -1:
+                q = self.queues[index]
+                q.put(api_data)
+                break
+            time.sleep(0.1)
+
+    def _choose_max_empty_site_strategy(self):
+        maximum = 0
+        index = -1
+        # 充分利用多卡资源，防止任务过多分配给前面的卡
+        _reverse = 1 if not self.reverse_sort else -1
+        for i, q in enumerate(self.queues[::_reverse]):
+            empty_site = self.capacity - q.qsize()
+            if empty_site > maximum:
+                maximum = empty_site
+                index = i
+        index = len(self.queues) - index - 1 if index != -1 and self.reverse_sort else index
+        self.reverse_sort = not self.reverse_sort
+        return index

msprobe/pytorch/api_accuracy_checker/tensor_transport_layer/server.py

@@ -0,0 +1,218 @@
+import os.path
+import struct
+import hashlib
+import time
+import io
+from threading import Thread
+
+from OpenSSL import SSL
+from twisted.internet import ssl, reactor, protocol, endpoints
+
+from msprobe.pytorch.common.utils import logger
+from msprobe.pytorch.api_accuracy_checker.tensor_transport_layer.ssl_config import cipher_list
+
+
+class TCPServer:
+    def __init__(self, port, shared_queue, check_sum=False, tls_path=None) -> None:
+        self.port = port
+        self.shared_queue = shared_queue
+        self.check_sum = check_sum
+        self.tls_path = tls_path
+        self.factory = MessageServerFactory()
+        self.reactor_thread = None
+
+    @staticmethod
+    def run_reactor():
+        reactor.run(installSignalHandlers=False)
+
+    def start(self):
+        self.factory.protocol = self.build_protocol
+
+        if self.tls_path:
+            server_key = os.path.join(self.tls_path, "server.key")
+            server_crt = os.path.join(self.tls_path, "server.crt")
+            server_context_factory = ssl.DefaultOpenSSLContextFactory(server_key, server_crt, SSL.TLSv1_2_METHOD)
+            server_context_ = server_context_factory.getContext()
+            server_context_.set_cipher_list(cipher_list)
+            server_context_.set_options(SSL.OP_NO_RENEGOTIATION)
+            endpoint = endpoints.SSL4ServerEndpoint(reactor, self.port, server_context_factory)
+        else:
+            endpoint = endpoints.TCP4ServerEndpoint(reactor, self.port)
+        endpoint.listen(self.factory)
+        self.reactor_thread = Thread(target=self.run_reactor, daemon=True)
+        self.reactor_thread.start()
+
+    def is_running(self):
+        return not self.factory.is_all_connection_closed()
+
+    def stop(self):
+        self.factory.doStop()
+        reactor.callFromThread(reactor.sigInt, 2)
+        self.reactor_thread.join()
+
+    def build_protocol(self):
+        return ServerProtocol(self.shared_queue, self.check_sum)
+
+
+class ServerProtocol(protocol.Protocol):
+    ACK_SUCCESS = b"OK___"
+    ACK_ERROR = b"ERROR"
+    ACK_BUSY = b"BUSY_"
+    ACK_STOP = b"STOP_"
+    ACK_STOP_CONFIRM = b"OVER_"
+    ACK_KILL_PROCESS = b"KILL_"
+
+    def __init__(self, shared_queue, check_sum=False):
+        self.start_time = None
+        self.buffer = io.BytesIO()
+        self.consumer_queue = shared_queue
+        self.check_sum = check_sum
+        self.length_width = 8
+        self.md5_width = 32
+        self.obj_length = None
+        self.tell = 0
+        self.obj_md5 = None
+        self.obj_body = None
+        self.sequence_number = -1
+        self.rank = -1
+        self.step = -1
+        self.sequence_number_dict = dict()
+
+    def connectionMade(self):
+        self.buffer = io.BytesIO()
+        self.obj_length = None
+        self.tell = 0
+        self.obj_md5 = None
+        self.obj_body = None
+        self.factory.transport_dict[self.transport] = 1
+        self.factory.transport_list.append(self.transport)
+        logger.info(f"Connected to {self.transport.getPeer()} successfully.")
+
+    def connectionLost(self, reason):
+        self.factory.transport_dict.pop(self.transport, None)
+        if len(self.factory.transport_dict) == 0:
+            self.consumer_queue.put(self.ACK_KILL_PROCESS)
+
+        logger.info(f"Lost connection with {self.transport.getPeer()}. Reason is: {reason} 与客户端 断开连接, "
+                    f"current connection number is: {len(self.factory.transport_dict)}")
+
+    def send_ack(self, ack_info):
+        ack_message = b"".join([
+            ack_info,
+            self.sequence_number.to_bytes(8, byteorder='big'),
+            self.rank.to_bytes(8, byteorder='big'),
+            self.step.to_bytes(8, byteorder='big')
+        ])
+        self.transport.write(ack_message)
+
+    def post_process(self):
+        send_busy_ack = False
+        while self.consumer_queue.full():
+            if not send_busy_ack:
+                self.send_ack(self.ACK_BUSY)
+                logger.debug("sending BUSY ACK")
+            send_busy_ack = True
+            time.sleep(0.1)
+
+        obj_key = str(self.sequence_number) + "_" + str(self.rank) + "_" + str(self.step)
+
+        recv_md5 = hashlib.md5(self.obj_body).hexdigest()
+        if self.check_sum and recv_md5 != self.obj_md5:
+            # when needs check md5 and check no pass, indicates received data error, send b"ERROR" to client.
+            logger.debug(f"Error:接收数据有问题，流水号{self.sequence_number}, expected {self.obj_md5}, but get {recv_md5}")
+            self.send_ack(self.ACK_ERROR)
+        else:
+            if self.obj_body == self.ACK_STOP:
+                self.handle_with_stop()
+            else:
+                self.send_ack(self.ACK_SUCCESS)
+            if obj_key in self.sequence_number_dict:
+                logger.debug(f"这是一次异常的重传，可以忽略。 {obj_key}, {self.sequence_number_dict}")
+            else:
+                self.sequence_number_dict[obj_key] = self.obj_md5
+                self.consumer_queue.put(self.obj_body, block=True)
+
+        self.reset_env()
+        finish_time = time.time()
+        logger.debug(f"finish_time: {finish_time - self.start_time}")
+
+    def handle_with_stop(self):
+        logger.debug(f"接收到停止传输信号 TCP{self.transport.getPeer()}")
+        self.send_ack(self.ACK_STOP_CONFIRM)
+        if len(self.factory.transport_dict) == 0:
+            _rank, _step, _sequence_number = 0, 0, 100000000
+            ack_kill = self.ACK_KILL_PROCESS + \
+                       _sequence_number.to_bytes(8, byteorder='big') + \
+                       _rank.to_bytes(8, byteorder='big') + \
+                       _step.to_bytes(8, byteorder='big')
+            for trans in self.factory.transport_list:
+                trans.write(ack_kill)
+            logger.debug(f"发送KILL信息给{self.transport.getPeer()}")
+            self.consumer_queue.put(self.ACK_KILL_PROCESS)
+            time.sleep(2)
+
+    def reset_env(self):
+        self.obj_length = None
+        self.sequence_number = -1
+        self.rank = -1
+        self.step = -1
+        self.obj_md5 = None
+        self.obj_body = None
+
+    def dataReceived(self, data):
+        self.buffer.seek(0, 2)
+        self.buffer.write(data)
+        self.buffer.seek(self.tell)
+
+        # The first data packet is packet header, it contains obj_length, sequence_number, rank, step
+        if self.obj_length is None and len(self.buffer.getvalue()) >= self.length_width * 4:
+            self.start_time = time.time()
+            self.obj_length = struct.unpack('!Q', self.buffer.read(self.length_width))[0]
+            self.sequence_number = struct.unpack('!Q', self.buffer.read(self.length_width))[0]
+            self.rank = struct.unpack('!Q', self.buffer.read(self.length_width))[0]
+            self.step = struct.unpack('!Q', self.buffer.read(self.length_width))[0]
+            self.tell += self.length_width * 4
+            logger.debug(
+                f"流水号: {self.sequence_number}; RANK: {self.rank}; STEP: {self.step}; Length: {self.obj_length}")
+
+        # If needs check md5 but not parse md5 yet, read 32b md5 values
+        check_sum_and_md5 = (self.check_sum
+                             and self.obj_length is not None
+                             and self.obj_md5 is None
+                             and len(self.buffer.getvalue()) - self.tell >= self.md5_width)
+        if check_sum_and_md5:
+            self.obj_md5 = self.buffer.read(self.md5_width).decode()
+            self.tell += self.md5_width
+            logger.debug(f"MD5: {self.obj_md5}")
+
+        current_length = len(self.buffer.getvalue()) - self.tell
+        if self.obj_length is not None and 0 < self.obj_length <= current_length:
+            # Current api data receive finished
+            self.obj_body = self.buffer.read(self.obj_length)
+
+            self.tell += self.obj_length
+            self.buffer = io.BytesIO(self.buffer.getvalue()[self.tell:])
+            self.buffer.seek(0)
+            self.tell = 0
+            recv_data_time = time.time()
+            logger.debug(f"self.sequence_number {self.sequence_number} "
+                         f"recv_data_time {recv_data_time - self.start_time}")
+
+            if self.obj_body == self.ACK_STOP:
+                # Indicates the current TCP link receives a STOP signal and remove from the transport_dict
+                _transport = self.factory.transport_dict.pop(self.transport, None)
+                logger.debug(f"接收到b'STOP_' self.sequence_number {self.sequence_number} ")
+            self.post_process()
+
+
+class MessageServerFactory(protocol.ServerFactory):
+    def __init__(self) -> None:
+        """
+        transport_dict: links that have not completed data transmission.
+        transport_list: Records all TCP links. Appends TCP link to the transport list when a new TCP link is established.
+        """
+        self.transport_dict = {}
+        self.transport_list = []
+
+    def is_all_connection_closed(self):
+        return len(self.transport_dict) == 0

msprobe/pytorch/api_accuracy_checker/tensor_transport_layer/ssl_config.py

@@ -0,0 +1,10 @@
+cipher_list = ":".join([
+    'ECDHE-ECDSA-AES128-GCM-SHA256',
+    'ECDHE-RSA-AES128-GCM-SHA256',
+    'ECDHE-ECDSA-AES256-GCM-SHA384',
+    'ECDHE-RSA-AES256-GCM-SHA384',
+    'ECDHE-ECDSA-CHACHA20-POLY1305',
+    'ECDHE-RSA-CHACHA20-POLY1305',
+    'DHE-RSA-AES128-GCM-SHA256',
+    'DHE-RSA-AES256-GCM-SHA384'
+]).encode()

msprobe/pytorch/bench_functions/__init__.py

@@ -0,0 +1,15 @@
+import os
+from pkgutil import iter_modules
+from importlib import import_module
+
+"""
+gpu and cpu not implement benchmark function, supplementary benchmarking function implementation
+"""
+
+package_path = os.path.dirname(os.path.realpath(__file__))
+for _, module_name, _ in iter_modules([package_path]):
+    module = import_module(f"{__name__}.{module_name}")
+    for attr_name in dir(module):
+        attr = getattr(module, attr_name)
+        if callable(attr) and "npu_custom" not in attr_name:
+            globals()[attr_name] = attr

msprobe/pytorch/bench_functions/apply_adam_w.py

@@ -0,0 +1,28 @@
+import torch
+
+
+def npu_apply_adam_w(beta1_power, beta2_power, lr, weight_decay,
+                     beta1, beta2, eps, grad, max_grad_norm, amsgrad, maximize, out):
+    var, m, v = out
+    if amsgrad:
+        max_grad_norm = (torch.rand(var.shape) * 10.0 - 5.0).to(var.dtype)
+    beta1_power_out = beta1_power * beta1
+    beta2_power_out = beta2_power * beta2
+    var_t = var * (1 + (-lr * weight_decay))
+    gt = -grad if maximize else grad
+    m_out = m * beta1 - (beta1 + (-1)) * gt
+    v_out = v * beta2 - (beta2 + (-1)) * gt * gt
+
+    if amsgrad:
+        max_grad_norm_out = torch.max(max_grad_norm, v_out)
+        if (1 - beta2_power_out) == 0:
+            beta2_power_out -= eps
+        denom = torch.sqrt(torch.div(max_grad_norm_out, (1 - beta2_power_out))) + eps
+    else:
+        vraintain = torch.div(v_out, (1 - beta2_power_out))
+        denom = torch.sqrt(vraintain) + eps
+
+    if (1 - beta1_power_out) == 0:
+        beta1_power_out -= eps
+    var_out = var_t + torch.div(-lr * m_out, (1 - beta1_power_out)).div(denom)
+    return var_out.cpu(), m_out.cpu(), v_out.cpu()

msprobe/pytorch/bench_functions/confusion_transpose.py

@@ -0,0 +1,19 @@
+def npu_confusion_transpose(data, perm, shape, transpose_first):
+    if transpose_first:
+        output = data.permute(*perm).contiguous().view(shape)
+    else:
+        output = data.view(shape).permute(*perm)
+    return output.cpu()
+
+
+def npu_confusion_transpose_backward(grad, perm, shape, transpose_first):
+    shape_cal = shape if transpose_first else [shape[perm_dim] for perm_dim in perm]
+    perm_cal = [0] * len(perm)
+    for i, perm_dim in enumerate(perm):
+        perm_cal[perm_dim] = i
+
+    if transpose_first:
+        result = grad.permute(*perm_cal).reshape(shape_cal)
+    else:
+        result = grad.reshape(shape_cal).permute(*perm_cal)
+    return result.cpu()

msprobe/pytorch/bench_functions/fast_gelu.py

@@ -0,0 +1,55 @@
+import torch
+
+
+def fast_gelu(input0):
+    attr = 1.702
+    const_0 = 0 - attr
+    const_1 = 1
+    const_2 = attr / 2
+
+    abs_x = torch.abs(input0)
+    mul_abs_x = abs_x * const_0
+    exp_abs_x = torch.exp(mul_abs_x)
+    div_down = exp_abs_x + const_1
+
+    pn_x = input0 - abs_x
+    mul_pn_x = pn_x * const_2
+    exp_pn_x = torch.exp(mul_pn_x)
+    div_up = input0 * exp_pn_x
+    div_down_rec = torch.reciprocal(div_down)
+    result = div_up * div_down_rec
+
+    return result.cpu()
+
+
+def npu_fast_gelu_backward(grad, input_x):
+    const_2 = 1.702
+    const_3 = 1.0
+    const_1 = 0.0 - const_2
+
+    # e^(-1.702x)
+    abs_x = torch.abs(input_x)
+    mul_abs_x = abs_x * const_1
+    exp_x = torch.exp(mul_abs_x)
+
+    # 1.702xe^(-1.702x)
+    add_2 = input_x * exp_x
+    add_2 = add_2 * const_2
+
+    # e^(1.702(x-|x|))
+    pn_x = input_x - abs_x
+    mul_pn_x = pn_x * const_2
+    exp_pn_x = torch.exp(mul_pn_x)
+
+    #  e^(-1.702x) + 1.702xe^(-1.702x) + e^(1.702(x-|x|))
+    div_up = exp_x + add_2
+    div_up = div_up + exp_pn_x
+
+    # (e^(-1.702x)+1)^2
+    div_down_i = exp_x + const_3
+    div_down = div_down_i * div_down_i
+    div_down_rec = torch.reciprocal(div_down)
+    result_temp = div_up * div_down_rec
+    result = grad * result_temp
+
+    return result.cpu()

msprobe/pytorch/bench_functions/layer_norm_eval.py

@@ -0,0 +1,6 @@
+import torch
+
+
+def npu_layer_norm_eval(data, normalized_shape):
+    result = torch.nn.functional.layer_norm(data, normalized_shape)
+    return result.cpu()

msprobe/pytorch/bench_functions/linear.py

@@ -0,0 +1,12 @@
+import torch
+
+
+def npu_linear(x, weight, bias):
+    output = torch.nn.functional.linear(x, weight, bias)
+    return output.cpu()
+
+
+def npu_linear_backward(grad, input_data, weight):
+    input_grad = torch.matmul(grad, weight)
+    weight_grad = torch.matmul(grad.t(), input_data)
+    return input_grad.cpu(), weight_grad.cpu()

msprobe/pytorch/bench_functions/matmul_backward.py

@@ -0,0 +1,48 @@
+import torch
+
+
+def matmul_backward(grad, self, other, mask):
+    grad_self, grad_other = None, None
+    dim_self = self.dim()
+    dim_other = other.dim()
+
+    size_grad = list(grad.size())
+    size_self = list(self.size())
+    size_other = list(other.size())
+    if dim_self == 1 and dim_other == 1:
+        grad_self = other.mul(grad) if mask[0] else grad_self
+        grad_other = self.mul(grad) if mask[1] else grad_other
+    elif dim_self == 2 and dim_other == 1:
+        grad_self = grad.unsqueeze(1).mm(other.unsqueeze(0)) if mask[0] else grad_self
+        grad_other = self.transpose(-1, -2).mm(grad.unsqueeze(1)).squeeze_(1) if mask[1] else grad_other
+    elif dim_self == 1 and dim_other == 2:
+        grad_self = grad.unsqueeze(0).mm(other.transpose(-1, -2)).squeeze_(0) if mask[0] else grad_self
+        grad_other = self.unsqueeze(1).mm(grad.unsqueeze(0)) if mask[1] else grad_other
+    elif dim_self >= 3 and (dim_other == 1 or dim_other == 2):
+        view_size = 1 if dim_other == 1 else size_grad[-1]
+        unfolded_grad = (grad.unsqueeze(-1) if dim_other == 1 else grad).contiguous().view(-1, view_size)
+        if mask[0]:
+            grad_self = unfolded_grad.mm(other.unsqueeze(0) if dim_other == 1 else other.transpose(-1, -2)) \
+                .view(size_self)
+        if mask[1]:
+            unfolded_self = self.contiguous().view([-1, size_self[-1]])
+            grad_other = unfolded_self.transpose(-1, -2).mm(unfolded_grad).view(size_other)
+    elif (dim_self == 1 or dim_self == 2) and dim_other >= 3:
+        view_size = 1 if dim_self == 1 else size_grad[-2]
+        unfolded_grad_T = grad.view([-1, view_size]) \
+            if dim_self == 1 else grad.transpose(-1, -2).contiguous().view([-1, view_size])
+        if mask[0]:
+            # create a 2D-matrix from other
+            unfolded_other_T = \
+                other.transpose(-1, -2).contiguous().view([-1, size_other[-2]]).transpose(-1, -2)
+            grad_self = unfolded_other_T.mm(unfolded_grad_T).transpose(-1, -2).view(size_self)
+        if mask[1]:
+            size_other_T = size_other[:-2]
+            size_other_T.extend(size_other[::-1][:2])
+            grad_other = \
+                unfolded_grad_T.mm(self.unsqueeze(0) if dim_self == 1 else self).view(size_other_T).transpose(-1, -2)
+    else:
+        grad_self = torch.matmul(grad, other.transpose(-1, -2)) if mask[0] else grad_self
+        grad_other = torch.matmul(self.transpose(-1, -2), grad) if mask[1] else grad_other
+
+    return grad_self.cpu(), grad_other.cpu()