已合并
perf: 修复 scatter reduce include_self 宽行场景性能回退（含 fp16/bf16） #10096
BartonFang创建于 9月8日
perf: 修复 scatter reduce include_self 宽行场景性能回退（含 fp16/bf16） #10096
已合并
共 1 个文件变更+24-5
| @@ -75,9 +75,12 @@ public: | |||
| 75 | inputLocal.SetValue(kIndex, updatesLocal.GetValue(k)); | 75 | inputLocal.SetValue(kIndex, updatesLocal.GetValue(k)); |
| 76 | return; | 76 | return; |
| 77 | } | 77 | } |
| 78 | - int hitCount = countLocal.GetValue(kIndex); | 78 | + bool useUpdateOnly = false; |
| 79 | - countLocal.SetValue(kIndex, hitCount + 1); | 79 | + if (NeedHitCount()) { |
| 80 | - bool useUpdateOnly = includeSelf == 0 && hitCount == 0; | 80 | + int hitCount = countLocal.GetValue(kIndex); |
| 81 | + countLocal.SetValue(kIndex, hitCount + 1); | ||
| 82 | + useUpdateOnly = includeSelf == 0 && hitCount == 0; | ||
| 83 | + } | ||
| 81 | if constexpr (IsCastFloatType()) { | 84 | if constexpr (IsCastFloatType()) { |
| 82 | if (IsCastFloat()) { | 85 | if (IsCastFloat()) { |
| 83 | float inputValue = inputTemp.GetValue(kIndex); | 86 | float inputValue = inputTemp.GetValue(kIndex); |
| @@ -400,13 +403,29 @@ private: | |||
| 400 | inQueueUpdates.FreeTensor(updatesLocal); | 403 | inQueueUpdates.FreeTensor(updatesLocal); |
| 401 | } | 404 | } |
| 402 | 405 | ||
| 403 | - __aicore__ inline bool NeedHitCount() const | 406 | + __aicore__ inline bool IsReduceMode() const |
| 404 | { | 407 | { |
| 405 | return mode >= ScatterElementsV2NS::SCATTER_MODE_REDUCTION_BEGIN && | 408 | return mode >= ScatterElementsV2NS::SCATTER_MODE_REDUCTION_BEGIN && |
| 406 | mode <= ScatterElementsV2NS::SCATTER_MODE_REDUCTION_END; | 409 | mode <= ScatterElementsV2NS::SCATTER_MODE_REDUCTION_END; |
| 407 | } | 410 | } |
| 408 | 411 | ||
| 409 | - __aicore__ inline bool IsCastFloat() const { return IsCastFloatType() && NeedHitCount(); } | 412 | + __aicore__ inline bool NeedHitCount() const |
| 413 | + { | ||
| 414 | + if (!IsReduceMode()) { | ||
| 415 | + return false; | ||
| 416 | + } | ||
| 417 | + // Hit count is required only where the result depends on the number of hits: | ||
| 418 | + // include_self=false (the first hit must replace the loaded self, later hits | ||
| 419 | + // accumulate) and mean (division by hit count). include_self=true add/mul/min/max | ||
| 420 | + // always combine over the already-loaded self, so skip the per-chunk count | ||
| 421 | + // bookkeeping that dominated wide-row performance. | ||
| 422 | + return includeSelf == 0 || mode == ScatterElementsV2NS::SCATTER_MODE_MEAN; | ||
| 423 | + } | ||
| 424 | + | ||
| 425 | + // fp16/bf16 reductions accumulate in fp32 and cast back once (CAST_RINT) for precision. | ||
| 426 | + // This is independent of hit-count bookkeeping: include_self=true add/mul/min/max need | ||
| 427 | + // no counts but still require the fp32 intermediate. | ||
| 428 | + __aicore__ inline bool IsCastFloat() const { return IsCastFloatType() && IsReduceMode(); } | ||
| 410 | 429 | ||
| 411 | __aicore__ static constexpr bool IsCastFloatType() | 430 | __aicore__ static constexpr bool IsCastFloatType() |
| 412 | { | 431 | { |