mindspore 2.4.1__cp39-none-any.whl → 2.5.0__cp39-none-any.whl

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/op_param.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2023-2024 Huawei Technologies Co., Ltd
+ * 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.
@@ -13,127 +13,172 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-#ifndef MS_KERNELS_INTERNAL_OP_PARAM_H_
-#define MS_KERNELS_INTERNAL_OP_PARAM_H_
-#include "asdops/op_desc.h"
-#include "asdops/params/matmul.h"
-#include "asdops/params/mix.h"
-#include "asdops/params/slice.h"
-#include "asdops/params/gather.h"
-#include "asdops/params/elewise.h"
-#include "asdops/params/activation.h"
-#include "asdops/params/concat.h"
-#include "asdops/params/transpose.h"
-#include "asdops/params/norm.h"
-#include "asdops/params/softmax.h"
-#include "asdops/params/split.h"
-#include "asdops/params/expand.h"
-#include "asdops/params/fill.h"
-#include "asdops/params/reduce.h"
-#include "asdops/params/sort.h"
-#include "asdops/params/transdata.h"
-#include <memory>
+
+#ifndef MS_KERNELS_INTERNAL_KERNEL_OP_PARAM_H_
+#define MS_KERNELS_INTERNAL_KERNEL_OP_PARAM_H_
+
+#include <stdint.h>
 #include <vector>
