mindspore 2.3.0rc1__cp39-none-any.whl → 2.3.0rc2__cp39-none-any.whl

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/apply_rotary_pos_emb/kernel/apply_rotary_pos_emb_fp16.h

@@ -0,0 +1,116 @@
+/**
+ * Copyright (c) Huawei Technologies Co., Ltd. 2024. 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.
+ */
+#ifndef ROTARY_POS_EMB_FP16
+#define ROTARY_POS_EMB_FP16
+#include "apply_rotary_pos_emb_base.h"
+template <typename QK_DTYPE, typename COS_DTYPE, bool IF_COS_BROADCAST>
+class RopeFp16 : public RopeBase<QK_DTYPE, COS_DTYPE, IF_COS_BROADCAST> {
+ public:
+  __aicore__ inline RopeFp16(RopeTilingData *tilingData) : RopeBase<QK_DTYPE, COS_DTYPE, IF_COS_BROADCAST>(tilingData) {
+    this->repeatSize_ = 128;  // 128 = 256B / sizeof(half)
+    this->maxProcessNum_ = this->tilingData_->maxUbSize / sizeof(uint16_t);
+    this->repeatTimesQ_ = (this->tilingData_->hiddenSizeQ + this->repeatSize_ - 1) / this->repeatSize_;
+    this->repeatTimesK_ = (this->tilingData_->hiddenSizeK + this->repeatSize_ - 1) / this->repeatSize_;
+    headDimAlign_ = ((this->tilingData_->headDim + ELE_NUM_FP16 - 1) / ELE_NUM_FP16) * ELE_NUM_FP16;
+    this->alignHalfHeadDim_ = (this->rotateStride_ * NUM_TWO) % ELE_NUM_FP16;
+    this->hiddenSizeAlign_ = ((this->hiddenSize_ + this->repeatSize_ - 1) / this->repeatSize_) * this->repeatSize_;
+
+    this->cosPad_ = 0;
+    this->sinPad_ = this->cosPad_ + this->hiddenSizeAlign_;
+    this->negOne_ = this->sinPad_ + this->hiddenSizeAlign_;
+    this->oriPos_ = this->negOne_ + this->hiddenSizeAlign_;
+    this->padBefore_ = this->oriPos_ + this->hiddenSizeAlign_;
+    this->removeBefore_ = this->padBefore_ + this->hiddenSizeAlign_;
+    sinResPos_ = this->removeBefore_ + this->hiddenSizeAlign_;
+    this->repeatTimes_ = this->hiddenSizeAlign_ / this->repeatSize_;
+
+    this->syncOffset_ =
+      (this->tilingData_->headDim % ELE_NUM_FP16 == 0) ? this->hiddenSizeAlign_ : this->headNum_ * headDimAlign_;
+    this->offsetExtraGm_ = NUM_TWO * block_idx * this->syncOffset_;
+    this->pipe_.InitBuffer(outQueueCO2_, 1, ((this->maxProcessNum_ - this->batchSize_ * NUM_TWO) * sizeof(QK_DTYPE)));
+    AscendC::LocalTensor<QK_DTYPE> cache_perloop_ub_ = outQueueCO2_.AllocTensor<QK_DTYPE>();
+    commonUbuf_ = (__ubuf__ QK_DTYPE *)cache_perloop_ub_.GetPhyAddr();
+  }
+
+  __aicore__ inline void Process(__gm__ uint8_t *extraGm) {
+    if (this->tilingData_->cosFormat == 1) {
+      pipe_barrier((PIPE_ALL));
+      this->ExpandCosSin(commonUbuf_, this->cosGm_, (__gm__ COS_DTYPE *)extraGm);
+      this->cosGm_ = (__gm__ COS_DTYPE *)extraGm;
+      pipe_barrier((PIPE_ALL));
+      this->ExpandCosSin(commonUbuf_, this->sinGm_,
+                         (__gm__ COS_DTYPE *)extraGm + this->tilingData_->ntokens * this->tilingData_->headDim);
+      this->sinGm_ = (__gm__ COS_DTYPE *)extraGm + this->tilingData_->ntokens * this->tilingData_->headDim;
+      extraGm =
+        extraGm + this->tilingData_->ntokens * this->tilingData_->headDim * 4;  // sizeof(uint8_t) * 2 = sizeof(half)
+      pipe_barrier((PIPE_ALL));
+    }
+
+    this->ExpandNeg(commonUbuf_, sinResPos_, this->headNum_, this->repeatTimes_);  // 根据是否对齐选择1 -1 還是 -1 0
+    for (uint32_t zz = 0; zz < this->dynamicRound_; ++zz) {
+      this->CosSinBroadcast(extraGm, zz, commonUbuf_, this->tilingData_->hiddenSizeQ);  // cos sin 和 QK 无关
+
+      this->QkComm(this->qGm_ + block_idx * this->nlCoreRun_ * this->tilingData_->hiddenSizeQ +
+                     zz * this->tilingData_->hiddenSizeQ,
+                   extraGm, this->tilingData_->hiddenSizeQ, commonUbuf_, this->tilingData_->headNumQ);
+
+      if (this->alignRotary_ == 0) {
+        pipe_barrier((PIPE_V));
+        this->CalcRopeAlign(commonUbuf_, this->repeatTimesQ_, this->oriPos_, this->removeBefore_, this->padBefore_);
+      } else {
+        set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
+        wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
+        this->CalcRope(commonUbuf_, this->repeatTimesQ_, this->oriPos_, this->removeBefore_, this->padBefore_,
+                       sinResPos_, this->padBefore_);
+      }
+      pipe_barrier((PIPE_ALL));  // 需要
+      copy_ubuf_to_gm(this->outQGm_ + block_idx * this->nlCoreRun_ * this->tilingData_->hiddenSizeQ +
+                        zz * this->tilingData_->hiddenSizeQ,
+                      commonUbuf_ + this->padBefore_, 0, 1, this->tilingData_->hiddenSizeQ / ELE_NUM_FP16, 0, 0);
+
+      set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1);
+      wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1);
+
+      this->QkComm(this->kGm_ + block_idx * this->nlCoreRun_ * this->tilingData_->hiddenSizeK +
+                     zz * this->tilingData_->hiddenSizeK,
+                   extraGm, this->tilingData_->hiddenSizeK, commonUbuf_, this->tilingData_->headNumK);
+
+      if (this->alignRotary_ == 0) {
+        pipe_barrier((PIPE_V));
+        this->CalcRopeAlign(commonUbuf_, this->repeatTimesK_, this->oriPos_, this->removeBefore_, this->padBefore_);
+      } else {
+        set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
+        wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1);
+        this->CalcRope(commonUbuf_, this->repeatTimesK_, this->oriPos_, this->removeBefore_, this->padBefore_,
+                       sinResPos_, this->padBefore_);
+      }
+      pipe_barrier((PIPE_ALL));  // 需要
+      copy_ubuf_to_gm(this->outKGm_ + block_idx * this->nlCoreRun_ * this->tilingData_->hiddenSizeK +
+                        zz * this->tilingData_->hiddenSizeK,
+                      commonUbuf_ + this->padBefore_, 0, 1, this->tilingData_->hiddenSizeK / ELE_NUM_FP16, 0, 0);
+      set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1);
+      wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1);
+    }
+  }
+
+ private:
+  AscendC::TQue<AscendC::QuePosition::VECIN, 1> outQueueCO2_;
+  __ubuf__ QK_DTYPE *commonUbuf_{nullptr};
+  uint32_t headDimAlign_;  // 对齐的headDim
+  uint32_t sinResPos_{0};  // fp32的buf中0 0 0 1 1 1的位置
+};
+
+#endif

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/apply_rotary_pos_emb/kernel/apply_rotary_pos_emb_tiling.h

