mindspore 2.1.0__cp38-none-any.whl → 2.2.10__cp38-none-any.whl

mindspore/lib/plugin/ascend/custom_aicore_ops/op_impl/vector_core/tbe/custom_aicore_ops_impl/kv_cache_mgr.py

@@ -0,0 +1,241 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+# Copyright 2023 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.
+# ============================================================================
+"""ascend custom op: kv_cache_mgr by tik"""
+
+import functools
+from tbe import tik
+import tbe.common.platform as tbe_platform
+from tbe.common.utils import para_check
+
+
+# 'pylint: disable=unused-argument,unused-variable,too-many-arguments,too-many-locals
+def check_supported(past, cur, index, out, kernel_name="kv_cache_mgr"):
+    """check data type and shape"""
+    # check data type
+    past_dtype = past.get("dtype").lower()
+    cur_dtype = cur.get("dtype").lower()
+    out_dtype = out.get("dtype").lower()
+
+    if past_dtype != cur_dtype or past_dtype != out_dtype:
+        reason = "past_dtype is %s, cur_dtype is %s, out_dtype is %s" % (past_dtype, cur_dtype, out_dtype)
+        return False, reason
+
+    support_dtype_list = ["float32", "int32", "uint32",
+                          "float16", "int16", "uint16",
+                          "int8", "uint8"]
+    if past_dtype not in support_dtype_list:
+        reason = "past_dtype(%s) is not support" % (past_dtype)
+        return False, reason
+
+    index_dtype = index.get("dtype").lower()
+    if index_dtype != "int32":
+        reason = "index_dtype is %s, not int32" % (index_dtype)
+        return False, reason
+
+    # check shape
+    past_shape = past.get("shape")
+    cur_shape = cur.get("shape")
+
+    if len(past_shape) != 4 or len(cur_shape) != 4:
+        reason = "len(past_shape) != 4 or len(cur_shape) != 4 "
+        return False, reason
+
+    # key_past shape: (bs, num_heads, size_per_head, seq_length)
+    # value_past shape: (bs, num_heads, seq_length, size_per_head)
+    # key shape: (bs, num_heads, 1, size_per_head)
+    # value shape: (bs, num_heads, 1, size_per_head)
+
+    if past_shape[0] != cur_shape[0] or past_shape[1] != cur_shape[1]:
+        reason = "past_shape[0] != cur_shape[0] or past_shape[1] != cur_shape[1] "
+        return False, reason
+
+    if past_shape[3] != cur_shape[3]:
+        reason = "past_shape[3] != cur_shape[3]"
+        return False, reason
+
+    return True, ""
+
+
+def ceil_div(dividend, divisor):
+    return (dividend + divisor - 1) // divisor
+
+
+def get_loop_info(total_num, each_loop_num):
+    loop_times = ceil_div(total_num, each_loop_num)
+    last_loop_num = total_num - each_loop_num * (loop_times - 1)
+    return loop_times, last_loop_num
+
+
+def elements_align(index_elements, data_size, align_size):
+    """Get element num align to align_size"""
+    total_size = index_elements * data_size
+    aligned_total_size = (total_size + align_size - 1) // align_size * align_size
+    return aligned_total_size // data_size
+
+
+class TilingHelper:
+    """Tiling parameter"""
+    def __init__(self, past, cur, index, out, kernel_name="kv_cache_mgr"):
+        self.kernel_name = kernel_name
+
+        # sys info
+        self.core_num = tbe_platform.get_soc_spec(tbe_platform.CORE_NUM)
+        self.ub_size = tbe_platform.get_soc_spec(tbe_platform.UB_SIZE)
+
+        self.past_shape = past.get("shape")
+        self.cur_shape = cur.get("shape")
+        self.index_shape = index.get("shape")
+
+        self.gm_type = past.get("dtype").lower()
+        self.ub_type = self.gm_type
+        self.index_ub_type = "int32"
+        self.int32_size = 4
+
+        self.gm_dtype_size = 2
+        if self.gm_type in ["int8", "uint8"]:
+            self.gm_dtype_size = 1
+        elif self.gm_type in ["float16", "int16", "uint16"]:
+            self.gm_dtype_size = 2
+        elif self.gm_type in ["float32", "int32", "uint32"]:
+            self.gm_dtype_size = 4
+
+        # tiling policy
+        self.seq_length = self.past_shape[2]
+        self.size_per_head = self.past_shape[3]
+        self.update_seq_length = self.cur_shape[2]
+
+        self.num_head = self.past_shape[1]
+
+        self.past_elements = functools.reduce(lambda a, b: a * b, self.past_shape)
+        self.cur_elements = functools.reduce(lambda a, b: a * b, self.cur_shape)
+
+        # The `burst` unit is 32B
+        index_elements = functools.reduce(lambda a, b: a * b, self.index_shape)
+        self.index_elements = elements_align(index_elements, self.int32_size, 32)
+
+        # split cur
+        self.cur_bs = self.cur_shape[0] * self.cur_shape[1]
+        self.each_core_bs_num = ceil_div(self.cur_bs, self.core_num)
+        self.core_num, self.last_core_bs_num = get_loop_info(self.cur_bs, self.each_core_bs_num)
+        self.cur_ub_elements = self.each_core_bs_num * self.update_seq_length * self.size_per_head
+        self.last_cure_ub_elements = self.last_core_bs_num * self.update_seq_length * self.size_per_head
+
+
+class KVCacheImpl(TilingHelper):
+    """KVCacheImpl"""
+    def __init__(self, past, cur, index, out, kernel_name):
+        super().__init__(past, cur, index, out, kernel_name)
+        # key_past or value_past shape: (bs, num_heads, seq_length, size_per_head)
+        # batch_valid_length
+        # cur update shape: (bs, num_heads, 1, size_per_head)
+
+        self.tik_inst = tik.Tik(disable_debug=True)
+        self.past_gm = self.tik_inst.Tensor(self.gm_type, (self.past_elements,), name="past_gm", scope=tik.scope_gm)
+        self.cur_gm = self.tik_inst.Tensor(self.gm_type, (self.cur_elements,), name="cur_gm", scope=tik.scope_gm)
+        self.index_gm = self.tik_inst.Tensor(self.index_ub_type, (self.index_elements,), name="index_gm",
+                                             scope=tik.scope_gm)
+        # we use is_atomic_add=True to set the out_gm zeros. But if inplace update out_gm, no need to set this flag.
+        self.out_gm = self.tik_inst.Tensor(self.gm_type, (self.past_elements,), name="out_gm", scope=tik.scope_gm)
+
+    def valid_cur_ub_load(self, core_idx):
+        """KVCacheImpl.valid_cur_ub_load"""
+        cur_ub = self.tik_inst.Tensor(self.ub_type, (self.cur_ub_elements,), name="valid_cur_ub",
+                                      scope=tik.scope_ubuf)
+        cur_gm_offset = core_idx * self.cur_ub_elements
+        with self.tik_inst.if_scope(core_idx != self.core_num -1):
+            self.tik_inst.data_move(cur_ub, self.cur_gm[cur_gm_offset:], 0, 1,
+                                    self.cur_ub_elements * self.gm_dtype_size // 32, 0, 0)
+        with self.tik_inst.else_scope():
+            self.tik_inst.data_move(cur_ub, self.cur_gm[cur_gm_offset:], 0, 1,
+                                    self.last_cure_ub_elements * self.gm_dtype_size // 32, 0, 0)
+        return cur_ub
+
+    def valid_index_ub_load(self):
+        """KVCacheImpl.valid_index_ub_load"""
+        index_ub = self.tik_inst.Tensor(self.index_ub_type, (self.index_elements,), name="valid_index_ub",
+                                        scope=tik.scope_ubuf)
+        self.tik_inst.data_move(index_ub, self.index_gm, 0, 1, self.index_elements * self.int32_size // 32, 0, 0)
+        return index_ub
+
+    def valid_pos_update(self, core_idx, cur_ub, index_ub, each_core_bs_num):
+        """KVCacheImpl.valid_pos_update"""
+        src_bs_stride = self.update_seq_length * self.size_per_head
+        dst_bs_stride = self.seq_length * self.size_per_head
+        burst_len = self.update_seq_length * self.size_per_head * self.gm_dtype_size // 32
+
+        valid_idx = self.tik_inst.Scalar(dtype="int32")
+        with self.tik_inst.for_range(0, each_core_bs_num) as each_core_bs_idx:
+            bs_idx = core_idx * self.each_core_bs_num + each_core_bs_idx
+            # because we fused bs * num_head, we need get the real bs_idx
+            valid_idx.set_as(index_ub[bs_idx // self.num_head])
+            with self.tik_inst.if_scope(valid_idx >= 0):
+                dst_offset = bs_idx * dst_bs_stride + valid_idx * self.size_per_head
+                src_offset = each_core_bs_idx * src_bs_stride
+                if burst_len < 65536:
+                    self.tik_inst.data_move(self.out_gm[dst_offset], cur_ub[src_offset],
+                                            0, 1, burst_len, 0, 0)
+                else:
+                    nburst = 1
+                    each_burst_len = burst_len
+                    while each_burst_len > 65535:
+                        nburst += 1
+                        each_burst_len = burst_len // nburst
+                    self.tik_inst.data_move(self.out_gm[dst_offset], cur_ub[src_offset], 0,
+                                            nburst, each_burst_len, 0, 0)
+
+    # 'pylint: disable=too-many-arguments
+    def compute_each_core(self, core_idx, core_bs_num):
+        """KVCacheImpl.compute_each_core"""
+        index_ub = self.valid_index_ub_load()
+        cur_ub = self.valid_cur_ub_load(core_idx)
+        self.valid_pos_update(core_idx, cur_ub, index_ub, core_bs_num)
+
+    def compute(self):
+        """KVCacheImpl.compute"""
+        if self.each_core_bs_num == self.last_core_bs_num:
+            with self.tik_inst.for_range(0, self.core_num, block_num=self.core_num) as core_index:
+                self.compute_each_core(core_idx=core_index, core_bs_num=self.each_core_bs_num)
+        else:
+            with self.tik_inst.for_range(0, self.core_num, block_num=self.core_num) as core_index:
+                with self.tik_inst.if_scope(core_index < self.core_num - 1):
+                    self.compute_each_core(core_idx=core_index, core_bs_num=self.each_core_bs_num)
+                with self.tik_inst.else_scope():
+                    self.compute_each_core(core_idx=core_index, core_bs_num=self.last_core_bs_num)
+
+        self.tik_inst.BuildCCE(kernel_name=self.kernel_name,
+                               inputs=[self.past_gm, self.cur_gm, self.index_gm],
+                               outputs=[self.out_gm],
+                               )
+        return self.tik_inst
+
+
+# 'pylint: disable = unused-argument
+# 'pylint: disable=too-many-arguments,too-many-locals
+@para_check.check_op_params(para_check.REQUIRED_INPUT, para_check.REQUIRED_INPUT,
+                            para_check.REQUIRED_INPUT, para_check.REQUIRED_OUTPUT,
+                            para_check.KERNEL_NAME)
+def kv_cache_mgr(past, cur, index, out, kernel_name="kv_cache_mgr"):
+    """
+    :param past: key_past or value_past. shape: (bs, num_head, seq_length, size_pre_head)
+    :param cur: key_current or value_current. shape: (bs, num_head, update_seq_length, size_pre_head)
+    :param index: which index to update. shape * len(dtype) need be multiples of 32. Option Input.
+    :param out: output shape: (bs, num_head, seq_length, size_pre_head)
+    :param kernel_name: the name of the op
+    :return:
+    """
+    obj = KVCacheImpl(past, cur, index, out, kernel_name)
+    return obj.compute()

mindspore/lib/plugin/ascend/custom_aicore_ops/op_impl/vector_core/tbe/custom_aicore_ops_impl/matmul_tik.py

@@ -0,0 +1,212 @@
+"""
+Copyright 2020 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.
+
+matmul_tik
+"""
+
+from tbe import tik
+from tbe.common.platform import get_soc_spec
+
+DTYPE_SIZE = {
+    'bool': 1,
+    'uint8': 1,
+    'int8': 1,
+    'uint16': 2,
+    'int16': 2,
+    'int24': 3,
+    'uint32': 4,
+    'int32': 4,
+    'float16': 2,
+    'float32': 4,
+    'int48': 6,
+    'int64': 8,
+    'uint64': 8,
+    'float64': 8
+}
+
+
+def MK_TO_K1MK0(tik_instance, mk_input_tensor, k1mk0_tensor, dtype, k1, m, k0):
+    """data move mk to k1mk0"""
+    src_ub = tik_instance.Tensor(dtype, (k1, m, k0), name='src_ub', scope=tik.scope_ubuf)
+
+    # data_move(m, k) ---> (k1, m, k0)
+    with tik_instance.for_range(0, k1) as i:
+        tik_instance.data_move(src_ub[i * m * k0:], mk_input_tensor[i * k0:], 0, m, k0 * DTYPE_SIZE[dtype] // 32,
+                               (k1 - 1) * k0 * DTYPE_SIZE[dtype] // 32, 0)
+
+    tik_instance.data_move(k1mk0_tensor, src_ub, 0, 1, k1 * m * k0 * DTYPE_SIZE[dtype] // 32, 0, 0)
+
+
+def KN_TO_K1NK0(tik_instance, kn_input_tensor, k1nk0_tensor, dtype, k1, n, k0):
+    """data move kn to k1nk0"""
+
+    with tik_instance.for_range(0, k1) as index:
+        k1nk0_ub = tik_instance.Tensor(dtype, (n, k0), tik.scope_ubuf, "k1nk0_ub")
+        src_ub = tik_instance.Tensor(dtype, (k0, n), tik.scope_ubuf, "src_ub")
+        burst_len = k0 * n * DTYPE_SIZE[dtype] // 32
+        tik_instance.data_move(src_ub, kn_input_tensor[index * k0 * n], 0, 1, burst_len, 0, 0)
+        dst_list = [k1nk0_ub[16 * i] for i in range(16)]
+        src_list = [src_ub[n * i] for i in range(16)]
+        rep_times = n // k0
+        dst_rep_stride = k0
+        src_rep_stride = 1
+        tik_instance.vec_trans_scatter(False, False, dst_list, src_list, rep_times, dst_rep_stride, src_rep_stride)
+        tik_instance.data_move(k1nk0_tensor[index * k0 * n], k1nk0_ub, 0, 1, burst_len, 0, 0)
+
+
+def N1MN0_TO_MN(tik_instance, mn_output_tensor, n1mn0_tensor, dtype, n1, m, n0):
+    """data move mn to n1mn0"""
+    src_ub = tik_instance.Tensor(dtype, (m, n1 * n0), name='src_ub', scope=tik.scope_ubuf)
+
+    # data_move(n1, m, n0) ---> (m, n)
+    with tik_instance.for_range(0, n1) as i:
+        tik_instance.data_move(src_ub[i * n0:], n1mn0_tensor[i * m * n0:], 0, m,
+                               n0 * DTYPE_SIZE[dtype] // 32, 0, (n1 - 1) * n0 * DTYPE_SIZE[dtype] // 32)
+
+    tik_instance.data_move(mn_output_tensor, src_ub, 0, 1, m * n1 * n0 * DTYPE_SIZE[dtype] // 32, 0, 0)
+
+
+def matmul_tik_compute(params, kernel_name):
+    """
+    matmul tik compute
+    @param params: matmul data
+    @param kernel_name: kernel name
+    @return: tik instance
+    """
+    tik_instance = tik.Tik()
+    if not isinstance(params, dict):
+        params = params.__dict__
+    m_size, k_size, n_size = params['M'], params['K'], params['N']
+    data_type = params["data_type"]
+    m_tiling_size = int(params["m_tiling_size"])
+    n_tiling_size = int(params["n_tiling_size"])
+    k_tiling_size = int(params['k_tiling_size'])
+
+    m_cycle_times = params["m_cycle_times"]
+    n_cycle_times = params["n_cycle_times"]
+    k_cycle_times = params["k_cycle_times"]
+
+    # Determine the output type
+    if data_type == "float16":
+        if get_soc_spec("SOC_VERSION") in ["SD3403", "OPTG", "Hi3796CV300CS", "TsnsC"]:
+            C_loc_out_type = "float16"
+        else:
+            C_loc_out_type = "float32"
+        K0 = 16
+    else:
+        C_loc_out_type = "int32"
+        K0 = 32
+    block_size = 16
+
+    n_thread_num = params['n_thread_num']
+    m_thread_num = params['m_thread_num']
+    k_thread_num = params['k_thread_num']
+
+    mk_gm_input = tik_instance.Tensor(data_type, (m_size, k_size), name="mk_input_gm", scope=tik.scope_gm)
+    kn_gm_input = tik_instance.Tensor(data_type, (k_size, n_size), name="kn_input_gm", scope=tik.scope_gm)
+
+    k1mk0_workspace = tik_instance.Tensor(data_type, (k_size // K0, m_size, K0), name="k1mk0_workspace",
+                                          scope=tik.scope_gm, is_workspace=True)
+
+    k1nk0_workspace = tik_instance.Tensor(data_type, (k_size // K0, n_size, K0), name="k1nk0_workspace",
+                                          scope=tik.scope_gm, is_workspace=True)
+
+    mn_gm_output = tik_instance.Tensor(C_loc_out_type, (m_size, n_size), tik.scope_gm, name="mn_output_gm")
+    nmk0_workspace = tik_instance.Tensor(C_loc_out_type, (n_size // block_size, m_size, block_size),
+                                         name="nmk0_workspace", scope=tik.scope_gm, is_workspace=True)
+
+    MK_TO_K1MK0(tik_instance, mk_gm_input, k1mk0_workspace, data_type, k_size // K0, m_size, K0)
+    KN_TO_K1NK0(tik_instance, kn_gm_input, k1nk0_workspace, data_type, k_size // K0, n_size, K0)
+
+    # Tiling is realized through the for_range() loop.
+    with tik_instance.for_range(0, 2, block_num=1) as core_id:
+        with tik_instance.for_range(0, n_cycle_times // 2, thread_num=n_thread_num) as n_idx:
+            with tik_instance.for_range(0, m_cycle_times, thread_num=m_thread_num) as m_idx:
+                dst_l0c = tik_instance.Tensor(C_loc_out_type, [n_tiling_size // 16, m_tiling_size, 16], name='dst_l0c',
+                                              scope=tik.scope_cbuf_out)
+                with tik_instance.for_range(0, k_cycle_times,
+                                            thread_num=k_thread_num) as k_idx:
+                    # Calculation result data transfer.
+                    inputa_l1 = tik_instance.Tensor(params['data_type'], [k_tiling_size // K0, m_tiling_size, K0],
+                                                    name="A_tiling_l1", scope=tik.scope_cbuf)
+                    tik_instance.data_move(inputa_l1,
+                                           k1mk0_workspace[k_idx * k_tiling_size // K0, m_idx * m_tiling_size, :],
+                                           0, k_tiling_size // K0, m_tiling_size, m_size - m_tiling_size, 0)
+                    inputb_l1 = tik_instance.Tensor(params["data_type"], [k_tiling_size // K0, n_tiling_size, K0],
+                                                    name="B_tiling_l1", scope=tik.scope_cbuf)
+                    if n_size - n_tiling_size > 65535:
+                        with tik_instance.for_range(0, k_tiling_size // K0) \
+                                as dma_k_idx:
+                            tik_instance.data_move(inputb_l1[dma_k_idx, :, :],
+                                                   k1nk0_workspace[k_idx * k_tiling_size // K0 + dma_k_idx,
+                                                                   (core_id * n_cycle_times // 2 + n_idx)
+                                                                   * n_tiling_size, :],
+                                                   0, 1, n_tiling_size, 0, 0)
+                    else:
+                        tik_instance.data_move(inputb_l1, k1nk0_workspace[k_idx * k_tiling_size // K0,
+                                                                          (core_id * n_cycle_times // 2 + n_idx)
+                                                                          * n_tiling_size, :],
+                                               0, k_tiling_size // K0, n_tiling_size, n_size - n_tiling_size, 0)
+                    # Call matmul API to matrix multiplication calculation.
+                    with tik_instance.if_scope(k_idx == 0):
+                        tik_instance.matmul(dst_l0c, inputa_l1, inputb_l1, m_tiling_size, k_tiling_size, n_tiling_size,
+                                            init_l1out=True)
+                    with tik_instance.else_scope():
+                        tik_instance.matmul(dst_l0c, inputa_l1, inputb_l1, m_tiling_size, k_tiling_size, n_tiling_size,
+                                            init_l1out=False)
+                tik_instance.fixpipe(nmk0_workspace[n_tiling_size // 16 * (core_id * n_cycle_times // 2 + n_idx),
+                                                    m_idx * m_tiling_size, :],
+                                     dst_l0c, n_tiling_size // 16,
+                                     m_tiling_size * 16 * DTYPE_SIZE[C_loc_out_type] // 32,
+                                     (m_size - m_tiling_size) * 16 * DTYPE_SIZE[C_loc_out_type] // 32, 0)
+
+    N1MN0_TO_MN(tik_instance, mn_gm_output, nmk0_workspace, C_loc_out_type, n_size // K0, m_size, K0)
+
+    tik_instance.BuildCCE(kernel_name=kernel_name, inputs=[mk_gm_input, kn_gm_input], outputs=[mn_gm_output])
+    return tik_instance
+
+
+def matmul_tik(input_x1, input_x2, output_y=None, kernel_name="simple_matmul"):
+    """
+    matmul_tik main func
+    Parameters
+    ----------
+    input_x1: input data 1
+    input_x2: input data 2
+    output_y: output dta
+    """
+    shape_a = input_x1.get("ori_shape")
+    shape_b = input_x2.get("ori_shape")
+    m = shape_a[0]
+    k = shape_a[1]
+    n = shape_b[1]
+    data_type = input_x1.get("dtype").lower()
+    params = {
+        'M': m,
+        'K': k,
+        'N': n,
+        'data_type': data_type,
+        'm_tiling_size': 16,
+        'm_cycle_times': 1,
+        'm_thread_num': 1,
+        'n_tiling_size': 64,
+        'n_cycle_times': 16,
+        'n_thread_num': 1,
+        'k_tiling_size': 32,
+        'k_cycle_times': 2,
+        'k_thread_num': 2,
+        'output_y': output_y
+    }
+    return matmul_tik_compute(params, kernel_name)