已合并
optimize repeat_interleave kernel 102 tailRepeat Copy Logic #2749
陈展熹创建于 3月16日
optimize repeat_interleave kernel 102 tailRepeat Copy Logic #2749
已合并
共 1 个文件变更+5-7
| @@ -220,12 +220,12 @@ __aicore__ inline void RepeatInterleaveImpl<T, U>::CopyXToMatchOut(int64_t start | |||
| 220 | xOutLocal = xOutQueue_.AllocTensor<T>(); | 220 | xOutLocal = xOutQueue_.AllocTensor<T>(); |
| 221 | xOutQueue_.EnQue(xOutLocal); | 221 | xOutQueue_.EnQue(xOutLocal); |
| 222 | } | 222 | } |
| 223 | - if (copyToMatchOutNum_ > tilingData_.ubFactor - tailRepeatTimes * tilingData_.mergedDims[2]) { | ||
| 224 | - CopyMatchOutToY(); | ||
| 225 | - xOutLocal = xOutQueue_.AllocTensor<T>(); | ||
| 226 | - xOutQueue_.EnQue(xOutLocal); | ||
| 227 | - } | ||
| 228 | CopyOneCpToRepeatOut(xInLocal, tailRepeatTimes); | 223 | CopyOneCpToRepeatOut(xInLocal, tailRepeatTimes); |
| 224 | + CopyMatchOutToY(); | ||
| 225 | + if (likely(repeatDimIdx < cpCount - 1)) { | ||
| 226 | + xOutLocal = xOutQueue_.AllocTensor<T>(); | ||
| 227 | + xOutQueue_.EnQue(xOutLocal); | ||
| 228 | + } | ||
| 229 | copyFromXNum_ += tilingData_.mergedDims[2]; | 229 | copyFromXNum_ += tilingData_.mergedDims[2]; |
| 230 | } | 230 | } |
| 231 | xInQueue_.FreeTensor(xInLocal); | 231 | xInQueue_.FreeTensor(xInLocal); |
| @@ -316,12 +316,10 @@ __aicore__ inline void RepeatInterleaveImpl<T, U>::ProcessCpMatchToUb(int64_t st | |||
| 316 | handleStartCpIdx = startCpIdx + loopIdx * mainCpNum; | 316 | handleStartCpIdx = startCpIdx + loopIdx * mainCpNum; |
| 317 | CopyInX(handleStartCpIdx, 0, mainCpNum * tilingData_.mergedDims[2]); | 317 | CopyInX(handleStartCpIdx, 0, mainCpNum * tilingData_.mergedDims[2]); |
| 318 | CopyXToMatchOut(handleStartCpIdx, mainCpNum); | 318 | CopyXToMatchOut(handleStartCpIdx, mainCpNum); |
| 319 | - CopyMatchOutToY(); | ||
| 320 | } | 319 | } |
| 321 | handleStartCpIdx = startCpIdx + (loopSize - 1) * mainCpNum; | 320 | handleStartCpIdx = startCpIdx + (loopSize - 1) * mainCpNum; |
| 322 | CopyInX(handleStartCpIdx, 0, tailCpNum * tilingData_.mergedDims[2]); | 321 | CopyInX(handleStartCpIdx, 0, tailCpNum * tilingData_.mergedDims[2]); |
| 323 | CopyXToMatchOut(handleStartCpIdx, tailCpNum); | 322 | CopyXToMatchOut(handleStartCpIdx, tailCpNum); |
| 324 | - CopyMatchOutToY(); | ||
| 325 | } | 323 | } |
| 326 | 324 | ||
| 327 | template <typename T, typename U> | 325 | template <typename T, typename U> |