已合并
perf(sort): 优化归并、非末轴排序及索引输出 #5672
黄晓彬创建于 9月14日
perf(sort): 优化归并、非末轴排序及索引输出 #5672
已合并
共 20 个文件变更+1864-538
| @@ -380,6 +380,21 @@ static const aclTensor* GetTensorWithValueZero(aclTensor* out, aclOpExecutor* ex | |||
| 380 | return viewCopyResult; | 380 | return viewCopyResult; |
| 381 | } | 381 | } |
| 382 | 382 | ||
| 383 | +static bool HasOnlyTrailingSingletons(const aclTensor* self, int64_t dim) | ||
| 384 | +{ | ||
| 385 | + const auto& shape = self->GetViewShape(); | ||
| 386 | + int64_t rank = static_cast<int64_t>(shape.GetDimNum()); | ||
| 387 | + if (rank > DIM_MAX || dim >= rank - 1) { | ||
| 388 | + return false; | ||
| 389 | + } | ||
| 390 | + for (int64_t i = dim + 1; i < rank; ++i) { | ||
| 391 | + if (shape.GetDim(i) != 1) { | ||
| 392 | + return false; | ||
| 393 | + } | ||
| 394 | + } | ||
| 395 | + return true; | ||
| 396 | +} | ||
| 397 | + | ||
| 383 | static aclnnStatus BuildSortGraph(const aclTensor* self, int64_t dimPositive, int64_t dimSize, bool stable, | 398 | static aclnnStatus BuildSortGraph(const aclTensor* self, int64_t dimPositive, int64_t dimSize, bool stable, |
| 384 | bool descending, aclTensor* valuesOut, aclTensor* indicesOut, aclOpExecutor* executor) | 399 | bool descending, aclTensor* valuesOut, aclTensor* indicesOut, aclOpExecutor* executor) |
| 385 | { | 400 | { |
| @@ -397,7 +412,22 @@ static aclnnStatus BuildSortGraph(const aclTensor* self, int64_t dimPositive, in | |||
| 397 | auto selfShapeDetail = GetTensorShape(selfContiguous, executor); | 412 | auto selfShapeDetail = GetTensorShape(selfContiguous, executor); |
| 398 | auto indicesType = indicesOut->GetDataType(); | 413 | auto indicesType = indicesOut->GetDataType(); |
| 399 | std::tuple<const aclTensor*, const aclTensor*> sortRes; | 414 | std::tuple<const aclTensor*, const aclTensor*> sortRes; |
| 400 | - if (UseNoTranspose(selfContiguous, dimPositive)) { | 415 | + if (IsRegBase() && HasOnlyTrailingSingletons(selfContiguous, dimPositive)) { |
| 416 | + // Removing trailing singleton dimensions preserves the contiguous element order. | ||
| 417 | + std::vector<int64_t> compactShape; | ||
| 418 | + for (int64_t i = 0; i <= dimPositive; ++i) { | ||
| 419 | + compactShape.push_back(selfContiguous->GetViewShape().GetDim(i)); | ||
| 420 | + } | ||
| 421 | + auto shapeArray = executor->AllocIntArray(compactShape.data(), compactShape.size()); | ||
| 422 | + CHECK_RET(shapeArray != nullptr && selfShapeDetail != nullptr, ACLNN_ERR_PARAM_NULLPTR); | ||
| 423 | + auto compactSelf = l0op::Reshape(selfContiguous, shapeArray, executor); | ||
| 424 | + CHECK_RET(compactSelf != nullptr, ACLNN_ERR_PARAM_NULLPTR); | ||
| 425 | + auto compactRes = l0op::Sort(compactSelf, -1, descending, stable, indicesType, executor); | ||
| 426 | + CHECK_RET(CheckTupleNullptr(compactRes), ACLNN_ERR_PARAM_NULLPTR); | ||
| 427 | + auto values = l0op::Reshape(std::get<0>(compactRes), selfShapeDetail, executor); | ||
| 428 | + auto indices = l0op::Reshape(std::get<1>(compactRes), selfShapeDetail, executor); | ||
| 429 | + sortRes = std::make_tuple(values, indices); | ||
| 430 | + } else if (UseNoTranspose(selfContiguous, dimPositive)) { | ||
| 401 | sortRes = l0op::Sort(selfContiguous, dimPositive, descending, stable, indicesType, executor); | 431 | sortRes = l0op::Sort(selfContiguous, dimPositive, descending, stable, indicesType, executor); |
| 402 | } else { | 432 | } else { |
| 403 | if (dimSize > DIM_MAX) { | 433 | if (dimSize > DIM_MAX) { |
| @@ -369,7 +369,7 @@ ge::graphStatus SetMergeMoreCoreTiling(gert::TilingContext* context, SortKthTile | |||
| 369 | OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "ComputeMergeMoreCoreTiling", "false", | 369 | OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "ComputeMergeMoreCoreTiling", "false", |
| 370 | "The value of ComputeMergeMoreCoreTiling must be true."), | 370 | "The value of ComputeMergeMoreCoreTiling must be true."), |
| 371 | return ge::GRAPH_FAILED); | 371 | return ge::GRAPH_FAILED); |
| 372 | - OP_LOGI("[mergeSort]", "maxDealingNum: %u", info.keyParams0); | 372 | + OP_LOGI("[mergeSort]", "maxDealingNum: %u, syncMergeBlockSize: %u", info.keyParams0, info.keyParams1); |
| 373 | return ge::GRAPH_SUCCESS; | 373 | return ge::GRAPH_SUCCESS; |
| 374 | } | 374 | } |
| 375 | 375 | ||
| @@ -437,12 +437,31 @@ void PrintTilingDataSort(gert::TilingContext* context, SortKthTileInfo& sortTile | |||
| 437 | return; | 437 | return; |
| 438 | } | 438 | } |
| 439 | 439 | ||
| 440 | -ge::graphStatus SetMergeSortTiling(gert::TilingContext* context, SortKthTileInfo& info) | 440 | +ge::graphStatus SetMergeSortTiling(gert::TilingContext* context, SortKthTileInfo& info, bool useUbCapacity = false) |
| 441 | { | 441 | { |
| 442 | - OP_CHECK_IF(!ComputeMergeSortTiling(context, info, info.y2DtypeSize), | 442 | + // A rejected candidate falls through to another schedule; it is not a tiling error. |
| 443 | - OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "ComputeMergeSortTiling", "false", | 443 | + if (!ComputeMergeSortTiling(context, info, info.y2DtypeSize, useUbCapacity)) { |
| 444 | - "The value of ComputeMergeSortTiling must be true."), | 444 | + return ge::GRAPH_FAILED; |
| 445 | - return ge::GRAPH_FAILED); | 445 | + } |
| 446 | + info.keyParams5 = 0U; | ||
| 447 | + // This is a measured batching policy, not a dtype/UB capability check. It must match | ||
| 448 | + // the FP16/int64-only SortBatchedRows instantiation; other types keep row-local proposals. | ||
| 449 | + if (info.dataType == ge::DT_FLOAT16 && info.y2DtypeSize == sizeof(int64_t) && | ||
| 450 | + info.keyParams0 >= SORT_BATCH_MERGE_MIN_ROWS) { | ||
| 451 | + uint64_t batchElements = static_cast<uint64_t>(info.keyParams0) * info.keyParams3; | ||
| 452 | + constexpr uint32_t proposalBytes = sizeof(float) + sizeof(uint32_t); | ||
| 453 | + uint64_t queueBytesPerElement = static_cast<uint64_t>(info.keyParams4) * | ||
| 454 | + (2U * info.dtypeSize + info.y2DtypeSize); | ||
| 455 | + // Batched merge keeps a B32 index sequence and two 8-byte proposal arrays | ||
| 456 | + // for every row, in addition to the existing double-buffered I/O queues. | ||
| 457 | + constexpr uint32_t maxSortElements = std::numeric_limits<uint8_t>::max() * SORT32_SMALL_AXIS_THRESHOLD; | ||
| 458 | + if (batchElements <= maxSortElements) { | ||
| 459 | + uint64_t requiredUb = batchElements * (queueBytesPerElement + sizeof(uint32_t) + 2U * proposalBytes); | ||
| 460 | + if (requiredUb <= info.ubSize) { | ||
| 461 | + info.keyParams5 = 1U; | ||
| 462 | + } | ||
| 463 | + } | ||
| 464 | + } | ||
| 446 | return ge::GRAPH_SUCCESS; | 465 | return ge::GRAPH_SUCCESS; |
| 447 | } | 466 | } |
| 448 | 467 | ||
| @@ -452,6 +471,23 @@ ge::graphStatus SetMergeIntraCoreTiling(gert::TilingContext* context, SortKthTil | |||
| 452 | OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "ComputeMergeIntraCoreTiling", "false", | 471 | OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "ComputeMergeIntraCoreTiling", "false", |
| 453 | "The value of ComputeMergeIntraCoreTiling must be true."), | 472 | "The value of ComputeMergeIntraCoreTiling must be true."), |
| 454 | return ge::GRAPH_FAILED); | 473 | return ge::GRAPH_FAILED); |
| 474 | + if (info.lastDimTileNum == SORT_RESIDENT_MERGE_MAX_BLOCKS) { | ||
| 475 | + constexpr uint32_t proposalBytes = sizeof(float) + sizeof(uint32_t); | ||
| 476 | + constexpr uint32_t scratchBytes = sizeof(float) + sizeof(uint32_t) + proposalBytes; | ||
| 477 | + constexpr uint32_t residentBytes = 2U * SORT_RESIDENT_MERGE_MIN_BLOCKS * proposalBytes + scratchBytes; | ||
| 478 | + constexpr uint32_t maxSortElements = std::numeric_limits<uint8_t>::max() * SORT32_SMALL_AXIS_THRESHOLD; | ||
| 479 | + uint64_t blockSize = Ops::Base::CeilAlign( | ||
| 480 | + Ops::Base::CeilDiv(static_cast<uint64_t>(info.lastAxis), | ||
| 481 | + static_cast<uint64_t>(SORT_RESIDENT_MERGE_MIN_BLOCKS)), | ||
| 482 | + static_cast<uint64_t>(SORT32_SMALL_AXIS_THRESHOLD)); | ||
| 483 | + if (blockSize > 0 && blockSize <= maxSortElements && blockSize * residentBytes <= info.ubSize) { | ||
| 484 | + // Resident two-list merging permits larger initial blocks; the old workspace remains an upper bound. | ||
| 485 | + info.numTileDataSize = static_cast<uint32_t>(blockSize); | ||
| 486 | + info.lastDimTileNum = SORT_RESIDENT_MERGE_MIN_BLOCKS; | ||
| 487 | + info.keyParams3 = info.numTileDataSize * SORT_RESIDENT_MERGE_MIN_BLOCKS; | ||
| 488 | + info.keyParams5 = static_cast<uint32_t>(INT32_MAX / blockSize); | ||
| 489 | + } | ||
| 490 | + } | ||
| 455 | OP_LOGI("MergeIntraCoreTiling", | 491 | OP_LOGI("MergeIntraCoreTiling", |
| 456 | "B %ld, N %ld, batchPerCore %u, actualCoreNum %u, blockSortSize %u, extractChunkSize %u, " | 492 | "B %ld, N %ld, batchPerCore %u, actualCoreNum %u, blockSortSize %u, extractChunkSize %u, " |
| 457 | "blocksPerRow %u, alignNum %u, ubSize %u", | 493 | "blocksPerRow %u, alignNum %u, ubSize %u", |
| @@ -518,9 +554,9 @@ static bool TryAlignBytesToUint32(uint64_t bytes, uint32_t alignBytes, uint32_t& | |||
| 518 | 554 | ||
| 519 | static bool AddBankRotationPadding(uint32_t blockBytes, uint32_t& rowBytes) | 555 | static bool AddBankRotationPadding(uint32_t blockBytes, uint32_t& rowBytes) |
| 520 | { | 556 | { |
| 521 | - // A merge row stride spanning a multiple of four 32-byte blocks repeatedly addresses the same bank groups. | 557 | + // A merge row stride spanning a multiple of two 32-byte blocks repeatedly addresses the same bank groups. |
| 522 | // One allocated block of padding rotates successive row starts without changing the logical Sort length. | 558 | // One allocated block of padding rotates successive row starts without changing the logical Sort length. |
| 523 | - constexpr uint32_t bankConcentrationStrideBlocks = 4U; | 559 | + constexpr uint32_t bankConcentrationStrideBlocks = 2U; |
| 524 | if (blockBytes == 0U || blockBytes > std::numeric_limits<uint32_t>::max() / bankConcentrationStrideBlocks) { | 560 | if (blockBytes == 0U || blockBytes > std::numeric_limits<uint32_t>::max() / bankConcentrationStrideBlocks) { |
| 525 | return false; | 561 | return false; |
| 526 | } | 562 | } |
| @@ -548,9 +584,7 @@ static bool ComputeSortNonLastUbLayout(const SortKthTileInfo& sortTileInfo, uint | |||
| 548 | uint32_t align = sortTileInfo.blockUbSize; | 584 | uint32_t align = sortTileInfo.blockUbSize; |
| 549 | uint32_t sortDtypeSize = GetNonLastSortDtypeSize(sortTileInfo.dtypeSize, useMergeSort, sortTileInfo.dataType); | 585 | uint32_t sortDtypeSize = GetNonLastSortDtypeSize(sortTileInfo.dtypeSize, useMergeSort, sortTileInfo.dataType); |
| 550 | uint64_t valueAxisRawBytes = static_cast<uint64_t>(sortCount) * sortDtypeSize; | 586 | uint64_t valueAxisRawBytes = static_cast<uint64_t>(sortCount) * sortDtypeSize; |
| 551 | - if (useMergeSort) { | 587 | + if (useMergeSort && sortTileInfo.keyParams5 == 0U) { |
| 552 | - // Advanced merge sort stores value/index pairs internally, so each aligned value row | ||
| 553 | - // must reserve at least sortCount * 8 bytes even for fp16/bf16 inputs. | ||
| 554 | valueAxisRawBytes = std::max(valueAxisRawBytes, static_cast<uint64_t>(sortCount) * SORT_STRUCT_BYTES); | 588 | valueAxisRawBytes = std::max(valueAxisRawBytes, static_cast<uint64_t>(sortCount) * SORT_STRUCT_BYTES); |
| 555 | } | 589 | } |
| 556 | if (!TryAlignBytesToUint32(static_cast<uint64_t>(innerChunk) * sortTileInfo.dtypeSize, align, | 590 | if (!TryAlignBytesToUint32(static_cast<uint64_t>(innerChunk) * sortTileInfo.dtypeSize, align, |
| @@ -607,9 +641,11 @@ static bool EstimateSortNonLastSmallAxisUb(SortKthTileInfo& sortTileInfo, uint32 | |||
| 607 | } | 641 | } |
| 608 | 642 | ||
| 609 | uint64_t sortedIndexBytes = static_cast<uint64_t>(innerChunk) * layout.indexAxisBytes; | 643 | uint64_t sortedIndexBytes = static_cast<uint64_t>(innerChunk) * layout.indexAxisBytes; |
| 644 | + uint64_t proposalBytes = sortTileInfo.keyParams5 != 0U ? static_cast<uint64_t>(sortCount) * SORT_STRUCT_BYTES : | ||
| 645 | + static_cast<uint64_t>(innerChunk) * layout.valueAxisBytes; | ||
| 610 | uint64_t sortPhaseUb = static_cast<uint64_t>(axisLen) * layout.inputRowBytes + | 646 | uint64_t sortPhaseUb = static_cast<uint64_t>(axisLen) * layout.inputRowBytes + |
| 611 | - static_cast<uint64_t>(innerChunk) * layout.valueAxisBytes * 2U + sortedIndexBytes + | 647 | + static_cast<uint64_t>(innerChunk) * layout.valueAxisBytes + proposalBytes + |
| 612 | - bf16CastBytes + static_cast<uint64_t>(sortTileInfo.tmpUbSize); | 648 | + sortedIndexBytes + bf16CastBytes + static_cast<uint64_t>(sortTileInfo.tmpUbSize); |
| 613 | peakUb = sortPhaseUb; | 649 | peakUb = sortPhaseUb; |
| 614 | if (innerChunk > 1U) { | 650 | if (innerChunk > 1U) { |
| 615 | uint64_t outputIndexBytes = static_cast<uint64_t>(axisLen) * layout.outputIndexRowBytes; | 651 | uint64_t outputIndexBytes = static_cast<uint64_t>(axisLen) * layout.outputIndexRowBytes; |
| @@ -671,6 +707,78 @@ bool ApplyNonLastSmallAxisResult(gert::TilingContext* context, SortKthTileInfo& | |||
| 671 | return true; | 707 | return true; |
| 672 | } | 708 | } |
| 673 | 709 | ||
| 710 | +static bool UseCompactNonLastMerge(gert::TilingContext* context, const SortKthTileInfo& sortTileInfo, bool useMergeSort) | ||
| 711 | +{ | ||
| 712 | + // Limit the compact NDDMA path to short floating axes with enough rows to | ||
黄 | |||
| 713 | + // amortize two full waves of column tiles. BF16 retains its cast/transpose path. | ||
| 714 | + // Start at the third Sort32 run; this crossover is in elements, not UB-block bytes. | ||
| 715 | + constexpr uint32_t compactMinAxis = 2U * SORT32_SMALL_AXIS_THRESHOLD + 1U; | ||
| 716 | + constexpr uint32_t compactMaxAxis = 256U; | ||
| 717 | + const uint32_t compactBatchColumns = Ops::Base::GetVRegSize(context) / sizeof(uint32_t); | ||
| 718 | + constexpr uint32_t compactMinWaves = 2U; | ||
| 719 | + return useMergeSort && sortTileInfo.dataType != ge::DT_BF16 && sortTileInfo.lastAxis >= compactMinAxis && | ||
| 720 | + sortTileInfo.lastAxis <= compactMaxAxis && sortTileInfo.innerSize >= compactBatchColumns && | ||
| 721 | + static_cast<uint64_t>(sortTileInfo.unsortedDim) >= | ||
| 722 | + static_cast<uint64_t>(sortTileInfo.maxCoreNum) * compactBatchColumns * compactMinWaves; | ||
| 723 | +} | ||
| 724 | + | ||
| 725 | +template <typename EstimateUb> | ||
| 726 | +static void RefineCompactNonLastPlan(gert::TilingContext* context, const SortKthTileInfo& sortTileInfo, | ||
| 727 | + uint64_t usableUb, const EstimateUb& estimateUb, NonLastSmallAxisCandidate& best, | ||
| 728 | + SortKthTileInfo& selectedInfo) | ||
| 729 | +{ | ||
| 730 | + const uint32_t compactBatchColumns = Ops::Base::GetVRegSize(context) / sizeof(uint32_t); | ||
| 731 | + if (sortTileInfo.keyParams5 != 0U) { | ||
| 732 | + // Compact value rows and the shared proposal buffer permit wider batches. | ||
| 733 | + // Cap at one B32 vector of output indices; retain the original core utilization. | ||
| 734 | + constexpr uint32_t chunkReduction = 2U; | ||
| 735 | + for (uint32_t chunk : {compactBatchColumns, compactBatchColumns / chunkReduction}) { | ||
| 736 | + if (chunk > sortTileInfo.innerSize) { | ||
| 737 | + continue; | ||
| 738 | + } | ||
| 739 | + NonLastSmallAxisCandidate cur; | ||
| 740 | + cur.innerChunk = chunk; | ||
| 741 | + if (!TryGetSortNonLastTileCount(sortTileInfo, chunk, cur.innerLoopNum, cur.tileCount)) { | ||
| 742 | + continue; | ||
| 743 | + } | ||
| 744 | + cur.activeCore = static_cast<uint32_t>(std::min<uint64_t>(sortTileInfo.maxCoreNum, cur.tileCount)); | ||
| 745 | + SortKthTileInfo candidateInfo = sortTileInfo; | ||
| 746 | + if (!estimateUb(candidateInfo, chunk, cur.peakUb, cur) || cur.peakUb > usableUb) { | ||
| 747 | + continue; | ||
| 748 | + } | ||
| 749 | + if (cur.activeCore > best.activeCore || | ||
| 750 | + (cur.activeCore == best.activeCore && cur.innerChunk > best.innerChunk)) { | ||
| 751 | + best = cur; | ||
| 752 | + selectedInfo = candidateInfo; | ||
| 753 | + } | ||
| 754 | + } | ||
| 755 | + } | ||
| 756 | + const uint32_t minBalancedColumns = compactBatchColumns / 2U; | ||
| 757 | + if (sortTileInfo.keyParams5 != 0U && sortTileInfo.outerSize == 1 && best.innerChunk >= minBalancedColumns && | ||
| 758 | + sortTileInfo.maxCoreNum > 0U && best.tileCount > 0U) { | ||
| 759 | + // Keep the same number of tile iterations on the busiest core, but reduce | ||
| 760 | + // each tile to the columns needed to distribute the row across all cores. | ||
| 761 | + uint64_t tilesPerCore = Ops::Base::CeilDiv(best.tileCount, static_cast<uint64_t>(sortTileInfo.maxCoreNum)); | ||
| 762 | + uint64_t columnsPerCore = Ops::Base::CeilDiv(static_cast<uint64_t>(sortTileInfo.innerSize), | ||
| 763 | + static_cast<uint64_t>(sortTileInfo.maxCoreNum)); | ||
| 764 | + uint64_t balancedColumns = Ops::Base::CeilDiv(columnsPerCore, tilesPerCore); | ||
| 765 | + if (balancedColumns >= minBalancedColumns && balancedColumns < best.innerChunk) { | ||
| 766 | + NonLastSmallAxisCandidate cur; | ||
| 767 | + cur.innerChunk = static_cast<uint32_t>(balancedColumns); | ||
| 768 | + SortKthTileInfo candidateInfo = sortTileInfo; | ||
| 769 | + if (TryGetSortNonLastTileCount(sortTileInfo, cur.innerChunk, cur.innerLoopNum, cur.tileCount)) { | ||
| 770 | + cur.activeCore = static_cast<uint32_t>(std::min<uint64_t>(sortTileInfo.maxCoreNum, cur.tileCount)); | ||
| 771 | + if (cur.activeCore == best.activeCore && | ||
| 772 | + Ops::Base::CeilDiv(cur.tileCount, static_cast<uint64_t>(sortTileInfo.maxCoreNum)) <= tilesPerCore && | ||
| 773 | + estimateUb(candidateInfo, cur.innerChunk, cur.peakUb, cur) && cur.peakUb <= usableUb) { | ||
| 774 | + best = cur; | ||
| 775 | + selectedInfo = candidateInfo; | ||
| 776 | + } | ||
| 777 | + } | ||
| 778 | + } | ||
| 779 | + } | ||
| 780 | +} | ||
| 781 | + | ||
| 674 | bool TryNonLastSmallAxis(gert::TilingContext* context, SortKthTileInfo& sortTileInfo, uint64_t& schId) | 782 | bool TryNonLastSmallAxis(gert::TilingContext* context, SortKthTileInfo& sortTileInfo, uint64_t& schId) |
| 675 | { | 783 | { |
| 676 | if (!sortTileInfo.isNonLastAxis) { | 784 | if (!sortTileInfo.isNonLastAxis) { |
| @@ -690,6 +798,8 @@ bool TryNonLastSmallAxis(gert::TilingContext* context, SortKthTileInfo& sortTile | |||
| 690 | 798 | ||
| 691 | uint32_t sortCount = GetNonLastSortCount(sortTileInfo.dataType, static_cast<uint32_t>(sortTileInfo.lastAxis)); | 799 | uint32_t sortCount = GetNonLastSortCount(sortTileInfo.dataType, static_cast<uint32_t>(sortTileInfo.lastAxis)); |
| 692 | bool useMergeSort = UseNonLastMergeSort(sortTileInfo.dataType, static_cast<uint32_t>(sortTileInfo.lastAxis)); | 800 | bool useMergeSort = UseNonLastMergeSort(sortTileInfo.dataType, static_cast<uint32_t>(sortTileInfo.lastAxis)); |
| 801 | + sortTileInfo.keyParams5 = UseCompactNonLastMerge(context, sortTileInfo, useMergeSort); | ||
| 802 | + | ||
| 693 | // Query Sort tmp with the same aligned sortCount and effective dtype that the | 803 | // Query Sort tmp with the same aligned sortCount and effective dtype that the |
| 694 | // kernel will use; BF16 merge sorts as fp32 after the UB cast. | 804 | // kernel will use; BF16 merge sorts as fp32 after the UB cast. |
| 695 | uint32_t tmpUb = 0; | 805 | uint32_t tmpUb = 0; |
| @@ -708,13 +818,87 @@ bool TryNonLastSmallAxis(gert::TilingContext* context, SortKthTileInfo& sortTile | |||
| 708 | OP_LOGI(context->GetNodeName(), "non-last small-axis no-transpose no valid innerChunk"); | 818 | OP_LOGI(context->GetNodeName(), "non-last small-axis no-transpose no valid innerChunk"); |
| 709 | return false; | 819 | return false; |
| 710 | } | 820 | } |
| 821 | + RefineCompactNonLastPlan(context, sortTileInfo, usableUb, estimateUb, best, selectedInfo); | ||
| 711 | sortTileInfo = selectedInfo; | 822 | sortTileInfo = selectedInfo; |
| 712 | return ApplyNonLastSmallAxisResult(context, sortTileInfo, schId, useMergeSort, best, usableUb); | 823 | return ApplyNonLastSmallAxisResult(context, sortTileInfo, schId, useMergeSort, best, usableUb); |
| 713 | } | 824 | } |
| 714 | 825 | ||
| 826 | +static bool PreferSortMergeMoreCore(const SortKthTileInfo& info) | ||
| 827 | +{ | ||
| 828 | + uint32_t dataSize = 0U; | ||
| 829 | + if (!IsMergeMoreCoreProfitable(info.dataType, info.lastAxis, info.unsortedDim, info.maxCoreNum) || | ||
| 830 | + !SelectMergeMoreCoreDataSize(info.lastAxis, info.unsortedDim, info.maxCoreNum, dataSize)) { | ||
| 831 | + return false; | ||
| 832 | + } | ||
| 833 | + // Native blocks have no local serial merge; keep this path when all rows fit. | ||
| 834 | + uint64_t coresPerRow = Ops::Base::CeilDiv(static_cast<uint64_t>(info.lastAxis), static_cast<uint64_t>(dataSize)); | ||
| 835 | + if (static_cast<uint64_t>(info.unsortedDim) <= info.maxCoreNum / coresPerRow) { | ||
| 836 | + return true; | ||
| 837 | + } | ||
| 838 | + // Sync blocks must fit all rows in one round. Prefer radix only when its | ||
| 839 | + // single-round plan gives each row more cores than the sync-merge plan. | ||
| 840 | + uint32_t syncBlockSize = 0U; | ||
| 841 | + if (!SelectMergeSyncMergeBlockSize(info.lastAxis, info.unsortedDim, info.maxCoreNum, syncBlockSize)) { | ||
| 842 | + return false; | ||
| 843 | + } | ||
| 844 | + SortKthTileInfo radix = info; | ||
| 845 | + uint64_t syncCoresPerRow = Ops::Base::CeilDiv(static_cast<uint64_t>(info.lastAxis), | ||
| 846 | + static_cast<uint64_t>(syncBlockSize)); | ||
| 847 | + return !FillRadixMoreCoreInfo(radix) || radix.sortLoopTimes > 1U || radix.lastDimNeedCore <= syncCoresPerRow; | ||
| 848 | +} | ||
| 849 | + | ||
| 850 | +static void TryParallelFinalMerge(gert::TilingContext* context, const SortKthTileInfo& candidate, | ||
| 851 | + bool useDirectSchedule, SortKthTileInfo& sortTileInfo, uint64_t& schId) | ||
| 852 | +{ | ||
| 853 | + // Keep two final runs for Merge Path. Retile nine to twelve native runs | ||
| 854 | + // into eight; their larger native blocks still fit the reserved UB budget. | ||
| 855 | + // Reserve SIMT scratch before recomputing the parallel final-merge plan. | ||
| 856 | + constexpr uint32_t minParallelRows = 4U; | ||
| 857 | + constexpr uint32_t minRetiledRows = 2U; | ||
| 858 | + constexpr uint32_t minInitialRuns = 6U; | ||
| 859 | + constexpr uint32_t maxInitialRuns = MERGE_SORT_LIST_NUM * 2U; | ||
| 860 | + constexpr uint32_t maxRetiledRuns = MERGE_SORT_LIST_NUM * 3U; | ||
| 861 | + constexpr uint32_t mergePathRankBytes = 32U; // two int64 co-ranks, UB-block aligned | ||
| 862 | + constexpr uint32_t mergePathBytesPerElem = MERGE_SORT_LIST_NUM * MERGE_SORT_DATA_BYTES * DOUBLE_BUFFER_NUM; | ||
| 863 | + if (useDirectSchedule && | ||
| 864 | + (candidate.unsortedDim >= minParallelRows || | ||
| 865 | + (candidate.unsortedDim >= minRetiledRows && candidate.lastDimNeedCore > maxInitialRuns)) && | ||
| 866 | + candidate.lastDimNeedCore >= minInitialRuns && candidate.lastDimNeedCore <= maxRetiledRuns && | ||
| 867 | + candidate.ubSize >= SIMT_UB + mergePathRankBytes + mergePathBytesPerElem * MERGE_INTRA_CORE_SORT_ALIGN) { | ||
| 868 | + SortKthTileInfo parallel = candidate; | ||
| 869 | + parallel.ubSize -= SIMT_UB; | ||
| 870 | + if (SetMergeMoreCoreTiling(context, parallel) == ge::GRAPH_SUCCESS) { | ||
| 871 | + if (parallel.lastDimNeedCore > maxInitialRuns) { | ||
| 872 | + // Direct scheduling already fits all rows. Reducing cores per row | ||
| 873 | + // preserves that plan and its row-workspace size. | ||
| 874 | + // Phase 1 is FP32: align32(last run) * (4 + 8 + 4 + 8) bytes. | ||
| 875 | + // At most 12 initial runs of <=4096 elements give axis <=49152. | ||
| 876 | + // After retile, last run = floor(axis/8) + axis%8 <=6151; | ||
| 877 | + // align32 <=6176, so phase 1 uses <=148224 bytes. Together | ||
| 878 | + // with 32768 SIMT bytes this fits the current 248-KiB usable UB. | ||
| 879 | + // ProcessDirect resets the pipe after SyncAll before merge allocation. | ||
| 880 | + // Recheck this bound if the run limits or phase-1 buffers change. | ||
| 881 | + parallel.lastDimNeedCore = maxInitialRuns; | ||
| 882 | + parallel.numTileDataSize = static_cast<uint32_t>(parallel.lastAxis / maxInitialRuns); | ||
| 883 | + parallel.coreNumNeed = static_cast<uint32_t>(parallel.unsortedDim) * maxInitialRuns; | ||
| 884 | + } | ||
| 885 | + // Fused output needs only input/output proposal windows and co-rank scratch. | ||
| 886 | + // Kernel rounds the per-run capacity down to a Sort-aligned element count. | ||
| 887 | + parallel.keyParams0 = (parallel.ubSize - mergePathRankBytes) / mergePathBytesPerElem; | ||
| 888 | + sortTileInfo = parallel; | ||
| 889 | + schId = SORT_SCHID_13; | ||
| 890 | + } | ||
| 891 | + } | ||
| 892 | +} | ||
| 893 | + | ||
| 715 | bool TryMerge(gert::TilingContext* context, SortKthTileInfo& sortTileInfo, uint64_t& schId) | 894 | bool TryMerge(gert::TilingContext* context, SortKthTileInfo& sortTileInfo, uint64_t& schId) |
| 716 | { | 895 | { |
| 717 | - if (IsMergeSortSupported(sortTileInfo.dataType, sortTileInfo.lastAxis)) { | 896 | + // Performance thresholds, independent of kernel buffer capacity. |
| 897 | + constexpr int64_t maxFp16MergeAxis = 2048; | ||
| 898 | + constexpr int64_t maxBf16MergeAxis = 1536; | ||
| 899 | + bool useB16Merge = (sortTileInfo.dataType == ge::DT_FLOAT16 && sortTileInfo.lastAxis <= maxFp16MergeAxis) || | ||
| 900 | + (sortTileInfo.dataType == ge::DT_BF16 && sortTileInfo.lastAxis <= maxBf16MergeAxis); | ||
| 901 | + if (useB16Merge || IsMergeSortSupported(sortTileInfo.dataType, sortTileInfo.lastAxis)) { | ||
| 718 | SortKthTileInfo candidate = sortTileInfo; | 902 | SortKthTileInfo candidate = sortTileInfo; |
| 719 | if (SetMergeSortTiling(context, candidate) == ge::GRAPH_SUCCESS) { | 903 | if (SetMergeSortTiling(context, candidate) == ge::GRAPH_SUCCESS) { |
| 720 | sortTileInfo = candidate; | 904 | sortTileInfo = candidate; |
| @@ -723,16 +907,28 @@ bool TryMerge(gert::TilingContext* context, SortKthTileInfo& sortTileInfo, uint6 | |||
| 723 | } | 907 | } |
| 724 | } | 908 | } |
| 725 | 909 | ||
| 726 | - if (IsMergeMoreCoreSupported(sortTileInfo.dataType, sortTileInfo.lastAxis, sortTileInfo.unsortedDim, | 910 | + if (PreferSortMergeMoreCore(sortTileInfo)) { |
| 727 | - sortTileInfo.maxCoreNum)) { | ||
| 728 | SortKthTileInfo candidate = sortTileInfo; | 911 | SortKthTileInfo candidate = sortTileInfo; |
| 729 | if (SetMergeMoreCoreTiling(context, candidate) == ge::GRAPH_SUCCESS) { | 912 | if (SetMergeMoreCoreTiling(context, candidate) == ge::GRAPH_SUCCESS) { |
| 730 | sortTileInfo = candidate; | 913 | sortTileInfo = candidate; |
| 731 | - schId = SORT_SCHID_3; | 914 | + const bool useDirectSchedule = sortTileInfo.sortLoopTimes == 1U && sortTileInfo.keyParams1 == 0U; |
| 915 | + schId = useDirectSchedule ? SORT_SCHID_12 : SORT_SCHID_3; | ||
| 916 | + TryParallelFinalMerge(context, candidate, useDirectSchedule, sortTileInfo, schId); | ||
| 732 | return true; | 917 | return true; |
| 733 | } | 918 | } |
| 734 | } | 919 | } |
| 735 | 920 | ||
| 921 | + // Preserve the large-FP32 policy and more-core priority, but keep every schId 0 | ||
| 922 | + // candidate here. Fitting in UB is necessary; it does not replace the throughput policy. | ||
| 923 | + if (IsMergeIntraCoreSupported(sortTileInfo.dataType, sortTileInfo.lastAxis, sortTileInfo.unsortedDim, | ||
| 924 | + sortTileInfo.maxCoreNum, sortTileInfo.ubSize)) { | ||
| 925 | + SortKthTileInfo candidate = sortTileInfo; | ||
| 926 | + if (SetMergeSortTiling(context, candidate, true) == ge::GRAPH_SUCCESS) { | ||
| 927 | + sortTileInfo = candidate; | ||
| 928 | + schId = static_cast<uint64_t>(0); | ||
| 929 | + return true; | ||
| 930 | + } | ||
| 931 | + } | ||
| 736 | return false; | 932 | return false; |
| 737 | } | 933 | } |
| 738 | 934 | ||
| @@ -781,8 +977,8 @@ ge::graphStatus SelectSortSchedule(gert::TilingContext* context, SortKthTileInfo | |||
| 781 | "The value of sortAxis must be the last axis or meet no-transpose schedule constraints."); | 977 | "The value of sortAxis must be the last axis or meet no-transpose schedule constraints."); |
| 782 | return ge::GRAPH_FAILED; | 978 | return ge::GRAPH_FAILED; |
| 783 | } | 979 | } |
| 784 | - if (TryMerge(context, sortTileInfo, schId) || TryRadixOneCore(context, sortTileInfo, schId) || | 980 | + if (TryMerge(context, sortTileInfo, schId) || TryMergeIntraCore(context, sortTileInfo, schId) || |
| 785 | - TryMergeIntraCore(context, sortTileInfo, schId)) { | 981 | + TryRadixOneCore(context, sortTileInfo, schId)) { |
| 786 | return ge::GRAPH_SUCCESS; | 982 | return ge::GRAPH_SUCCESS; |
| 787 | } | 983 | } |
| 788 | 984 | ||
| @@ -10,6 +10,7 @@ | |||
| 10 | 10 | ||
| 11 | #include "sort_tiling_common.h" | 11 | #include "sort_tiling_common.h" |
| 12 | 12 | ||
| 13 | +#include <algorithm> | ||
| 13 | #include <cmath> | 14 | #include <cmath> |
| 14 | #include <functional> | 15 | #include <functional> |
| 15 | #include <string> | 16 | #include <string> |
| @@ -267,8 +268,14 @@ bool GetNonLastSortTmpSize(ge::DataType dataType, uint32_t sortCount, bool useMe | |||
| 267 | uint32_t minValue = 0; | 268 | uint32_t minValue = 0; |
| 268 | AscendC::GetSortMaxMinTmpSize(srcShape, GetNonLastSortDtype(dataType, useMergeSort), ge::DT_UINT32, true, config, | 269 | AscendC::GetSortMaxMinTmpSize(srcShape, GetNonLastSortDtype(dataType, useMergeSort), ge::DT_UINT32, true, config, |
| 269 | maxValue, minValue); | 270 | maxValue, minValue); |
| 270 | - tmpUbSize = maxValue; | 271 | + // The unified merge path reuses tmp as the pong proposal buffer. Each proposal is |
| 271 | - return maxValue > 0; | 272 | + // value/index packed into 8 bytes, independent of the value storage dtype. |
| 273 | + uint64_t proposalBytes = static_cast<uint64_t>(sortCount) * MERGE_SORT_DATA_BYTES; | ||
| 274 | + if (proposalBytes > static_cast<uint64_t>(UINT32_MAX)) { | ||
| 275 | + return false; | ||
| 276 | + } | ||
| 277 | + tmpUbSize = useMergeSort ? std::max(maxValue, static_cast<uint32_t>(proposalBytes)) : maxValue; | ||
| 278 | + return tmpUbSize > 0; | ||
| 272 | } | 279 | } |
| 273 | 280 | ||
| 274 | bool SearchNonLastSmallAxisPlan( | 281 | bool SearchNonLastSmallAxisPlan( |
| @@ -409,7 +416,8 @@ static bool AdjustBSharedSingleHTile(int64_t axisLen, int64_t unsortedDim, uint3 | |||
| 409 | if (hCore == 0U) { | 416 | if (hCore == 0U) { |
| 410 | return false; | 417 | return false; |
| 411 | } | 418 | } |
| 412 | - uint64_t hTileData64 = static_cast<uint64_t>(axisLen) / hCore; | 419 | + // Cover the full axis with hCore tiles. Floor division may introduce an extra tail tile after alignment. |
| 420 | + uint64_t hTileData64 = Ops::Base::CeilDiv(static_cast<uint64_t>(axisLen), static_cast<uint64_t>(hCore)); | ||
| 413 | if (hTileData64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { | 421 | if (hTileData64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { |
| 414 | return false; | 422 | return false; |
| 415 | } | 423 | } |
| @@ -426,7 +434,8 @@ static bool AdjustMultiTileHWithBSharing(int64_t axisLen, int64_t unsortedDim, u | |||
| 426 | uint32_t& tileData, uint32_t lastDimTileNum, uint32_t& tmpUbSize, | 434 | uint32_t& tileData, uint32_t lastDimTileNum, uint32_t& tmpUbSize, |
| 427 | bool& adjusted) | 435 | bool& adjusted) |
| 428 | { | 436 | { |
| 429 | - uint64_t newTileData64 = static_cast<uint64_t>(axisLen) / static_cast<uint64_t>(lastDimTileNum); | 437 | + // Keep the adjusted split within lastDimTileNum; floor division can recreate a one-element tail tile. |
| 438 | + uint64_t newTileData64 = Ops::Base::CeilDiv(static_cast<uint64_t>(axisLen), static_cast<uint64_t>(lastDimTileNum)); | ||
| 430 | if (!CeilAlignUint32(newTileData64, BIN_NUM, tileData)) { | 439 | if (!CeilAlignUint32(newTileData64, BIN_NUM, tileData)) { |
| 431 | return false; | 440 | return false; |
| 432 | } | 441 | } |
| @@ -705,14 +714,57 @@ bool ComputeMergeSortPlan(int64_t axisLen, int64_t unsortedDim, uint32_t blockUb | |||
| 705 | return true; | 714 | return true; |
| 706 | } | 715 | } |
| 707 | 716 | ||
| 708 | -bool FillMergeSortInfo(SortKthTileInfo& info, uint32_t indexDtypeSize, uint32_t concatTmpSize) | 717 | +static bool ComputeMergeSortBatchCapacity(const SortKthTileInfo& info, uint32_t indexDtypeSize, uint32_t& tileDataNum) |
| 709 | { | 718 | { |
| 710 | - constexpr uint32_t TILE_DATA_NUM = 4096; | 719 | + if (info.lastAxis <= 0 || info.blockUbSize == 0U || info.dtypeSize == 0U || indexDtypeSize == 0U) { |
| 711 | - MergeSortPlan plan; | ||
| 712 | - if (!ComputeMergeSortPlan(info.lastAxis, info.unsortedDim, info.blockUbSize, TILE_DATA_NUM, info.maxCoreNum, | ||
| 713 | - plan)) { | ||
| 714 | return false; | 720 | return false; |
| 715 | } | 721 | } |
| 722 | + uint64_t alignedAxis = Ops::Base::CeilAlign(static_cast<uint64_t>(info.lastAxis), | ||
| 723 | + static_cast<uint64_t>(info.blockUbSize)); | ||
| 724 | + uint32_t bufferNum = info.lastAxis > ONE_CORE_DATA_SIZE ? 1U : DOUBLE_BUFFER_NUM; | ||
| 725 | + // Both kernels keep one index vector and two 8-byte proposal buffers per row axis. | ||
| 726 | + uint64_t fixedBytes = alignedAxis * (sizeof(uint32_t) + 2U * MERGE_SORT_DATA_BYTES); | ||
| 727 | + uint32_t castBytes = (info.dataType == ge::DT_BF16 || info.dataType == ge::DT_INT16) ? sizeof(float) : 0U; | ||
| 728 | + uint64_t rowBytes = alignedAxis * ((2U * info.dtypeSize + indexDtypeSize) * bufferNum + castBytes); | ||
| 729 | + rowBytes += (info.dtypeSize + sizeof(int64_t)) * bufferNum; | ||
| 730 | + // KthValue additionally aligns two compact output queues independently per batch. | ||
| 731 | + fixedBytes += 2U * bufferNum * (info.blockUbSize - 1U); | ||
| 732 | + if (fixedBytes >= info.ubSize || rowBytes == 0U) { | ||
| 733 | + return false; | ||
| 734 | + } | ||
| 735 | + uint64_t rows = (info.ubSize - fixedBytes) / rowBytes; | ||
| 736 | + // ComputeMergeSortPlan consumes half of tileDataNum as the batch element budget. | ||
| 737 | + uint64_t tileElements = rows * alignedAxis * 2U; | ||
| 738 | + if (rows == 0U || tileElements > std::numeric_limits<uint32_t>::max()) { | ||
| 739 | + return false; | ||
| 740 | + } | ||
| 741 | + tileDataNum = static_cast<uint32_t>(tileElements); | ||
| 742 | + return true; | ||
| 743 | +} | ||
| 744 | + | ||
| 745 | +bool FillMergeSortInfo(SortKthTileInfo& info, uint32_t indexDtypeSize, bool useUbCapacity) | ||
| 746 | +{ | ||
| 747 | + // The fixed budget keeps small Sort batches pipelined; the UB-derived budget also | ||
| 748 | + // rebalances rows across rounds. Switching policies is not only a capacity check. | ||
| 749 | + constexpr uint32_t SORT_TILE_DATA_NUM = 4096; | ||
| 750 | + uint32_t tileDataNum = SORT_TILE_DATA_NUM; | ||
| 751 | + if (useUbCapacity && !ComputeMergeSortBatchCapacity(info, indexDtypeSize, tileDataNum)) { | ||
| 752 | + return false; | ||
| 753 | + } | ||
| 754 | + MergeSortPlan plan; | ||
| 755 | + if (!ComputeMergeSortPlan(info.lastAxis, info.unsortedDim, info.blockUbSize, tileDataNum, info.maxCoreNum, plan)) { | ||
| 756 | + return false; | ||
| 757 | + } | ||
| 758 | + if (useUbCapacity) { | ||
| 759 | + // Balance the per-core row budget instead of leaving an oversized final batch on a few cores. | ||
| 760 | + uint64_t rowsPerCore = Ops::Base::CeilDiv(static_cast<uint64_t>(info.unsortedDim), | ||
| 761 | + static_cast<uint64_t>(plan.coreNumNeed)); | ||
| 762 | + uint64_t rounds = Ops::Base::CeilDiv(rowsPerCore, static_cast<uint64_t>(plan.oneCoreRowNum)); | ||
| 763 | + plan.oneCoreRowNum = static_cast<uint32_t>(Ops::Base::CeilDiv(rowsPerCore, rounds)); | ||
| 764 | + uint64_t rowsPerRound = static_cast<uint64_t>(plan.coreNumNeed) * plan.oneCoreRowNum; | ||
| 765 | + plan.sortLoopTimes = static_cast<uint32_t>( | ||
| 766 | + Ops::Base::CeilDiv(static_cast<uint64_t>(info.unsortedDim), rowsPerRound)); | ||
| 767 | + } | ||
| 716 | info.sortLoopTimes = plan.sortLoopTimes; | 768 | info.sortLoopTimes = plan.sortLoopTimes; |
| 717 | info.lastDimTileNum = 1; | 769 | info.lastDimTileNum = 1; |
| 718 | info.unsortedDimParallel = plan.coreNumNeed; | 770 | info.unsortedDimParallel = plan.coreNumNeed; |
| @@ -724,24 +776,15 @@ bool FillMergeSortInfo(SortKthTileInfo& info, uint32_t indexDtypeSize, uint32_t | |||
| 724 | info.keyParams2 = plan.alignNum * plan.oneCoreRowNum * indexDtypeSize; | 776 | info.keyParams2 = plan.alignNum * plan.oneCoreRowNum * indexDtypeSize; |
| 725 | info.keyParams3 = plan.alignNum; | 777 | info.keyParams3 = plan.alignNum; |
| 726 | info.keyParams4 = info.lastAxis > ONE_CORE_DATA_SIZE ? 1 : DOUBLE_BUFFER_NUM; | 778 | info.keyParams4 = info.lastAxis > ONE_CORE_DATA_SIZE ? 1 : DOUBLE_BUFFER_NUM; |
| 727 | - info.tmpUbSize = std::max(concatTmpSize, info.blockUbSize); | 779 | + info.tmpUbSize = 0; |
| 728 | info.workspaceSize = WORK_SPACE_SIZE; | 780 | info.workspaceSize = WORK_SPACE_SIZE; |
| 729 | return true; | 781 | return true; |
| 730 | } | 782 | } |
| 731 | 783 | ||
| 732 | -bool ComputeMergeSortTiling(gert::TilingContext* context, SortKthTileInfo& info, uint32_t indexDtypeSize) | 784 | +bool ComputeMergeSortTiling(gert::TilingContext* context, SortKthTileInfo& info, uint32_t indexDtypeSize, |
| 785 | + bool useUbCapacity) | ||
| 733 | { | 786 | { |
| 734 | - constexpr uint32_t TILE_DATA_NUM = 4096; | 787 | + if (!FillMergeSortInfo(info, indexDtypeSize, useUbCapacity)) { |
| 735 | - MergeSortPlan plan; | ||
| 736 | - if (!ComputeMergeSortPlan(info.lastAxis, info.unsortedDim, info.blockUbSize, TILE_DATA_NUM, info.maxCoreNum, | ||
| 737 | - plan)) { | ||
| 738 | - return false; | ||
| 739 | - } | ||
| 740 | - auto platformInfo = context->GetPlatformInfo(); | ||
| 741 | - auto plat = platform_ascendc::PlatformAscendC(platformInfo); | ||
| 742 | - uint32_t maxTypeSize = (info.dataType == ge::DT_BF16) ? static_cast<uint32_t>(sizeof(float)) : info.dtypeSize; | ||
| 743 | - uint32_t concatTmpSize = AscendC::GetConcatTmpSize(plat, plan.alignNum, maxTypeSize); | ||
| 744 | - if (!FillMergeSortInfo(info, indexDtypeSize, concatTmpSize)) { | ||
| 745 | return false; | 788 | return false; |
| 746 | } | 789 | } |
| 747 | size_t* userWorkSpaceSize = context->GetWorkspaceSizes(1); | 790 | size_t* userWorkSpaceSize = context->GetWorkspaceSizes(1); |
| @@ -752,51 +795,167 @@ bool ComputeMergeSortTiling(gert::TilingContext* context, SortKthTileInfo& info, | |||
| 752 | // ============================================================================= | 795 | // ============================================================================= |
| 753 | // Merge sort — multi-core (more-core) | 796 | // Merge sort — multi-core (more-core) |
| 754 | // ============================================================================= | 797 | // ============================================================================= |
| 755 | -bool IsMergeMoreCoreSupported(ge::DataType dataType, int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum) | 798 | +// Select per-core data chunk size (2048 or 4096) for merge-more-core. Prefers the smaller base size when direct |
| 799 | +// parallel fits within maxCoreNum; otherwise falls back to the larger size to reduce horizontal core count. | ||
| 800 | +bool SelectMergeMoreCoreDataSize(int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum, uint32_t& dataSize) | ||
| 756 | { | 801 | { |
| 757 | - if (dataType != ge::DT_FLOAT || axisLen <= static_cast<int64_t>(MERGE_SORT_MAX_AXIS_FP32) || | 802 | + if (axisLen <= 0 || unsortedDim <= 0 || maxCoreNum == 0U) { |
| 758 | - axisLen > static_cast<int64_t>(MULTI_CORE_MERGE_SORT_MAX_AXIS)) { | ||
| 759 | return false; | 803 | return false; |
| 760 | } | 804 | } |
| 761 | - uint64_t hCoreNum = (static_cast<uint64_t>(axisLen) + ONE_CORE_DATA_SIZE - 1U) / ONE_CORE_DATA_SIZE; | 805 | + uint32_t baseHCoreNum; |
| 762 | - return hCoreNum > 0 && static_cast<uint64_t>(unsortedDim) * hCoreNum <= maxCoreNum; | 806 | + if (!CeilDivUint32(static_cast<uint64_t>(axisLen), MERGE_MORE_CORE_DATA_SIZE_BASE, baseHCoreNum) || |
| 807 | + baseHCoreNum == 0U || baseHCoreNum > maxCoreNum) { | ||
| 808 | + return false; | ||
| 809 | + } | ||
| 810 | + if (static_cast<uint64_t>(unsortedDim) * baseHCoreNum <= maxCoreNum) { | ||
| 811 | + dataSize = MERGE_MORE_CORE_DATA_SIZE_BASE; | ||
| 812 | + return true; | ||
| 813 | + } | ||
| 814 | + uint32_t largeHCoreNum; | ||
| 815 | + if (!CeilDivUint32(static_cast<uint64_t>(axisLen), MERGE_MORE_CORE_DATA_SIZE_LARGE, largeHCoreNum) || | ||
| 816 | + largeHCoreNum == 0U || largeHCoreNum > maxCoreNum) { | ||
| 817 | + return false; | ||
| 818 | + } | ||
| 819 | + dataSize = MERGE_MORE_CORE_DATA_SIZE_LARGE; | ||
| 820 | + return true; | ||
| 763 | } | 821 | } |
| 764 | 822 | ||
| 765 | -bool ComputeMergeMoreCorePlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t mergeBytesPerElem, | 823 | +// Select a sync-merge block size (aligned to 2048) that groups multiple 4096-element chunks into one block, |
| 766 | - MergeMoreCorePlan& plan) | 824 | +// reducing horizontal core count so that unsortedDim * hCoreNum fits within maxCoreNum. |
| 825 | +bool SelectMergeSyncMergeBlockSize(int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum, uint32_t& blockSize) | ||
| 826 | +{ | ||
| 827 | + if (axisLen <= static_cast<int64_t>(MERGE_MORE_CORE_DATA_SIZE_LARGE) || unsortedDim <= 0 || maxCoreNum == 0U) { | ||
| 828 | + return false; | ||
| 829 | + } | ||
| 830 | + uint32_t maxHCoreByCore = static_cast<uint32_t>(static_cast<uint64_t>(maxCoreNum) / | ||
| 831 | + static_cast<uint64_t>(unsortedDim)); | ||
| 832 | + uint32_t rawChunkNum = 0; | ||
| 833 | + if (!CeilDivUint32(static_cast<uint64_t>(axisLen), MERGE_MORE_CORE_DATA_SIZE_LARGE, rawChunkNum) || | ||
| 834 | + rawChunkNum <= 1U || maxHCoreByCore <= 1U) { | ||
| 835 | + return false; | ||
| 836 | + } | ||
| 837 | + // A sync-merge block is larger than one 4096-element Sort block. Use the available per-row core budget, | ||
| 838 | + // capped below the direct 4096 split, then align to the existing 2048-element more-core granularity. | ||
| 839 | + uint32_t targetHCore = std::min(rawChunkNum - 1U, maxHCoreByCore); | ||
| 840 | + if (targetHCore <= 1U) { | ||
| 841 | + return false; | ||
| 842 | + } | ||
| 843 | + uint32_t candidate = 0; | ||
| 844 | + if (!CeilAlignUint32(Ops::Base::CeilDiv(static_cast<uint64_t>(axisLen), static_cast<uint64_t>(targetHCore)), | ||
| 845 | + MERGE_MORE_CORE_DATA_SIZE_BASE, candidate) || | ||
| 846 | + candidate <= MERGE_MORE_CORE_DATA_SIZE_LARGE) { | ||
| 847 | + return false; | ||
| 848 | + } | ||
| 849 | + uint32_t hCoreNum = 0; | ||
| 850 | + if (!CeilDivUint32(static_cast<uint64_t>(axisLen), candidate, hCoreNum) || hCoreNum <= 1U || | ||
| 851 | + hCoreNum > maxCoreNum || | ||
| 852 | + static_cast<uint64_t>(unsortedDim) * static_cast<uint64_t>(hCoreNum) > static_cast<uint64_t>(maxCoreNum)) { | ||
| 853 | + return false; | ||
| 854 | + } | ||
| 855 | + blockSize = candidate; | ||
| 856 | + return true; | ||
| 857 | +} | ||
| 858 | + | ||
| 859 | +bool IsMergeMoreCoreProfitable(ge::DataType dataType, int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum) | ||
| 860 | +{ | ||
| 861 | + if (dataType != ge::DT_FLOAT || axisLen <= static_cast<int64_t>(MERGE_SORT_MAX_AXIS_FP32) || | ||
| 862 | + axisLen > static_cast<int64_t>(MULTI_CORE_MERGE_SORT_MAX_AXIS) || unsortedDim <= 0 || maxCoreNum == 0U) { | ||
| 863 | + return false; | ||
| 864 | + } | ||
| 865 | + uint32_t dataSize; | ||
| 866 | + if (!SelectMergeMoreCoreDataSize(axisLen, unsortedDim, maxCoreNum, dataSize)) { | ||
| 867 | + return false; | ||
| 868 | + } | ||
| 869 | + uint64_t hCoreNum = (static_cast<uint64_t>(axisLen) + dataSize - 1U) / dataSize; | ||
| 870 | + if (static_cast<uint64_t>(unsortedDim) * hCoreNum <= maxCoreNum) { | ||
| 871 | + return true; | ||
| 872 | + } | ||
| 873 | + uint32_t syncBlockSize = 0; | ||
| 874 | + if (SelectMergeSyncMergeBlockSize(axisLen, unsortedDim, maxCoreNum, syncBlockSize)) { | ||
| 875 | + return true; | ||
| 876 | + } | ||
| 877 | + return hCoreNum >= MERGE_MORE_CORE_MIN_CORES_PER_ROW_FOR_MULTI_ROUND; | ||
| 878 | +} | ||
| 879 | + | ||
| 880 | +namespace { | ||
| 881 | + | ||
| 882 | +bool BuildMergeMoreCoreRoundPlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t maxCoreNum, | ||
| 883 | + uint32_t mergeBytesPerElem, uint32_t hCoreNum, uint32_t numTileDataSize, | ||
| 884 | + uint32_t syncBlockSize, MergeMoreCorePlan& plan) | ||
| 885 | +{ | ||
| 886 | + uint64_t rowsPerRound64 = std::min(static_cast<uint64_t>(maxCoreNum) / hCoreNum, | ||
| 887 | + static_cast<uint64_t>(unsortedDim)); | ||
| 888 | + if (rowsPerRound64 == 0U || rowsPerRound64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { | ||
| 889 | + return false; | ||
| 890 | + } | ||
| 891 | + uint64_t coreNumNeed64 = rowsPerRound64 * hCoreNum; | ||
| 892 | + uint64_t sortLoopTimes64 = Ops::Base::CeilDiv(static_cast<uint64_t>(unsortedDim), rowsPerRound64); | ||
| 893 | + if (static_cast<uint64_t>(axisLen) > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) || | ||
| 894 | + static_cast<uint64_t>(unsortedDim) > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) || | ||
| 895 | + coreNumNeed64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) || | ||
| 896 | + sortLoopTimes64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { | ||
| 897 | + return false; | ||
| 898 | + } | ||
| 899 | + plan.lastDimTileNum = static_cast<uint32_t>(axisLen); | ||
| 900 | + plan.lastDimNeedCore = hCoreNum; | ||
| 901 | + plan.numTileDataSize = numTileDataSize; | ||
| 902 | + plan.unsortedDimParallel = static_cast<uint32_t>(rowsPerRound64); | ||
| 903 | + plan.sortLoopTimes = static_cast<uint32_t>(sortLoopTimes64); | ||
| 904 | + plan.coreNumNeed = static_cast<uint32_t>(coreNumNeed64); | ||
| 905 | + plan.keyParams0 = ubSize / mergeBytesPerElem; | ||
| 906 | + plan.keyParams1 = syncBlockSize; | ||
| 907 | + return true; | ||
| 908 | +} | ||
| 909 | + | ||
| 910 | +bool BuildMergeSyncBlockPlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t maxCoreNum, | ||
| 911 | + uint32_t mergeBytesPerElem, uint32_t syncBlockSize, MergeMoreCorePlan& plan) | ||
| 912 | +{ | ||
| 913 | + uint32_t syncHCoreNum = 0; | ||
| 914 | + if (!CeilDivUint32(static_cast<uint64_t>(axisLen), syncBlockSize, syncHCoreNum) || syncHCoreNum <= 1U || | ||
| 915 | + syncHCoreNum > maxCoreNum) { | ||
| 916 | + return false; | ||
| 917 | + } | ||
| 918 | + return BuildMergeMoreCoreRoundPlan(axisLen, unsortedDim, ubSize, maxCoreNum, mergeBytesPerElem, syncHCoreNum, | ||
| 919 | + MERGE_MORE_CORE_DATA_SIZE_LARGE, syncBlockSize, plan); | ||
| 920 | +} | ||
| 921 | + | ||
| 922 | +} // namespace | ||
| 923 | + | ||
| 924 | +// Try direct, sync-block, then multi-round scheduling in that order. | ||
| 925 | +bool ComputeMergeMoreCorePlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t maxCoreNum, | ||
| 926 | + uint32_t mergeBytesPerElem, MergeMoreCorePlan& plan) | ||
| 767 | { | 927 | { |
| 768 | if (axisLen <= 0 || unsortedDim <= 0 || mergeBytesPerElem == 0U) { | 928 | if (axisLen <= 0 || unsortedDim <= 0 || mergeBytesPerElem == 0U) { |
| 769 | return false; | 929 | return false; |
| 770 | } | 930 | } |
| 771 | - uint32_t hCoreNum; | 931 | + uint32_t dataSize = 0; |
| 772 | - if (!CeilDivUint32(static_cast<uint64_t>(axisLen), ONE_CORE_DATA_SIZE, hCoreNum)) { | 932 | + uint32_t hCoreNum = 0; |
| 933 | + if (!SelectMergeMoreCoreDataSize(axisLen, unsortedDim, maxCoreNum, dataSize) || | ||
| 934 | + !CeilDivUint32(static_cast<uint64_t>(axisLen), dataSize, hCoreNum) || hCoreNum == 0U || hCoreNum > maxCoreNum) { | ||
| 773 | return false; | 935 | return false; |
| 774 | } | 936 | } |
| 775 | - uint64_t coreNumNeed64 = static_cast<uint64_t>(unsortedDim) * hCoreNum; | 937 | + bool directFit = static_cast<uint64_t>(unsortedDim) * hCoreNum <= maxCoreNum; |
| 776 | - if (static_cast<uint64_t>(axisLen) > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) || | 938 | + uint32_t syncBlockSize = 0; |
| 777 | - static_cast<uint64_t>(unsortedDim) > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) || | 939 | + if (!directFit && SelectMergeSyncMergeBlockSize(axisLen, unsortedDim, maxCoreNum, syncBlockSize)) { |
| 778 | - coreNumNeed64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { | 940 | + return BuildMergeSyncBlockPlan(axisLen, unsortedDim, ubSize, maxCoreNum, mergeBytesPerElem, syncBlockSize, |
| 941 | + plan); | ||
| 942 | + } | ||
| 943 | + if (!directFit && hCoreNum < MERGE_MORE_CORE_MIN_CORES_PER_ROW_FOR_MULTI_ROUND) { | ||
| 779 | return false; | 944 | return false; |
| 780 | } | 945 | } |
| 781 | uint64_t numTileDataSize64 = static_cast<uint64_t>(axisLen) / hCoreNum; | 946 | uint64_t numTileDataSize64 = static_cast<uint64_t>(axisLen) / hCoreNum; |
| 782 | if (numTileDataSize64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { | 947 | if (numTileDataSize64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { |
| 783 | return false; | 948 | return false; |
| 784 | } | 949 | } |
| 785 | - plan.lastDimTileNum = static_cast<uint32_t>(static_cast<uint64_t>(axisLen)); | 950 | + return BuildMergeMoreCoreRoundPlan(axisLen, unsortedDim, ubSize, maxCoreNum, mergeBytesPerElem, hCoreNum, |
| 786 | - plan.lastDimNeedCore = hCoreNum; | 951 | + static_cast<uint32_t>(numTileDataSize64), 0U, plan); |
| 787 | - plan.numTileDataSize = static_cast<uint32_t>(numTileDataSize64); | ||
| 788 | - plan.unsortedDimParallel = static_cast<uint32_t>(static_cast<uint64_t>(unsortedDim)); | ||
| 789 | - plan.sortLoopTimes = 1; | ||
| 790 | - plan.coreNumNeed = static_cast<uint32_t>(coreNumNeed64); | ||
| 791 | - // keyParams0 is max elements handled per UB round by the merge more-core pipeline. | ||
| 792 | - plan.keyParams0 = ubSize / mergeBytesPerElem; | ||
| 793 | - return true; | ||
| 794 | } | 952 | } |
| 795 | 953 | ||
| 796 | bool FillMergeMoreCoreInfo(SortKthTileInfo& info, uint32_t mergeBytesPerElem) | 954 | bool FillMergeMoreCoreInfo(SortKthTileInfo& info, uint32_t mergeBytesPerElem) |
| 797 | { | 955 | { |
| 798 | MergeMoreCorePlan plan; | 956 | MergeMoreCorePlan plan; |
| 799 | - if (!ComputeMergeMoreCorePlan(info.lastAxis, info.unsortedDim, info.ubSize, mergeBytesPerElem, plan)) { | 957 | + if (!ComputeMergeMoreCorePlan(info.lastAxis, info.unsortedDim, info.ubSize, info.maxCoreNum, mergeBytesPerElem, |
| 958 | + plan)) { | ||
| 800 | return false; | 959 | return false; |
| 801 | } | 960 | } |
| 802 | info.lastDimTileNum = plan.lastDimTileNum; | 961 | info.lastDimTileNum = plan.lastDimTileNum; |
| @@ -806,6 +965,7 @@ bool FillMergeMoreCoreInfo(SortKthTileInfo& info, uint32_t mergeBytesPerElem) | |||
| 806 | info.sortLoopTimes = plan.sortLoopTimes; | 965 | info.sortLoopTimes = plan.sortLoopTimes; |
| 807 | info.coreNumNeed = plan.coreNumNeed; | 966 | info.coreNumNeed = plan.coreNumNeed; |
| 808 | info.keyParams0 = plan.keyParams0; | 967 | info.keyParams0 = plan.keyParams0; |
| 968 | + info.keyParams1 = plan.keyParams1; | ||
| 809 | uint64_t wsBytes = static_cast<uint64_t>(MERGE_SORT_WORKSPACE_PARAM) * info.lastAxis * info.unsortedDim * | 969 | uint64_t wsBytes = static_cast<uint64_t>(MERGE_SORT_WORKSPACE_PARAM) * info.lastAxis * info.unsortedDim * |
| 810 | sizeof(int32_t); | 970 | sizeof(int32_t); |
| 811 | info.workspaceSize = static_cast<size_t>(wsBytes + WORK_SPACE_SIZE); | 971 | info.workspaceSize = static_cast<size_t>(wsBytes + WORK_SPACE_SIZE); |
| @@ -833,6 +993,19 @@ uint32_t ComputeMergeIntraCoreBlockSortSize(uint32_t ubSize) | |||
| 833 | return (blockSortSize / MERGE_INTRA_CORE_SORT_ALIGN) * MERGE_INTRA_CORE_SORT_ALIGN; | 993 | return (blockSortSize / MERGE_INTRA_CORE_SORT_ALIGN) * MERGE_INTRA_CORE_SORT_ALIGN; |
| 834 | } | 994 | } |
| 835 | 995 | ||
| 996 | +// Max elements a single core can sort in one pass (single-block mode). | ||
| 997 | +// Constrained by the tighter of phase1 (initial sort) and phase3 (extract/merge) UB capacity. | ||
| 998 | +uint32_t ComputeMergeIntraCoreSingleBlockSortSize(uint32_t ubSize) | ||
| 999 | +{ | ||
| 1000 | + constexpr uint32_t MERGE_INTRA_PHASE_BUFFER_NUM = 2; | ||
| 1001 | + constexpr uint32_t PHASE1_BYTES_PER_ELEM = MERGE_INTRA_PHASE_BUFFER_NUM * sizeof(float) + SORT_STRUCT_BYTES + | ||
| 1002 | + MERGE_INTRA_PHASE_BUFFER_NUM * SORT_STRUCT_BYTES + sizeof(uint32_t); | ||
| 1003 | + uint32_t phase1Size = ubSize / PHASE1_BYTES_PER_ELEM; | ||
| 1004 | + uint32_t phase3Size = ComputeMergeIntraCoreExtractChunkSize(ubSize); | ||
| 1005 | + uint32_t singleBlockSize = std::min(phase1Size, phase3Size); | ||
| 1006 | + return (singleBlockSize / MERGE_INTRA_CORE_SORT_ALIGN) * MERGE_INTRA_CORE_SORT_ALIGN; | ||
| 1007 | +} | ||
| 1008 | + | ||
| 836 | uint32_t ComputeMergeIntraCoreExtractChunkSize(uint32_t ubSize) | 1009 | uint32_t ComputeMergeIntraCoreExtractChunkSize(uint32_t ubSize) |
| 837 | { | 1010 | { |
| 838 | constexpr uint32_t PHASE3_BYTES_PER_ELEM = (SORT_STRUCT_BYTES + sizeof(float) + sizeof(int32_t) + sizeof(int64_t)) * | 1011 | constexpr uint32_t PHASE3_BYTES_PER_ELEM = (SORT_STRUCT_BYTES + sizeof(float) + sizeof(int32_t) + sizeof(int64_t)) * |
| @@ -853,11 +1026,13 @@ bool IsMergeIntraCoreSupported(ge::DataType dataType, int64_t axisLen, int64_t u | |||
| 853 | return false; | 1026 | return false; |
| 854 | } | 1027 | } |
| 855 | uint32_t blockSortSize = ComputeMergeIntraCoreBlockSortSize(ubSize); | 1028 | uint32_t blockSortSize = ComputeMergeIntraCoreBlockSortSize(ubSize); |
| 1029 | + uint32_t singleBlockSortSize = ComputeMergeIntraCoreSingleBlockSortSize(ubSize); | ||
| 856 | uint32_t extractChunkSize = ComputeMergeIntraCoreExtractChunkSize(ubSize); | 1030 | uint32_t extractChunkSize = ComputeMergeIntraCoreExtractChunkSize(ubSize); |
| 857 | - if (blockSortSize == 0U || extractChunkSize == 0U) { | 1031 | + if (blockSortSize == 0U || singleBlockSortSize == 0U || extractChunkSize == 0U) { |
| 858 | return false; | 1032 | return false; |
| 859 | } | 1033 | } |
| 860 | - return axisLen <= static_cast<int64_t>(blockSortSize) * MERGE_INTRA_CORE_MAX_BLOCKS; | 1034 | + return axisLen <= static_cast<int64_t>(singleBlockSortSize) || |
| 1035 | + axisLen <= static_cast<int64_t>(blockSortSize) * MERGE_INTRA_CORE_MAX_BLOCKS; | ||
| 861 | } | 1036 | } |
| 862 | 1037 | ||
| 863 | bool ComputeMergeIntraCorePlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t maxCoreNum, | 1038 | bool ComputeMergeIntraCorePlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t maxCoreNum, |
| @@ -876,14 +1051,26 @@ bool ComputeMergeIntraCorePlan(int64_t axisLen, int64_t unsortedDim, uint32_t ub | |||
| 876 | return false; | 1051 | return false; |
| 877 | } | 1052 | } |
| 878 | plan.actualCoreNum = static_cast<uint32_t>(actualCoreNum64); | 1053 | plan.actualCoreNum = static_cast<uint32_t>(actualCoreNum64); |
| 879 | - plan.blockSortSize = ComputeMergeIntraCoreBlockSortSize(ubSize); | ||
| 880 | - if (plan.blockSortSize == 0U) { | ||
| 881 | - return false; | ||
| 882 | - } | ||
| 883 | plan.extractChunkSize = ComputeMergeIntraCoreExtractChunkSize(ubSize); | 1054 | plan.extractChunkSize = ComputeMergeIntraCoreExtractChunkSize(ubSize); |
| 884 | if (plan.extractChunkSize == 0U) { | 1055 | if (plan.extractChunkSize == 0U) { |
| 885 | return false; | 1056 | return false; |
| 886 | } | 1057 | } |
| 1058 | + uint32_t singleBlockSortSize = ComputeMergeIntraCoreSingleBlockSortSize(ubSize); | ||
| 1059 | + if (axisLen <= static_cast<int64_t>(singleBlockSortSize)) { | ||
| 1060 | + uint64_t blockSortSize64 = Ops::Base::CeilAlign(static_cast<uint64_t>(axisLen), | ||
| 1061 | + static_cast<uint64_t>(MERGE_INTRA_CORE_SORT_ALIGN)); | ||
| 1062 | + if (blockSortSize64 == 0U || blockSortSize64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { | ||
| 1063 | + return false; | ||
| 1064 | + } | ||
| 1065 | + plan.blockSortSize = static_cast<uint32_t>(blockSortSize64); | ||
| 1066 | + plan.blocksPerRow = 1; | ||
| 1067 | + plan.alignNum = plan.blockSortSize; | ||
| 1068 | + return true; | ||
| 1069 | + } | ||
| 1070 | + plan.blockSortSize = ComputeMergeIntraCoreBlockSortSize(ubSize); | ||
| 1071 | + if (plan.blockSortSize == 0U) { | ||
| 1072 | + return false; | ||
| 1073 | + } | ||
| 887 | uint64_t blocksPerRow64 = (static_cast<uint64_t>(axisLen) + plan.blockSortSize - 1U) / plan.blockSortSize; | 1074 | uint64_t blocksPerRow64 = (static_cast<uint64_t>(axisLen) + plan.blockSortSize - 1U) / plan.blockSortSize; |
| 888 | if (blocksPerRow64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { | 1075 | if (blocksPerRow64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { |
| 889 | return false; | 1076 | return false; |
| @@ -1257,13 +1444,14 @@ static bool EstimateSmallAxisTwoStageBatching(const SortKthTileInfo& info, uint3 | |||
| 1257 | // ============================================================================= | 1444 | // ============================================================================= |
| 1258 | static bool SelectSmallAxisRouteImpl(const SortKthTileInfo& info, uint32_t batchSizeCap, | 1445 | static bool SelectSmallAxisRouteImpl(const SortKthTileInfo& info, uint32_t batchSizeCap, |
| 1259 | std::function<bool(uint32_t, uint32_t&)> computeBatchNum, | 1446 | std::function<bool(uint32_t, uint32_t&)> computeBatchNum, |
| 1260 | - bool preferFullCoreBatching, SmallAxisRoutePlan& plan) | 1447 | + bool preferFullCoreBatching, SmallAxisRoutePlan& plan, |
| 1448 | + const SmallAxisRule* ruleOverride = nullptr) | ||
| 1261 | { | 1449 | { |
| 1262 | uint32_t axisLen = static_cast<uint32_t>(info.lastAxis); | 1450 | uint32_t axisLen = static_cast<uint32_t>(info.lastAxis); |
| 1263 | if (axisLen <= 1U) { | 1451 | if (axisLen <= 1U) { |
| 1264 | return false; | 1452 | return false; |
| 1265 | } | 1453 | } |
| 1266 | - const SmallAxisRule* rule = FindSmallAxisRule(info.dataType); | 1454 | + const SmallAxisRule* rule = ruleOverride == nullptr ? FindSmallAxisRule(info.dataType) : ruleOverride; |
| 1267 | if (rule == nullptr) { | 1455 | if (rule == nullptr) { |
| 1268 | return false; | 1456 | return false; |
| 1269 | } | 1457 | } |
| @@ -1291,7 +1479,7 @@ static bool SelectSmallAxisRouteImpl(const SortKthTileInfo& info, uint32_t batch | |||
| 1291 | return true; | 1479 | return true; |
| 1292 | } | 1480 | } |
| 1293 | 1481 | ||
| 1294 | -bool SelectSmallAxisRoute(const SortKthTileInfo& info, SmallAxisRoutePlan& plan) | 1482 | +bool SelectSmallAxisRoute(const SortKthTileInfo& info, SmallAxisRoutePlan& plan, const SmallAxisRule* ruleOverride) |
| 1295 | { | 1483 | { |
| 1296 | uint32_t fullCoreSegs = 0; | 1484 | uint32_t fullCoreSegs = 0; |
| 1297 | if (!CeilDivUint32(static_cast<uint64_t>(info.unsortedDim), static_cast<uint64_t>(info.maxCoreNum), fullCoreSegs)) { | 1485 | if (!CeilDivUint32(static_cast<uint64_t>(info.unsortedDim), static_cast<uint64_t>(info.maxCoreNum), fullCoreSegs)) { |
| @@ -1300,7 +1488,7 @@ bool SelectSmallAxisRoute(const SortKthTileInfo& info, SmallAxisRoutePlan& plan) | |||
| 1300 | auto computeBatchNum = [&info](uint32_t batchSize, uint32_t& batchNum) -> bool { | 1488 | auto computeBatchNum = [&info](uint32_t batchSize, uint32_t& batchNum) -> bool { |
| 1301 | return CeilDivUint32(static_cast<uint64_t>(info.unsortedDim), batchSize, batchNum); | 1489 | return CeilDivUint32(static_cast<uint64_t>(info.unsortedDim), batchSize, batchNum); |
| 1302 | }; | 1490 | }; |
| 1303 | - return SelectSmallAxisRouteImpl(info, fullCoreSegs, computeBatchNum, false, plan); | 1491 | + return SelectSmallAxisRouteImpl(info, fullCoreSegs, computeBatchNum, false, plan, ruleOverride); |
| 1304 | } | 1492 | } |
| 1305 | 1493 | ||
| 1306 | static bool SelectNonLastSmallAxisRouteImpl(const SortKthTileInfo& info, SmallAxisRoutePlan& plan, | 1494 | static bool SelectNonLastSmallAxisRouteImpl(const SortKthTileInfo& info, SmallAxisRoutePlan& plan, |
| @@ -41,6 +41,8 @@ constexpr int64_t NON_LAST_SMALL_AXIS_MIN_AXIS_LEN = 2; | |||
| 41 | constexpr int64_t NON_LAST_SMALL_AXIS_THRESHOLD = 2048; | 41 | constexpr int64_t NON_LAST_SMALL_AXIS_THRESHOLD = 2048; |
| 42 | constexpr int64_t ONE_CORE_DATA_SIZE = 2048; | 42 | constexpr int64_t ONE_CORE_DATA_SIZE = 2048; |
| 43 | constexpr uint32_t DOUBLE_BUFFER_NUM = 2; | 43 | constexpr uint32_t DOUBLE_BUFFER_NUM = 2; |
| 44 | +constexpr uint32_t MERGE_MORE_CORE_DATA_SIZE_BASE = 2048; | ||
| 45 | +constexpr uint32_t MERGE_MORE_CORE_DATA_SIZE_LARGE = 4096; | ||
| 44 | constexpr uint32_t SORT32_SMALL_AXIS_THRESHOLD = 32; | 46 | constexpr uint32_t SORT32_SMALL_AXIS_THRESHOLD = 32; |
| 45 | constexpr uint32_t SMALL_AXIS_MAX_DATACOPY_BLOCK_COUNT = 4095; // DataCopy hardware limit for blockCount | 47 | constexpr uint32_t SMALL_AXIS_MAX_DATACOPY_BLOCK_COUNT = 4095; // DataCopy hardware limit for blockCount |
| 46 | constexpr uint32_t SORT_STRUCT_BYTES = 8; // fp32 sort struct size (index + value) | 48 | constexpr uint32_t SORT_STRUCT_BYTES = 8; // fp32 sort struct size (index + value) |
| @@ -216,7 +218,8 @@ bool SearchNonLastSmallAxisPlan( | |||
| 216 | std::function<bool(SortKthTileInfo&, uint32_t, uint64_t&, NonLastSmallAxisCandidate&)> estimateUb, | 218 | std::function<bool(SortKthTileInfo&, uint32_t, uint64_t&, NonLastSmallAxisCandidate&)> estimateUb, |
| 217 | NonLastSmallAxisCandidate& best, SortKthTileInfo* selectedInfo = nullptr); | 219 | NonLastSmallAxisCandidate& best, SortKthTileInfo* selectedInfo = nullptr); |
| 218 | 220 | ||
| 219 | -bool SelectSmallAxisRoute(const SortKthTileInfo& info, SmallAxisRoutePlan& plan); | 221 | +bool SelectSmallAxisRoute(const SortKthTileInfo& info, SmallAxisRoutePlan& plan, |
| 222 | + const SmallAxisRule* ruleOverride = nullptr); | ||
| 220 | bool SelectNonLastSmallAxisRoute(const SortKthTileInfo& info, SmallAxisRoutePlan& plan); | 223 | bool SelectNonLastSmallAxisRoute(const SortKthTileInfo& info, SmallAxisRoutePlan& plan); |
| 221 | bool SelectSortNonLastSmallAxisRoute(const SortKthTileInfo& info, SmallAxisRoutePlan& plan); | 224 | bool SelectSortNonLastSmallAxisRoute(const SortKthTileInfo& info, SmallAxisRoutePlan& plan); |
| 222 | 225 | ||
| @@ -265,6 +268,9 @@ constexpr uint32_t MERGE_SORT_LIST_NUM = 4; | |||
| 265 | constexpr uint32_t MERGE_SORT_DATA_BYTES = 8; | 268 | constexpr uint32_t MERGE_SORT_DATA_BYTES = 8; |
| 266 | constexpr int64_t MERGE_SORT_WORKSPACE_PARAM = 5; | 269 | constexpr int64_t MERGE_SORT_WORKSPACE_PARAM = 5; |
| 267 | constexpr uint32_t MULTI_CORE_MERGE_SORT_MAX_AXIS = 32768; | 270 | constexpr uint32_t MULTI_CORE_MERGE_SORT_MAX_AXIS = 32768; |
| 271 | +// Multi-round merge-more-core executes global barriers for every merge level. Use this route only when the sort axis | ||
| 272 | +// can occupy at least eight cores per row; smaller horizontal splits fall back to a route with less barrier overhead. | ||
| 273 | +constexpr uint32_t MERGE_MORE_CORE_MIN_CORES_PER_ROW_FOR_MULTI_ROUND = 8; | ||
| 268 | constexpr uint32_t MERGE_INTRA_CORE_SORT_ALIGN = 32; // Sort/Extract API alignment requirement (elements) | 274 | constexpr uint32_t MERGE_INTRA_CORE_SORT_ALIGN = 32; // Sort/Extract API alignment requirement (elements) |
| 269 | // Beyond 4 merge rounds (>256 blocks), radix sort has little performance disadvantage. | 275 | // Beyond 4 merge rounds (>256 blocks), radix sort has little performance disadvantage. |
| 270 | constexpr uint32_t MERGE_INTRA_CORE_MAX_BLOCKS = 256; | 276 | constexpr uint32_t MERGE_INTRA_CORE_MAX_BLOCKS = 256; |
| @@ -280,7 +286,7 @@ struct MergeSortPlan { | |||
| 280 | 286 | ||
| 281 | bool ComputeMergeSortPlan(int64_t axisLen, int64_t unsortedDim, uint32_t blockUbSize, uint32_t tileDataNum, | 287 | bool ComputeMergeSortPlan(int64_t axisLen, int64_t unsortedDim, uint32_t blockUbSize, uint32_t tileDataNum, |
| 282 | uint32_t maxCoreNum, MergeSortPlan& plan); | 288 | uint32_t maxCoreNum, MergeSortPlan& plan); |
| 283 | -bool FillMergeSortInfo(SortKthTileInfo& info, uint32_t indexDtypeSize, uint32_t concatTmpSize); | 289 | +bool FillMergeSortInfo(SortKthTileInfo& info, uint32_t indexDtypeSize, bool useUbCapacity = false); |
| 284 | 290 | ||
| 285 | // ============================================================================= | 291 | // ============================================================================= |
| 286 | // Merge sort — multi-core (more-core) | 292 | // Merge sort — multi-core (more-core) |
| @@ -293,17 +299,23 @@ struct MergeMoreCorePlan { | |||
| 293 | uint32_t sortLoopTimes = 1; | 299 | uint32_t sortLoopTimes = 1; |
| 294 | uint32_t coreNumNeed = 0; | 300 | uint32_t coreNumNeed = 0; |
| 295 | uint32_t keyParams0 = 0; | 301 | uint32_t keyParams0 = 0; |
| 302 | + uint32_t keyParams1 = 0; | ||
| 296 | }; | 303 | }; |
| 297 | 304 | ||
| 298 | -bool IsMergeMoreCoreSupported(ge::DataType dataType, int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum); | 305 | +// Shared performance eligibility; the tiling builder still validates resource feasibility. |
| 299 | -bool ComputeMergeMoreCorePlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t mergeBytesPerElem, | 306 | +// Callers may impose operator-specific priorities before choosing this route. |
| 300 | - MergeMoreCorePlan& plan); | 307 | +bool IsMergeMoreCoreProfitable(ge::DataType dataType, int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum); |
| 308 | +bool SelectMergeMoreCoreDataSize(int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum, uint32_t& dataSize); | ||
| 309 | +bool SelectMergeSyncMergeBlockSize(int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum, uint32_t& blockSize); | ||
| 310 | +bool ComputeMergeMoreCorePlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t maxCoreNum, | ||
| 311 | + uint32_t mergeBytesPerElem, MergeMoreCorePlan& plan); | ||
| 301 | bool FillMergeMoreCoreInfo(SortKthTileInfo& info, uint32_t mergeBytesPerElem); | 312 | bool FillMergeMoreCoreInfo(SortKthTileInfo& info, uint32_t mergeBytesPerElem); |
| 302 | 313 | ||
| 303 | // ============================================================================= | 314 | // ============================================================================= |
| 304 | // Merge sort — intra-core | 315 | // Merge sort — intra-core |
| 305 | // ============================================================================= | 316 | // ============================================================================= |
| 306 | uint32_t ComputeMergeIntraCoreBlockSortSize(uint32_t ubSize); | 317 | uint32_t ComputeMergeIntraCoreBlockSortSize(uint32_t ubSize); |
| 318 | +uint32_t ComputeMergeIntraCoreSingleBlockSortSize(uint32_t ubSize); | ||
| 307 | uint32_t ComputeMergeIntraCoreExtractChunkSize(uint32_t ubSize); | 319 | uint32_t ComputeMergeIntraCoreExtractChunkSize(uint32_t ubSize); |
| 308 | 320 | ||
| 309 | struct MergeIntraCorePlan { | 321 | struct MergeIntraCorePlan { |
| @@ -370,7 +382,8 @@ void PlanToTilingData(const SortKthTileInfo& info, TilingData* tiling) | |||
| 370 | tiling->indexAxisBytes = info.indexAxisBytes; | 382 | tiling->indexAxisBytes = info.indexAxisBytes; |
| 371 | } | 383 | } |
| 372 | 384 | ||
| 373 | -bool ComputeMergeSortTiling(gert::TilingContext* context, SortKthTileInfo& info, uint32_t indexDtypeSize); | 385 | +bool ComputeMergeSortTiling(gert::TilingContext* context, SortKthTileInfo& info, uint32_t indexDtypeSize, |
| 386 | + bool useUbCapacity = false); | ||
| 374 | bool ComputeMergeMoreCoreTiling(gert::TilingContext* context, SortKthTileInfo& info, uint32_t mergeBytesPerElem); | 387 | bool ComputeMergeMoreCoreTiling(gert::TilingContext* context, SortKthTileInfo& info, uint32_t mergeBytesPerElem); |
| 375 | bool ComputeMergeIntraCoreTiling(gert::TilingContext* context, SortKthTileInfo& info); | 388 | bool ComputeMergeIntraCoreTiling(gert::TilingContext* context, SortKthTileInfo& info); |
| 376 | 389 | ||
| @@ -25,12 +25,12 @@ | |||
| 25 | #include "op_kernel/math_util.h" | 25 | #include "op_kernel/math_util.h" |
| 26 | #include "op_kernel/platform_util.h" | 26 | #include "op_kernel/platform_util.h" |
| 27 | #include "merge_sort_constants.h" | 27 | #include "merge_sort_constants.h" |
| 28 | +#include "ping_pong_merge_sort.h" | ||
| 28 | 29 | ||
| 29 | namespace MergeIntraCoreCommon { | 30 | namespace MergeIntraCoreCommon { |
| 30 | 31 | ||
| 31 | using namespace AscendC; | 32 | using namespace AscendC; |
| 32 | 33 | ||
| 33 | -using MergeSortConstants::DEALING_CONCAT_NUM_ONCE; | ||
| 34 | using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE; | 34 | using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE; |
| 35 | using MergeSortConstants::DEALING_SORT_NUM_ONCE; | 35 | using MergeSortConstants::DEALING_SORT_NUM_ONCE; |
| 36 | using MergeSortConstants::MERGE_INTRA_BUFFER_NUM; | 36 | using MergeSortConstants::MERGE_INTRA_BUFFER_NUM; |
| @@ -83,13 +83,12 @@ public: | |||
| 83 | uint32_t lastBlockSize_ = 0; // Actual element count of the last block | 83 | uint32_t lastBlockSize_ = 0; // Actual element count of the last block |
| 84 | uint32_t sortBufferSize_ = 0; // blockSortLen_ * sizeof(ValueType) | 84 | uint32_t sortBufferSize_ = 0; // blockSortLen_ * sizeof(ValueType) |
| 85 | uint32_t sortRepeatTimes_ = 0; // blockSortSize_ / DEALING_SORT_NUM_ONCE | 85 | uint32_t sortRepeatTimes_ = 0; // blockSortSize_ / DEALING_SORT_NUM_ONCE |
| 86 | - uint32_t concatRepeatTimes_ = 0; // blockSortSize_ / DEALING_CONCAT_NUM_ONCE | ||
| 87 | uint32_t maxMergeIterations_ = 0; // INT32_MAX / blockSortSize_ | 86 | uint32_t maxMergeIterations_ = 0; // INT32_MAX / blockSortSize_ |
| 88 | 87 | ||
| 89 | // Queues and buffers for UB sorting (Phase 1) | 88 | // Queues and buffers for UB sorting (Phase 1) |
| 90 | TQue<QuePosition::VECIN, MERGE_INTRA_BUFFER_NUM> inQueueX_; | 89 | TQue<QuePosition::VECIN, MERGE_INTRA_BUFFER_NUM> inQueueX_; |
| 91 | TQue<QuePosition::VECOUT, MERGE_INTRA_BUFFER_NUM> outValueQueue_, outIdxQueue_, outIdxInt64Queue_; | 90 | TQue<QuePosition::VECOUT, MERGE_INTRA_BUFFER_NUM> outValueQueue_, outIdxQueue_, outIdxInt64Queue_; |
| 92 | - TBuf<TPosition::VECCALC> concatTmpBuf_, sortTmpBuf_, indexTmpBuf_; | 91 | + TBuf<TPosition::VECCALC> sortTmpBuf_, indexTmpBuf_; |
| 93 | TQue<QuePosition::VECOUT, MERGE_INTRA_BUFFER_NUM> sortedOutQueue_; | 92 | TQue<QuePosition::VECOUT, MERGE_INTRA_BUFFER_NUM> sortedOutQueue_; |
| 94 | 93 | ||
| 95 | // Merge queues (Phase 2) | 94 | // Merge queues (Phase 2) |
| @@ -142,16 +141,21 @@ __aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDesce | |||
| 142 | return; | 141 | return; |
| 143 | } | 142 | } |
| 144 | 143 | ||
| 145 | - for (int64_t batchIdx = startBatch; batchIdx < endBatch; batchIdx++) { | 144 | + int64_t batchOffset = startBatch * this->sortAxisNum_; |
| 145 | + for (int64_t batchIdx = startBatch; batchIdx < endBatch; batchIdx++, batchOffset += this->sortAxisNum_) { | ||
| 146 | // Phase 1: derived owns the phase-specific UB allocation; base owns phase ordering. | 146 | // Phase 1: derived owns the phase-specific UB allocation; base owns phase ordering. |
| 147 | static_cast<Derived*>(this)->InitPhase1Buffers(); | 147 | static_cast<Derived*>(this)->InitPhase1Buffers(); |
| 148 | - this->SortSingleBatchInUb(this->inputXGm_, batchIdx * this->sortAxisNum_); | 148 | + this->SortSingleBatchInUb(this->inputXGm_, batchOffset); |
| 149 | PhaseBarrierAndReset(); | 149 | PhaseBarrierAndReset(); |
| 150 | 150 | ||
| 151 | // Phase 2: merge sorted blocks in the per-core ping-pong cache. | 151 | // Phase 2: merge sorted blocks in the per-core ping-pong cache. |
| 152 | - static_cast<Derived*>(this)->InitPhase2Buffers(); | 152 | + // A single block is already a complete sorted row in cache region 0, so skip the merge phase and its UB. |
| 153 | - uint32_t resultRegion = this->MergeSingleBatch(); | 153 | + uint32_t resultRegion = 0; |
| 154 | - PhaseBarrierAndReset(); | 154 | + if (this->blocksPerRow_ > 1) { |
| 155 | + static_cast<Derived*>(this)->InitPhase2Buffers(); | ||
| 156 | + resultRegion = this->MergeSingleBatch(); | ||
| 157 | + PhaseBarrierAndReset(); | ||
| 158 | + } | ||
| 155 | 159 | ||
| 156 | // Phase 3: derived writes either the full sorted row or the kth pair. | 160 | // Phase 3: derived writes either the full sorted row or the kth pair. |
| 157 | static_cast<Derived*>(this)->InitPhase3Buffers(); | 161 | static_cast<Derived*>(this)->InitPhase3Buffers(); |
| @@ -175,8 +179,8 @@ template <typename Derived, typename ValueType, typename IndexType, bool IsDesce | |||
| 175 | __aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDescend>::SortSingleBatchInUb( | 179 | __aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDescend>::SortSingleBatchInUb( |
| 176 | GlobalTensor<ValueType> inputX, int64_t batchOffset) | 180 | GlobalTensor<ValueType> inputX, int64_t batchOffset) |
| 177 | { | 181 | { |
| 178 | - for (uint32_t blockId = 0; blockId < this->blocksPerRow_; blockId++) { | 182 | + for (uint32_t blockId = 0, elemOffset = 0; blockId < this->blocksPerRow_; |
| 179 | - uint32_t elemOffset = blockId * this->blockSortSize_; | 183 | + blockId++, elemOffset += this->blockSortSize_) { |
| 180 | uint32_t actualElem = (blockId < this->blocksPerRow_ - 1) ? this->blockSortSize_ : this->lastBlockSize_; | 184 | uint32_t actualElem = (blockId < this->blocksPerRow_ - 1) ? this->blockSortSize_ : this->lastBlockSize_; |
| 181 | 185 | ||
| 182 | LocalTensor<ValueType> xLocal = this->inQueueX_.template AllocTensor<ValueType>(); | 186 | LocalTensor<ValueType> xLocal = this->inQueueX_.template AllocTensor<ValueType>(); |
| @@ -201,7 +205,7 @@ __aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDesce | |||
| 201 | { | 205 | { |
| 202 | ValueType defaultValue = IsDescend ? static_cast<ValueType>(-INFINITY) : static_cast<ValueType>(NAN); | 206 | ValueType defaultValue = IsDescend ? static_cast<ValueType>(-INFINITY) : static_cast<ValueType>(NAN); |
| 203 | 207 | ||
| 204 | - // For the last block (actualElem < blockSortSize_), Sort/Concat APIs process up to | 208 | + // For the last block (actualElem < blockSortSize_), Sort API processes up to |
| 205 | // CeilAlign(actualElem, UB_BLOCK_BYTES) elements. DataCopyPad only pads to | 209 | // CeilAlign(actualElem, UB_BLOCK_BYTES) elements. DataCopyPad only pads to |
| 206 | // CeilAlign(actualElem, UB_BLOCK_BYTES/sizeof(ValueType)) which may be smaller. | 210 | // CeilAlign(actualElem, UB_BLOCK_BYTES/sizeof(ValueType)) which may be smaller. |
| 207 | // Pre-fill with Duplicate to cover the full aligned region, then overwrite valid data. | 211 | // Pre-fill with Duplicate to cover the full aligned region, then overwrite valid data. |
| @@ -230,11 +234,6 @@ __aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDesce | |||
| 230 | uint32_t sortRepeatTimes = (actualElem == this->blockSortSize_) ? | 234 | uint32_t sortRepeatTimes = (actualElem == this->blockSortSize_) ? |
| 231 | this->sortRepeatTimes_ : | 235 | this->sortRepeatTimes_ : |
| 232 | Ops::Base::CeilDiv(alignSize, DEALING_SORT_NUM_ONCE); | 236 | Ops::Base::CeilDiv(alignSize, DEALING_SORT_NUM_ONCE); |
| 233 | - uint32_t concatRepeatTimes = (actualElem == this->blockSortSize_) ? | ||
| 234 | - this->concatRepeatTimes_ : | ||
| 235 | - Ops::Base::CeilDiv(alignSize, DEALING_CONCAT_NUM_ONCE); | ||
| 236 | - | ||
| 237 | - LocalTensor<ValueType> concatTmpLocal = this->concatTmpBuf_.template Get<ValueType>(); | ||
| 238 | LocalTensor<ValueType> sortTmpLocal = this->sortTmpBuf_.template Get<ValueType>(); | 237 | LocalTensor<ValueType> sortTmpLocal = this->sortTmpBuf_.template Get<ValueType>(); |
| 239 | LocalTensor<uint32_t> indexTmpLocal = this->indexTmpBuf_.template Get<uint32_t>(); | 238 | LocalTensor<uint32_t> indexTmpLocal = this->indexTmpBuf_.template Get<uint32_t>(); |
| 240 | ArithProgression<int32_t>(indexTmpLocal.template ReinterpretCast<int32_t>(), baseOffset, 1, actualElem); | 239 | ArithProgression<int32_t>(indexTmpLocal.template ReinterpretCast<int32_t>(), baseOffset, 1, actualElem); |
| @@ -247,9 +246,11 @@ __aicore__ inline void MergeIntraCoreBase<Derived, ValueType, IndexType, IsDesce | |||
| 247 | Adds(castTensor, castTensor, 0x80000000, alignSize); // Flip sign bit for ascending float order | 246 | Adds(castTensor, castTensor, 0x80000000, alignSize); // Flip sign bit for ascending float order |
| 248 | } | 247 | } |
| 249 | 248 | ||
| 250 | - LocalTensor<ValueType> concatLocal; | 249 | + bool resultInSorted = PingPongMergeSortCommon::SortToProposal(sortedLocal, sortTmpLocal, xLocal, indexTmpLocal, |
| 251 | - Concat(concatLocal, xLocal, concatTmpLocal, concatRepeatTimes); | 250 | + sortRepeatTimes); |
| 252 | - AscendC::Sort<ValueType, true>(sortedLocal, concatLocal, indexTmpLocal, sortTmpLocal, sortRepeatTimes); | 251 | + if (!resultInSorted) { |
| 252 | + DataCopy(sortedLocal, sortTmpLocal, AscendC::GetSortLen<ValueType>(alignSize)); | ||
| 253 | + } | ||
| 253 | } | 254 | } |
| 254 | 255 | ||
| 255 | template <typename Derived, typename ValueType, typename IndexType, bool IsDescend> | 256 | template <typename Derived, typename ValueType, typename IndexType, bool IsDescend> |
| @@ -20,13 +20,13 @@ | |||
| 20 | #include "kernel_operator.h" | 20 | #include "kernel_operator.h" |
| 21 | #include "op_kernel/platform_util.h" | 21 | #include "op_kernel/platform_util.h" |
| 22 | #include "merge_sort_constants.h" | 22 | #include "merge_sort_constants.h" |
| 23 | +#include "ping_pong_merge_sort.h" | ||
| 23 | #include "util_type_simd.h" | 24 | #include "util_type_simd.h" |
| 24 | 25 | ||
| 25 | namespace MergeMoreCoreCommon { | 26 | namespace MergeMoreCoreCommon { |
| 26 | 27 | ||
| 27 | using namespace AscendC; | 28 | using namespace AscendC; |
| 28 | 29 | ||
| 29 | -using MergeSortConstants::DEALING_CONCAT_NUM_ONCE; | ||
| 30 | using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE; | 30 | using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE; |
| 31 | using MergeSortConstants::DEALING_SORT_NUM_ONCE; | 31 | using MergeSortConstants::DEALING_SORT_NUM_ONCE; |
| 32 | using MergeSortConstants::FP32_DTYPE_BYTES; | 32 | using MergeSortConstants::FP32_DTYPE_BYTES; |
| @@ -77,27 +77,28 @@ public: | |||
| 77 | // ===== Local Buffers ===== | 77 | // ===== Local Buffers ===== |
| 78 | TBuf<QuePosition::VECCALC> sortedValueUb_; | 78 | TBuf<QuePosition::VECCALC> sortedValueUb_; |
| 79 | TBuf<QuePosition::VECCALC> sortedValueIndexUb_; | 79 | TBuf<QuePosition::VECCALC> sortedValueIndexUb_; |
| 80 | - TBuf<QuePosition::VECCALC> concatTempBuf_; | ||
| 81 | TBuf<QuePosition::VECCALC> sortTempBuf_; | 80 | TBuf<QuePosition::VECCALC> sortTempBuf_; |
| 82 | - TBuf<QuePosition::VECCALC> sortedValueLocalCastTbuf_; | ||
| 83 | 81 | ||
| 84 | // ===== Scalar Members ===== | 82 | // ===== Scalar Members ===== |
| 85 | TPipe* pipe_ = nullptr; | 83 | TPipe* pipe_ = nullptr; |
| 86 | uint32_t blockIdx_ = 0; | 84 | uint32_t blockIdx_ = 0; |
| 87 | uint32_t numTileData_ = 0; | 85 | uint32_t numTileData_ = 0; |
| 88 | - uint32_t sortLoopRound_ = 0; | ||
| 89 | - uint32_t platformCoreNum_ = 0; | ||
| 90 | uint32_t outputLastDimValue_ = 0; | 86 | uint32_t outputLastDimValue_ = 0; |
| 91 | uint32_t frontCoreNum_ = 0; | 87 | uint32_t frontCoreNum_ = 0; |
| 88 | + uint32_t rowGroupIdx_ = 0; | ||
| 92 | uint32_t rowIdx_ = 0; | 89 | uint32_t rowIdx_ = 0; |
| 93 | uint32_t rowCoreIdx_ = 0; | 90 | uint32_t rowCoreIdx_ = 0; |
| 91 | + uint32_t rowsPerRound_ = 0; | ||
| 92 | + uint32_t sortLoopTimes_ = 1; | ||
| 93 | + uint32_t logicalBlockSize_ = 0; | ||
| 94 | + bool syncMergeSortMode_ = false; | ||
| 95 | + int64_t batchNum_ = 0; | ||
| 94 | uint32_t vfLenFp32_ = Ops::Base::GetVRegSize() / FP32_DTYPE_BYTES; | 96 | uint32_t vfLenFp32_ = Ops::Base::GetVRegSize() / FP32_DTYPE_BYTES; |
| 95 | int64_t rowDataOffset_ = 0; | 97 | int64_t rowDataOffset_ = 0; |
| 96 | int64_t rowWorkspaceOffset_ = 0; | 98 | int64_t rowWorkspaceOffset_ = 0; |
| 97 | 99 | ||
| 98 | // ===== Merge-sort state ===== | 100 | // ===== Merge-sort state ===== |
| 99 | int64_t listNum_{0}; | 101 | int64_t listNum_{0}; |
| 100 | - int64_t flag_ = 0; | ||
| 101 | int64_t remainListNum_{0}; | 102 | int64_t remainListNum_{0}; |
| 102 | int64_t outOffset_{0}; | 103 | int64_t outOffset_{0}; |
| 103 | int64_t offsets_[4] = {0}; | 104 | int64_t offsets_[4] = {0}; |
| @@ -112,49 +113,173 @@ public: | |||
| 112 | uint16_t validBitTail_; | 113 | uint16_t validBitTail_; |
| 113 | uint32_t listSortedNums_[4] = {0}; | 114 | uint32_t listSortedNums_[4] = {0}; |
| 114 | uint32_t workSpaceFlag_ = 0; | 115 | uint32_t workSpaceFlag_ = 0; |
| 115 | - LocalTensor<CONVERT_TYPE> ubInputs_[4]; | ||
| 116 | - LocalTensor<CONVERT_TYPE> ubMainInput_; | ||
| 117 | 116 | ||
| 118 | // ======================== Shared Member Functions ======================== | 117 | // ======================== Shared Member Functions ======================== |
| 119 | 118 | ||
| 120 | __aicore__ inline void Process() | 119 | __aicore__ inline void Process() |
| 120 | + { | ||
| 121 | + if (this->frontCoreNum_ == 0U || this->rowsPerRound_ == 0U || this->sortLoopTimes_ == 0U) { | ||
| 122 | + return; | ||
| 123 | + } | ||
| 124 | + for (uint32_t loopIdx = 0; loopIdx < this->sortLoopTimes_; loopIdx++) { | ||
| 125 | + this->rowIdx_ = loopIdx * this->rowsPerRound_ + this->rowGroupIdx_; | ||
| 126 | + // In multi-round mode, the last round may have fewer rows than rowsPerRound_. | ||
| 127 | + // Cores assigned to out-of-range rows skip real work but still join SyncAll to avoid deadlock. | ||
| 128 | + bool rowActive = static_cast<int64_t>(this->rowIdx_) < this->batchNum_; | ||
| 129 | + if (rowActive) { | ||
| 130 | + static_cast<Derived*>(this)->PrepareRowWorkspace(); | ||
| 131 | + static_cast<Derived*>(this)->InitSortBuffers(); | ||
| 132 | + } | ||
| 133 | + | ||
| 134 | + int64_t offsetPerCore = 0; | ||
| 135 | + uint32_t tileNum = 0; | ||
| 136 | + if (rowActive) { | ||
| 137 | + uint32_t phase1BlockSize = this->syncMergeSortMode_ ? this->logicalBlockSize_ : this->numTileData_; | ||
| 138 | + offsetPerCore = static_cast<int64_t>(phase1BlockSize) * static_cast<int64_t>(this->rowCoreIdx_); | ||
| 139 | + if (this->rowCoreIdx_ < this->frontCoreNum_ - 1) { | ||
| 140 | + tileNum = phase1BlockSize; | ||
| 141 | + } else if (this->rowCoreIdx_ == this->frontCoreNum_ - 1) { | ||
| 142 | + tileNum = this->outputLastDimValue_ - phase1BlockSize * (this->frontCoreNum_ - 1); | ||
| 143 | + } | ||
| 144 | + } | ||
| 145 | + if (rowActive && tileNum > 0U) { | ||
| 146 | + if (this->syncMergeSortMode_) { | ||
| 147 | + this->SortLogicalBlock(tileNum, offsetPerCore); | ||
| 148 | + } else { | ||
| 149 | + this->SortInSingleCore(tileNum, offsetPerCore); | ||
| 150 | + } | ||
| 151 | + } | ||
| 152 | + SyncAll(); | ||
| 153 | + this->pipe_->Reset(); | ||
| 154 | + | ||
| 155 | + if (rowActive) { | ||
| 156 | + static_cast<Derived*>(this)->InitMergeBuffers(); | ||
| 157 | + if (this->syncMergeSortMode_ && tileNum > 0U) { | ||
| 158 | + this->workspaceInput_ = this->workspaceGm_[0]; | ||
| 159 | + this->workspaceOutput_ = this->workspaceGm_[1]; | ||
| 160 | + this->MergeLogicalBlock(tileNum, offsetPerCore); | ||
| 161 | + this->ClearCache(); | ||
| 162 | + } | ||
| 163 | + } | ||
| 164 | + if (this->syncMergeSortMode_) { | ||
| 165 | + SyncAll(); | ||
| 166 | + } | ||
| 167 | + | ||
| 168 | + InitMergeState(this->syncMergeSortMode_ ? this->logicalBlockSize_ : this->numTileData_, | ||
| 169 | + this->syncMergeSortMode_ ? 1U : 0U); | ||
| 170 | + MergeAcrossCores(rowActive); | ||
| 171 | + this->pipe_->Reset(); | ||
| 172 | + } | ||
| 173 | + } | ||
| 174 | + | ||
| 175 | + // The one-round, non-sync schedule has one active row per core group. Keeping it in a dedicated | ||
| 176 | + // kernel instance removes the generic row-activity and sync-mode branches from this hot orchestration path. | ||
| 177 | + __aicore__ inline void ProcessDirect() | ||
| 121 | { | 178 | { |
| 122 | if (this->frontCoreNum_ == 0U) { | 179 | if (this->frontCoreNum_ == 0U) { |
| 123 | return; | 180 | return; |
| 124 | } | 181 | } |
| 125 | - int64_t offsetPerCore = 0; | 182 | + int64_t offsetPerCore = static_cast<int64_t>(this->numTileData_) * this->rowCoreIdx_; |
| 126 | - if (this->rowCoreIdx_ < this->frontCoreNum_ - 1) { | 183 | + uint32_t tileNum = this->rowCoreIdx_ < this->frontCoreNum_ - 1 ? |
| 127 | - offsetPerCore = this->numTileData_ * this->rowCoreIdx_; | 184 | + this->numTileData_ : |
| 128 | - this->SortInSingleCore(this->numTileData_, offsetPerCore); | 185 | + this->outputLastDimValue_ - this->numTileData_ * (this->frontCoreNum_ - 1); |
| 129 | - } else if (this->rowCoreIdx_ == this->frontCoreNum_ - 1) { | 186 | + SortInSingleCore(tileNum, offsetPerCore); |
| 130 | - uint32_t tailNum = this->outputLastDimValue_ - this->numTileData_ * (this->frontCoreNum_ - 1); | ||
| 131 | - offsetPerCore = this->numTileData_ * this->rowCoreIdx_; | ||
| 132 | - this->SortInSingleCore(tailNum, offsetPerCore); | ||
| 133 | - } | ||
| 134 | SyncAll(); | 187 | SyncAll(); |
| 135 | this->pipe_->Reset(); | 188 | this->pipe_->Reset(); |
| 136 | 189 | ||
| 137 | static_cast<Derived*>(this)->InitMergeBuffers(); | 190 | static_cast<Derived*>(this)->InitMergeBuffers(); |
| 191 | + InitMergeState(this->numTileData_, 0U); | ||
| 192 | + MergeAcrossCores(true); | ||
| 193 | + } | ||
| 138 | 194 | ||
| 195 | + __aicore__ inline void InitMergeState(int64_t blockElements, uint32_t workspaceFlag) | ||
| 196 | + { | ||
| 197 | + this->workSpaceFlag_ = workspaceFlag; | ||
| 139 | this->listNum_ = this->frontCoreNum_; | 198 | this->listNum_ = this->frontCoreNum_; |
| 140 | - this->currentElements_ = this->numTileData_; | 199 | + this->currentElements_ = blockElements; |
| 141 | - this->currentTailElements_ = this->outputLastDimValue_ - this->numTileData_ * (this->frontCoreNum_ - 1); | 200 | + this->currentTailElements_ = this->outputLastDimValue_ - blockElements * (this->frontCoreNum_ - 1); |
| 142 | - uint32_t currentCoreNum; | 201 | + } |
| 143 | - uint32_t remainListNum; | 202 | + |
| 203 | + // Inactive row groups advance the same merge levels and participate in the | ||
| 204 | + // barriers inside SortInMultiCore/FinalSortAndCopy, without accessing row workspace. | ||
| 205 | + __aicore__ inline void MergeAcrossCores(bool rowActive) | ||
| 206 | + { | ||
| 144 | while (this->listNum_ > MERGE_LIST_MAX_NUM) { | 207 | while (this->listNum_ > MERGE_LIST_MAX_NUM) { |
| 145 | - this->workspaceInput_ = this->workspaceGm_[this->workSpaceFlag_]; | 208 | + if (rowActive) { |
| 146 | - this->workspaceOutput_ = this->workspaceGm_[1 - this->workSpaceFlag_]; | 209 | + this->workspaceInput_ = this->workspaceGm_[this->workSpaceFlag_]; |
| 147 | - this->SortInMultiCore(); | 210 | + this->workspaceOutput_ = this->workspaceGm_[1U - this->workSpaceFlag_]; |
| 148 | - currentCoreNum = Ops::Base::CeilDiv(this->listNum_, static_cast<int64_t>(MERGE_LIST_MAX_NUM)); | 211 | + } |
| 149 | - remainListNum = this->listNum_ - (currentCoreNum - 1) * MERGE_LIST_MAX_NUM; | 212 | + SortInMultiCore(rowActive); |
| 150 | - this->currentTailElements_ = this->currentElements_ * (remainListNum - 1) + this->currentTailElements_; | 213 | + uint32_t currentCoreNum = Ops::Base::CeilDiv(this->listNum_, static_cast<int64_t>(MERGE_LIST_MAX_NUM)); |
| 214 | + uint32_t remainListNum = this->listNum_ - (currentCoreNum - 1U) * MERGE_LIST_MAX_NUM; | ||
| 215 | + this->currentTailElements_ = this->currentElements_ * (remainListNum - 1U) + this->currentTailElements_; | ||
| 151 | this->listNum_ = currentCoreNum; | 216 | this->listNum_ = currentCoreNum; |
| 152 | - this->currentElements_ = this->currentElements_ * MERGE_LIST_MAX_NUM; | 217 | + this->currentElements_ *= MERGE_LIST_MAX_NUM; |
| 153 | - this->workSpaceFlag_ = (this->workSpaceFlag_ + 1) % MERGE_WORKSPACE_BUFFER_NUM; | 218 | + this->workSpaceFlag_ = 1U - this->workSpaceFlag_; |
| 219 | + } | ||
| 220 | + if (rowActive) { | ||
| 221 | + this->workspaceInput_ = this->workspaceGm_[this->workSpaceFlag_]; | ||
| 222 | + this->workspaceOutput_ = this->workspaceGm_[1U - this->workSpaceFlag_]; | ||
| 223 | + } | ||
| 224 | + static_cast<Derived*>(this)->FinalSortAndCopy(rowActive); | ||
| 225 | + } | ||
| 226 | + | ||
| 227 | + // Sync-merge mode: sort a logical block (larger than numTileData_) in multiple passes, | ||
| 228 | + // each pass handling one numTileData_-sized chunk and flushing to workspaceGm_[0]. | ||
| 229 | + __aicore__ inline void SortLogicalBlock(uint32_t tileNum, int64_t offsetPerCore) | ||
| 230 | + { | ||
| 231 | + uint32_t sortedNum = 0; | ||
| 232 | + while (sortedNum < tileNum) { | ||
| 233 | + uint32_t curTileNum = this->numTileData_; | ||
| 234 | + if (curTileNum > tileNum - sortedNum) { | ||
| 235 | + curTileNum = tileNum - sortedNum; | ||
| 236 | + } | ||
| 237 | + this->SortInSingleCore(curTileNum, offsetPerCore + static_cast<int64_t>(sortedNum)); | ||
| 238 | + event_t eventIdMTE3ToS = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE3_S)); | ||
| 239 | + SetFlag<HardEvent::MTE3_S>(eventIdMTE3ToS); | ||
| 240 | + WaitFlag<HardEvent::MTE3_S>(eventIdMTE3ToS); | ||
| 241 | + sortedNum += curTileNum; | ||
| 242 | + } | ||
| 243 | + } | ||
| 244 | + | ||
| 245 | + // Sync-merge mode: merge the locally sorted chunks (produced by SortLogicalBlock) into one | ||
| 246 | + // contiguous sorted logical block, using the streaming merge loop (CopyIn → MrgSort → CopyOut). | ||
| 247 | + __aicore__ inline void MergeLogicalBlock(uint32_t tileNum, int64_t offsetPerCore) | ||
| 248 | + { | ||
| 249 | + this->listNum_ = Ops::Base::CeilDiv(tileNum, this->numTileData_); | ||
| 250 | + this->currentElements_ = this->numTileData_; | ||
| 251 | + this->currentTailElements_ = tileNum - this->numTileData_ * (this->listNum_ - 1); | ||
| 252 | + this->InitLogicalBlockMerge(tileNum, offsetPerCore); | ||
| 253 | + for (; this->allRemainElements_ > 0;) { | ||
| 254 | + CopyInMultiCore(); | ||
| 255 | + UpdateMrgParam(); | ||
| 256 | + DealingMergeSort(); | ||
| 257 | + UpdateSortInfo(); | ||
| 258 | + CopyOutMultiCore(); | ||
| 259 | + } | ||
| 260 | + } | ||
| 261 | + | ||
| 262 | + // Set up per-run offsets and remaining element counts for the logical block merge. | ||
| 263 | + // Each run corresponds to one numTileData_-sized chunk sorted in SortLogicalBlock. | ||
| 264 | + __aicore__ inline void InitLogicalBlockMerge(uint32_t tileNum, int64_t offsetPerCore) | ||
| 265 | + { | ||
| 266 | + this->outOffset_ = GetSortLen<CONVERT_TYPE>(offsetPerCore); | ||
| 267 | + for (int64_t i = 0; i < MERGE_LIST_MAX_NUM; i++) { | ||
| 268 | + if (i < this->listNum_ - 1) { | ||
| 269 | + this->listRemainElements_[i] = this->numTileData_; | ||
| 270 | + this->offsets_[i] = GetSortOffset<CONVERT_TYPE>(offsetPerCore + i * this->numTileData_); | ||
| 271 | + this->allRemainElements_ += this->listRemainElements_[i]; | ||
| 272 | + } else if (i == this->listNum_ - 1) { | ||
| 273 | + this->listRemainElements_[i] = this->currentTailElements_; | ||
| 274 | + this->offsets_[i] = GetSortOffset<CONVERT_TYPE>(offsetPerCore + i * this->numTileData_); | ||
| 275 | + this->allRemainElements_ += this->currentTailElements_; | ||
| 276 | + } else { | ||
| 277 | + this->listRemainElements_[i] = 0; | ||
| 278 | + } | ||
| 279 | + } | ||
| 280 | + if (this->listNum_ == 1) { | ||
| 281 | + this->currentTailElements_ = tileNum; | ||
| 154 | } | 282 | } |
| 155 | - this->workspaceInput_ = this->workspaceGm_[this->workSpaceFlag_]; | ||
| 156 | - this->workspaceOutput_ = this->workspaceGm_[1 - this->workSpaceFlag_]; | ||
| 157 | - this->FinalSortAndCopy(); | ||
| 158 | } | 283 | } |
| 159 | 284 | ||
| 160 | __aicore__ inline void SortInSingleCore(uint32_t tileNum, int64_t offsetPerCore) | 285 | __aicore__ inline void SortInSingleCore(uint32_t tileNum, int64_t offsetPerCore) |
| @@ -172,6 +297,8 @@ public: | |||
| 172 | this->inputQueue_.FreeTensor(inputLocal); | 297 | this->inputQueue_.FreeTensor(inputLocal); |
| 173 | } | 298 | } |
| 174 | 299 | ||
| 300 | + // Copy tileNum elements from GM to UB, padding the tail to alignment with defaultValue | ||
| 301 | + // (−INF for descending, NAN for ascending) so Sort API processes a full aligned block. | ||
| 175 | __aicore__ inline void CopyInData(uint32_t tileNum, int64_t offsetPerCore) | 302 | __aicore__ inline void CopyInData(uint32_t tileNum, int64_t offsetPerCore) |
| 176 | { | 303 | { |
| 177 | LocalTensor<T> inputLocal = this->inputQueue_.template AllocTensor<T>(); | 304 | LocalTensor<T> inputLocal = this->inputQueue_.template AllocTensor<T>(); |
| @@ -213,22 +340,18 @@ public: | |||
| 213 | LocalTensor<uint32_t> sortedValueIndexLocal) | 340 | LocalTensor<uint32_t> sortedValueIndexLocal) |
| 214 | { | 341 | { |
| 215 | AscendC::LocalTensor<CONVERT_TYPE> sortTempLocal = this->sortTempBuf_.template Get<CONVERT_TYPE>(); | 342 | AscendC::LocalTensor<CONVERT_TYPE> sortTempLocal = this->sortTempBuf_.template Get<CONVERT_TYPE>(); |
| 216 | - AscendC::LocalTensor<CONVERT_TYPE> concatTempLocal = this->concatTempBuf_.template Get<CONVERT_TYPE>(); | ||
| 217 | - AscendC::LocalTensor<CONVERT_TYPE> sortedValueLocalCast = this->sortedValueLocalCastTbuf_ | ||
| 218 | - .template Get<CONVERT_TYPE>(); | ||
| 219 | - | ||
| 220 | uint32_t aglinTileNum = ROUND_UP_AGLIN(tileNum); | 343 | uint32_t aglinTileNum = ROUND_UP_AGLIN(tileNum); |
| 221 | uint32_t sortRepeatTimes = Ops::Base::CeilDiv(aglinTileNum, DEALING_SORT_NUM_ONCE); | 344 | uint32_t sortRepeatTimes = Ops::Base::CeilDiv(aglinTileNum, DEALING_SORT_NUM_ONCE); |
| 222 | - uint32_t concatRepeatTimes = Ops::Base::CeilDiv(aglinTileNum, DEALING_CONCAT_NUM_ONCE); | ||
| 223 | 345 | ||
| 224 | static_cast<Derived*>(this)->PrepareInputForSort(inputLocal, aglinTileNum); | 346 | static_cast<Derived*>(this)->PrepareInputForSort(inputLocal, aglinTileNum); |
| 225 | if constexpr (!IS_DESCEND) { | 347 | if constexpr (!IS_DESCEND) { |
| 226 | FlipSignBit(inputLocal, aglinTileNum); | 348 | FlipSignBit(inputLocal, aglinTileNum); |
| 227 | } | 349 | } |
| 228 | - AscendC::LocalTensor<CONVERT_TYPE> concatLocal; | 350 | + bool resultInSorted = PingPongMergeSortCommon::SortToProposal(sortedValueLocal, sortTempLocal, inputLocal, |
| 229 | - AscendC::Concat(concatLocal, inputLocal, concatTempLocal, concatRepeatTimes); | 351 | + sortedValueIndexLocal, sortRepeatTimes); |
| 230 | - AscendC::Sort<CONVERT_TYPE, true>(sortedValueLocal, concatLocal, sortedValueIndexLocal, sortTempLocal, | 352 | + if (!resultInSorted) { |
| 231 | - sortRepeatTimes); | 353 | + DataCopy(sortedValueLocal, sortTempLocal, GetSortOffset<CONVERT_TYPE>(aglinTileNum)); |
| 354 | + } | ||
| 232 | } | 355 | } |
| 233 | 356 | ||
| 234 | __aicore__ inline void FlipSignBit(LocalTensor<CONVERT_TYPE> xLocal, uint32_t aglinTileNum) | 357 | __aicore__ inline void FlipSignBit(LocalTensor<CONVERT_TYPE> xLocal, uint32_t aglinTileNum) |
| @@ -272,12 +395,12 @@ public: | |||
| 272 | this->sortedQueue_.template FreeTensor<CONVERT_TYPE>(sortTempBuffer); | 395 | this->sortedQueue_.template FreeTensor<CONVERT_TYPE>(sortTempBuffer); |
| 273 | } | 396 | } |
| 274 | 397 | ||
| 275 | - __aicore__ inline void SortInMultiCore() | 398 | + __aicore__ inline void SortInMultiCore(bool rowActive) |
| 276 | { | 399 | { |
| 277 | // Phase 2: every active row core merges up to 4 sorted runs. The merged output is written to the alternate | 400 | // Phase 2: every active row core merges up to 4 sorted runs. The merged output is written to the alternate |
| 278 | // workspace buffer, and Process/derived code swaps the input/output workspace for the next merge level. | 401 | // workspace buffer, and Process/derived code swaps the input/output workspace for the next merge level. |
| 279 | uint32_t needCoreNum = Ops::Base::CeilDiv(this->listNum_, static_cast<int64_t>(MERGE_LIST_MAX_NUM)); | 402 | uint32_t needCoreNum = Ops::Base::CeilDiv(this->listNum_, static_cast<int64_t>(MERGE_LIST_MAX_NUM)); |
| 280 | - if (this->rowCoreIdx_ < needCoreNum) { | 403 | + if (rowActive && this->rowCoreIdx_ < needCoreNum) { |
| 281 | MultiCoreInit(); | 404 | MultiCoreInit(); |
| 282 | for (; this->allRemainElements_ > 0;) { | 405 | for (; this->allRemainElements_ > 0;) { |
| 283 | CopyInMultiCore(); | 406 | CopyInMultiCore(); |
| @@ -291,11 +414,11 @@ public: | |||
| 291 | SyncAll(); | 414 | SyncAll(); |
| 292 | } | 415 | } |
| 293 | 416 | ||
| 294 | - __aicore__ inline void FinalSortAndCopy() | 417 | + __aicore__ inline void FinalSortAndCopy(bool rowActive) |
| 295 | { | 418 | { |
| 296 | // Phase 3: once at most 4 runs remain for the row, rowCoreIdx_ 0 performs the final merge and lets the | 419 | // Phase 3: once at most 4 runs remain for the row, rowCoreIdx_ 0 performs the final merge and lets the |
| 297 | // derived kernel extract values/indices directly to the operator outputs. | 420 | // derived kernel extract values/indices directly to the operator outputs. |
| 298 | - if (this->rowCoreIdx_ == 0) { | 421 | + if (rowActive && this->rowCoreIdx_ == 0) { |
| 299 | FinalSortInit(); | 422 | FinalSortInit(); |
| 300 | for (; this->allRemainElements_ > 0;) { | 423 | for (; this->allRemainElements_ > 0;) { |
| 301 | CopyInMultiCore(); | 424 | CopyInMultiCore(); |
| @@ -395,7 +518,26 @@ public: | |||
| 395 | } | 518 | } |
| 396 | } | 519 | } |
| 397 | 520 | ||
| 521 | + // Perform MrgSort on 2/3/4 source runs. Unused source-list slots are filled with a duplicate | ||
| 522 | + // of tmpUbInputs[0] because MrgSort requires valid tensor addresses for all four parameters; | ||
| 523 | + // validBitTail_ controls which lists actually participate. | ||
| 398 | __aicore__ inline void DealingMergeSort() | 524 | __aicore__ inline void DealingMergeSort() |
| 525 | + { | ||
| 526 | + // A full window needs no refill. Consume every run in one merge rather | ||
| 527 | + // than stopping when the first source is exhausted. | ||
| 528 | + bool allLoaded = true; | ||
| 529 | + for (uint32_t i = 0; i < MERGE_LIST_MAX_NUM; ++i) { | ||
| 530 | + allLoaded = allLoaded && this->dealLengths_[i] == this->listRemainElements_[i]; | ||
| 531 | + } | ||
| 532 | + if (allLoaded) { | ||
| 533 | + MergeLoadedLists<false>(); | ||
| 534 | + return; | ||
| 535 | + } | ||
| 536 | + MergeLoadedLists<true>(); | ||
| 537 | + } | ||
| 538 | + | ||
| 539 | + template <bool SuspendOnExhaustion> | ||
| 540 | + __aicore__ inline void MergeLoadedLists() | ||
| 399 | { | 541 | { |
| 400 | LocalTensor<CONVERT_TYPE> sortTempBuffer = this->sortedQueue_.template AllocTensor<CONVERT_TYPE>(); | 542 | LocalTensor<CONVERT_TYPE> sortTempBuffer = this->sortedQueue_.template AllocTensor<CONVERT_TYPE>(); |
| 401 | LocalTensor<CONVERT_TYPE> ubMainInput = this->copyInQueue_.template DeQue<CONVERT_TYPE>(); | 543 | LocalTensor<CONVERT_TYPE> ubMainInput = this->copyInQueue_.template DeQue<CONVERT_TYPE>(); |
| @@ -410,24 +552,31 @@ public: | |||
| 410 | // Unused source-list arguments are duplicated because validBitTail controls which lists participate. | 552 | // Unused source-list arguments are duplicated because validBitTail controls which lists participate. |
| 411 | MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], | 553 | MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[0], |
| 412 | tmpUbInputs[0]); | 554 | tmpUbInputs[0]); |
| 413 | - MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, this->elementCountList_, this->listSortedNums_, | 555 | + MrgSort<CONVERT_TYPE, SuspendOnExhaustion>(sortTempBuffer, sortListTail, this->elementCountList_, |
| 414 | - this->validBitTail_, 1); | 556 | + this->listSortedNums_, this->validBitTail_, 1); |
| 415 | } else if (this->remainListNum_ == THREE_WAY_MERGE_LIST_NUM) { | 557 | } else if (this->remainListNum_ == THREE_WAY_MERGE_LIST_NUM) { |
| 416 | MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[2], | 558 | MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[2], |
| 417 | tmpUbInputs[0]); | 559 | tmpUbInputs[0]); |
| 418 | - MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, this->elementCountList_, this->listSortedNums_, | 560 | + MrgSort<CONVERT_TYPE, SuspendOnExhaustion>(sortTempBuffer, sortListTail, this->elementCountList_, |
| 419 | - this->validBitTail_, 1); | 561 | + this->listSortedNums_, this->validBitTail_, 1); |
| 420 | } else if (this->remainListNum_ == MERGE_LIST_MAX_NUM) { | 562 | } else if (this->remainListNum_ == MERGE_LIST_MAX_NUM) { |
| 421 | MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[2], | 563 | MrgSortSrcList sortListTail = MrgSortSrcList(tmpUbInputs[0], tmpUbInputs[1], tmpUbInputs[2], |
| 422 | tmpUbInputs[3]); | 564 | tmpUbInputs[3]); |
| 423 | - MrgSort<CONVERT_TYPE, true>(sortTempBuffer, sortListTail, this->elementCountList_, this->listSortedNums_, | 565 | + MrgSort<CONVERT_TYPE, SuspendOnExhaustion>(sortTempBuffer, sortListTail, this->elementCountList_, |
| 424 | - this->validBitTail_, 1); | 566 | + this->listSortedNums_, this->validBitTail_, 1); |
| 425 | } else { | 567 | } else { |
| 426 | AscendC::Copy(sortTempBuffer, tmpUbInputs[0], | 568 | AscendC::Copy(sortTempBuffer, tmpUbInputs[0], |
| 427 | ROUND_UP_AGLIN(GetSortLen<CONVERT_TYPE>(this->elementCountList_[0]) * sizeof(CONVERT_TYPE)) / | 569 | ROUND_UP_AGLIN(GetSortLen<CONVERT_TYPE>(this->elementCountList_[0]) * sizeof(CONVERT_TYPE)) / |
| 428 | sizeof(CONVERT_TYPE)); | 570 | sizeof(CONVERT_TYPE)); |
| 429 | this->listSortedNums_[0] = this->elementCountList_[0]; | 571 | this->listSortedNums_[0] = this->elementCountList_[0]; |
| 430 | } | 572 | } |
| 573 | + if constexpr (!SuspendOnExhaustion) { | ||
| 574 | + // Non-suspending MrgSort does not report per-list consumption. All | ||
| 575 | + // remaining records were loaded, so their input lengths are exact. | ||
| 576 | + for (uint32_t i = 0; i < this->remainListNum_; ++i) { | ||
| 577 | + this->listSortedNums_[i] = this->elementCountList_[i]; | ||
| 578 | + } | ||
| 579 | + } | ||
| 431 | this->sortedQueue_.EnQue(sortTempBuffer); | 580 | this->sortedQueue_.EnQue(sortTempBuffer); |
| 432 | this->copyInQueue_.FreeTensor(ubMainInput); | 581 | this->copyInQueue_.FreeTensor(ubMainInput); |
| 433 | } | 582 | } |
| @@ -24,6 +24,8 @@ | |||
| 24 | #include "simt_api/asc_simt.h" | 24 | #include "simt_api/asc_simt.h" |
| 25 | #include "signed_zero_sort_utils.h" | 25 | #include "signed_zero_sort_utils.h" |
| 26 | #include "util_type_simd.h" | 26 | #include "util_type_simd.h" |
| 27 | +#include "merge_sort_constants.h" | ||
| 28 | +#include "ping_pong_merge_sort.h" | ||
| 27 | 29 | ||
| 28 | namespace SmallAxisCommon { | 30 | namespace SmallAxisCommon { |
| 29 | using namespace AscendC; | 31 | using namespace AscendC; |
| @@ -61,6 +63,9 @@ public: | |||
| 61 | const uint64_t startTile = static_cast<uint64_t>(this->blockIdx_) * tilesPerCore; | 63 | const uint64_t startTile = static_cast<uint64_t>(this->blockIdx_) * tilesPerCore; |
| 62 | uint64_t endTile = startTile + tilesPerCore; | 64 | uint64_t endTile = startTile + tilesPerCore; |
| 63 | endTile = endTile > tileCount ? tileCount : endTile; | 65 | endTile = endTile > tileCount ? tileCount : endTile; |
| 66 | + // The derived class may place this tensor in shared UB instead of sortedValueBuf_. | ||
| 67 | + // Keep its allocated address across tiles: merge Extract redirects sortedValue_ to sortInput_. | ||
| 68 | + LocalTensor<SortT> proposalPing = this->sortedValue_; | ||
| 64 | for (uint64_t tileId = startTile; tileId < endTile; ++tileId) { | 69 | for (uint64_t tileId = startTile; tileId < endTile; ++tileId) { |
| 65 | const uint64_t outerId = tileId / this->innerLoopNum_; | 70 | const uint64_t outerId = tileId / this->innerLoopNum_; |
| 66 | const uint32_t innerTileId = static_cast<uint32_t>(tileId - | 71 | const uint32_t innerTileId = static_cast<uint32_t>(tileId - |
| @@ -76,7 +81,7 @@ public: | |||
| 76 | int64_t outputOffset = static_cast<int64_t>(outerId) * this->innerSize_ + innerStart; | 81 | int64_t outputOffset = static_cast<int64_t>(outerId) * this->innerSize_ + innerStart; |
| 77 | this->LoadTile(inputOffset, curInnerChunk); | 82 | this->LoadTile(inputOffset, curInnerChunk); |
| 78 | this->TransposeToSortMajor(curInnerChunk); | 83 | this->TransposeToSortMajor(curInnerChunk); |
| 79 | - this->SortRows(curInnerChunk); | 84 | + this->SortRows(curInnerChunk, proposalPing); |
| 80 | static_cast<Derived*>(this)->StoreTile(inputOffset, outputOffset, curInnerChunk); | 85 | static_cast<Derived*>(this)->StoreTile(inputOffset, outputOffset, curInnerChunk); |
| 81 | } | 86 | } |
| 82 | } | 87 | } |
| @@ -85,6 +90,7 @@ protected: | |||
| 85 | static constexpr SortConfig sortConfig_{UseMergeSort ? SortType::MERGE_SORT : SortType::RADIX_SORT, IsDescend}; | 90 | static constexpr SortConfig sortConfig_{UseMergeSort ? SortType::MERGE_SORT : SortType::RADIX_SORT, IsDescend}; |
| 86 | 91 | ||
| 87 | TPipe* pipe_ = nullptr; | 92 | TPipe* pipe_ = nullptr; |
| 93 | + bool compactMerge_ = false; | ||
| 88 | 94 | ||
| 89 | GlobalTensor<T> inputGm_; | 95 | GlobalTensor<T> inputGm_; |
| 90 | 96 | ||
| @@ -148,6 +154,31 @@ protected: | |||
| 148 | 154 | ||
| 149 | __aicore__ inline void LoadTile(int64_t baseOffset, uint32_t curInnerChunk) | 155 | __aicore__ inline void LoadTile(int64_t baseOffset, uint32_t curInnerChunk) |
| 150 | { | 156 | { |
| 157 | + if constexpr (UseMergeSort && !IsBf16Merge) { | ||
| 158 | + if (this->compactMerge_) { | ||
| 159 | + // Load GM [axis, inner] directly into compact UB [inner, axis]. | ||
| 160 | + // The aligned data stride adds at most one block beyond Sort32 padding, | ||
| 161 | + // so the complete row padding fits NDDMA's B8 padding field. | ||
| 162 | + // Loop order is [inner, axis]: GM = base + inner + axis * innerSize, | ||
| 163 | + // UB = inner * valueAxisElems + axis. Strides and sizes count T elements. | ||
| 164 | + // Only the axis dimension is right-padded; curInnerChunk excludes tail columns. | ||
| 165 | + constexpr uint32_t transposeDims = 2U; | ||
| 166 | + NdDmaLoopInfo<transposeDims> loops{{1U, static_cast<uint64_t>(this->innerSize_)}, | ||
| 167 | + {this->valueAxisElems_, 1U}, | ||
| 168 | + {curInnerChunk, this->axisLen_}, | ||
| 169 | + {0U, 0U}, | ||
| 170 | + {0U, static_cast<uint8_t>(this->valueAxisElems_ - this->axisLen_)}}; | ||
| 171 | + T fill = IsDescend ? static_cast<T>(-INFINITY) : static_cast<T>(NAN); | ||
| 172 | + NdDmaParams<T, transposeDims> params{loops, fill}; | ||
| 173 | + static constexpr NdDmaConfig config; | ||
| 174 | + NdDmaDci(); | ||
| 175 | + DataCopy<T, transposeDims, config>(this->sortInput_, this->inputGm_[baseOffset], params); | ||
| 176 | + event_t ready = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE2_V)); | ||
| 177 | + SetFlag<HardEvent::MTE2_V>(ready); | ||
| 178 | + WaitFlag<HardEvent::MTE2_V>(ready); | ||
| 179 | + return; | ||
| 180 | + } | ||
| 181 | + } | ||
| 151 | uint32_t curBytes = curInnerChunk * sizeof(T); | 182 | uint32_t curBytes = curInnerChunk * sizeof(T); |
| 152 | uint32_t curAlignedBytes = ROUND_UP_AGLIN(curBytes); | 183 | uint32_t curAlignedBytes = ROUND_UP_AGLIN(curBytes); |
| 153 | uint32_t rightPadding = (curAlignedBytes - curBytes) / sizeof(T); | 184 | uint32_t rightPadding = (curAlignedBytes - curBytes) / sizeof(T); |
| @@ -174,6 +205,11 @@ protected: | |||
| 174 | 205 | ||
| 175 | __aicore__ inline void TransposeToSortMajor(uint32_t curInnerChunk) | 206 | __aicore__ inline void TransposeToSortMajor(uint32_t curInnerChunk) |
| 176 | { | 207 | { |
| 208 | + if constexpr (UseMergeSort && !IsBf16Merge) { | ||
| 209 | + if (this->compactMerge_) { | ||
| 210 | + return; | ||
| 211 | + } | ||
| 212 | + } | ||
| 177 | if constexpr (UseMergeSort) { | 213 | if constexpr (UseMergeSort) { |
| 178 | SortT defaultValue = IsDescend ? static_cast<SortT>(-INFINITY) : static_cast<SortT>(NAN); | 214 | SortT defaultValue = IsDescend ? static_cast<SortT>(-INFINITY) : static_cast<SortT>(NAN); |
| 179 | Duplicate(this->sortInput_, defaultValue, curInnerChunk * this->valueAxisElems_); | 215 | Duplicate(this->sortInput_, defaultValue, curInnerChunk * this->valueAxisElems_); |
| @@ -244,13 +280,30 @@ protected: | |||
| 244 | } | 280 | } |
| 245 | } | 281 | } |
| 246 | 282 | ||
| 247 | - __aicore__ inline void SortRows(uint32_t curInnerChunk) | 283 | + __aicore__ inline void SortRows(uint32_t curInnerChunk, LocalTensor<SortT> proposalPing) |
| 248 | { | 284 | { |
| 285 | + LocalTensor<uint32_t> baseIndex; | ||
| 286 | + if constexpr (UseMergeSort) { | ||
| 287 | + // The last row's indices stay unused until its Extract, after every Sort32 has read this sequence. | ||
| 288 | + baseIndex = this->sortedIndex_[(curInnerChunk - 1U) * this->indexAxisElems_]; | ||
| 289 | + ArithProgression<int32_t>(baseIndex.template ReinterpretCast<int32_t>(), 0, 1, this->sortCount_); | ||
| 290 | + } | ||
| 291 | + // Keep one batched flip, but exclude trailing UB storage that Sort32 never reads. | ||
| 292 | + // sortCount_ includes the padding needed by the final Sort32 group. | ||
| 293 | + if constexpr (UseMergeSort && !IsDescend) { | ||
| 294 | + FlipMergeSignBit(this->sortInput_, (curInnerChunk - 1U) * this->valueAxisElems_ + this->sortCount_); | ||
| 295 | + } | ||
| 249 | for (uint32_t inner = 0; inner < curInnerChunk; ++inner) { | 296 | for (uint32_t inner = 0; inner < curInnerChunk; ++inner) { |
| 250 | LocalTensor<SortT> src = this->sortInput_[inner * this->valueAxisElems_]; | 297 | LocalTensor<SortT> src = this->sortInput_[inner * this->valueAxisElems_]; |
| 251 | - LocalTensor<SortT> dst = this->sortedValue_[inner * this->valueAxisElems_]; | 298 | + // Merge proposals are consumed by Extract before the next row reuses the buffer. |
| 299 | + LocalTensor<SortT> dst = this->compactMerge_ ? proposalPing : proposalPing[inner * this->valueAxisElems_]; | ||
| 252 | LocalTensor<uint32_t> dstIndex = this->sortedIndex_[inner * this->indexAxisElems_]; | 300 | LocalTensor<uint32_t> dstIndex = this->sortedIndex_[inner * this->indexAxisElems_]; |
| 253 | - if constexpr (NormalizeSignedZero) { | 301 | + if constexpr (UseMergeSort) { |
| 302 | + LocalTensor<SortT> pong = this->tmp_.template ReinterpretCast<SortT>(); | ||
| 303 | + uint32_t repeatTimes = this->sortCount_ / MergeSortConstants::DEALING_SORT_NUM_ONCE; | ||
| 304 | + bool resultInPing = PingPongMergeSortCommon::SortToProposal(dst, pong, src, baseIndex, repeatTimes); | ||
| 305 | + Extract(src, dstIndex, resultInPing ? dst : pong, repeatTimes); | ||
| 306 | + } else if constexpr (NormalizeSignedZero) { | ||
| 254 | SignedZeroSortCommon::PrepareSignedZeroKeysVec(src, this->sourceIndex_, this->sortCount_); | 307 | SignedZeroSortCommon::PrepareSignedZeroKeysVec(src, this->sourceIndex_, this->sortCount_); |
| 255 | AscendC::Sort<SortT, true, sortConfig_>(dst, dstIndex, src, this->tmp_, this->sortCount_); | 308 | AscendC::Sort<SortT, true, sortConfig_>(dst, dstIndex, src, this->tmp_, this->sortCount_); |
| 256 | SignedZeroSortCommon::RestoreSignedZeroValuesByIndexVec(dst, dstIndex, this->sourceIndex_, | 309 | SignedZeroSortCommon::RestoreSignedZeroValuesByIndexVec(dst, dstIndex, this->sourceIndex_, |
| @@ -259,6 +312,25 @@ protected: | |||
| 259 | AscendC::Sort<SortT, true, sortConfig_>(dst, dstIndex, src, this->tmp_, this->sortCount_); | 312 | AscendC::Sort<SortT, true, sortConfig_>(dst, dstIndex, src, this->tmp_, this->sortCount_); |
| 260 | } | 313 | } |
| 261 | } | 314 | } |
| 315 | + if constexpr (UseMergeSort) { | ||
| 316 | + if constexpr (!IsDescend) { | ||
| 317 | + FlipMergeSignBit(this->sortInput_, (curInnerChunk - 1U) * this->valueAxisElems_ + this->sortCount_); | ||
| 318 | + } | ||
| 319 | + // Extract writes the sorted values back to sortInput_ after its original input is dead, | ||
| 320 | + // so the input buffer is intentionally reused as the value output without extra UB. | ||
| 321 | + this->sortedValue_ = this->sortInput_; | ||
| 322 | + } | ||
| 323 | + } | ||
| 324 | + | ||
| 325 | + __aicore__ inline void FlipMergeSignBit(LocalTensor<SortT> tensor, uint32_t count) | ||
| 326 | + { | ||
| 327 | + if constexpr (std::is_same_v<SortT, float>) { | ||
| 328 | + LocalTensor<int32_t> bits = tensor.template ReinterpretCast<int32_t>(); | ||
| 329 | + Adds(bits, bits, MergeSortConstants::XOR_OP_VALUE_FP, count); | ||
| 330 | + } else if constexpr (std::is_same_v<SortT, half>) { | ||
| 331 | + LocalTensor<int16_t> bits = tensor.template ReinterpretCast<int16_t>(); | ||
| 332 | + Adds(bits, bits, MergeSortConstants::XOR_OP_VALUE_HALF, count); | ||
| 333 | + } | ||
| 262 | } | 334 | } |
| 263 | }; | 335 | }; |
| 264 | 336 | ||
| @@ -0,0 +1,115 @@ | |||
| 1 | +/** | ||
| 2 | + * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | + * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | + * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | + * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | + */ | ||
| 10 | + | ||
| 11 | +#ifndef PING_PONG_MERGE_SORT_H | ||
| 12 | +#define PING_PONG_MERGE_SORT_H | ||
| 13 | + | ||
| 14 | +#include "kernel_operator.h" | ||
| 15 | +#include "merge_sort_constants.h" | ||
| 16 | + | ||
| 17 | +namespace PingPongMergeSortCommon { | ||
| 18 | +using namespace AscendC; | ||
| 19 | + | ||
| 20 | +using MergeSortConstants::DEALING_SORT_NUM_ONCE; | ||
| 21 | +using MergeSortConstants::MERGE_LIST_MAX_NUM; | ||
| 22 | + | ||
| 23 | +// One round of 4-way merge: src has sorted runs of length queueLength, merge every 4 adjacent runs into dst. | ||
| 24 | +// Full groups (4 complete runs) are batched in one MrgSort call; an incomplete final group is handled separately. | ||
| 25 | +// Current arch35 callers keep both proposal buffers in UB (>=8 bytes/element | ||
| 26 | +// each). Even allowing all 256 KiB, totalElements <=16384. Queue lengths are | ||
| 27 | +// 32*4^s; while queueCount > 1 the largest is 8192, so uint16 counts do not | ||
| 28 | +// truncate. Recheck this bound if larger UB or a different proposal layout is used. | ||
| 29 | +template <typename T> | ||
| 30 | +__aicore__ inline void MergeStage(LocalTensor<T> dst, LocalTensor<T> src, uint32_t totalElements, uint32_t queueLength, | ||
| 31 | + uint32_t queueCount) | ||
| 32 | +{ | ||
| 33 | + constexpr uint16_t allListsValid = (1U << MERGE_LIST_MAX_NUM) - 1U; | ||
| 34 | + // Full groups: each group merges 4 runs of queueLength elements. | ||
| 35 | + uint32_t groupSpan = queueLength * MERGE_LIST_MAX_NUM; | ||
| 36 | + uint32_t fullGroupCount = queueCount / MERGE_LIST_MAX_NUM; | ||
| 37 | + uint32_t tailQueueCount = queueCount % MERGE_LIST_MAX_NUM; | ||
| 38 | + // A divisible queue count can still end in a partial run. Leave that final | ||
| 39 | + // group to the tail path so its last element count is described exactly. | ||
| 40 | + if (tailQueueCount == 0U && totalElements != queueCount * queueLength) { | ||
| 41 | + --fullGroupCount; | ||
| 42 | + tailQueueCount = MERGE_LIST_MAX_NUM; | ||
| 43 | + } | ||
| 44 | + if (fullGroupCount > 0U) { | ||
| 45 | + uint16_t elementCounts[MERGE_LIST_MAX_NUM] = { | ||
| 46 | + static_cast<uint16_t>(queueLength), static_cast<uint16_t>(queueLength), static_cast<uint16_t>(queueLength), | ||
| 47 | + static_cast<uint16_t>(queueLength)}; | ||
| 48 | + uint32_t sourceStride = GetSortOffset<T>(queueLength); | ||
| 49 | + MrgSortSrcList<T> sourceList(src, src[sourceStride], src[sourceStride * 2U], src[sourceStride * 3U]); | ||
| 50 | + uint32_t sortedCounts[MERGE_LIST_MAX_NUM]; | ||
| 51 | + // false disables exhausted-list suspension: consume every valid list. | ||
| 52 | + MrgSort<T, false>(dst, sourceList, elementCounts, sortedCounts, allListsValid, | ||
| 53 | + static_cast<int32_t>(fullGroupCount)); | ||
| 54 | + } | ||
| 55 | + | ||
| 56 | + // Tail: 1-4 remaining runs; a fourth run can be shorter than queueLength. | ||
| 57 | + uint32_t tailOffset = fullGroupCount * groupSpan; | ||
| 58 | + uint32_t tailElements = totalElements - tailOffset; | ||
| 59 | + if (tailElements == 0U) { | ||
| 60 | + return; | ||
| 61 | + } | ||
| 62 | + uint32_t proposalTailOffset = GetSortOffset<T>(tailOffset); | ||
| 63 | + if (tailQueueCount == 1U) { | ||
| 64 | + DataCopy(dst[proposalTailOffset], src[proposalTailOffset], GetSortOffset<T>(tailElements)); | ||
| 65 | + return; | ||
| 66 | + } | ||
| 67 | + | ||
| 68 | + uint16_t elementCounts[MERGE_LIST_MAX_NUM] = {0, 0, 0, 0}; | ||
| 69 | + uint32_t sourceOffsets[MERGE_LIST_MAX_NUM] = {tailOffset, tailOffset, tailOffset, tailOffset}; | ||
| 70 | + for (uint32_t list = 0; list < tailQueueCount; ++list) { | ||
| 71 | + uint32_t sourceOffset = tailOffset + list * queueLength; | ||
| 72 | + uint32_t sourceRemaining = totalElements - sourceOffset; | ||
| 73 | + sourceOffsets[list] = sourceOffset; | ||
| 74 | + elementCounts[list] = static_cast<uint16_t>(sourceRemaining < queueLength ? sourceRemaining : queueLength); | ||
| 75 | + } | ||
| 76 | + MrgSortSrcList<T> sourceList(src[GetSortOffset<T>(sourceOffsets[0])], src[GetSortOffset<T>(sourceOffsets[1])], | ||
| 77 | + src[GetSortOffset<T>(sourceOffsets[2])], src[GetSortOffset<T>(sourceOffsets[3])]); | ||
| 78 | + uint32_t sortedCounts[MERGE_LIST_MAX_NUM]; | ||
| 79 | + uint16_t validBit = static_cast<uint16_t>((1U << tailQueueCount) - 1U); | ||
| 80 | + MrgSort<T, false>(dst[proposalTailOffset], sourceList, elementCounts, sortedCounts, validBit, 1); | ||
| 81 | +} | ||
| 82 | + | ||
| 83 | +// Ping-pong merge sort in UB: | ||
| 84 | +// 1. Sort32 produces repeatTimes sorted runs of DEALING_SORT_NUM_ONCE elements into ping. | ||
| 85 | +// 2. Iteratively 4-way merge adjacent runs, alternating src/dst between ping and pong each round. | ||
| 86 | +// After each round, run length grows by 4x until only one sorted run remains. | ||
| 87 | +// Returns true when the final sorted result is in ping, false when in pong. | ||
| 88 | +template <typename T> | ||
| 89 | +__aicore__ inline bool SortToProposal(LocalTensor<T> ping, LocalTensor<T> pong, LocalTensor<T> srcValue, | ||
| 90 | + LocalTensor<uint32_t> baseIndex, uint32_t repeatTimes) | ||
| 91 | +{ | ||
| 92 | + Sort32<T>(ping, srcValue, baseIndex, repeatTimes); | ||
| 93 | + if (repeatTimes <= 1U) { | ||
| 94 | + return true; | ||
| 95 | + } | ||
| 96 | + uint32_t totalElements = repeatTimes * DEALING_SORT_NUM_ONCE; | ||
| 97 | + uint32_t queueLength = DEALING_SORT_NUM_ONCE; | ||
| 98 | + uint32_t queueCount = repeatTimes; | ||
| 99 | + bool sourceIsPing = true; | ||
| 100 | + while (queueCount > 1U) { | ||
| 101 | + if (sourceIsPing) { | ||
| 102 | + MergeStage(pong, ping, totalElements, queueLength, queueCount); | ||
| 103 | + } else { | ||
| 104 | + MergeStage(ping, pong, totalElements, queueLength, queueCount); | ||
| 105 | + } | ||
| 106 | + sourceIsPing = !sourceIsPing; | ||
| 107 | + queueLength *= MERGE_LIST_MAX_NUM; | ||
| 108 | + queueCount = (queueCount + MERGE_LIST_MAX_NUM - 1U) / MERGE_LIST_MAX_NUM; | ||
| 109 | + } | ||
| 110 | + return sourceIsPing; | ||
| 111 | +} | ||
| 112 | + | ||
| 113 | +} // namespace PingPongMergeSortCommon | ||
| 114 | + | ||
| 115 | +#endif // PING_PONG_MERGE_SORT_H | ||
| @@ -23,9 +23,106 @@ | |||
| 23 | #include "common/merge_more_core_base.h" | 23 | #include "common/merge_more_core_base.h" |
| 24 | 24 | ||
| 25 | using namespace AscendC; | 25 | using namespace AscendC; |
| 26 | +#include "simt_api/asc_simt.h" | ||
| 27 | +namespace SortMergePath { | ||
| 28 | +constexpr uint32_t WARP_SIZE = 32U; | ||
| 29 | +constexpr uint32_t PARTITION_THREADS = WARP_SIZE * 2U; | ||
| 30 | +constexpr uint32_t PROPOSAL_WORDS = 2U; | ||
| 31 | +constexpr uint32_t FLOAT_SIGN = 0x80000000U; | ||
| 32 | +constexpr uint32_t FLOAT_MAGNITUDE = 0x7fffffffU; | ||
| 33 | +constexpr uint32_t FLOAT_INFINITY = 0x7f800000U; | ||
| 34 | +constexpr uint32_t CANONICAL_NAN = 0x7fc00000U; | ||
| 35 | +constexpr uint32_t OUTPUT_THREADS = 256U; | ||
| 36 | +template <bool Descend> | ||
| 37 | +__simt_callee__ __aicore__ inline uint32_t Key(__gm__ uint32_t* input, int64_t pos) | ||
| 38 | +{ | ||
| 39 | + // Proposals contain the sign-flipped value and its original uint32 index. | ||
| 40 | + // NaNs are canonicalized before sorting; signed zeros compare equal. | ||
| 41 | + uint32_t bits = input[pos * PROPOSAL_WORDS]; | ||
| 42 | + if constexpr (!Descend) { | ||
| 43 | + bits ^= FLOAT_SIGN; | ||
| 44 | + } | ||
| 45 | + uint32_t magnitude = bits & FLOAT_MAGNITUDE; | ||
| 46 | + if (magnitude == 0U) { | ||
| 47 | + return FLOAT_SIGN; | ||
| 48 | + } | ||
| 49 | + return (bits & FLOAT_SIGN) ? ~bits : (bits ^ FLOAT_SIGN); | ||
| 50 | +} | ||
| 51 | +template <bool Descend> | ||
| 52 | +__simt_callee__ __aicore__ inline bool Before(__gm__ uint32_t* input, int64_t a, int64_t b) | ||
| 53 | +{ | ||
| 54 | + uint32_t ka = Key<Descend>(input, a), kb = Key<Descend>(input, b); | ||
| 55 | + return Descend ? ka > kb : ka < kb; | ||
| 56 | +} | ||
| 57 | +template <bool Descend> | ||
| 58 | +__simt_vf__ LAUNCH_BOUND(PARTITION_THREADS) __aicore__ | ||
| 59 | + void Partition(__gm__ uint32_t* input, __ubuf__ int64_t* ranks, int64_t a, int64_t b, int64_t first, int64_t end) | ||
| 60 | +{ | ||
| 61 | + uint32_t lane = threadIdx.x; | ||
| 62 | + if (lane != 0U && lane != WARP_SIZE) { | ||
| 63 | + return; | ||
| 64 | + } | ||
| 65 | + uint32_t query = lane / WARP_SIZE; | ||
| 66 | + int64_t diagonal = query == 0 ? first : end; | ||
| 67 | + int64_t low = diagonal > b ? diagonal - b : 0; | ||
| 68 | + int64_t high = diagonal < a ? diagonal : a; | ||
| 69 | + while (low <= high) { | ||
| 70 | + int64_t i = (low + high) / 2; | ||
| 71 | + int64_t j = diagonal - i; | ||
| 72 | + if (i > 0 && j < b && Before<Descend>(input, a + j, i - 1)) { | ||
| 73 | + high = i - 1; | ||
| 74 | + } else if (j > 0 && i < a && !Before<Descend>(input, a + j - 1, i)) { | ||
| 75 | + low = i + 1; | ||
| 76 | + } else { | ||
| 77 | + ranks[query] = i; | ||
| 78 | + return; | ||
| 79 | + } | ||
| 80 | + } | ||
| 81 | + ranks[query] = low; | ||
| 82 | +} | ||
| 83 | +template <bool Descend, typename Index> | ||
| 84 | +__simt_vf__ LAUNCH_BOUND(OUTPUT_THREADS) __aicore__ | ||
| 85 | + void WriteOutput(__ubuf__ uint32_t* proposals, __gm__ uint32_t* input, __gm__ uint32_t* values, | ||
| 86 | + __gm__ Index* indices, uint32_t count) | ||
| 87 | +{ | ||
| 88 | + for (uint32_t i = threadIdx.x; i < count; i += OUTPUT_THREADS) { | ||
| 89 | + uint32_t bits = proposals[PROPOSAL_WORDS * i]; | ||
| 90 | + uint32_t index = proposals[PROPOSAL_WORDS * i + 1U]; | ||
| 91 | + if constexpr (!Descend) { | ||
| 92 | + bits ^= FLOAT_SIGN; | ||
| 93 | + } | ||
| 94 | + if ((bits & FLOAT_MAGNITUDE) > FLOAT_INFINITY) { | ||
| 95 | + bits = input[index]; | ||
| 96 | + } | ||
| 97 | + values[i] = bits; | ||
| 98 | + indices[i] = static_cast<Index>(index); | ||
| 99 | + } | ||
| 100 | +} | ||
| 101 | +__aicore__ inline void CanonicalizeNan(LocalTensor<float> input, uint32_t count) | ||
| 102 | +{ | ||
| 103 | + constexpr uint32_t vl = Ops::Base::GetVRegSize() / sizeof(uint32_t); | ||
| 104 | + uint16_t repeats = Ops::Base::CeilDiv(count, vl); | ||
| 105 | + __ubuf__ uint32_t* address = (__ubuf__ uint32_t*)input.GetPhyAddr(); | ||
| 106 | + __VEC_SCOPE__ | ||
| 107 | + { | ||
| 108 | + Reg::RegTensor<uint32_t> value, magnitude, magnitudeMask, canonical, result; | ||
| 109 | + Reg::MaskReg full = Reg::CreateMask<uint32_t, Reg::MaskPattern::ALL>(); | ||
| 110 | + Reg::Duplicate(magnitudeMask, FLOAT_MAGNITUDE, full); | ||
| 111 | + Reg::Duplicate(canonical, CANONICAL_NAN, full); | ||
| 112 | + for (uint16_t i = 0; i < repeats; ++i) { | ||
| 113 | + Reg::MaskReg valid = Reg::UpdateMask<uint32_t>(count), nan; | ||
| 114 | + Reg::LoadAlign(value, address + i * vl); | ||
| 115 | + Reg::And(magnitude, value, magnitudeMask, valid); | ||
| 116 | + Reg::Compares<uint32_t, CMPMODE::GT>(nan, magnitude, FLOAT_INFINITY, valid); | ||
| 117 | + Reg::Select(result, canonical, value, nan); | ||
| 118 | + Reg::StoreAlign(address + i * vl, result, valid); | ||
| 119 | + } | ||
| 120 | + } | ||
| 121 | +} | ||
| 122 | + | ||
| 123 | +} // namespace SortMergePath | ||
| 26 | 124 | ||
| 27 | // Import shared constants from MergeSortConstants namespace | 125 | // Import shared constants from MergeSortConstants namespace |
| 28 | -using MergeSortConstants::DEALING_CONCAT_NUM_ONCE; | ||
| 29 | using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE; | 126 | using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE; |
| 30 | using MergeSortConstants::DEALING_SORT_NUM_ONCE; | 127 | using MergeSortConstants::DEALING_SORT_NUM_ONCE; |
| 31 | using MergeSortConstants::MERGE_LIST_MAX_NUM; | 128 | using MergeSortConstants::MERGE_LIST_MAX_NUM; |
| @@ -34,25 +131,91 @@ using MergeSortConstants::UB_BLOCK_BYTES; | |||
| 34 | using MergeSortConstants::XOR_OP_VALUE_FP; | 131 | using MergeSortConstants::XOR_OP_VALUE_FP; |
| 35 | using MergeSortConstants::XOR_OP_VALUE_HALF; | 132 | using MergeSortConstants::XOR_OP_VALUE_HALF; |
| 36 | 133 | ||
| 37 | -template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE> | 134 | +template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool DIRECT_SCHEDULE = false, |
| 135 | + bool USE_MERGE_PATH = false> | ||
| 38 | struct MergeSortBigSize | 136 | struct MergeSortBigSize |
| 39 | - : public MergeMoreCoreCommon::MergeMoreCoreBase<MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>, T, | 137 | + : public MergeMoreCoreCommon::MergeMoreCoreBase< |
| 40 | - CONVERT_TYPE, IS_DESCEND, INDEX_TYPE> { | 138 | + MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, DIRECT_SCHEDULE, USE_MERGE_PATH>, T, CONVERT_TYPE, |
| 41 | - using Base = MergeMoreCoreCommon::MergeMoreCoreBase<MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>, T, | 139 | + IS_DESCEND, INDEX_TYPE> { |
| 42 | - CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>; | 140 | + using Base = MergeMoreCoreCommon::MergeMoreCoreBase< |
| 141 | + MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, DIRECT_SCHEDULE, USE_MERGE_PATH>, T, CONVERT_TYPE, | ||
| 142 | + IS_DESCEND, INDEX_TYPE>; | ||
| 43 | friend Base; | 143 | friend Base; |
| 44 | 144 | ||
| 45 | __aicore__ inline MergeSortBigSize() {} | 145 | __aicore__ inline MergeSortBigSize() {} |
| 46 | __aicore__ inline void Init(GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace, | 146 | __aicore__ inline void Init(GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace, |
| 47 | const SortRegBaseTilingData* tilingData, TPipe* pipe); | 147 | const SortRegBaseTilingData* tilingData, TPipe* pipe); |
| 148 | + __aicore__ inline void Process(); | ||
| 149 | + __aicore__ inline void PrepareRowWorkspace(); | ||
| 150 | + __aicore__ inline void InitSortBuffers(); | ||
| 48 | __aicore__ inline void InitMergeBuffers(); | 151 | __aicore__ inline void InitMergeBuffers(); |
| 49 | __aicore__ inline void ExtractAndCopyOut(); | 152 | __aicore__ inline void ExtractAndCopyOut(); |
| 153 | + __aicore__ inline void CopyMergePathOutput(); | ||
| 50 | __aicore__ inline void OnInputLoaded(LocalTensor<T>, uint32_t) {} | 154 | __aicore__ inline void OnInputLoaded(LocalTensor<T>, uint32_t) {} |
| 51 | - __aicore__ inline void PrepareInputForSort(LocalTensor<T>, uint32_t) {} | 155 | + __aicore__ inline void PrepareInputForSort(LocalTensor<T> input, uint32_t count) |
| 156 | + { | ||
| 157 | + if constexpr (USE_MERGE_PATH) { | ||
| 158 | + // Native proposals must have the same NaN ordering as the co-rank search. | ||
| 159 | + // Output gathers restore the original NaN payloads through their stable indices. | ||
| 160 | + SortMergePath::CanonicalizeNan(input, count); | ||
| 161 | + } | ||
| 162 | + } | ||
| 163 | + | ||
| 164 | + __aicore__ inline void FinalSortAndCopy(bool rowActive) | ||
| 165 | + { | ||
| 166 | + if constexpr (!USE_MERGE_PATH) { | ||
| 167 | + Base::FinalSortAndCopy(rowActive); | ||
| 168 | + return; | ||
| 169 | + } else { | ||
| 170 | + if (this->listNum_ != MergeSortConstants::TWO_WAY_MERGE_LIST_NUM || | ||
| 171 | + this->frontCoreNum_ < MergeSortConstants::TWO_WAY_MERGE_LIST_NUM) { | ||
| 172 | + Base::FinalSortAndCopy(rowActive); | ||
| 173 | + return; | ||
| 174 | + } | ||
| 175 | + if (rowActive) { | ||
| 176 | + // Each core owns disjoint output ranks. Co-ranking chooses equal keys from | ||
| 177 | + // the earlier input run first, preserving stable indices across partitions. | ||
| 178 | + int64_t first = static_cast<int64_t>(this->outputLastDimValue_) * this->rowCoreIdx_ / | ||
| 179 | + this->frontCoreNum_; | ||
| 180 | + int64_t end = static_cast<int64_t>(this->outputLastDimValue_) * (this->rowCoreIdx_ + 1U) / | ||
| 181 | + this->frontCoreNum_; | ||
| 182 | + LocalTensor<int64_t> ranks = this->castIndexQueue_.template AllocTensor<int64_t>(); | ||
| 183 | + asc_vf_call<SortMergePath::Partition<IS_DESCEND>>( | ||
| 184 | + dim3(SortMergePath::PARTITION_THREADS), (__gm__ uint32_t*)this->workspaceInput_.GetPhyAddr(), | ||
| 185 | + (__ubuf__ int64_t*)ranks.GetPhyAddr(), this->currentElements_, this->currentTailElements_, first, | ||
| 186 | + end); | ||
| 187 | + event_t ready = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_S)); | ||
| 188 | + SetFlag<HardEvent::V_S>(ready); | ||
| 189 | + WaitFlag<HardEvent::V_S>(ready); | ||
| 190 | + int64_t startA = ranks.GetValue(0), endA = ranks.GetValue(1); | ||
| 191 | + this->castIndexQueue_.FreeTensor(ranks); | ||
| 192 | + this->ClearCache(); | ||
| 193 | + this->outOffset_ = first; | ||
| 194 | + this->offsets_[0] = GetSortOffset<CONVERT_TYPE>(startA); | ||
| 195 | + this->offsets_[1] = GetSortOffset<CONVERT_TYPE>(this->currentElements_ + first - startA); | ||
| 196 | + this->listRemainElements_[0] = endA - startA; | ||
| 197 | + this->listRemainElements_[1] = end - first - (endA - startA); | ||
| 198 | + this->allRemainElements_ = end - first; | ||
| 199 | + while (this->allRemainElements_ > 0) { | ||
| 200 | + this->CopyInMultiCore(); | ||
| 201 | + this->UpdateMrgParam(); | ||
| 202 | + this->DealingMergeSort(); | ||
| 203 | + this->UpdateSortInfo(); | ||
| 204 | + ExtractAndCopyOut(); | ||
| 205 | + } | ||
| 206 | + this->ClearCache(); | ||
| 207 | + } | ||
| 208 | + SyncAll(); | ||
| 209 | + } | ||
| 210 | + } | ||
| 211 | + | ||
| 212 | + __gm__ CONVERT_TYPE* workspaceBase_ = nullptr; | ||
| 213 | + uint64_t rowWorkspaceElements_ = 0; | ||
| 52 | }; | 214 | }; |
| 53 | 215 | ||
| 54 | -template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE> | 216 | +template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool DIRECT_SCHEDULE, |
| 55 | -__aicore__ inline void MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::Init( | 217 | + bool USE_MERGE_PATH> |
| 218 | +__aicore__ inline void MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, DIRECT_SCHEDULE, USE_MERGE_PATH>::Init( | ||
| 56 | GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace, const SortRegBaseTilingData* tilingData, | 219 | GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace, const SortRegBaseTilingData* tilingData, |
| 57 | TPipe* pipe) | 220 | TPipe* pipe) |
| 58 | { | 221 | { |
| @@ -64,45 +227,105 @@ __aicore__ inline void MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE> | |||
| 64 | if (this->frontCoreNum_ == 0U) { | 227 | if (this->frontCoreNum_ == 0U) { |
| 65 | return; | 228 | return; |
| 66 | } | 229 | } |
| 230 | + if constexpr (DIRECT_SCHEDULE) { | ||
| 231 | + this->syncMergeSortMode_ = false; | ||
| 232 | + } else { | ||
| 233 | + this->rowsPerRound_ = tilingData->unsortedDimParallel == 0U ? | ||
| 234 | + static_cast<uint32_t>(tilingData->unsortedDimNum) : | ||
| 235 | + tilingData->unsortedDimParallel; | ||
| 236 | + this->sortLoopTimes_ = tilingData->sortLoopTimes == 0U ? 1U : tilingData->sortLoopTimes; | ||
| 237 | + this->batchNum_ = tilingData->unsortedDimNum; | ||
| 238 | + this->logicalBlockSize_ = tilingData->keyParams1; | ||
| 239 | + this->syncMergeSortMode_ = this->logicalBlockSize_ > this->numTileData_ && this->numTileData_ > 0U; | ||
| 240 | + if (this->rowsPerRound_ == 0U || this->sortLoopTimes_ == 0U) { | ||
| 241 | + return; | ||
| 242 | + } | ||
| 243 | + } | ||
| 67 | uint32_t sortBufferSize = 8; | 244 | uint32_t sortBufferSize = 8; |
| 68 | - this->rowIdx_ = this->blockIdx_ / this->frontCoreNum_; | 245 | + this->rowGroupIdx_ = this->blockIdx_ / this->frontCoreNum_; |
| 69 | this->rowCoreIdx_ = this->blockIdx_ % this->frontCoreNum_; | 246 | this->rowCoreIdx_ = this->blockIdx_ % this->frontCoreNum_; |
| 70 | - this->rowDataOffset_ = static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(this->outputLastDimValue_); | ||
| 71 | // Per-row workspace stores Sort API sort-struct data. This capacity uses sortBufferSize bytes per | 247 | // Per-row workspace stores Sort API sort-struct data. This capacity uses sortBufferSize bytes per |
| 72 | // original element and UB-block byte alignment; it must cover later GetSortLen-based accesses. | 248 | // original element and UB-block byte alignment; it must cover later GetSortLen-based accesses. |
| 73 | uint64_t rowWorkspaceBytes = ROUND_UP_AGLIN_UINT64(static_cast<uint64_t>(this->outputLastDimValue_) * | 249 | uint64_t rowWorkspaceBytes = ROUND_UP_AGLIN_UINT64(static_cast<uint64_t>(this->outputLastDimValue_) * |
| 74 | sortBufferSize); | 250 | sortBufferSize); |
| 75 | - uint64_t rowWorkspaceElements = rowWorkspaceBytes / sizeof(CONVERT_TYPE); | 251 | + this->rowWorkspaceElements_ = rowWorkspaceBytes / sizeof(CONVERT_TYPE); |
| 76 | - this->rowWorkspaceOffset_ = static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(rowWorkspaceElements) * 2; | ||
| 77 | this->onceMaxElements_ = tilingData->keyParams0 / DEALING_SORT_NUM_ONCE * DEALING_SORT_NUM_ONCE; | 252 | this->onceMaxElements_ = tilingData->keyParams0 / DEALING_SORT_NUM_ONCE * DEALING_SORT_NUM_ONCE; |
| 78 | 253 | ||
| 79 | this->inputValueGm_.SetGlobalBuffer((__gm__ T*)(inputValue)); | 254 | this->inputValueGm_.SetGlobalBuffer((__gm__ T*)(inputValue)); |
| 80 | this->outValueGm_.SetGlobalBuffer((__gm__ T*)(value)); | 255 | this->outValueGm_.SetGlobalBuffer((__gm__ T*)(value)); |
| 81 | this->outIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(indices)); | 256 | this->outIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(indices)); |
| 82 | - this->workspaceGm_[0].SetGlobalBuffer((__gm__ CONVERT_TYPE*)(workSpace) + this->rowWorkspaceOffset_, | 257 | + this->workspaceBase_ = (__gm__ CONVERT_TYPE*)workSpace; |
| 83 | - rowWorkspaceElements); | 258 | + if constexpr (DIRECT_SCHEDULE) { |
| 84 | - this->workspaceGm_[1].SetGlobalBuffer( | 259 | + this->rowIdx_ = this->rowGroupIdx_; |
| 85 | - (__gm__ CONVERT_TYPE*)(workSpace) + this->rowWorkspaceOffset_ + rowWorkspaceElements, rowWorkspaceElements); | 260 | + PrepareRowWorkspace(); |
| 261 | + InitSortBuffers(); | ||
| 262 | + } | ||
| 263 | +} | ||
| 86 | 264 | ||
| 87 | - uint32_t tailNum = this->outputLastDimValue_ - (this->frontCoreNum_ - 1) * this->numTileData_; | 265 | +template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool DIRECT_SCHEDULE, |
| 88 | - uint32_t alignTile = ROUND_UP_AGLIN(tailNum); | 266 | + bool USE_MERGE_PATH> |
| 267 | +__aicore__ inline void | ||
| 268 | +MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, DIRECT_SCHEDULE, USE_MERGE_PATH>::Process() | ||
| 269 | +{ | ||
| 270 | + if constexpr (DIRECT_SCHEDULE) { | ||
| 271 | + this->ProcessDirect(); | ||
| 272 | + } else { | ||
| 273 | + Base::Process(); | ||
| 274 | + } | ||
| 275 | +} | ||
| 276 | + | ||
| 277 | +template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool DIRECT_SCHEDULE, | ||
| 278 | + bool USE_MERGE_PATH> | ||
| 279 | +__aicore__ inline void | ||
| 280 | +MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, DIRECT_SCHEDULE, USE_MERGE_PATH>::PrepareRowWorkspace() | ||
| 281 | +{ | ||
| 282 | + this->rowDataOffset_ = static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(this->outputLastDimValue_); | ||
| 283 | + this->rowWorkspaceOffset_ = static_cast<int64_t>(this->rowIdx_) * | ||
| 284 | + static_cast<int64_t>(this->rowWorkspaceElements_) * 2; | ||
| 285 | + this->workspaceGm_[0].SetGlobalBuffer(this->workspaceBase_ + this->rowWorkspaceOffset_, | ||
| 286 | + this->rowWorkspaceElements_); | ||
| 287 | + this->workspaceGm_[1].SetGlobalBuffer( | ||
| 288 | + this->workspaceBase_ + this->rowWorkspaceOffset_ + this->rowWorkspaceElements_, this->rowWorkspaceElements_); | ||
| 289 | +} | ||
| 290 | + | ||
| 291 | +template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool DIRECT_SCHEDULE, | ||
| 292 | + bool USE_MERGE_PATH> | ||
| 293 | +__aicore__ inline void | ||
| 294 | +MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, DIRECT_SCHEDULE, USE_MERGE_PATH>::InitSortBuffers() | ||
| 295 | +{ | ||
| 296 | + uint32_t sortBufferSize = 8; | ||
| 297 | + uint32_t sortTileNum; | ||
| 298 | + if constexpr (DIRECT_SCHEDULE) { | ||
| 299 | + sortTileNum = this->outputLastDimValue_ - (this->frontCoreNum_ - 1) * this->numTileData_; | ||
| 300 | + } else { | ||
| 301 | + sortTileNum = this->syncMergeSortMode_ ? | ||
| 302 | + this->numTileData_ : | ||
| 303 | + this->outputLastDimValue_ - (this->frontCoreNum_ - 1) * this->numTileData_; | ||
| 304 | + } | ||
| 305 | + uint32_t alignTile = ROUND_UP_AGLIN(sortTileNum); | ||
| 89 | this->pipe_->InitBuffer(this->inputQueue_, MERGE_MORE_BUFFER_NUM, alignTile * sizeof(T)); | 306 | this->pipe_->InitBuffer(this->inputQueue_, MERGE_MORE_BUFFER_NUM, alignTile * sizeof(T)); |
| 90 | 307 | ||
| 91 | this->pipe_->InitBuffer(this->sortedValueUb_, alignTile * sortBufferSize); | 308 | this->pipe_->InitBuffer(this->sortedValueUb_, alignTile * sortBufferSize); |
| 92 | this->pipe_->InitBuffer(this->sortedValueIndexUb_, alignTile * sizeof(uint32_t)); | 309 | this->pipe_->InitBuffer(this->sortedValueIndexUb_, alignTile * sizeof(uint32_t)); |
| 93 | - this->pipe_->InitBuffer(this->concatTempBuf_, alignTile * sortBufferSize); | ||
| 94 | this->pipe_->InitBuffer(this->sortTempBuf_, alignTile * sortBufferSize); | 310 | this->pipe_->InitBuffer(this->sortTempBuf_, alignTile * sortBufferSize); |
| 95 | - this->pipe_->InitBuffer(this->sortedValueLocalCastTbuf_, alignTile * sortBufferSize); | ||
| 96 | } | 311 | } |
| 97 | 312 | ||
| 98 | -template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE> | 313 | +template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool DIRECT_SCHEDULE, |
| 99 | -__aicore__ inline void MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::InitMergeBuffers() | 314 | + bool USE_MERGE_PATH> |
| 315 | +__aicore__ inline void | ||
| 316 | +MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, DIRECT_SCHEDULE, USE_MERGE_PATH>::InitMergeBuffers() | ||
| 100 | { | 317 | { |
| 101 | uint32_t sortBufferSize = 8; | 318 | uint32_t sortBufferSize = 8; |
| 102 | this->pipe_->InitBuffer(this->sortedQueue_, MERGE_MORE_BUFFER_NUM, | 319 | this->pipe_->InitBuffer(this->sortedQueue_, MERGE_MORE_BUFFER_NUM, |
| 103 | MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sortBufferSize); | 320 | MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sortBufferSize); |
| 104 | this->pipe_->InitBuffer(this->copyInQueue_, MERGE_MORE_BUFFER_NUM, | 321 | this->pipe_->InitBuffer(this->copyInQueue_, MERGE_MORE_BUFFER_NUM, |
| 105 | MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sortBufferSize); | 322 | MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sortBufferSize); |
| 323 | + if constexpr (USE_MERGE_PATH) { | ||
| 324 | + // Fused SIMT output only needs proposals. Reuse this queue for the two | ||
| 325 | + // int64 co-ranks, rounded up to one UB block; no cast/output buffers. | ||
| 326 | + this->pipe_->InitBuffer(this->castIndexQueue_, MERGE_MORE_BUFFER_NUM, UB_BLOCK_BYTES); | ||
| 327 | + return; | ||
| 328 | + } | ||
| 106 | this->pipe_->InitBuffer(this->castValueQueue_, MERGE_MORE_BUFFER_NUM, | 329 | this->pipe_->InitBuffer(this->castValueQueue_, MERGE_MORE_BUFFER_NUM, |
| 107 | MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sizeof(CONVERT_TYPE)); | 330 | MERGE_LIST_MAX_NUM * this->onceMaxElements_ * sizeof(CONVERT_TYPE)); |
| 108 | this->pipe_->InitBuffer(this->castIndexQueue_, MERGE_MORE_BUFFER_NUM, | 331 | this->pipe_->InitBuffer(this->castIndexQueue_, MERGE_MORE_BUFFER_NUM, |
| @@ -113,9 +336,34 @@ __aicore__ inline void MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE> | |||
| 113 | } | 336 | } |
| 114 | } | 337 | } |
| 115 | 338 | ||
| 116 | -template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE> | 339 | +template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool DIRECT_SCHEDULE, |
| 117 | -__aicore__ inline void MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE>::ExtractAndCopyOut() | 340 | + bool USE_MERGE_PATH> |
| 341 | +__aicore__ inline void | ||
| 342 | +MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, DIRECT_SCHEDULE, USE_MERGE_PATH>::CopyMergePathOutput() | ||
| 118 | { | 343 | { |
| 344 | + LocalTensor<CONVERT_TYPE> proposals = this->sortedQueue_.template DeQue<CONVERT_TYPE>(); | ||
| 345 | + asc_vf_call<SortMergePath::WriteOutput<IS_DESCEND, INDEX_TYPE>>( | ||
| 346 | + dim3(SortMergePath::OUTPUT_THREADS), (__ubuf__ uint32_t*)proposals.GetPhyAddr(), | ||
| 347 | + (__gm__ uint32_t*)this->inputValueGm_[this->rowDataOffset_].GetPhyAddr(), | ||
| 348 | + (__gm__ uint32_t*)this->outValueGm_[this->rowDataOffset_ + this->outOffset_].GetPhyAddr(), | ||
| 349 | + (__gm__ INDEX_TYPE*)this->outIndexGm_[this->rowDataOffset_ + this->outOffset_].GetPhyAddr(), | ||
| 350 | + static_cast<uint32_t>(this->curLoopSortedNum_)); | ||
| 351 | + event_t done = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_S)); | ||
| 352 | + SetFlag<HardEvent::V_S>(done); | ||
| 353 | + WaitFlag<HardEvent::V_S>(done); | ||
| 354 | + this->sortedQueue_.FreeTensor(proposals); | ||
| 355 | + this->outOffset_ += this->curLoopSortedNum_; | ||
| 356 | +} | ||
| 357 | + | ||
| 358 | +template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool DIRECT_SCHEDULE, | ||
| 359 | + bool USE_MERGE_PATH> | ||
| 360 | +__aicore__ inline void | ||
| 361 | +MergeSortBigSize<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, DIRECT_SCHEDULE, USE_MERGE_PATH>::ExtractAndCopyOut() | ||
| 362 | +{ | ||
| 363 | + if constexpr (USE_MERGE_PATH) { | ||
| 364 | + CopyMergePathOutput(); | ||
| 365 | + return; | ||
| 366 | + } | ||
| 119 | LocalTensor<INDEX_TYPE> ubOutput2; | 367 | LocalTensor<INDEX_TYPE> ubOutput2; |
| 120 | if constexpr (std::is_same<int64_t, INDEX_TYPE>::value) { | 368 | if constexpr (std::is_same<int64_t, INDEX_TYPE>::value) { |
| 121 | ubOutput2 = this->outIndexQueue_.template AllocTensor<INDEX_TYPE>(); | 369 | ubOutput2 = this->outIndexQueue_.template AllocTensor<INDEX_TYPE>(); |
| @@ -12,7 +12,8 @@ | |||
| 12 | * \file sort_merge_intra_core.h | 12 | * \file sort_merge_intra_core.h |
| 13 | * \brief Intra-core block merge sort for fp32 | 13 | * \brief Intra-core block merge sort for fp32 |
| 14 | * \details Each core independently sorts blocks and merges them via 4-way MrgSort. | 14 | * \details Each core independently sorts blocks and merges them via 4-way MrgSort. |
| 15 | - * No cross-core synchronization needed in merge phase. | 15 | + * Two or three blocks remain in UB through the final merge and output. |
| 16 | + * Larger rows use workspace without cross-core merge synchronization. | ||
| 16 | * Handles fp32, N > 4096, blocksPerRow <= 256. | 17 | * Handles fp32, N > 4096, blocksPerRow <= 256. |
| 17 | */ | 18 | */ |
| 18 | 19 | ||
| @@ -24,6 +25,7 @@ | |||
| 24 | #include "op_kernel/math_util.h" | 25 | #include "op_kernel/math_util.h" |
| 25 | #include "op_kernel/platform_util.h" | 26 | #include "op_kernel/platform_util.h" |
| 26 | #include "kernel_tiling/kernel_tiling.h" | 27 | #include "kernel_tiling/kernel_tiling.h" |
| 28 | +#include "sort_tiling_data.h" | ||
| 27 | #include "common/merge_sort_constants.h" | 29 | #include "common/merge_sort_constants.h" |
| 28 | #include "common/merge_intra_core_base.h" | 30 | #include "common/merge_intra_core_base.h" |
| 29 | 31 | ||
| @@ -31,7 +33,6 @@ namespace Sort { | |||
| 31 | using namespace AscendC; | 33 | using namespace AscendC; |
| 32 | 34 | ||
| 33 | // Import shared constants from MergeSortConstants namespace | 35 | // Import shared constants from MergeSortConstants namespace |
| 34 | -using MergeSortConstants::DEALING_CONCAT_NUM_ONCE; | ||
| 35 | using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE; | 36 | using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE; |
| 36 | using MergeSortConstants::DEALING_SORT_NUM_ONCE; | 37 | using MergeSortConstants::DEALING_SORT_NUM_ONCE; |
| 37 | using MergeSortConstants::MERGE_INTRA_BUFFER_NUM; | 38 | using MergeSortConstants::MERGE_INTRA_BUFFER_NUM; |
| @@ -58,6 +59,23 @@ public: | |||
| 58 | __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR indices, GM_ADDR workspace, | 59 | __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR indices, GM_ADDR workspace, |
| 59 | const SortRegBaseTilingData* tilingData, TPipe* pipe); | 60 | const SortRegBaseTilingData* tilingData, TPipe* pipe); |
| 60 | 61 | ||
| 62 | + __aicore__ inline void Process() | ||
| 63 | + { | ||
| 64 | + if (this->blockSortSize_ == 0U) { | ||
| 65 | + return; | ||
| 66 | + } | ||
| 67 | + if (this->blocksPerRow_ < SORT_RESIDENT_MERGE_MIN_BLOCKS || | ||
| 68 | + this->blocksPerRow_ > SORT_RESIDENT_MERGE_MAX_BLOCKS) { | ||
| 69 | + Base::Process(); | ||
| 70 | + return; | ||
| 71 | + } | ||
| 72 | + int64_t first = static_cast<int64_t>(this->blockIdx_) * this->batchPerCore_; | ||
| 73 | + int64_t end = first + this->batchPerCore_ < this->batchNum_ ? first + this->batchPerCore_ : this->batchNum_; | ||
| 74 | + for (int64_t row = first; row < end; ++row) { | ||
| 75 | + ProcessResidentBatch(row); | ||
| 76 | + } | ||
| 77 | + } | ||
| 78 | + | ||
| 61 | protected: | 79 | protected: |
| 62 | __aicore__ inline void InitPhase1Buffers(); | 80 | __aicore__ inline void InitPhase1Buffers(); |
| 63 | __aicore__ inline void InitPhase2Buffers(); | 81 | __aicore__ inline void InitPhase2Buffers(); |
| @@ -67,6 +85,8 @@ protected: | |||
| 67 | uint32_t elemProcessed, uint32_t elemCount); | 85 | uint32_t elemProcessed, uint32_t elemCount); |
| 68 | 86 | ||
| 69 | private: | 87 | private: |
| 88 | + __aicore__ inline void ProcessResidentBatch(int64_t row); | ||
| 89 | + __aicore__ inline void CopyResidentOutputs(int64_t row, LocalTensor<ValueType> proposal); | ||
| 70 | // Phase 3 only | 90 | // Phase 3 only |
| 71 | __aicore__ inline void ExtractAndCopyOut(int64_t batchIdx, uint32_t resultRegion); | 91 | __aicore__ inline void ExtractAndCopyOut(int64_t batchIdx, uint32_t resultRegion); |
| 72 | }; | 92 | }; |
| @@ -102,7 +122,6 @@ __aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::Init | |||
| 102 | this->batchSortLen_ = AscendC::GetSortLen<ValueType>(this->alignNum_); | 122 | this->batchSortLen_ = AscendC::GetSortLen<ValueType>(this->alignNum_); |
| 103 | this->sortBufferSize_ = this->blockSortLen_ * sizeof(ValueType); | 123 | this->sortBufferSize_ = this->blockSortLen_ * sizeof(ValueType); |
| 104 | this->sortRepeatTimes_ = this->blockSortSize_ / DEALING_SORT_NUM_ONCE; | 124 | this->sortRepeatTimes_ = this->blockSortSize_ / DEALING_SORT_NUM_ONCE; |
| 105 | - this->concatRepeatTimes_ = this->blockSortSize_ / DEALING_CONCAT_NUM_ONCE; | ||
| 106 | this->lastBlockSize_ = static_cast<uint32_t>(this->sortAxisNum_ - | 125 | this->lastBlockSize_ = static_cast<uint32_t>(this->sortAxisNum_ - |
| 107 | static_cast<int64_t>(this->blocksPerRow_ - 1) * this->blockSortSize_); | 126 | static_cast<int64_t>(this->blocksPerRow_ - 1) * this->blockSortSize_); |
| 108 | 127 | ||
| @@ -127,7 +146,6 @@ template <typename ValueType, typename IndexType, bool IsDescend> | |||
| 127 | __aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::InitPhase1Buffers() | 146 | __aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::InitPhase1Buffers() |
| 128 | { | 147 | { |
| 129 | this->pipe_->InitBuffer(this->inQueueX_, MERGE_INTRA_BUFFER_NUM, this->blockSortSize_ * sizeof(ValueType)); | 148 | this->pipe_->InitBuffer(this->inQueueX_, MERGE_INTRA_BUFFER_NUM, this->blockSortSize_ * sizeof(ValueType)); |
| 130 | - this->pipe_->InitBuffer(this->concatTmpBuf_, this->sortBufferSize_); | ||
| 131 | this->pipe_->InitBuffer(this->sortTmpBuf_, this->sortBufferSize_); | 149 | this->pipe_->InitBuffer(this->sortTmpBuf_, this->sortBufferSize_); |
| 132 | this->pipe_->InitBuffer(this->sortedOutQueue_, MERGE_INTRA_BUFFER_NUM, this->sortBufferSize_); | 150 | this->pipe_->InitBuffer(this->sortedOutQueue_, MERGE_INTRA_BUFFER_NUM, this->sortBufferSize_); |
| 133 | this->pipe_->InitBuffer(this->indexTmpBuf_, this->blockSortSize_ * sizeof(uint32_t)); | 151 | this->pipe_->InitBuffer(this->indexTmpBuf_, this->blockSortSize_ * sizeof(uint32_t)); |
| @@ -227,6 +245,85 @@ __aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::Extr | |||
| 227 | this->outValueQueue_.FreeTensor(valueLocal); | 245 | this->outValueQueue_.FreeTensor(valueLocal); |
| 228 | } | 246 | } |
| 229 | 247 | ||
| 248 | +template <typename ValueType, typename IndexType, bool IsDescend> | ||
| 249 | +__aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::ProcessResidentBatch(int64_t row) | ||
| 250 | +{ | ||
| 251 | + // Two proposal arrays plus one block each of input, temporary proposals and source indices | ||
| 252 | + // use (16 * blocksPerRow + 16) bytes per block element, at most the existing 64-byte budget. | ||
| 253 | + this->pipe_->InitBuffer(this->inQueueX_, 1, this->blockSortSize_ * sizeof(ValueType)); | ||
| 254 | + this->pipe_->InitBuffer(this->sortTmpBuf_, this->sortBufferSize_); | ||
| 255 | + this->pipe_->InitBuffer(this->indexTmpBuf_, this->blockSortSize_ * sizeof(uint32_t)); | ||
| 256 | + this->pipe_->InitBuffer(this->sortedOutQueue_, 1, this->blocksPerRow_ * this->sortBufferSize_); | ||
| 257 | + this->pipe_->InitBuffer(this->mergeOutQueue_, 1, this->blocksPerRow_ * this->sortBufferSize_); | ||
| 258 | + LocalTensor<ValueType> sortedBlocks = this->sortedOutQueue_.template AllocTensor<ValueType>(); | ||
| 259 | + uint16_t counts[MERGE_LIST_MAX_NUM] = {0, 0, 0, 0}; | ||
| 260 | + for (uint32_t i = 0; i < this->blocksPerRow_; ++i) { | ||
| 261 | + counts[i] = i + 1U == this->blocksPerRow_ ? this->lastBlockSize_ : this->blockSortSize_; | ||
| 262 | + LocalTensor<ValueType> block = this->inQueueX_.template AllocTensor<ValueType>(); | ||
| 263 | + this->CopyInBlock(this->inputXGm_[row * this->sortAxisNum_], block, i * this->blockSortSize_, counts[i]); | ||
| 264 | + this->inQueueX_.EnQue(block); | ||
| 265 | + block = this->inQueueX_.template DeQue<ValueType>(); | ||
| 266 | + this->SortBlockToStruct(block, sortedBlocks[i * this->blockSortLen_], counts[i], i * this->blockSortSize_); | ||
| 267 | + this->inQueueX_.FreeTensor(block); | ||
| 268 | + } | ||
| 269 | + LocalTensor<ValueType> proposal = this->mergeOutQueue_.template AllocTensor<ValueType>(); | ||
| 270 | + LocalTensor<ValueType> lists[MERGE_LIST_MAX_NUM]; | ||
| 271 | + for (uint32_t i = 0; i < MERGE_LIST_MAX_NUM; ++i) { | ||
| 272 | + lists[i] = sortedBlocks[i < this->blocksPerRow_ ? i * this->blockSortLen_ : 0U]; | ||
| 273 | + } | ||
| 274 | + MrgSortSrcList<ValueType> sources(lists[0], lists[1], lists[2], lists[3]); | ||
| 275 | + uint32_t consumed[MERGE_LIST_MAX_NUM]; | ||
| 276 | + MrgSort<ValueType, false>(proposal, sources, counts, consumed, | ||
| 277 | + static_cast<uint16_t>((1U << this->blocksPerRow_) - 1U), 1); | ||
| 278 | + CopyResidentOutputs(row, proposal); | ||
| 279 | + this->sortedOutQueue_.FreeTensor(sortedBlocks); | ||
| 280 | + this->mergeOutQueue_.FreeTensor(proposal); | ||
| 281 | + event_t flushed = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE3_MTE2)); | ||
| 282 | + SetFlag<HardEvent::MTE3_MTE2>(flushed); | ||
| 283 | + WaitFlag<HardEvent::MTE3_MTE2>(flushed); | ||
| 284 | + this->pipe_->Reset(); | ||
| 285 | +} | ||
| 286 | + | ||
| 287 | +template <typename ValueType, typename IndexType, bool IsDescend> | ||
| 288 | +__aicore__ inline void SortMergeIntraCore<ValueType, IndexType, IsDescend>::CopyResidentOutputs( | ||
| 289 | + int64_t row, LocalTensor<ValueType> proposal) | ||
| 290 | +{ | ||
| 291 | + // Sorting is complete: reuse input, source-index and temporary-proposal buffers for outputs. | ||
| 292 | + LocalTensor<ValueType> values = this->inQueueX_.template AllocTensor<ValueType>(); | ||
| 293 | + LocalTensor<uint32_t> indices = this->indexTmpBuf_.template Get<uint32_t>(); | ||
| 294 | + LocalTensor<int64_t> wideIndices = this->sortTmpBuf_.template Get<int64_t>(); | ||
| 295 | + for (uint32_t offset = 0; offset < this->sortAxisNum_; offset += this->blockSortSize_) { | ||
| 296 | + uint32_t count = this->sortAxisNum_ - offset < this->blockSortSize_ ? this->sortAxisNum_ - offset : | ||
| 297 | + this->blockSortSize_; | ||
| 298 | + Extract(values, indices, proposal[GetSortLen<ValueType>(offset)], | ||
| 299 | + Ops::Base::CeilDiv(count, DEALING_EXTRACT_NUM_ONCE)); | ||
| 300 | + if constexpr (!IsDescend) { | ||
| 301 | + Adds(values.template ReinterpretCast<int32_t>(), values.template ReinterpretCast<int32_t>(), | ||
| 302 | + static_cast<int32_t>(MergeSortConstants::XOR_OP_VALUE_FP), count); | ||
| 303 | + } | ||
| 304 | + if constexpr (IsSameType<int64_t, IndexType>::value) { | ||
| 305 | + Cast(wideIndices, indices.template ReinterpretCast<int32_t>(), RoundMode::CAST_NONE, | ||
| 306 | + Ops::Base::CeilAlign(count, static_cast<uint32_t>(Ops::Base::GetUbBlockSize() / sizeof(int64_t)))); | ||
| 307 | + } | ||
| 308 | + event_t ready = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_MTE3)); | ||
| 309 | + SetFlag<HardEvent::V_MTE3>(ready); | ||
| 310 | + WaitFlag<HardEvent::V_MTE3>(ready); | ||
| 311 | + DataCopyExtParams copy{1, static_cast<uint32_t>(count * sizeof(ValueType)), 0, 0, 0}; | ||
| 312 | + DataCopyPad(this->outValueGm_[row * this->sortAxisNum_ + offset], values, copy); | ||
| 313 | + copy.blockLen = count * sizeof(IndexType); | ||
| 314 | + if constexpr (IsSameType<int64_t, IndexType>::value) { | ||
| 315 | + DataCopyPad(this->outIdxGm_[row * this->sortAxisNum_ + offset], wideIndices, copy); | ||
| 316 | + } else { | ||
| 317 | + DataCopyPad(this->outIdxGm_[row * this->sortAxisNum_ + offset], indices.template ReinterpretCast<int32_t>(), | ||
| 318 | + copy); | ||
| 319 | + } | ||
| 320 | + event_t reused = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE3_V)); | ||
| 321 | + SetFlag<HardEvent::MTE3_V>(reused); | ||
| 322 | + WaitFlag<HardEvent::MTE3_V>(reused); | ||
| 323 | + } | ||
| 324 | + this->inQueueX_.FreeTensor(values); | ||
| 325 | +} | ||
| 326 | + | ||
| 230 | } // namespace Sort | 327 | } // namespace Sort |
| 231 | 328 | ||
| 232 | #endif // SORT_MERGE_INTRA_CORE_H | 329 | #endif // SORT_MERGE_INTRA_CORE_H |
| @@ -8,346 +8,408 @@ | |||
| 8 | * See LICENSE in the root of the software repository for the full text of the License. | 8 | * See LICENSE in the root of the software repository for the full text of the License. |
| 9 | */ | 9 | */ |
| 10 | 10 | ||
| 11 | -/* ! | 11 | +#ifndef SORT_MERGE_SORT_H |
| 12 | - * \file sort_merge_sort.h | 12 | +#define SORT_MERGE_SORT_H |
| 13 | - * \brief | ||
| 14 | - */ | ||
| 15 | - | ||
| 16 | -#ifndef MERGE_SORT_H | ||
| 17 | -#define MERGE_SORT_H | ||
| 18 | 13 | ||
| 19 | #include <cmath> | 14 | #include <cmath> |
| 20 | #include "kernel_operator.h" | 15 | #include "kernel_operator.h" |
| 21 | #include "op_kernel/platform_util.h" | 16 | #include "op_kernel/platform_util.h" |
| 22 | #include "kernel_tiling/kernel_tiling.h" | 17 | #include "kernel_tiling/kernel_tiling.h" |
| 23 | #include "common/merge_sort_constants.h" | 18 | #include "common/merge_sort_constants.h" |
| 19 | +#include "common/ping_pong_merge_sort.h" | ||
| 24 | 20 | ||
| 25 | namespace Sort { | 21 | namespace Sort { |
| 26 | using namespace AscendC; | 22 | using namespace AscendC; |
| 27 | 23 | ||
| 24 | +using MergeSortConstants::DEALING_SORT_NUM_ONCE; | ||
| 25 | +using MergeSortConstants::MERGE_LIST_MAX_NUM; | ||
| 28 | using MergeSortConstants::XOR_OP_VALUE_FP; | 26 | using MergeSortConstants::XOR_OP_VALUE_FP; |
| 29 | using MergeSortConstants::XOR_OP_VALUE_HALF; | 27 | using MergeSortConstants::XOR_OP_VALUE_HALF; |
| 30 | 28 | ||
| 31 | -const uint32_t UB_AGLIN_VALUE = Ops::Base::GetUbBlockSize(); | 29 | +constexpr uint32_t SORT_PING_PONG_UB_BLOCK_BYTES = Ops::Base::GetUbBlockSize(); |
| 32 | -const uint32_t CONCAT_AGLIN_VALUE = 16; | 30 | +constexpr uint32_t SORT_PROPOSAL_BYTES = 8; |
| 33 | -// T1输入x dtype T2输出Idx dtype UT无符号的数据类型 | 31 | + |
| 34 | -template <typename T1, typename T2, typename CONVERT_TYPE, uint64_t isDescend, uint64_t isSort32SmallAxis = 0> | 32 | +template <typename T1, typename T2, typename CONVERT_TYPE, uint64_t IsDescend> |
| 35 | class MergeSort { | 33 | class MergeSort { |
| 36 | public: | 34 | public: |
| 37 | - __aicore__ inline MergeSort(){}; | 35 | + __aicore__ inline MergeSort() {} |
| 38 | __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace, | 36 | __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace, |
| 39 | const SortRegBaseTilingData* __restrict tilingData, TPipe* pipe); | 37 | const SortRegBaseTilingData* __restrict tilingData, TPipe* pipe); |
| 40 | __aicore__ inline void Process(); | 38 | __aicore__ inline void Process(); |
| 41 | 39 | ||
| 42 | private: | 40 | private: |
| 43 | - __aicore__ inline void ProcessSingleBlockSort(GlobalTensor<T1> inputX, int32_t sortLoopRound); | 41 | + __aicore__ inline void ParseTilingData(); |
| 44 | - __aicore__ inline void ParserTilingData(); | 42 | + __aicore__ inline uint32_t GetRowCount(uint32_t sortLoopRound) const; |
| 45 | - __aicore__ inline void CopyDataIn(GlobalTensor<T1> inputX, uint64_t tileOffset, uint32_t currTileSize, | 43 | + __aicore__ inline void ProcessCurrentBlock(LocalTensor<T1> xLocal, uint32_t rowCount); |
| 46 | - uint32_t oneCoreRowNum); | 44 | + __aicore__ inline void CopyOutCurrentBlock(uint32_t sortLoopRound, uint32_t rowCount); |
| 47 | - __aicore__ inline void VbsMergeSortBf16(LocalTensor<bfloat16_t> xLocal, LocalTensor<T1> sortedValueLocal, | 45 | + __aicore__ inline void ProcessSingleBlock(GlobalTensor<T1> inputX, int32_t sortLoopRound); |
| 48 | - LocalTensor<uint32_t> sortedValueIndexLocal, uint32_t currTileSize, | 46 | + __aicore__ inline void CopyIn(GlobalTensor<T1> inputX, uint64_t tileOffset, uint32_t rowCount); |
| 49 | - uint32_t nowCoreRealRowNum); | 47 | + __aicore__ inline void CopyOut(uint64_t gmOffset, uint64_t tileOffset, uint32_t rowCount); |
| 50 | - __aicore__ inline void VbsMergeSort(LocalTensor<T1> xLocal, LocalTensor<T1> sortedValueLocal, | 48 | + __aicore__ inline void SortRows(LocalTensor<T1> xLocal, LocalTensor<T1> valueLocal, |
| 51 | - LocalTensor<uint32_t> sortedValueIndexLocal, uint32_t currTileSize, | 49 | + LocalTensor<uint32_t> indexLocal, uint32_t rowCount); |
| 52 | - uint32_t nowCoreRealRowNum); | 50 | + __aicore__ inline void SortBatchedRows(LocalTensor<T1> xLocal, LocalTensor<T1> valueLocal, |
| 53 | - __aicore__ inline void CopyValue2Gm(uint64_t gmOffset, uint64_t tileOffset, uint32_t outputLastDimValue, | 51 | + LocalTensor<uint32_t> indexLocal, uint32_t rowCount); |
| 54 | - uint32_t oneCoreRowNum); | 52 | + __aicore__ inline void SortRowsBf16(LocalTensor<bfloat16_t> xLocal, LocalTensor<T1> valueLocal, |
| 55 | - __aicore__ inline void flipSignBit(LocalTensor<CONVERT_TYPE> xLocal, uint32_t offsetOneRow, uint32_t aglinTileSize); | 53 | + LocalTensor<uint32_t> indexLocal, uint32_t rowCount); |
| 54 | + __aicore__ inline void SortOneRow(LocalTensor<CONVERT_TYPE> dstValue, LocalTensor<uint32_t> dstIndex, | ||
| 55 | + LocalTensor<CONVERT_TYPE> srcValue, uint32_t repeatTimes); | ||
| 56 | + __aicore__ inline void FlipSignBit(LocalTensor<CONVERT_TYPE> tensor, uint32_t count); | ||
| 56 | 57 | ||
| 57 | - GlobalTensor<T1> inputXGm_; | 58 | + GlobalTensor<T1> inputGm_; |
| 58 | - GlobalTensor<T1> outValueGm_; | 59 | + GlobalTensor<T1> outputValueGm_; |
| 59 | - GlobalTensor<T2> outIdxGm_; | 60 | + GlobalTensor<T2> outputIndexGm_; |
| 60 | 61 | ||
| 61 | - TPipe* pipe_; | 62 | + TPipe* pipe_ = nullptr; |
| 62 | - // 599040 5120tile | 63 | + const SortRegBaseTilingData* tilingData_ = nullptr; |
| 63 | - const SortRegBaseTilingData* tilingData_; | 64 | + // InitBuffer selects one or two physical buffers; queue events protect reuse across batches. |
| 64 | - TQue<QuePosition::VECIN, 1> inQueueX_; | 65 | + TQue<QuePosition::VECIN, 2> inputQueue_; |
| 65 | - TBuf<TPosition::VECCALC> tmpUb_; | 66 | + TQue<QuePosition::VECOUT, 2> outputIndexQueue_; |
| 66 | - TQue<QuePosition::VECOUT, 1> outIdxQueue_; | 67 | + TQue<QuePosition::VECOUT, 2> outputValueQueue_; |
| 67 | - TQue<QuePosition::VECOUT, 1> outValueQueue_; | 68 | + TBuf<TPosition::VECCALC> indexBuffer_; |
| 68 | - // merg sort | 69 | + TBuf<TPosition::VECCALC> proposalPingBuffer_; |
| 69 | - TBuf<TPosition::VECCALC> contCatTmpTbuf_; | 70 | + TBuf<TPosition::VECCALC> proposalPongBuffer_; |
| 70 | - TBuf<TPosition::VECCALC> indeXLocalTbuf_; | 71 | + TBuf<TPosition::VECCALC> castInputBuffer_; |
| 71 | - TBuf<TPosition::VECCALC> sortedTmpLocalTbuf_; | 72 | + |
| 72 | - TBuf<TPosition::VECCALC> sortedLocalResTbuf_; | 73 | + LocalTensor<uint32_t> baseIndex_; |
| 73 | - TBuf<TPosition::VECCALC> xLocalCastTbuf_; | ||
| 74 | - TBuf<TPosition::VECCALC> sortedValueLocalCastTbuf_; | ||
| 75 | - LocalTensor<uint32_t> indexLocal_; | ||
| 76 | uint32_t blockIdx_ = 0; | 74 | uint32_t blockIdx_ = 0; |
| 77 | - uint32_t oneCoreRowNum_ = 0; | 75 | + uint32_t rowsPerBatch_ = 0; |
| 78 | - uint32_t tmpUbSize_ = 0; | 76 | + uint32_t axisLength_ = 0; |
| 79 | - uint32_t numTileData_ = 0; | 77 | + int64_t totalRows_ = 0; |
| 80 | - int64_t unsortedDimNum_ = 0; | 78 | + uint32_t parallelRows_ = 0; |
| 81 | - uint32_t unsortedDimParallel_ = 0; | 79 | + uint32_t loopCount_ = 0; |
| 82 | - uint32_t sortLoopTimes_ = 0; | 80 | + uint32_t alignedAxis_ = 0; |
| 83 | - int64_t outputLastDimValue_ = 0; | 81 | + bool batchMerge_ = false; |
| 84 | - uint32_t alignSize_ = 0; | ||
| 85 | }; | 82 | }; |
| 86 | 83 | ||
| 87 | -template <typename T1, typename T2, typename CONVERT_TYPE, uint64_t isDescend, uint64_t isSort32SmallAxis> | 84 | +template <typename T1, typename T2, typename CONVERT_TYPE, uint64_t IsDescend> |
| 88 | -__aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, isDescend, isSort32SmallAxis>::Init( | 85 | +__aicore__ inline void MergeSort<T1, T2, CONVERT_TYPE, IsDescend>::Init( |
| 89 | GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace, const SortRegBaseTilingData* __restrict tilingData, | 86 | GM_ADDR x, GM_ADDR value, GM_ADDR sortIndex, GM_ADDR workspace, const SortRegBaseTilingData* __restrict tilingData, |
| 90 | TPipe* pipe) | 87 | TPipe* pipe) |
| 91 | { | 88 | { |
| 89 | + (void)workspace; | ||
| 92 | blockIdx_ = GetBlockIdx(); | 90 | blockIdx_ = GetBlockIdx(); |
| 93 | pipe_ = pipe; | 91 | pipe_ = pipe; |
| 94 | tilingData_ = tilingData; | 92 | tilingData_ = tilingData; |
| 95 | - ParserTilingData(); | 93 | + ParseTilingData(); |
| 96 | - inputXGm_.SetGlobalBuffer((__gm__ T1*)x); | 94 | + |
| 97 | - outValueGm_.SetGlobalBuffer((__gm__ T1*)value); | 95 | + inputGm_.SetGlobalBuffer((__gm__ T1*)x); |
| 98 | - outIdxGm_.SetGlobalBuffer((__gm__ T2*)sortIndex); | 96 | + outputValueGm_.SetGlobalBuffer((__gm__ T1*)value); |
| 97 | + outputIndexGm_.SetGlobalBuffer((__gm__ T2*)sortIndex); | ||
| 99 | 98 | ||
| @@ -104,6 +104,8 @@ public: | |||
| 104 | __aicore__ inline void ParseTilingData(); | 104 | __aicore__ inline void ParseTilingData(); |
| 105 | __aicore__ inline void BuildOutputs(uint32_t curInnerChunk); | 105 | __aicore__ inline void BuildOutputs(uint32_t curInnerChunk); |
| 106 | __aicore__ inline void GatherOutputValues(uint32_t curInnerChunk); | 106 | __aicore__ inline void GatherOutputValues(uint32_t curInnerChunk); |
| 107 | + __aicore__ inline void BuildOutputIndices(uint32_t curInnerChunk); | ||
| 108 | + __aicore__ inline void BuildOutputIndicesVector(uint32_t curInnerChunk); | ||
| 107 | __aicore__ inline void StoreTileOutput(int64_t baseOffset, uint32_t curInnerChunk); | 109 | __aicore__ inline void StoreTileOutput(int64_t baseOffset, uint32_t curInnerChunk); |
| 108 | __aicore__ inline void StoreSingleInnerTile(int64_t baseOffset); | 110 | __aicore__ inline void StoreSingleInnerTile(int64_t baseOffset); |
| 109 | __aicore__ inline void StoreConvertedTile(int64_t baseOffset, uint32_t curInnerChunk); | 111 | __aicore__ inline void StoreConvertedTile(int64_t baseOffset, uint32_t curInnerChunk); |
| @@ -157,11 +159,15 @@ __aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort> | |||
| 157 | } | 159 | } |
| 158 | uint32_t inputCastBytes = IsBf16Merge_ ? this->innerChunk_ * this->inputValueAxisBytes_ : 0U; | 160 | uint32_t inputCastBytes = IsBf16Merge_ ? this->innerChunk_ * this->inputValueAxisBytes_ : 0U; |
| 159 | uint32_t sortValueBytes = this->innerChunk_ * this->valueAxisBytes_; | 161 | uint32_t sortValueBytes = this->innerChunk_ * this->valueAxisBytes_; |
| 162 | + // The native proposal remains value/index packed into eight bytes. Only one | ||
| 163 | + // row of proposals is live; extracted values stay in compact dtype-sized rows. | ||
| 164 | + constexpr uint32_t proposalElemBytes = sizeof(float) + sizeof(uint32_t); | ||
| 165 | + uint32_t proposalBytes = this->compactMerge_ ? this->sortCount_ * proposalElemBytes : sortValueBytes; | ||
| 160 | uint32_t sortedIndexBytes = this->innerChunk_ * this->indexAxisBytes_; | 166 | uint32_t sortedIndexBytes = this->innerChunk_ * this->indexAxisBytes_; |
| 161 | uint32_t outputIndexBytes = this->axisLen_ * outputIndexRowBytes_; | 167 | uint32_t outputIndexBytes = this->axisLen_ * outputIndexRowBytes_; |
| 162 | useSharedUb_ = this->tilingData_->keyParams4 == 1U; | 168 | useSharedUb_ = this->tilingData_->keyParams4 == 1U; |
| 163 | if (useSharedUb_) { | 169 | if (useSharedUb_) { |
| 164 | - uint32_t sortPrefixBytes = inputTileBytes + inputCastBytes + sortValueBytes * 2U + this->tmpUbSize_; | 170 | + uint32_t sortPrefixBytes = inputTileBytes + inputCastBytes + sortValueBytes + proposalBytes + this->tmpUbSize_; |
| 165 | uint32_t sortPhaseBytes = sortPrefixBytes + sortedIndexBytes; | 171 | uint32_t sortPhaseBytes = sortPrefixBytes + sortedIndexBytes; |
| 166 | uint32_t indexPhaseBytes = outputIndexBytes + sortedIndexBytes; | 172 | uint32_t indexPhaseBytes = outputIndexBytes + sortedIndexBytes; |
| 167 | uint32_t sharedUbSize = sortPhaseBytes > indexPhaseBytes ? sortPhaseBytes : indexPhaseBytes; | 173 | uint32_t sharedUbSize = sortPhaseBytes > indexPhaseBytes ? sortPhaseBytes : indexPhaseBytes; |
| @@ -177,7 +183,7 @@ __aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort> | |||
| 177 | this->sortInput_ = sharedUb[offset].template ReinterpretCast<SortT_>(); | 183 | this->sortInput_ = sharedUb[offset].template ReinterpretCast<SortT_>(); |
| 178 | offset += sortValueBytes; | 184 | offset += sortValueBytes; |
| 179 | this->sortedValue_ = sharedUb[offset].template ReinterpretCast<SortT_>(); | 185 | this->sortedValue_ = sharedUb[offset].template ReinterpretCast<SortT_>(); |
| 180 | - offset += sortValueBytes; | 186 | + offset += proposalBytes; |
| 181 | this->tmp_ = sharedUb[offset].template ReinterpretCast<uint8_t>(); | 187 | this->tmp_ = sharedUb[offset].template ReinterpretCast<uint8_t>(); |
| 182 | this->sortedIndex_ = sharedUb[sharedUbSize - sortedIndexBytes].template ReinterpretCast<uint32_t>(); | 188 | this->sortedIndex_ = sharedUb[sharedUbSize - sortedIndexBytes].template ReinterpretCast<uint32_t>(); |
| 183 | outputIndex_ = sharedUb.template ReinterpretCast<OutIdxT>(); | 189 | outputIndex_ = sharedUb.template ReinterpretCast<OutIdxT>(); |
| @@ -188,7 +194,7 @@ __aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort> | |||
| 188 | this->inputCast_ = this->inputCastBuf_.template Get<T>(); | 194 | this->inputCast_ = this->inputCastBuf_.template Get<T>(); |
| 189 | } | 195 | } |
| 190 | this->pipe_->InitBuffer(this->sortInputBuf_, sortValueBytes); | 196 | this->pipe_->InitBuffer(this->sortInputBuf_, sortValueBytes); |
| 191 | - this->pipe_->InitBuffer(this->sortedValueBuf_, sortValueBytes); | 197 | + this->pipe_->InitBuffer(this->sortedValueBuf_, proposalBytes); |
| 192 | this->pipe_->InitBuffer(this->sortedIndexBuf_, sortedIndexBytes); | 198 | this->pipe_->InitBuffer(this->sortedIndexBuf_, sortedIndexBytes); |
| 193 | if (this->innerChunk_ > 1) { | 199 | if (this->innerChunk_ > 1) { |
| 194 | this->pipe_->InitBuffer(this->outputIndexBuf_, outputIndexBytes); | 200 | this->pipe_->InitBuffer(this->outputIndexBuf_, outputIndexBytes); |
| @@ -210,6 +216,7 @@ template <typename T, typename OutIdxT, bool IsDescend, bool UseMergeSort> | |||
| 210 | __aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort>::ParseTilingData() | 216 | __aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort>::ParseTilingData() |
| 211 | { | 217 | { |
| 212 | this->axisLen_ = static_cast<uint32_t>(this->tilingData_->lastAxisNum); | 218 | this->axisLen_ = static_cast<uint32_t>(this->tilingData_->lastAxisNum); |
| 219 | + this->compactMerge_ = UseMergeSort && !IsBf16Merge_ && this->tilingData_->keyParams5 != 0U; | ||
| 213 | this->outerSize_ = this->tilingData_->outerSize; | 220 | this->outerSize_ = this->tilingData_->outerSize; |
| 214 | this->innerSize_ = this->tilingData_->innerSize; | 221 | this->innerSize_ = this->tilingData_->innerSize; |
| 215 | this->innerLoopNum_ = this->tilingData_->innerLoopNum; | 222 | this->innerLoopNum_ = this->tilingData_->innerLoopNum; |
| @@ -238,12 +245,74 @@ template <typename T, typename OutIdxT, bool IsDescend, bool UseMergeSort> | |||
| 238 | __aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort>::BuildOutputs(uint32_t curInnerChunk) | 245 | __aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort>::BuildOutputs(uint32_t curInnerChunk) |
| 239 | { | 246 | { |
| 240 | GatherOutputValues(curInnerChunk); | 247 | GatherOutputValues(curInnerChunk); |
| 248 | + BuildOutputIndices(curInnerChunk); | ||
| 249 | +} | ||
| 250 | + | ||
| 251 | +template <typename T, typename OutIdxT, bool IsDescend, bool UseMergeSort> | ||
| 252 | +__aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort>::BuildOutputIndices( | ||
| 253 | + uint32_t curInnerChunk) | ||
| 254 | +{ | ||
| 255 | + if constexpr (UseMergeSort && (std::is_same_v<T, half> || std::is_same_v<T, float>)) { | ||
| 256 | + // Below one B32 UB block, the per-axis vector loop does not amortize. | ||
| 257 | + constexpr uint32_t minVectorInner = Ops::Base::GetUbBlockSize() / sizeof(uint32_t); | ||
| 258 | + if (curInnerChunk >= minVectorInner) { | ||
| 259 | + BuildOutputIndicesVector(curInnerChunk); | ||
| 260 | + return; | ||
| 261 | + } | ||
| 262 | + } | ||
| 241 | asc_vf_call<BuildOutputIndexTile<OutIdxT>>( | 263 | asc_vf_call<BuildOutputIndexTile<OutIdxT>>( |
| 242 | dim3(SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM), this->axisLen_, curInnerChunk, | 264 | dim3(SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM), this->axisLen_, curInnerChunk, |
| 243 | SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM, this->indexAxisElems_, outputIndexRowElems_, | 265 | SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM, this->indexAxisElems_, outputIndexRowElems_, |
| 244 | (__ubuf__ uint32_t*)this->sortedIndex_.GetPhyAddr(), (__ubuf__ OutIdxT*)outputIndex_.GetPhyAddr()); | 266 | (__ubuf__ uint32_t*)this->sortedIndex_.GetPhyAddr(), (__ubuf__ OutIdxT*)outputIndex_.GetPhyAddr()); |
| 245 | } | 267 | } |
| 246 | 268 | ||
| 269 | +template <typename T, typename OutIdxT, bool IsDescend, bool UseMergeSort> | ||
| 270 | +__aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort>::BuildOutputIndicesVector( | ||
| 271 | + uint32_t curInnerChunk) | ||
| 272 | +{ | ||
| 273 | + uint32_t axisLen = this->axisLen_; | ||
| 274 | + uint32_t innerChunk = curInnerChunk; | ||
| 275 | + uint32_t indexAxisElems = this->indexAxisElems_; | ||
| 276 | + uint32_t outputIndexRowElems = outputIndexRowElems_; | ||
| 277 | + __ubuf__ uint32_t* sortedIndex = (__ubuf__ uint32_t*)this->sortedIndex_.GetPhyAddr(); | ||
| 278 | + __ubuf__ OutIdxT* outputIndex = (__ubuf__ OutIdxT*)outputIndex_.GetPhyAddr(); | ||
| 279 | + // Source indices remain live outside the reused output prefix. Interleave | ||
| 280 | + // with zero widens the nonnegative B32 indices to B64 without scalar stores. | ||
| 281 | + constexpr uint32_t lanes = Ops::Base::GetVRegSize() / sizeof(uint32_t); | ||
| 282 | + uint16_t repeats = static_cast<uint16_t>((innerChunk + lanes - 1U) / lanes); | ||
| 283 | + __VEC_SCOPE__ | ||
| 284 | + { | ||
| 285 | + Reg::RegTensor<int32_t> laneIds; | ||
| 286 | + Reg::RegTensor<uint32_t> offsets, indices, zero; | ||
| 287 | + Reg::RegTensor<OutIdxT> lo, hi; | ||
| 288 | + Reg::MaskReg full = Reg::CreateMask<uint32_t>(); | ||
| 289 | + Reg::Arange(laneIds, 0); | ||
| 290 | + Reg::Muls(offsets, (Reg::RegTensor<uint32_t>&)laneIds, indexAxisElems, full); | ||
| 291 | + Reg::Duplicate(zero, 0U); | ||
| 292 | + for (uint16_t repeat = 0; repeat < repeats; ++repeat) { | ||
| 293 | + uint32_t count = innerChunk - repeat * lanes; | ||
| 294 | + Reg::MaskReg valid = Reg::UpdateMask<uint32_t>(count); | ||
| 295 | + // UpdateMask consumes count64: B64 low gets up to 32 elements; | ||
| 296 | + // high gets only the remainder (zero for tails of at most 32). | ||
| 297 | + uint32_t count64 = innerChunk - repeat * lanes; | ||
| 298 | + Reg::MaskReg lowMask = Reg::UpdateMask<OutIdxT>(count64); | ||
| 299 | + Reg::MaskReg highMask = Reg::UpdateMask<OutIdxT>(count64); | ||
| 300 | + for (uint16_t axis = 0; axis < axisLen; ++axis) { | ||
| 301 | + Reg::Gather(indices, sortedIndex + repeat * lanes * indexAxisElems + axis, offsets, valid); | ||
| 302 | + if constexpr (sizeof(OutIdxT) == sizeof(uint64_t)) { | ||
| 303 | + Reg::Interleave((Reg::RegTensor<uint32_t>&)lo, (Reg::RegTensor<uint32_t>&)hi, indices, zero); | ||
| 304 | + Reg::StoreAlign(outputIndex + axis * outputIndexRowElems + repeat * lanes, lo, lowMask); | ||
| 305 | + Reg::StoreAlign(outputIndex + axis * outputIndexRowElems + repeat * lanes + lanes / 2U, hi, | ||
| 306 | + highMask); | ||
| 307 | + } else { | ||
| 308 | + Reg::StoreAlign(outputIndex + axis * outputIndexRowElems + repeat * lanes, | ||
| 309 | + (Reg::RegTensor<OutIdxT>&)indices, valid); | ||
| 310 | + } | ||
| 311 | + } | ||
| 312 | + } | ||
| 313 | + } | ||
| 314 | +} | ||
| 315 | + | ||
| 247 | template <typename T, typename OutIdxT, bool IsDescend, bool UseMergeSort> | 316 | template <typename T, typename OutIdxT, bool IsDescend, bool UseMergeSort> |
| 248 | __aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort>::GatherOutputValues( | 317 | __aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort>::GatherOutputValues( |
| 249 | uint32_t curInnerChunk) | 318 | uint32_t curInnerChunk) |
| @@ -309,10 +378,7 @@ __aicore__ inline void SortNonLastSmallAxis<T, OutIdxT, IsDescend, UseMergeSort> | |||
| 309 | event_t eventIdMte3ToV = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE3_V)); | 378 | event_t eventIdMte3ToV = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE3_V)); |
| 310 | SetFlag<HardEvent::MTE3_V>(eventIdMte3ToV); | 379 | SetFlag<HardEvent::MTE3_V>(eventIdMte3ToV); |
| 311 | WaitFlag<HardEvent::MTE3_V>(eventIdMte3ToV); | 380 | WaitFlag<HardEvent::MTE3_V>(eventIdMte3ToV); |
| 312 | - asc_vf_call<BuildOutputIndexTile<OutIdxT>>( | 381 | + BuildOutputIndices(curInnerChunk); |
| 313 | - dim3(SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM), this->axisLen_, curInnerChunk, | ||
| 314 | - SmallAxisCommon::NON_LAST_TRANSPOSE_THREAD_NUM, this->indexAxisElems_, outputIndexRowElems_, | ||
| 315 | - (__ubuf__ uint32_t*)this->sortedIndex_.GetPhyAddr(), (__ubuf__ OutIdxT*)outputIndex_.GetPhyAddr()); | ||
| 316 | eventIdVToMte3 = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_MTE3)); | 382 | eventIdVToMte3 = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::V_MTE3)); |
| 317 | SetFlag<HardEvent::V_MTE3>(eventIdVToMte3); | 383 | SetFlag<HardEvent::V_MTE3>(eventIdVToMte3); |
| 318 | WaitFlag<HardEvent::V_MTE3>(eventIdVToMte3); | 384 | WaitFlag<HardEvent::V_MTE3>(eventIdVToMte3); |
| @@ -116,6 +116,9 @@ private: | |||
| 116 | uint32_t innerSize = static_cast<uint32_t>(innerSize_); | 116 | uint32_t innerSize = static_cast<uint32_t>(innerSize_); |
| 117 | uint32_t outerElems = segmentLen_ * innerSize; | 117 | uint32_t outerElems = segmentLen_ * innerSize; |
| 118 | // Map contiguous GM [outer, axis, inner] slices directly to dense UB [outer, inner, axis]. | 118 | // Map contiguous GM [outer, axis, inner] slices directly to dense UB [outer, inner, axis]. |
| 119 | + // Loop order [inner, axis, outer] gives GM = outer*outerElems + axis*innerSize + inner | ||
| 120 | + // and UB = outer*outerElems + inner*segmentLen + axis, all in T elements. | ||
| 121 | + // validOuter limits the final batch; no padded outer slice is read. | ||
| 119 | NdDmaLoopInfo<NDDMA_DIM_NUM> loopInfo{{1, innerSize, static_cast<uint64_t>(outerElems)}, | 122 | NdDmaLoopInfo<NDDMA_DIM_NUM> loopInfo{{1, innerSize, static_cast<uint64_t>(outerElems)}, |
| 120 | {segmentLen_, 1, outerElems}, | 123 | {segmentLen_, 1, outerElems}, |
| 121 | {innerSize, segmentLen_, validOuter}, | 124 | {innerSize, segmentLen_, validOuter}, |
| @@ -16,6 +16,10 @@ | |||
| 16 | #ifndef _SORT_REGBASE_TILING_DATA_H_ | 16 | #ifndef _SORT_REGBASE_TILING_DATA_H_ |
| 17 | #define _SORT_REGBASE_TILING_DATA_H_ | 17 | #define _SORT_REGBASE_TILING_DATA_H_ |
| 18 | 18 | ||
| 19 | +constexpr uint32_t SORT_BATCH_MERGE_MIN_ROWS = 4; | ||
| 20 | +constexpr uint32_t SORT_RESIDENT_MERGE_MIN_BLOCKS = 2U; | ||
| 21 | +constexpr uint32_t SORT_RESIDENT_MERGE_MAX_BLOCKS = 3U; | ||
| 22 | + | ||
| 19 | struct SortRegBaseTilingData { | 23 | struct SortRegBaseTilingData { |
| 20 | uint32_t numTileDataSize; // h轴ub一次处理个数 | 24 | uint32_t numTileDataSize; // h轴ub一次处理个数 |
| 21 | uint32_t unsortedDimParallel; // b轴使用的核数 | 25 | uint32_t unsortedDimParallel; // b轴使用的核数 |
| @@ -30,6 +34,7 @@ struct SortRegBaseTilingData { | |||
| 30 | // radix: 清零 excusiveBinsGmWk_ 的核 | 34 | // radix: 清零 excusiveBinsGmWk_ 的核 |
| 31 | // radix_one_core: y2OutQue 需要的 ub 大小 | 35 | // radix_one_core: y2OutQue 需要的 ub 大小 |
| 32 | // merge: xQue ub 大小 | 36 | // merge: xQue ub 大小 |
| 37 | + // merge more-core(sch3): sort_sync_merge_sort 逻辑块大小，0 表示普通 more-core | ||
| 33 | uint32_t keyParams1; | 38 | uint32_t keyParams1; |
| 34 | // radix: 清零的一次 ub 数据量 | 39 | // radix: 清零的一次 ub 数据量 |
| 35 | // radix_one_core: 输出 int64 时，一半的 ub 偏移 | 40 | // radix_one_core: 输出 int64 时，一半的 ub 偏移 |
| @@ -47,6 +52,8 @@ struct SortRegBaseTilingData { | |||
| 47 | // non-last small-axis radix(sch10): phase-shared UB layout flag | 52 | // non-last small-axis radix(sch10): phase-shared UB layout flag |
| 48 | // non-last small-axis two-stage(sch11): outer slices per batch (0 means ordinary mapping) | 53 | // non-last small-axis two-stage(sch11): outer slices per batch (0 means ordinary mapping) |
| 49 | uint32_t keyParams4; | 54 | uint32_t keyParams4; |
| 55 | + // merge_sort(sch0): enable row-batched proposal buffers and merge stages | ||
黄 | |||
| 56 | + // non-last merge(sch9): 0 = ordinary layout, 1 = compact value rows/shared proposals | ||
| 50 | uint32_t keyParams5; // radix：清零chunk大小；intra_core(sch4)：最大归并迭代次数 | 57 | uint32_t keyParams5; // radix：清零chunk大小；intra_core(sch4)：最大归并迭代次数 |
| 51 | uint32_t tmpUbSize; // sort高级api需要的临时ub大小 | 58 | uint32_t tmpUbSize; // sort高级api需要的临时ub大小 |
| 52 | int64_t lastAxisNum; // h轴大小 | 59 | int64_t lastAxisNum; // h轴大小 |
| @@ -29,11 +29,13 @@ | |||
| 29 | #define SORT_SCHID_9 9 | 29 | #define SORT_SCHID_9 9 |
| 30 | #define SORT_SCHID_10 10 | 30 | #define SORT_SCHID_10 10 |
| 31 | #define SORT_SCHID_11 11 | 31 | #define SORT_SCHID_11 11 |
| 32 | +#define SORT_SCHID_12 12 | ||
| 33 | +#define SORT_SCHID_13 13 | ||
| 32 | 34 | ||
| 33 | #define SORT_TPL_KEY_DECL() \ | 35 | #define SORT_TPL_KEY_DECL() \ |
| 34 | ASCENDC_TPL_UINT_DECL(schId, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, SORT_SCHID_0, SORT_SCHID_1, SORT_SCHID_2, \ | 36 | ASCENDC_TPL_UINT_DECL(schId, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, SORT_SCHID_0, SORT_SCHID_1, SORT_SCHID_2, \ |
| 35 | SORT_SCHID_3, SORT_SCHID_4, SORT_SCHID_5, SORT_SCHID_6, SORT_SCHID_7, SORT_SCHID_8, \ | 37 | SORT_SCHID_3, SORT_SCHID_4, SORT_SCHID_5, SORT_SCHID_6, SORT_SCHID_7, SORT_SCHID_8, \ |
| 36 | - SORT_SCHID_9, SORT_SCHID_10, SORT_SCHID_11), \ | 38 | + SORT_SCHID_9, SORT_SCHID_10, SORT_SCHID_11, SORT_SCHID_12, SORT_SCHID_13), \ |
| 37 | ASCENDC_TPL_UINT_DECL(isInt32, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, 0, 1), \ | 39 | ASCENDC_TPL_UINT_DECL(isInt32, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, 0, 1), \ |
| 38 | ASCENDC_TPL_UINT_DECL(isDescend, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, 0, 1) | 40 | ASCENDC_TPL_UINT_DECL(isDescend, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, 0, 1) |
| 39 | 41 | ||
| @@ -52,9 +54,9 @@ | |||
| 52 | ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1), \ | 54 | ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1), \ |
| 53 | ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1) | 55 | ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1) |
| 54 | 56 | ||
| 55 | -#define SORT_TPL_MERGE_BIG_SIZE_KEY_SEL() \ | 57 | +#define SORT_TPL_MERGE_BIG_SIZE_KEY_SEL() \ |
| 56 | - ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_3), \ | 58 | + ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, SORT_SCHID_3, SORT_SCHID_12, SORT_SCHID_13), \ |
| 57 | - ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1), \ | 59 | + ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1), \ |
| 58 | ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1) | 60 | ASCENDC_TPL_UINT_SEL(isDescend, ASCENDC_TPL_UI_LIST, 0, 1) |
| 59 | 61 | ||
| 60 | #define SORT_TPL_MERGE_INTRA_CORE_KEY_SEL() \ | 62 | #define SORT_TPL_MERGE_INTRA_CORE_KEY_SEL() \ |
| @@ -40,6 +40,10 @@ using namespace Sort; | |||
| 40 | #define SORT_RADIX_MORE_CORE_UINT32_DESCENDING_TILING_KEY 65794 | 40 | #define SORT_RADIX_MORE_CORE_UINT32_DESCENDING_TILING_KEY 65794 |
| 41 | #define SORT_MERGE_MORE_CORE_ASCENDING_TILING_KEY 259 | 41 | #define SORT_MERGE_MORE_CORE_ASCENDING_TILING_KEY 259 |
| 42 | #define SORT_MERGE_MORE_CORE_DESCENDING_TILING_KEY 65795 | 42 | #define SORT_MERGE_MORE_CORE_DESCENDING_TILING_KEY 65795 |
| 43 | +#define SORT_MERGE_DIRECT_ASCENDING_TILING_KEY 268 | ||
| 44 | +#define SORT_MERGE_DIRECT_DESCENDING_TILING_KEY 65804 | ||
| 45 | +#define SORT_MERGE_PATH_ASCENDING_TILING_KEY 269 | ||
| 46 | +#define SORT_MERGE_PATH_DESCENDING_TILING_KEY 65805 | ||
| 43 | 47 | ||
| 44 | template <typename Op> | 48 | template <typename Op> |
| 45 | __aicore__ inline void LaunchSortKernel(GM_ADDR input, GM_ADDR values, GM_ADDR indices, GM_ADDR userWorkspace, | 49 | __aicore__ inline void LaunchSortKernel(GM_ADDR input, GM_ADDR values, GM_ADDR indices, GM_ADDR userWorkspace, |
| @@ -64,18 +68,17 @@ __aicore__ inline void LaunchRadixMoreCore(GM_ADDR radixInput, GM_ADDR radixValu | |||
| 64 | radixInput, radixValues, radixIndices, radixWorkspace, radixTiling, radixPipeline); | 68 | radixInput, radixValues, radixIndices, radixWorkspace, radixTiling, radixPipeline); |
| 65 | } | 69 | } |
| 66 | 70 | ||
| 67 | -template <uint64_t schId, uint64_t isDescend> | 71 | +template <uint64_t isDescend> |
| 68 | __aicore__ inline void LaunchMergeSortRoute(GM_ADDR mergeInput, GM_ADDR mergeValues, GM_ADDR mergeIndices, | 72 | __aicore__ inline void LaunchMergeSortRoute(GM_ADDR mergeInput, GM_ADDR mergeValues, GM_ADDR mergeIndices, |
| 69 | GM_ADDR mergeWorkspace, const SortRegBaseTilingData* mergeTiling, | 73 | GM_ADDR mergeWorkspace, const SortRegBaseTilingData* mergeTiling, |
| 70 | TPipe* mergePipeline) | 74 | TPipe* mergePipeline) |
| 71 | { | 75 | { |
| 72 | - constexpr bool isSort32SmallAxis = (schId == SORT_SCHID_8); | ||
| 73 | if constexpr (std::is_same_v<bfloat16_t, DTYPE_X>) { | 76 | if constexpr (std::is_same_v<bfloat16_t, DTYPE_X>) { |
| 74 | - LaunchSortKernel<MergeSort<DTYPE_X, DTYPE_Y2, float, isDescend, isSort32SmallAxis>>( | 77 | + LaunchSortKernel<MergeSort<DTYPE_X, DTYPE_Y2, float, isDescend>>(mergeInput, mergeValues, mergeIndices, |
| 75 | - mergeInput, mergeValues, mergeIndices, mergeWorkspace, mergeTiling, mergePipeline); | 78 | + mergeWorkspace, mergeTiling, mergePipeline); |
| 76 | } else if constexpr (std::is_same_v<half, DTYPE_X> || std::is_same_v<float, DTYPE_X>) { | 79 | } else if constexpr (std::is_same_v<half, DTYPE_X> || std::is_same_v<float, DTYPE_X>) { |
| 77 | - LaunchSortKernel<MergeSort<DTYPE_X, DTYPE_Y2, DTYPE_X, isDescend, isSort32SmallAxis>>( | 80 | + LaunchSortKernel<MergeSort<DTYPE_X, DTYPE_Y2, DTYPE_X, isDescend>>(mergeInput, mergeValues, mergeIndices, |
| 78 | - mergeInput, mergeValues, mergeIndices, mergeWorkspace, mergeTiling, mergePipeline); | 81 | + mergeWorkspace, mergeTiling, mergePipeline); |
| 79 | } | 82 | } |
| 80 | } | 83 | } |
| 81 | 84 | ||
| @@ -94,6 +97,10 @@ __global__ __aicore__ void sort(GM_ADDR input, GM_ADDR sortedValues, GM_ADDR sor | |||
| 94 | KERNEL_TASK_TYPE(SORT_RADIX_MORE_CORE_UINT32_DESCENDING_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | 97 | KERNEL_TASK_TYPE(SORT_RADIX_MORE_CORE_UINT32_DESCENDING_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); |
| 95 | KERNEL_TASK_TYPE(SORT_MERGE_MORE_CORE_ASCENDING_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | 98 | KERNEL_TASK_TYPE(SORT_MERGE_MORE_CORE_ASCENDING_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); |
| 96 | KERNEL_TASK_TYPE(SORT_MERGE_MORE_CORE_DESCENDING_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | 99 | KERNEL_TASK_TYPE(SORT_MERGE_MORE_CORE_DESCENDING_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); |
| 100 | + KERNEL_TASK_TYPE(SORT_MERGE_DIRECT_ASCENDING_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | ||
| 101 | + KERNEL_TASK_TYPE(SORT_MERGE_DIRECT_DESCENDING_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | ||
| 102 | + KERNEL_TASK_TYPE(SORT_MERGE_PATH_ASCENDING_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | ||
| 103 | + KERNEL_TASK_TYPE(SORT_MERGE_PATH_DESCENDING_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | ||
| 97 | TPipe sortPipeline; | 104 | TPipe sortPipeline; |
| 98 | constexpr bool isDescending = (isDescend != 0); | 105 | constexpr bool isDescending = (isDescend != 0); |
| 99 | if constexpr (schId == SORT_SCHID_7) { | 106 | if constexpr (schId == SORT_SCHID_7) { |
| @@ -106,10 +113,13 @@ __global__ __aicore__ void sort(GM_ADDR input, GM_ADDR sortedValues, GM_ADDR sor | |||
| 106 | LaunchSortKernel<SortRadixOneCore<DTYPE_X, DTYPE_Y2, isDescending>>( | 113 | LaunchSortKernel<SortRadixOneCore<DTYPE_X, DTYPE_Y2, isDescending>>( |
| 107 | input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData, &sortPipeline); | 114 | input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData, &sortPipeline); |
| 108 | } else if constexpr (schId == SORT_SCHID_0 || schId == SORT_SCHID_8) { | 115 | } else if constexpr (schId == SORT_SCHID_0 || schId == SORT_SCHID_8) { |
| 109 | - LaunchMergeSortRoute<schId, isDescend>(input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData, | 116 | + LaunchMergeSortRoute<isDescend>(input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData, |
| 110 | - &sortPipeline); | 117 | + &sortPipeline); |
| 111 | - } else if constexpr (schId == SORT_SCHID_3 && std::is_same_v<float, DTYPE_X>) { | 118 | + } else if constexpr ((schId == SORT_SCHID_3 || schId == SORT_SCHID_12 || schId == SORT_SCHID_13) && |
黄 [P1] 新增 schId=12/13 跨核归并路径需要补齐 MIX kernel 类型注册 这两个分支都进入 ProcessDirect()，在 common/merge_more_core_base.h:187 调用硬件 SyncAll()，后续归并也依赖核间同步。但本文件第 90–96 行仍只为旧 schId=3 的 key 259/65795 注册 KERNEL_TYPE_MIX_AIV_1_0，新 key 268/65804、269/65805 会落入 KERNEL_TYPE_AIV_ONLY 默认配置。 AscendC SyncAll 的硬件同步约束要求纯 Vector 算子使用 MIX 类型；CrossCoreSetFlag 的约束说明指出 AIV_ONLY 不开启核间同步所需的硬件调度模块。host 的 SetScheduleMode(1) 不替代 kernel 类型配置，因此这些新增路径存在运行异常或挂起风险。例如 FP32 [2,4097] 已由本 PR 的 UT 期望分派到 key 268，[8,12185] 可进入 schId=13。 建议为四个新增 tiling key 补齐 KERNEL_TASK_TYPE(..., KERNEL_TYPE_MIX_AIV_1_0)，并分别验证两个路径的升序/降序多核用例。 验证范围：基于当前提交 087bbcc418bfa20a4506fe06eb0d9b42c6c653a4 的静态调用链及 AscendC API 约束核对，未做 NPU 实跑。   | |||
| 112 | - LaunchSortKernel<MergeSortBigSize<DTYPE_X, DTYPE_X, isDescending, DTYPE_Y2>>( | 119 | + std::is_same_v<float, DTYPE_X>) { |
| 120 | + constexpr bool useDirectSchedule = (schId == SORT_SCHID_12 || schId == SORT_SCHID_13); | ||
| 121 | + constexpr bool useMergePath = (schId == SORT_SCHID_13); | ||
| 122 | + LaunchSortKernel<MergeSortBigSize<DTYPE_X, DTYPE_X, isDescending, DTYPE_Y2, useDirectSchedule, useMergePath>>( | ||
| 113 | input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData, &sortPipeline); | 123 | input, sortedValues, sortedIndices, sortWorkspace, &sortTilingData, &sortPipeline); |
| 114 | } else if constexpr (schId == SORT_SCHID_4 && std::is_same_v<float, DTYPE_X>) { | 124 | } else if constexpr (schId == SORT_SCHID_4 && std::is_same_v<float, DTYPE_X>) { |
| 115 | LaunchSortKernel<Sort::SortMergeIntraCore<DTYPE_X, DTYPE_Y2, isDescending>>( | 125 | LaunchSortKernel<Sort::SortMergeIntraCore<DTYPE_X, DTYPE_Y2, isDescending>>( |
| @@ -4,3 +4,6 @@ aclnn_sort_002,aclnnSort,"('bfloat16','bfloat16','int64')","('ND', 'ND', 'ND')", | |||
| 4 | aclnn_sort_003,aclnnSort,"('float32','float32','int64')","('ND', 'ND', 'ND')","{ 'stable': True,'dim': -1, 'descending': False}","((3, 32768), (3, 32768), (3,32768))","(1,2)","((-200,200),)",1 | 4 | aclnn_sort_003,aclnnSort,"('float32','float32','int64')","('ND', 'ND', 'ND')","{ 'stable': True,'dim': -1, 'descending': False}","((3, 32768), (3, 32768), (3,32768))","(1,2)","((-200,200),)",1 |
| 5 | aclnn_sort_004,aclnnSort,"('int32','int32','int64')","('ND', 'ND', 'ND')","{ 'stable': True,'dim': -1, 'descending': True}","((3, 32768), (3, 32768), (3,32768))","(1,2)","((-200,200),)",1 | 5 | aclnn_sort_004,aclnnSort,"('int32','int32','int64')","('ND', 'ND', 'ND')","{ 'stable': True,'dim': -1, 'descending': True}","((3, 32768), (3, 32768), (3,32768))","(1,2)","((-200,200),)",1 |
| 6 | aclnn_sort_005,aclnnSort,"('int16','int16','int64')","('ND', 'ND', 'ND')","{ 'stable': True,'dim': -1, 'descending': False}","((3, 32768), (3, 32768), (3,32768))","(1,2)","((-200,200),)",1 | 6 | aclnn_sort_005,aclnnSort,"('int16','int16','int64')","('ND', 'ND', 'ND')","{ 'stable': True,'dim': -1, 'descending': False}","((3, 32768), (3, 32768), (3,32768))","(1,2)","((-200,200),)",1 |
| 7 | +aclnn_sort_trailing_singletons_fp32,aclnnSort,"('float32', 'float32', 'int64')","('ND','ND','ND')","{'stable': True, 'dim': 1, 'descending': False}","((8, 12185, 1, 1), (8, 12185, 1, 1), (8, 12185, 1, 1))","(1,2)","((-2,2),)",1 | ||
| 8 | +aclnn_sort_trailing_singletons_bf16,aclnnSort,"('bfloat16', 'bfloat16', 'int64')","('ND','ND','ND')","{'stable': True, 'dim': -3, 'descending': True}","((2, 65, 1, 1), (2, 65, 1, 1), (2, 65, 1, 1))","(1,2)","((-2,2),)",1 | ||
| 9 | +aclnn_sort_non_singleton_tail,aclnnSort,"('float32', 'float32', 'int64')","('ND','ND','ND')","{'stable': True, 'dim': 1, 'descending': False}","((2, 65, 1, 2), (2, 65, 1, 2), (2, 65, 1, 2))","(1,2)","((-2,2),)",1 | ||
| @@ -2,7 +2,7 @@ testcase_name,network_name,op_name,input_shapes,input_dtypes,input_formats,outpu | |||
| 2 | sort_sch00_f32,,sort,"((2, 1024),)","('float32',)","('ND',)","((2, 1024), (2, 1024))","('float32', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': False, 'stable': True, 'y2_dtype': 9}","((-100, 100),)",,,(),(),True,schId=0 merge_sort,,0,,(),() | 2 | sort_sch00_f32,,sort,"((2, 1024),)","('float32',)","('ND',)","((2, 1024), (2, 1024))","('float32', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': False, 'stable': True, 'y2_dtype': 9}","((-100, 100),)",,,(),(),True,schId=0 merge_sort,,0,,(),() |
| 3 | sort_sch01_i32,,sort,"((128, 1024),)","('int32',)","('ND',)","((128, 1024), (128, 1024))","('int32', 'int32')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': False, 'stable': True, 'y2_dtype': 3}","((-100, 100),)",,,(),(),True,schId=1 radix_one_core,,0,,(),() | 3 | sort_sch01_i32,,sort,"((128, 1024),)","('int32',)","('ND',)","((128, 1024), (128, 1024))","('int32', 'int32')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': False, 'stable': True, 'y2_dtype': 3}","((-100, 100),)",,,(),(),True,schId=1 radix_one_core,,0,,(),() |
| 4 | sort_sch02_f16,,sort,"((2, 32768),)","('float16',)","('ND',)","((2, 32768), (2, 32768))","('float16', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': True, 'stable': True, 'y2_dtype': 9}","((-100, 100),)",,,(),(),True,schId=2 radix_more_core,,0,,(),() | 4 | sort_sch02_f16,,sort,"((2, 32768),)","('float16',)","('ND',)","((2, 32768), (2, 32768))","('float16', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': True, 'stable': True, 'y2_dtype': 9}","((-100, 100),)",,,(),(),True,schId=2 radix_more_core,,0,,(),() |
| 5 | -sort_sch03_f32,,sort,"((2, 4097),)","('float32',)","('ND',)","((2, 4097), (2, 4097))","('float32', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': False, 'stable': True, 'y2_dtype': 9}","((-100, 100),)",,,(),(),True,schId=3 merge_more_core,,0,,(),() | 5 | +sort_sch12_f32,,sort,"((2, 4097),)","('float32',)","('ND',)","((2, 4097), (2, 4097))","('float32', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': False, 'stable': True, 'y2_dtype': 9}","((-100, 100),)",,,(),(),True,schId=12 merge_more_core direct schedule,,0,,(),() |
| 6 | sort_sch04_f32,,sort,"((64, 16384),)","('float32',)","('ND',)","((64, 16384), (64, 16384))","('float32', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': True, 'stable': True, 'y2_dtype': 9}","((-100, 100),)",,,(),(),True,schId=4 merge_intra_core,,0,,(),() | 6 | sort_sch04_f32,,sort,"((64, 16384),)","('float32',)","('ND',)","((64, 16384), (64, 16384))","('float32', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': True, 'stable': True, 'y2_dtype': 9}","((-100, 100),)",,,(),(),True,schId=4 merge_intra_core,,0,,(),() |
| 7 | sort_sch05_i64,,sort,"((256, 12),)","('int64',)","('ND',)","((256, 12), (256, 12))","('int64', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': True, 'stable': True, 'y2_dtype': 9}","((-100, 100),)",,,(),(),True,schId=5 insertion,,0,,(),() | 7 | sort_sch05_i64,,sort,"((256, 12),)","('int64',)","('ND',)","((256, 12), (256, 12))","('int64', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': True, 'stable': True, 'y2_dtype': 9}","((-100, 100),)",,,(),(),True,schId=5 insertion,,0,,(),() |
| 8 | sort_sch06_i32,,sort,"((256, 13),)","('int32',)","('ND',)","((256, 13), (256, 13))","('int32', 'int32')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': False, 'stable': True, 'y2_dtype': 3}","((-100, 100),)",,,(),(),True,schId=6 two_stage,,0,,(),() | 8 | sort_sch06_i32,,sort,"((256, 13),)","('int32',)","('ND',)","((256, 13), (256, 13))","('int32', 'int32')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': False, 'stable': True, 'y2_dtype': 3}","((-100, 100),)",,,(),(),True,schId=6 two_stage,,0,,(),() |
| @@ -24,3 +24,13 @@ sort_nan,,sort,"((2, 8),)","('float32',)","('ND',)","((2, 8), (2, 8))","('float3 | |||
| 24 | sort_infinity,,sort,"((2, 8),)","('float32',)","('ND',)","((2, 8), (2, 8))","('float32', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': True, 'stable': True, 'y2_dtype': 9}","((-4, 4),)",,,(),(),True,infinity descending,,0,,(),() | 24 | sort_infinity,,sort,"((2, 8),)","('float32',)","('ND',)","((2, 8), (2, 8))","('float32', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': True, 'stable': True, 'y2_dtype': 9}","((-4, 4),)",,,(),(),True,infinity descending,,0,,(),() |
| 25 | sort_ins_asc_nan,,sort,"((1024, 4),)","('float32',)","('ND',)","((1024, 4), (1024, 4))","('float32', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': False, 'stable': True, 'y2_dtype': 9}","((-4, 4),)",,,(),(),True,schId=5 ascending NaN,,0,,(),() | 25 | sort_ins_asc_nan,,sort,"((1024, 4),)","('float32',)","('ND',)","((1024, 4), (1024, 4))","('float32', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': False, 'stable': True, 'y2_dtype': 9}","((-4, 4),)",,,(),(),True,schId=5 ascending NaN,,0,,(),() |
| 26 | sort_ins_desc_nan,,sort,"((1024, 4),)","('float32',)","('ND',)","((1024, 4), (1024, 4))","('float32', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': True, 'stable': True, 'y2_dtype': 9}","((-4, 4),)",,,(),(),True,schId=5 descending NaN,,0,,(),() | 26 | sort_ins_desc_nan,,sort,"((1024, 4),)","('float32',)","('ND',)","((1024, 4), (1024, 4))","('float32', 'int64')","('ND', 'ND')",,,,,"{'axis': -1, 'descending': True, 'stable': True, 'y2_dtype': 9}","((-4, 4),)",,,(),(),True,schId=5 descending NaN,,0,,(),() |
| 27 | +sort_resident_random_43_10560_int64_False,,sort,"((43, 10560),)","('float32',)","('ND',)","((43, 10560), (43, 10560))","('float32', 'int64')","('ND','ND')","((43, 10560),)",,"((43, 10560), (43, 10560))",,"{'axis': 1, 'descending': False, 'stable': True, 'y2_dtype': 9}","((-100,100),)","((0,0),(0,0))",0,(),(),True,"FP32 resident merge capacity boundary, stable indices",,,,(),() | ||
| 28 | +sort_resident_random_43_10561_int32_True,,sort,"((43, 10561),)","('float32',)","('ND',)","((43, 10561), (43, 10561))","('float32', 'int32')","('ND','ND')","((43, 10561),)",,"((43, 10561), (43, 10561))",,"{'axis': 1, 'descending': True, 'stable': True, 'y2_dtype': 3}","((-100,100),)","((0,0),(0,0))",0,(),(),True,"FP32 resident merge capacity boundary, stable indices",,,,(),() | ||
| 29 | +sort_mp_random_8_12185_int64,,sort,"((8, 12185),)","('float32',)","('ND',)","((8, 12185), (8, 12185))","('float32', 'int64')","('ND','ND')","((8, 12185),)",,"((8, 12185), (8, 12185))",,"{'axis': 1, 'descending': False, 'stable': True, 'y2_dtype': 9}","((-100,100),)","((0,0),(0,0))",0,(),(),True,Parallel two-run final merge; disjoint output partitions and stable indices,,,,(),() | ||
| 30 | +sort_mp_random_4_14001_int32,,sort,"((4, 14001),)","('float32',)","('ND',)","((4, 14001), (4, 14001))","('float32', 'int32')","('ND','ND')","((4, 14001),)",,"((4, 14001), (4, 14001))",,"{'axis': 1, 'descending': True, 'stable': True, 'y2_dtype': 3}","((-100,100),)","((0,0),(0,0))",0,(),(),True,Parallel two-run final merge; disjoint output partitions and stable indices,,,,(),() | ||
| 31 | +sort_mp_retile_2_19176_int64,,sort,"((2, 19176),)","('float32',)","('ND',)","((2, 19176), (2, 19176))","('float32', 'int64')","('ND','ND')",,,,,"{'axis': 1, 'descending': True, 'stable': True, 'y2_dtype': 9}","((-100,100),)","((0,0),(0,0))",0,(),(),True,Retile three final runs into two; stable parallel merge,,,,(),() | ||
| 32 | +sort_mp_retile_4_24576_int32,,sort,"((4, 24576),)","('float32',)","('ND',)","((4, 24576), (4, 24576))","('float32', 'int32')","('ND','ND')",,,,,"{'axis': 1, 'descending': False, 'stable': True, 'y2_dtype': 3}","((-100,100),)","((0,0),(0,0))",0,(),(),True,Retile three final runs into two; stable parallel merge,,,,(),() | ||
| 33 | +sort_mp_partial_window_7_32768_int64,,sort,"((7, 32768),)","('float32',)","('ND',)","((7, 32768), (7, 32768))","('float32', 'int64')","('ND','ND')",,,,,"{'axis': 1, 'descending': True, 'stable': True, 'y2_dtype': 9}","((0,0),)","((0,0),(0,0))",0,(),(),True,Partial-window refill and exhausted-run handling; stable duplicate indices,,,,(),() | ||
| 34 | +sort_full_window_sync_13_20001_int32,,sort,"((13, 20001),)","('float32',)","('ND',)","((13, 20001), (13, 20001))","('float32', 'int32')","('ND','ND')",,,,,"{'axis': 1, 'descending': True, 'stable': True, 'y2_dtype': 3}","((0,0),)","((0,0),(0,0))",0,(),(),True,Complete loaded four-run windows; stable duplicate indices,,,,(),() | ||
| 35 | +sort_nonlast_compact_fp16_shared_tail,,sort,"((211, 8193),)","('float16',)","('ND',)","((211, 8193), (211, 8193))","('float16', 'int64')","('ND','ND')",,,,,"{'axis': 0, 'descending': False, 'stable': True, 'y2_dtype': 9}","((-1,1),)","((0,0),(0,0))",0,(),(),True,"Compact non-last merge; padded input, shared proposals and partial inner tile",,,,(),() | ||
| 36 | +sort_nonlast_compact_fp32_index32,,sort,"((119, 8193),)","('float32',)","('ND',)","((119, 8193), (119, 8193))","('float32', 'int32')","('ND','ND')",,,,,"{'axis': 0, 'descending': True, 'stable': True, 'y2_dtype': 3}","((-1,1),)","((0,0),(0,0))",0,(),(),True,"Compact non-last merge; padded input, shared proposals and partial inner tile",,,,(),() | ||
| @@ -146,20 +146,67 @@ TEST_F(SortTilingTest, test_sort_merge_big_size_fp32_boundary_4097) | |||
| 146 | { | 146 | { |
| 147 | auto tilingContextPara = MakeSortTilingContext({{2, 4097}, {2, 4097}}, ge::DT_FLOAT, ge::DT_INT64); | 147 | auto tilingContextPara = MakeSortTilingContext({{2, 4097}, {2, 4097}}, ge::DT_FLOAT, ge::DT_INT64); |
| 148 | 148 | ||
| 149 | - uint64_t expectTilingKey = 259; | 149 | + uint64_t expectTilingKey = 268; |
| 150 | std::vector<size_t> expectWorkspaces = {16941096}; | 150 | std::vector<size_t> expectWorkspaces = {16941096}; |
| 151 | ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectWorkspaces); | 151 | ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectWorkspaces); |
| 152 | } | 152 | } |
| 153 | 153 | ||
| 154 | -TEST_F(SortTilingTest, test_sort_radix_one_core_precedes_merge_intra_core_fp32_128x8192) | 154 | +TEST_F(SortTilingTest, test_sort_merge_intra_core_precedes_radix_one_core_fp32_128x8192) |
| 155 | { | 155 | { |
| 156 | auto tilingContextPara = MakeSortTilingContext({{128, 8192}, {128, 8192}}, ge::DT_FLOAT, ge::DT_INT64); | 156 | auto tilingContextPara = MakeSortTilingContext({{128, 8192}, {128, 8192}}, ge::DT_FLOAT, ge::DT_INT64); |
| 157 | 157 | ||
| 158 | - uint64_t expectTilingKey = 257; | 158 | + uint64_t expectTilingKey = 260; |
| 159 | - std::vector<size_t> expectWorkspaces = {WORK_SPACE_SIZE}; | 159 | + std::vector<size_t> expectWorkspaces = {25165824}; |
| 160 | ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectWorkspaces); | 160 | ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectWorkspaces); |
| 161 | } | 161 | } |
| 162 | 162 | ||
| 163 | +TEST_F(SortTilingTest, test_sort_falls_back_to_sync_merge_when_radix_plan_is_unavailable) | ||
| 164 | +{ | ||
| 165 | + auto tilingContextPara = MakeSortTilingContext({{9, 32768}, {9, 32768}}, ge::DT_FLOAT, ge::DT_INT64); | ||
| 166 | + | ||
| 167 | + TilingInfo tilingInfo; | ||
| 168 | + ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo)); | ||
| 169 | + EXPECT_EQ(tilingInfo.tilingKey, 259); | ||
| 170 | + EXPECT_EQ(tilingInfo.blockNum, 54); | ||
| 171 | + ASSERT_GE(tilingInfo.tilingDataSize, sizeof(SortRegBaseTilingData)); | ||
| 172 | + const auto* tilingData = reinterpret_cast<const SortRegBaseTilingData*>(tilingInfo.tilingData.get()); | ||
| 173 | + EXPECT_EQ(tilingData->lastDimNeedCore, 6U); | ||
| 174 | + EXPECT_EQ(tilingData->unsortedDimParallel, 9U); | ||
| 175 | + EXPECT_EQ(tilingData->sortLoopTimes, 1U); | ||
| 176 | + EXPECT_EQ(tilingData->keyParams1, 6144U); | ||
| 177 | +} | ||
| 178 | + | ||
| 179 | +TEST_F(SortTilingTest, test_sort_prefers_row_parallel_merge_over_multi_round_merge) | ||
| 180 | +{ | ||
| 181 | + auto tilingContextPara = MakeSortTilingContext({{65, 32768}, {65, 32768}}, ge::DT_FLOAT, ge::DT_INT64); | ||
| 182 | + | ||
| 183 | + TilingInfo tilingInfo; | ||
| 184 | + ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo)); | ||
| 185 | + EXPECT_EQ(tilingInfo.tilingKey, 260); | ||
| 186 | + EXPECT_EQ(tilingInfo.blockNum, 33); | ||
| 187 | + ASSERT_GE(tilingInfo.tilingDataSize, sizeof(SortRegBaseTilingData)); | ||
| 188 | + const auto* tilingData = reinterpret_cast<const SortRegBaseTilingData*>(tilingInfo.tilingData.get()); | ||
| 189 | + EXPECT_EQ(tilingData->keyParams0, 2U); | ||
| 190 | + EXPECT_EQ(tilingData->numTileDataSize, 4096U); | ||
| 191 | + EXPECT_EQ(tilingData->unsortedDimParallel, 33U); | ||
| 192 | + ASSERT_EQ(tilingInfo.workspaceSizes.size(), 1); | ||
| 193 | + EXPECT_EQ(tilingInfo.workspaceSizes[0], 34078720); | ||
| 194 | +} | ||
| 195 | + | ||
| 196 | +TEST_F(SortTilingTest, test_sort_keeps_row_parallel_merge_across_multi_round_plan_boundary) | ||
| 197 | +{ | ||
| 198 | + auto belowThreshold = MakeSortTilingContext({{65, 7 * 4096}, {65, 7 * 4096}}, ge::DT_FLOAT, ge::DT_INT64); | ||
| 199 | + auto atThreshold = MakeSortTilingContext({{65, 7 * 4096 + 1}, {65, 7 * 4096 + 1}}, ge::DT_FLOAT, ge::DT_INT64); | ||
| 200 | + | ||
| 201 | + TilingInfo belowThresholdInfo; | ||
| 202 | + ASSERT_TRUE(ExecuteTiling(belowThreshold, belowThresholdInfo)); | ||
| 203 | + EXPECT_EQ(belowThresholdInfo.tilingKey, 260); | ||
| 204 | + | ||
| 205 | + TilingInfo atThresholdInfo; | ||
| 206 | + ASSERT_TRUE(ExecuteTiling(atThreshold, atThresholdInfo)); | ||
| 207 | + EXPECT_EQ(atThresholdInfo.tilingKey, 260); | ||
| 208 | +} | ||
| 209 | + | ||
| 163 | TEST_F(SortTilingTest, test_sort_merge_intra_core_fp32_64x16384) | 210 | TEST_F(SortTilingTest, test_sort_merge_intra_core_fp32_64x16384) |
| 164 | { | 211 | { |
| 165 | auto tilingContextPara = MakeSortTilingContext({{64, 16384}, {64, 16384}}, ge::DT_FLOAT, ge::DT_INT64); | 212 | auto tilingContextPara = MakeSortTilingContext({{64, 16384}, {64, 16384}}, ge::DT_FLOAT, ge::DT_INT64); |
| @@ -34,10 +34,10 @@ template <typename ValueType, typename IndexType, bool IsDescend> | |||
| 34 | class TopKMergeSortIntraCore : public Sort::SortMergeIntraCore<ValueType, IndexType, IsDescend> { | 34 | class TopKMergeSortIntraCore : public Sort::SortMergeIntraCore<ValueType, IndexType, IsDescend> { |
| 35 | public: | 35 | public: |
| 36 | using Base = Sort::SortMergeIntraCore<ValueType, IndexType, IsDescend>; | 36 | using Base = Sort::SortMergeIntraCore<ValueType, IndexType, IsDescend>; |
| 37 | - | 37 | + |
| 38 | __aicore__ inline TopKMergeSortIntraCore() : Base() {} | 38 | __aicore__ inline TopKMergeSortIntraCore() : Base() {} |
| 39 | __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR indices, GM_ADDR workspace, | 39 | __aicore__ inline void Init(GM_ADDR x, GM_ADDR value, GM_ADDR indices, GM_ADDR workspace, |
| 40 | - const TopKV2TilingDataSimd* tilingData, TPipe* pipe); | 40 | + const TopKV2TilingDataSimd* tilingData, TPipe* pipe); |
| 41 | __aicore__ inline void Process(); | 41 | __aicore__ inline void Process(); |
| 42 | 42 | ||
| 43 | private: | 43 | private: |
| @@ -45,26 +45,25 @@ private: | |||
| 45 | 45 | ||
| 46 | __aicore__ inline uint32_t MergeSingleBatch(); | 46 | __aicore__ inline uint32_t MergeSingleBatch(); |
| 47 | __aicore__ inline void DoIncrementalMerge(int64_t dstOffsetBase, typename Base::MergeListContext& ctx); | 47 | __aicore__ inline void DoIncrementalMerge(int64_t dstOffsetBase, typename Base::MergeListContext& ctx); |
| 48 | - __aicore__ inline void MergeOneGroup(uint32_t groupStart, uint32_t groupBlockCount, | 48 | + __aicore__ inline void MergeOneGroup(uint32_t groupStart, uint32_t groupBlockCount, uint32_t fullBlockElemCount, |
| 49 | - uint32_t fullBlockElemCount, uint32_t fullBlockSortLen, uint32_t lastBlockElemCount, | 49 | + uint32_t fullBlockSortLen, uint32_t lastBlockElemCount, uint32_t numBlocks, |
| 50 | - uint32_t numBlocks, uint32_t pingPongFlag, uint32_t& cumulativeOffset, | 50 | + uint32_t pingPongFlag, uint32_t& cumulativeOffset, |
| 51 | - uint32_t& mergedGroupElemCount); | 51 | + uint32_t& mergedGroupElemCount); |
| 52 | - __aicore__ inline void CopyRemainingList(typename Base::MergeListContext& ctx, | 52 | + __aicore__ inline void CopyRemainingList(typename Base::MergeListContext& ctx, int64_t dstOffsetBase, |
| 53 | - int64_t dstOffsetBase, uint32_t& dstCumulativeOffset, uint32_t& dstElemCount); | 53 | + uint32_t& dstCumulativeOffset, uint32_t& dstElemCount); |
| 54 | __aicore__ inline void ExtractAndCopyOut(int64_t batchIdx, uint32_t resultRegion); | 54 | __aicore__ inline void ExtractAndCopyOut(int64_t batchIdx, uint32_t resultRegion); |
| 55 | }; | 55 | }; |
| 56 | 56 | ||
| 57 | template <typename ValueType, typename IndexType, bool IsDescend> | 57 | template <typename ValueType, typename IndexType, bool IsDescend> |
| 58 | __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::Init( | 58 | __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::Init( |
| 59 | - GM_ADDR x, GM_ADDR value, GM_ADDR indices, GM_ADDR workspace, | 59 | + GM_ADDR x, GM_ADDR value, GM_ADDR indices, GM_ADDR workspace, const TopKV2TilingDataSimd* tilingData, TPipe* pipe) |
| 60 | - const TopKV2TilingDataSimd* tilingData, TPipe* pipe) | ||
| 61 | { | 60 | { |
| 62 | if (tilingData == nullptr || pipe == nullptr) { | 61 | if (tilingData == nullptr || pipe == nullptr) { |
| 63 | return; | 62 | return; |
| 64 | } | 63 | } |
| 65 | - | 64 | + |
| 66 | topKValue_ = static_cast<uint32_t>(tilingData->topKRealValue); | 65 | topKValue_ = static_cast<uint32_t>(tilingData->topKRealValue); |
| 67 | - | 66 | + |
| 68 | SortRegBaseTilingData sortTilingData; | 67 | SortRegBaseTilingData sortTilingData; |
| 69 | sortTilingData.unsortedDimNum = tilingData->unsortedDimNum; | 68 | sortTilingData.unsortedDimNum = tilingData->unsortedDimNum; |
| 70 | sortTilingData.lastAxisNum = tilingData->lastAxisNum; | 69 | sortTilingData.lastAxisNum = tilingData->lastAxisNum; |
| @@ -75,7 +74,7 @@ __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>:: | |||
| 75 | sortTilingData.keyParams3 = tilingData->keyParams3; | 74 | sortTilingData.keyParams3 = tilingData->keyParams3; |
| 76 | sortTilingData.keyParams4 = tilingData->keyParams4; | 75 | sortTilingData.keyParams4 = tilingData->keyParams4; |
| 77 | sortTilingData.keyParams5 = tilingData->keyParams5; | 76 | sortTilingData.keyParams5 = tilingData->keyParams5; |
| 78 | - | 77 | + |
| 79 | Base::Init(x, value, indices, workspace, &sortTilingData, pipe); | 78 | Base::Init(x, value, indices, workspace, &sortTilingData, pipe); |
| 80 | } | 79 | } |
| 81 | 80 | ||
| @@ -84,16 +83,16 @@ __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>:: | |||
| 84 | { | 83 | { |
| 85 | // Calculate batch range for this core | 84 | // Calculate batch range for this core |
| 86 | int64_t startBatch = static_cast<int64_t>(this->blockIdx_) * this->batchPerCore_; | 85 | int64_t startBatch = static_cast<int64_t>(this->blockIdx_) * this->batchPerCore_; |
| 87 | - int64_t endBatch = (startBatch + this->batchPerCore_ < this->batchNum_) ? | 86 | + int64_t endBatch = (startBatch + this->batchPerCore_ < this->batchNum_) ? (startBatch + this->batchPerCore_) : |
| 88 | - (startBatch + this->batchPerCore_) : this->batchNum_; | 87 | + this->batchNum_; |
| 89 | 88 | ||
| 90 | if (startBatch >= this->batchNum_) { | 89 | if (startBatch >= this->batchNum_) { |
| 91 | - return; // This core has no work | 90 | + return; // This core has no work |
| 92 | } | 91 | } |
| 93 | 92 | ||
| 94 | // Process each batch through all 3 phases before moving to the next batch. | 93 | // Process each batch through all 3 phases before moving to the next batch. |
| 95 | - // This allows cache (workspace) to be reused across batches within the same core, | 94 | + // This allows cache (workspace) to be reused across batches within the same core. |
| 96 | - | 95 | + |
| 97 | // Precompute buffer sizes (loop-invariant) | 96 | // Precompute buffer sizes (loop-invariant) |
| 98 | uint32_t mergeBufferSize = Sort::MERGE_LIST_MAX_NUM * this->sortBufferSize_; | 97 | uint32_t mergeBufferSize = Sort::MERGE_LIST_MAX_NUM * this->sortBufferSize_; |
| 99 | uint32_t extractInSize = AscendC::GetSortLen<ValueType>(this->extractChunkSize_) * sizeof(ValueType); | 98 | uint32_t extractInSize = AscendC::GetSortLen<ValueType>(this->extractChunkSize_) * sizeof(ValueType); |
| @@ -101,7 +100,6 @@ __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>:: | |||
| 101 | for (int64_t batchIdx = startBatch; batchIdx < endBatch; batchIdx++) { | 100 | for (int64_t batchIdx = startBatch; batchIdx < endBatch; batchIdx++) { |
| 102 | // ========== Phase 1: Sort blocks in UB ========== | 101 | // ========== Phase 1: Sort blocks in UB ========== |
| 103 | this->pipe_->InitBuffer(this->inQueueX_, DOUBLE_BUFFER, this->blockSortSize_ * sizeof(ValueType)); | 102 | this->pipe_->InitBuffer(this->inQueueX_, DOUBLE_BUFFER, this->blockSortSize_ * sizeof(ValueType)); |
| 104 | - this->pipe_->InitBuffer(this->concatTmpBuf_, this->sortBufferSize_); | ||
| 105 | this->pipe_->InitBuffer(this->sortTmpBuf_, this->sortBufferSize_); | 103 | this->pipe_->InitBuffer(this->sortTmpBuf_, this->sortBufferSize_); |
| 106 | this->pipe_->InitBuffer(this->sortedOutQueue_, DOUBLE_BUFFER, this->sortBufferSize_); | 104 | this->pipe_->InitBuffer(this->sortedOutQueue_, DOUBLE_BUFFER, this->sortBufferSize_); |
| 107 | this->pipe_->InitBuffer(this->indexTmpBuf_, this->blockSortSize_ * sizeof(uint32_t)); | 105 | this->pipe_->InitBuffer(this->indexTmpBuf_, this->blockSortSize_ * sizeof(uint32_t)); |
| @@ -158,23 +156,23 @@ __aicore__ inline uint32_t TopKMergeSortIntraCore<ValueType, IndexType, IsDescen | |||
| 158 | uint32_t newLastBlockElemCount = 0; | 156 | uint32_t newLastBlockElemCount = 0; |
| 159 | 157 | ||
| 160 | for (uint32_t i = 0; i < numBlocks; i += Sort::MERGE_LIST_MAX_NUM) { | 158 | for (uint32_t i = 0; i < numBlocks; i += Sort::MERGE_LIST_MAX_NUM) { |
| 161 | - uint32_t groupBlockCount = (i + Sort::MERGE_LIST_MAX_NUM <= numBlocks) ? | 159 | + uint32_t groupBlockCount = (i + Sort::MERGE_LIST_MAX_NUM <= numBlocks) ? Sort::MERGE_LIST_MAX_NUM : |
| 162 | - Sort::MERGE_LIST_MAX_NUM : (numBlocks - i); | 160 | + (numBlocks - i); |
| 163 | 161 | ||
| 164 | uint32_t mergedGroupElemCount = 0; | 162 | uint32_t mergedGroupElemCount = 0; |
| 165 | - MergeOneGroup(i, groupBlockCount, fullBlockElemCount, fullBlockSortLen, lastBlockElemCount, | 163 | + MergeOneGroup(i, groupBlockCount, fullBlockElemCount, fullBlockSortLen, lastBlockElemCount, numBlocks, |
| 166 | - numBlocks, pingPongFlag, cumulativeOffset, mergedGroupElemCount); | 164 | + pingPongFlag, cumulativeOffset, mergedGroupElemCount); |
| 167 | if (newNumBlocks == 0) { | 165 | if (newNumBlocks == 0) { |
| 168 | newFullBlockElemCount = mergedGroupElemCount; | 166 | newFullBlockElemCount = mergedGroupElemCount; |
| 169 | } | 167 | } |
| 170 | newLastBlockElemCount = mergedGroupElemCount; | 168 | newLastBlockElemCount = mergedGroupElemCount; |
| 171 | newNumBlocks++; | 169 | newNumBlocks++; |
| 172 | } | 170 | } |
| 173 | - | 171 | + |
| 174 | if (newNumBlocks == 0 || newNumBlocks >= numBlocks) { | 172 | if (newNumBlocks == 0 || newNumBlocks >= numBlocks) { |
| 175 | break; | 173 | break; |
| 176 | } | 174 | } |
| 177 | - | 175 | + |
| 178 | numBlocks = newNumBlocks; | 176 | numBlocks = newNumBlocks; |
| 179 | if (numBlocks == 1) { | 177 | if (numBlocks == 1) { |
| 180 | fullBlockElemCount = newLastBlockElemCount; | 178 | fullBlockElemCount = newLastBlockElemCount; |
| @@ -195,14 +193,13 @@ __aicore__ inline uint32_t TopKMergeSortIntraCore<ValueType, IndexType, IsDescen | |||
| 195 | 193 | ||
| 196 | template <typename ValueType, typename IndexType, bool IsDescend> | 194 | template <typename ValueType, typename IndexType, bool IsDescend> |
| 197 | __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::MergeOneGroup( | 195 | __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::MergeOneGroup( |
| 198 | - uint32_t groupStart, uint32_t groupBlockCount, | 196 | + uint32_t groupStart, uint32_t groupBlockCount, uint32_t fullBlockElemCount, uint32_t fullBlockSortLen, |
| 199 | - uint32_t fullBlockElemCount, uint32_t fullBlockSortLen, uint32_t lastBlockElemCount, | 197 | + uint32_t lastBlockElemCount, uint32_t numBlocks, uint32_t pingPongFlag, uint32_t& cumulativeOffset, |
| 200 | - uint32_t numBlocks, uint32_t pingPongFlag, uint32_t& cumulativeOffset, | ||
| 201 | uint32_t& mergedGroupElemCount) | 198 | uint32_t& mergedGroupElemCount) |
| 202 | { | 199 | { |
| 203 | typename Base::MergeListContext ctx; | 200 | typename Base::MergeListContext ctx; |
| 204 | ctx.listCount = groupBlockCount; | 201 | ctx.listCount = groupBlockCount; |
| 205 | - | 202 | + |
| 206 | int64_t srcRegionOffset = (pingPongFlag == 0) ? 0 : this->batchSortLen_; | 203 | int64_t srcRegionOffset = (pingPongFlag == 0) ? 0 : this->batchSortLen_; |
| 207 | int64_t dstRegionOffset = (pingPongFlag == 0) ? this->batchSortLen_ : 0; | 204 | int64_t dstRegionOffset = (pingPongFlag == 0) ? this->batchSortLen_ : 0; |
| 208 | 205 | ||
| @@ -219,18 +216,20 @@ __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>:: | |||
| 219 | for (uint32_t j = 0; j < groupBlockCount; j++) { | 216 | for (uint32_t j = 0; j < groupBlockCount; j++) { |
| 220 | totalElem += ctx.elemCounts[j]; | 217 | totalElem += ctx.elemCounts[j]; |
| 221 | } | 218 | } |
| 222 | - mergedGroupElemCount = (totalElem > topKValue_) ? topKValue_ : totalElem; // TopK limitation | 219 | + mergedGroupElemCount = (totalElem > topKValue_) ? topKValue_ : totalElem; // TopK limitation |
| 223 | cumulativeOffset += AscendC::GetSortLen<ValueType>(mergedGroupElemCount); | 220 | cumulativeOffset += AscendC::GetSortLen<ValueType>(mergedGroupElemCount); |
| 224 | } else if (groupBlockCount == 1) { | 221 | } else if (groupBlockCount == 1) { |
| 225 | - mergedGroupElemCount = (ctx.elemCounts[0] > topKValue_) ? topKValue_ : ctx.elemCounts[0]; // TopK limitation | 222 | + mergedGroupElemCount = (ctx.elemCounts[0] > topKValue_) ? topKValue_ : ctx.elemCounts[0]; // TopK limitation |
| 226 | if (ctx.elemCounts[0] > 0) { | 223 | if (ctx.elemCounts[0] > 0) { |
| 227 | int64_t srcOffset = ctx.srcOffsets[0]; | 224 | int64_t srcOffset = ctx.srcOffsets[0]; |
| 228 | uint32_t remainElems = mergedGroupElemCount; | 225 | uint32_t remainElems = mergedGroupElemCount; |
| 229 | uint32_t srcChunkOffset = 0; | 226 | uint32_t srcChunkOffset = 0; |
| 230 | while (remainElems > 0) { | 227 | while (remainElems > 0) { |
| 231 | uint32_t chunkSize = (remainElems > this->blockSortSize_) ? this->blockSortSize_ : remainElems; | 228 | uint32_t chunkSize = (remainElems > this->blockSortSize_) ? this->blockSortSize_ : remainElems; |
| 232 | - if (chunkSize == 0) break; | 229 | + if (chunkSize == 0) |
| 233 | - this->CopyBlockChunk(srcOffset + AscendC::GetSortLen<ValueType>(srcChunkOffset), | 230 | + break; |
| 231 | + this->CopyBlockChunk( | ||
| 232 | + srcOffset + AscendC::GetSortLen<ValueType>(srcChunkOffset), | ||
| 234 | dstRegionOffset + cumulativeOffset + AscendC::GetSortLen<ValueType>(srcChunkOffset), chunkSize); | 233 | dstRegionOffset + cumulativeOffset + AscendC::GetSortLen<ValueType>(srcChunkOffset), chunkSize); |
| 235 | remainElems -= chunkSize; | 234 | remainElems -= chunkSize; |
| 236 | srcChunkOffset += chunkSize; | 235 | srcChunkOffset += chunkSize; |
| @@ -244,19 +243,22 @@ template <typename ValueType, typename IndexType, bool IsDescend> | |||
| 244 | __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::DoIncrementalMerge( | 243 | __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::DoIncrementalMerge( |
| 245 | int64_t dstOffsetBase, typename Base::MergeListContext& ctx) | 244 | int64_t dstOffsetBase, typename Base::MergeListContext& ctx) |
| 246 | { | 245 | { |
| 247 | - if (ctx.listCount > Sort::MERGE_LIST_MAX_NUM) return; | 246 | + if (ctx.listCount > Sort::MERGE_LIST_MAX_NUM) |
| 247 | + return; | ||
| 248 | 248 | ||
| 249 | for (uint32_t i = 0; i < ctx.listCount; i++) { | 249 | for (uint32_t i = 0; i < ctx.listCount; i++) { |
| 250 | ctx.remains[i] = ctx.elemCounts[i]; | 250 | ctx.remains[i] = ctx.elemCounts[i]; |
| 251 | } | 251 | } |
| 252 | - | 252 | + |
| 253 | uint32_t dstCumulativeOffset = 0, dstElemCount = 0, loopGuard = 0; | 253 | uint32_t dstCumulativeOffset = 0, dstElemCount = 0, loopGuard = 0; |
| 254 | - while (dstElemCount < topKValue_) { // TopK: early termination | 254 | + while (dstElemCount < topKValue_) { // TopK: early termination |
| 255 | uint32_t activeLists = 0; | 255 | uint32_t activeLists = 0; |
| 256 | for (uint32_t i = 0; i < ctx.listCount; i++) { | 256 | for (uint32_t i = 0; i < ctx.listCount; i++) { |
| 257 | - if (ctx.remains[i] > 0) activeLists++; | 257 | + if (ctx.remains[i] > 0) |
| 258 | + activeLists++; | ||
| 258 | } | 259 | } |
| 259 | - if (activeLists <= 1) break; | 260 | + if (activeLists <= 1) |
| 261 | + break; | ||
| 260 | 262 | ||
| 261 | LocalTensor<ValueType> ubMainInput = this->mergeInQueue_.template AllocTensor<ValueType>(); | 263 | LocalTensor<ValueType> ubMainInput = this->mergeInQueue_.template AllocTensor<ValueType>(); |
| 262 | uint16_t elementCountList[Sort::MERGE_LIST_MAX_NUM] = {0, 0, 0, 0}; | 264 | uint16_t elementCountList[Sort::MERGE_LIST_MAX_NUM] = {0, 0, 0, 0}; |
| @@ -267,8 +269,8 @@ __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>:: | |||
| 267 | LocalTensor<ValueType> dstLocal = this->mergeOutQueue_.template AllocTensor<ValueType>(); | 269 | LocalTensor<ValueType> dstLocal = this->mergeOutQueue_.template AllocTensor<ValueType>(); |
| 268 | 270 | ||
| 269 | uint32_t listSortedNums[Sort::MERGE_LIST_MAX_NUM] = {0, 0, 0, 0}; | 271 | uint32_t listSortedNums[Sort::MERGE_LIST_MAX_NUM] = {0, 0, 0, 0}; |
| 270 | - uint32_t mergedCount = this->ExecuteMrgSort(dstLocal, ubMainInputCalc, | 272 | + uint32_t mergedCount = this->ExecuteMrgSort(dstLocal, ubMainInputCalc, elementCountList, listSortedNums, |
| 271 | - elementCountList, listSortedNums, remainListNum); | 273 | + remainListNum); |
| 272 | if (mergedCount == 0) { | 274 | if (mergedCount == 0) { |
| 273 | this->mergeInQueue_.FreeTensor(ubMainInputCalc); | 275 | this->mergeInQueue_.FreeTensor(ubMainInputCalc); |
| 274 | this->mergeOutQueue_.FreeTensor(dstLocal); | 276 | this->mergeOutQueue_.FreeTensor(dstLocal); |
| @@ -280,10 +282,11 @@ __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>:: | |||
| 280 | 282 | ||
| 281 | LocalTensor<ValueType> dstLocalOut = this->mergeOutQueue_.template DeQue<ValueType>(); | 283 | LocalTensor<ValueType> dstLocalOut = this->mergeOutQueue_.template DeQue<ValueType>(); |
| 282 | uint32_t copyCount = mergedCount; | 284 | uint32_t copyCount = mergedCount; |
| 283 | - if (dstElemCount + copyCount > topKValue_) { // TopK: limit output | 285 | + if (dstElemCount + copyCount > topKValue_) { // TopK: limit output |
| 284 | copyCount = topKValue_ - dstElemCount; | 286 | copyCount = topKValue_ - dstElemCount; |
| 285 | } | 287 | } |
| 286 | - DataCopyExtParams outCopyParams{1, static_cast<uint32_t>(AscendC::GetSortLen<ValueType>(copyCount) * sizeof(ValueType)), 0, 0, 0}; | 288 | + DataCopyExtParams outCopyParams{ |
| 289 | + 1, static_cast<uint32_t>(AscendC::GetSortLen<ValueType>(copyCount) * sizeof(ValueType)), 0, 0, 0}; | ||
| 287 | DataCopyPad(this->cacheGm_[dstOffsetBase + dstCumulativeOffset], dstLocalOut, outCopyParams); | 290 | DataCopyPad(this->cacheGm_[dstOffsetBase + dstCumulativeOffset], dstLocalOut, outCopyParams); |
| 288 | this->mergeOutQueue_.FreeTensor(dstLocalOut); | 291 | this->mergeOutQueue_.FreeTensor(dstLocalOut); |
| 289 | 292 | ||
| @@ -297,7 +300,8 @@ __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>:: | |||
| 297 | } | 300 | } |
| 298 | dstCumulativeOffset += AscendC::GetSortLen<ValueType>(copyCount); | 301 | dstCumulativeOffset += AscendC::GetSortLen<ValueType>(copyCount); |
| 299 | dstElemCount += copyCount; | 302 | dstElemCount += copyCount; |
| 300 | - if (++loopGuard > this->maxMergeIterations_) break; | 303 | + if (++loopGuard > this->maxMergeIterations_) |
| 304 | + break; | ||
| 301 | } | 305 | } |
| 302 | 306 | ||
| 303 | CopyRemainingList(ctx, dstOffsetBase, dstCumulativeOffset, dstElemCount); | 307 | CopyRemainingList(ctx, dstOffsetBase, dstCumulativeOffset, dstElemCount); |
| @@ -307,18 +311,20 @@ template <typename ValueType, typename IndexType, bool IsDescend> | |||
| 307 | __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::CopyRemainingList( | 311 | __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::CopyRemainingList( |
| 308 | typename Base::MergeListContext& ctx, int64_t dstOffsetBase, uint32_t& dstCumulativeOffset, uint32_t& dstElemCount) | 312 | typename Base::MergeListContext& ctx, int64_t dstOffsetBase, uint32_t& dstCumulativeOffset, uint32_t& dstElemCount) |
| 309 | { | 313 | { |
| 310 | - for (uint32_t listIdx = 0; listIdx < ctx.listCount && dstElemCount < topKValue_; listIdx++) { // TopK: early stop | 314 | + for (uint32_t listIdx = 0; listIdx < ctx.listCount && dstElemCount < topKValue_; listIdx++) { // TopK: early stop |
| 311 | while (ctx.remains[listIdx] > 0) { | 315 | while (ctx.remains[listIdx] > 0) { |
| 312 | if (dstElemCount >= topKValue_) { | 316 | if (dstElemCount >= topKValue_) { |
| 313 | break; | 317 | break; |
| 314 | } | 318 | } |
| 315 | - uint32_t loadCount = (ctx.remains[listIdx] > this->blockSortSize_) ? this->blockSortSize_ : ctx.remains[listIdx]; | 319 | + uint32_t loadCount = (ctx.remains[listIdx] > this->blockSortSize_) ? this->blockSortSize_ : |
| 320 | + ctx.remains[listIdx]; | ||
| 316 | uint32_t remainingTopK = topKValue_ - dstElemCount; | 321 | uint32_t remainingTopK = topKValue_ - dstElemCount; |
| 317 | - loadCount = (loadCount > remainingTopK) ? remainingTopK : loadCount; // TopK: limit chunk | 322 | + loadCount = (loadCount > remainingTopK) ? remainingTopK : loadCount; // TopK: limit chunk |
| 318 | - if (loadCount == 0) break; | 323 | + if (loadCount == 0) |
| 319 | - | 324 | + break; |
| 320 | - this->CopyBlockChunk(ctx.srcOffsets[listIdx] + ctx.gmOffsets[listIdx], | 325 | + |
| 321 | - dstOffsetBase + dstCumulativeOffset, loadCount); | 326 | + this->CopyBlockChunk(ctx.srcOffsets[listIdx] + ctx.gmOffsets[listIdx], dstOffsetBase + dstCumulativeOffset, |
| 327 | + loadCount); | ||
| 322 | 328 | ||
| 323 | ctx.gmOffsets[listIdx] += AscendC::GetSortLen<ValueType>(loadCount); | 329 | ctx.gmOffsets[listIdx] += AscendC::GetSortLen<ValueType>(loadCount); |
| 324 | ctx.remains[listIdx] -= loadCount; | 330 | ctx.remains[listIdx] -= loadCount; |
| @@ -329,21 +335,22 @@ __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>:: | |||
| 329 | } | 335 | } |
| 330 | 336 | ||
| 331 | template <typename ValueType, typename IndexType, bool IsDescend> | 337 | template <typename ValueType, typename IndexType, bool IsDescend> |
| 332 | -__aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::ExtractAndCopyOut( | 338 | +__aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>::ExtractAndCopyOut(int64_t batchIdx, |
| 333 | - int64_t batchIdx, uint32_t resultRegion) | 339 | + uint32_t resultRegion) |
| 334 | { | 340 | { |
| 335 | - int64_t outputOffset = batchIdx * static_cast<int64_t>(topKValue_); // TopK: output offset | 341 | + int64_t outputOffset = batchIdx * static_cast<int64_t>(topKValue_); // TopK: output offset |
| 336 | 342 | ||
| 337 | int64_t cacheBatchOffset = (resultRegion == 1) ? this->batchSortLen_ : 0; | 343 | int64_t cacheBatchOffset = (resultRegion == 1) ? this->batchSortLen_ : 0; |
| 338 | 344 | ||
| 339 | - uint32_t totalElem = topKValue_; // TopK: limit output count | 345 | + uint32_t totalElem = topKValue_; // TopK: limit output count |
| 340 | uint32_t elemProcessed = 0; | 346 | uint32_t elemProcessed = 0; |
| 341 | uint32_t cacheOffset = 0; | 347 | uint32_t cacheOffset = 0; |
| 342 | 348 | ||
| 343 | while (elemProcessed < totalElem) { | 349 | while (elemProcessed < totalElem) { |
| 344 | - uint32_t elemCount = (elemProcessed + this->extractChunkSize_ <= totalElem) ? | 350 | + uint32_t elemCount = (elemProcessed + this->extractChunkSize_ <= totalElem) ? this->extractChunkSize_ : |
| 345 | - this->extractChunkSize_ : (totalElem - elemProcessed); | 351 | + (totalElem - elemProcessed); |
| 346 | - if (elemCount == 0) break; | 352 | + if (elemCount == 0) |
| 353 | + break; | ||
| 347 | 354 | ||
| 348 | this->ExtractAndCopyChunk(cacheBatchOffset, cacheOffset, outputOffset, elemProcessed, elemCount); | 355 | this->ExtractAndCopyChunk(cacheBatchOffset, cacheOffset, outputOffset, elemProcessed, elemCount); |
| 349 | 356 | ||
| @@ -354,4 +361,4 @@ __aicore__ inline void TopKMergeSortIntraCore<ValueType, IndexType, IsDescend>:: | |||
| 354 | 361 | ||
| 355 | } // namespace topkV2 | 362 | } // namespace topkV2 |
| 356 | 363 | ||
| 357 | -#endif // TOP_K_MERGE_SORT_INTRA_CORE_H | 364 | +#endif // TOP_K_MERGE_SORT_INTRA_CORE_H |
【函数规模】TryNonLastSmallAxis 由修改前 34 行增至 89 行 NBNC，超过本轮检视采用的 50 行阈值。
建议把 compact 路径判定、UB tile 计算以及最终 tiling 字段填写分别封装，避免策略与资源计算继续集中在一个函数。
统计范围为当前文件第 708–809 行；剔除空行/注释，保留签名、大括号及函数内预处理指令，不计函数前独立 template/修饰宏行。此为静态 NBNC 复核，非正式 CodeCheck 输出。