已合并
【9.1.0同步】fix: foreach_norm修复大tensor场景下calcBuf越界问题 #6589
east_yang创建于 6月25日
【9.1.0同步】fix: foreach_norm修复大tensor场景下calcBuf越界问题 #6589
已合并
共 1 个文件变更+35-12
| @@ -255,17 +255,17 @@ public: | |||
| 255 | if (tensorDataCountList[i] == 0) { | 255 | if (tensorDataCountList[i] == 0) { |
| 256 | continue; | 256 | continue; |
| 257 | } | 257 | } |
| 258 | - int64_t cursorStart_5 = 0; | 258 | + int64_t cursorStart = 0; |
| 259 | int64_t cursorEnd = tensorDataCountList[i] - 1; | 259 | int64_t cursorEnd = tensorDataCountList[i] - 1; |
| 260 | int64_t dataCount = 0; | 260 | int64_t dataCount = 0; |
| 261 | if (i == tensorStart) { | 261 | if (i == tensorStart) { |
| 262 | - cursorStart_5 = tensorStartOffset; | 262 | + cursorStart = tensorStartOffset; |
| 263 | } | 263 | } |
| 264 | if (i == tensorEnd) { | 264 | if (i == tensorEnd) { |
| 265 | cursorEnd = tensorEndOffset; | 265 | cursorEnd = tensorEndOffset; |
| 266 | } | 266 | } |
| 267 | - dataCount = cursorEnd - cursorStart_5 + 1; | 267 | + dataCount = cursorEnd - cursorStart + 1; |
| 268 | - inTensorGM.SetGlobalBuffer(GetTensorAddr(i, inTensorPtr) + cursorStart_5); | 268 | + inTensorGM.SetGlobalBuffer(GetTensorAddr(i, inTensorPtr) + cursorStart); |
| 269 | 269 | ||
| 270 | // coreMiddleOffset : describe this core's offset for middle value of tensor | 270 | // coreMiddleOffset : describe this core's offset for middle value of tensor |
| 271 | SingleTensorProcess(dataCount, coreMiddleOffset + i - tensorStart); | 271 | SingleTensorProcess(dataCount, coreMiddleOffset + i - tensorStart); |
| @@ -312,24 +312,47 @@ private: | |||
| 312 | __aicore__ inline void SingleTensorProcess(int64_t dataCount, uint16_t offset) | 312 | __aicore__ inline void SingleTensorProcess(int64_t dataCount, uint16_t offset) |
| 313 | { | 313 | { |
| 314 | // Batch handling and calculation. | 314 | // Batch handling and calculation. |
| 315 | - uint32_t copyTimes = dataCount / maxDataCount; | 315 | + uint64_t copyTimes = dataCount / maxDataCount; |
| 316 | uint32_t datacountRemainder = dataCount % maxDataCount; | 316 | uint32_t datacountRemainder = dataCount % maxDataCount; |
| 317 | 317 | ||
| 318 | if (datacountRemainder > 0) { | 318 | if (datacountRemainder > 0) { |
| 319 | copyTimes++; | 319 | copyTimes++; |
| 320 | } | 320 | } |
| 321 | - LocalTensor<P> tempLocal = calcBuf.Get<P>(CeilA2B(copyTimes, BYTE_BLOCK / sizeof(P)) * BYTE_BLOCK / sizeof(P)); | 321 | + // Keep slot 0 as the running accumulator; cache partial sums from a 32B-aligned offset. |
| 322 | + uint16_t tempLocalCount = byteLen / sizeof(P); | ||
| 323 | + uint16_t partialStartOffset = BYTE_BLOCK / sizeof(P); | ||
| 324 | + uint16_t cachedPartialCountMax = tempLocalCount - partialStartOffset; | ||
| 325 | + uint16_t cachedPartialCount = 0; | ||
| 326 | + bool hasAccumulator = false; | ||
| 327 | + LocalTensor<P> tempLocal = calcBuf.Get<P>(tempLocalCount); | ||
| 322 | uint32_t tempDataCount = maxDataCount; | 328 | uint32_t tempDataCount = maxDataCount; |
| 323 | - for (uint32_t i = 0; i < copyTimes; i++) { | 329 | + for (uint64_t i = 0; i < copyTimes; i++) { |
| 324 | if (i == copyTimes - 1 && datacountRemainder > 0) { | 330 | if (i == copyTimes - 1 && datacountRemainder > 0) { |
| 325 | tempDataCount = datacountRemainder; | 331 | tempDataCount = datacountRemainder; |
| 326 | } | 332 | } |
| 327 | CopyInStage1(i, tempDataCount); | 333 | CopyInStage1(i, tempDataCount); |
| 328 | - SquareAndReduceRound1(i, tempDataCount, tempLocal); | 334 | + if (!hasAccumulator) { |
| 335 | + SquareAndReduceRound1(0, tempDataCount, tempLocal); | ||
| 336 | + hasAccumulator = true; | ||
| 337 | + continue; | ||
| 338 | + } | ||
| 339 | + | ||
| 340 | + cachedPartialCount++; | ||
| 341 | + SquareAndReduceRound1(partialStartOffset + cachedPartialCount - 1, tempDataCount, tempLocal); | ||
| 342 | + if (cachedPartialCount == cachedPartialCountMax || i == copyTimes - 1) { | ||
| 343 | + if (cachedPartialCount > 1) { | ||
| 344 | + PipeBarrier<PIPE_V>(); | ||
| 345 | + ReduceSum<P>( | ||
| 346 | + tempLocal[partialStartOffset], tempLocal[partialStartOffset], tempLocal[partialStartOffset], | ||
| 347 | + cachedPartialCount); | ||
| 348 | + } | ||
| 349 | + PipeBarrier<PIPE_V>(); | ||
| 350 | + Add(tempLocal, tempLocal, tempLocal[partialStartOffset], 1); | ||
| 351 | + PipeBarrier<PIPE_V>(); | ||
| 352 | + cachedPartialCount = 0; | ||
| 353 | + } | ||
| 329 | } | 354 | } |
| 330 | 355 | ||
| 331 | - PipeBarrier<PIPE_V>(); | ||
| 332 | - ReduceSum<P>(tempLocal, tempLocal, tempLocal, copyTimes); | ||
| 333 | PipeBarrier<PIPE_V>(); | 356 | PipeBarrier<PIPE_V>(); |
| 334 | 357 | ||
| 335 | event_t eventIDVToMTE3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); | 358 | event_t eventIDVToMTE3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); |
| @@ -346,7 +369,7 @@ private: | |||
| 346 | } | 369 | } |
| 347 | 370 | ||
| 348 | // CopyIn, Compute and CopyOut | 371 | // CopyIn, Compute and CopyOut |
| 349 | - __aicore__ inline void CopyInStage1(uint16_t index, int64_t dataCount) | 372 | + __aicore__ inline void CopyInStage1(uint64_t index, int64_t dataCount) |
| 350 | { | 373 | { |
| 351 | LocalTensor<T> dataLocal = dataQueue.AllocTensor<T>(); | 374 | LocalTensor<T> dataLocal = dataQueue.AllocTensor<T>(); |
| 352 | 375 | ||
| @@ -485,4 +508,4 @@ private: | |||
| 485 | 508 | ||
| 486 | } // namespace ForeachNorm | 509 | } // namespace ForeachNorm |
| 487 | 510 | ||
| 488 | -#endif // FOREACH_NORM_N_D_H | 511 | +#endif // FOREACH_NORM_N_D_H |