@@ -14,38 +14,30 @@
  * limitations under the License.
  */
 
-#ifndef MS_KERNELS_INTERNAL_KERNEL_ASCENDC_APPLY_ROTARY_POS_EMB_TILING_DATA_H_
-#define MS_KERNELS_INTERNAL_KERNEL_ASCENDC_APPLY_ROTARY_POS_EMB_TILING_DATA_H_
+#ifndef MS_KERNELS_INTERNAL_KERNEL_ASCENDC_ROPE_TILING_DATA_H_
+#define MS_KERNELS_INTERNAL_KERNEL_ASCENDC_ROPE_TILING_DATA_H_
 
 #include <stdint.h>
 
-struct ApplyRotaryPosEmbTilingData {
+struct RopeTilingData {
   uint32_t hiddenSizeQ{16};
   uint32_t hiddenSizeK{16};
-  uint32_t headDim{1};
+  uint32_t headDim{1};  // qk头长度的最大值
   uint32_t headNumQ{1};
   uint32_t headNumK{1};
-  uint32_t rotaryCoeff{4};
-  uint32_t ntokens{1};
-  uint32_t ropeFormat{0};
-  uint32_t realCore{0};
-  uint32_t cosFormat{0};
-  uint32_t batch{32};
-  uint32_t highPrecision{0};
-  uint32_t maxUbSize{0};
+  uint32_t rotaryCoeff{4};  // 旋转系数
+  uint32_t ntokens{1};      // 总token数
+  uint32_t realCore{0};     // 实际用到核数
+  uint32_t cosFormat{0};    // 是否复用cos sin
+  uint32_t batch{32};       // 几个batch
+  uint32_t maxUbSize{0};    // 最大UB内存
+  uint32_t tilingId{0};
 
-  int32_t ndim;
-  int32_t qkDtype;   // 0=fp16 1=bf16 2=fp32
-  int32_t posDtype;  // 0=i64 1=u64 2=i32 3=u32
-
-  // int32_t batch;
-  int32_t numHeadQ;
-  int32_t numHeadK;
-  // int32_t hiddenDim;
-  int32_t seqLen;
-  int32_t maxSeqLen;
-
-  int32_t posSize;  // seqLen==1 ? batch : seqLen
+  uint32_t seqLen;
+  uint32_t broadCastCos{0};
+  uint32_t posDtype;
+  uint32_t posSize;
+  uint32_t maxSeqLen;
 };
 
 #endif

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/apply_rotary_pos_emb/kernel/apply_rotary_pos_emb_value.h

@@ -0,0 +1,27 @@
+/**
+ * Copyright (c) Huawei Technologies Co., Ltd. 2024. 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.
+ */
+#ifndef COMMON_VAL_H
+#define COMMON_VAL_H
+const constexpr uint32_t NUM_TWO = 2; // 2
+const constexpr uint32_t BLK_SIZE = 32; // 一个block字节数
+const constexpr uint32_t ELE_NUM_FP16 = 16; // 一个block fp16元素个数
+const constexpr uint32_t ELE_NUM_FP32 = 8; // 一个block字节数 fp32元素个数
+const constexpr uint32_t MAX_LEN_FP16 = 8192; // 非fp16情况下最大长度（hiddensize）
+const constexpr uint8_t DEFAULT_REPEAT_STRIDE = 8; // 默认stride, 8 * 32 = 256
+const constexpr int64_t REG_910B = 48; // 饱和模式寄存器位置
+const constexpr int64_t REG_310P = 53; // 饱和模式寄存器位置
+const constexpr int64_t SLICE_SIZE = 4096; // 切片大小
+#endif

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/asdop/asd_op_impl.h

@@ -41,8 +41,6 @@ class AsdOpsImpl : public InternelKernelImpl {
   std::vector<uint64_t> GetWorkSpaceSize() override;
   int InferShape(const std::vector<DIMS> &input_shapes, std::vector<DIMS> &output_shapes) override;
 
-  void SetCacheInfo(const CacheInfo &cache_info);
-
  private:
   AsdOps::Tactic *InitAndGetTactic();
 
@@ -52,8 +50,6 @@ class AsdOpsImpl : public InternelKernelImpl {
   AsdOps::LaunchParam launch_param_;
   AsdOps::OpDesc op_desc_;
   bool validated_ = false;
-
-  RunInfo run_info_;
 };
 
 }  // namespace internal

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/{attention/FlashAttentionScore_impl.h → flash_attention_score/flash_attention_score_impl.h}

@@ -24,7 +24,7 @@
 #include "asdops/tensor.h"
 
 #include "internal_kernel.h"
-
+#include "param/attention_param.h"
 #include "acl_rt.h"
 
 #include <unordered_map>
@@ -49,6 +49,7 @@ class FlashAttentionScoreImpl : public InternelKernelImpl {
 
  private:
   uint64_t B, N, Q_S, KV_S, D, G, CORE_NUM;
+  int inner_precise, pre_tokens, next_tokens, sparse_mode;
   bool BFLOAT16, BSH, ALIBI, AMASK;
   const std::vector<Tensor *> *inputs_;
   const std::vector<Tensor *> *outputs_;

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/{attention/bs_attention_tiling.h → flash_attention_score/flash_attention_score_tiling.h}

@@ -18,16 +18,12 @@ typedef struct {
 #pragma pack()
 
 #define MAX_CORE_NUM 25
-#define ATTENTION_DEBUG false // 开启时会对S/P写入调试数据
-#define BUFFER_NUM 2 // 核间流水数，暂不支持修改
-
-#if PA
-#define INC true
-#define OP_NAME PagedAttention
-#else
-#define INC false
+#define ATTENTION_DEBUG false   // 开启时会对S/P写入调试数据
+#define ROWMAX true
 #define OP_NAME FlashAttentionScore
-#endif
+#define BUFFER_NUM 2            // 核间流水数，暂不支持修改
+constexpr uint64_t WORKSPACE_MAX_SEQLEN = 16384;    // max seqlen
+constexpr uint64_t WORKSPACE_SIZE = 128 * WORKSPACE_MAX_SEQLEN;
 
 #if BFLOAT16
 #define TYPE_NAME _bf16
@@ -41,14 +37,16 @@ typedef struct {
 #define LAYOUT_NAME _BNSD
 #endif
 
-#define CONCAT_(A, B, C, D) A##B##C##D
-#define CONCAT(A, B, C, D) CONCAT_(A, B, C, D)
-#define FUNC_NAME_AIC CONCAT(OP_NAME, TYPE_NAME, LAYOUT_NAME, _mix_aic)
-#define FUNC_NAME_AIV CONCAT(OP_NAME, TYPE_NAME, LAYOUT_NAME, _mix_aiv)
-
-#if INC
-#define CORE_PER_KV_HEAD 4 // 增量推理时开启，每个kv_head切分成多少个任务计算
+#if LOWER_TRIANGLE
+#define TRI_NAME _tri
+#else
+#define TRI_NAME _full
 #endif
+ 
+#define CONCAT_(A, B, C, D, E) A##B##C##D##E
+#define CONCAT(A, B, C, D, E) CONCAT_(A, B, C, D, E)
+#define FUNC_NAME_AIC CONCAT(OP_NAME, TYPE_NAME, LAYOUT_NAME, TRI_NAME, _mix_aic)
+#define FUNC_NAME_AIV CONCAT(OP_NAME, TYPE_NAME, LAYOUT_NAME, TRI_NAME, _mix_aiv)
 
 // **************mask patten模式**************//
 // 第一种：下三角，开启LOWER_TRIANGLE时会直接采用下三角，不依赖mask
@@ -63,9 +61,7 @@ typedef struct {
 // 第四种：全矩阵，LOWER_TRIANGLE、BLOCK_SPARSE和AMASK如果全部关闭，则此attention采用全矩阵运算，不抑制S中的元素
 // *******************************************//
 
-constexpr uint64_t WORKSPACE_MAX_SEQLEN = 10240;    // max seqlen: 10240
-constexpr uint64_t WORKSPACE_MAX_SEQLEN_BLOCK = WORKSPACE_MAX_SEQLEN / 16;    // max seqlen: 10240
-constexpr uint64_t WORKSPACE_SIZE = 128 * WORKSPACE_MAX_SEQLEN; 
+constexpr uint64_t WORKSPACE_MAX_SEQLEN_BLOCK = WORKSPACE_MAX_SEQLEN / 16;
 constexpr uint64_t BUFFER_SIZE = MAX_CORE_NUM * WORKSPACE_SIZE * sizeof(uint16_t);
 
 #endif

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/gelu/tiling/gelu_tiling.h

@@ -26,15 +26,13 @@ struct GeLUTilingData {
   uint32_t tailBlockTileNum{0};
 };
 static std::ostream &operator<<(std::ostream &os, const GeLUTilingData &dt) {
-  os << "blockDims:" << dt.blockDims << std::endl;
-  os << "totalLength:" << dt.totalLength << std::endl;
-  os << "blockLength:" << dt.blockLength << std::endl;
-  os << "tileLength:" << dt.tileLength << std::endl;
-  os << "tileNum:" << dt.tileNum << std::endl;
-  os << "tailBlockTileNum:" << dt.tailBlockTileNum << std::endl;
-  os << "tilingKey:" << dt.tilingKey << std::endl;
-  // os << "axisDim:" << dt.axisDim << std::endl;
-  // os << "splitNum:" << dt.splitNum << std::endl;
+  os << "blockDims:" << dt.blockDims;
+  os << ", totalLength:" << dt.totalLength;
+  os << ", blockLength:" << dt.blockLength;
+  os << ", tileLength:" << dt.tileLength;
+  os << ", tileNum:" << dt.tileNum;
+  os << ", tailBlockTileNum:" << dt.tailBlockTileNum;
+  os << ", tilingKey:" << dt.tilingKey;
   return os;
 }
 #endif  // MS_KERNELS_INTERNAL_ASCENDC_GELU_TILING_H

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/lccl/lccl_wrapper.h

@@ -0,0 +1,58 @@
+/**
+ * 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.
+ */
+#ifndef LCCL_WRAPPER_H_
+#define LCCL_WRAPPER_H_
+
+#include <memory>
+#include "lccl.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+using namespace Lcal;
+using LcclComm = std::shared_ptr<Lccl>;
+enum class LcclResult {
+    LCAL_SUCCESS = 0,
+    LCAL_ERROR_NOT_INITIALIZED = -1,
+    LCAL_ERROR_ASDRT = -2,
+    LCAL_ERROR_PARA_CHECK_FAIL = -3,
+    LCAL_ERROR_INTERNAL = -4,
+    LCAL_ERROR_TIMEOUT = -5,
+    LCCL_ERROR_INIT_HCCL_FAILED = -6
+};
+
+extern LcclResult LcclCommInitRank(uint32_t nRanks, uint32_t rank, LcclComm *comm);
+
+extern LcclResult LcclAllReduce(void *sendBuff, void *recvBuff, int64_t count, HcclDataType dataType,
+                                HcclReduceOp op, aclrtStream stream);
+
+extern LcclResult LcclReduceScatter(void *sendBuff, void *recvBuff, int64_t count, HcclDataType dataType,
+                                    HcclReduceOp op, aclrtStream stream);
+
+extern LcclResult LcclAllGather(void *sendBuff, void *recvBuff, int64_t count, HcclDataType dataType, aclrtStream stream);
+
+extern LcclResult LcclAll2All(void *sendBuff, void *recvBuff, int64_t count, HcclDataType dataType, aclrtStream stream);
+
+extern LcclResult LcclBroadcast(void *buff, int64_t count, HcclDataType dataType, int32_t root, aclrtStream stream);
+
+extern LcclResult LcclCommDestroy(LcclComm comm);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  // LCCL_WRAPPER_H_

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/matmul/matmul_impl.h

@@ -24,11 +24,13 @@
 #include "asdops/tensor.h"
 
 #include "utils.h"
-#include "pp_matmul_info.h"
-#include "tiling_data.h"
-#include "pp_matmul_common_tiling.h"
+#include "backend_param.h"
+#include "param/matmul_ext_param.h"
+#include "matmul_common/pp_matmul_info.h"
+#include "matmul_common/tiling_utils.h"
+#include "matmul_common/tiling_data.h"
+#include "matmul_common/pp_matmul_common_tiling.h"
 #include "tune_repo/utils.h"
-
 #include "internal_kernel.h"
 
 #include "acl_rt.h"
@@ -37,6 +39,8 @@
 namespace mindspore {
 namespace internal {
 
+using namespace tiling;
+
 enum class MatMulAlgo { PP = 0, LLM_CUSTOM = 1 };
 
 class MatMulImpl : public InternelKernelImpl {
@@ -48,12 +52,15 @@ class MatMulImpl : public InternelKernelImpl {
   int Launch() override;
   size_t GetTilingBufSize() override;
   int Tiling(HostRawBuf &tilingBuf) override;
-  int TilingPp(HostRawBuf &tilingBuf, uint64_t tilingId);
-  int TilingLLMCustom(HostRawBuf &tilingBuf, uint64_t tilingId);
+  void TilingBasicFromPp(uint32_t &blockDim, PpTilingData &tilingdata);
+  int TilingPp(HostRawBuf &tilingBuf, uint64_t tilingId, const uint32_t &blockDim, const PpTilingData &tilingdata);
+  int TilingLLMCustom(HostRawBuf &tilingBuf, uint64_t tilingId, const uint32_t &blockDim,
+                      const PpTilingData &tilingdata);
   std::vector<uint64_t> GetWorkSpaceSize() override;
   int InferShape(const std::vector<DIMS> &input_shapes, std::vector<DIMS> &output_shapes) override;
-  bool CanUseLLMCustom();
-  int GetCoresFromValidLLMCustomTile(const std::vector<int> &cand);
+  bool UseCustomMatMul();
+  void GetTunedKey();
+  void SetTunedValueCustom(const std::vector<int> &tuned_config);
 
  private:
   uint32_t m_, k_, n_;
@@ -61,7 +68,11 @@ class MatMulImpl : public InternelKernelImpl {
   MatMulAlgo algo_ = MatMulAlgo::PP;
   DeviceRawBuf tiling_addr_;
   std::string soc_{"Ascend910B2"};
+  HardwareInfo hwInfo_;
+  CustomMatmulTilingData t_;
+  std::vector<int> tune_key_;
   REPO tuningTable_;
+  REPO tuningTableCustom_;
   TensorDType input_dtype_;
   TensorDType output_dtype_;
   int block_dim_ = 0;

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/{matmul → matmul_common}/pp_matmul_common_tiling.h

@@ -17,6 +17,8 @@
 #ifndef MATMUL_COMMMON_TILING_H
 #define MATMUL_COMMMON_TILING_H
 
+#include <cmath>
+#include <iostream>
 #include "pp_matmul_info.h"
 
 namespace mindspore {
@@ -100,9 +102,12 @@ inline __attribute__((always_inline)) float CostFunc(const HardwareType &hwInfo,
 template <bool PRI_FLAG, typename OpShareType, typename TilingType, typename HardwareType, typename MatMulInfoType>
 void TilingFunc(OpShareType &opShape, TilingType &tilingParam, const HardwareType &hwInfo, const MatMulInfoType &mmInfo,
                 bool compressFlag = false, const uint32_t tilingN = 1) {
+  using namespace std;
   float costMin = 1;
-  uint32_t priAxes = Round<16>(PRI_FLAG ? opShape.m : opShape.n);
-  uint32_t axes = Round<16>(PRI_FLAG ? opShape.n : opShape.m);
+  const uint32_t CONST_16 = 16;
+  uint32_t roundBase = pow(2, ceil(log(CeilDiv(PRI_FLAG ? opShape.n : opShape.m, CONST_16)))) * CONST_16;
+  uint32_t priAxes = RoundUp(PRI_FLAG ? opShape.m : opShape.n, CONST_16);
+  uint32_t axes = RoundUp(PRI_FLAG ? opShape.n : opShape.m, roundBase);
   float axes0Max = static_cast<float>(AXES_ALIGN_SIZE) / mmInfo.inDtype;
 
   uint32_t n0TilingInit =
@@ -129,10 +134,15 @@ void TilingFunc(OpShareType &opShape, TilingType &tilingParam, const HardwareTyp
       }
       opShape.m0 = PRI_FLAG ? priAxes0 : axes0;
       opShape.n0 = PRI_FLAG ? axes0 : priAxes0;
+      if ((mmInfo.qkv_n0 + mmInfo.qkv_n1 + mmInfo.qkv_n2 != 0) &&
+          (mmInfo.qkv_n0 < opShape.n0 || mmInfo.qkv_n1 < opShape.n0 ||
+           (mmInfo.qkv_n2 < opShape.n0 && mmInfo.qkv_n2 > 1))) {
+        continue;
+      }
       float cost = CostFunc<HardwareType, OpShareType>(hwInfo, opShape);
       if (cost < costMin) {
         costMin = cost;
-        tilingParam.SetBaseOp(hwInfo.coreNum, opShape.m0, opShape.n0);
+        tilingParam.SetBaseOp(hwInfo.coreNum, opShape.m0, opShape.n0, mmInfo.qkv_n0, mmInfo.qkv_n1, mmInfo.qkv_n2);
       }
     }
   }
@@ -140,7 +150,7 @@ void TilingFunc(OpShareType &opShape, TilingType &tilingParam, const HardwareTyp
 
 template <typename PpTilingDataType>
 uint32_t Swizzl(PpTilingDataType &tilingData) {
-  uint32_t swizzleDirect = 0;
+  uint32_t swizzlDirect = 0;
   uint32_t swizzlCount = 1;
   float m0 = tilingData.opShape.m0;
   float n0 = tilingData.opShape.n0;
@@ -154,14 +164,14 @@ uint32_t Swizzl(PpTilingDataType &tilingData) {
     float cost;
     // B0 + A < A0 + B
     if (i * n0 + m < m0 * c + n) {
-      swizzleDirect = 1;  // Nz
+      swizzlDirect = 1;  // Nz
       cost = n0 * i + m0 * c;
       if (cost <= mincost) {
         mincost = cost;
         swizzlCount = i;
       }
     } else {
-      swizzleDirect = 0;  // Zn
+      swizzlDirect = 0;  // Zn
       cost = m0 * i + n0 * c;
       if (cost < mincost) {
         mincost = cost;
@@ -169,9 +179,9 @@ uint32_t Swizzl(PpTilingDataType &tilingData) {
       }
     }
   }
-  tilingData.swizzleDirect = swizzleDirect;
+  tilingData.swizzlDirect = swizzlDirect;
   tilingData.swizzlCount = swizzlCount;
-  return swizzleDirect;
+  return swizzlDirect;
 }
 
 }  // namespace tiling

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/{matmul → matmul_common}/pp_matmul_info.h

@@ -26,7 +26,10 @@ namespace internal {
 namespace tiling {
 struct MatMulInfo {
   uint32_t batchSize{0};
-  uint32_t m{0};       // 实际输入的 m
+  uint32_t m{0};  // 实际输入的 m
+  uint32_t qkv_n0{0};
+  uint32_t qkv_n1{0};
+  uint32_t qkv_n2{0};
   uint32_t n{0};       // 实际输入的 n
   uint32_t k{0};       // 实际输入的 k
   bool transA{0};      // false: 0, true: 1
@@ -57,12 +60,12 @@ struct PpTilingData {
   uint32_t swizzlCount{1};
   uint32_t tilingKey{0};
   uint32_t blockDim{1};
-  uint32_t swizzleDirect{0};
+  uint32_t swizzlDirect{0};
   uint32_t splitk{0};
 
   void SetBaseShape(uint32_t batchSize, uint32_t m, uint32_t k, uint32_t n);
-  void SetBaseOp(uint32_t coreNum, uint32_t mBase, uint32_t nBase);
-  void SetTilingKey(const MatMulInfo &mmInfo, uint32_t swizzleDirect, uint32_t enSplitK);
+  void SetBaseOp(uint32_t coreNum, uint32_t mBase, uint32_t nBase, uint32_t qkv_n0, uint32_t qkv_n1, uint32_t qkv_n2);
+  void SetTilingKey(const MatMulInfo &mmInfo, uint32_t swizzlDirect, uint32_t enSplitK);
   uint32_t End(const MatMulInfo &mmInfo);
 };
 }  // namespace tiling

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/{matmul → matmul_common}/tiling_data.h

@@ -18,7 +18,6 @@
 #define MATMUL_TILING_DATA_H
 
 #include <stdint.h>
-#include <algorithm>
 
 namespace mindspore {
 namespace internal {
@@ -38,8 +37,16 @@ struct PpMatmulTilingData {
   uint32_t swizzlCount{0};
   uint32_t tilingKey{0};
   uint32_t blockDim{1};
-  uint32_t swizzleDirect{0};
+  uint32_t swizzlDirect{0};
   uint32_t splitk{0};
+  uint32_t enShuffleK{0};
+  uint32_t unused0{0};
+  uint32_t unused1{0};
+  uint32_t unused2{0};
+  uint32_t unused3{0};
+  uint32_t unused4{0};
+  uint32_t unused5{0};
+  uint32_t unused6{0};
   uint32_t tilingId{0};
 };
 
@@ -60,18 +67,49 @@ struct CustomMatmulTilingData {
   uint32_t BaseKNum{0};
   uint32_t BaseNNum{0};
   uint32_t MmadM{0};
+  uint32_t MmadK{0};
+  uint32_t MmadN{0};
+  uint32_t fractal_k_num{0};
+  uint32_t FractalKInBlockNum{0};
+  uint32_t PartKInMmad{0};
+  uint32_t TransA{0};
+  uint32_t TransB{0};
+  uint32_t shuffleFlag{0};
   uint32_t tilingId{0};
+};
+
+constexpr size_t maxTilingBufSize = sizeof(CustomMatmulTilingData);
+
+struct MatmulStridedSliceFusionTilingData {
+  uint32_t tilingId{0};
+  uint32_t BlockDimM{0};
+  uint32_t BlockDimN{0};
+  uint32_t BlockTotal{0};
+  uint32_t M{0};
+  uint32_t K{0};
+  uint32_t N{0};
+  uint32_t N0{0};
+  uint32_t N1{0};
+  uint32_t N2{0};
+  uint32_t BaseM{0};
+  uint32_t BaseK{0};
+  uint32_t BaseN{0};
+  uint32_t BlockLenM{0};
+  uint32_t BlockLenK{0};
+  uint32_t BlockLenN{0};
+  uint32_t BaseMNum{0};
+  uint32_t BaseKNum{0};
+  uint32_t BaseNNum{0};
+  uint32_t MmadM{0};
   uint32_t MmadK{0};
   uint32_t MmadN{0};
-  uint32_t FractalKNum{0};
   uint32_t FractalKInBlockNum{0};
-  uint32_t PartKInL0A{2};
+  uint32_t PartKInMmad{2};
   uint32_t TransA{0};
   uint32_t TransB{1};
+  uint32_t shuffleFlag{0};
 };
 
-constexpr size_t maxTilingBufSize = std::max(sizeof(PpMatmulTilingData), sizeof(CustomMatmulTilingData));
-
 }  // namespace tiling
 }  // namespace internal
 }  // namespace mindspore