Merged
TopKV2算子代码优化 #6036
cy_hwAdd on Sep 29
TopKV2算子代码优化 #6036
Merged
Total: 1 files changed+11-10
| @@ -264,13 +264,13 @@ uint32_t ComputeTopkMergeIntraCoreExtractChunkSize(uint64_t ubSizePlatForm, ge:: | |||
| 264 | 264 | ||
| 265 | // ==================== NonLastSmallAxis Helpers ==================== | 265 | // ==================== NonLastSmallAxis Helpers ==================== |
| 266 | 266 | ||
| 267 | -uint32_t GetTopkPreferredInnerChunk(ge::DataType dataType, uint32_t index); | 267 | +uint32_t GetTopkPreferredInnerChunk(ge::DataType dtype, uint32_t candidateIdx); |
| 268 | -bool UseTopkNonLastMergeSort(ge::DataType dataType, uint32_t axisLen); | 268 | +bool UseTopkNonLastMergeSort(ge::DataType dtype, uint32_t axisLength); |
| 269 | -ge::DataType GetTopkNonLastSortDtype(ge::DataType dataType, bool useMergeSort); | 269 | +ge::DataType GetTopkNonLastSortDtype(ge::DataType dtype, bool mergeSortOn); |
| 270 | -uint32_t GetTopkNonLastSortDtypeSize(uint32_t dtypeSize, bool useMergeSort, ge::DataType dataType); | 270 | +uint32_t GetTopkNonLastSortDtypeSize(uint32_t elemSize, bool mergeSortOn, ge::DataType dtype); |
| 271 | -bool GetTopkNonLastSortTmpSize(ge::DataType dataType, uint32_t sortCount, bool useMergeSort, bool isDescend, | 271 | +bool GetTopkNonLastSortTmpSize(ge::DataType dtype, uint32_t sortElemCount, bool mergeSortOn, bool descendOn, |
| 272 | - uint32_t& tmpUbSize); | 272 | + uint32_t& tempUbSize); |
| 273 | -void ComputeTopkAxisDimProducts(const gert::Shape& shape, int64_t axis, TopkNonLastSmallAxisTileInfo& info); | 273 | +void ComputeTopkAxisDimProducts(const gert::Shape& shape, int64_t sortAxisIdx, TopkNonLastSmallAxisTileInfo& info); |
| 274 | 274 | ||
| 275 | // ==================== TopK API Buffer Calculation ==================== | 275 | // ==================== TopK API Buffer Calculation ==================== |
| 276 | 276 | ||
| @@ -307,9 +307,10 @@ bool IsBitonicSmallTopkMode(int64_t kValue, int64_t sortPolicy, bool isSort = tr | |||
| 307 | // ==================== NonLastSmallAxis Calculation Helpers ==================== | 307 | // ==================== NonLastSmallAxis Calculation Helpers ==================== |
| 308 | 308 | ||
| 309 | bool SearchTopkNonLastSmallAxisPlan( | 309 | bool SearchTopkNonLastSmallAxisPlan( |
| 310 | - const TopkNonLastSmallAxisTileInfo& info, uint64_t usableUb, | 310 | + const TopkNonLastSmallAxisTileInfo& info, uint64_t availableUb, |
| 311 | - std::function<bool(TopkNonLastSmallAxisTileInfo&, uint32_t, uint64_t&, TopkNonLastSmallAxisCandidate&)> estimateUb, | 311 | + std::function<bool(TopkNonLastSmallAxisTileInfo&, uint32_t, uint64_t&, TopkNonLastSmallAxisCandidate&)> |
| 312 | - TopkNonLastSmallAxisCandidate& best, TopkNonLastSmallAxisTileInfo* selectedInfo = nullptr); | 312 | + estimateUbFn, |
| 313 | + TopkNonLastSmallAxisCandidate& bestCand, TopkNonLastSmallAxisTileInfo* chosenInfo = nullptr); | ||
| 313 | bool ComputeTopkNonLastLayout(const TopkNonLastSmallAxisTileInfo& info, uint32_t kValue, uint32_t innerChunk, | 314 | bool ComputeTopkNonLastLayout(const TopkNonLastSmallAxisTileInfo& info, uint32_t kValue, uint32_t innerChunk, |
| 314 | bool useMergeSort, topkV2DataInfo::NonLastSmallAxisTopkLayout& layout); | 315 | bool useMergeSort, topkV2DataInfo::NonLastSmallAxisTopkLayout& layout); |
| 315 | bool EstimateTopkNonLastSmallAxisUb(TopkNonLastSmallAxisTileInfo& info, uint32_t kValue, uint32_t innerChunk, | 316 | bool EstimateTopkNonLastSmallAxisUb(TopkNonLastSmallAxisTileInfo& info, uint32_t kValue, uint32_t innerChunk, |