Line data Source code
1 : /**
2 : * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3 : * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4 : * CANN Open Software License Agreement Version 2.0 (the "License").
5 : * Please refer to the License for details. You may not use this file except in compliance with the License.
6 : * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7 : * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8 : * See LICENSE in the root of the software repository for the full text of the License.
9 : */
10 :
11 : /*!
12 : * \file compressor_tiling.cpp
13 : * \file compressor_tiling.cpp
14 : * \brief
15 : */
16 :
17 : #include <functional>
18 : #include <algorithm>
19 : #include <unordered_map>
20 : #include <graph/utils/type_utils.h>
21 : #include "log/log.h"
22 : #include "register/op_def_registry.h"
23 : #include "compressor_tiling_arch22.h"
24 :
25 : using namespace ge;
26 : using namespace AscendC;
27 : namespace optiling {
28 : namespace {
29 :
30 : void CompressorTiling::ConvertRequiredParams(gert::TilingContext &context, CompressorContext &compressorContext)
31 : {
32 : compressorContext.x.desc = context.GetRequiredInputDesc(TOKEN_X_INPUT_INDEX);
33 : compressorContext.x.shape = context.GetRequiredInputShape(TOKEN_X_INPUT_INDEX);
34 : compressorContext.wkv.desc = context.GetRequiredInputDesc(WEIGHT_KV_INPUT_INDEX);
35 : compressorContext.wkv.shape = context.GetRequiredInputShape(WEIGHT_KV_INPUT_INDEX);
36 : compressorContext.wgate.desc = context.GetRequiredInputDesc(WEIGHT_WGATE_INPUT_INDEX);
37 : compressorContext.wgate.shape = context.GetRequiredInputShape(WEIGHT_WGATE_INPUT_INDEX);
38 : compressorContext.stateCache.desc = context.GetRequiredInputDesc(STATE_CACHE_INPUT_INDEX);
39 : compressorContext.stateCache.shape = context.GetRequiredInputShape(STATE_CACHE_INPUT_INDEX);
40 : compressorContext.ape.desc = context.GetRequiredInputDesc(APE_INPUT_INDEX);
41 : compressorContext.ape.shape = context.GetRequiredInputShape(APE_INPUT_INDEX);
42 :
43 : compressorContext.cmpKv.desc = context.GetOutputDesc(CMP_KV_OUTPUT_INDEX);
44 : compressorContext.cmpKv.shape = context.GetOutputShape(CMP_KV_OUTPUT_INDEX);
45 :
46 : compressorContext.dtype = compressorContext.x.desc->GetDataType();
47 : auto xDimNum = compressorContext.x.shape->GetStorageShape().GetDimNum();
48 : if (xDimNum == COMPRESSOR_DIM_NUM_3) {
49 : compressorContext.layout = LayoutType::LAYOUT_BSH;
50 : } else if (xDimNum == COMPRESSOR_DIM_NUM_2) {
51 : compressorContext.layout = LayoutType::LAYOUT_TH;
52 : }
53 : }
54 :
55 : void CompressorTiling::ConvertOptionalParams(gert::TilingContext &context, CompressorContext &compressorContext)
56 : {
57 : compressorContext.stateBlockTable.desc = context.GetOptionalInputDesc(STATE_BLOCK_TABLE_INPUT_INDEX);
58 : compressorContext.stateBlockTable.shape = context.GetOptionalInputShape(STATE_BLOCK_TABLE_INPUT_INDEX);
59 : compressorContext.cuSeqlens.desc = context.GetOptionalInputDesc(CU_SEQ_LEN_INPUT_INDEX);
60 : compressorContext.cuSeqlens.shape = context.GetOptionalInputShape(CU_SEQ_LEN_INPUT_INDEX);
61 : compressorContext.seqUsed.desc = context.GetOptionalInputDesc(SEQ_USED_INPUT_INDEX);
62 : compressorContext.seqUsed.shape = context.GetOptionalInputShape(SEQ_USED_INPUT_INDEX);
63 : compressorContext.startPos.desc = context.GetOptionalInputDesc(START_POS_INPUT_INDEX);
64 : compressorContext.startPos.shape = context.GetOptionalInputShape(START_POS_INPUT_INDEX);
65 : }
66 :
67 : ge::graphStatus CompressorTiling::ConvertContext(gert::TilingContext &context, CompressorContext &compressorContext)
68 : {
69 : if (context.GetNodeName() == nullptr) {
70 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON("Compressor", "opName", "got from TilingContext is nullptr");
71 : return ge::GRAPH_FAILED;
72 : }
73 :
74 : OP_LOGI("Getting Context");
75 :
76 : compressorContext.opName = context.GetNodeName();
77 : compressorContext.opType = context.GetNodeType();
78 : compressorContext.platformInfo = context.GetPlatformInfo();
79 : ConvertRequiredParams(context, compressorContext);
80 : ConvertOptionalParams(context, compressorContext);
81 :
82 : auto attrs = context.GetAttrs();
83 : OP_CHECK_IF(attrs == nullptr,
84 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.GetNodeName(), "attrs", "got from ge is nullptr"),
85 : return ge::GRAPH_FAILED);
86 : compressorContext.coff = attrs->GetAttrPointer<int>(COFF_ATTR_INDEX);
87 : compressorContext.cmpRatio = attrs->GetAttrPointer<int>(CMP_RATIO_ATTR_INDEX);
88 : compressorContext.cacheMode = attrs->GetAttrPointer<int>(CACHE_MODE_ATTR_INDEX);
89 : compressorContext.stateCacheStrideDim0 = attrs->GetAttrPointer<int>(STATE_CACHE_STRIDE_DIM0_ATTR_INDEX);
90 22 : compressorContext.gradEnabled = attrs->GetAttrPointer<bool>(GRAD_ENABLED_ATTR_INDEX);
91 :
92 : OP_CHECK_IF(
93 : context.GetWorkspaceSizes(1) == nullptr,
94 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context.GetNodeName(), "workSpaceSize", "got from ge is nullptr"),
95 : return ge::GRAPH_FAILED);
96 : compressorContext.workSpaces = context.GetWorkspaceSizes(1);
97 :
98 : // Validate state_cache axis 1/2 contiguity via GetInputStride.
99 : // Only axis 0 non-contiguous is supported; axis 1 and 2 must be contiguous.
100 : {
101 : auto *stateCacheStride = context.GetInputStride(STATE_CACHE_INPUT_INDEX);
102 : const auto &stateCacheShape = compressorContext.stateCache.shape->GetStorageShape();
103 : if (stateCacheStride != nullptr && stateCacheStride->GetDimNum() == stateCacheShape.GetDimNum()) {
104 : uint64_t expectedStride = 1;
105 : for (int64_t i = static_cast<int64_t>(stateCacheShape.GetDimNum()) - 1; i >= 1; --i) {
106 : uint64_t actualStride = static_cast<uint64_t>(stateCacheStride->GetStride(static_cast<size_t>(i)));
107 : if (actualStride != expectedStride) {
108 : OP_LOGE(context.GetNodeName(),
109 : "state_cache must be contiguous on axis 1 and 2, "
110 : "only axis 0 non-contiguous is supported. "
111 : "axis %ld: actual stride=%lu, expected=%lu.",
112 : i, actualStride, expectedStride);
113 : return ge::GRAPH_FAILED;
114 : }
115 : expectedStride *= static_cast<uint64_t>(stateCacheShape.GetDim(static_cast<size_t>(i)));
116 : }
117 : }
118 : }
119 :
120 : return ge::GRAPH_SUCCESS;
121 : }
122 :
123 : ge::graphStatus CompressorTiling::GetNpuInfo()
124 : {
125 : OP_CHECK_IF(context_->platformInfo == nullptr,
126 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "platformInfo", "is nullptr"),
127 : return ge::GRAPH_FAILED);
128 :
129 : auto ascendcPlatform = platform_ascendc::PlatformAscendC(context_->platformInfo);
130 : socVersion_ = ascendcPlatform.GetSocVersion();
131 :
132 : libapiSize_ = ascendcPlatform.GetLibApiWorkSpaceSize();
133 :
134 : ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize_);
135 : ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L1, l1Size_);
136 : ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_C, l0cSize_);
137 : ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::L0_B, l0bSize_);
138 :
139 : aivNum_ = ascendcPlatform.GetCoreNumAiv();
140 : aicNum_ = ascendcPlatform.GetCoreNumAic();
141 :
142 : OP_CHECK_IF(
143 : aicNum_ == 0 || aivNum_ == 0,
144 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "aicNum/aivNum", "num of core obtained is 0"),
145 : return GRAPH_FAILED);
146 :
147 : return ge::GRAPH_SUCCESS;
148 : }
149 :
150 : ge::graphStatus CompressorTiling::SetBaseInfo()
151 : {
152 : if (context_->x.shape->GetStorageShape().GetDimNum() == COMPRESSOR_DIM_NUM_3) {
153 : baseParams_->batchSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0);
154 : baseParams_->seqSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1);
155 : baseParams_->hiddenSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_2);
156 : baseParams_->tokenSize = baseParams_->batchSize * baseParams_->seqSize;
157 : } else {
158 : baseParams_->batchSize = context_->cuSeqlens.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0) - 1;
159 : baseParams_->tokenSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0);
160 : baseParams_->hiddenSize = context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1);
161 : }
162 :
163 : coff = static_cast<uint8_t>(*context_->coff);
164 : baseParams_->headDim = context_->wkv.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0) / coff;
165 : baseParams_->cmpRatio = static_cast<uint32_t>(*context_->cmpRatio);
166 : baseParams_->csSize = baseParams_->seqSize - (baseParams_->seqSize % baseParams_->cmpRatio);
167 : baseParams_->stateCacheStrideDim0 = static_cast<uint64_t>(*context_->stateCacheStrideDim0);
168 : baseParams_->nSize = 2; // 2:每个核处理两个基本块后做全核同步
169 : baseParams_->usedCoreNum = aicNum_;
170 :
171 : OP_LOGI(context_->opName, "[TILING] bSize:%u tSize:%u cmpRatio:%u coff:%u", baseParams_->batchSize,
172 : baseParams_->tokenSize, baseParams_->cmpRatio, coff);
173 :
174 : return ge::GRAPH_SUCCESS;
175 : }
176 :
177 : ge::graphStatus CompressorTiling::SetPageAttentionInfo()
178 : {
179 : pageAttentionParams_->blockNum = context_->stateCache.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0);
180 : pageAttentionParams_->blockSize = context_->stateCache.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1);
181 : if (static_cast<uint8_t>(*context_->cacheMode) == static_cast<uint8_t>(CACHE_MODE::LINEAR_BUFFER)) {
182 : pageAttentionParams_->maxBlockNumPerBatch =
183 : context_->stateBlockTable.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1);
184 : }
185 :
186 : return ge::GRAPH_SUCCESS;
187 : }
188 :
189 : ge::graphStatus CompressorTiling::SetWorkSpaceInfo()
190 : {
191 : workspaceParams_->dbWorkspaceRatio = 2;
192 : workspaceParams_->mm1KvResSize = innerSplitParams_->mBaseSize * baseParams_->headDim * coff;
193 : workspaceParams_->mm1ScoreResSize = innerSplitParams_->mBaseSize * baseParams_->headDim * coff;
194 : if (coff == 2) {
195 : workspaceParams_->vec1TailCacheSize = baseParams_->cmpRatio * baseParams_->headDim;
196 : }
197 : workspaceParams_->vec1ResSize = innerSplitParams_->mBaseSize * baseParams_->headDim * baseParams_->nSize;
198 :
199 : return ge::GRAPH_SUCCESS;
200 : }
201 :
202 8 : ge::graphStatus CompressorTiling::SetScenarioInfo()
203 : {
204 8 : return ge::GRAPH_SUCCESS;
205 : }
206 :
207 : ge::graphStatus CompressorTiling::SetTemplateId()
208 : {
209 : if (context_->templateId == TemplateId::EMPTY_X) {
210 : return ge::GRAPH_SUCCESS;
211 : }
212 : // 设置高性能模板
213 : if (context_->layout == LayoutType::LAYOUT_BSH && baseParams_->seqSize <= 4 && baseParams_->tokenSize <= 128) {
214 : context_->templateId = TemplateId::FULL_LOAD;
215 : }
216 : return ge::GRAPH_SUCCESS;
217 : }
218 :
219 : ge::graphStatus CompressorTiling::SetInnerSplitInfo()
220 : {
221 : if (context_->templateId == TemplateId::FULL_LOAD) {
222 : uint32_t kAlignNum = baseParams_->hiddenSize / 128;
223 : innerSplitParams_->mBaseSize = 128; // 256:核间切分，M轴基本块大小
224 : innerSplitParams_->dBaseSize = 256 / (coff * 2); // nBase = dBase * coff * 2
225 : uint32_t dBaseNum = baseParams_->headDim / innerSplitParams_->dBaseSize;
226 : uint32_t mBaseNum = (baseParams_->tokenSize + innerSplitParams_->mBaseSize - 1) / innerSplitParams_->mBaseSize;
227 : baseParams_->coreGroupNum = baseParams_->usedCoreNum / dBaseNum;
228 : baseParams_->kBaseNum = 1;
229 : baseParams_->kBaseSize = baseParams_->hiddenSize;
230 : if ((dBaseNum * mBaseNum) < baseParams_->usedCoreNum) {
231 : baseParams_->kBaseNum = baseParams_->usedCoreNum / dBaseNum;
232 : baseParams_->kBaseSize = kAlignNum / baseParams_->kBaseNum * 128;
233 : }
234 : for (uint32_t i = 0; i < baseParams_->usedCoreNum; i++) {
235 : baseParams_->splitCoreParam[i].nStart = (i % dBaseNum) * innerSplitParams_->dBaseSize;
236 : baseParams_->splitCoreParam[i].nEnd = baseParams_->splitCoreParam[i].nStart + innerSplitParams_->dBaseSize;
237 : if (baseParams_->kBaseNum > 1) {
238 : uint32_t kStartIdx = i / dBaseNum;
239 : uint32_t dealKSize = baseParams_->kBaseSize;
240 : if (kStartIdx < kAlignNum % baseParams_->kBaseNum) {
241 : dealKSize += 128;
242 : baseParams_->splitCoreParam[i].kStart = kStartIdx * dealKSize;
243 : } else if (kStartIdx < baseParams_->kBaseNum) {
244 : baseParams_->splitCoreParam[i].kStart =
245 : kStartIdx * baseParams_->kBaseSize + (kAlignNum % baseParams_->kBaseNum) * 128;
246 : } else {
247 : dealKSize = 0;
248 : baseParams_->splitCoreParam[i].kStart = 0;
249 : }
250 : baseParams_->splitCoreParam[i].kEnd = baseParams_->splitCoreParam[i].kStart + dealKSize;
251 : baseParams_->splitCoreParam[i].mStart = 0;
252 : baseParams_->splitCoreParam[i].mEnd = baseParams_->tokenSize;
253 : baseParams_->mLoopNum = 1;
254 : } else {
255 : baseParams_->splitCoreParam[i].kStart = 0;
256 : baseParams_->splitCoreParam[i].kEnd = baseParams_->hiddenSize;
257 : baseParams_->splitCoreParam[i].mStart = (i / dBaseNum) * innerSplitParams_->mBaseSize;
258 : baseParams_->splitCoreParam[i].mEnd =
259 : baseParams_->splitCoreParam[i].mStart + innerSplitParams_->mBaseSize;
260 : baseParams_->mLoopNum = mBaseNum / baseParams_->coreGroupNum;
261 : }
262 : }
263 : } else {
264 : if (coff == 2) {
265 : innerSplitParams_->mBaseSize = 128;
266 : } else {
267 : innerSplitParams_->mBaseSize = 256;
268 : }
269 : innerSplitParams_->dBaseSize = 64;
270 : }
271 : return ge::GRAPH_SUCCESS;
272 : }
273 :
274 : ge::graphStatus CompressorTiling::CalcWorkSpace()
275 : {
276 : constexpr uint32_t MM1_RES_ELEM_SIZE = 4; // 4: fp32
277 : constexpr uint32_t V1_RES_ELEM_SIZE = 4; // 4: fp32
278 : uint32_t maxGroupNum = aicNum_ / (baseParams_->headDim / innerSplitParams_->dBaseSize);
279 : workspaceSize_ = libapiSize_;
280 : workspaceSize_ +=
281 : workspaceParams_->mm1KvResSize * maxGroupNum * MM1_RES_ELEM_SIZE * workspaceParams_->dbWorkspaceRatio;
282 : workspaceSize_ +=
283 : workspaceParams_->mm1ScoreResSize * maxGroupNum * MM1_RES_ELEM_SIZE * workspaceParams_->dbWorkspaceRatio;
284 : workspaceSize_ +=
285 : workspaceParams_->vec1TailCacheSize * MM1_RES_ELEM_SIZE * workspaceParams_->dbWorkspaceRatio * 2; // 2 kv和score
286 : workspaceSize_ +=
287 : workspaceParams_->vec1ResSize * maxGroupNum * V1_RES_ELEM_SIZE * workspaceParams_->dbWorkspaceRatio;
288 :
289 : if (context_->workSpaces) {
290 : context_->workSpaces[0] = workspaceSize_;
291 : }
292 :
293 : OP_LOGI(context_->opName, "Tiling info: workspaceSize_ = %zu", workspaceSize_);
294 : return ge::GRAPH_SUCCESS;
295 : }
296 :
297 : ge::graphStatus CompressorTiling::CheckEmptyTensor() const
298 : {
299 : if ((context_->layout == LayoutType::LAYOUT_BSH &&
300 : context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0) == 0) ||
301 : (context_->layout == LayoutType::LAYOUT_BSH &&
302 : context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1) == 0) ||
303 : (context_->layout == LayoutType::LAYOUT_TH &&
304 : context_->x.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0) == 0)) {
305 : context_->templateId = TemplateId::EMPTY_X;
306 : } else {
307 : if (context_->x.shape->GetStorageShape().GetShapeSize() == 0 ||
308 : context_->wkv.shape->GetStorageShape().GetShapeSize() == 0 ||
309 : context_->wgate.shape->GetStorageShape().GetShapeSize() == 0 ||
310 : context_->stateCache.shape->GetStorageShape().GetShapeSize() == 0 ||
311 : context_->ape.shape->GetStorageShape().GetShapeSize() == 0 ||
312 : context_->stateBlockTable.shape->GetStorageShape().GetShapeSize() == 0) {
313 : OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->opName, "x", "0",
314 : "Only input tensor x dim B or S or T supports to be 0");
315 : return ge::GRAPH_FAILED;
316 : }
317 : context_->templateId = TemplateId::NORMAL;
318 : }
319 : return ge::GRAPH_SUCCESS;
320 : }
321 :
322 : ge::graphStatus CompressorTiling::RunBigKernelTiling(CompressorTilingData *tilingData)
323 : {
324 : this->baseParams_ = &tilingData->baseParams;
325 : this->pageAttentionParams_ = &tilingData->pageAttentionParams;
326 : this->innerSplitParams_ = &tilingData->innerSplitParams;
327 : this->workspaceParams_ = &tilingData->workspaceParams;
328 : using StatusFunction = std::function<ge::graphStatus()>;
329 : std::vector<StatusFunction> requiredTilingFuncs{std::bind(&CompressorTiling::GetNpuInfo, this),
330 : std::bind(&CompressorTiling::CheckRequiredParaExistence, this),
331 : std::bind(&CompressorTiling::CheckEmptyTensor, this),
332 : std::bind(&CompressorTiling::CheckSinglePara, this),
333 : std::bind(&CompressorTiling::SetBaseInfo, this),
334 : std::bind(&CompressorTiling::SetPageAttentionInfo, this),
335 : std::bind(&CompressorTiling::CheckFeature, this),
336 : std::bind(&CompressorTiling::CheckMultiParaConsistency, this),
337 : std::bind(&CompressorTiling::CheckBlockDimConstrain, this),
338 : std::bind(&CompressorTiling::SetTemplateId, this),
339 : std::bind(&CompressorTiling::SetInnerSplitInfo, this),
340 : std::bind(&CompressorTiling::SetWorkSpaceInfo, this),
341 : std::bind(&CompressorTiling::SetScenarioInfo, this)};
342 : for (const auto &func : requiredTilingFuncs) {
343 : if (func() != ge::GRAPH_SUCCESS) {
344 : return ge::GRAPH_FAILED;
345 : }
346 : }
347 :
348 : if (context_->templateId == TemplateId::EMPTY_X) {
349 : workspaceSize_ = libapiSize_;
350 : if (context_->workSpaces) {
351 : context_->workSpaces[0] = workspaceSize_;
352 : }
353 : GenTilingKey();
354 : context_->blockDim = 1U;
355 : return ge::GRAPH_SUCCESS;
356 : }
357 : std::vector<StatusFunction> optionalTilingFuncs{std::bind(&CompressorTiling::CalcWorkSpace, this),
358 : std::bind(&CompressorTiling::GenTilingKey, this)};
359 : for (const auto &func : optionalTilingFuncs) {
360 : if (func() != ge::GRAPH_SUCCESS) {
361 : return ge::GRAPH_FAILED;
362 : }
363 : }
364 :
365 : context_->blockDim = aicNum_;
366 :
367 : OP_LOGI("Run big kernel");
368 :
369 : return ge::GRAPH_SUCCESS;
370 : }
371 :
372 : ge::graphStatus CompressorTiling::GenTilingKey() const
373 : {
374 : // 0:BF16, 1:FP16
375 : uint8_t dtype = 0;
376 : // 0: BSH 1:TH
377 : uint8_t layout = 0;
378 : uint8_t templateId = static_cast<uint8_t>(context_->templateId);
379 : uint8_t cacheMode = static_cast<uint8_t>(*context_->cacheMode);
380 :
381 : auto xDtype = context_->x.desc->GetDataType();
382 : if (xDtype == ge::DT_BF16) {
383 : dtype = 0;
384 : } else if (xDtype == ge::DT_FLOAT16) {
385 : dtype = 1;
386 : }
387 : auto xDimNum = context_->x.shape->GetStorageShape().GetDimNum();
388 : if (xDimNum == COMPRESSOR_DIM_NUM_3) {
389 : layout = 0;
390 : } else {
391 : layout = 1;
392 : }
393 :
394 : context_->tilingKey = GET_TPL_TILING_KEY(layout, dtype, coff, cacheMode, templateId, 0);
395 : OP_LOGI(context_->opName, "Compressor dtype:%hhu layout:%hhu coff:%hhu, cacheMode: %u, template_id:%hhu", dtype,
396 : layout, coff, cacheMode, templateId);
397 : OP_LOGI(context_->opName, "Compressor tilingKey:%lu", context_->tilingKey);
398 :
399 : return ge::GRAPH_SUCCESS;
400 : }
401 :
402 : ge::graphStatus CompressorTiling::CheckSinglePara() const
403 : {
404 : if (ge::GRAPH_SUCCESS != CheckSingleParaX() || ge::GRAPH_SUCCESS != CheckSingleParaWkv() ||
405 : ge::GRAPH_SUCCESS != CheckSingleParaWgate() || ge::GRAPH_SUCCESS != CheckSingleParaStateCache() ||
406 : ge::GRAPH_SUCCESS != CheckSingleParaApe() || ge::GRAPH_SUCCESS != CheckSingleParaStateBlockTable() ||
407 : ge::GRAPH_SUCCESS != CheckSingleParaCuSeqlens() || ge::GRAPH_SUCCESS != CheckSingleParaSeqused() ||
408 : ge::GRAPH_SUCCESS != CheckSingleParaStartPos() || ge::GRAPH_SUCCESS != CheckSingleParaCmpKv() ||
409 : ge::GRAPH_SUCCESS != CheckSingleParaCmpRatio() || ge::GRAPH_SUCCESS != CheckSingleParaCoff() ||
410 42 : ge::GRAPH_SUCCESS != CheckSingleParaCacheMode() || ge::GRAPH_SUCCESS != CheckSingleParaGradEnabled()) {
411 : return ge::GRAPH_FAILED;
412 : }
413 : return ge::GRAPH_SUCCESS;
414 : }
415 :
416 : template <typename T>
417 : ge::graphStatus CompressorTiling::CheckFeatureValueSupport(const T *featureValue,
418 : const std::vector<T> &expectFeatureValList,
419 : const std::string &name) const
420 : {
421 : if (std::find(expectFeatureValList.begin(), expectFeatureValList.end(), *featureValue) ==
422 : expectFeatureValList.end()) {
423 : LogErrorNumberSupport(expectFeatureValList, *featureValue, name, "feature value");
424 : return ge::GRAPH_FAILED;
425 : }
426 : return ge::GRAPH_SUCCESS;
427 : }
428 :
429 : template <typename T>
430 : ge::graphStatus CompressorTiling::CheckAttrValueSupport(const T *attrValue, const std::vector<T> &expectAttrValList,
431 : const std::string &name) const
432 : {
433 : if (attrValue == nullptr) {
434 : return ge::GRAPH_SUCCESS;
435 : }
436 :
437 : if (std::find(expectAttrValList.begin(), expectAttrValList.end(), *attrValue) == expectAttrValList.end()) {
438 : LogErrorNumberSupport(expectAttrValList, *attrValue, name, "attr value");
439 : return ge::GRAPH_FAILED;
440 : }
441 :
442 : return ge::GRAPH_SUCCESS;
443 : }
444 :
445 : template <typename T>
446 : std::string to_string(const T &value)
447 : {
448 : if (std::is_same_v<T, bool>) {
449 : return value ? "true" : "false";
450 : } else {
451 : return std::to_string(value);
452 : }
453 : }
454 :
455 : template <typename T>
456 : void CompressorTiling::LogErrorNumberSupport(const std::vector<T> &expectNumberList, const T &actualValue,
457 : const std::string &name, const std::string subName) const
458 : {
459 : std::ostringstream oss;
460 : for (size_t i = 0; i < expectNumberList.size(); ++i) {
461 : oss << to_string(expectNumberList[i]);
462 : if (i < expectNumberList.size() - 1) {
463 : oss << ", ";
464 : }
465 : }
466 :
467 : OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->opName, name, to_string(actualValue),
468 : subName + " only supports " + oss.str());
469 : }
470 :
471 : static std::string LayoutTypeToStr(LayoutType layout)
472 : {
473 : switch (layout) {
474 : case LayoutType::LAYOUT_BSH:
475 : return "BSH";
476 : case LayoutType::LAYOUT_TH:
477 : return "TH";
478 : default:
479 : return "UNKNOWN_LAYOUT";
480 : }
481 : }
482 :
483 : ge::graphStatus CompressorTiling::CheckDimNumInLayoutSupport(const std::string &layout, const gert::StorageShape *shape,
484 : const std::string &name) const
485 : {
486 : const auto &dimIt = LAYOUT_DIM_MAP.find(layout);
487 : OP_CHECK_IF(
488 : shape->GetStorageShape().GetDimNum() != dimIt->second,
489 : OP_LOGE_FOR_INVALID_SHAPEDIM(context_->opName, name, std::to_string(shape->GetStorageShape().GetDimNum()),
490 : std::to_string(dimIt->second)),
491 : return ge::GRAPH_FAILED);
492 : return ge::GRAPH_SUCCESS;
493 : }
494 :
495 : ge::graphStatus CompressorTiling::CheckDtypeSupport(const gert::CompileTimeTensorDesc *desc,
496 : const std::string &name) const
497 : {
498 : if (desc != nullptr) {
499 : const auto &it = DTYPE_SUPPORT_MAP.find(name);
500 : OP_CHECK_IF(it == DTYPE_SUPPORT_MAP.end(),
501 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(
502 : context_->opName, name, "datatype support list should be specify in DTYPE_SUPPORT_MAP"),
503 : return ge::GRAPH_FAILED);
504 : auto &expectDtypeList = it->second;
505 : OP_CHECK_IF(
506 : std::find(expectDtypeList.begin(), expectDtypeList.end(), desc->GetDataType()) == expectDtypeList.end(),
507 : LogErrorDtypeSupport(expectDtypeList, desc->GetDataType(), name), return ge::GRAPH_FAILED);
508 : }
509 : return ge::GRAPH_SUCCESS;
510 : }
511 :
512 : void CompressorTiling::LogErrorDtypeSupport(const std::vector<ge::DataType> &expectDtypeList,
513 : const ge::DataType &actualDtype, const std::string &name) const
514 : {
515 : std::ostringstream oss;
516 : for (size_t i = 0; i < expectDtypeList.size(); ++i) {
517 : oss << DataTypeToSerialString(expectDtypeList[i]);
518 : if (i < expectDtypeList.size() - 1) {
519 : oss << ", ";
520 : }
521 : }
522 : OP_LOGE_FOR_INVALID_DTYPE(context_->opName, name, DataTypeToSerialString(actualDtype), oss.str());
523 : }
524 :
525 : static std::string DataTypeToSerialString(ge::DataType type)
526 : {
527 : const auto it = DATATYPE_TO_STRING_MAP.find(type);
528 : if (it != DATATYPE_TO_STRING_MAP.end()) {
529 : return it->second;
530 : } else {
531 : OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON("Compressor", "datatype", std::to_string(static_cast<int32_t>(type)),
532 : "not support");
533 : return "UNDEFINED";
534 : }
535 : }
536 :
537 : ge::graphStatus CompressorTiling::CheckDimNumSupport(const gert::StorageShape *shape, const std::string &name) const
538 : {
539 : if (shape == nullptr) {
540 : return ge::GRAPH_SUCCESS;
541 : }
542 : const auto &it = DIM_NUM_MAP.find(name);
543 : OP_CHECK_IF(it == DIM_NUM_MAP.end(),
544 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, name,
545 : "dim number support list should be specify in DIM_NUM_MAP"),
546 : return ge::GRAPH_FAILED);
547 : auto &expectDimNumList = it->second;
548 : OP_CHECK_IF(
549 : std::find(expectDimNumList.begin(), expectDimNumList.end(), shape->GetStorageShape().GetDimNum()) ==
550 : expectDimNumList.end(),
551 : [&]() {
552 : std::ostringstream oss;
553 : for (size_t i = 0; i < expectDimNumList.size(); ++i) {
554 : oss << expectDimNumList[i];
555 : if (i < expectDimNumList.size() - 1) {
556 : oss << " or ";
557 : }
558 : }
559 : OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(context_->opName, name,
560 : std::to_string(shape->GetStorageShape().GetDimNum()),
561 : name + " dimension should be " + oss.str());
562 : }(),
563 : return ge::GRAPH_FAILED);
564 : return ge::GRAPH_SUCCESS;
565 : }
566 :
567 : ge::graphStatus CompressorTiling::CheckSingleParaX() const
568 : {
569 : if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->x.desc, X_NAME) ||
570 : ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->x.shape, X_NAME) ||
571 : ge::GRAPH_SUCCESS != CheckDimNumInLayoutSupport(LayoutTypeToStr(context_->layout), context_->x.shape, X_NAME)) {
572 : return ge::GRAPH_FAILED;
573 : }
574 : OP_CHECK_IF(context_->x.shape->GetStorageShape().GetDim(context_->x.shape->GetStorageShape().GetDimNum() - 1) >
575 : MAX_HIDDEN_SIZE ||
576 : context_->x.shape->GetStorageShape().GetDim(context_->x.shape->GetStorageShape().GetDimNum() - 1) <
577 : MIN_HIDDEN_SIZE ||
578 : context_->x.shape->GetStorageShape().GetDim(context_->x.shape->GetStorageShape().GetDimNum() - 1) %
579 : ALIGN_FACTOR_HIDDEN_SIZE !=
580 : 0,
581 : OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
582 : context_->opName, X_NAME,
583 : "dim " + std::to_string(context_->x.shape->GetStorageShape().GetDimNum() - 1) + "=" +
584 : std::to_string(context_->x.shape->GetStorageShape().GetDim(
585 : context_->x.shape->GetStorageShape().GetDimNum() - 1)),
586 : "x should be within [" + std::to_string(MIN_HIDDEN_SIZE) + ", " + std::to_string(MAX_HIDDEN_SIZE) +
587 : "] and be 512-aligned"),
588 : return ge::GRAPH_FAILED);
589 : return ge::GRAPH_SUCCESS;
590 : }
591 :
592 : ge::graphStatus CompressorTiling::CheckSingleParaWkv() const
593 : {
594 : if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->wkv.desc, WKV_NAME) ||
595 : ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->wkv.shape, WKV_NAME)) {
596 : return ge::GRAPH_FAILED;
597 : }
598 : uint32_t coffVal = static_cast<uint32_t>(*context_->coff);
599 : uint32_t headDim = context_->wkv.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_0) / coffVal;
600 : if (ge::GRAPH_SUCCESS != CheckFeatureValueSupport(&headDim, HEAD_DIM, WKV_NAME)) {
601 : return ge::GRAPH_FAILED;
602 : }
603 : return ge::GRAPH_SUCCESS;
604 : }
605 :
606 : ge::graphStatus CompressorTiling::CheckSingleParaWgate() const
607 : {
608 : if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->wgate.desc, WGATE_NAME) ||
609 : ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->wgate.shape, WGATE_NAME)) {
610 : return ge::GRAPH_FAILED;
611 : }
612 : return ge::GRAPH_SUCCESS;
613 : }
614 :
615 : ge::graphStatus CompressorTiling::CheckSingleParaStateCache() const
616 : {
617 : if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->stateCache.desc, STATE_CACHE_NAME) ||
618 : ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->stateCache.shape, STATE_CACHE_NAME)) {
619 : return ge::GRAPH_FAILED;
620 : }
621 : uint32_t blockSize = context_->stateCache.shape->GetStorageShape().GetDim(COMPRESSOR_DIM_INDEX_1);
622 : OP_CHECK_IF(
623 : blockSize > MAX_BLOCK_SIZE || blockSize < MIN_BLOCK_SIZE,
624 : OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->opName, STATE_CACHE_NAME, "dim 1=" + std::to_string(blockSize),
625 : "state_cache dim 1 should be within [" + std::to_string(MIN_BLOCK_SIZE) +
626 : ", " + std::to_string(MAX_BLOCK_SIZE) + "]"),
627 : return ge::GRAPH_FAILED);
628 : uint64_t contiguousStride0 =
629 : context_->stateCache.shape->GetShape().GetDim(1) * context_->stateCache.shape->GetShape().GetDim(2);
630 : uint64_t stateCacheStrideDim0 = static_cast<uint64_t>(*context_->stateCacheStrideDim0);
631 : OP_CHECK_IF(stateCacheStrideDim0 < contiguousStride0,
632 : OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
633 : context_->opName, STATE_CACHE_NAME, "stride0=" + std::to_string(stateCacheStrideDim0),
634 : "state_cache stride0 must be greater than or equal to contiguous stride " +
635 : std::to_string(contiguousStride0) + ", only axis 0 non-contiguous is supported"),
636 : return ge::GRAPH_FAILED);
637 : return ge::GRAPH_SUCCESS;
638 : }
639 :
640 : ge::graphStatus CompressorTiling::CheckSingleParaApe() const
641 : {
642 : if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->ape.desc, APE_NAME) ||
643 : ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->ape.shape, APE_NAME)) {
644 : return ge::GRAPH_FAILED;
645 : }
646 : return ge::GRAPH_SUCCESS;
647 : }
648 :
649 : ge::graphStatus CompressorTiling::CheckSingleParaStateBlockTable() const
650 : {
651 : if (context_->stateBlockTable.desc == nullptr) {
652 : return ge::GRAPH_SUCCESS;
653 : }
654 : if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->stateBlockTable.desc, STATE_BLOCK_TABLE_NAME) ||
655 : ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->stateBlockTable.shape, STATE_BLOCK_TABLE_NAME)) {
656 : return ge::GRAPH_FAILED;
657 : }
658 : return ge::GRAPH_SUCCESS;
659 : }
660 :
661 : ge::graphStatus CompressorTiling::CheckSingleParaCuSeqlens() const
662 : {
663 : if (context_->cuSeqlens.desc == nullptr) {
664 : return ge::GRAPH_SUCCESS;
665 : }
666 : if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->cuSeqlens.desc, CU_SEQLENS_NAME) ||
667 : ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->cuSeqlens.shape, CU_SEQLENS_NAME)) {
668 : return ge::GRAPH_FAILED;
669 : }
670 : return ge::GRAPH_SUCCESS;
671 : }
672 :
673 : ge::graphStatus CompressorTiling::CheckSingleParaSeqused() const
674 : {
675 : if (context_->seqUsed.desc == nullptr) {
676 : return ge::GRAPH_SUCCESS;
677 : }
678 : if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->seqUsed.desc, SEQUSED_NAME) ||
679 : ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->seqUsed.shape, SEQUSED_NAME)) {
680 : return ge::GRAPH_FAILED;
681 : }
682 : return ge::GRAPH_SUCCESS;
683 : }
684 :
685 : ge::graphStatus CompressorTiling::CheckSingleParaStartPos() const
686 : {
687 : if (context_->startPos.desc == nullptr) {
688 : return ge::GRAPH_SUCCESS;
689 : }
690 : if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->startPos.desc, START_POS_NAME) ||
691 : ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->startPos.shape, START_POS_NAME)) {
692 : return ge::GRAPH_FAILED;
693 : }
694 : return ge::GRAPH_SUCCESS;
695 : }
696 :
697 : ge::graphStatus CompressorTiling::CheckSingleParaCmpKv() const
698 : {
699 : if (context_->cmpKv.desc == nullptr) {
700 : return ge::GRAPH_SUCCESS;
701 : }
702 : if (ge::GRAPH_SUCCESS != CheckDtypeSupport(context_->cmpKv.desc, CMP_KV_NAME) ||
703 : ge::GRAPH_SUCCESS != CheckDimNumSupport(context_->cmpKv.shape, CMP_KV_NAME)) {
704 : return ge::GRAPH_FAILED;
705 : }
706 : return ge::GRAPH_SUCCESS;
707 : }
708 :
709 : ge::graphStatus CompressorTiling::CheckSingleParaCmpRatio() const
710 : {
711 : if (ge::GRAPH_SUCCESS != CheckAttrValueSupport(context_->cmpRatio, CMP_RATIO, CMP_RATIO_NAME)) {
712 : return ge::GRAPH_FAILED;
713 : }
714 : return ge::GRAPH_SUCCESS;
715 : }
716 :
717 : ge::graphStatus CompressorTiling::CheckSingleParaCoff() const
718 : {
719 : if (ge::GRAPH_SUCCESS != CheckAttrValueSupport(context_->coff, COFF, COFF_NAME)) {
720 : return ge::GRAPH_FAILED;
721 : }
722 : return ge::GRAPH_SUCCESS;
723 : }
724 :
725 : ge::graphStatus CompressorTiling::CheckSingleParaCacheMode() const
726 : {
727 : // A3 does not support the ring buffer mode (cache_mode=2).
728 : if (ge::GRAPH_SUCCESS != CheckAttrValueSupport(context_->cacheMode, CACHE_MODE, CACHE_MODE_NAME)) {
729 : return ge::GRAPH_FAILED;
730 : }
731 : return ge::GRAPH_SUCCESS;
732 : }
733 :
734 11 : ge::graphStatus CompressorTiling::CheckSingleParaGradEnabled() const
735 : {
736 : // A3 only supports grad_enabled=false because arch22 has no backward-output path.
737 12 : OP_CHECK_IF(context_->gradEnabled != nullptr && *context_->gradEnabled,
738 : OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->opName, GRAD_ENABLED_NAME, "true",
739 : "grad_enabled only supports false"),
740 : return ge::GRAPH_FAILED);
741 : return ge::GRAPH_SUCCESS;
742 : }
743 :
744 : ge::graphStatus CompressorTiling::CheckRequiredParaExistence() const
745 : {
746 : if (CheckRequiredInOutExistence() != ge::GRAPH_SUCCESS || CheckRequiredAttrExistence() != ge::GRAPH_SUCCESS) {
747 : return ge::GRAPH_FAILED;
748 : }
749 : return ge::GRAPH_SUCCESS;
750 : }
751 :
752 : ge::graphStatus CompressorTiling::CheckRequiredInOutExistence() const
753 : {
754 : OP_CHECK_IF(context_->x.shape == nullptr,
755 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "x", "shape is nullptr"),
756 : return ge::GRAPH_FAILED);
757 : OP_CHECK_IF(context_->x.desc == nullptr,
758 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "x", "desc is nullptr"),
759 : return ge::GRAPH_FAILED);
760 : OP_CHECK_IF(context_->wkv.shape == nullptr,
761 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "wkv", "shape is nullptr"),
762 : return ge::GRAPH_FAILED);
763 : OP_CHECK_IF(context_->wkv.desc == nullptr,
764 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "wkv", "desc is nullptr"),
765 : return ge::GRAPH_FAILED);
766 : OP_CHECK_IF(context_->wgate.shape == nullptr,
767 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "wgate", "shape is nullptr"),
768 : return ge::GRAPH_FAILED);
769 : OP_CHECK_IF(context_->wgate.desc == nullptr,
770 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "wgate", "desc is nullptr"),
771 : return ge::GRAPH_FAILED);
772 : OP_CHECK_IF(context_->stateCache.shape == nullptr,
773 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "state_cache", "shape is nullptr"),
774 : return ge::GRAPH_FAILED);
775 : OP_CHECK_IF(context_->stateCache.desc == nullptr,
776 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "state_cache", "desc is nullptr"),
777 : return ge::GRAPH_FAILED);
778 : OP_CHECK_IF(context_->ape.shape == nullptr,
779 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "ape", "shape is nullptr"),
780 : return ge::GRAPH_FAILED);
781 : OP_CHECK_IF(context_->ape.desc == nullptr,
782 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "ape", "desc is nullptr"),
783 : return ge::GRAPH_FAILED);
784 : OP_CHECK_IF(context_->stateBlockTable.shape == nullptr,
785 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "state_block_table", "shape is nullptr"),
786 : return ge::GRAPH_FAILED);
787 : OP_CHECK_IF(context_->stateBlockTable.desc == nullptr,
788 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "state_block_table", "desc is nullptr"),
789 : return ge::GRAPH_FAILED);
790 : OP_CHECK_IF(context_->cmpKv.shape == nullptr,
791 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "cmp_kv", "shape is nullptr"),
792 : return ge::GRAPH_FAILED);
793 : OP_CHECK_IF(context_->cmpKv.desc == nullptr,
794 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "cmp_kv", "desc is nullptr"),
795 : return ge::GRAPH_FAILED);
796 : if (context_->layout == LayoutType::LAYOUT_TH) {
797 : OP_CHECK_IF(context_->cuSeqlens.desc == nullptr,
798 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "cu_seqlens",
799 : "cu_seqlens should not be nullptr in TH layout"),
800 : return ge::GRAPH_FAILED);
801 : OP_CHECK_IF(context_->cuSeqlens.shape == nullptr,
802 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "cu_seqlens",
803 : "cu_seqlens should not be nullptr in TH layout"),
804 : return ge::GRAPH_FAILED);
805 : } else {
806 : OP_CHECK_IF(context_->cuSeqlens.desc != nullptr,
807 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "cu_seqlens",
808 : "cu_seqlens must be nullptr in BSH layout"),
809 : return ge::GRAPH_FAILED);
810 : OP_CHECK_IF(context_->cuSeqlens.shape != nullptr,
811 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "cu_seqlens",
812 : "cu_seqlens must be nullptr in BSH layout"),
813 : return ge::GRAPH_FAILED);
814 : }
815 : return ge::GRAPH_SUCCESS;
816 : }
817 :
818 : ge::graphStatus CompressorTiling::CheckRequiredAttrExistence() const
819 : {
820 : OP_CHECK_IF(context_->cmpRatio == nullptr,
821 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context_->opName, "cmp_ratio", "attr is nullptr"),
822 : return ge::GRAPH_FAILED);
823 : return ge::GRAPH_SUCCESS;
824 : }
825 :
826 10 : ge::graphStatus CompressorTiling::CheckFeature() const
827 : {
828 10 : return ge::GRAPH_SUCCESS;
829 : }
830 :
831 : ge::graphStatus CompressorTiling::LogErrorShapeConsistency(const std::string &name, const gert::StorageShape *shape,
832 : const uint32_t &dimNum, const std::string &subName,
833 : const uint32_t &expectNum) const
834 : {
835 : if (shape == nullptr) {
836 : return ge::GRAPH_SUCCESS;
837 : }
838 :
839 : const uint32_t actualNum = shape->GetStorageShape().GetDim(dimNum);
840 : OP_CHECK_IF(actualNum != expectNum,
841 : OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
842 : context_->opName, name, "dim " + std::to_string(dimNum) + "=" + std::to_string(actualNum),
843 : name + " should be equal to " + subName + ": " + std::to_string(expectNum)),
844 : return ge::GRAPH_FAILED);
845 :
846 : return ge::GRAPH_SUCCESS;
847 : }
848 :
849 : ge::graphStatus CompressorTiling::CheckShapeConsistency() const
850 : {
851 : auto coffD = coff * baseParams_->headDim;
852 : uint32_t stateNum = 2;
853 : if (ge::GRAPH_SUCCESS != LogErrorShapeConsistency("stateBlockTable", context_->stateBlockTable.shape,
854 : COMPRESSOR_DIM_INDEX_0, "batchSize", baseParams_->batchSize) ||
855 : ge::GRAPH_SUCCESS != LogErrorShapeConsistency("cuSeqlens", context_->cuSeqlens.shape, COMPRESSOR_DIM_INDEX_0,
856 : "batchSize+1", baseParams_->batchSize + 1) ||
857 : ge::GRAPH_SUCCESS != LogErrorShapeConsistency("seqUsed", context_->seqUsed.shape, COMPRESSOR_DIM_INDEX_0,
858 : "batchSize", baseParams_->batchSize) ||
859 : ge::GRAPH_SUCCESS != LogErrorShapeConsistency("startPos", context_->startPos.shape, COMPRESSOR_DIM_INDEX_0,
860 : "batchSize", baseParams_->batchSize) ||
861 : ge::GRAPH_SUCCESS != LogErrorShapeConsistency("wkv", context_->wkv.shape, COMPRESSOR_DIM_INDEX_1, "x",
862 : baseParams_->hiddenSize) ||
863 : ge::GRAPH_SUCCESS != LogErrorShapeConsistency("wgate", context_->wgate.shape, COMPRESSOR_DIM_INDEX_1, "x",
864 : baseParams_->hiddenSize) ||
865 : ge::GRAPH_SUCCESS != LogErrorShapeConsistency("wkv", context_->wkv.shape, COMPRESSOR_DIM_INDEX_0,
866 : "coff*headDim", static_cast<uint32_t>(coffD)) ||
867 : ge::GRAPH_SUCCESS != LogErrorShapeConsistency("wgate", context_->wgate.shape, COMPRESSOR_DIM_INDEX_0,
868 : "coff*headDim", static_cast<uint32_t>(coffD)) ||
869 : ge::GRAPH_SUCCESS != LogErrorShapeConsistency("stateCache", context_->stateCache.shape, COMPRESSOR_DIM_INDEX_2,
870 : "2*coff*headDim", stateNum * static_cast<uint32_t>(coffD)) ||
871 : ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ape", context_->ape.shape, COMPRESSOR_DIM_INDEX_1,
872 : "coff*headDim", static_cast<uint32_t>(coffD)) ||
873 : ge::GRAPH_SUCCESS != LogErrorShapeConsistency("ape", context_->ape.shape, COMPRESSOR_DIM_INDEX_0, "cmpRatio",
874 : baseParams_->cmpRatio)) {
875 : return ge::GRAPH_FAILED;
876 : }
877 : if (static_cast<uint8_t>(*context_->cacheMode) == static_cast<uint8_t>(CACHE_MODE::LINEAR_BUFFER) &&
878 : (ge::GRAPH_SUCCESS != LogErrorShapeConsistency("stateCache", context_->stateCache.shape, COMPRESSOR_DIM_INDEX_0,
879 : "blockNum", pageAttentionParams_->blockNum) ||
880 : ge::GRAPH_SUCCESS != LogErrorShapeConsistency("stateCache", context_->stateCache.shape, COMPRESSOR_DIM_INDEX_1,
881 : "blockSize", pageAttentionParams_->blockSize))) {
882 : return ge::GRAPH_FAILED;
883 : }
884 : return ge::GRAPH_SUCCESS;
885 : }
886 :
887 : ge::graphStatus CompressorTiling::CheckDtypeConsistencyX(const gert::CompileTimeTensorDesc *desc,
888 : const std::string &name) const
889 : {
890 : const auto actualDtype = desc->GetDataType();
891 : OP_CHECK_IF(actualDtype != context_->dtype,
892 : OP_LOGE_FOR_INVALID_DTYPE_WITH_REASON(
893 : context_->opName, name, DataTypeToSerialString(actualDtype),
894 : name + " should be same with x: " + DataTypeToSerialString(context_->dtype)),
895 : return ge::GRAPH_FAILED);
896 : return ge::GRAPH_SUCCESS;
897 : }
898 :
899 : ge::graphStatus CompressorTiling::CheckDtypeConsistency() const
900 : {
901 : if (CheckDtypeConsistencyX(context_->wkv.desc, WKV_NAME) != ge::GRAPH_SUCCESS ||
902 : CheckDtypeConsistencyX(context_->wgate.desc, WGATE_NAME) != ge::GRAPH_SUCCESS ||
903 : CheckDtypeConsistencyX(context_->cmpKv.desc, CMP_KV_NAME) != ge::GRAPH_SUCCESS) {
904 : return ge::GRAPH_FAILED;
905 : }
906 : return ge::GRAPH_SUCCESS;
907 : }
908 :
909 : ge::graphStatus CompressorTiling::CheckDimNumConsistency() const
910 : {
911 : auto xDimNum = context_->x.shape->GetStorageShape().GetDimNum();
912 : OP_CHECK_IF(
913 : xDimNum != context_->cmpKv.shape->GetStorageShape().GetDimNum(),
914 : OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(
915 : context_->opName, "cmp_kv, x",
916 : std::to_string(context_->cmpKv.shape->GetStorageShape().GetDimNum()) + ", " + std::to_string(xDimNum),
917 : "dim num of cmp_kv should be equal to x"),
918 : return ge::GRAPH_FAILED);
919 : return ge::GRAPH_SUCCESS;
920 : }
921 :
922 : ge::graphStatus CompressorTiling::CheckScenarioConsistency() const
923 : {
924 : auto curCmpratio = baseParams_->cmpRatio;
925 : auto curHeaddim = baseParams_->headDim;
926 : auto curCoff = static_cast<uint8_t>(*context_->coff);
927 : std::vector<uint32_t> curScenario{curCmpratio, curCoff, curHeaddim};
928 : const std::vector<std::vector<uint32_t>> allowdScenarios = {{4, 2, 512}, {4, 2, 128}, {128, 1, 512}};
929 :
930 : OP_CHECK_IF(std::find(allowdScenarios.begin(), allowdScenarios.end(), curScenario) == allowdScenarios.end(),
931 : OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(
932 : context_->opName, "cmpratio,coff,headdim",
933 : "cmpratio=" + std::to_string(curCmpratio) + ",coff=" + std::to_string(curCoff) +
934 : ",headdim=" + std::to_string(curHeaddim),
935 : "cmpratio,coff,headdim should be equal to {4, 2, 512}, {4, 2, 128}, {128, 1, 512}"),
936 : return ge::GRAPH_FAILED);
937 : return ge::GRAPH_SUCCESS;
938 : }
939 :
940 : ge::graphStatus CompressorTiling::CheckBlockDimConstrain() const
941 : {
942 : uint32_t minBlockNum = baseParams_->headDim / 64; // 64 is the largest dBaseSize
943 : OP_CHECK_IF(aicNum_ < minBlockNum,
944 : OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->opName, "aicNum", std::to_string(aicNum_),
945 : "aicNum should not be less than " + std::to_string(minBlockNum)),
946 : return ge::GRAPH_FAILED);
947 : return ge::GRAPH_SUCCESS;
948 : }
949 :
950 : ge::graphStatus CompressorTiling::CheckMultiParaConsistency() const
951 : {
952 : if (CheckShapeConsistency() != ge::GRAPH_SUCCESS || CheckDtypeConsistency() != ge::GRAPH_SUCCESS ||
953 : CheckDimNumConsistency() != ge::GRAPH_SUCCESS) {
954 : return ge::GRAPH_FAILED;
955 : }
956 : #ifdef DAY0_SCOPE
957 : if (CheckScenarioConsistency() != ge::GRAPH_SUCCESS) {
958 : return ge::GRAPH_FAILED;
959 : }
960 : #endif
961 : return ge::GRAPH_SUCCESS;
962 : }
963 :
964 : } // namespace
965 :
966 : CMP_EXTERN_C ge::graphStatus TilingCompressorArch22(gert::TilingContext *context)
967 : {
968 : OP_CHECK_IF(context == nullptr, OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON("Compressor", "context", "is nullptr"),
969 : return ge::GRAPH_FAILED);
970 :
971 : OP_LOGI("Getting Tiling");
972 :
973 : CompressorContext compressorContext{};
974 : if (CompressorTiling::ConvertContext(*context, compressorContext) != ge::GRAPH_SUCCESS) {
975 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(context->GetNodeName(), "context",
976 : "error occurred while converting tilingContext to Compressor context");
977 : return ge::GRAPH_FAILED;
978 : }
979 : CompressorTiling compressorTiling(&compressorContext);
980 : CompressorTilingData *tilingData = context->GetTilingData<CompressorTilingData>();
981 : OP_CHECK_IF(tilingData == nullptr,
982 : OP_LOGE_FOR_INVALID_ARGUMENT_WITH_REASON(compressorContext.opName, "tilingData", "is nullptr"),
983 : return ge::GRAPH_FAILED);
984 : // 使用SyncAll，需要设置为batchmode模式，所有核同时启动，否则多流方式下执行可能会卡死
985 : context->SetScheduleMode(BATCH_MODE_SCHEDULE);
986 : if (compressorTiling.RunBigKernelTiling(tilingData) != ge::GRAPH_SUCCESS) {
987 : return ge::GRAPH_FAILED;
988 : }
989 : context->SetTilingKey(compressorContext.tilingKey);
990 : context->SetBlockDim(compressorContext.blockDim);
991 : OP_LOGI(compressorContext.opName, "block dim: %u.", compressorContext.blockDim);
992 : return ge::GRAPH_SUCCESS;
993 : }
994 :
995 : } // namespace optiling
|
