已合并
cleancode fix #3601
yangjinwen创建于 6月27日
cleancode fix #3601
已合并
yangjinwen创建于 6月27日
15 个文件变更+253-363
@@ -26,17 +26,27 @@ namespace brcto {
26static constexpr int64_t DIM_NUM_THRESHOLD_FOR_R4_SIZE = 5; // 触发尾轴4D尺寸计算的维度数阈值26static constexpr int64_t DIM_NUM_THRESHOLD_FOR_R4_SIZE = 5; // 触发尾轴4D尺寸计算的维度数阈值
27static constexpr int64_t TRAILING_DIM_NUM_FOR_R4_SIZE = 4; // 用于4D尺寸计算的尾轴维度数27static constexpr int64_t TRAILING_DIM_NUM_FOR_R4_SIZE = 4; // 用于4D尺寸计算的尾轴维度数
28 28 
29-ge::graphStatus GetABFlag(const gert::TilingContext* context, const gert::Shape& inShape, const gert::Shape& outShape,29+ge::graphStatus CheckSameDimNum(
30- std::array<bool, MAX_DIM_NUM>& abInfo)30+ const gert::TilingContext* context, const gert::Shape& inShape, const gert::Shape& outShape, size_t& dimNum)
31{31{
32- auto inDimNum = inShape.GetDimNum();32+ dimNum = inShape.GetDimNum();
33- if (inDimNum != outShape.GetDimNum()) {33+ if (dimNum != outShape.GetDimNum()) {
34- std::string dimMsg = std::to_string(inDimNum) + " and " + std::to_string(outShape.GetDimNum());34+ std::string dimMsg = std::to_string(dimNum) + " and " + std::to_string(outShape.GetDimNum());
35 std::string reasonMsg = "The input and output shape dim num should be equal.";35 std::string reasonMsg = "The input and output shape dim num should be equal.";
36 OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(36 OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(
37 context->GetNodeName(), "x and y", dimMsg.c_str(), reasonMsg.c_str());37 context->GetNodeName(), "x and y", dimMsg.c_str(), reasonMsg.c_str());
38 return ge::GRAPH_FAILED;38 return ge::GRAPH_FAILED;
39 }39 }
40+ return ge::GRAPH_SUCCESS;
41+}
42+ 
43+ge::graphStatus GetABFlag(const gert::TilingContext* context, const gert::Shape& inShape, const gert::Shape& outShape,
44+ std::array<bool, MAX_DIM_NUM>& abInfo)
45+{
46+ size_t inDimNum;
47+ if (CheckSameDimNum(context, inShape, outShape, inDimNum) != ge::GRAPH_SUCCESS) {
48+ return ge::GRAPH_FAILED;
49+ }
40 50 
41 for (size_t idx = 0; idx < inDimNum; idx++) {51 for (size_t idx = 0; idx < inDimNum; idx++) {
42 abInfo[idx] = (inShape[idx] != outShape[idx]);52 abInfo[idx] = (inShape[idx] != outShape[idx]);
@@ -47,12 +57,8 @@ ge::graphStatus GetABFlag(const gert::TilingContext* context, const gert::Shape&
47 57 
48ge::graphStatus MergeAxis(const gert::TilingContext* context, gert::Shape& inShape, gert::Shape& outShape)58ge::graphStatus MergeAxis(const gert::TilingContext* context, gert::Shape& inShape, gert::Shape& outShape)
49{59{
50- auto dimNum = inShape.GetDimNum();60+ size_t dimNum;
51- if (dimNum != outShape.GetDimNum()) {61+ if (CheckSameDimNum(context, inShape, outShape, dimNum) != ge::GRAPH_SUCCESS) {
52- std::string dimMsg = std::to_string(dimNum) + " and " + std::to_string(outShape.GetDimNum());
53- std::string reasonMsg = "The input and output shape dim num should be equal.";
54- OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(
55- context->GetNodeName(), "x and y", dimMsg.c_str(), reasonMsg.c_str());
56 return ge::GRAPH_FAILED;62 return ge::GRAPH_FAILED;
57 }63 }
58 64 
@@ -89,12 +95,8 @@ ge::graphStatus MergeAxis(const gert::TilingContext* context, gert::Shape& inSha
89 95 
90ge::graphStatus DeleteOneSizeAxis(const gert::TilingContext* context, gert::Shape& inShape, gert::Shape& outShape)96ge::graphStatus DeleteOneSizeAxis(const gert::TilingContext* context, gert::Shape& inShape, gert::Shape& outShape)
91{97{
92- auto dimNum = inShape.GetDimNum();98+ size_t dimNum;
93- if (dimNum != outShape.GetDimNum()) {99+ if (CheckSameDimNum(context, inShape, outShape, dimNum) != ge::GRAPH_SUCCESS) {
94- std::string dimMsg = std::to_string(dimNum) + " and " + std::to_string(outShape.GetDimNum());
95- std::string reasonMsg = "The input and output shape dim num should be equal.";
96- OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(
97- context->GetNodeName(), "x and y", dimMsg.c_str(), reasonMsg.c_str());
98 return ge::GRAPH_FAILED;100 return ge::GRAPH_FAILED;
99 }101 }
100 102 
@@ -143,16 +145,12 @@ void AdjustShapesToSameDimNum(gert::Shape& inShape, size_t outDimNum)
143static ge::graphStatus CheckBroadcastRule(const gert::TilingContext* context, const gert::Shape& inShape,145static ge::graphStatus CheckBroadcastRule(const gert::TilingContext* context, const gert::Shape& inShape,
144 const gert::Shape& outShape)146 const gert::Shape& outShape)
145{147{
146- auto outDimNum = outShape.GetDimNum();148+ size_t dimNum;
147- if (inShape.GetDimNum() != outDimNum) {149+ if (CheckSameDimNum(context, inShape, outShape, dimNum) != ge::GRAPH_SUCCESS) {
148- std::string dimMsg = std::to_string(inShape.GetDimNum()) + " and " + std::to_string(outDimNum);
149- std::string reasonMsg = "The input and output shape dim num should be equal.";
150- OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(
151- context->GetNodeName(), "x and y", dimMsg.c_str(), reasonMsg.c_str());
152 return ge::GRAPH_FAILED;150 return ge::GRAPH_FAILED;
153 }151 }
154 152 
155- for (size_t i = 0; i < outDimNum; i++) {153+ for (size_t i = 0; i < dimNum; i++) {
156 if (inShape[i] != 1 && outShape[i] != inShape[i]) {154 if (inShape[i] != 1 && outShape[i] != inShape[i]) {
157 return ge::GRAPH_FAILED;155 return ge::GRAPH_FAILED;
158 }156 }
@@ -82,6 +82,8 @@ constexpr float coreFactor = 0.75;
82constexpr int64_t LAST_DIM_GATE = 8;82constexpr int64_t LAST_DIM_GATE = 8;
83 83 
84ge::graphStatus GetShapeInfo(const gert::TilingContext* context, gert::Shape& inShape, gert::Shape& outShape);84ge::graphStatus GetShapeInfo(const gert::TilingContext* context, gert::Shape& inShape, gert::Shape& outShape);
85+ge::graphStatus CheckSameDimNum(
86+ const gert::TilingContext* context, const gert::Shape& inShape, const gert::Shape& outShape, size_t& dimNum);
85ge::graphStatus GetABFlag(const gert::TilingContext* context, const gert::Shape& inShape, const gert::Shape& outShape,87ge::graphStatus GetABFlag(const gert::TilingContext* context, const gert::Shape& inShape, const gert::Shape& outShape,
86 std::array<bool, MAX_DIM_NUM>& abInfo);88 std::array<bool, MAX_DIM_NUM>& abInfo);
87void AdjustShapesToSameDimNum(gert::Shape& inShape, size_t outDimNum);89void AdjustShapesToSameDimNum(gert::Shape& inShape, size_t outDimNum);
@@ -101,11 +101,11 @@ uint64_t PadACTiling::GetSizeOfBlockAlign(uint64_t inputSize, uint64_t alignBloc
101 return (inputSize + alignBlockSize - 1) / alignBlockSize * alignBlockSize;101 return (inputSize + alignBlockSize - 1) / alignBlockSize * alignBlockSize;
102}102}
103 103 
104-void PadACTiling::DoFindSplitAxisByInput(bool isBigLastDim)104+void PadACTiling::FindSplitAxisLoop(bool isBigLastDim, uint64_t& dimSizeInUb, uint64_t& dimSizeInLast4Axis)
105{105{
106- OP_LOGD(context_, "Start PadACTiling CalculateTilingKey DoFindSplitAxis.");106+ OP_LOGD(context_, "Start PadACTiling CalculateTilingKey FindSplitAxisLoop.");
107- uint64_t dimSizeInUb = dtypeBytes_;107+ dimSizeInUb = dtypeBytes_;
108- uint64_t dimSizeInLast4Axis = dimSizeInUb;108+ dimSizeInLast4Axis = dimSizeInUb;
109 // 找到切分轴109 // 找到切分轴
110 for (int64_t i = dimNum_ - 1; i >= 0; i--) {110 for (int64_t i = dimNum_ - 1; i >= 0; i--) {
111 if (isBigLastDim && i == static_cast<int64_t>(dimNum_ - 1)) {111 if (isBigLastDim && i == static_cast<int64_t>(dimNum_ - 1)) {
@@ -122,6 +122,13 @@ void PadACTiling::DoFindSplitAxisByInput(bool isBigLastDim)
122 break;122 break;
123 }123 }
124 }124 }
125+}
126+ 
127+void PadACTiling::DoFindSplitAxisByInput(bool isBigLastDim)
128+{
129+ uint64_t dimSizeInUb;
130+ uint64_t dimSizeInLast4Axis;
131+ FindSplitAxisLoop(isBigLastDim, dimSizeInUb, dimSizeInLast4Axis);
125 // 维度超过4,满载后4个轴132 // 维度超过4,满载后4个轴
126 if (dimNum_ - ubAxis_ > PAD_DIM_INDEX_FOURTH) {133 if (dimNum_ - ubAxis_ > PAD_DIM_INDEX_FOURTH) {
127 ubAxis_ = dimNum_ - PAD_DIM_INDEX_FOURTH;134 ubAxis_ = dimNum_ - PAD_DIM_INDEX_FOURTH;
@@ -146,25 +153,9 @@ void PadACTiling::DoFindSplitAxisByInput(bool isBigLastDim)
146 153 
147void PadACTiling::DoFindSplitAxis(bool isBigLastDim)154void PadACTiling::DoFindSplitAxis(bool isBigLastDim)
148{155{
149- OP_LOGD(context_, "Start PadACTiling CalculateTilingKey DoFindSplitAxis.");156+ uint64_t dimSizeInUb;
150- uint64_t dimSizeInUb = dtypeBytes_;157+ uint64_t dimSizeInLast4Axis;
151- uint64_t dimSizeInLast4Axis = dimSizeInUb;158+ FindSplitAxisLoop(isBigLastDim, dimSizeInUb, dimSizeInLast4Axis);
152- // 找到切分轴
153- for (int64_t i = dimNum_ - 1; i >= 0; i--) {
154- if (isBigLastDim && i == static_cast<int64_t>(dimNum_ - 1)) {
155- dimSizeInUb = GetSizeOfBlockAlign(dimSizeInUb * tilingData_->outShape[i], blockSize_);
156- } else {
157- dimSizeInUb *= tilingData_->outShape[i];
158- }
159- // 切分超过4根轴时只切最后4根轴,记录下最后4根轴的大小
160- if (i == dimNum_ - PAD_DIM_INDEX_FOURTH) {
161- dimSizeInLast4Axis = dimSizeInUb;
162- }
163- if (dimSizeInUb >= bufferSize_) {
164- ubAxis_ = i;
165- break;
166- }
167- }
168 159 
169 // 维度超过4,满载后4个轴160 // 维度超过4,满载后4个轴
170 if (dimNum_ - ubAxis_ > PAD_DIM_INDEX_FOURTH) {161 if (dimNum_ - ubAxis_ > PAD_DIM_INDEX_FOURTH) {
@@ -103,6 +103,7 @@ private:
103 uint64_t GetSizeOfBlockAlign(uint64_t inputSize, uint64_t alignBlockSize);103 uint64_t GetSizeOfBlockAlign(uint64_t inputSize, uint64_t alignBlockSize);
104 void DoFindSplitAxis(bool isBigLastDim);104 void DoFindSplitAxis(bool isBigLastDim);
105 void DoFindSplitAxisByInput(bool isBigLastDim);105 void DoFindSplitAxisByInput(bool isBigLastDim);
106+ void FindSplitAxisLoop(bool isBigLastDim, uint64_t& dimSizeInUb, uint64_t& dimSizeInLast4Axis);
106 void CalculateGatherOrScatter();107 void CalculateGatherOrScatter();
107 void CaculateTilingParams();108 void CaculateTilingParams();
108 void CircularOnlyLastTiling(uint64_t lastShapeSizeAlign);109 void CircularOnlyLastTiling(uint64_t lastShapeSizeAlign);
@@ -371,30 +371,47 @@ private:
371 }371 }
372 }372 }
373 373 
374- __aicore__ inline void CopyOutLeftPadBw(374+ __aicore__ inline void CopyOutPadBwInnerLoop(
375- uint16_t ubAxisInCopyNum, const uint64_t* inIndex, LocalTensor<T>& outLocalBw, uint32_t inLeftPadNum,375+ LocalTensor<T>& outLocalBwReal, LocalTensor<T>& outLocalBwTmp,
376- uint32_t inLeftPadStart, OutIndicesSet* totalOutIdx)376+ uint32_t copyStartOffset, uint32_t copyOutNum, uint32_t alignOffset,
377+ DataCopyExtParams& copyOutParams, DataCopyExtParams& outParamAlign,
378+ OutIndicesSet* totalOutIdx, uint64_t baseOffset)
377 {379 {
378- // reflect:380+ for (int32_t o3 = 0; o3 < totalOutIdx[CONST3].count; o3++) {
379- // 左pad: outIdx=leftNum-inIdx381+ uint64_t outAddr = baseOffset + totalOutIdx[CONST3].outIdx[o3] * tdPtr_->outStride[CONST3];
380- // symmetric:382+ if (copyOutNum > 0) {
381- // 左pad: outIdx=leftNum-1-inIdx383+ DataCopyPad(outputGm_[outAddr + alignOffset], outLocalBwReal[copyStartOffset], copyOutParams);
382- 384+ }
383- // ub切分轴,leftpad385+ if (alignOffset > 0) {
384- uint64_t inIdx = inIndex[ubAxis_] + ubAxisInCopyNum - 1 - inLeftPadStart;386+ DataCopyPad(outputGm_[outAddr], outLocalBwTmp, outParamAlign);
385- if constexpr (IS_REFLECT) {387+ }
386- totalOutIdx[ubAxis_].outIdx[0] = tdPtr_->leftPad[ubAxis_] - inIdx;
387- totalOutIdx[ubAxis_].count = 1;
388- } else {
389- totalOutIdx[ubAxis_].outIdx[0] = tdPtr_->leftPad[ubAxis_] - 1 - inIdx;
390- totalOutIdx[ubAxis_].count = 1;
391 }388 }
389+ }
392 390 
393- LocalTensor<T> outLocalBwReal = outLocalBw[BLOCK_NUM * CONST2];391+ __aicore__ inline void CopyOutPadBwLoopImpl(
394- LocalTensor<T> outLocalBwTmp = outLocalBw[0];392+ LocalTensor<T>& outLocalBwReal, LocalTensor<T>& outLocalBwTmp,
393+ uint32_t copyStartOffset, uint32_t copyOutNum, uint32_t alignOffset,
394+ DataCopyExtParams& copyOutParams, DataCopyExtParams& outParamAlign,
395+ OutIndicesSet* totalOutIdx)
396+ {
397+ for (int32_t o0 = 0; o0 < totalOutIdx[0].count; o0++) {
398+ uint64_t o0Offset = totalOutIdx[0].outIdx[o0] * tdPtr_->outStride[0];
399+ for (int32_t o1 = 0; o1 < totalOutIdx[1].count; o1++) {
400+ uint64_t o1Offset = totalOutIdx[1].outIdx[o1] * tdPtr_->outStride[1];
401+ for (int32_t o2 = 0; o2 < totalOutIdx[CONST2].count; o2++) {
402+ uint64_t baseOffset = o0Offset + o1Offset +
403+ totalOutIdx[CONST2].outIdx[o2] * tdPtr_->outStride[CONST2];
404+ CopyOutPadBwInnerLoop(outLocalBwReal, outLocalBwTmp, copyStartOffset,
405+ copyOutNum, alignOffset, copyOutParams, outParamAlign, totalOutIdx, baseOffset);
406+ }
407+ }
408+ }
409+ }
395 410 
396- uint32_t copyOutNum = inLeftPadNum * tdPtr_->outStride[ubAxis_];411+ __aicore__ inline void CopyOutPadBwImpl(
397- uint32_t copyStartOffset = inLeftPadStart * tdPtr_->outStride[ubAxis_];412+ LocalTensor<T>& outLocalBwReal, LocalTensor<T>& outLocalBwTmp,
413+ uint32_t copyStartOffset, uint32_t copyOutNum, OutIndicesSet* totalOutIdx)
414+ {
398 uint32_t alignRed = copyStartOffset % BLOCK_NUM;415 uint32_t alignRed = copyStartOffset % BLOCK_NUM;
399 uint32_t alignOffset = 0;416 uint32_t alignOffset = 0;
400 if (alignRed != 0) {417 if (alignRed != 0) {
@@ -418,26 +435,34 @@ private:
418 DataCopyExtParams copyOutParams = {1u, static_cast<uint32_t>(copyOutNum * sizeof(T)), 0, 0, 0};435 DataCopyExtParams copyOutParams = {1u, static_cast<uint32_t>(copyOutNum * sizeof(T)), 0, 0, 0};
419 DataCopyExtParams outParamAlign = {1u, static_cast<uint32_t>(alignOffset * sizeof(T)), 0, 0, 0};436 DataCopyExtParams outParamAlign = {1u, static_cast<uint32_t>(alignOffset * sizeof(T)), 0, 0, 0};
420 437 
421- for (int32_t o0 = 0; o0 < totalOutIdx[0].count; o0++) {438+ CopyOutPadBwLoopImpl(outLocalBwReal, outLocalBwTmp, copyStartOffset, copyOutNum, alignOffset,
422- uint64_t o0Offset = totalOutIdx[0].outIdx[o0] * tdPtr_->outStride[0];439+ copyOutParams, outParamAlign, totalOutIdx);
423- for (int32_t o1 = 0; o1 < totalOutIdx[1].count; o1++) {440+ }
424- uint64_t o1Offset = totalOutIdx[1].outIdx[o1] * tdPtr_->outStride[1];441+ 
425- for (int32_t o2 = 0; o2 < totalOutIdx[CONST2].count; o2++) {442+ __aicore__ inline void CopyOutLeftPadBw(
426- uint64_t o2Offset = totalOutIdx[CONST2].outIdx[o2] * tdPtr_->outStride[CONST2];443+ uint16_t ubAxisInCopyNum, const uint64_t* inIndex, LocalTensor<T>& outLocalBw, uint32_t inLeftPadNum,
427- for (int32_t o3 = 0; o3 < totalOutIdx[CONST3].count; o3++) {444+ uint32_t inLeftPadStart, OutIndicesSet* totalOutIdx)
428- uint64_t o3Offset = totalOutIdx[CONST3].outIdx[o3] * tdPtr_->outStride[CONST3];445+ {
429- uint64_t outAddr = o0Offset + o1Offset + o2Offset + o3Offset;446+ // reflect:
430- if (copyOutNum > 0) {447+ // 左pad: outIdx=leftNum-inIdx
431- DataCopyPad(448+ // symmetric:
432- outputGm_[outAddr + alignOffset], outLocalBwReal[copyStartOffset], copyOutParams);449+ // 左pad: outIdx=leftNum-1-inIdx
433- }450+ 
434- if (alignOffset > 0) {451+ // ub切分轴,leftpad
435- DataCopyPad(outputGm_[outAddr], outLocalBwTmp, outParamAlign);452+ uint64_t inIdx = inIndex[ubAxis_] + ubAxisInCopyNum - 1 - inLeftPadStart;
436- }453+ if constexpr (IS_REFLECT) {
437- }454+ totalOutIdx[ubAxis_].outIdx[0] = tdPtr_->leftPad[ubAxis_] - inIdx;
438- }455+ } else {
439- }456+ totalOutIdx[ubAxis_].outIdx[0] = tdPtr_->leftPad[ubAxis_] - 1 - inIdx;
440 }457 }
458+ totalOutIdx[ubAxis_].count = 1;
459+ 
460+ LocalTensor<T> outLocalBwReal = outLocalBw[BLOCK_NUM * CONST2];
461+ LocalTensor<T> outLocalBwTmp = outLocalBw[0];
462+ uint32_t copyOutNum = inLeftPadNum * tdPtr_->outStride[ubAxis_];
463+ uint32_t copyStartOffset = inLeftPadStart * tdPtr_->outStride[ubAxis_];
464+ 
465+ CopyOutPadBwImpl(outLocalBwReal, outLocalBwTmp, copyStartOffset, copyOutNum, totalOutIdx);
441 }466 }
442 467 
443 __aicore__ inline void CopyOutRightPadBw(468 __aicore__ inline void CopyOutRightPadBw(
@@ -453,60 +478,17 @@ private:
453 uint64_t inIdx = inIndex[ubAxis_] + ubAxisInCopyNum - 1 - inRightPadStart;478 uint64_t inIdx = inIndex[ubAxis_] + ubAxisInCopyNum - 1 - inRightPadStart;
454 if constexpr (IS_REFLECT) {479 if constexpr (IS_REFLECT) {
455 totalOutIdx[ubAxis_].outIdx[0] = tdPtr_->leftPad[ubAxis_] + CONST2 * (tdPtr_->inShape[ubAxis_] - 1) - inIdx;480 totalOutIdx[ubAxis_].outIdx[0] = tdPtr_->leftPad[ubAxis_] + CONST2 * (tdPtr_->inShape[ubAxis_] - 1) - inIdx;
456- totalOutIdx[ubAxis_].count = 1;
457 } else {481 } else {
458 totalOutIdx[ubAxis_].outIdx[0] = tdPtr_->leftPad[ubAxis_] + CONST2 * tdPtr_->inShape[ubAxis_] - 1 - inIdx;482 totalOutIdx[ubAxis_].outIdx[0] = tdPtr_->leftPad[ubAxis_] + CONST2 * tdPtr_->inShape[ubAxis_] - 1 - inIdx;
459- totalOutIdx[ubAxis_].count = 1;
460 }483 }
484+ totalOutIdx[ubAxis_].count = 1;
461 485 
462 LocalTensor<T> outLocalBwReal = outLocalBw[BLOCK_NUM * CONST2];486 LocalTensor<T> outLocalBwReal = outLocalBw[BLOCK_NUM * CONST2];
463 LocalTensor<T> outLocalBwTmp = outLocalBw[BLOCK_NUM];487 LocalTensor<T> outLocalBwTmp = outLocalBw[BLOCK_NUM];
464- 
465 uint32_t copyOutNum = inRightPadNum * tdPtr_->outStride[ubAxis_];488 uint32_t copyOutNum = inRightPadNum * tdPtr_->outStride[ubAxis_];
466 uint32_t copyStartOffset = inRightPadStart * tdPtr_->outStride[ubAxis_];489 uint32_t copyStartOffset = inRightPadStart * tdPtr_->outStride[ubAxis_];
467- uint32_t alignRed = copyStartOffset % BLOCK_NUM;
468- uint32_t alignOffset = 0;
469- if (alignRed != 0) {
470- __local_mem__ T* inAddrTmp = (__local_mem__ T*)outLocalBwReal.GetPhyAddr() + copyStartOffset;
471- __local_mem__ T* outAddrTmp = (__local_mem__ T*)outLocalBwTmp.GetPhyAddr();
472 490 
473- alignOffset = BLOCK_NUM - alignRed;491+ CopyOutPadBwImpl(outLocalBwReal, outLocalBwTmp, copyStartOffset, copyOutNum, totalOutIdx);
474- copyStartOffset = copyStartOffset + alignOffset;
475- if (copyOutNum > alignOffset) {
476- copyOutNum = (copyOutNum - alignOffset);
477- } else {
478- alignOffset = copyOutNum;
479- copyOutNum = 0;
480- }
481- 
482- CopyTmpUnAlign(inAddrTmp, outAddrTmp, alignOffset);
483- 
484- SetWaitEvent<HardEvent::V_MTE3>(HardEvent::V_MTE3);
485- }
486- 
487- DataCopyExtParams copyOutParams = {1u, static_cast<uint32_t>(copyOutNum * sizeof(T)), 0, 0, 0};
488- DataCopyExtParams outParamAlign = {1u, static_cast<uint32_t>(alignOffset * sizeof(T)), 0, 0, 0};
489- 
490- for (int32_t o0 = 0; o0 < totalOutIdx[0].count; o0++) {
491- uint64_t o0Offset = totalOutIdx[0].outIdx[o0] * tdPtr_->outStride[0];
492- for (int32_t o1 = 0; o1 < totalOutIdx[1].count; o1++) {
493- uint64_t o1Offset = totalOutIdx[1].outIdx[o1] * tdPtr_->outStride[1];
494- for (int32_t o2 = 0; o2 < totalOutIdx[CONST2].count; o2++) {
495- uint64_t o2Offset = totalOutIdx[CONST2].outIdx[o2] * tdPtr_->outStride[CONST2];
496- for (int32_t o3 = 0; o3 < totalOutIdx[CONST3].count; o3++) {
497- uint64_t o3Offset = totalOutIdx[CONST3].outIdx[o3] * tdPtr_->outStride[CONST3];
498- uint64_t outAddr = o0Offset + o1Offset + o2Offset + o3Offset;
499- if (copyOutNum > 0) {
500- DataCopyPad(
501- outputGm_[outAddr + alignOffset], outLocalBwReal[copyStartOffset], copyOutParams);
502- }
503- if (alignOffset > 0) {
504- DataCopyPad(outputGm_[outAddr], outLocalBwTmp, outParamAlign);
505- }
506- }
507- }
508- }
509- }
510 }492 }
511 493 
512 __aicore__ inline void CopyTmpUnAlign(__local_mem__ T* inAddrTmp, __local_mem__ T* outAddrTmp, uint32_t alignOffset)494 __aicore__ inline void CopyTmpUnAlign(__local_mem__ T* inAddrTmp, __local_mem__ T* outAddrTmp, uint32_t alignOffset)
@@ -276,15 +276,8 @@ void PadV3GradACTiling::CalculateTilingKeyMirror()
276 uint64_t lastShapeSizeAlign = GetSizeOfBlockAlign(tilingData_->inShape[dimNum_ - 1], alignNum);276 uint64_t lastShapeSizeAlign = GetSizeOfBlockAlign(tilingData_->inShape[dimNum_ - 1], alignNum);
277 277 
278 bufferSize_ = GetSizeOfBlockAlign(ubSize_ / (CONST2 * dtypeBytes_ + CONST4 * FP32_SIZE) - alignNum, alignNum);278 bufferSize_ = GetSizeOfBlockAlign(ubSize_ / (CONST2 * dtypeBytes_ + CONST4 * FP32_SIZE) - alignNum, alignNum);
279- if (bufferSize_ > UB_MAX_DATA_SIZE_PER_BUFFER / dtypeBytes_) {279+ if (ClampBufferSizeAndCheckBigShape(lastShapeSizeAlign)) {
280- bufferSize_ = UB_MAX_DATA_SIZE_PER_BUFFER / dtypeBytes_;280+ return;
281- }
282- if (lastShapeSizeAlign > bufferSize_) {
283- cutMode_ = TPL_SIMD_BIG;
284- ubAxis_ = dimNum_ - 1;
285- ubFactor_ = bufferSize_;
286- outTileSize_ = bufferSize_;
287- return TilingInfoTune();
288 }281 }
289 // 不切w,但是倒数第二根轴只能切1,此时也走切W分支282 // 不切w,但是倒数第二根轴只能切1,此时也走切W分支
290 // 不切w, 但是只有一根轴 & w > 128B,也走切w分支283 // 不切w, 但是只有一根轴 & w > 128B,也走切w分支
@@ -331,6 +324,23 @@ void PadV3GradACTiling::CalculateTilingKeyMirror()
331 TilingInfoTuneForNormal(lastShapeSizeAlign);324 TilingInfoTuneForNormal(lastShapeSizeAlign);
332 }325 }
333}326}
327+ 
328+bool PadV3GradACTiling::ClampBufferSizeAndCheckBigShape(uint64_t lastShapeSizeAlign)
329+{
330+ if (bufferSize_ > UB_MAX_DATA_SIZE_PER_BUFFER / dtypeBytes_) {
331+ bufferSize_ = UB_MAX_DATA_SIZE_PER_BUFFER / dtypeBytes_;
332+ }
333+ if (lastShapeSizeAlign > bufferSize_) {
334+ cutMode_ = TPL_SIMD_BIG;
335+ ubAxis_ = dimNum_ - 1;
336+ ubFactor_ = bufferSize_;
337+ outTileSize_ = bufferSize_;
338+ TilingInfoTune();
339+ return true;
340+ }
341+ return false;
342+}
343+ 
334void PadV3GradACTiling::CalculateTilingKeyCircular()344void PadV3GradACTiling::CalculateTilingKeyCircular()
335{345{
336 OP_LOGD(context_, "Start PadV3GradACTiling CalculateTilingKeyCircular.");346 OP_LOGD(context_, "Start PadV3GradACTiling CalculateTilingKeyCircular.");
@@ -343,15 +353,8 @@ void PadV3GradACTiling::CalculateTilingKeyCircular()
343 uint64_t lastShapeSizeAlign = GetSizeOfBlockAlign(tilingData_->outShape[dimNum_ - 1], alignNum);353 uint64_t lastShapeSizeAlign = GetSizeOfBlockAlign(tilingData_->outShape[dimNum_ - 1], alignNum);
344 354 
345 bufferSize_ = GetSizeOfBlockAlign(ubSize_ / (CONST2 * dtypeBytes_ + CONST2 * FP32_SIZE) - alignNum, alignNum);355 bufferSize_ = GetSizeOfBlockAlign(ubSize_ / (CONST2 * dtypeBytes_ + CONST2 * FP32_SIZE) - alignNum, alignNum);
346- if (bufferSize_ > UB_MAX_DATA_SIZE_PER_BUFFER / dtypeBytes_) {356+ if (ClampBufferSizeAndCheckBigShape(lastShapeSizeAlign)) {
347- bufferSize_ = UB_MAX_DATA_SIZE_PER_BUFFER / dtypeBytes_;357+ return;
348- }
349- if (lastShapeSizeAlign > bufferSize_) {
350- cutMode_ = TPL_SIMD_BIG;
351- ubAxis_ = dimNum_ - 1;
352- ubFactor_ = bufferSize_;
353- outTileSize_ = bufferSize_;
354- return TilingInfoTune();
355 }358 }
356 DoTilingWithSIMTCircular();359 DoTilingWithSIMTCircular();
357}360}
@@ -94,6 +94,7 @@ private:
94 void TilingInfoTuneForNormal(uint64_t lastShapeSizeAlign);94 void TilingInfoTuneForNormal(uint64_t lastShapeSizeAlign);
95 void CalculateTilingKeyMirror();95 void CalculateTilingKeyMirror();
96 void CalculateTilingKeyCircular();96 void CalculateTilingKeyCircular();
97+ bool ClampBufferSizeAndCheckBigShape(uint64_t lastShapeSizeAlign);
97 void DoTilingWithSIMDMirror();98 void DoTilingWithSIMDMirror();
98 void DoTilingWithSIMDCircular();99 void DoTilingWithSIMDCircular();
99 100 
@@ -259,6 +259,24 @@ private:
259 __local_mem__ GatherRangeType* midIdxAddr, __local_mem__ GatherRangeType* leftIdxAddr,259 __local_mem__ GatherRangeType* midIdxAddr, __local_mem__ GatherRangeType* leftIdxAddr,
260 __local_mem__ GatherRangeType* rightIdxAddr);260 __local_mem__ GatherRangeType* rightIdxAddr);
261 261 
262+ // 由 padRow 计算 H 轴 pad 折叠的 srcRowOffset / dstRowOffset
263+ __aicore__ inline void CalcHAxisPadRowOffset(
264+ int64_t padRow, int64_t padCount, bool isLeftPad, int64_t leftPadH, int64_t outH, int64_t width,
265+ int64_t& srcRowOffset, int64_t& dstRowOffset)
266+ {
267+ int64_t srcRow;
268+ int64_t dstRow;
269+ if (isLeftPad) {
270+ srcRow = padCount - 1 - padRow;
271+ dstRow = IS_REFLECT ? (2 * leftPadH - srcRow) : (2 * leftPadH - 1 - srcRow);
272+ } else {
273+ srcRow = leftPadH + outH + padRow;
274+ dstRow = IS_REFLECT ? (leftPadH + outH - 2 - padRow) : (leftPadH + outH - 1 - padRow);
275+ }
276+ srcRowOffset = srcRow * width;
277+ dstRowOffset = dstRow * width;
278+ }
279+ 
262 /**280 /**
263 * @brief 折叠 H 轴的 pad 数据(GatherAddScatter 方式)281 * @brief 折叠 H 轴的 pad 数据(GatherAddScatter 方式)
264 *282 *
@@ -1791,16 +1809,8 @@ __aicore__ inline void PadV3GradGather<T,modeName>::FoldHAxisPad(
1791 if (nPerVF < 1) nPerVF = 1;1809 if (nPerVF < 1) nPerVF = 1;
1792 1810 
1793 for (int64_t padRow = 0; padRow < padCount; ++padRow) {1811 for (int64_t padRow = 0; padRow < padCount; ++padRow) {
1794- int64_t srcRow, dstRow;1812+ int64_t srcRowOffset, dstRowOffset;
1795- if (isLeftPad) {1813+ CalcHAxisPadRowOffset(padRow, padCount, isLeftPad, leftPadH, outH, width, srcRowOffset, dstRowOffset);
1796- srcRow = padCount - 1 - padRow;
1797- dstRow = IS_REFLECT ? (2 * leftPadH - srcRow) : (2 * leftPadH - 1 - srcRow);
1798- } else {
1799- srcRow = leftPadH + outH + padRow;
1800- dstRow = IS_REFLECT ? (leftPadH + outH - 2 - padRow) : (leftPadH + outH - 1 - padRow);
1801- }
1802- const int64_t srcRowOffset = srcRow * width;
1803- const int64_t dstRowOffset = dstRow * width;
1804 1814 
1805 for (int64_t nBase = 0; nBase < nFactor; nBase += nPerVF) {1815 for (int64_t nBase = 0; nBase < nFactor; nBase += nPerVF) {
1806 int64_t curN = Std::min<int64_t>(nPerVF, nFactor - nBase);1816 int64_t curN = Std::min<int64_t>(nPerVF, nFactor - nBase);
@@ -1892,16 +1902,8 @@ __aicore__ inline void PadV3GradGather<T,modeName>::FoldHAxisPad(
1892 __local_mem__ PromoteDataT* nBaseAddr = baseAddr + n * static_cast<int64_t>(nStride);1902 __local_mem__ PromoteDataT* nBaseAddr = baseAddr + n * static_cast<int64_t>(nStride);
1893 1903 
1894 for (int64_t padRow = 0; padRow < padCount; ++padRow) {1904 for (int64_t padRow = 0; padRow < padCount; ++padRow) {
1895- int64_t srcRow, dstRow;1905+ int64_t srcRowOffset, dstRowOffset;
1896- if (isLeftPad) {1906+ CalcHAxisPadRowOffset(padRow, padCount, isLeftPad, leftPadH, outH, width, srcRowOffset, dstRowOffset);
1897- srcRow = padCount - 1 - padRow;
1898- dstRow = IS_REFLECT ? (2 * leftPadH - srcRow) : (2 * leftPadH - 1 - srcRow);
1899- } else {
1900- srcRow = leftPadH + outH + padRow;
1901- dstRow = IS_REFLECT ? (leftPadH + outH - 2 - padRow) : (leftPadH + outH - 1 - padRow);
1902- }
1903- const int64_t srcRowOffset = srcRow * width;
1904- const int64_t dstRowOffset = dstRow * width;
1905 1907 
1906 int64_t remainC = factor;1908 int64_t remainC = factor;
1907 int64_t cBase = 0;1909 int64_t cBase = 0;
@@ -97,22 +97,29 @@ static bool GetPermutePermAttr(const GNode& permuteNode, std::vector<int64_t>& p
97 return false;97 return false;
98}98}
99 99 
100+// Resolve the TensorDesc of a subgraph input's first matched input.
101+// Fallback: if the input desc's shape is empty, try to get it from the source data node's output desc instead.
102+static TensorDesc ResolveSubgraphInputTensorDesc(const SubgraphInput& subgraphInput)
103+{
104+ auto matchNode = subgraphInput.GetAllInputs().at(0);
105+ TensorDesc tensorDesc;
106+ matchNode.node.GetInputDesc(matchNode.index, tensorDesc);
107+ if (tensorDesc.GetShape().GetDims().empty()) {
108+ auto srcInfo = matchNode.node.GetInDataNodesAndPortIndexs(0);
109+ if (srcInfo.first != nullptr) {
110+ GNode srcNode = *srcInfo.first;
111+ srcNode.GetOutputDesc(srcInfo.second, tensorDesc);
112+ }
113+ }
114+ return tensorDesc;
115+}
116+ 
100static void GetInputsInfo(117static void GetInputsInfo(
101 const std::vector<SubgraphInput>& subgraphInputs, std::vector<Shape>& inputShapes,118 const std::vector<SubgraphInput>& subgraphInputs, std::vector<Shape>& inputShapes,
102 std::vector<DataType>& inputDtypes, std::vector<Format>& inputFormats)119 std::vector<DataType>& inputDtypes, std::vector<Format>& inputFormats)
103{120{
104 for (const auto& subgraphInput : subgraphInputs) {121 for (const auto& subgraphInput : subgraphInputs) {
105- auto matchNode = subgraphInput.GetAllInputs().at(0);122+ TensorDesc tensorDesc = ResolveSubgraphInputTensorDesc(subgraphInput);
106- TensorDesc tensorDesc;
107- matchNode.node.GetInputDesc(matchNode.index, tensorDesc);
108- // Fallback: if shape is empty, try to get from source data node's output desc
109- if (tensorDesc.GetShape().GetDims().empty()) {
110- auto srcInfo = matchNode.node.GetInDataNodesAndPortIndexs(0);
111- if (srcInfo.first != nullptr) {
112- GNode srcNode = *srcInfo.first;
113- srcNode.GetOutputDesc(srcInfo.second, tensorDesc);
114- }
115- }
116 inputShapes.emplace_back(tensorDesc.GetShape());123 inputShapes.emplace_back(tensorDesc.GetShape());
117 inputDtypes.emplace_back(tensorDesc.GetDataType());124 inputDtypes.emplace_back(tensorDesc.GetDataType());
118 inputFormats.emplace_back(tensorDesc.GetFormat());125 inputFormats.emplace_back(tensorDesc.GetFormat());
@@ -124,17 +131,7 @@ static Status InferShape(const GraphUniqPtr& replaceGraph, const std::vector<Sub
124 OP_LOGD(kFusionPassName.c_str(), "Begin infershape for replacement.");131 OP_LOGD(kFusionPassName.c_str(), "Begin infershape for replacement.");
125 std::vector<Shape> inputShapes;132 std::vector<Shape> inputShapes;
126 for (const auto& subgraphInput : subgraphInputs) {133 for (const auto& subgraphInput : subgraphInputs) {
127- auto matchNode = subgraphInput.GetAllInputs().at(0);134+ TensorDesc tensorDesc = ResolveSubgraphInputTensorDesc(subgraphInput);
128- TensorDesc tensorDesc;
129- matchNode.node.GetInputDesc(matchNode.index, tensorDesc);
130- // Fallback: if shape is empty, try to get from source data node's output desc
131- if (tensorDesc.GetShape().GetDims().empty()) {
132- auto srcInfo = matchNode.node.GetInDataNodesAndPortIndexs(0);
133- if (srcInfo.first != nullptr) {
134- GNode srcNode = *srcInfo.first;
135- srcNode.GetOutputDesc(srcInfo.second, tensorDesc);
136- }
137- }
138 inputShapes.emplace_back(tensorDesc.GetShape());135 inputShapes.emplace_back(tensorDesc.GetShape());
139 }136 }
140 return GeUtils::InferShape(*replaceGraph, inputShapes);137 return GeUtils::InferShape(*replaceGraph, inputShapes);
@@ -54,6 +54,27 @@ protected:
54 };54 };
55};55};
56 56 
57+// Compute this core's processing range from multicore tiling params.
58+// Returns false when this core has no work (blockIdx >= realCoreNum).
59+__aicore__ inline bool ParseMultiCoreRange(
60+ int64_t blockIdx, int64_t realCoreNum, int64_t blkFactor, int64_t blkTailFactor, int64_t& blkProcessNum,
61+ int64_t& blkProcessIdxStart, int64_t& blkProcessIdxEnd)
62+{
63+ if (blockIdx >= realCoreNum) {
64+ return false;
65+ }
66+ blkProcessNum = blkFactor;
67+ blkProcessIdxStart = blockIdx * blkFactor;
68+ if (blockIdx < blkTailFactor) {
69+ blkProcessNum += 1;
70+ blkProcessIdxStart += blockIdx;
71+ } else {
72+ blkProcessIdxStart += blkTailFactor;
73+ }
74+ blkProcessIdxEnd = blkProcessIdxStart + blkProcessNum;
75+ return true;
76+}
77+ 
57} // namespace Transpose78} // namespace Transpose
58 79 
59#endif // TRANSPOSE_BASE_H80#endif // TRANSPOSE_BASE_H
@@ -126,18 +126,11 @@ __aicore__ inline void TransposeBigDim<T>::SetLoopInfo(MultiCopyLoopInfo<NDDMA_M
126template <typename T>126template <typename T>
127__aicore__ inline void TransposeBigDim<T>::Process()127__aicore__ inline void TransposeBigDim<T>::Process()
128{128{
129- if (blockIdx_ >= tiling_->realCoreNum) {129+ if (!ParseMultiCoreRange(
130+ blockIdx_, tiling_->realCoreNum, tiling_->blkFactor, tiling_->blkTailFactor, blkProcessNum_,
131+ blkProcessIdxStart_, blkProcessIdxEnd_)) {
130 return;132 return;
131 }133 }
132- blkProcessNum_ = tiling_->blkFactor;
133- blkProcessIdxStart_ = blockIdx_ * tiling_->blkFactor;
134- if (blockIdx_ < tiling_->blkTailFactor) {
135- blkProcessNum_ += 1;
136- blkProcessIdxStart_ += blockIdx_;
137- } else {
138- blkProcessIdxStart_ += tiling_->blkTailFactor;
139- }
140- blkProcessIdxEnd_ = blkProcessIdxStart_ + blkProcessNum_;
141 134 
142 MultiCopyLoopInfo<NDDMA_MAX_DIM_NUM> loopInfo;135 MultiCopyLoopInfo<NDDMA_MAX_DIM_NUM> loopInfo;
143 SetLoopInfo(loopInfo);136 SetLoopInfo(loopInfo);
@@ -99,18 +99,11 @@ __aicore__ inline void TransposeCutOneAxis<T>::Init(
99template <typename T>99template <typename T>
100__aicore__ inline void TransposeCutOneAxis<T>::Process()100__aicore__ inline void TransposeCutOneAxis<T>::Process()
101{101{
102- if (blockIdx_ >= tiling_->realCoreNum) {102+ if (!ParseMultiCoreRange(
103+ blockIdx_, tiling_->realCoreNum, tiling_->blkFactor, tiling_->blkTailFactor, blkProcessNum_,
104+ blkProcessIdxStart_, blkProcessIdxEnd_)) {
103 return;105 return;
104 }106 }
105- blkProcessNum_ = tiling_->blkFactor;
106- blkProcessIdxStart_ = blockIdx_ * tiling_->blkFactor;
107- if (blockIdx_ < tiling_->blkTailFactor) {
108- blkProcessNum_ += 1;
109- blkProcessIdxStart_ += blockIdx_;
110- } else {
111- blkProcessIdxStart_ += tiling_->blkTailFactor;
112- }
113- blkProcessIdxEnd_ = blkProcessIdxStart_ + blkProcessNum_;
114 ProcessPerCore();107 ProcessPerCore();
115}108}
116 109 
@@ -72,6 +72,10 @@ private:
72 };72 };
73 73 
74 __aicore__ inline void ParseTilingData();74 __aicore__ inline void ParseTilingData();
75+ __aicore__ inline int64_t FindPermIndex(int64_t j);
76+ __aicore__ inline void GetLoopAddressOffsetImpl(
77+ bool isSrc, int64_t loopNumArray[], int64_t loopShapeSizeArray[], int64_t cutIndex1, int64_t cutIndex2,
78+ const int64_t expandedShape[], const int64_t inUbShape[]);
75 __aicore__ inline void GetMainLoopAddressOffset(79 __aicore__ inline void GetMainLoopAddressOffset(
76 bool isSrc, int64_t loopNumArray[], int64_t loopShapeSizeArray[], const int64_t expandedShape[],80 bool isSrc, int64_t loopNumArray[], int64_t loopShapeSizeArray[], const int64_t expandedShape[],
77 const int64_t inUbShape[]);81 const int64_t inUbShape[]);
@@ -176,75 +180,19 @@ __aicore__ inline void TransposeCutTwoAxis<T>::ParseTilingData()
176}180}
177 181 
178template <typename T>182template <typename T>
179-__aicore__ inline void TransposeCutTwoAxis<T>::GetMainLoopAddressOffset(183+__aicore__ inline int64_t TransposeCutTwoAxis<T>::FindPermIndex(int64_t j)
180- bool isSrc, int64_t loopNumArray[], int64_t loopShapeSizeArray[], const int64_t expandedShape[],
181- const int64_t inUbShape[])
182{184{
183- int64_t startIndex = NDDMA_MAX_DIM_NUM - 1;185+ for (int64_t permIndex = 0; permIndex < NDDMA_MAX_DIM_NUM; permIndex++) {
184- for (int64_t i = 0; i < NDDMA_MAX_DIM_NUM; i++) {186+ if (tiling_->expandedPerm[permIndex] == j) {
185- for (int64_t j = startIndex; j >= 0; j--) {187+ return permIndex;
186- int64_t idx = j;
187- if (!isSrc) {
188- for (int64_t permIndex = 0; permIndex < NDDMA_MAX_DIM_NUM; permIndex++) {
189- if (tiling_->expandedPerm[permIndex] == j) {
190- idx = permIndex;
191- break;
192- }
193- }
194- }
195- int64_t loopNum = expandedShape[idx] / inUbShape[idx];
196- int64_t loopSize = inUbShape[idx];
197- for (int64_t k = idx + 1; k < NDDMA_MAX_DIM_NUM; k++) {
198- loopSize *= expandedShape[k];
199- }
200- if (loopNum > 1) {
201- startIndex = j - 1;
202- loopNumArray[i] = loopNum;
203- loopShapeSizeArray[i] = loopSize;
204- break;
205- }
206 }188 }
207 }189 }
190+ return j;
208}191}
209 192 
210template <typename T>193template <typename T>
211-__aicore__ inline void TransposeCutTwoAxis<T>::GetTailLoopAddressOffset(194+__aicore__ inline void TransposeCutTwoAxis<T>::GetLoopAddressOffsetImpl(
212- bool isSrc, int64_t loopNumArray[], int64_t loopShapeSizeArray[], int64_t cutIndex, const int64_t expandedShape[],195+ bool isSrc, int64_t loopNumArray[], int64_t loopShapeSizeArray[], int64_t cutIndex1, int64_t cutIndex2,
213- const int64_t inUbShape[])
214-{
215- int64_t startIndex = NDDMA_MAX_DIM_NUM - 1;
216- for (int64_t i = 0; i < NDDMA_MAX_DIM_NUM; i++) {
217- for (int64_t j = startIndex; j >= 0; j--) {
218- int64_t idx = j;
219- if (!isSrc) {
220- for (int64_t permIndex = 0; permIndex < NDDMA_MAX_DIM_NUM; permIndex++) {
221- if (tiling_->expandedPerm[permIndex] == j) {
222- idx = permIndex;
223- break;
224- }
225- }
226- }
227- if (idx == cutIndex) {
228- continue;
229- }
230- int64_t loopNum = expandedShape[idx] / inUbShape[idx];
231- int64_t loopSize = inUbShape[idx];
232- for (int64_t k = idx + 1; k < NDDMA_MAX_DIM_NUM; k++) {
233- loopSize *= expandedShape[k];
234- }
235- if (loopNum > 1) {
236- startIndex = j - 1;
237- loopNumArray[i] = loopNum;
238- loopShapeSizeArray[i] = loopSize;
239- break;
240- }
241- }
242- }
243-}
244- 
245-template <typename T>
246-__aicore__ inline void TransposeCutTwoAxis<T>::GetTailTailLoopAddressOffset(
247- bool isSrc, int64_t loopNumArray[], int64_t loopShapeSizeArray[], int64_t inputCutIndex, int64_t outputCutIndex,
248 const int64_t expandedShape[], const int64_t inUbShape[])196 const int64_t expandedShape[], const int64_t inUbShape[])
249{197{
250 int64_t startIndex = NDDMA_MAX_DIM_NUM - 1;198 int64_t startIndex = NDDMA_MAX_DIM_NUM - 1;
@@ -252,14 +200,9 @@ __aicore__ inline void TransposeCutTwoAxis<T>::GetTailTailLoopAddressOffset(
252 for (int64_t j = startIndex; j >= 0; j--) {200 for (int64_t j = startIndex; j >= 0; j--) {
253 int64_t idx = j;201 int64_t idx = j;
254 if (!isSrc) {202 if (!isSrc) {
255- for (int64_t permIndex = 0; permIndex < NDDMA_MAX_DIM_NUM; permIndex++) {203+ idx = FindPermIndex(j);
256- if (tiling_->expandedPerm[permIndex] == j) {
257- idx = permIndex;
258- break;
259- }
260- }
261 }204 }
262- if (idx == inputCutIndex || idx == outputCutIndex) {205+ if (idx == cutIndex1 || idx == cutIndex2) {
263 continue;206 continue;
264 }207 }
265 int64_t loopNum = expandedShape[idx] / inUbShape[idx];208 int64_t loopNum = expandedShape[idx] / inUbShape[idx];
@@ -277,6 +220,31 @@ __aicore__ inline void TransposeCutTwoAxis<T>::GetTailTailLoopAddressOffset(
277 }220 }
278}221}
279 222 
223+template <typename T>
224+__aicore__ inline void TransposeCutTwoAxis<T>::GetMainLoopAddressOffset(
225+ bool isSrc, int64_t loopNumArray[], int64_t loopShapeSizeArray[], const int64_t expandedShape[],
226+ const int64_t inUbShape[])
227+{
228+ GetLoopAddressOffsetImpl(isSrc, loopNumArray, loopShapeSizeArray, -1, -1, expandedShape, inUbShape);
229+}
230+ 
231+template <typename T>
232+__aicore__ inline void TransposeCutTwoAxis<T>::GetTailLoopAddressOffset(
233+ bool isSrc, int64_t loopNumArray[], int64_t loopShapeSizeArray[], int64_t cutIndex,
234+ const int64_t expandedShape[], const int64_t inUbShape[])
235+{
236+ GetLoopAddressOffsetImpl(isSrc, loopNumArray, loopShapeSizeArray, cutIndex, -1, expandedShape, inUbShape);
237+}
238+ 
239+template <typename T>
240+__aicore__ inline void TransposeCutTwoAxis<T>::GetTailTailLoopAddressOffset(
241+ bool isSrc, int64_t loopNumArray[], int64_t loopShapeSizeArray[], int64_t inputCutIndex, int64_t outputCutIndex,
242+ const int64_t expandedShape[], const int64_t inUbShape[])
243+{
244+ GetLoopAddressOffsetImpl(isSrc, loopNumArray, loopShapeSizeArray, inputCutIndex, outputCutIndex,
245+ expandedShape, inUbShape);
246+}
247+ 
280template <typename T>248template <typename T>
281__aicore__ inline MultiCopyLoopInfo<NDDMA_MAX_DIM_NUM> TransposeCutTwoAxis<T>::SetupLoopInfo(249__aicore__ inline MultiCopyLoopInfo<NDDMA_MAX_DIM_NUM> TransposeCutTwoAxis<T>::SetupLoopInfo(
282 const int64_t inUbSrcShape[], const int64_t inUbDstShape[])250 const int64_t inUbSrcShape[], const int64_t inUbDstShape[])
@@ -598,33 +566,21 @@ __aicore__ inline void TransposeCutTwoAxis<T>::ProcessTail(int64_t loopidxStart,
598 int64_t srcAddressOffset = tailSrcAddressOffsetBase;566 int64_t srcAddressOffset = tailSrcAddressOffsetBase;
599 int64_t dstAddressOffset = tailDstAddressOffsetBase;567 int64_t dstAddressOffset = tailDstAddressOffsetBase;
600 568 
601- int64_t srcAddressOffsetMixedBase[NDDMA_MAX_DIM_NUM] = {0};569+ int64_t tailLoopLocalIdx =
602- DecimalToMixedBase(
603 loopidx -570 loopidx -
604- ((tiling_->expandedInputShape[expandedInputCutIndex_] /571+ ((tiling_->expandedInputShape[expandedInputCutIndex_] / tiling_->inUbMainSrcShape[expandedInputCutIndex_]) *
605- tiling_->inUbMainSrcShape[expandedInputCutIndex_]) *572+ (tiling_->expandedInputShape[inputOutputCutIndex_] / tiling_->inUbMainSrcShape[inputOutputCutIndex_]) *
606- (tiling_->expandedInputShape[inputOutputCutIndex_] / tiling_->inUbMainSrcShape[inputOutputCutIndex_]) *573+ outUbLoop_) -
607- outUbLoop_) -574+ ((tiling_->expandedInputShape[expandedInputCutIndex_] / tiling_->inUbMainSrcShape[expandedInputCutIndex_]) *
608- ((tiling_->expandedInputShape[expandedInputCutIndex_] /575+ outUbLoop_) -
609- tiling_->inUbMainSrcShape[expandedInputCutIndex_]) *576+ ((tiling_->expandedInputShape[inputOutputCutIndex_] / tiling_->inUbMainSrcShape[inputOutputCutIndex_]) *
610- outUbLoop_) -577+ outUbLoop_);
611- ((tiling_->expandedInputShape[inputOutputCutIndex_] / tiling_->inUbMainSrcShape[inputOutputCutIndex_]) *578+ 
612- outUbLoop_),579+ int64_t srcAddressOffsetMixedBase[NDDMA_MAX_DIM_NUM] = {0};
613- tailSrcLoopNumArray, srcAddressOffsetMixedBase);580+ DecimalToMixedBase(tailLoopLocalIdx, tailSrcLoopNumArray, srcAddressOffsetMixedBase);
614 581 
615 int64_t dstAddressOffsetMixedBase[NDDMA_MAX_DIM_NUM] = {0};582 int64_t dstAddressOffsetMixedBase[NDDMA_MAX_DIM_NUM] = {0};
616- DecimalToMixedBase(583+ DecimalToMixedBase(tailLoopLocalIdx, tailDstLoopNumArray, dstAddressOffsetMixedBase);
617- loopidx -
618- ((tiling_->expandedInputShape[expandedInputCutIndex_] /
619- tiling_->inUbMainSrcShape[expandedInputCutIndex_]) *
620- (tiling_->expandedInputShape[inputOutputCutIndex_] / tiling_->inUbMainSrcShape[inputOutputCutIndex_]) *
621- outUbLoop_) -
622- ((tiling_->expandedInputShape[expandedInputCutIndex_] /
623- tiling_->inUbMainSrcShape[expandedInputCutIndex_]) *
624- outUbLoop_) -
625- ((tiling_->expandedInputShape[inputOutputCutIndex_] / tiling_->inUbMainSrcShape[inputOutputCutIndex_]) *
626- outUbLoop_),
627- tailDstLoopNumArray, dstAddressOffsetMixedBase);
628 for (int64_t i = 0; i < NDDMA_MAX_DIM_NUM; i++) {584 for (int64_t i = 0; i < NDDMA_MAX_DIM_NUM; i++) {
629 if (tailSrcLoopShapeSizeArray[i] == 0) {585 if (tailSrcLoopShapeSizeArray[i] == 0) {
630 break;586 break;
@@ -92,18 +92,11 @@ __aicore__ inline void TransposeNLast<T>::Init(
92template <typename T>92template <typename T>
93__aicore__ inline void TransposeNLast<T>::Process()93__aicore__ inline void TransposeNLast<T>::Process()
94{94{
95- if (blockIdx_ >= tiling_->realCoreNum) {95+ if (!ParseMultiCoreRange(
96+ blockIdx_, tiling_->realCoreNum, tiling_->blkFactor, tiling_->blkTailFactor, blkProcessNum_,
97+ blkProcessIdxStart_, blkProcessIdxEnd_)) {
96 return;98 return;
97 }99 }
98- blkProcessNum_ = tiling_->blkFactor;
99- blkProcessIdxStart_ = blockIdx_ * tiling_->blkFactor;
100- if (blockIdx_ < tiling_->blkTailFactor) {
101- blkProcessNum_ += 1;
102- blkProcessIdxStart_ += blockIdx_;
103- } else {
104- blkProcessIdxStart_ += tiling_->blkTailFactor;
105- }
106- blkProcessIdxEnd_ = blkProcessIdxStart_ + blkProcessNum_;
107 ProcessPerCore();100 ProcessPerCore();
108}101}
109 102 
@@ -61,50 +61,11 @@ void AdjustShapesToSameDimNum(gert::Shape& inShape, size_t outDimNum)
61 inShape = newShape;61 inShape = newShape;
62}62}
63 63 
64-ge::graphStatus DeleteOneSizeAxis(const gert::TilingContext* context, gert::Shape& inShape, gert::Shape& outShape)
65-{
66- auto dimNum = inShape.GetDimNum();
67- if (dimNum != outShape.GetDimNum()) {
68- std::string dimMsg = std::to_string(dimNum) + " and " + std::to_string(outShape.GetDimNum());
69- std::string reasonMsg = "The input and output shape dim num should be equal.";
70- OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(
71- context->GetNodeName(), "x and y", dimMsg.c_str(), reasonMsg.c_str());
72- return ge::GRAPH_FAILED;
73- }
74- 
75- if (dimNum == 1) {
76- return ge::GRAPH_SUCCESS;
77- }
78- 
79- size_t mIdx = 0;
80- for (size_t oIdx = 0; oIdx < dimNum; oIdx++) {
81- if (outShape[oIdx] != 1) {
82- inShape[mIdx] = inShape[oIdx];
83- outShape[mIdx] = outShape[oIdx];
84- mIdx += size_t(1);
85- }
86- }
87- 
88- if (mIdx == size_t(0)) {
89- inShape[0] = 1;
90- outShape[0] = 1;
91- mIdx += size_t(1);
92- }
93- inShape.SetDimNum(mIdx);
94- outShape.SetDimNum(mIdx);
95- 
96- return ge::GRAPH_SUCCESS;
97-}
98- 
99ge::graphStatus GetABFlag(const gert::TilingContext* context, const gert::Shape& inShape, const gert::Shape& outShape,64ge::graphStatus GetABFlag(const gert::TilingContext* context, const gert::Shape& inShape, const gert::Shape& outShape,
100 std::array<bool, MAX_DIM_NUM>& abInfo)65 std::array<bool, MAX_DIM_NUM>& abInfo)
101{66{
102- auto inDimNum = inShape.GetDimNum();67+ size_t inDimNum;
103- if (inDimNum != outShape.GetDimNum()) {68+ if (brcto::CheckSameDimNum(context, inShape, outShape, inDimNum) != ge::GRAPH_SUCCESS) {
104- std::string dimMsg = std::to_string(inDimNum) + " and " + std::to_string(outShape.GetDimNum());
105- std::string reasonMsg = "The input and output shape dim num should be equal.";
106- OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(
107- context->GetNodeName(), "x and y", dimMsg.c_str(), reasonMsg.c_str());
108 return ge::GRAPH_FAILED;69 return ge::GRAPH_FAILED;
109 }70 }
110 71 
@@ -117,12 +78,8 @@ ge::graphStatus GetABFlag(const gert::TilingContext* context, const gert::Shape&
117 78 
118ge::graphStatus MergeAxis(const gert::TilingContext* context, gert::Shape& inShape, gert::Shape& outShape)79ge::graphStatus MergeAxis(const gert::TilingContext* context, gert::Shape& inShape, gert::Shape& outShape)
119{80{
120- auto dimNum = inShape.GetDimNum();81+ size_t dimNum;
121- if (dimNum != outShape.GetDimNum()) {82+ if (brcto::CheckSameDimNum(context, inShape, outShape, dimNum) != ge::GRAPH_SUCCESS) {
122- std::string dimMsg = std::to_string(dimNum) + " and " + std::to_string(outShape.GetDimNum());
123- std::string reasonMsg = "The input and output shape dim num should be equal.";
124- OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(
125- context->GetNodeName(), "x and y", dimMsg.c_str(), reasonMsg.c_str());
126 return ge::GRAPH_FAILED;83 return ge::GRAPH_FAILED;
127 }84 }
128 85 
@@ -295,7 +252,7 @@ ge::graphStatus GetShapeInfo(const gert::TilingContext* context, gert::Shape& in
295 OP_LOGD(context->GetNodeName(), "%s",252 OP_LOGD(context->GetNodeName(), "%s",
296 ConcatString("input0 and input1 infer output, output shape is ",253 ConcatString("input0 and input1 infer output, output shape is ",
297 Ops::Base::ToString(outShape).c_str()).c_str());254 Ops::Base::ToString(outShape).c_str()).c_str());
298- if (DeleteOneSizeAxis(context, inShape, outShape) != ge::GRAPH_SUCCESS) {255+ if (brcto::DeleteOneSizeAxis(context, inShape, outShape) != ge::GRAPH_SUCCESS) {
299 std::string shapeMsg = "unknown";256 std::string shapeMsg = "unknown";
300 std::string reasonMsg = "Failed to delete one size axes.";257 std::string reasonMsg = "Failed to delete one size axes.";
301 OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(258 OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(