-#include "ms_int_types.h"
+
 namespace mindspore {
 namespace internal {
-struct DtypesParam {
-  int op_id_ = 0;
-  std::vector<int64_t> in_dtypes_;
-  std::vector<int64_t> out_dtypes_;
+// matmul fused op
+constexpr auto kInternalMatMulOpName = "MatMul";
+constexpr auto kInternalMultiWeightMatmulOpName = "MultiWeightMatmul";
+constexpr auto kInternalMatMulAddRmsNormOpName = "MatMulAddRmsNorm";
+// attention fused op
+constexpr auto kInternalFlashAttentionScoreOpName = "FlashAttentionScore";
+constexpr auto kInternalPagedAttentionOpName = "PagedAttention";
+constexpr auto kInternalReshapeAndCacheOpName = "ReshapeAndCache";
+constexpr auto kInternalReshapeAndCacheNzOpName = "ReshapeAndCacheNz";
+constexpr auto kInternalApplyRotaryPosEmbOpName = "ApplyRotaryPosEmb";
+constexpr auto kInternalApplyRotaryPosEmbNzOpName = "ApplyRotaryPosEmbNz";
+// norm fused op
+constexpr auto kInternalAddLayerNormOpName = "AddLayerNorm";
+constexpr auto kInternalRmsNormOpName = "RmsNorm";
+constexpr auto kInternalAddRmsNormOpName = "AddRmsNorm";
+constexpr auto kInternalRmsNormQuantOpName = "RmsNormQuant";
+constexpr auto kInternalAddRmsNormQuantOpName = "AddRmsNormQuantV2";
+constexpr auto kInternalAddRmsNormDynamicQuantOpName = "AddRmsNormDynamicQuant";
+// activation
+constexpr auto kInternalReluOpName = "Relu";
+constexpr auto kInternalGeLUOpName = "GeLU";
+constexpr auto kInternalFastGeLUOpName = "FastGeLU";
+constexpr auto kInternalSwishOpName = "Swish";
+constexpr auto kInternalSwiGLUOpName = "SwiGLU";
+// elewise unary
+constexpr auto kInternalCastOpName = "Cast";
+constexpr auto kInternalExpOpName = "Exp";
+constexpr auto kInternalLnOpName = "Ln";
+constexpr auto kInternalRsqrtOpName = "Rsqrt";
+constexpr auto kInternalSqrtOpName = "Sqrt";
+constexpr auto kInternalAbsOpName = "Abs";
+constexpr auto kInternalReciprocalOpName = "Reciprocal";
+// elewise binary
+constexpr auto kInternalAddOpName = "Add";
+constexpr auto kInternalSubOpName = "Sub";
+constexpr auto kInternalMulOpName = "Mul";
+constexpr auto kInternalDivOpName = "Div";
+constexpr auto kInternalRealDivOpName = "RealDiv";
+constexpr auto kInternalMaxOpName = "Max";
+constexpr auto kInternalMinOpName = "Min";
+constexpr auto kInternalNotOpName = "Not";
+constexpr auto kInternalOrOpName = "Or";
+constexpr auto kInternalAndOpName = "And";
+constexpr auto kInternalEqualOpName = "Equal";
+constexpr auto kInternalNotEqualOpName = "NotEqual";
+constexpr auto kInternalLessOpName = "Less";
+constexpr auto kInternalLessEqualOpName = "LessEqual";
+constexpr auto kInternalGreaterOpName = "Greater";
+constexpr auto kInternalGreaterEqualOpName = "GreaterEqual";
+constexpr auto kInternalLogicalNotOpName = "LogicalNot";
+// others
+constexpr auto kInternalGatherOpName = "Gather";
+constexpr auto kInternalTransposeOpName = "Transpose";
+constexpr auto kInternalTransDataOpName = "TransData";
+constexpr auto kInternalQuantPerChannelOpName = "QuantPerChannel";
+constexpr auto kInternalSoftmaxOpName = "Softmax";
+constexpr auto kInternalReduceSumOpName = "ReduceSum";
+constexpr auto kInternalQuantLinearSparseOpName = "QuantLinearSparse";
+
+struct AxesParam {
+  std::vector<int64_t> axes;
+};
+
+using TransposeParam = AxesParam;
+using SoftmaxParam = AxesParam;
+using ReduceSumParam = AxesParam;
+
+struct GatherParam {
+  int64_t batch_dims;
+  std::vector<int64_t> axes;
+};
+
+struct SwiGLUParam {
+  int64_t axis;
+};
+
+struct MatmulParam {
+  bool transpose_a{false};
+  bool transpose_b{false};
+  bool enable_dequant{false};
+  bool with_relu{false};
+  bool with_gelu{false};
+  bool with_fastgelu{false};
+  bool with_bias{false};
+  bool with_bias_fastgelu{false};
+  bool enable_shuffle{false};
+  uint32_t tilingN = 0;  // 压缩算法透传参数, 单压缩块 n 方向的基块数
+  uint32_t tilingK = 0;  // 压缩算法透传参数, 单压缩块 k 方向的基块数
 };
-struct OpParam : public AsdOps::OpDesc {
-  int dtype_ = 0;
-  std::vector<int64_t> in_dtypes_;
-  std::vector<int64_t> out_dtypes_;
-  std::string op_fullname_;
+
+struct MatmulAddRmsNormParam {
+  bool transpose_a{false};
+  bool transpose_b{false};
+  float eps{1e-6};
 };
-enum OpId : int {
-  MatMul,
-  ReshapeAndCache,
-  ReshapeAndCacheNz,
-  Slice,
-  Gather,
-  ApplyRotaryPosEmb,
-  ApplyRotaryPosEmbNz,
-  Add,
-  Sub,
-  Exp,
-  Relu,
-  FlashAttentionScore,
-  PagedAttention,
-  Cast,
-  Gelu,
-  Transpose,
-  Equal,
-  NotEqual,
-  LogicalNot,
-  Less,
-  LessEqual,
-  Greater,
-  GreaterEqual,
-  Mul,
-  RealDiv,
-  QuantPerChannel,
-  LayerNorm,
-  AddLayerNorm,
-  RmsNorm,
-  AddRmsNorm,
-  RmsNormQuant,
-  AddRmsNormQuant,
-  MatmulAddRmsNorm,
-  Softmax,
-  Split,
-  Swish,
-  SwiGLU,
-  Concat,
-  MatmulQkv,
-  MaskedFill,
-  BroadcastTo,
-  ReduceSum,
-  TopK,
-  Tile,
-  GroupedMatmul,
-  OpId_END,
-  FastGeLU,
-  TransData,
-  QuantLinearSparse,
+
+struct MultiWeightMatmulParam {
+  uint32_t n0_len{0};
+  uint32_t n1_len{0};
+  uint32_t n2_len{0};
+  bool transpose_a;
+  bool transpose_b;
+  int32_t silu_position{-1};
+  bool with_bias{false};
 };
-using MatMulParam = AsdOps::OpParam::MatMul;
-using MixParam = AsdOps::OpParam::Mix;
-using GatherParam = AsdOps::OpParam::Gather;
-using ElewiseParam = AsdOps::OpParam::Elewise;
-using SliceParam = AsdOps::OpParam::Slice;
-using ActivationParam = AsdOps::OpParam::Activation;
-using TransposeParam = AsdOps::OpParam::Transpose;
-using NormParam = AsdOps::OpParam::Norm;
-using SoftmaxParam = AsdOps::OpParam::Softmax;
-using SplitParam = AsdOps::OpParam::Split;
-using ConcatParam = AsdOps::OpParam::Concat;
-using MaskedFillParam = AsdOps::OpParam::Fill;
-using BroadcastToParam = AsdOps::OpParam::Expand;
-using ReduceParam = AsdOps::OpParam::Reduce;
-using SortParam = AsdOps::OpParam::Sort;
-using ExpandParam = AsdOps::OpParam::Expand;
-using TransDataParam = AsdOps::OpParam::Transdata;
-
-struct AddLayerNormParam {
+
+struct NormParam {
   float eps;
-  bool operator==(const AddLayerNormParam &other) const { return this->eps == other.eps; }
+  bool operator==(const NormParam &other) const { return this->eps == other.eps; }
 };
 
 struct ApplyRotaryPosEmbParam {
-  // cosFormat=0  shape是[maxSeqLen, headDim]，    cos/sin不交替
-  // cosFormat=1  shape是[maxSeqLen, headDim]，    cos/sin交替
-  // cosFormat=2  shape是[batch*seqLen, headDim]， cos/sin不交替
-  // cosFormat=3  shape是[batch*seqLen, headDim]， cos/sin交替
-  int32_t cosFormat{0};
+  // cos_format=0  shape是[maxSeqLen, headDim]，    cos/sin不交替
+  // cos_format=1  shape是[maxSeqLen, headDim]，    cos/sin交替
+  // cos_format=2  shape是[batch*seqLen, headDim]， cos/sin不交替
+  // cos_format=3  shape是[batch*seqLen, headDim]， cos/sin交替
+  int32_t cos_format{0};
+  int32_t rotary_coeff{-1};
+  std::vector<int32_t> batch_valid_length;
 };
 
-struct AddRmsNormParam {
-  float eps;
-  bool operator==(const AddRmsNormParam &other) const { return this->eps == other.eps; }
+struct TransDataParam {
+  enum TransdataType { UNDEFINED = 0, FRACTAL_NZ_TO_ND, ND_TO_FRACTAL_NZ };
+  TransdataType transdataType = UNDEFINED;
+  enum SpecialType { NORMAL = 0, ATTENTION_INPUT_QKV, ATTENTION_INPUT_MASK };
+  int64_t specialTransdata = NORMAL;
+};
+
+struct FlashAttentionScoreParam {
+  int32_t head_num = 0;
+  int32_t inner_precise = 0;
+  int32_t pre_tokens = 2147483647;
+  int32_t next_tokens = 0;
+  int32_t sparse_mode = 0;
+  int32_t mask_dtype = 0;
+  int32_t input_layout = 0;
+  std::vector<int64_t> mask_dims;
+  std::vector<int32_t> kv_seq_len;
+  std::vector<int32_t> q_seq_len;
+  float tor = 0;
+
+  enum InputLayoutMode : int64_t { BSH = 0, BNSD = 1, SBH = 2, BSND = 3, TND = 4, TH = 5, NSD = 6, SH = 7 };
+};
+
+struct PagedAttentionParam {
+  int32_t inner_precise = 0;
+  int32_t head_num = 0;
+  int32_t kv_head_num = 0;
+  std::vector<int32_t> kv_seq_len;
+  std::vector<int32_t> q_seq_len;
+  float tor = 0;
+
+  enum MaskType : uint32_t { kMaskTypeNone = 0, kMaskTypeAlibi = 1, kMaskTypeLookAhead = 2 };
+  MaskType mask_type = kMaskTypeNone;
+  int32_t kv_cache_quant_mode = 0;
 };
 }  // namespace internal
 }  // namespace mindspore
-#endif
+
+#endif  // MS_KERNELS_INTERNAL_KERNEL_OP_PARAM_H_

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/{acme/include/tiling_info.h → tiling_info.h}

@@ -14,21 +14,22 @@
  * limitations under the License.
  */
 
-#ifndef MS_KERNELS_INTERNAL_KERNEL_ACME_TILING_INFO_H_
-#define MS_KERNELS_INTERNAL_KERNEL_ACME_TILING_INFO_H_
+#ifndef MS_KERNELS_INTERNAL_KERNEL_TILING_INFO_H_
+#define MS_KERNELS_INTERNAL_KERNEL_TILING_INFO_H_
 
 #include <sstream>
 
-#include "acme/include/base_type.h"
+#include "include/base_type.h"
 
 namespace mindspore {
-namespace acme {
+namespace internal {
 class HostRunInfo {
  public:
   HostRunInfo() = default;
   virtual ~HostRunInfo() = default;
   void SetWorkSpaceSize(const std::vector<size_t> &workSpaceSize) { this->ws_size_ = workSpaceSize; }
   std::vector<size_t> GetWorkSpaceSize() const { return this->ws_size_; }
+  virtual uint32_t GetBlockDim() const = 0;
 
  private:
   std::vector<size_t> ws_size_;
@@ -37,10 +38,12 @@ using HostRunInfoPtr = std::shared_ptr<HostRunInfo>;
 
 class HostRunInfoComm : public HostRunInfo {
  public:
-  HostRunInfoComm() = default;
-  explicit HostRunInfoComm(size_t block_dim) : block_dims_(block_dim) {}
+  explicit HostRunInfoComm(uint32_t block_dim);
   ~HostRunInfoComm() = default;
-  uint32_t block_dims_{0};
+
+  uint32_t GetBlockDim() const override;
+
+  uint32_t block_dim_{0};
   uint64_t any_value0_{0};
   uint64_t any_value1_{0};
   uint64_t any_value2_{0};
@@ -61,7 +64,7 @@ class TilingInfo {
   HostRunInfoPtr host_run_info_{nullptr};
 };
 using TilingInfoPtr = std::shared_ptr<TilingInfo>;
-}  // namespace acme
+}  // namespace internal
 }  // namespace mindspore
 
-#endif  // MS_KERNELS_INTERNAL_KERNEL_ACME_TILING_INFO_H_
+#endif  // MS_KERNELS_INTERNAL_KERNEL_TILING_INFO_H_

mindspore/lib/plugin/ascend/ms_kernels_internal/internal_kernel/include/tiling_utils.h

@@ -0,0 +1,178 @@
+/**
+ * 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 MS_KERNELS_INTERNAL_KERNEL_TILING_UTILS_H_
+#define MS_KERNELS_INTERNAL_KERNEL_TILING_UTILS_H_
+
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <sstream>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "include/internal_op.h"
+
+namespace mindspore {
+namespace internal {
+void Encrypt(char **src, size_t str_len, char *dest, bool offset_src = false);
+void Decrypt(char **src, size_t str_len, char *dest, bool offset_src = true);
+
+struct ScopeData {
+  uint32_t begin;
+  uint32_t end;
+  uint32_t offset;
+  uint32_t num;
+};
+
+struct ScopeInfo {
+  std::vector<uint16_t> types;
+  std::vector<ScopeData> scopes;
+};
+
+struct TilingData {
+  uint16_t version;
+  uint32_t base_offset;
+  uint16_t key_len;
+  uint16_t value_len;
+  uint32_t item_num;
+  uint16_t type_len;
+  std::vector<ScopeInfo> scope_infos;
+};
+
+template <typename T>
+struct ArrayHasher {
+  std::size_t operator()(const std::vector<T> &arr) const {
+    std::size_t hash = 0;
+    for (T num : arr) {
+      hash ^= std::hash<T>{}(num) + 0x9e3779b9 + (hash << 6) + (hash >> 2);
+    }
+    return hash;
+  }
+};
+
+using Repo = std::unordered_map<std::vector<int>, std::vector<int>, ArrayHasher<int>>;
+using RepoMap = std::unordered_map<std::vector<uint16_t>, Repo, ArrayHasher<uint16_t>>;
+
+class TilingLoadUtil {
+ public:
+  explicit TilingLoadUtil(const std::string &db_file, bool is_full_path = false)
+      : db_file_(db_file), is_full_(is_full_path) {}
+  ~TilingLoadUtil();
+  InternalStatus LoadTilingData(TilingData *tiling_data_out);
+  InternalStatus LoadScopeTilings(uint32_t base, uint16_t key_len, uint16_t value_len, const ScopeData &scope, Repo *repo);
+
+ private:
+  std::string GetDataFileFullPath(const std::string &db_name) const;
+  InternalStatus InitData();
+  InternalStatus ReadTilingData(TilingData *tiling_data);
+  InternalStatus ReadScopeData(TilingData *tiling_data);
+
+  std::string db_file_;
+  bool is_full_;
+  int f_{-1};
+  size_t memory_size_{0};
+  char *data_{nullptr};
+};
+
+using KeyScope = std::pair<int, int>;
+
+class TilingRepo {
+ public:
+  TilingRepo(const TilingRepo &flags) = delete;
+  TilingRepo(TilingRepo &&flags) = delete;
+  TilingRepo &operator=(const TilingRepo &flags) = delete;
+  TilingRepo &operator=(TilingRepo &&flags) = delete;
+  ~TilingRepo() = default;
+
+  static TilingRepo &Instance();
+  void Register(const std::string &name, const std::string &db_name);
+  InternalStatus GetTiling(const std::string &name, const std::vector<uint16_t> &types, const std::vector<int> &key,
+                       std::vector<int> *out,
+                       const std::function<KeyScope(const std::vector<int> &)> &gen_scope_func = nullptr);
+  using LoadRecord = std::unordered_map<std::vector<uint16_t>, std::vector<KeyScope>, ArrayHasher<uint16_t>>;
+
+ private:
+  TilingRepo() = default;
+  bool IsRegistered(const std::string &name) const;
+  std::vector<int> Get(const std::string &name, const std::vector<uint16_t> &types, const std::vector<int> &key) const;
+  InternalStatus LoadScopeData(const std::string &name, const std::vector<uint16_t> &types,
+                           const KeyScope &scope = std::make_pair(0, 0));
+
+  std::unordered_map<std::string, size_t> name_idx_map_;
+  std::unordered_map<std::string, size_t> db_name_idx_map_;
+  std::vector<std::string> db_names_;
+  std::vector<TilingData> tiling_datas_;
+  std::vector<RepoMap> repo_maps_;
+  std::vector<LoadRecord> loaded_;
+};
+
+class TilingDBLoad {
+ public:
+  TilingDBLoad(const std::string &op_name, const std::string &db_name) {
+    TilingRepo::Instance().Register(op_name, db_name);
+  }
+};
+
+struct RuningInfo {
+  internal::ShapeInfoList input_shapes;
+  internal::InputsImmutableInfoList input_infos;
+  internal::ShapeInfoList output_shapes;
+  internal::InputsImmutableInfoList output_infos;
+};
+
+class Tunable {
+ public:
+  Tunable() = default;
+  virtual ~Tunable() = default;
+  virtual InternalOpPtr CreateOpByKey(const std::vector<int64_t> &key) = 0;
+  virtual RuningInfo GetRuningInfo(const std::vector<int64_t> &key) const = 0;
+};
+using TunablePtr = std::shared_ptr<Tunable>;
+
+using TunableCreator = std::function<TunablePtr()>;
+class TunableBuilder {
+ public:
+  ~TunableBuilder() = default;
+  TunableBuilder(const TunableBuilder &) = delete;
+  TunableBuilder &operator=(const TunableBuilder &) = delete;
+  static TunableBuilder &Instance();
+
+  void Register(const std::string &op_name, TunableCreator &&creator);
+  TunablePtr Create(const std::string &op_name) const;
+
+ private:
+  TunableBuilder() = default;
+  std::unordered_map<std::string, TunableCreator> tunable_creators_;
+};
+
+class TuneRegister {
+ public:
+  TuneRegister(const std::string &op_name, TunableCreator creator) noexcept {
+    TunableBuilder::Instance().Register(op_name, std::move(creator));
+  }
+  ~TuneRegister() = default;
+};
+
+#define REG_OP_TUNABLE(op_name, TargetClass)                                                                  \
+  static_assert(std::is_base_of<Tunable, TargetClass>::value, #TargetClass " must be derived from Tunable!"); \
+  static const TuneRegister g_##op_name##_tunable_reg(#op_name,                                               \
+                                                      []() -> TunablePtr { return std::make_shared<TargetClass>(); })
+}  // namespace internal
+}  // namespace mindspore
+
+#endif  // MS_KERNELS_INTERNAL_KERNEL_TILING_UTILS_H_