已合并
perf(index): 优化 KthValue、Median 与 NanMedian 选择及归并路径 #10411
黄晓彬创建于 9月14日
perf(index): 优化 KthValue、Median 与 NanMedian 选择及归并路径 #10411
已合并
共 27 个文件变更+2133-412
| @@ -353,9 +353,9 @@ static ge::graphStatus SetRadixMoreCoreTiling(gert::TilingContext* context, Sort | |||
| 353 | } | 353 | } |
| 354 | 354 | ||
| 355 | static ge::graphStatus SetKthValueMergeSortTiling(gert::TilingContext* context, SortKthTileInfo& info, | 355 | static ge::graphStatus SetKthValueMergeSortTiling(gert::TilingContext* context, SortKthTileInfo& info, |
| 356 | - uint32_t& blockDim, uint64_t& schId) | 356 | + uint32_t& blockDim, uint64_t& schId, bool useUbCapacity) |
| 357 | { | 357 | { |
| 358 | - OP_CHECK_IF(!ComputeMergeSortTiling(context, info, static_cast<uint32_t>(sizeof(uint32_t))), | 358 | + OP_CHECK_IF(!ComputeMergeSortTiling(context, info, static_cast<uint32_t>(sizeof(uint32_t)), useUbCapacity), |
| 359 | OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "ComputeMergeSortTiling", "false", | 359 | OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "ComputeMergeSortTiling", "false", |
| 360 | "The value of ComputeMergeSortTiling must be true."), | 360 | "The value of ComputeMergeSortTiling must be true."), |
| 361 | return ge::GRAPH_FAILED); | 361 | return ge::GRAPH_FAILED); |
| @@ -377,7 +377,8 @@ static ge::graphStatus SetMergeMoreCoreTiling(gert::TilingContext* context, Sort | |||
| 377 | "The value of ComputeMergeMoreCoreTiling must be true."), | 377 | "The value of ComputeMergeMoreCoreTiling must be true."), |
| 378 | return ge::GRAPH_FAILED); | 378 | return ge::GRAPH_FAILED); |
| 379 | blockDim = info.coreNumNeed; | 379 | blockDim = info.coreNumNeed; |
| 380 | - OP_LOGI("KthValueMergeMoreCoreTiling", "maxDealingNum: %u", info.keyParams0); | 380 | + OP_LOGI("KthValueMergeMoreCoreTiling", "maxDealingNum: %u, syncMergeBlockSize: %u", info.keyParams0, |
| 381 | + info.keyParams1); | ||
| 381 | return ge::GRAPH_SUCCESS; | 382 | return ge::GRAPH_SUCCESS; |
| 382 | } | 383 | } |
| 383 | 384 | ||
| @@ -589,7 +590,76 @@ static bool TryRadixOneCore(gert::TilingContext* context, const SortKthTileInfo& | |||
| 589 | return true; | 590 | return true; |
| 590 | } | 591 | } |
| 591 | 592 | ||
| 592 | -static bool IsRadixSelectProfitable(const SortKthTileInfo& info, int64_t kthIndex) | 593 | +// Empirical route policy, including the power-of-two crossovers below. These are not |
| 594 | +// hardware capacity guarantees; each tiling builder independently checks its UB footprint. | ||
| 595 | +namespace KthRoutePolicy { | ||
| 596 | +constexpr int64_t RESIDENT_MIN_AXIS = 4096; | ||
| 597 | +constexpr int64_t RESIDENT_MAX_AXIS = 16384; | ||
| 598 | +constexpr uint64_t RADIX_SELECT_MAX_TILE_ELEMS = 32768UL; | ||
| 599 | +constexpr int64_t INT32_RESIDENT_MIN_AXIS = 1024; | ||
| 600 | +constexpr int64_t FLOAT_SELECT_MIN_AXIS = 1024; | ||
| 601 | +constexpr int64_t RESIDENT_HISTOGRAM_MAX_AXIS = 2048; | ||
| 602 | +constexpr int64_t RESIDENT_HISTOGRAM_SINGLE_WAVE_MAX_AXIS = 8192; | ||
| 603 | +constexpr int64_t RESIDENT_HISTOGRAM_MIN_ROWS_PER_CORE = 8; | ||
| 604 | +// Keep the validated resident-half range and tile/core split stable. Growing this cap | ||
| 605 | +// changes both route eligibility and cores per row, even when the extra elements fit UB. | ||
| 606 | +constexpr uint64_t HALF_SELECT_MAX_TILE_ELEMS = 49152UL; | ||
| 607 | +// Measured merge/radix crossover, not an alignment constraint. Cast + merge is not | ||
| 608 | +// selected for the shorter rows handled by the earlier narrow/byte policies. | ||
| 609 | +constexpr int64_t INT16_MERGE_MIN_AXIS = 192; | ||
| 610 | +constexpr int64_t INT16_MERGE_MAX_AXIS = 512; | ||
| 611 | +// KthValue batches wider byte rows than Sort. Keep its policy explicit rather | ||
| 612 | +// than copying and patching the shared Sort rule at runtime. | ||
| 613 | +// Signed and unsigned byte keys use the same one-byte radix pipeline. | ||
| 614 | +constexpr SmallAxisRule BYTE_SMALL_AXIS_RULES[] = { | ||
| 615 | + {ge::DT_INT8, 8, 384, {{4, 2}, {8, 7}, {0, 0}}, {{3, 8}, {64, 7}, {384, 16}, {0, 0}}}, | ||
| 616 | + {ge::DT_UINT8, 8, 384, {{4, 2}, {8, 7}, {0, 0}}, {{3, 8}, {64, 7}, {384, 16}, {0, 0}}}, | ||
| 617 | +}; | ||
| 618 | +static_assert(HALF_SELECT_MAX_TILE_ELEMS <= std::numeric_limits<uint16_t>::max(), | ||
| 619 | + "A per-tile cumulative histogram must not overflow uint16_t"); | ||
| 620 | +static_assert(RADIX_SELECT_MAX_TILE_ELEMS <= std::numeric_limits<uint16_t>::max(), | ||
| 621 | + "A per-tile cumulative histogram must not overflow uint16_t"); | ||
| 622 | +} // namespace KthRoutePolicy | ||
| 623 | + | ||
| 624 | +static bool IsResidentIntegerSelect(const SortKthTileInfo& info) | ||
| 625 | +{ | ||
| 626 | + using namespace KthRoutePolicy; | ||
| 627 | + if (info.dataType == ge::DT_INT32 && info.lastAxis >= INT32_RESIDENT_MIN_AXIS && | ||
| 628 | + info.lastAxis < RESIDENT_MIN_AXIS) { | ||
| 629 | + return true; | ||
| 630 | + } | ||
| 631 | + // Int16 rows can use the full existing tile capacity; the tiling builder | ||
| 632 | + // still limits allocation by available UB before selecting the kernel. | ||
| 633 | + int64_t maxResidentAxis = info.dataType == ge::DT_INT16 ? static_cast<int64_t>(RADIX_SELECT_MAX_TILE_ELEMS) : | ||
| 634 | + RESIDENT_MAX_AXIS; | ||
| 635 | + return (info.dataType == ge::DT_INT16 || info.dataType == ge::DT_INT32 || info.dataType == ge::DT_UINT32 || | ||
| 636 | + info.dataType == ge::DT_INT64 || info.dataType == ge::DT_UINT64) && | ||
| 637 | + info.lastAxis >= RESIDENT_MIN_AXIS && info.lastAxis <= maxResidentAxis; | ||
| 638 | +} | ||
| 639 | + | ||
| 640 | +static bool IsResidentFloatSelect(const SortKthTileInfo& info, uint32_t medianMode) | ||
| 641 | +{ | ||
| 642 | + // Resident rows avoid full-sort traffic. NaN propagation keeps its existing multi-core schedule. | ||
| 643 | + using namespace KthRoutePolicy; | ||
| 644 | + return info.dataType == ge::DT_FLOAT && (info.unsortedDim == 1 || medianMode != MEDIAN_MODE_PROPAGATE_NAN) && | ||
| 645 | + info.lastAxis > RESIDENT_MIN_AXIS && info.lastAxis <= RESIDENT_MAX_AXIS; | ||
| 646 | +} | ||
| 647 | + | ||
| 648 | +static bool IsKthFloatSelect(const SortKthTileInfo& info, uint32_t medianMode) | ||
| 649 | +{ | ||
| 650 | + using namespace KthRoutePolicy; | ||
| 651 | + if (medianMode != MEDIAN_MODE_STATIC || info.lastAxis < FLOAT_SELECT_MIN_AXIS) { | ||
| 652 | + return false; | ||
| 653 | + } | ||
| 654 | + // Extend small FP32 rows and resident 16-bit float rows; retain median policies. | ||
| 655 | + if (info.dataType == ge::DT_FLOAT) { | ||
黄 | |||
| 656 | + return info.lastAxis <= RESIDENT_MIN_AXIS; | ||
| 657 | + } | ||
| 658 | + return (info.dataType == ge::DT_FLOAT16 || info.dataType == ge::DT_BF16) && | ||
| 659 | + info.lastAxis <= static_cast<int64_t>(HALF_SELECT_MAX_TILE_ELEMS); | ||
| 660 | +} | ||
| 661 | + | ||
| 662 | +static bool IsRadixSelectProfitable(const SortKthTileInfo& info, int64_t kthIndex, uint32_t medianMode) | ||
| 593 | { | 663 | { |
| 594 | if (kthIndex < 0 || info.lastAxis <= 0 || info.unsortedDim <= 0) { | 664 | if (kthIndex < 0 || info.lastAxis <= 0 || info.unsortedDim <= 0) { |
| 595 | return false; | 665 | return false; |
| @@ -600,6 +670,12 @@ static bool IsRadixSelectProfitable(const SortKthTileInfo& info, int64_t kthInde | |||
| 600 | return false; | 670 | return false; |
| 601 | } | 671 | } |
| 602 | 672 | ||
| 673 | + // Retained rows can select in UB without scattering a complete sorted row. | ||
| 674 | + if (IsResidentIntegerSelect(info) || IsResidentFloatSelect(info, medianMode) || | ||
| 675 | + IsKthFloatSelect(info, medianMode)) { | ||
| 676 | + return true; | ||
| 677 | + } | ||
| 678 | + | ||
| 603 | constexpr uint64_t radixSelectMinAxis = 65536UL; | 679 | constexpr uint64_t radixSelectMinAxis = 65536UL; |
| 604 | constexpr uint64_t radixSelectHugeAxis = 1000000UL; | 680 | constexpr uint64_t radixSelectHugeAxis = 1000000UL; |
| 605 | constexpr uint64_t radixSelectInt64SmallAxis = 4096UL; | 681 | constexpr uint64_t radixSelectInt64SmallAxis = 4096UL; |
| @@ -665,26 +741,33 @@ static uint64_t GetRadixSelectFixedBytes() | |||
| 665 | ((reservedRawBytes + radixSelectReserveAlign - 1UL) / radixSelectReserveAlign) * radixSelectReserveAlign; | 741 | ((reservedRawBytes + radixSelectReserveAlign - 1UL) / radixSelectReserveAlign) * radixSelectReserveAlign; |
| 666 | } | 742 | } |
| 667 | 743 | ||
| 668 | -static bool ComputeRadixSelectTileElems(const SortKthTileInfo& info, uint64_t& tileElems) | 744 | +static bool ComputeRadixSelectTileElems(const SortKthTileInfo& info, uint32_t medianMode, uint64_t& tileElems) |
| 669 | { | 745 | { |
| 670 | - constexpr uint64_t radixSelectMaxTileElems = 32768UL; | ||
| 671 | constexpr uint64_t radixSelectMinAlignElems = 256UL; | 746 | constexpr uint64_t radixSelectMinAlignElems = 256UL; |
| 672 | uint64_t usableUb = ComputeUbAfterSimtReserve(info.ubSize); | 747 | uint64_t usableUb = ComputeUbAfterSimtReserve(info.ubSize); |
| 673 | - uint64_t bytesPerElem = static_cast<uint64_t>(info.dtypeSize) * 2UL; | 748 | + // Retain the established slice/core plan outside the resident integer range. |
| 749 | + // A larger tile there can reduce cores per row and lose parallelism. | ||
| 750 | + uint64_t bytesPerElem = static_cast<uint64_t>(info.dtypeSize) * (IsResidentIntegerSelect(info) ? 1UL : 2UL); | ||
| 674 | uint64_t fixedBytes = GetRadixSelectFixedBytes(); | 751 | uint64_t fixedBytes = GetRadixSelectFixedBytes(); |
| 675 | if (bytesPerElem == 0U || usableUb <= fixedBytes) { | 752 | if (bytesPerElem == 0U || usableUb <= fixedBytes) { |
| 676 | return false; | 753 | return false; |
| 677 | } | 754 | } |
| 678 | - tileElems = std::min<uint64_t>((usableUb - fixedBytes) / bytesPerElem, radixSelectMaxTileElems); | 755 | + // Only the newly enabled resident half/bfloat range needs the larger tile. |
| 756 | + // Both caps (32768/49152) are below 65536: each tile clears its uint16 | ||
| 757 | + // cumulative histogram, then widens it before accumulating across tiles. | ||
| 758 | + uint64_t maxTileElems = IsKthFloatSelect(info, medianMode) && info.dataType != ge::DT_FLOAT ? | ||
| 759 | + KthRoutePolicy::HALF_SELECT_MAX_TILE_ELEMS : | ||
| 760 | + KthRoutePolicy::RADIX_SELECT_MAX_TILE_ELEMS; | ||
| 761 | + tileElems = std::min<uint64_t>((usableUb - fixedBytes) / bytesPerElem, maxTileElems); | ||
| 679 | uint64_t alignElems = std::max<uint64_t>(radixSelectMinAlignElems, info.blockUbSize / info.dtypeSize); | 762 | uint64_t alignElems = std::max<uint64_t>(radixSelectMinAlignElems, info.blockUbSize / info.dtypeSize); |
| 680 | tileElems = tileElems / alignElems * alignElems; | 763 | tileElems = tileElems / alignElems * alignElems; |
| 681 | tileElems = std::min<uint64_t>(tileElems, static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())); | 764 | tileElems = std::min<uint64_t>(tileElems, static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())); |
| 682 | return tileElems != 0U; | 765 | return tileElems != 0U; |
| 683 | } | 766 | } |
| 684 | 767 | ||
| 685 | -static bool ComputeRadixSelectPlan(const SortKthTileInfo& info, RadixSelectPlan& plan) | 768 | +static bool ComputeRadixSelectPlan(const SortKthTileInfo& info, uint32_t medianMode, RadixSelectPlan& plan) |
| 686 | { | 769 | { |
| 687 | - if (!ComputeRadixSelectTileElems(info, plan.tileElems)) { | 770 | + if (!ComputeRadixSelectTileElems(info, medianMode, plan.tileElems)) { |
| 688 | return false; | 771 | return false; |
| 689 | } | 772 | } |
| 690 | plan.rowsParallel = static_cast<uint32_t>( | 773 | plan.rowsParallel = static_cast<uint32_t>( |
| @@ -723,12 +806,20 @@ static bool ComputeRadixSelectWorkspace(const SortKthTileInfo& info, RadixSelect | |||
| 723 | static bool TryRadixSelect(gert::TilingContext* context, SortKthTileInfo& info, KthValueTilingData* tilingData, | 806 | static bool TryRadixSelect(gert::TilingContext* context, SortKthTileInfo& info, KthValueTilingData* tilingData, |
| 724 | uint32_t& blockDim, uint64_t& schId) | 807 | uint32_t& blockDim, uint64_t& schId) |
| 725 | { | 808 | { |
| 726 | - if (!IsRadixSelectProfitable(info, tilingData->kthIndex) || info.isNonLastAxis || | 809 | + if (info.isNonLastAxis || info.lastAxis <= static_cast<int64_t>(SMALL_AXIS_THRESHOLD) || info.maxCoreNum == 0U) { |
| 727 | - info.lastAxis <= static_cast<int64_t>(SMALL_AXIS_THRESHOLD) || info.maxCoreNum == 0U) { | ||
| 728 | return false; | 810 | return false; |
| 729 | } | 811 | } |
| 730 | RadixSelectPlan plan; | 812 | RadixSelectPlan plan; |
| 731 | - if (!ComputeRadixSelectPlan(info, plan) || !ComputeRadixSelectWorkspace(info, plan)) { | 813 | + if (!ComputeRadixSelectPlan(info, tilingData->medianMode, plan) || !ComputeRadixSelectWorkspace(info, plan)) { |
| 814 | + return false; | ||
| 815 | + } | ||
| 816 | + // A single B64 row can skip common prefix bytes across its cooperating cores when | ||
| 817 | + // every slice remains in UB. Streaming rows retain the established routing policy. | ||
| 818 | + // ParseKthValueShapeInfo has already checked 1 <= k <= lastAxis before | ||
| 819 | + // kthIndex = k - 1 is initialized; this exception only bypasses profitability. | ||
| 820 | + bool retainedB64Row = (info.dataType == ge::DT_INT64 || info.dataType == ge::DT_UINT64) && info.unsortedDim == 1 && | ||
| 821 | + plan.coresPerRow > 1U && plan.tileCount <= plan.coresPerRow; | ||
| 822 | + if (!IsRadixSelectProfitable(info, tilingData->kthIndex, tilingData->medianMode) && !retainedB64Row) { | ||
| 732 | return false; | 823 | return false; |
| 733 | } | 824 | } |
| 734 | size_t* userWorkspaceSize = context->GetWorkspaceSizes(1); | 825 | size_t* userWorkspaceSize = context->GetWorkspaceSizes(1); |
| @@ -750,7 +841,18 @@ static bool TryRadixSelect(gert::TilingContext* context, SortKthTileInfo& info, | |||
| 750 | info = candidate; | 841 | info = candidate; |
| 751 | blockDim = plan.blockDim; | 842 | blockDim = plan.blockDim; |
| 752 | PlanToTilingData(info, tilingData); | 843 | PlanToTilingData(info, tilingData); |
| 753 | - schId = KTH_VALUE_SCHID_RADIX_SELECT; | 844 | + // Use local histograms for batched medium rows or longer rows that fit in one core wave. |
| 845 | + bool batchedResidentHistogram = info.lastAxis >= KthRoutePolicy::FLOAT_SELECT_MIN_AXIS && | ||
| 846 | + info.lastAxis <= KthRoutePolicy::RESIDENT_HISTOGRAM_MAX_AXIS && | ||
| 847 | + info.unsortedDim >= static_cast<int64_t>(info.maxCoreNum) * | ||
| 848 | + KthRoutePolicy::RESIDENT_HISTOGRAM_MIN_ROWS_PER_CORE; | ||
| 849 | + bool singleWaveResidentHistogram = info.lastAxis > KthRoutePolicy::RESIDENT_MIN_AXIS && | ||
| 850 | + info.lastAxis <= KthRoutePolicy::RESIDENT_HISTOGRAM_SINGLE_WAVE_MAX_AXIS && | ||
| 851 | + info.unsortedDim <= static_cast<int64_t>(info.maxCoreNum); | ||
| 852 | + bool useResidentHistogram = tilingData->medianMode == MEDIAN_MODE_STATIC && info.dataType == ge::DT_FLOAT16 && | ||
| 853 | + (batchedResidentHistogram || singleWaveResidentHistogram) && plan.coresPerRow == 1U && | ||
| 854 | + plan.tileCount == 1U; | ||
| 855 | + schId = useResidentHistogram ? KTH_VALUE_SCHID_RESIDENT_HISTOGRAM : KTH_VALUE_SCHID_RADIX_SELECT; | ||
| 754 | userWorkspaceSize[0] = static_cast<size_t>(WORK_SPACE_SIZE + plan.workspace); | 856 | userWorkspaceSize[0] = static_cast<size_t>(WORK_SPACE_SIZE + plan.workspace); |
| 755 | context->SetScheduleMode(1); | 857 | context->SetScheduleMode(1); |
| 756 | return true; | 858 | return true; |
| @@ -788,8 +890,61 @@ static bool TrySmallAxisShortRankSelect(const SortKthTileInfo& info, KthValueTil | |||
| 788 | return true; | 890 | return true; |
| 789 | } | 891 | } |
| 790 | 892 | ||
| 893 | +static bool SelectKthSmallAxisRoute(const SortKthTileInfo& info, SmallAxisRoutePlan& plan) | ||
| 894 | +{ | ||
| 895 | + for (const auto& rule : KthRoutePolicy::BYTE_SMALL_AXIS_RULES) { | ||
| 896 | + if (info.dataType == rule.dtype) { | ||
| 897 | + return SelectSmallAxisRoute(info, plan, &rule); | ||
| 898 | + } | ||
| 899 | + } | ||
| 900 | + return SelectSmallAxisRoute(info, plan); | ||
| 901 | +} | ||
| 902 | + | ||
| 903 | +static bool SetNarrowSelectBatch(const SortKthTileInfo& info, uint32_t medianMode, SmallAxisRoutePlan& plan) | ||
| 904 | +{ | ||
| 905 | + constexpr int64_t maxByteAxis = 384; | ||
| 906 | + constexpr int64_t maxSignedByteAxis = 512; | ||
| 907 | + constexpr int64_t minFloatAxis = 32; | ||
| 908 | + constexpr int64_t maxFloatAxis = 256; | ||
| 909 | + constexpr uint64_t rowsPerWave = 32UL; | ||
| 910 | + constexpr uint64_t maxInputBytes = 131072UL; | ||
| 911 | + bool isByte = info.dataType == ge::DT_INT8 || info.dataType == ge::DT_UINT8; | ||
| 912 | + bool isFloat = info.dataType == ge::DT_FLOAT16 || info.dataType == ge::DT_BF16; | ||
| 913 | + if (info.isNonLastAxis || medianMode != MEDIAN_MODE_STATIC || info.lastAxis <= 0 || info.maxCoreNum == 0U || | ||
| 914 | + (!isByte && !isFloat) || info.unsortedDim < static_cast<int64_t>(info.maxCoreNum * rowsPerWave)) { | ||
| 915 | + return false; | ||
| 916 | + } | ||
| 917 | + int64_t maxAxis = isFloat ? maxFloatAxis : (info.dataType == ge::DT_INT8 ? maxSignedByteAxis : maxByteAxis); | ||
| 918 | + if (info.lastAxis > maxAxis || (isFloat && info.lastAxis < minFloatAxis) || | ||
| 919 | + (isByte && info.lastAxis <= maxByteAxis && plan.kind != SmallAxisRouteKind::TWO_STAGE)) { | ||
| 920 | + return false; | ||
| 921 | + } | ||
| 922 | + uint64_t elementBytes = isByte ? sizeof(uint8_t) : sizeof(uint16_t); | ||
| 923 | + if (info.dtypeSize != elementBytes) { | ||
| 924 | + return false; | ||
| 925 | + } | ||
| 926 | + uint64_t inputBytes = std::min<uint64_t>(ComputeUbAfterSimtReserve(info.ubSize), maxInputBytes); | ||
| 927 | + inputBytes = inputBytes / rowsPerWave * rowsPerWave; | ||
| 928 | + uint64_t capacity = inputBytes / (static_cast<uint64_t>(info.lastAxis) * elementBytes); | ||
| 929 | + if (capacity == 0U) { | ||
| 930 | + return false; | ||
| 931 | + } | ||
| 932 | + uint64_t rows = Ops::Base::CeilDiv(static_cast<uint64_t>(info.unsortedDim), static_cast<uint64_t>(info.maxCoreNum)); | ||
| 933 | + rows = std::min<uint64_t>(Ops::Base::CeilAlign(rows, rowsPerWave), capacity); | ||
| 934 | + uint64_t batches = Ops::Base::CeilDiv(static_cast<uint64_t>(info.unsortedDim), rows); | ||
| 935 | + if (batches > std::numeric_limits<uint32_t>::max()) { | ||
| 936 | + return false; | ||
| 937 | + } | ||
| 938 | + plan.batchSize = static_cast<uint32_t>(rows); | ||
| 939 | + plan.batchNum = static_cast<uint32_t>(batches); | ||
| 940 | + plan.blockDim = std::min<uint32_t>(plan.batchNum, info.maxCoreNum); | ||
| 941 | + plan.tmpUbSize = 0; | ||
| 942 | + plan.kind = SmallAxisRouteKind::TWO_STAGE; | ||
| 943 | + return true; | ||
| 944 | +} | ||
| 945 | + | ||
| 791 | static bool TrySmallAxis(gert::TilingContext* context, SortKthTileInfo& info, KthValueTilingData* tilingData, | 946 | static bool TrySmallAxis(gert::TilingContext* context, SortKthTileInfo& info, KthValueTilingData* tilingData, |
| 792 | - uint32_t& blockDim, uint64_t& schId) | 947 | + uint32_t& blockDim, uint64_t& schId, bool isMedianOp) |
| 793 | { | 948 | { |
| 794 | if (info.lastAxis > static_cast<int64_t>(SMALL_AXIS_THRESHOLD)) { | 949 | if (info.lastAxis > static_cast<int64_t>(SMALL_AXIS_THRESHOLD)) { |
| 795 | return false; | 950 | return false; |
| @@ -804,8 +959,9 @@ static bool TrySmallAxis(gert::TilingContext* context, SortKthTileInfo& info, Kt | |||
| 804 | return true; | 959 | return true; |
| 805 | } | 960 | } |
| 806 | SmallAxisRoutePlan plan; | 961 | SmallAxisRoutePlan plan; |
| 807 | - bool selected = info.isNonLastAxis ? SelectNonLastSmallAxisRoute(info, plan) : SelectSmallAxisRoute(info, plan); | 962 | + bool selected = info.isNonLastAxis ? SelectNonLastSmallAxisRoute(info, plan) : SelectKthSmallAxisRoute(info, plan); |
| 808 | - if (!selected) { | 963 | + bool narrowSelect = !isMedianOp && SetNarrowSelectBatch(info, tilingData->medianMode, plan); |
| 964 | + if (!selected && !narrowSelect) { | ||
| 809 | return false; | 965 | return false; |
| 810 | } | 966 | } |
| 811 | if (info.isNonLastAxis) { | 967 | if (info.isNonLastAxis) { |
| @@ -814,24 +970,78 @@ static bool TrySmallAxis(gert::TilingContext* context, SortKthTileInfo& info, Kt | |||
| 814 | } | 970 | } |
| 815 | } else { | 971 | } else { |
| 816 | FillSmallAxisTiling(tilingData, plan, static_cast<uint32_t>(info.lastAxis), blockDim); | 972 | FillSmallAxisTiling(tilingData, plan, static_cast<uint32_t>(info.lastAxis), blockDim); |
| 973 | + if (narrowSelect) { | ||
| 974 | + // Small-axis keyParams3 selects the narrow-key value-domain kernel and its input-only UB layout. | ||
| 975 | + constexpr uint32_t narrowSelectMode = 1U; | ||
| 976 | + tilingData->keyParams3 = narrowSelectMode; | ||
| 977 | + } | ||
| 817 | } | 978 | } |
| 818 | schId = plan.kind == SmallAxisRouteKind::TWO_STAGE ? KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE : | 979 | schId = plan.kind == SmallAxisRouteKind::TWO_STAGE ? KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE : |
| 819 | KTH_VALUE_SCHID_SMALL_AXIS_INSERTION; | 980 | KTH_VALUE_SCHID_SMALL_AXIS_INSERTION; |
| 820 | return true; | 981 | return true; |
| 821 | } | 982 | } |
| 822 | 983 | ||
| 984 | +static bool FitsInt16MergeUb(const SortKthTileInfo& info) | ||
| 985 | +{ | ||
| 986 | + uint64_t rows = info.keyParams0; | ||
| 987 | + uint64_t alignedAxis = info.keyParams3; | ||
| 988 | + uint64_t buffers = info.keyParams4 == 0U ? 1U : info.keyParams4; | ||
| 989 | + uint64_t compactValues = Ops::Base::CeilAlign(rows * sizeof(int16_t), static_cast<uint64_t>(info.blockUbSize)); | ||
| 990 | + uint64_t compactIndices = Ops::Base::CeilAlign(rows * sizeof(int64_t), static_cast<uint64_t>(info.blockUbSize)); | ||
| 991 | + // Input/value/index queues and compact outputs, two proposal buffers, source | ||
| 992 | + // indices, and the exact int16-to-float cast batch mirror Kernel::Init. | ||
| 993 | + uint64_t queues = buffers * (2UL * info.keyParams1 + info.keyParams2 + compactValues + compactIndices); | ||
| 994 | + uint64_t proposalsAndIndices = alignedAxis * (2UL * MERGE_SORT_DATA_BYTES + sizeof(uint32_t)); | ||
| 995 | + uint64_t castBatch = alignedAxis * rows * sizeof(float); | ||
| 996 | + uint64_t fullOutputUb = queues + proposalsAndIndices + castBatch; | ||
| 997 | + uint64_t mergeBatchRows = alignedAxis >= KTH_INT16_MERGE_BATCH_MIN_AXIS && | ||
| 998 | + alignedAxis <= KTH_INT16_MERGE_BATCH_MAX_AXIS && | ||
| 999 | + rows >= KTH_INT16_MERGE_BATCH_ROWS ? | ||
| 1000 | + KTH_INT16_MERGE_BATCH_ROWS : | ||
| 1001 | + 1U; | ||
| 1002 | + // Non-median int16 keeps only one Sort32 block per output row. The freed | ||
| 1003 | + // space holds per-row proposals and repeated source indices for batched merging. | ||
| 1004 | + uint64_t compactRowBytes = SORT32_SMALL_AXIS_THRESHOLD * (sizeof(int16_t) + sizeof(uint32_t)); | ||
| 1005 | + uint64_t compactQueues = buffers * (info.keyParams1 + rows * compactRowBytes + compactValues + compactIndices); | ||
| 1006 | + uint64_t compactOutputUb = compactQueues + proposalsAndIndices * mergeBatchRows + castBatch; | ||
| 1007 | + // Median retains full outputs, so validate both kernel layouts before accepting the route. | ||
| 1008 | + return std::max(fullOutputUb, compactOutputUb) <= info.ubSize; | ||
| 1009 | +} | ||
| 1010 | + | ||
| 1011 | +static bool PreferKthMergeMoreCore(const SortKthTileInfo& info) | ||
| 1012 | +{ | ||
| 1013 | + // Splitting every row needs at least two cores per row. Beyond that budget, | ||
| 1014 | + // prefer the existing row-parallel plan when its UB and occupancy checks pass. | ||
| 1015 | + constexpr uint32_t minCoresPerSplitRow = 2U; | ||
| 1016 | + bool canSplitEveryRow = info.unsortedDim <= static_cast<int64_t>(info.maxCoreNum / minCoresPerSplitRow); | ||
| 1017 | + return IsMergeMoreCoreProfitable(info.dataType, info.lastAxis, info.unsortedDim, info.maxCoreNum) && | ||
| 1018 | + (canSplitEveryRow || | ||
| 1019 | + !IsMergeIntraCoreSupported(info.dataType, info.lastAxis, info.unsortedDim, info.maxCoreNum, info.ubSize)); | ||
| 1020 | +} | ||
| 1021 | + | ||
| 823 | static bool TryMerge(gert::TilingContext* context, SortKthTileInfo& info, KthValueTilingData* tilingData, | 1022 | static bool TryMerge(gert::TilingContext* context, SortKthTileInfo& info, KthValueTilingData* tilingData, |
| 824 | uint32_t& blockDim, uint64_t& schId) | 1023 | uint32_t& blockDim, uint64_t& schId) |
| 825 | { | 1024 | { |
| 826 | - if (IsMergeSortSupported(info.dataType, info.lastAxis)) { | 1025 | + // Dispatch, compact extraction and GatherKthInt16 implement signed int16 only. |
| 1026 | + // uint16 needs its own dispatch/cast validation; byte types already have a one-byte/narrow route. | ||
| 1027 | + // int16 is exactly representable in float. Batch medium short rows through | ||
| 1028 | + // the merge hardware instead of paying two radix passes and scalar waits per row. | ||
| 1029 | + bool useInt16Merge = info.dataType == ge::DT_INT16 && info.lastAxis > KthRoutePolicy::INT16_MERGE_MIN_AXIS && | ||
| 1030 | + info.lastAxis <= KthRoutePolicy::INT16_MERGE_MAX_AXIS && | ||
| 1031 | + info.unsortedDim >= static_cast<int64_t>(info.maxCoreNum); | ||
| 1032 | + if (useInt16Merge || IsMergeSortSupported(info.dataType, info.lastAxis)) { | ||
| 827 | SortKthTileInfo candidate = info; | 1033 | SortKthTileInfo candidate = info; |
| 828 | - if (SetKthValueMergeSortTiling(context, candidate, blockDim, schId) == ge::GRAPH_SUCCESS) { | 1034 | + if (SetKthValueMergeSortTiling(context, candidate, blockDim, schId, |
| 1035 | + tilingData->medianMode == MEDIAN_MODE_STATIC) == ge::GRAPH_SUCCESS) { | ||
| 1036 | + if (useInt16Merge && !FitsInt16MergeUb(candidate)) { | ||
| 1037 | + return false; | ||
| 1038 | + } | ||
| 829 | info = candidate; | 1039 | info = candidate; |
| 830 | PlanToTilingData(info, tilingData); | 1040 | PlanToTilingData(info, tilingData); |
| 831 | return true; | 1041 | return true; |
| 832 | } | 1042 | } |
| 833 | } | 1043 | } |
| 834 | - if (IsMergeMoreCoreSupported(info.dataType, info.lastAxis, info.unsortedDim, info.maxCoreNum)) { | 1044 | + if (PreferKthMergeMoreCore(info)) { |
| 835 | SortKthTileInfo candidate = info; | 1045 | SortKthTileInfo candidate = info; |
| 836 | if (SetMergeMoreCoreTiling(context, candidate, blockDim) == ge::GRAPH_SUCCESS) { | 1046 | if (SetMergeMoreCoreTiling(context, candidate, blockDim) == ge::GRAPH_SUCCESS) { |
| 837 | info = candidate; | 1047 | info = candidate; |
| @@ -873,12 +1083,13 @@ static bool TryNonLastSmallAxis(gert::TilingContext* context, const SortKthTileI | |||
| 873 | // Route selection and finalization | 1083 | // Route selection and finalization |
| 874 | // ============================================================================= | 1084 | // ============================================================================= |
| 875 | static ge::graphStatus SelectKthValueRoute(gert::TilingContext* context, SortKthTileInfo& info, | 1085 | static ge::graphStatus SelectKthValueRoute(gert::TilingContext* context, SortKthTileInfo& info, |
| 876 | - KthValueTilingData* tilingData, uint32_t& blockDim, uint64_t& schId) | 1086 | + KthValueTilingData* tilingData, uint32_t& blockDim, uint64_t& schId, |
| 1087 | + bool isMedianOp) | ||
| 877 | { | 1088 | { |
| 878 | if (TrySmallAxisShortRankSelect(info, tilingData, blockDim, schId)) { | 1089 | if (TrySmallAxisShortRankSelect(info, tilingData, blockDim, schId)) { |
| 879 | return ge::GRAPH_SUCCESS; | 1090 | return ge::GRAPH_SUCCESS; |
| 880 | } | 1091 | } |
| 881 | - if (TrySmallAxis(context, info, tilingData, blockDim, schId)) { | 1092 | + if (TrySmallAxis(context, info, tilingData, blockDim, schId, isMedianOp)) { |
| 882 | return ge::GRAPH_SUCCESS; | 1093 | return ge::GRAPH_SUCCESS; |
| 883 | } | 1094 | } |
| 884 | if (TryNonLastSmallAxis(context, info, tilingData, blockDim, schId)) { | 1095 | if (TryNonLastSmallAxis(context, info, tilingData, blockDim, schId)) { |
| @@ -890,6 +1101,13 @@ static ge::graphStatus SelectKthValueRoute(gert::TilingContext* context, SortKth | |||
| 890 | "The value of sortAxis must be the last axis or meet no-transpose schedule constraints."); | 1101 | "The value of sortAxis must be the last axis or meet no-transpose schedule constraints."); |
| 891 | return ge::GRAPH_FAILED; | 1102 | return ge::GRAPH_FAILED; |
| 892 | } | 1103 | } |
| 1104 | + // Priority: small/non-last axes, resident select, merge, radix-one, | ||
| 1105 | + // general select, intra-core merge, then radix-more as the fallback. | ||
| 1106 | + if ((IsResidentIntegerSelect(info) || IsResidentFloatSelect(info, tilingData->medianMode) || | ||
| 1107 | + IsKthFloatSelect(info, tilingData->medianMode)) && | ||
| 1108 | + TryRadixSelect(context, info, tilingData, blockDim, schId)) { | ||
| 1109 | + return ge::GRAPH_SUCCESS; | ||
| 1110 | + } | ||
| 893 | if (TryMerge(context, info, tilingData, blockDim, schId) || TryRadixOneCore(context, info, tilingData, schId) || | 1111 | if (TryMerge(context, info, tilingData, blockDim, schId) || TryRadixOneCore(context, info, tilingData, schId) || |
| 894 | TryRadixSelect(context, info, tilingData, blockDim, schId) || | 1112 | TryRadixSelect(context, info, tilingData, blockDim, schId) || |
| 895 | TryMergeIntraCore(context, info, tilingData, blockDim, schId)) { | 1113 | TryMergeIntraCore(context, info, tilingData, blockDim, schId)) { |
| @@ -932,7 +1150,8 @@ static ge::graphStatus FinalizeKthValueRoute(gert::TilingContext* context, | |||
| 932 | OP_CHECK_NULL_WITH_CONTEXT(context, userWorkspaceSize); | 1150 | OP_CHECK_NULL_WITH_CONTEXT(context, userWorkspaceSize); |
| 933 | if (schId != KTH_VALUE_SCHID_MERGE_MORE_CORE && schId != KTH_VALUE_SCHID_MERGE_INTRA_CORE && | 1151 | if (schId != KTH_VALUE_SCHID_MERGE_MORE_CORE && schId != KTH_VALUE_SCHID_MERGE_INTRA_CORE && |
| 934 | schId != KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS && schId != KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX && | 1152 | schId != KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS && schId != KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX && |
| 935 | - schId != KTH_VALUE_SCHID_RADIX_MORE_CORE && schId != KTH_VALUE_SCHID_RADIX_SELECT) { | 1153 | + schId != KTH_VALUE_SCHID_RADIX_MORE_CORE && schId != KTH_VALUE_SCHID_RADIX_SELECT && |
| 1154 | + schId != KTH_VALUE_SCHID_RESIDENT_HISTOGRAM) { | ||
| 936 | userWorkspaceSize[0] = WORK_SPACE_SIZE; | 1155 | userWorkspaceSize[0] = WORK_SPACE_SIZE; |
| 937 | } | 1156 | } |
| 938 | return ge::GRAPH_SUCCESS; | 1157 | return ge::GRAPH_SUCCESS; |
| @@ -946,8 +1165,8 @@ static void SetKthValueTilingContext(gert::TilingContext* context, uint64_t schI | |||
| 946 | context->SetBlockDim(blockDim); | 1165 | context->SetBlockDim(blockDim); |
| 947 | if (schId == KTH_VALUE_SCHID_SMALL_AXIS_INSERTION || schId == KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE || | 1166 | if (schId == KTH_VALUE_SCHID_SMALL_AXIS_INSERTION || schId == KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE || |
| 948 | schId == KTH_VALUE_SCHID_SMALL_AXIS_SHORT_RANK_SELECT || schId == KTH_VALUE_SCHID_RADIX_MORE_CORE || | 1167 | schId == KTH_VALUE_SCHID_SMALL_AXIS_SHORT_RANK_SELECT || schId == KTH_VALUE_SCHID_RADIX_MORE_CORE || |
| 949 | - schId == KTH_VALUE_SCHID_RADIX_SELECT || schId == KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS || | 1168 | + schId == KTH_VALUE_SCHID_RADIX_SELECT || schId == KTH_VALUE_SCHID_RESIDENT_HISTOGRAM || |
| 950 | - schId == KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX) { | 1169 | + schId == KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS || schId == KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX) { |
| 951 | context->SetLocalMemorySize(info.ubSize - SIMT_UB); | 1170 | context->SetLocalMemorySize(info.ubSize - SIMT_UB); |
| 952 | } else { | 1171 | } else { |
| 953 | context->SetLocalMemorySize(info.ubSize); | 1172 | context->SetLocalMemorySize(info.ubSize); |
| @@ -1017,6 +1236,44 @@ static ge::graphStatus AddMedianWorkspace(gert::TilingContext* context, const So | |||
| 1017 | return ge::GRAPH_SUCCESS; | 1236 | return ge::GRAPH_SUCCESS; |
| 1018 | } | 1237 | } |
| 1019 | 1238 | ||
| 1239 | +// Serialize the selected layout only after routing and workspace planning have completed. | ||
| 1240 | +static void StoreKthValueTiling(gert::TilingContext* context, uint64_t schId, const KthValueTilingData& data) | ||
| 1241 | +{ | ||
| 1242 | + if (schId == KTH_VALUE_SCHID_MERGE_SORT || schId == KTH_VALUE_SCHID_SORT32_SMALL_AXIS) { | ||
| 1243 | + KthValueMergeOneCoreTilingData compact{}; | ||
| 1244 | + compact.kthIndex = data.kthIndex; | ||
| 1245 | + compact.unsortedDimNum = data.unsortedDimNum; | ||
| 1246 | + compact.numTileDataSize = data.numTileDataSize; | ||
| 1247 | + compact.unsortedDimParallel = data.unsortedDimParallel; | ||
| 1248 | + compact.sortLoopTimes = data.sortLoopTimes; | ||
| 1249 | + compact.keyParams0 = data.keyParams0; | ||
| 1250 | + compact.keyParams1 = data.keyParams1; | ||
| 1251 | + compact.keyParams2 = data.keyParams2; | ||
| 1252 | + compact.keyParams3 = data.keyParams3; | ||
| 1253 | + compact.keyParams4 = data.keyParams4; | ||
| 1254 | + compact.medianMode = data.medianMode; | ||
| 1255 | + // Capacity was already checked for the larger common layout. | ||
| 1256 | + *context->GetTilingData<KthValueMergeOneCoreTilingData>() = compact; | ||
| 1257 | + } else if (schId == KTH_VALUE_SCHID_RADIX_ONE_CORE) { | ||
| 1258 | + KthValueRadixOneCoreTilingData compact{}; | ||
| 1259 | + compact.kthIndex = data.kthIndex; | ||
| 1260 | + compact.lastAxisNum = data.lastAxisNum; | ||
| 1261 | + compact.unsortedDimNum = data.unsortedDimNum; | ||
| 1262 | + compact.numTileDataSize = data.numTileDataSize; | ||
| 1263 | + compact.unsortedDimParallel = data.unsortedDimParallel; | ||
| 1264 | + compact.keyParams0 = data.keyParams0; | ||
| 1265 | + compact.keyParams1 = data.keyParams1; | ||
| 1266 | + compact.keyParams3 = data.keyParams3; | ||
| 1267 | + compact.keyParams4 = data.keyParams4; | ||
| 1268 | + compact.tmpUbSize = data.tmpUbSize; | ||
| 1269 | + compact.medianMode = data.medianMode; | ||
| 1270 | + // Capacity was already checked for the larger common layout. | ||
| 1271 | + *context->GetTilingData<KthValueRadixOneCoreTilingData>() = compact; | ||
| 1272 | + } else { | ||
| 1273 | + *context->GetTilingData<KthValueTilingData>() = data; | ||
| 1274 | + } | ||
| 1275 | +} | ||
| 1276 | + | ||
| 1020 | static ge::graphStatus SelectAndFinalizeKthValueRoute(gert::TilingContext* context, | 1277 | static ge::graphStatus SelectAndFinalizeKthValueRoute(gert::TilingContext* context, |
| 1021 | const platform_ascendc::PlatformAscendC& ascendcPlatform, | 1278 | const platform_ascendc::PlatformAscendC& ascendcPlatform, |
| 1022 | SortKthTileInfo& info, KthValueTilingData* tilingData, | 1279 | SortKthTileInfo& info, KthValueTilingData* tilingData, |
| @@ -1025,10 +1282,11 @@ static ge::graphStatus SelectAndFinalizeKthValueRoute(gert::TilingContext* conte | |||
| 1025 | KthValueTilingData candidateTilingData = *tilingData; | 1282 | KthValueTilingData candidateTilingData = *tilingData; |
| 1026 | uint64_t schId = KTH_VALUE_SCHID_RADIX_MORE_CORE; | 1283 | uint64_t schId = KTH_VALUE_SCHID_RADIX_MORE_CORE; |
| 1027 | uint32_t blockDim = 1; | 1284 | uint32_t blockDim = 1; |
| 1028 | - OP_CHECK_IF((SelectKthValueRoute(context, info, &candidateTilingData, blockDim, schId) != ge::GRAPH_SUCCESS), | 1285 | + OP_CHECK_IF( |
| 1029 | - OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "SelectKthValueRoute", "GRAPH_FAILED", | 1286 | + (SelectKthValueRoute(context, info, &candidateTilingData, blockDim, schId, isMedianOp) != ge::GRAPH_SUCCESS), |
| 1030 | - "The value of SelectKthValueRoute must be GRAPH_SUCCESS."), | 1287 | + OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "SelectKthValueRoute", "GRAPH_FAILED", |
| 1031 | - return ge::GRAPH_FAILED); | 1288 | + "The value of SelectKthValueRoute must be GRAPH_SUCCESS."), |
| 1289 | + return ge::GRAPH_FAILED); | ||
| 1032 | OP_CHECK_IF((FinalizeKthValueRoute(context, ascendcPlatform, info, &candidateTilingData, schId, blockDim) != | 1290 | OP_CHECK_IF((FinalizeKthValueRoute(context, ascendcPlatform, info, &candidateTilingData, schId, blockDim) != |
| 1033 | ge::GRAPH_SUCCESS), | 1291 | ge::GRAPH_SUCCESS), |
| 1034 | OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "FinalizeKthValueRoute", "GRAPH_FAILED", | 1292 | OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "FinalizeKthValueRoute", "GRAPH_FAILED", |
| @@ -1039,7 +1297,7 @@ static ge::graphStatus SelectAndFinalizeKthValueRoute(gert::TilingContext* conte | |||
| 1039 | OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "AddMedianWorkspace", "GRAPH_FAILED", | 1297 | OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "AddMedianWorkspace", "GRAPH_FAILED", |
| 1040 | "The value of AddMedianWorkspace must be GRAPH_SUCCESS."), | 1298 | "The value of AddMedianWorkspace must be GRAPH_SUCCESS."), |
| 1041 | return ge::GRAPH_FAILED); | 1299 | return ge::GRAPH_FAILED); |
| 1042 | - *tilingData = candidateTilingData; | 1300 | + StoreKthValueTiling(context, schId, candidateTilingData); |
| 1043 | OP_LOGI(context->GetNodeName(), | 1301 | OP_LOGI(context->GetNodeName(), |
| 1044 | "KthValueTiling: schId=%lu, blockDim=%u, lastAxis=%ld, unsortedDim=%ld, " | 1302 | "KthValueTiling: schId=%lu, blockDim=%u, lastAxis=%ld, unsortedDim=%ld, " |
| 1045 | "isNonLastAxis=%d, dtypeSize=%u", | 1303 | "isNonLastAxis=%d, dtypeSize=%u", |
| @@ -10,6 +10,7 @@ | |||
| 10 | 10 | ||
| 11 | #include "kth_value_tiling_common.h" | 11 | #include "kth_value_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( |
| @@ -410,7 +417,8 @@ static bool AdjustBSharedSingleHTile(int64_t axisLen, int64_t unsortedDim, uint3 | |||
| 410 | if (hCore == 0U) { | 417 | if (hCore == 0U) { |
| 411 | return false; | 418 | return false; |
| 412 | } | 419 | } |
| 413 | - uint64_t hTileData64 = static_cast<uint64_t>(axisLen) / hCore; | 420 | + // Cover the full axis with hCore tiles. Floor division may introduce an extra tail tile after alignment. |
| 421 | + uint64_t hTileData64 = Ops::Base::CeilDiv(static_cast<uint64_t>(axisLen), static_cast<uint64_t>(hCore)); | ||
| 414 | if (hTileData64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { | 422 | if (hTileData64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { |
| 415 | return false; | 423 | return false; |
| 416 | } | 424 | } |
| @@ -427,7 +435,8 @@ static bool AdjustMultiTileHWithBSharing(int64_t axisLen, int64_t unsortedDim, u | |||
| 427 | uint32_t& tileData, uint32_t lastDimTileNum, uint32_t& tmpUbSize, | 435 | uint32_t& tileData, uint32_t lastDimTileNum, uint32_t& tmpUbSize, |
| 428 | bool& adjusted) | 436 | bool& adjusted) |
| 429 | { | 437 | { |
| 430 | - uint64_t newTileData64 = static_cast<uint64_t>(axisLen) / static_cast<uint64_t>(lastDimTileNum); | 438 | + // Keep the adjusted split within lastDimTileNum; floor division can recreate a one-element tail tile. |
| 439 | + uint64_t newTileData64 = Ops::Base::CeilDiv(static_cast<uint64_t>(axisLen), static_cast<uint64_t>(lastDimTileNum)); | ||
| 431 | if (!CeilAlignUint32(newTileData64, BIN_NUM, tileData)) { | 440 | if (!CeilAlignUint32(newTileData64, BIN_NUM, tileData)) { |
| 432 | return false; | 441 | return false; |
| 433 | } | 442 | } |
| @@ -706,14 +715,57 @@ bool ComputeMergeSortPlan(int64_t axisLen, int64_t unsortedDim, uint32_t blockUb | |||
| 706 | return true; | 715 | return true; |
| 707 | } | 716 | } |
| 708 | 717 | ||
| 709 | -bool FillMergeSortInfo(SortKthTileInfo& info, uint32_t indexDtypeSize, uint32_t concatTmpSize) | 718 | +static bool ComputeMergeSortBatchCapacity(const SortKthTileInfo& info, uint32_t indexDtypeSize, uint32_t& tileDataNum) |
| 710 | { | 719 | { |
| 711 | - constexpr uint32_t TILE_DATA_NUM = 4096; | 720 | + if (info.lastAxis <= 0 || info.blockUbSize == 0U || info.dtypeSize == 0U || indexDtypeSize == 0U) { |
| 712 | - MergeSortPlan plan; | ||
| 713 | - if (!ComputeMergeSortPlan(info.lastAxis, info.unsortedDim, info.blockUbSize, TILE_DATA_NUM, info.maxCoreNum, | ||
| 714 | - plan)) { | ||
| 715 | return false; | 721 | return false; |
| 716 | } | 722 | } |
| 723 | + uint64_t alignedAxis = Ops::Base::CeilAlign(static_cast<uint64_t>(info.lastAxis), | ||
| 724 | + static_cast<uint64_t>(info.blockUbSize)); | ||
| 725 | + uint32_t bufferNum = info.lastAxis > ONE_CORE_DATA_SIZE ? 1U : DOUBLE_BUFFER_NUM; | ||
| 726 | + // Both kernels keep one index vector and two 8-byte proposal buffers per row axis. | ||
| 727 | + uint64_t fixedBytes = alignedAxis * (sizeof(uint32_t) + 2U * MERGE_SORT_DATA_BYTES); | ||
| 728 | + uint32_t castBytes = (info.dataType == ge::DT_BF16 || info.dataType == ge::DT_INT16) ? sizeof(float) : 0U; | ||
| 729 | + uint64_t rowBytes = alignedAxis * ((2U * info.dtypeSize + indexDtypeSize) * bufferNum + castBytes); | ||
| 730 | + rowBytes += (info.dtypeSize + sizeof(int64_t)) * bufferNum; | ||
| 731 | + // KthValue additionally aligns two compact output queues independently per batch. | ||
| 732 | + fixedBytes += 2U * bufferNum * (info.blockUbSize - 1U); | ||
| 733 | + if (fixedBytes >= info.ubSize || rowBytes == 0U) { | ||
| 734 | + return false; | ||
| 735 | + } | ||
| 736 | + uint64_t rows = (info.ubSize - fixedBytes) / rowBytes; | ||
| 737 | + // ComputeMergeSortPlan consumes half of tileDataNum as the batch element budget. | ||
| 738 | + uint64_t tileElements = rows * alignedAxis * 2U; | ||
| 739 | + if (rows == 0U || tileElements > std::numeric_limits<uint32_t>::max()) { | ||
| 740 | + return false; | ||
| 741 | + } | ||
| 742 | + tileDataNum = static_cast<uint32_t>(tileElements); | ||
| 743 | + return true; | ||
| 744 | +} | ||
| 745 | + | ||
| 746 | +bool FillMergeSortInfo(SortKthTileInfo& info, uint32_t indexDtypeSize, bool useUbCapacity) | ||
| 747 | +{ | ||
| 748 | + // The fixed budget keeps small Sort batches pipelined; the UB-derived budget also | ||
| 749 | + // rebalances rows across rounds. Switching policies is not only a capacity check. | ||
| 750 | + constexpr uint32_t SORT_TILE_DATA_NUM = 4096; | ||
| 751 | + uint32_t tileDataNum = SORT_TILE_DATA_NUM; | ||
| 752 | + if (useUbCapacity && !ComputeMergeSortBatchCapacity(info, indexDtypeSize, tileDataNum)) { | ||
| 753 | + return false; | ||
| 754 | + } | ||
| 755 | + MergeSortPlan plan; | ||
| 756 | + if (!ComputeMergeSortPlan(info.lastAxis, info.unsortedDim, info.blockUbSize, tileDataNum, info.maxCoreNum, plan)) { | ||
| 757 | + return false; | ||
| 758 | + } | ||
| 759 | + if (useUbCapacity) { | ||
| 760 | + // Balance the per-core row budget instead of leaving an oversized final batch on a few cores. | ||
| 761 | + uint64_t rowsPerCore = Ops::Base::CeilDiv(static_cast<uint64_t>(info.unsortedDim), | ||
| 762 | + static_cast<uint64_t>(plan.coreNumNeed)); | ||
| 763 | + uint64_t rounds = Ops::Base::CeilDiv(rowsPerCore, static_cast<uint64_t>(plan.oneCoreRowNum)); | ||
| 764 | + plan.oneCoreRowNum = static_cast<uint32_t>(Ops::Base::CeilDiv(rowsPerCore, rounds)); | ||
| 765 | + uint64_t rowsPerRound = static_cast<uint64_t>(plan.coreNumNeed) * plan.oneCoreRowNum; | ||
| 766 | + plan.sortLoopTimes = static_cast<uint32_t>( | ||
| 767 | + Ops::Base::CeilDiv(static_cast<uint64_t>(info.unsortedDim), rowsPerRound)); | ||
| 768 | + } | ||
| 717 | info.sortLoopTimes = plan.sortLoopTimes; | 769 | info.sortLoopTimes = plan.sortLoopTimes; |
| 718 | info.lastDimTileNum = 1; | 770 | info.lastDimTileNum = 1; |
| 719 | info.unsortedDimParallel = plan.coreNumNeed; | 771 | info.unsortedDimParallel = plan.coreNumNeed; |
| @@ -725,24 +777,15 @@ bool FillMergeSortInfo(SortKthTileInfo& info, uint32_t indexDtypeSize, uint32_t | |||
| 725 | info.keyParams2 = plan.alignNum * plan.oneCoreRowNum * indexDtypeSize; | 777 | info.keyParams2 = plan.alignNum * plan.oneCoreRowNum * indexDtypeSize; |
| 726 | info.keyParams3 = plan.alignNum; | 778 | info.keyParams3 = plan.alignNum; |
| 727 | info.keyParams4 = info.lastAxis > ONE_CORE_DATA_SIZE ? 1 : DOUBLE_BUFFER_NUM; | 779 | info.keyParams4 = info.lastAxis > ONE_CORE_DATA_SIZE ? 1 : DOUBLE_BUFFER_NUM; |
| 728 | - info.tmpUbSize = std::max(concatTmpSize, info.blockUbSize); | 780 | + info.tmpUbSize = 0; |
| 729 | info.workspaceSize = WORK_SPACE_SIZE; | 781 | info.workspaceSize = WORK_SPACE_SIZE; |
| 730 | return true; | 782 | return true; |
| 731 | } | 783 | } |
| 732 | 784 | ||
| 733 | -bool ComputeMergeSortTiling(gert::TilingContext* context, SortKthTileInfo& info, uint32_t indexDtypeSize) | 785 | +bool ComputeMergeSortTiling(gert::TilingContext* context, SortKthTileInfo& info, uint32_t indexDtypeSize, |
| 786 | + bool useUbCapacity) | ||
| 734 | { | 787 | { |
| 735 | - constexpr uint32_t TILE_DATA_NUM = 4096; | 788 | + if (!FillMergeSortInfo(info, indexDtypeSize, useUbCapacity)) { |
| 736 | - MergeSortPlan plan; | ||
| 737 | - if (!ComputeMergeSortPlan(info.lastAxis, info.unsortedDim, info.blockUbSize, TILE_DATA_NUM, info.maxCoreNum, | ||
| 738 | - plan)) { | ||
| 739 | - return false; | ||
| 740 | - } | ||
| 741 | - auto platformInfo = context->GetPlatformInfo(); | ||
| 742 | - auto plat = platform_ascendc::PlatformAscendC(platformInfo); | ||
| 743 | - uint32_t maxTypeSize = (info.dataType == ge::DT_BF16) ? static_cast<uint32_t>(sizeof(float)) : info.dtypeSize; | ||
| 744 | - uint32_t concatTmpSize = AscendC::GetConcatTmpSize(plat, plan.alignNum, maxTypeSize); | ||
| 745 | - if (!FillMergeSortInfo(info, indexDtypeSize, concatTmpSize)) { | ||
| 746 | return false; | 789 | return false; |
| 747 | } | 790 | } |
| 748 | size_t* userWorkSpaceSize = context->GetWorkspaceSizes(1); | 791 | size_t* userWorkSpaceSize = context->GetWorkspaceSizes(1); |
| @@ -753,51 +796,167 @@ bool ComputeMergeSortTiling(gert::TilingContext* context, SortKthTileInfo& info, | |||
| 753 | // ============================================================================= | 796 | // ============================================================================= |
| 754 | // Merge sort — multi-core (more-core) | 797 | // Merge sort — multi-core (more-core) |
| 755 | // ============================================================================= | 798 | // ============================================================================= |
| 756 | -bool IsMergeMoreCoreSupported(ge::DataType dataType, int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum) | 799 | +// Select per-core data chunk size (2048 or 4096) for merge-more-core. Prefers the smaller base size when direct |
| 800 | +// parallel fits within maxCoreNum; otherwise falls back to the larger size to reduce horizontal core count. | ||
| 801 | +bool SelectMergeMoreCoreDataSize(int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum, uint32_t& dataSize) | ||
| 757 | { | 802 | { |
| 758 | - if (dataType != ge::DT_FLOAT || axisLen <= static_cast<int64_t>(MERGE_SORT_MAX_AXIS_FP32) || | 803 | + if (axisLen <= 0 || unsortedDim <= 0 || maxCoreNum == 0U) { |
| 759 | - axisLen > static_cast<int64_t>(MULTI_CORE_MERGE_SORT_MAX_AXIS)) { | ||
| 760 | return false; | 804 | return false; |
| 761 | } | 805 | } |
| 762 | - uint64_t hCoreNum = (static_cast<uint64_t>(axisLen) + ONE_CORE_DATA_SIZE - 1U) / ONE_CORE_DATA_SIZE; | 806 | + uint32_t baseHCoreNum; |
| 763 | - return hCoreNum > 0 && static_cast<uint64_t>(unsortedDim) * hCoreNum <= maxCoreNum; | 807 | + if (!CeilDivUint32(static_cast<uint64_t>(axisLen), MERGE_MORE_CORE_DATA_SIZE_BASE, baseHCoreNum) || |
| 808 | + baseHCoreNum == 0U || baseHCoreNum > maxCoreNum) { | ||
| 809 | + return false; | ||
| 810 | + } | ||
| 811 | + if (static_cast<uint64_t>(unsortedDim) * baseHCoreNum <= maxCoreNum) { | ||
| 812 | + dataSize = MERGE_MORE_CORE_DATA_SIZE_BASE; | ||
| 813 | + return true; | ||
| 814 | + } | ||
| 815 | + uint32_t largeHCoreNum; | ||
| 816 | + if (!CeilDivUint32(static_cast<uint64_t>(axisLen), MERGE_MORE_CORE_DATA_SIZE_LARGE, largeHCoreNum) || | ||
| 817 | + largeHCoreNum == 0U || largeHCoreNum > maxCoreNum) { | ||
| 818 | + return false; | ||
| 819 | + } | ||
| 820 | + dataSize = MERGE_MORE_CORE_DATA_SIZE_LARGE; | ||
| 821 | + return true; | ||
| 764 | } | 822 | } |
| 765 | 823 | ||
| 766 | -bool ComputeMergeMoreCorePlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t mergeBytesPerElem, | 824 | +// Select a sync-merge block size (aligned to 2048) that groups multiple 4096-element chunks into one block, |
| 767 | - MergeMoreCorePlan& plan) | 825 | +// reducing horizontal core count so that unsortedDim * hCoreNum fits within maxCoreNum. |
| 826 | +bool SelectMergeSyncMergeBlockSize(int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum, uint32_t& blockSize) | ||
| 827 | +{ | ||
| 828 | + if (axisLen <= static_cast<int64_t>(MERGE_MORE_CORE_DATA_SIZE_LARGE) || unsortedDim <= 0 || maxCoreNum == 0U) { | ||
| 829 | + return false; | ||
| 830 | + } | ||
| 831 | + uint32_t maxHCoreByCore = static_cast<uint32_t>(static_cast<uint64_t>(maxCoreNum) / | ||
| 832 | + static_cast<uint64_t>(unsortedDim)); | ||
| 833 | + uint32_t rawChunkNum = 0; | ||
| 834 | + if (!CeilDivUint32(static_cast<uint64_t>(axisLen), MERGE_MORE_CORE_DATA_SIZE_LARGE, rawChunkNum) || | ||
| 835 | + rawChunkNum <= 1U || maxHCoreByCore <= 1U) { | ||
| 836 | + return false; | ||
| 837 | + } | ||
| 838 | + // A sync-merge block is larger than one 4096-element Sort block. Use the available per-row core budget, | ||
| 839 | + // capped below the direct 4096 split, then align to the existing 2048-element more-core granularity. | ||
| 840 | + uint32_t targetHCore = std::min(rawChunkNum - 1U, maxHCoreByCore); | ||
| 841 | + if (targetHCore <= 1U) { | ||
| 842 | + return false; | ||
| 843 | + } | ||
| 844 | + uint32_t candidate = 0; | ||
| 845 | + if (!CeilAlignUint32(Ops::Base::CeilDiv(static_cast<uint64_t>(axisLen), static_cast<uint64_t>(targetHCore)), | ||
| 846 | + MERGE_MORE_CORE_DATA_SIZE_BASE, candidate) || | ||
| 847 | + candidate <= MERGE_MORE_CORE_DATA_SIZE_LARGE) { | ||
| 848 | + return false; | ||
| 849 | + } | ||
| 850 | + uint32_t hCoreNum = 0; | ||
| 851 | + if (!CeilDivUint32(static_cast<uint64_t>(axisLen), candidate, hCoreNum) || hCoreNum <= 1U || | ||
| 852 | + hCoreNum > maxCoreNum || | ||
| 853 | + static_cast<uint64_t>(unsortedDim) * static_cast<uint64_t>(hCoreNum) > static_cast<uint64_t>(maxCoreNum)) { | ||
| 854 | + return false; | ||
| 855 | + } | ||
| 856 | + blockSize = candidate; | ||
| 857 | + return true; | ||
| 858 | +} | ||
| 859 | + | ||
| 860 | +bool IsMergeMoreCoreProfitable(ge::DataType dataType, int64_t axisLen, int64_t unsortedDim, uint32_t maxCoreNum) | ||
| 861 | +{ | ||
| 862 | + if (dataType != ge::DT_FLOAT || axisLen <= static_cast<int64_t>(MERGE_SORT_MAX_AXIS_FP32) || | ||
| 863 | + axisLen > static_cast<int64_t>(MULTI_CORE_MERGE_SORT_MAX_AXIS) || unsortedDim <= 0 || maxCoreNum == 0U) { | ||
| 864 | + return false; | ||
| 865 | + } | ||
| 866 | + uint32_t dataSize; | ||
| 867 | + if (!SelectMergeMoreCoreDataSize(axisLen, unsortedDim, maxCoreNum, dataSize)) { | ||
| 868 | + return false; | ||
| 869 | + } | ||
| 870 | + uint64_t hCoreNum = (static_cast<uint64_t>(axisLen) + dataSize - 1U) / dataSize; | ||
| 871 | + if (static_cast<uint64_t>(unsortedDim) * hCoreNum <= maxCoreNum) { | ||
| 872 | + return true; | ||
| 873 | + } | ||
| 874 | + uint32_t syncBlockSize = 0; | ||
| 875 | + if (SelectMergeSyncMergeBlockSize(axisLen, unsortedDim, maxCoreNum, syncBlockSize)) { | ||
| 876 | + return true; | ||
| 877 | + } | ||
| 878 | + return hCoreNum >= MERGE_MORE_CORE_MIN_CORES_PER_ROW_FOR_MULTI_ROUND; | ||
| 879 | +} | ||
| 880 | + | ||
| 881 | +namespace { | ||
| 882 | + | ||
| 883 | +bool BuildMergeMoreCoreRoundPlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t maxCoreNum, | ||
| 884 | + uint32_t mergeBytesPerElem, uint32_t hCoreNum, uint32_t numTileDataSize, | ||
| 885 | + uint32_t syncBlockSize, MergeMoreCorePlan& plan) | ||
| 886 | +{ | ||
| 887 | + uint64_t rowsPerRound64 = std::min(static_cast<uint64_t>(maxCoreNum) / hCoreNum, | ||
| 888 | + static_cast<uint64_t>(unsortedDim)); | ||
| 889 | + if (rowsPerRound64 == 0U || rowsPerRound64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { | ||
| 890 | + return false; | ||
| 891 | + } | ||
| 892 | + uint64_t coreNumNeed64 = rowsPerRound64 * hCoreNum; | ||
| 893 | + uint64_t sortLoopTimes64 = Ops::Base::CeilDiv(static_cast<uint64_t>(unsortedDim), rowsPerRound64); | ||
| 894 | + if (static_cast<uint64_t>(axisLen) > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) || | ||
| 895 | + static_cast<uint64_t>(unsortedDim) > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) || | ||
| 896 | + coreNumNeed64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) || | ||
| 897 | + sortLoopTimes64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { | ||
| 898 | + return false; | ||
| 899 | + } | ||
| 900 | + plan.lastDimTileNum = static_cast<uint32_t>(axisLen); | ||
| 901 | + plan.lastDimNeedCore = hCoreNum; | ||
| 902 | + plan.numTileDataSize = numTileDataSize; | ||
| 903 | + plan.unsortedDimParallel = static_cast<uint32_t>(rowsPerRound64); | ||
| 904 | + plan.sortLoopTimes = static_cast<uint32_t>(sortLoopTimes64); | ||
| 905 | + plan.coreNumNeed = static_cast<uint32_t>(coreNumNeed64); | ||
| 906 | + plan.keyParams0 = ubSize / mergeBytesPerElem; | ||
| 907 | + plan.keyParams1 = syncBlockSize; | ||
| 908 | + return true; | ||
| 909 | +} | ||
| 910 | + | ||
| 911 | +bool BuildMergeSyncBlockPlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t maxCoreNum, | ||
| 912 | + uint32_t mergeBytesPerElem, uint32_t syncBlockSize, MergeMoreCorePlan& plan) | ||
| 913 | +{ | ||
| 914 | + uint32_t syncHCoreNum = 0; | ||
| 915 | + if (!CeilDivUint32(static_cast<uint64_t>(axisLen), syncBlockSize, syncHCoreNum) || syncHCoreNum <= 1U || | ||
| 916 | + syncHCoreNum > maxCoreNum) { | ||
| 917 | + return false; | ||
| 918 | + } | ||
| 919 | + return BuildMergeMoreCoreRoundPlan(axisLen, unsortedDim, ubSize, maxCoreNum, mergeBytesPerElem, syncHCoreNum, | ||
| 920 | + MERGE_MORE_CORE_DATA_SIZE_LARGE, syncBlockSize, plan); | ||
| 921 | +} | ||
| 922 | + | ||
| 923 | +} // namespace | ||
| 924 | + | ||
| 925 | +// Try direct, sync-block, then multi-round scheduling in that order. | ||
| 926 | +bool ComputeMergeMoreCorePlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t maxCoreNum, | ||
| 927 | + uint32_t mergeBytesPerElem, MergeMoreCorePlan& plan) | ||
| 768 | { | 928 | { |
| 769 | if (axisLen <= 0 || unsortedDim <= 0 || mergeBytesPerElem == 0U) { | 929 | if (axisLen <= 0 || unsortedDim <= 0 || mergeBytesPerElem == 0U) { |
| 770 | return false; | 930 | return false; |
| 771 | } | 931 | } |
| 772 | - uint32_t hCoreNum; | 932 | + uint32_t dataSize = 0; |
| 773 | - if (!CeilDivUint32(static_cast<uint64_t>(axisLen), ONE_CORE_DATA_SIZE, hCoreNum)) { | 933 | + uint32_t hCoreNum = 0; |
| 934 | + if (!SelectMergeMoreCoreDataSize(axisLen, unsortedDim, maxCoreNum, dataSize) || | ||
| 935 | + !CeilDivUint32(static_cast<uint64_t>(axisLen), dataSize, hCoreNum) || hCoreNum == 0U || hCoreNum > maxCoreNum) { | ||
| 774 | return false; | 936 | return false; |
| 775 | } | 937 | } |
| 776 | - uint64_t coreNumNeed64 = static_cast<uint64_t>(unsortedDim) * hCoreNum; | 938 | + bool directFit = static_cast<uint64_t>(unsortedDim) * hCoreNum <= maxCoreNum; |
| 777 | - if (static_cast<uint64_t>(axisLen) > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) || | 939 | + uint32_t syncBlockSize = 0; |
| 778 | - static_cast<uint64_t>(unsortedDim) > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()) || | 940 | + if (!directFit && SelectMergeSyncMergeBlockSize(axisLen, unsortedDim, maxCoreNum, syncBlockSize)) { |
| 779 | - coreNumNeed64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { | 941 | + return BuildMergeSyncBlockPlan(axisLen, unsortedDim, ubSize, maxCoreNum, mergeBytesPerElem, syncBlockSize, |
| 942 | + plan); | ||
| 943 | + } | ||
| 944 | + if (!directFit && hCoreNum < MERGE_MORE_CORE_MIN_CORES_PER_ROW_FOR_MULTI_ROUND) { | ||
| 780 | return false; | 945 | return false; |
| 781 | } | 946 | } |
| 782 | uint64_t numTileDataSize64 = static_cast<uint64_t>(axisLen) / hCoreNum; | 947 | uint64_t numTileDataSize64 = static_cast<uint64_t>(axisLen) / hCoreNum; |
| 783 | if (numTileDataSize64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { | 948 | if (numTileDataSize64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { |
| 784 | return false; | 949 | return false; |
| 785 | } | 950 | } |
| 786 | - plan.lastDimTileNum = static_cast<uint32_t>(static_cast<uint64_t>(axisLen)); | 951 | + return BuildMergeMoreCoreRoundPlan(axisLen, unsortedDim, ubSize, maxCoreNum, mergeBytesPerElem, hCoreNum, |
| 787 | - plan.lastDimNeedCore = hCoreNum; | 952 | + static_cast<uint32_t>(numTileDataSize64), 0U, plan); |
| 788 | - plan.numTileDataSize = static_cast<uint32_t>(numTileDataSize64); | ||
| 789 | - plan.unsortedDimParallel = static_cast<uint32_t>(static_cast<uint64_t>(unsortedDim)); | ||
| 790 | - plan.sortLoopTimes = 1; | ||
| 791 | - plan.coreNumNeed = static_cast<uint32_t>(coreNumNeed64); | ||
| 792 | - // keyParams0 is max elements handled per UB round by the merge more-core pipeline. | ||
| 793 | - plan.keyParams0 = ubSize / mergeBytesPerElem; | ||
| 794 | - return true; | ||
| 795 | } | 953 | } |
| 796 | 954 | ||
| 797 | bool FillMergeMoreCoreInfo(SortKthTileInfo& info, uint32_t mergeBytesPerElem) | 955 | bool FillMergeMoreCoreInfo(SortKthTileInfo& info, uint32_t mergeBytesPerElem) |
| 798 | { | 956 | { |
| 799 | MergeMoreCorePlan plan; | 957 | MergeMoreCorePlan plan; |
| 800 | - if (!ComputeMergeMoreCorePlan(info.lastAxis, info.unsortedDim, info.ubSize, mergeBytesPerElem, plan)) { | 958 | + if (!ComputeMergeMoreCorePlan(info.lastAxis, info.unsortedDim, info.ubSize, info.maxCoreNum, mergeBytesPerElem, |
| 959 | + plan)) { | ||
| 801 | return false; | 960 | return false; |
| 802 | } | 961 | } |
| 803 | info.lastDimTileNum = plan.lastDimTileNum; | 962 | info.lastDimTileNum = plan.lastDimTileNum; |
| @@ -807,6 +966,7 @@ bool FillMergeMoreCoreInfo(SortKthTileInfo& info, uint32_t mergeBytesPerElem) | |||
| 807 | info.sortLoopTimes = plan.sortLoopTimes; | 966 | info.sortLoopTimes = plan.sortLoopTimes; |
| 808 | info.coreNumNeed = plan.coreNumNeed; | 967 | info.coreNumNeed = plan.coreNumNeed; |
| 809 | info.keyParams0 = plan.keyParams0; | 968 | info.keyParams0 = plan.keyParams0; |
| 969 | + info.keyParams1 = plan.keyParams1; | ||
| 810 | uint64_t wsBytes = static_cast<uint64_t>(MERGE_SORT_WORKSPACE_PARAM) * info.lastAxis * info.unsortedDim * | 970 | uint64_t wsBytes = static_cast<uint64_t>(MERGE_SORT_WORKSPACE_PARAM) * info.lastAxis * info.unsortedDim * |
| 811 | sizeof(int32_t); | 971 | sizeof(int32_t); |
| 812 | info.workspaceSize = static_cast<size_t>(wsBytes + WORK_SPACE_SIZE); | 972 | info.workspaceSize = static_cast<size_t>(wsBytes + WORK_SPACE_SIZE); |
| @@ -834,6 +994,19 @@ uint32_t ComputeMergeIntraCoreBlockSortSize(uint32_t ubSize) | |||
| 834 | return (blockSortSize / MERGE_INTRA_CORE_SORT_ALIGN) * MERGE_INTRA_CORE_SORT_ALIGN; | 994 | return (blockSortSize / MERGE_INTRA_CORE_SORT_ALIGN) * MERGE_INTRA_CORE_SORT_ALIGN; |
| 835 | } | 995 | } |
| 836 | 996 | ||
| 997 | +// Max elements a single core can sort in one pass (single-block mode). | ||
| 998 | +// Constrained by the tighter of phase1 (initial sort) and phase3 (extract/merge) UB capacity. | ||
| 999 | +uint32_t ComputeMergeIntraCoreSingleBlockSortSize(uint32_t ubSize) | ||
| 1000 | +{ | ||
| 1001 | + constexpr uint32_t MERGE_INTRA_PHASE_BUFFER_NUM = 2; | ||
| 1002 | + constexpr uint32_t PHASE1_BYTES_PER_ELEM = MERGE_INTRA_PHASE_BUFFER_NUM * sizeof(float) + SORT_STRUCT_BYTES + | ||
| 1003 | + MERGE_INTRA_PHASE_BUFFER_NUM * SORT_STRUCT_BYTES + sizeof(uint32_t); | ||
| 1004 | + uint32_t phase1Size = ubSize / PHASE1_BYTES_PER_ELEM; | ||
| 1005 | + uint32_t phase3Size = ComputeMergeIntraCoreExtractChunkSize(ubSize); | ||
| 1006 | + uint32_t singleBlockSize = std::min(phase1Size, phase3Size); | ||
| 1007 | + return (singleBlockSize / MERGE_INTRA_CORE_SORT_ALIGN) * MERGE_INTRA_CORE_SORT_ALIGN; | ||
| 1008 | +} | ||
| 1009 | + | ||
| 837 | uint32_t ComputeMergeIntraCoreExtractChunkSize(uint32_t ubSize) | 1010 | uint32_t ComputeMergeIntraCoreExtractChunkSize(uint32_t ubSize) |
| 838 | { | 1011 | { |
| 839 | constexpr uint32_t PHASE3_BYTES_PER_ELEM = (SORT_STRUCT_BYTES + sizeof(float) + sizeof(int32_t) + sizeof(int64_t)) * | 1012 | constexpr uint32_t PHASE3_BYTES_PER_ELEM = (SORT_STRUCT_BYTES + sizeof(float) + sizeof(int32_t) + sizeof(int64_t)) * |
| @@ -854,11 +1027,13 @@ bool IsMergeIntraCoreSupported(ge::DataType dataType, int64_t axisLen, int64_t u | |||
| 854 | return false; | 1027 | return false; |
| 855 | } | 1028 | } |
| 856 | uint32_t blockSortSize = ComputeMergeIntraCoreBlockSortSize(ubSize); | 1029 | uint32_t blockSortSize = ComputeMergeIntraCoreBlockSortSize(ubSize); |
| 1030 | + uint32_t singleBlockSortSize = ComputeMergeIntraCoreSingleBlockSortSize(ubSize); | ||
| 857 | uint32_t extractChunkSize = ComputeMergeIntraCoreExtractChunkSize(ubSize); | 1031 | uint32_t extractChunkSize = ComputeMergeIntraCoreExtractChunkSize(ubSize); |
| 858 | - if (blockSortSize == 0U || extractChunkSize == 0U) { | 1032 | + if (blockSortSize == 0U || singleBlockSortSize == 0U || extractChunkSize == 0U) { |
| 859 | return false; | 1033 | return false; |
| 860 | } | 1034 | } |
| 861 | - return axisLen <= static_cast<int64_t>(blockSortSize) * MERGE_INTRA_CORE_MAX_BLOCKS; | 1035 | + return axisLen <= static_cast<int64_t>(singleBlockSortSize) || |
| 1036 | + axisLen <= static_cast<int64_t>(blockSortSize) * MERGE_INTRA_CORE_MAX_BLOCKS; | ||
| 862 | } | 1037 | } |
| 863 | 1038 | ||
| 864 | bool ComputeMergeIntraCorePlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t maxCoreNum, | 1039 | bool ComputeMergeIntraCorePlan(int64_t axisLen, int64_t unsortedDim, uint32_t ubSize, uint32_t maxCoreNum, |
| @@ -877,14 +1052,26 @@ bool ComputeMergeIntraCorePlan(int64_t axisLen, int64_t unsortedDim, uint32_t ub | |||
| 877 | return false; | 1052 | return false; |
| 878 | } | 1053 | } |
| 879 | plan.actualCoreNum = static_cast<uint32_t>(actualCoreNum64); | 1054 | plan.actualCoreNum = static_cast<uint32_t>(actualCoreNum64); |
| 880 | - plan.blockSortSize = ComputeMergeIntraCoreBlockSortSize(ubSize); | ||
| 881 | - if (plan.blockSortSize == 0U) { | ||
| 882 | - return false; | ||
| 883 | - } | ||
| 884 | plan.extractChunkSize = ComputeMergeIntraCoreExtractChunkSize(ubSize); | 1055 | plan.extractChunkSize = ComputeMergeIntraCoreExtractChunkSize(ubSize); |
| 885 | if (plan.extractChunkSize == 0U) { | 1056 | if (plan.extractChunkSize == 0U) { |
| 886 | return false; | 1057 | return false; |
| 887 | } | 1058 | } |
| 1059 | + uint32_t singleBlockSortSize = ComputeMergeIntraCoreSingleBlockSortSize(ubSize); | ||
| 1060 | + if (axisLen <= static_cast<int64_t>(singleBlockSortSize)) { | ||
| 1061 | + uint64_t blockSortSize64 = Ops::Base::CeilAlign(static_cast<uint64_t>(axisLen), | ||
| 1062 | + static_cast<uint64_t>(MERGE_INTRA_CORE_SORT_ALIGN)); | ||
| 1063 | + if (blockSortSize64 == 0U || blockSortSize64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { | ||
| 1064 | + return false; | ||
| 1065 | + } | ||
| 1066 | + plan.blockSortSize = static_cast<uint32_t>(blockSortSize64); | ||
| 1067 | + plan.blocksPerRow = 1; | ||
| 1068 | + plan.alignNum = plan.blockSortSize; | ||
| 1069 | + return true; | ||
| 1070 | + } | ||
| 1071 | + plan.blockSortSize = ComputeMergeIntraCoreBlockSortSize(ubSize); | ||
| 1072 | + if (plan.blockSortSize == 0U) { | ||
| 1073 | + return false; | ||
| 1074 | + } | ||
| 888 | uint64_t blocksPerRow64 = (static_cast<uint64_t>(axisLen) + plan.blockSortSize - 1U) / plan.blockSortSize; | 1075 | uint64_t blocksPerRow64 = (static_cast<uint64_t>(axisLen) + plan.blockSortSize - 1U) / plan.blockSortSize; |
| 889 | if (blocksPerRow64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { | 1076 | if (blocksPerRow64 > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max())) { |
| 890 | return false; | 1077 | return false; |
| @@ -1258,13 +1445,14 @@ static bool EstimateSmallAxisTwoStageBatching(const SortKthTileInfo& info, uint3 | |||
| 1258 | // ============================================================================= | 1445 | // ============================================================================= |
| 1259 | static bool SelectSmallAxisRouteImpl(const SortKthTileInfo& info, uint32_t batchSizeCap, | 1446 | static bool SelectSmallAxisRouteImpl(const SortKthTileInfo& info, uint32_t batchSizeCap, |
| 1260 | std::function<bool(uint32_t, uint32_t&)> computeBatchNum, | 1447 | std::function<bool(uint32_t, uint32_t&)> computeBatchNum, |
| 1261 | - bool preferFullCoreBatching, SmallAxisRoutePlan& plan) | 1448 | + bool preferFullCoreBatching, SmallAxisRoutePlan& plan, |
| 1449 | + const SmallAxisRule* ruleOverride = nullptr) | ||
| 1262 | { | 1450 | { |
| 1263 | uint32_t axisLen = static_cast<uint32_t>(info.lastAxis); | 1451 | uint32_t axisLen = static_cast<uint32_t>(info.lastAxis); |
| 1264 | if (axisLen <= 1U) { | 1452 | if (axisLen <= 1U) { |
| 1265 | return false; | 1453 | return false; |
| 1266 | } | 1454 | } |
| 1267 | - const SmallAxisRule* rule = FindSmallAxisRule(info.dataType); | 1455 | + const SmallAxisRule* rule = ruleOverride == nullptr ? FindSmallAxisRule(info.dataType) : ruleOverride; |
| 1268 | if (rule == nullptr) { | 1456 | if (rule == nullptr) { |
| 1269 | return false; | 1457 | return false; |
| 1270 | } | 1458 | } |
| @@ -1292,7 +1480,7 @@ static bool SelectSmallAxisRouteImpl(const SortKthTileInfo& info, uint32_t batch | |||
| 1292 | return true; | 1480 | return true; |
| 1293 | } | 1481 | } |
| 1294 | 1482 | ||
| 1295 | -bool SelectSmallAxisRoute(const SortKthTileInfo& info, SmallAxisRoutePlan& plan) | 1483 | +bool SelectSmallAxisRoute(const SortKthTileInfo& info, SmallAxisRoutePlan& plan, const SmallAxisRule* ruleOverride) |
| 1296 | { | 1484 | { |
| 1297 | uint32_t fullCoreSegs = 0; | 1485 | uint32_t fullCoreSegs = 0; |
| 1298 | if (!CeilDivUint32(static_cast<uint64_t>(info.unsortedDim), static_cast<uint64_t>(info.maxCoreNum), fullCoreSegs)) { | 1486 | if (!CeilDivUint32(static_cast<uint64_t>(info.unsortedDim), static_cast<uint64_t>(info.maxCoreNum), fullCoreSegs)) { |
| @@ -1301,7 +1489,7 @@ bool SelectSmallAxisRoute(const SortKthTileInfo& info, SmallAxisRoutePlan& plan) | |||
| 1301 | auto computeBatchNum = [&info](uint32_t batchSize, uint32_t& batchNum) -> bool { | 1489 | auto computeBatchNum = [&info](uint32_t batchSize, uint32_t& batchNum) -> bool { |
| 1302 | return CeilDivUint32(static_cast<uint64_t>(info.unsortedDim), batchSize, batchNum); | 1490 | return CeilDivUint32(static_cast<uint64_t>(info.unsortedDim), batchSize, batchNum); |
| 1303 | }; | 1491 | }; |
| 1304 | - return SelectSmallAxisRouteImpl(info, fullCoreSegs, computeBatchNum, false, plan); | 1492 | + return SelectSmallAxisRouteImpl(info, fullCoreSegs, computeBatchNum, false, plan, ruleOverride); |
| 1305 | } | 1493 | } |
| 1306 | 1494 | ||
| 1307 | static bool SelectNonLastSmallAxisRouteImpl(const SortKthTileInfo& info, SmallAxisRoutePlan& plan, | 1495 | 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 | + 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,6 +518,9 @@ 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() |
| 399 | { | 525 | { |
| 400 | LocalTensor<CONVERT_TYPE> sortTempBuffer = this->sortedQueue_.template AllocTensor<CONVERT_TYPE>(); | 526 | LocalTensor<CONVERT_TYPE> sortTempBuffer = this->sortedQueue_.template AllocTensor<CONVERT_TYPE>(); |
| @@ -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 | } |
| @@ -247,13 +252,29 @@ protected: | |||
| 247 | } | 252 | } |
| 248 | } | 253 | } |
| 249 | 254 | ||
| 250 | - __aicore__ inline void SortRows(uint32_t curInnerChunk) | 255 | + __aicore__ inline void SortRows(uint32_t curInnerChunk, LocalTensor<SortT> proposalPing) |
| 251 | { | 256 | { |
| 257 | + LocalTensor<uint32_t> baseIndex; | ||
| 258 | + if constexpr (UseMergeSort) { | ||
| 259 | + // The last row's indices stay unused until its Extract, after every Sort32 has read this sequence. | ||
| 260 | + baseIndex = this->sortedIndex_[(curInnerChunk - 1U) * this->indexAxisElems_]; | ||
| 261 | + ArithProgression<int32_t>(baseIndex.template ReinterpretCast<int32_t>(), 0, 1, this->sortCount_); | ||
| 262 | + } | ||
| 263 | + // Keep one batched flip, but exclude trailing UB storage that Sort32 never reads. | ||
| 264 | + // sortCount_ includes the padding needed by the final Sort32 group. | ||
| 265 | + if constexpr (UseMergeSort && !IsDescend) { | ||
| 266 | + FlipMergeSignBit(this->sortInput_, (curInnerChunk - 1U) * this->valueAxisElems_ + this->sortCount_); | ||
| 267 | + } | ||
| 252 | for (uint32_t inner = 0; inner < curInnerChunk; ++inner) { | 268 | for (uint32_t inner = 0; inner < curInnerChunk; ++inner) { |
| 253 | LocalTensor<SortT> src = this->sortInput_[inner * this->valueAxisElems_]; | 269 | LocalTensor<SortT> src = this->sortInput_[inner * this->valueAxisElems_]; |
| 254 | - LocalTensor<SortT> dst = this->sortedValue_[inner * this->valueAxisElems_]; | 270 | + LocalTensor<SortT> dst = proposalPing[inner * this->valueAxisElems_]; |
| 255 | LocalTensor<uint32_t> dstIndex = this->sortedIndex_[inner * this->indexAxisElems_]; | 271 | LocalTensor<uint32_t> dstIndex = this->sortedIndex_[inner * this->indexAxisElems_]; |
| 256 | - if constexpr (NormalizeSignedZero) { | 272 | + if constexpr (UseMergeSort) { |
| 273 | + LocalTensor<SortT> pong = this->tmp_.template ReinterpretCast<SortT>(); | ||
| 274 | + uint32_t repeatTimes = this->sortCount_ / MergeSortConstants::DEALING_SORT_NUM_ONCE; | ||
| 275 | + bool resultInPing = PingPongMergeSortCommon::SortToProposal(dst, pong, src, baseIndex, repeatTimes); | ||
| 276 | + Extract(src, dstIndex, resultInPing ? dst : pong, repeatTimes); | ||
| 277 | + } else if constexpr (NormalizeSignedZero) { | ||
| 257 | SignedZeroSortCommon::PrepareSignedZeroKeysVec(src, this->sourceIndex_, this->sortCount_); | 278 | SignedZeroSortCommon::PrepareSignedZeroKeysVec(src, this->sourceIndex_, this->sortCount_); |
| 258 | AscendC::Sort<SortT, true, sortConfig_>(dst, dstIndex, src, this->tmp_, this->sortCount_); | 279 | AscendC::Sort<SortT, true, sortConfig_>(dst, dstIndex, src, this->tmp_, this->sortCount_); |
| 259 | SignedZeroSortCommon::RestoreSignedZeroValuesByIndexVec(dst, dstIndex, this->sourceIndex_, | 280 | SignedZeroSortCommon::RestoreSignedZeroValuesByIndexVec(dst, dstIndex, this->sourceIndex_, |
| @@ -262,6 +283,25 @@ protected: | |||
| 262 | AscendC::Sort<SortT, true, sortConfig_>(dst, dstIndex, src, this->tmp_, this->sortCount_); | 283 | AscendC::Sort<SortT, true, sortConfig_>(dst, dstIndex, src, this->tmp_, this->sortCount_); |
| 263 | } | 284 | } |
| 264 | } | 285 | } |
| 286 | + if constexpr (UseMergeSort) { | ||
| 287 | + if constexpr (!IsDescend) { | ||
| 288 | + FlipMergeSignBit(this->sortInput_, (curInnerChunk - 1U) * this->valueAxisElems_ + this->sortCount_); | ||
| 289 | + } | ||
| 290 | + // Extract writes the sorted values back to sortInput_ after its original input is dead, | ||
| 291 | + // so the input buffer is intentionally reused as the value output without extra UB. | ||
| 292 | + this->sortedValue_ = this->sortInput_; | ||
| 293 | + } | ||
| 294 | + } | ||
| 295 | + | ||
| 296 | + __aicore__ inline void FlipMergeSignBit(LocalTensor<SortT> tensor, uint32_t count) | ||
| 297 | + { | ||
| 298 | + if constexpr (std::is_same_v<SortT, float>) { | ||
| 299 | + LocalTensor<int32_t> bits = tensor.template ReinterpretCast<int32_t>(); | ||
| 300 | + Adds(bits, bits, MergeSortConstants::XOR_OP_VALUE_FP, count); | ||
| 301 | + } else if constexpr (std::is_same_v<SortT, half>) { | ||
| 302 | + LocalTensor<int16_t> bits = tensor.template ReinterpretCast<int16_t>(); | ||
| 303 | + Adds(bits, bits, MergeSortConstants::XOR_OP_VALUE_HALF, count); | ||
| 304 | + } | ||
| 265 | } | 305 | } |
| 266 | }; | 306 | }; |
| 267 | 307 | ||
| @@ -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 | ||
| @@ -130,7 +130,7 @@ __aicore__ inline void TwiddleInFp16Impl(LocalTensor<T1> inputX, LocalTensor<UT> | |||
| 130 | Reg::Duplicate(xorMaskReg, LOWEST_KEY_VALUE_B16, maskB16); | 130 | Reg::Duplicate(xorMaskReg, LOWEST_KEY_VALUE_B16, maskB16); |
| 131 | Reg::Duplicate(vandMask, XOR_OP_VALUE_B16, maskB16); | 131 | Reg::Duplicate(vandMask, XOR_OP_VALUE_B16, maskB16); |
| 132 | Reg::Duplicate(twiddledZeroReg, TWIDDLED_ZERO_BITS_FP16, maskB16); | 132 | Reg::Duplicate(twiddledZeroReg, TWIDDLED_ZERO_BITS_FP16, maskB16); |
| 133 | - // Median mode only: precompute the canonical NaN key and the abs-bits mask used below. | 133 | + // Precompute the canonical NaN key and abs-bits mask when normalization is enabled. |
| 134 | if constexpr (CanonicalizeNan) { | 134 | if constexpr (CanonicalizeNan) { |
| 135 | Reg::Duplicate(canonicalNanReg, CANONICAL_NAN_KEY_B16, maskB16); | 135 | Reg::Duplicate(canonicalNanReg, CANONICAL_NAN_KEY_B16, maskB16); |
| 136 | Reg::Duplicate(absMaskReg, FLOAT_ABS_MASK_B16, maskB16); | 136 | Reg::Duplicate(absMaskReg, FLOAT_ABS_MASK_B16, maskB16); |
| @@ -161,7 +161,7 @@ __aicore__ inline void TwiddleInFp16Impl(LocalTensor<T1> inputX, LocalTensor<UT> | |||
| 161 | // change -0.0 to +0.0 | 161 | // change -0.0 to +0.0 |
| 162 | Reg::RegTensor<uint16_t> resultReg; | 162 | Reg::RegTensor<uint16_t> resultReg; |
| 163 | Reg::Select(resultReg, twiddledZeroReg, xorVectorOne, minusZeroMask); | 163 | Reg::Select(resultReg, twiddledZeroReg, xorVectorOne, minusZeroMask); |
| 164 | - // Median mode only: detect NaN from the raw input bits (|bits| > inf, sign-independent) | 164 | + // Detect NaN from the raw input bits (|bits| > inf, sign-independent) |
| 165 | // and overwrite the twiddled key with the canonical NaN key, so every NaN lands in one | 165 | // and overwrite the twiddled key with the canonical NaN key, so every NaN lands in one |
| 166 | // shared max bucket. Only the output key changes; the input data stays untouched. | 166 | // shared max bucket. Only the output key changes; the input data stays untouched. |
| 167 | if constexpr (CanonicalizeNan) { | 167 | if constexpr (CanonicalizeNan) { |
| @@ -20,6 +20,7 @@ using namespace AscendC; | |||
| 20 | #define KTH_VALUE_RADIX_MORE_CORE_UINT32_TILING_KEY 258 | 20 | #define KTH_VALUE_RADIX_MORE_CORE_UINT32_TILING_KEY 258 |
| 21 | #define KTH_VALUE_MERGE_MORE_CORE_TILING_KEY 259 | 21 | #define KTH_VALUE_MERGE_MORE_CORE_TILING_KEY 259 |
| 22 | #define KTH_VALUE_RADIX_SELECT_TILING_KEY 267 | 22 | #define KTH_VALUE_RADIX_SELECT_TILING_KEY 267 |
| 23 | +#define KTH_VALUE_RESIDENT_HISTOGRAM_TILING_KEY 269 | ||
| 23 | 24 | ||
| 24 | template <uint64_t schId, uint64_t isInt32> | 25 | template <uint64_t schId, uint64_t isInt32> |
| 25 | __global__ __aicore__ void kth_value(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, GM_ADDR tiling) | 26 | __global__ __aicore__ void kth_value(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, GM_ADDR tiling) |
| @@ -29,8 +30,22 @@ __global__ __aicore__ void kth_value(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR | |||
| 29 | KERNEL_TASK_TYPE(KTH_VALUE_RADIX_MORE_CORE_UINT32_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | 30 | KERNEL_TASK_TYPE(KTH_VALUE_RADIX_MORE_CORE_UINT32_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); |
| 30 | KERNEL_TASK_TYPE(KTH_VALUE_MERGE_MORE_CORE_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | 31 | KERNEL_TASK_TYPE(KTH_VALUE_MERGE_MORE_CORE_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); |
| 31 | KERNEL_TASK_TYPE(KTH_VALUE_RADIX_SELECT_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | 32 | KERNEL_TASK_TYPE(KTH_VALUE_RADIX_SELECT_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); |
| 33 | + KERNEL_TASK_TYPE(KTH_VALUE_RESIDENT_HISTOGRAM_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | ||
| 32 | REGISTER_TILING_DEFAULT(KthValueTilingData); | 34 | REGISTER_TILING_DEFAULT(KthValueTilingData); |
| 33 | - GET_TILING_DATA_WITH_STRUCT(KthValueTilingData, tilingData, tiling); | ||
| 34 | TPipe pipe; | 35 | TPipe pipe; |
| 35 | - KthValue::Dispatch<false, schId, isInt32>(x, y1, y2, workspace, &tilingData, &pipe); | 36 | + // Host StoreKthValueTiling serializes a different layout for these single-core keys. |
| 37 | + // Register and decode the matching compact struct before Dispatch; a full-struct read | ||
| 38 | + // would use wrong field offsets. if constexpr removes the other layouts from each binary. | ||
| 39 | + if constexpr (schId == KTH_VALUE_SCHID_MERGE_SORT || schId == KTH_VALUE_SCHID_SORT32_SMALL_AXIS) { | ||
| 40 | + REGISTER_TILING_FOR_TILINGKEY(KTH_VALUE_MERGE_ONE_CORE_TILING_CONDITION, KthValueMergeOneCoreTilingData); | ||
| 41 | + GET_TILING_DATA_WITH_STRUCT(KthValueMergeOneCoreTilingData, tilingData, tiling); | ||
| 42 | + KthValue::Dispatch<false, schId, isInt32>(x, y1, y2, workspace, &tilingData, &pipe); | ||
| 43 | + } else if constexpr (schId == KTH_VALUE_SCHID_RADIX_ONE_CORE) { | ||
| 44 | + REGISTER_TILING_FOR_TILINGKEY(KTH_VALUE_RADIX_ONE_CORE_TILING_CONDITION, KthValueRadixOneCoreTilingData); | ||
| 45 | + GET_TILING_DATA_WITH_STRUCT(KthValueRadixOneCoreTilingData, tilingData, tiling); | ||
| 46 | + KthValue::Dispatch<false, schId, isInt32>(x, y1, y2, workspace, &tilingData, &pipe); | ||
| 47 | + } else { | ||
| 48 | + GET_TILING_DATA_WITH_STRUCT(KthValueTilingData, tilingData, tiling); | ||
| 49 | + KthValue::Dispatch<false, schId, isInt32>(x, y1, y2, workspace, &tilingData, &pipe); | ||
| 50 | + } | ||
| 36 | } | 51 | } |
| @@ -30,10 +30,11 @@ namespace KthValue { | |||
| 30 | using namespace AscendC; | 30 | using namespace AscendC; |
| 31 | 31 | ||
| 32 | template <bool EnableMedian, uint64_t schId> | 32 | template <bool EnableMedian, uint64_t schId> |
| 33 | -__aicore__ inline void RunMergeSortRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueTilingData* tilingData, TPipe* pipe) | 33 | +__aicore__ inline void RunMergeSortRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueMergeOneCoreTilingData* tilingData, |
| 34 | + TPipe* pipe) | ||
| 34 | { | 35 | { |
| 35 | constexpr uint64_t isSort32SmallAxis = (schId == KTH_VALUE_SCHID_SORT32_SMALL_AXIS); | 36 | constexpr uint64_t isSort32SmallAxis = (schId == KTH_VALUE_SCHID_SORT32_SMALL_AXIS); |
| 36 | - if constexpr (IsSameType<bfloat16_t, DTYPE_X>::value) { | 37 | + if constexpr (IsSameType<bfloat16_t, DTYPE_X>::value || IsSameType<int16_t, DTYPE_X>::value) { |
| 37 | KthValueMergeSortOneCore<DTYPE_X, float, isSort32SmallAxis, EnableMedian> op; | 38 | KthValueMergeSortOneCore<DTYPE_X, float, isSort32SmallAxis, EnableMedian> op; |
| 38 | op.Init(x, y1, y2, tilingData, pipe); | 39 | op.Init(x, y1, y2, tilingData, pipe); |
| 39 | op.Process(); | 40 | op.Process(); |
| @@ -158,8 +159,8 @@ __aicore__ inline void RunAxisOneCopyRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM | |||
| 158 | } | 159 | } |
| 159 | 160 | ||
| 160 | template <bool EnableMedian> | 161 | template <bool EnableMedian> |
| 161 | -__aicore__ inline void RunRadixOneCoreRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, KthValueTilingData* tilingData, | 162 | +__aicore__ inline void RunRadixOneCoreRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, |
| 162 | - TPipe* pipe) | 163 | + KthValueRadixOneCoreTilingData* tilingData, TPipe* pipe) |
| 163 | { | 164 | { |
| 164 | constexpr bool enableNanMode = EnableMedian && IS_MEDIAN_FLOAT_TYPE<DTYPE_X>; | 165 | constexpr bool enableNanMode = EnableMedian && IS_MEDIAN_FLOAT_TYPE<DTYPE_X>; |
| 165 | KthValueRadixOneCore<DTYPE_X, enableNanMode> op; | 166 | KthValueRadixOneCore<DTYPE_X, enableNanMode> op; |
| @@ -167,7 +168,7 @@ __aicore__ inline void RunRadixOneCoreRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, K | |||
| 167 | op.Process(); | 168 | op.Process(); |
| 168 | } | 169 | } |
| 169 | 170 | ||
| 170 | -template <bool EnableMedian> | 171 | +template <bool EnableMedian, bool UseResidentHistogram = false> |
| 171 | __aicore__ inline void RunRadixSelectRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, | 172 | __aicore__ inline void RunRadixSelectRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, |
| 172 | KthValueTilingData* tilingData, TPipe* pipe) | 173 | KthValueTilingData* tilingData, TPipe* pipe) |
| 173 | { | 174 | { |
| @@ -175,19 +176,19 @@ __aicore__ inline void RunRadixSelectRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM | |||
| 175 | if constexpr (sizeof(DTYPE_X) == 1) { | 176 | if constexpr (sizeof(DTYPE_X) == 1) { |
| 176 | KthValueRadixSelect<DTYPE_X, uint8_t, enableNanMode> op; | 177 | KthValueRadixSelect<DTYPE_X, uint8_t, enableNanMode> op; |
| 177 | op.Init(x, y1, y2, workspace, tilingData, pipe); | 178 | op.Init(x, y1, y2, workspace, tilingData, pipe); |
| 178 | - op.Process(); | 179 | + op.template Process<UseResidentHistogram>(); |
| 179 | } else if constexpr (sizeof(DTYPE_X) == 2) { | 180 | } else if constexpr (sizeof(DTYPE_X) == 2) { |
| 180 | KthValueRadixSelect<DTYPE_X, uint16_t, enableNanMode> op; | 181 | KthValueRadixSelect<DTYPE_X, uint16_t, enableNanMode> op; |
| 181 | op.Init(x, y1, y2, workspace, tilingData, pipe); | 182 | op.Init(x, y1, y2, workspace, tilingData, pipe); |
| 182 | - op.Process(); | 183 | + op.template Process<UseResidentHistogram>(); |
| 183 | } else if constexpr (sizeof(DTYPE_X) == 4) { | 184 | } else if constexpr (sizeof(DTYPE_X) == 4) { |
| 184 | KthValueRadixSelect<DTYPE_X, uint32_t, enableNanMode> op; | 185 | KthValueRadixSelect<DTYPE_X, uint32_t, enableNanMode> op; |
| 185 | op.Init(x, y1, y2, workspace, tilingData, pipe); | 186 | op.Init(x, y1, y2, workspace, tilingData, pipe); |
| 186 | - op.Process(); | 187 | + op.template Process<UseResidentHistogram>(); |
| 187 | } else if constexpr (sizeof(DTYPE_X) == 8) { | 188 | } else if constexpr (sizeof(DTYPE_X) == 8) { |
| 188 | KthValueRadixSelect<DTYPE_X, uint64_t, enableNanMode> op; | 189 | KthValueRadixSelect<DTYPE_X, uint64_t, enableNanMode> op; |
| 189 | op.Init(x, y1, y2, workspace, tilingData, pipe); | 190 | op.Init(x, y1, y2, workspace, tilingData, pipe); |
| 190 | - op.Process(); | 191 | + op.template Process<UseResidentHistogram>(); |
| 191 | } | 192 | } |
| 192 | } | 193 | } |
| 193 | 194 | ||
| @@ -232,8 +233,8 @@ __aicore__ inline bool TryRunSmallAxisRoute(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, K | |||
| 232 | 233 | ||
| 233 | // EnableMedian selects the shared Median/NanMedian behavior. tilingData->medianMode | 234 | // EnableMedian selects the shared Median/NanMedian behavior. tilingData->medianMode |
| 234 | // distinguishes propagate-NaN Median from ignore-NaN NanMedian at runtime. | 235 | // distinguishes propagate-NaN Median from ignore-NaN NanMedian at runtime. |
| 235 | -template <bool EnableMedian, uint64_t schId, uint64_t isInt32> | 236 | +template <bool EnableMedian, uint64_t schId, uint64_t isInt32, typename TilingData> |
| 236 | -__aicore__ inline void Dispatch(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, KthValueTilingData* tilingData, | 237 | +__aicore__ inline void Dispatch(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR workspace, TilingData* tilingData, |
| 237 | TPipe* pipe) | 238 | TPipe* pipe) |
| 238 | { | 239 | { |
| 239 | GM_ADDR userWorkspace = AscendC::GetUserWorkspace(workspace); | 240 | GM_ADDR userWorkspace = AscendC::GetUserWorkspace(workspace); |
| @@ -245,6 +246,8 @@ __aicore__ inline void Dispatch(GM_ADDR x, GM_ADDR y1, GM_ADDR y2, GM_ADDR works | |||
| 245 | RunRadixOneCoreRoute<EnableMedian>(x, y1, y2, tilingData, pipe); | 246 | RunRadixOneCoreRoute<EnableMedian>(x, y1, y2, tilingData, pipe); |
| 246 | } else if constexpr (schId == KTH_VALUE_SCHID_RADIX_MORE_CORE) { | 247 | } else if constexpr (schId == KTH_VALUE_SCHID_RADIX_MORE_CORE) { |
| 247 | RunRadixMoreCoreRoute<EnableMedian, isInt32>(x, y1, y2, userWorkspace, tilingData, pipe); | 248 | RunRadixMoreCoreRoute<EnableMedian, isInt32>(x, y1, y2, userWorkspace, tilingData, pipe); |
| 249 | + } else if constexpr (schId == KTH_VALUE_SCHID_RESIDENT_HISTOGRAM) { | ||
| 250 | + RunRadixSelectRoute<EnableMedian, true>(x, y1, y2, userWorkspace, tilingData, pipe); | ||
| 248 | } else if constexpr (schId == KTH_VALUE_SCHID_RADIX_SELECT) { | 251 | } else if constexpr (schId == KTH_VALUE_SCHID_RADIX_SELECT) { |
| 249 | RunRadixSelectRoute<EnableMedian>(x, y1, y2, userWorkspace, tilingData, pipe); | 252 | RunRadixSelectRoute<EnableMedian>(x, y1, y2, userWorkspace, tilingData, pipe); |
| 250 | } else if constexpr (schId == KTH_VALUE_SCHID_SMALL_AXIS_INSERTION || | 253 | } else if constexpr (schId == KTH_VALUE_SCHID_SMALL_AXIS_INSERTION || |
| @@ -32,7 +32,6 @@ namespace KthValue { | |||
| 32 | using namespace AscendC; | 32 | using namespace AscendC; |
| 33 | 33 | ||
| 34 | // Import shared constants from MergeSortConstants namespace | 34 | // Import shared constants from MergeSortConstants namespace |
| 35 | -using MergeSortConstants::DEALING_CONCAT_NUM_ONCE; | ||
| 36 | using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE; | 35 | using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE; |
| 37 | using MergeSortConstants::DEALING_SORT_NUM_ONCE; | 36 | using MergeSortConstants::DEALING_SORT_NUM_ONCE; |
| 38 | using MergeSortConstants::MERGE_INTRA_BUFFER_NUM; | 37 | using MergeSortConstants::MERGE_INTRA_BUFFER_NUM; |
| @@ -117,7 +116,6 @@ __aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend, E | |||
| 117 | this->batchSortLen_ = AscendC::GetSortLen<ValueType>(this->alignNum_); | 116 | this->batchSortLen_ = AscendC::GetSortLen<ValueType>(this->alignNum_); |
| 118 | this->sortBufferSize_ = this->blockSortLen_ * sizeof(ValueType); | 117 | this->sortBufferSize_ = this->blockSortLen_ * sizeof(ValueType); |
| 119 | this->sortRepeatTimes_ = this->blockSortSize_ / DEALING_SORT_NUM_ONCE; | 118 | this->sortRepeatTimes_ = this->blockSortSize_ / DEALING_SORT_NUM_ONCE; |
| 120 | - this->concatRepeatTimes_ = this->blockSortSize_ / DEALING_CONCAT_NUM_ONCE; | ||
| 121 | this->lastBlockSize_ = static_cast<uint32_t>(this->sortAxisNum_ - | 119 | this->lastBlockSize_ = static_cast<uint32_t>(this->sortAxisNum_ - |
| 122 | static_cast<int64_t>(this->blocksPerRow_ - 1) * this->blockSortSize_); | 120 | static_cast<int64_t>(this->blocksPerRow_ - 1) * this->blockSortSize_); |
| 123 | 121 | ||
| @@ -134,7 +132,6 @@ template <typename ValueType, typename IndexType, bool IsDescend, bool EnableMed | |||
| 134 | __aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend, EnableMedian>::InitPhase1Buffers() | 132 | __aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend, EnableMedian>::InitPhase1Buffers() |
| 135 | { | 133 | { |
| 136 | this->pipe_->InitBuffer(this->inQueueX_, MERGE_INTRA_BUFFER_NUM, this->blockSortSize_ * sizeof(ValueType)); | 134 | this->pipe_->InitBuffer(this->inQueueX_, MERGE_INTRA_BUFFER_NUM, this->blockSortSize_ * sizeof(ValueType)); |
| 137 | - this->pipe_->InitBuffer(this->concatTmpBuf_, this->sortBufferSize_); | ||
| 138 | this->pipe_->InitBuffer(this->sortTmpBuf_, this->sortBufferSize_); | 135 | this->pipe_->InitBuffer(this->sortTmpBuf_, this->sortBufferSize_); |
| 139 | this->pipe_->InitBuffer(this->sortedOutQueue_, MERGE_INTRA_BUFFER_NUM, this->sortBufferSize_); | 136 | this->pipe_->InitBuffer(this->sortedOutQueue_, MERGE_INTRA_BUFFER_NUM, this->sortBufferSize_); |
| 140 | this->pipe_->InitBuffer(this->indexTmpBuf_, this->blockSortSize_ * sizeof(uint32_t)); | 137 | this->pipe_->InitBuffer(this->indexTmpBuf_, this->blockSortSize_ * sizeof(uint32_t)); |
| @@ -178,7 +175,13 @@ __aicore__ inline void KthValueMergeIntraCore<ValueType, IndexType, IsDescend, E | |||
| 178 | { | 175 | { |
| 179 | int64_t outputOffset = batchIdx; | 176 | int64_t outputOffset = batchIdx; |
| 180 | if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<ValueType>) { | 177 | if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<ValueType>) { |
| 181 | - kthIndex_ = ResolveBatchK(batchIdx); | 178 | + if constexpr (KTH_VALUE_ENABLE_STATIC_MEDIAN_FAST_PATH) { |
| 179 | + if (medianMode_ != MEDIAN_MODE_STATIC) { | ||
| 180 | + kthIndex_ = ResolveBatchK(batchIdx); | ||
| 181 | + } | ||
| 182 | + } else { | ||
| 183 | + kthIndex_ = ResolveBatchK(batchIdx); | ||
| 184 | + } | ||
| 182 | } | 185 | } |
| 183 | 186 | ||
| 184 | // resultRegion: 0 = Ping (offset 0), 1 = Pong (offset batchSortLen_) | 187 | // resultRegion: 0 = Ping (offset 0), 1 = Pong (offset batchSortLen_) |
| @@ -26,7 +26,6 @@ namespace KthValue { | |||
| 26 | using namespace AscendC; | 26 | using namespace AscendC; |
| 27 | 27 | ||
| 28 | // Import shared constants from MergeSortConstants namespace | 28 | // Import shared constants from MergeSortConstants namespace |
| 29 | -using MergeSortConstants::DEALING_CONCAT_NUM_ONCE; | ||
| 30 | using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE; | 29 | using MergeSortConstants::DEALING_EXTRACT_NUM_ONCE; |
| 31 | using MergeSortConstants::DEALING_SORT_NUM_ONCE; | 30 | using MergeSortConstants::DEALING_SORT_NUM_ONCE; |
| 32 | using MergeSortConstants::MERGE_LIST_MAX_NUM; | 31 | using MergeSortConstants::MERGE_LIST_MAX_NUM; |
| @@ -48,6 +47,8 @@ struct KthValueMergeSortMoreCore : public MergeMoreCoreCommon::MergeMoreCoreBase | |||
| 48 | __aicore__ inline KthValueMergeSortMoreCore() {} | 47 | __aicore__ inline KthValueMergeSortMoreCore() {} |
| 49 | __aicore__ inline void Init(GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace, | 48 | __aicore__ inline void Init(GM_ADDR inputValue, GM_ADDR value, GM_ADDR indices, GM_ADDR workSpace, |
| 50 | const KthValueTilingData* tilingData, TPipe* pipe); | 49 | const KthValueTilingData* tilingData, TPipe* pipe); |
| 50 | + __aicore__ inline void PrepareRowWorkspace(); | ||
| 51 | + __aicore__ inline void InitSortBuffers(); | ||
| 51 | __aicore__ inline void InitMergeBuffers(); | 52 | __aicore__ inline void InitMergeBuffers(); |
| 52 | __aicore__ inline void ExtractAndCopyOut(); | 53 | __aicore__ inline void ExtractAndCopyOut(); |
| 53 | __aicore__ inline void OnInputLoaded(LocalTensor<T> inputLocal, uint32_t tileNum); | 54 | __aicore__ inline void OnInputLoaded(LocalTensor<T> inputLocal, uint32_t tileNum); |
| @@ -56,8 +57,13 @@ struct KthValueMergeSortMoreCore : public MergeMoreCoreCommon::MergeMoreCoreBase | |||
| 56 | 57 | ||
| 57 | // KthValue-specific member | 58 | // KthValue-specific member |
| 58 | uint32_t kthIndex_ = 0; | 59 | uint32_t kthIndex_ = 0; |
| 60 | + uint32_t staticKthIndex_ = 0; | ||
| 59 | uint32_t medianMode_ = 0; | 61 | uint32_t medianMode_ = 0; |
| 62 | + uint32_t nonNanCount_ = 0; | ||
| 63 | + TBuf<QuePosition::VECCALC> medianScratchBuf_; | ||
| 60 | GlobalTensor<uint32_t> medianCountGm_; | 64 | GlobalTensor<uint32_t> medianCountGm_; |
| 65 | + __gm__ CONVERT_TYPE* workspaceBase_ = nullptr; | ||
| 66 | + uint64_t rowWorkspaceElements_ = 0; | ||
| 61 | }; | 67 | }; |
| 62 | 68 | ||
| 63 | template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian> | 69 | template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian> |
| @@ -73,45 +79,72 @@ __aicore__ inline void KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, IN | |||
| 73 | this->outputLastDimValue_ = tilingData->lastDimTileNum; | 79 | this->outputLastDimValue_ = tilingData->lastDimTileNum; |
| 74 | this->numTileData_ = tilingData->numTileDataSize; | 80 | this->numTileData_ = tilingData->numTileDataSize; |
| 75 | this->frontCoreNum_ = tilingData->lastDimNeedCore; | 81 | this->frontCoreNum_ = tilingData->lastDimNeedCore; |
| 82 | + this->rowsPerRound_ = tilingData->unsortedDimParallel; | ||
| 83 | + this->sortLoopTimes_ = tilingData->sortLoopTimes; | ||
| 84 | + this->batchNum_ = tilingData->unsortedDimNum; | ||
| 85 | + this->logicalBlockSize_ = tilingData->keyParams1; | ||
| 86 | + this->syncMergeSortMode_ = this->logicalBlockSize_ > this->numTileData_ && this->numTileData_ > 0U; | ||
| 76 | kthIndex_ = tilingData->kthIndex; | 87 | kthIndex_ = tilingData->kthIndex; |
| 88 | + staticKthIndex_ = kthIndex_; | ||
| 77 | medianMode_ = tilingData->medianMode; | 89 | medianMode_ = tilingData->medianMode; |
| 78 | - if (this->frontCoreNum_ == 0U) { | 90 | + if (this->frontCoreNum_ == 0U || this->rowsPerRound_ == 0U || this->sortLoopTimes_ == 0U) { |
| 79 | return; | 91 | return; |
| 80 | } | 92 | } |
| 81 | uint32_t sortBufferSize = 8; | 93 | uint32_t sortBufferSize = 8; |
| 82 | - this->rowIdx_ = this->blockIdx_ / this->frontCoreNum_; | 94 | + this->rowGroupIdx_ = this->blockIdx_ / this->frontCoreNum_; |
| 83 | this->rowCoreIdx_ = this->blockIdx_ % this->frontCoreNum_; | 95 | this->rowCoreIdx_ = this->blockIdx_ % this->frontCoreNum_; |
| 84 | - this->rowDataOffset_ = static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(this->outputLastDimValue_); | ||
| 85 | // Per-row workspace stores Sort API sort-struct data. This capacity uses sortBufferSize bytes per | 96 | // Per-row workspace stores Sort API sort-struct data. This capacity uses sortBufferSize bytes per |
| 86 | // original element and UB-block byte alignment; it must cover later GetSortLen-based accesses. | 97 | // original element and UB-block byte alignment; it must cover later GetSortLen-based accesses. |
| 87 | uint64_t rowWorkspaceBytes = ROUND_UP_AGLIN_UINT64(static_cast<uint64_t>(this->outputLastDimValue_) * | 98 | uint64_t rowWorkspaceBytes = ROUND_UP_AGLIN_UINT64(static_cast<uint64_t>(this->outputLastDimValue_) * |
| 88 | sortBufferSize); | 99 | sortBufferSize); |
| 89 | - uint64_t rowWorkspaceElements = rowWorkspaceBytes / sizeof(CONVERT_TYPE); | 100 | + this->rowWorkspaceElements_ = rowWorkspaceBytes / sizeof(CONVERT_TYPE); |
| 90 | - this->rowWorkspaceOffset_ = static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(rowWorkspaceElements) * 2; | ||
| 91 | this->onceMaxElements_ = tilingData->keyParams0 / DEALING_SORT_NUM_ONCE * DEALING_SORT_NUM_ONCE; | 101 | this->onceMaxElements_ = tilingData->keyParams0 / DEALING_SORT_NUM_ONCE * DEALING_SORT_NUM_ONCE; |
| 92 | 102 | ||
| 93 | this->inputValueGm_.SetGlobalBuffer((__gm__ T*)(inputValue)); | 103 | this->inputValueGm_.SetGlobalBuffer((__gm__ T*)(inputValue)); |
| 94 | this->outValueGm_.SetGlobalBuffer((__gm__ T*)(value)); | 104 | this->outValueGm_.SetGlobalBuffer((__gm__ T*)(value)); |
| 95 | this->outIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(indices)); | 105 | this->outIndexGm_.SetGlobalBuffer((__gm__ INDEX_TYPE*)(indices)); |
| 96 | - this->workspaceGm_[0].SetGlobalBuffer((__gm__ CONVERT_TYPE*)(workSpace) + this->rowWorkspaceOffset_, | 106 | + this->workspaceBase_ = (__gm__ CONVERT_TYPE*)workSpace; |
| 97 | - rowWorkspaceElements); | ||
| 98 | - this->workspaceGm_[1].SetGlobalBuffer( | ||
| 99 | - (__gm__ CONVERT_TYPE*)(workSpace) + this->rowWorkspaceOffset_ + rowWorkspaceElements, rowWorkspaceElements); | ||
| 100 | uint64_t medianCountOffset = ROUND_UP_AGLIN_UINT64(static_cast<uint64_t>(tilingData->lastAxisNum) * | 107 | uint64_t medianCountOffset = ROUND_UP_AGLIN_UINT64(static_cast<uint64_t>(tilingData->lastAxisNum) * |
| 101 | static_cast<uint64_t>(tilingData->unsortedDimNum) * | 108 | static_cast<uint64_t>(tilingData->unsortedDimNum) * |
| 102 | MERGE_SORT_WORKSPACE_PARAM * sizeof(uint32_t)); | 109 | MERGE_SORT_WORKSPACE_PARAM * sizeof(uint32_t)); |
| 103 | medianCountGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint32_t*>(workSpace + medianCountOffset), | 110 | medianCountGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint32_t*>(workSpace + medianCountOffset), |
| 104 | tilingData->lastDimNeedCore * tilingData->unsortedDimNum); | 111 | tilingData->lastDimNeedCore * tilingData->unsortedDimNum); |
| 112 | +} | ||
| 105 | 113 | ||
| 106 | - uint32_t tailNum = this->outputLastDimValue_ - (this->frontCoreNum_ - 1) * this->numTileData_; | 114 | +// Called by the base Process() loop before each row's sort phase so the ping-pong |
| 107 | - uint32_t alignTile = ROUND_UP_AGLIN(tailNum); | 115 | +// workspace is rebound to the current row's GM region. |
| 116 | +template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian> | ||
| 117 | +__aicore__ inline void | ||
| 118 | +KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, EnableMedian>::PrepareRowWorkspace() | ||
| 119 | +{ | ||
| 120 | + this->rowDataOffset_ = static_cast<int64_t>(this->rowIdx_) * static_cast<int64_t>(this->outputLastDimValue_); | ||
| 121 | + this->rowWorkspaceOffset_ = static_cast<int64_t>(this->rowIdx_) * | ||
| 122 | + static_cast<int64_t>(this->rowWorkspaceElements_) * 2; | ||
| 123 | + this->workspaceGm_[0].SetGlobalBuffer(this->workspaceBase_ + this->rowWorkspaceOffset_, | ||
| 124 | + this->rowWorkspaceElements_); | ||
| 125 | + this->workspaceGm_[1].SetGlobalBuffer( | ||
| 126 | + this->workspaceBase_ + this->rowWorkspaceOffset_ + this->rowWorkspaceElements_, this->rowWorkspaceElements_); | ||
| 127 | + nonNanCount_ = 0U; | ||
| 128 | +} | ||
| 129 | + | ||
| 130 | +template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian> | ||
| 131 | +__aicore__ inline void | ||
| 132 | +KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, EnableMedian>::InitSortBuffers() | ||
| 133 | +{ | ||
| 134 | + uint32_t sortBufferSize = 8; | ||
| 135 | + uint32_t sortTileNum = this->syncMergeSortMode_ ? | ||
| 136 | + this->numTileData_ : | ||
| 137 | + this->outputLastDimValue_ - (this->frontCoreNum_ - 1) * this->numTileData_; | ||
| 138 | + uint32_t alignTile = ROUND_UP_AGLIN(sortTileNum); | ||
| 108 | this->pipe_->InitBuffer(this->inputQueue_, MERGE_MORE_BUFFER_NUM, alignTile * sizeof(T)); | 139 | this->pipe_->InitBuffer(this->inputQueue_, MERGE_MORE_BUFFER_NUM, alignTile * sizeof(T)); |
| 109 | 140 | ||
| 110 | this->pipe_->InitBuffer(this->sortedValueUb_, alignTile * sortBufferSize); | 141 | this->pipe_->InitBuffer(this->sortedValueUb_, alignTile * sortBufferSize); |
| 111 | this->pipe_->InitBuffer(this->sortedValueIndexUb_, alignTile * sizeof(uint32_t)); | 142 | this->pipe_->InitBuffer(this->sortedValueIndexUb_, alignTile * sizeof(uint32_t)); |
| 112 | - this->pipe_->InitBuffer(this->concatTempBuf_, alignTile * sortBufferSize); | ||
| 113 | this->pipe_->InitBuffer(this->sortTempBuf_, alignTile * sortBufferSize); | 143 | this->pipe_->InitBuffer(this->sortTempBuf_, alignTile * sortBufferSize); |
| 114 | - this->pipe_->InitBuffer(this->sortedValueLocalCastTbuf_, alignTile * sortBufferSize); | 144 | + if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) { |
| 145 | + uint32_t medianScratchBytes = ROUND_UP_AGLIN(this->frontCoreNum_ * sizeof(uint32_t)); | ||
| 146 | + this->pipe_->InitBuffer(medianScratchBuf_, medianScratchBytes); | ||
| 147 | + } | ||
| 115 | } | 148 | } |
| 116 | 149 | ||
| 117 | template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian> | 150 | template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian> |
| @@ -119,10 +152,15 @@ __aicore__ inline void KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, IN | |||
| 119 | LocalTensor<T> inputLocal, uint32_t tileNum) | 152 | LocalTensor<T> inputLocal, uint32_t tileNum) |
| 120 | { | 153 | { |
| 121 | if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) { | 154 | if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) { |
| 122 | - LocalTensor<float> scratch = this->concatTempBuf_.template Get<float>(); | 155 | + if constexpr (KTH_VALUE_ENABLE_STATIC_MEDIAN_FAST_PATH) { |
| 123 | - uint32_t nonNanCount = CountNonNan(inputLocal, tileNum, scratch, this->pipe_); | 156 | + if (medianMode_ == MEDIAN_MODE_STATIC) { |
| 157 | + return; | ||
| 158 | + } | ||
| 159 | + } | ||
| 160 | + LocalTensor<float> scratch = medianScratchBuf_.template Get<float>(); | ||
| 161 | + nonNanCount_ += CountNonNan(inputLocal, tileNum, scratch, this->pipe_); | ||
| 124 | LocalTensor<uint32_t> countLocal = scratch.template ReinterpretCast<uint32_t>(); | 162 | LocalTensor<uint32_t> countLocal = scratch.template ReinterpretCast<uint32_t>(); |
| 125 | - countLocal.SetValue(0, nonNanCount); | 163 | + countLocal.SetValue(0, nonNanCount_); |
| 126 | event_t eventId = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::S_MTE3)); | 164 | event_t eventId = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::S_MTE3)); |
| 127 | SetFlag<HardEvent::S_MTE3>(eventId); | 165 | SetFlag<HardEvent::S_MTE3>(eventId); |
| 128 | WaitFlag<HardEvent::S_MTE3>(eventId); | 166 | WaitFlag<HardEvent::S_MTE3>(eventId); |
| @@ -150,8 +188,9 @@ template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYP | |||
| 150 | __aicore__ inline uint32_t | 188 | __aicore__ inline uint32_t |
| 151 | KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, EnableMedian>::ResolveRowK() | 189 | KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, EnableMedian>::ResolveRowK() |
| 152 | { | 190 | { |
| 153 | - LocalTensor<uint32_t> countLocal = this->concatTempBuf_.template Get<uint32_t>(); | 191 | + LocalTensor<uint32_t> countLocal = medianScratchBuf_.template Get<uint32_t>(); |
| 154 | - uint32_t countOffset = this->rowIdx_ * this->frontCoreNum_; | 192 | + // Each round finishes before the next one starts, so the physical row-group count slots can be reused. |
| 193 | + uint32_t countOffset = this->rowGroupIdx_ * this->frontCoreNum_; | ||
| 155 | DataCopyExtParams copyParams{1, this->frontCoreNum_ * static_cast<uint32_t>(sizeof(uint32_t)), 0, 0, 0}; | 194 | DataCopyExtParams copyParams{1, this->frontCoreNum_ * static_cast<uint32_t>(sizeof(uint32_t)), 0, 0, 0}; |
| 156 | DataCopyPad(countLocal, medianCountGm_[countOffset], copyParams, {false, 0, 0, 0}); | 195 | DataCopyPad(countLocal, medianCountGm_[countOffset], copyParams, {false, 0, 0, 0}); |
| 157 | event_t eventId = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE2_S)); | 196 | event_t eventId = static_cast<event_t>(this->pipe_->FetchEventID(HardEvent::MTE2_S)); |
| @@ -161,7 +200,7 @@ KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, EnableMedian> | |||
| 161 | for (uint32_t core = 0U; core < this->frontCoreNum_; ++core) { | 200 | for (uint32_t core = 0U; core < this->frontCoreNum_; ++core) { |
| 162 | nonNanCount += countLocal.GetValue(core); | 201 | nonNanCount += countLocal.GetValue(core); |
| 163 | } | 202 | } |
| 164 | - return ResolveMedianK(kthIndex_, this->outputLastDimValue_, nonNanCount, medianMode_); | 203 | + return ResolveMedianK(staticKthIndex_, this->outputLastDimValue_, nonNanCount, medianMode_); |
| 165 | } | 204 | } |
| 166 | 205 | ||
| 167 | template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian> | 206 | template <typename T, typename CONVERT_TYPE, bool IS_DESCEND, typename INDEX_TYPE, bool EnableMedian> |
| @@ -186,8 +225,14 @@ __aicore__ inline void | |||
| 186 | KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, EnableMedian>::ExtractAndCopyOut() | 225 | KthValueMergeSortMoreCore<T, CONVERT_TYPE, IS_DESCEND, INDEX_TYPE, EnableMedian>::ExtractAndCopyOut() |
| 187 | { | 226 | { |
| 188 | if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) { | 227 | if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) { |
| 189 | - if (this->outOffset_ == 0) { | 228 | + if constexpr (KTH_VALUE_ENABLE_STATIC_MEDIAN_FAST_PATH) { |
| 190 | - kthIndex_ = ResolveRowK(); | 229 | + if (medianMode_ != MEDIAN_MODE_STATIC && this->outOffset_ == 0) { |
| 230 | + kthIndex_ = ResolveRowK(); | ||
| 231 | + } | ||
| 232 | + } else { | ||
| 233 | + if (this->outOffset_ == 0) { | ||
| 234 | + kthIndex_ = ResolveRowK(); | ||
| 235 | + } | ||
| 191 | } | 236 | } |
| 192 | } | 237 | } |
| 193 | LocalTensor<CONVERT_TYPE> sortTempBuffer = this->sortedQueue_.template DeQue<CONVERT_TYPE>(); | 238 | LocalTensor<CONVERT_TYPE> sortTempBuffer = this->sortedQueue_.template DeQue<CONVERT_TYPE>(); |
| @@ -12,28 +12,80 @@ | |||
| 12 | #define KTH_VALUE_MERGE_SORT_ONE_CORE_H | 12 | #define KTH_VALUE_MERGE_SORT_ONE_CORE_H |
| 13 | 13 | ||
| 14 | #include <cmath> | 14 | #include <cmath> |
| 15 | +#include <limits> | ||
| 15 | #include "kernel_operator.h" | 16 | #include "kernel_operator.h" |
| 16 | #include "op_kernel/platform_util.h" | 17 | #include "op_kernel/platform_util.h" |
| 17 | #include "kth_value_median_utils.h" | 18 | #include "kth_value_median_utils.h" |
| 18 | #include "kth_value_tiling_data.h" | 19 | #include "kth_value_tiling_data.h" |
| 19 | #include "common/merge_sort_constants.h" | 20 | #include "common/merge_sort_constants.h" |
| 21 | +#include "common/ping_pong_merge_sort.h" | ||
| 20 | #include "common/util_type_simd.h" | 22 | #include "common/util_type_simd.h" |
| 21 | 23 | ||
| 22 | namespace KthValue { | 24 | namespace KthValue { |
| 23 | using namespace AscendC; | 25 | using namespace AscendC; |
| 24 | 26 | ||
| 25 | -using MergeSortConstants::DEALING_CONCAT_NUM_ONCE; | ||
| 26 | using MergeSortConstants::DEALING_SORT_NUM_ONCE; | 27 | using MergeSortConstants::DEALING_SORT_NUM_ONCE; |
| 27 | using MergeSortConstants::XOR_OP_VALUE_FP; | 28 | using MergeSortConstants::XOR_OP_VALUE_FP; |
| 28 | using MergeSortConstants::XOR_OP_VALUE_HALF; | 29 | using MergeSortConstants::XOR_OP_VALUE_HALF; |
| 29 | 30 | ||
| 30 | constexpr uint32_t UB_BLOCK_SIZE = Ops::Base::GetUbBlockSize(); | 31 | constexpr uint32_t UB_BLOCK_SIZE = Ops::Base::GetUbBlockSize(); |
| 31 | 32 | ||
| 33 | +// Gather one fixed-rank element per row and widen its nonnegative index. | ||
| 34 | +// Rebase each value block so its relative gather offsets fit in uint16_t. | ||
| 35 | +__aicore__ inline void GatherKthInt16(LocalTensor<int16_t>& values, LocalTensor<uint32_t>& indices, | ||
| 36 | + LocalTensor<int16_t>& outValues, LocalTensor<int64_t>& outIndices, uint32_t rows, | ||
| 37 | + uint32_t stride, uint32_t kth) | ||
| 38 | +{ | ||
| 39 | + constexpr uint32_t lanes = Ops::Base::GetVRegSize() / sizeof(uint32_t); | ||
| 40 | + constexpr uint32_t wideLanes = Ops::Base::GetVRegSize() / sizeof(int64_t); | ||
| 41 | + __ubuf__ int16_t* valuePtr = reinterpret_cast<__ubuf__ int16_t*>(values.GetPhyAddr()); | ||
| 42 | + __ubuf__ uint32_t* indexPtr = reinterpret_cast<__ubuf__ uint32_t*>(indices.GetPhyAddr()); | ||
| 43 | + __ubuf__ int16_t* valueOut = reinterpret_cast<__ubuf__ int16_t*>(outValues.GetPhyAddr()); | ||
| 44 | + __ubuf__ int64_t* indexOut = reinterpret_cast<__ubuf__ int64_t*>(outIndices.GetPhyAddr()); | ||
| 45 | + uint16_t repeats = static_cast<uint16_t>(CeilDivision(rows, lanes)); | ||
| 46 | + uint32_t remain = rows; | ||
| 47 | + __VEC_SCOPE__ | ||
| 48 | + { | ||
| 49 | + Reg::RegTensor<int32_t> arange; | ||
| 50 | + Reg::RegTensor<uint32_t> offsets, selected, zero; | ||
| 51 | + Reg::RegTensor<int64_t> low, high; | ||
| 52 | + Reg::RegTensor<int16_t> result; | ||
| 53 | + Reg::RegTensor<uint16_t> narrowOffsets, highWords; | ||
| 54 | + Reg::MaskReg full = Reg::CreateMask<uint32_t>(); | ||
| 55 | + Reg::Arange(arange, 0); | ||
| 56 | + Reg::Muls(arange, arange, static_cast<int32_t>(stride), full); | ||
| 57 | + Reg::Adds(arange, arange, static_cast<int32_t>(kth), full); | ||
| 58 | + offsets = (Reg::RegTensor<uint32_t>&)arange; | ||
| 59 | + Reg::Duplicate(zero, 0U); | ||
| 60 | + Reg::DeInterleave(narrowOffsets, highWords, (Reg::RegTensor<uint16_t>&)offsets, | ||
| 61 | + (Reg::RegTensor<uint16_t>&)offsets); | ||
| 62 | + for (uint16_t rep = 0; rep < repeats; ++rep) { | ||
| 63 | + uint32_t count = remain < lanes ? remain : lanes; | ||
| 64 | + uint32_t valueCount = count; | ||
| 65 | + uint32_t lowCount = count < wideLanes ? count : wideLanes; | ||
| 66 | + uint32_t highCount = count > wideLanes ? count - wideLanes : 0; | ||
| 67 | + Reg::MaskReg valid = Reg::UpdateMask<uint32_t>(remain); | ||
| 68 | + Reg::MaskReg valueMask = Reg::UpdateMask<int16_t>(valueCount); | ||
| 69 | + Reg::MaskReg lowMask = Reg::UpdateMask<int64_t>(lowCount); | ||
| 70 | + Reg::MaskReg highMask = Reg::UpdateMask<int64_t>(highCount); | ||
| 71 | + Reg::Gather(result, valuePtr + rep * lanes * stride, narrowOffsets, valueMask); | ||
| 72 | + Reg::Gather(selected, indexPtr, offsets, valid); | ||
| 73 | + Reg::Interleave((Reg::RegTensor<uint32_t>&)low, (Reg::RegTensor<uint32_t>&)high, selected, zero); | ||
| 74 | + Reg::StoreAlign(valueOut + rep * lanes, result, valueMask); | ||
| 75 | + Reg::StoreAlign(indexOut + rep * lanes, low, lowMask); | ||
| 76 | + if (count > wideLanes) { | ||
| 77 | + Reg::StoreAlign(indexOut + rep * lanes + wideLanes, high, highMask); | ||
| 78 | + } | ||
| 79 | + Reg::Adds(offsets, offsets, stride * lanes, full); | ||
| 80 | + } | ||
| 81 | + } | ||
| 82 | +} | ||
| 83 | + | ||
| 32 | template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis = 0, bool EnableMedian = false> | 84 | template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis = 0, bool EnableMedian = false> |
| 33 | class KthValueMergeSortOneCore { | 85 | class KthValueMergeSortOneCore { |
| 34 | public: | 86 | public: |
| 35 | - __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling, | 87 | + __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, |
| 36 | - TPipe* pipe); | 88 | + const KthValueMergeOneCoreTilingData* tiling, TPipe* pipe); |
| 37 | __aicore__ inline void Process(); | 89 | __aicore__ inline void Process(); |
| 38 | 90 | ||
| 39 | private: | 91 | private: |
| @@ -43,27 +95,34 @@ private: | |||
| 43 | __aicore__ inline void FlipSignBit(LocalTensor<CONVERT_TYPE> xLocal, uint32_t offset, uint32_t count); | 95 | __aicore__ inline void FlipSignBit(LocalTensor<CONVERT_TYPE> xLocal, uint32_t offset, uint32_t count); |
| 44 | __aicore__ inline void SortRows(LocalTensor<T> xLocal, LocalTensor<T> sortedValueLocal, | 96 | __aicore__ inline void SortRows(LocalTensor<T> xLocal, LocalTensor<T> sortedValueLocal, |
| 45 | LocalTensor<uint32_t> sortedIndexLocal, uint32_t rowNum); | 97 | LocalTensor<uint32_t> sortedIndexLocal, uint32_t rowNum); |
| 98 | + __aicore__ inline void PadConvertedInt16(LocalTensor<CONVERT_TYPE> xSortLocal, uint32_t rowNum); | ||
| 99 | + __aicore__ inline void SortCompactBatches(LocalTensor<CONVERT_TYPE> xSortLocal, | ||
| 100 | + LocalTensor<CONVERT_TYPE> sortedValueCast, | ||
| 101 | + LocalTensor<uint32_t> sortedIndexLocal, uint32_t rowNum, | ||
| 102 | + uint32_t outputStride); | ||
| 103 | + __aicore__ inline void SelectScalarRows(LocalTensor<T> sortedValueLocal, LocalTensor<uint32_t> sortedIndexLocal, | ||
| 104 | + LocalTensor<T> compactValue, LocalTensor<int64_t> compactIndex, | ||
| 105 | + uint32_t rowNum, uint32_t outputStride, uint32_t outputK); | ||
| 46 | __aicore__ inline void CopyKthToGm(uint64_t outputOffset, uint32_t rowNum); | 106 | __aicore__ inline void CopyKthToGm(uint64_t outputOffset, uint32_t rowNum); |
| 47 | __aicore__ inline void ProcessSingleRound(uint32_t round); | 107 | __aicore__ inline void ProcessSingleRound(uint32_t round); |
| 48 | 108 | ||
| 49 | static constexpr uint32_t SORT_STRUCT_BYTES = 8; | 109 | static constexpr uint32_t SORT_STRUCT_BYTES = 8; |
| 50 | 110 | ||
| 111 | + uint32_t mergeBatchRows_ = 1U; | ||
| 51 | GlobalTensor<T> xGm_; | 112 | GlobalTensor<T> xGm_; |
| 52 | GlobalTensor<T> valuesGm_; | 113 | GlobalTensor<T> valuesGm_; |
| 53 | GlobalTensor<int64_t> indicesGm_; | 114 | GlobalTensor<int64_t> indicesGm_; |
| 54 | 115 | ||
| 55 | TPipe* pipe_{nullptr}; | 116 | TPipe* pipe_{nullptr}; |
| 56 | - const KthValueTilingData* tiling_{nullptr}; | 117 | + const KthValueMergeOneCoreTilingData* tiling_{nullptr}; |
| 57 | TQue<QuePosition::VECIN, 1> inQueueX_; | 118 | TQue<QuePosition::VECIN, 1> inQueueX_; |
| 58 | TQue<QuePosition::VECOUT, 1> outValueQueue_; | 119 | TQue<QuePosition::VECOUT, 1> outValueQueue_; |
| 59 | TQue<QuePosition::VECOUT, 1> outIndexQueue_; | 120 | TQue<QuePosition::VECOUT, 1> outIndexQueue_; |
| 60 | TQue<QuePosition::VECOUT, 1> compactValueQueue_; | 121 | TQue<QuePosition::VECOUT, 1> compactValueQueue_; |
| 61 | TQue<QuePosition::VECOUT, 1> compactIndexQueue_; | 122 | TQue<QuePosition::VECOUT, 1> compactIndexQueue_; |
| 62 | - TBuf<TPosition::VECCALC> concatTmpBuf_; | ||
| 63 | TBuf<TPosition::VECCALC> sortTmpBuf_; | 123 | TBuf<TPosition::VECCALC> sortTmpBuf_; |
| 64 | TBuf<TPosition::VECCALC> sortedLocalBuf_; | 124 | TBuf<TPosition::VECCALC> sortedLocalBuf_; |
| 65 | TBuf<TPosition::VECCALC> xCastBuf_; | 125 | TBuf<TPosition::VECCALC> xCastBuf_; |
| 66 | - TBuf<TPosition::VECCALC> sortedValueCastBuf_; | ||
| 67 | TBuf<TPosition::VECCALC> indexLocalBuf_; | 126 | TBuf<TPosition::VECCALC> indexLocalBuf_; |
| 68 | LocalTensor<uint32_t> indexLocal_; | 127 | LocalTensor<uint32_t> indexLocal_; |
| 69 | 128 | ||
| @@ -80,7 +139,7 @@ private: | |||
| 80 | 139 | ||
| 81 | template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian> | 140 | template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian> |
| 82 | __aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::Init( | 141 | __aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::Init( |
| 83 | - GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling, TPipe* pipe) | 142 | + GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueMergeOneCoreTilingData* tiling, TPipe* pipe) |
| 84 | { | 143 | { |
| 85 | if (tiling == nullptr || pipe == nullptr) { | 144 | if (tiling == nullptr || pipe == nullptr) { |
| 86 | return; | 145 | return; |
| @@ -96,22 +155,34 @@ __aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAx | |||
| 96 | 155 | ||
| 97 | uint32_t bufferNum = tiling_->keyParams4 == 0U ? 1U : tiling_->keyParams4; | 156 | uint32_t bufferNum = tiling_->keyParams4 == 0U ? 1U : tiling_->keyParams4; |
| 98 | pipe_->InitBuffer(inQueueX_, bufferNum, tiling_->keyParams1); | 157 | pipe_->InitBuffer(inQueueX_, bufferNum, tiling_->keyParams1); |
| 99 | - pipe_->InitBuffer(outValueQueue_, bufferNum, tiling_->keyParams1); | 158 | + constexpr bool compact = !EnableMedian && IsSameType<T, int16_t>::value; |
| 100 | - pipe_->InitBuffer(outIndexQueue_, bufferNum, tiling_->keyParams2); | 159 | + if constexpr (compact) { |
| 160 | + if (alignSize_ >= KTH_INT16_MERGE_BATCH_MIN_AXIS && alignSize_ <= KTH_INT16_MERGE_BATCH_MAX_AXIS && | ||
| 161 | + oneCoreRowNum_ >= KTH_INT16_MERGE_BATCH_ROWS) { | ||
| 162 | + mergeBatchRows_ = KTH_INT16_MERGE_BATCH_ROWS; | ||
| 163 | + } | ||
| 164 | + } | ||
| 165 | + uint32_t outputElements = DEALING_SORT_NUM_ONCE * oneCoreRowNum_; | ||
| 166 | + pipe_->InitBuffer(outValueQueue_, bufferNum, compact ? outputElements * sizeof(T) : tiling_->keyParams1); | ||
| 167 | + pipe_->InitBuffer(outIndexQueue_, bufferNum, compact ? outputElements * sizeof(uint32_t) : tiling_->keyParams2); | ||
| 101 | pipe_->InitBuffer(compactValueQueue_, bufferNum, ROUND_UP_AGLIN(oneCoreRowNum_ * sizeof(T))); | 168 | pipe_->InitBuffer(compactValueQueue_, bufferNum, ROUND_UP_AGLIN(oneCoreRowNum_ * sizeof(T))); |
| 102 | pipe_->InitBuffer(compactIndexQueue_, bufferNum, ROUND_UP_AGLIN(oneCoreRowNum_ * sizeof(int64_t))); | 169 | pipe_->InitBuffer(compactIndexQueue_, bufferNum, ROUND_UP_AGLIN(oneCoreRowNum_ * sizeof(int64_t))); |
| 103 | 170 | ||
| 104 | - uint32_t sortStructElems = alignSize_ * SORT_STRUCT_BYTES; | 171 | + uint32_t sortStructBytes = alignSize_ * SORT_STRUCT_BYTES * mergeBatchRows_; |
| 105 | - pipe_->InitBuffer(indexLocalBuf_, alignSize_ * sizeof(uint32_t)); | 172 | + pipe_->InitBuffer(indexLocalBuf_, alignSize_ * sizeof(uint32_t) * mergeBatchRows_); |
| 106 | - pipe_->InitBuffer(concatTmpBuf_, tiling_->tmpUbSize); | 173 | + pipe_->InitBuffer(sortTmpBuf_, sortStructBytes); |
| 107 | - pipe_->InitBuffer(sortTmpBuf_, sortStructElems * sizeof(CONVERT_TYPE)); | 174 | + pipe_->InitBuffer(sortedLocalBuf_, sortStructBytes); |
| 108 | - pipe_->InitBuffer(sortedLocalBuf_, sortStructElems * sizeof(CONVERT_TYPE)); | ||
| 109 | if constexpr (!IsSameType<T, CONVERT_TYPE>::value) { | 175 | if constexpr (!IsSameType<T, CONVERT_TYPE>::value) { |
| 110 | pipe_->InitBuffer(xCastBuf_, alignSize_ * oneCoreRowNum_ * sizeof(CONVERT_TYPE)); | 176 | pipe_->InitBuffer(xCastBuf_, alignSize_ * oneCoreRowNum_ * sizeof(CONVERT_TYPE)); |
| 111 | - pipe_->InitBuffer(sortedValueCastBuf_, alignSize_ * oneCoreRowNum_ * sizeof(CONVERT_TYPE)); | ||
| 112 | } | 177 | } |
| 113 | indexLocal_ = indexLocalBuf_.AllocTensor<uint32_t>(); | 178 | indexLocal_ = indexLocalBuf_.AllocTensor<uint32_t>(); |
| 114 | InitIndexLocal(); | 179 | InitIndexLocal(); |
| 180 | + // Each row keeps its own zero-based indices even when Sort32 processes multiple rows. | ||
| 181 | + if (mergeBatchRows_ > 1U) { | ||
| 182 | + for (uint32_t row = 1; row < mergeBatchRows_; ++row) { | ||
| 183 | + DataCopy(indexLocal_[row * alignSize_], indexLocal_, alignSize_); | ||
| 184 | + } | ||
| 185 | + } | ||
| 115 | } | 186 | } |
| 116 | 187 | ||
| 117 | template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian> | 188 | template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian> |
| @@ -152,7 +223,11 @@ __aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAx | |||
| 152 | uint64_t tileOffset, uint32_t currTileSize, uint32_t rowNum) | 223 | uint64_t tileOffset, uint32_t currTileSize, uint32_t rowNum) |
| 153 | { | 224 | { |
| 154 | LocalTensor<T> xLocal = inQueueX_.AllocTensor<T>(); | 225 | LocalTensor<T> xLocal = inQueueX_.AllocTensor<T>(); |
| 155 | - Duplicate(xLocal, static_cast<T>(NAN), alignSize_ * rowNum); | 226 | + T padding = 0; |
| 227 | + if constexpr (!IsSameType<T, int16_t>::value) { | ||
| 228 | + padding = static_cast<T>(NAN); | ||
| 229 | + } | ||
| 230 | + Duplicate(xLocal, padding, alignSize_ * rowNum); | ||
| 156 | event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_MTE2)); | 231 | event_t eventId = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_MTE2)); |
| 157 | SetFlag<HardEvent::V_MTE2>(eventId); | 232 | SetFlag<HardEvent::V_MTE2>(eventId); |
| 158 | WaitFlag<HardEvent::V_MTE2>(eventId); | 233 | WaitFlag<HardEvent::V_MTE2>(eventId); |
| @@ -160,8 +235,7 @@ __aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAx | |||
| 160 | uint32_t currTileSizeAlign = (currTileSize * sizeof(T) + UB_BLOCK_SIZE - 1U) / UB_BLOCK_SIZE * UB_BLOCK_SIZE / | 235 | uint32_t currTileSizeAlign = (currTileSize * sizeof(T) + UB_BLOCK_SIZE - 1U) / UB_BLOCK_SIZE * UB_BLOCK_SIZE / |
| 161 | sizeof(T); | 236 | sizeof(T); |
| 162 | uint32_t dstStride = ((alignSize_ - currTileSizeAlign) * sizeof(T)) / UB_BLOCK_SIZE; | 237 | uint32_t dstStride = ((alignSize_ - currTileSizeAlign) * sizeof(T)) / UB_BLOCK_SIZE; |
| 163 | - DataCopyPadExtParams<T> padParams{true, 0, static_cast<uint8_t>(currTileSizeAlign - currTileSize), | 238 | + DataCopyPadExtParams<T> padParams{true, 0, static_cast<uint8_t>(currTileSizeAlign - currTileSize), padding}; |
| 164 | - static_cast<T>(NAN)}; | ||
| 165 | DataCopyExtParams copyParam{static_cast<uint16_t>(rowNum), static_cast<uint32_t>(currTileSize * sizeof(T)), 0, | 239 | DataCopyExtParams copyParam{static_cast<uint16_t>(rowNum), static_cast<uint32_t>(currTileSize * sizeof(T)), 0, |
| 166 | dstStride, 0}; | 240 | dstStride, 0}; |
| 167 | DataCopyPad(xLocal, xGm_[tileOffset], copyParam, padParams); | 241 | DataCopyPad(xLocal, xGm_[tileOffset], copyParam, padParams); |
| @@ -181,47 +255,168 @@ __aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAx | |||
| 181 | } | 255 | } |
| 182 | } | 256 | } |
| 183 | 257 | ||
| 258 | +template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian> | ||
| 259 | +__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::PadConvertedInt16( | ||
| 260 | + LocalTensor<CONVERT_TYPE> xSortLocal, uint32_t rowNum) | ||
| 261 | +{ | ||
| 262 | + // Integer padding cannot represent infinity. Mark it after the exact | ||
| 263 | + // conversion so it sorts after every real value, including INT16_MAX. | ||
| 264 | + constexpr uint32_t vectorElems = Ops::Base::GetVRegSize() / sizeof(float); | ||
| 265 | + uint32_t tailStart = numTileData_ / vectorElems * vectorElems; | ||
| 266 | + uint32_t firstPadLane = numTileData_ - tailStart; | ||
| 267 | + uint32_t tailElems = alignSize_ - tailStart; | ||
| 268 | + uint32_t rowStride = alignSize_; | ||
| 269 | + __ubuf__ float* castPtr = reinterpret_cast<__ubuf__ float*>(xSortLocal.GetPhyAddr()); | ||
| 270 | + __VEC_SCOPE__ | ||
| 271 | + { | ||
| 272 | + Reg::RegTensor<float> infinity; | ||
| 273 | + Reg::RegTensor<int32_t> lanes; | ||
| 274 | + Reg::MaskReg valid = Reg::UpdateMask<float>(tailElems); | ||
| 275 | + Reg::MaskReg paddingMask; | ||
| 276 | + Reg::Arange(lanes, 0); | ||
| 277 | + Reg::Compares<int32_t, CMPMODE::GE>(paddingMask, lanes, static_cast<int32_t>(firstPadLane), valid); | ||
| 278 | + Reg::Duplicate(infinity, static_cast<float>(INFINITY)); | ||
| 279 | + for (uint16_t row = 0; row < static_cast<uint16_t>(rowNum); ++row) { | ||
| 280 | + Reg::StoreAlign<float, Reg::StoreDist::DIST_NORM_B32>(castPtr + row * rowStride + tailStart, infinity, | ||
| 281 | + paddingMask); | ||
| 282 | + } | ||
| 283 | + } | ||
| 284 | +} | ||
| 285 | + | ||
| 286 | +template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian> | ||
| 287 | +__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::SortCompactBatches( | ||
| 288 | + LocalTensor<CONVERT_TYPE> xSortLocal, LocalTensor<CONVERT_TYPE> sortedValueCast, | ||
| 289 | + LocalTensor<uint32_t> sortedIndexLocal, uint32_t rowNum, uint32_t outputStride) | ||
| 290 | +{ | ||
| 291 | + uint32_t sortRepeatTimes = alignSize_ / DEALING_SORT_NUM_ONCE; | ||
| 292 | + LocalTensor<CONVERT_TYPE> sortedLocal = sortedLocalBuf_.Get<CONVERT_TYPE>(); | ||
| 293 | + LocalTensor<CONVERT_TYPE> sortTmp = sortTmpBuf_.Get<CONVERT_TYPE>(); | ||
| 294 | + for (uint32_t first = 0; first < rowNum; first += mergeBatchRows_) { | ||
| 295 | + uint32_t rows = rowNum - first < mergeBatchRows_ ? rowNum - first : mergeBatchRows_; | ||
| 296 | + Sort32(sortedLocal, xSortLocal[first * alignSize_], indexLocal_, sortRepeatTimes * rows); | ||
| 297 | + bool inPing = true; | ||
| 298 | + uint32_t runCount = sortRepeatTimes; | ||
| 299 | + for (uint32_t run = DEALING_SORT_NUM_ONCE; runCount > 1U; run *= MergeSortConstants::MERGE_LIST_MAX_NUM) { | ||
| 300 | + for (uint32_t row = 0; row < rows; ++row) { | ||
| 301 | + uint32_t offset = GetSortOffset<CONVERT_TYPE>(row * alignSize_); | ||
| 302 | + if (inPing) { | ||
| 303 | + PingPongMergeSortCommon::MergeStage(sortTmp[offset], sortedLocal[offset], alignSize_, run, | ||
| 304 | + runCount); | ||
| 305 | + } else { | ||
| 306 | + PingPongMergeSortCommon::MergeStage(sortedLocal[offset], sortTmp[offset], alignSize_, run, | ||
| 307 | + runCount); | ||
| 308 | + } | ||
| 309 | + } | ||
| 310 | + inPing = !inPing; | ||
| 311 | + runCount = (runCount + MergeSortConstants::MERGE_LIST_MAX_NUM - 1U) / | ||
| 312 | + MergeSortConstants::MERGE_LIST_MAX_NUM; | ||
| 313 | + } | ||
| 314 | + LocalTensor<CONVERT_TYPE> proposal = inPing ? sortedLocal : sortTmp; | ||
| 315 | + __ubuf__ uint32_t* packed = (__ubuf__ uint32_t*)proposal.GetPhyAddr(); | ||
| 316 | + __ubuf__ uint32_t* valueOut = (__ubuf__ uint32_t*)sortedValueCast.GetPhyAddr() + first * outputStride; | ||
| 317 | + __ubuf__ uint32_t* indexOut = (__ubuf__ uint32_t*)sortedIndexLocal.GetPhyAddr() + first * outputStride; | ||
| 318 | + // Each FP32 proposal stores the value followed by its original B32 index. | ||
| 319 | + // Gather only k, keeping one result per row instead of extracting a Sort32 block. | ||
| 320 | + constexpr uint32_t proposalWords = SORT_STRUCT_BYTES / sizeof(uint32_t); | ||
| 321 | + constexpr uint32_t proposalIndexWord = 1U; | ||
| 322 | + static_assert(KTH_INT16_MERGE_BATCH_ROWS <= Ops::Base::GetVRegSize() / sizeof(uint32_t)); | ||
| 323 | + uint32_t proposalStride = alignSize_ * proposalWords; | ||
| 324 | + uint32_t selected = kthIndex_ * proposalWords; | ||
| 325 | + __VEC_SCOPE__ | ||
| 326 | + { | ||
| 327 | + Reg::RegTensor<uint32_t> lanes, srcIdx, values, indices; | ||
| 328 | + Reg::MaskReg mask = Reg::UpdateMask<uint32_t>(rows); | ||
| 329 | + Reg::Arange((Reg::RegTensor<int32_t>&)lanes, 0); | ||
| 330 | + Reg::Muls(srcIdx, lanes, proposalStride, mask); | ||
| 331 | + Reg::Adds(srcIdx, srcIdx, selected, mask); | ||
| 332 | + Reg::Gather(values, packed, srcIdx, mask); | ||
| 333 | + Reg::Adds(srcIdx, srcIdx, proposalIndexWord, mask); | ||
| 334 | + Reg::Gather(indices, packed, srcIdx, mask); | ||
| 335 | + Reg::StoreAlign(valueOut, values, mask); | ||
| 336 | + Reg::StoreAlign(indexOut, indices, mask); | ||
| 337 | + } | ||
| 338 | + } | ||
| 339 | +} | ||
| 340 | + | ||
| 184 | template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian> | 341 | template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian> |
| 185 | __aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::SortRows( | 342 | __aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::SortRows( |
| 186 | LocalTensor<T> xLocal, LocalTensor<T> sortedValueLocal, LocalTensor<uint32_t> sortedIndexLocal, uint32_t rowNum) | 343 | LocalTensor<T> xLocal, LocalTensor<T> sortedValueLocal, LocalTensor<uint32_t> sortedIndexLocal, uint32_t rowNum) |
| 187 | { | 344 | { |
| 188 | uint32_t sortRepeatTimes = alignSize_ / DEALING_SORT_NUM_ONCE; | 345 | uint32_t sortRepeatTimes = alignSize_ / DEALING_SORT_NUM_ONCE; |
| 189 | - uint32_t concatRepeatTimes = alignSize_ / DEALING_CONCAT_NUM_ONCE; | ||
| 190 | - LocalTensor<CONVERT_TYPE> concatTmp = concatTmpBuf_.Get<CONVERT_TYPE>(); | ||
| 191 | LocalTensor<CONVERT_TYPE> sortedLocal = sortedLocalBuf_.Get<CONVERT_TYPE>(); | 346 | LocalTensor<CONVERT_TYPE> sortedLocal = sortedLocalBuf_.Get<CONVERT_TYPE>(); |
| 192 | LocalTensor<CONVERT_TYPE> sortTmp = sortTmpBuf_.Get<CONVERT_TYPE>(); | 347 | LocalTensor<CONVERT_TYPE> sortTmp = sortTmpBuf_.Get<CONVERT_TYPE>(); |
| 193 | LocalTensor<CONVERT_TYPE> xSortLocal; | 348 | LocalTensor<CONVERT_TYPE> xSortLocal; |
| 194 | LocalTensor<CONVERT_TYPE> sortedValueCast; | 349 | LocalTensor<CONVERT_TYPE> sortedValueCast; |
| 195 | if constexpr (!IsSameType<T, CONVERT_TYPE>::value) { | 350 | if constexpr (!IsSameType<T, CONVERT_TYPE>::value) { |
| 196 | xSortLocal = xCastBuf_.Get<CONVERT_TYPE>(); | 351 | xSortLocal = xCastBuf_.Get<CONVERT_TYPE>(); |
| 197 | - sortedValueCast = sortedValueCastBuf_.Get<CONVERT_TYPE>(); | 352 | + // SortToProposal consumes each source row before Extract, so the cast input can hold the result. |
| 353 | + sortedValueCast = xSortLocal; | ||
| 198 | Cast(xSortLocal, xLocal, RoundMode::CAST_NONE, alignSize_ * rowNum); | 354 | Cast(xSortLocal, xLocal, RoundMode::CAST_NONE, alignSize_ * rowNum); |
| 355 | + if constexpr (IsSameType<T, int16_t>::value) { | ||
| 356 | + PadConvertedInt16(xSortLocal, rowNum); | ||
| 357 | + } | ||
| 199 | } else { | 358 | } else { |
| 200 | xSortLocal = xLocal; | 359 | xSortLocal = xLocal; |
| 201 | sortedValueCast = sortedValueLocal; | 360 | sortedValueCast = sortedValueLocal; |
| 202 | } | 361 | } |
| 203 | - for (uint32_t row = 0; row < rowNum; ++row) { | 362 | + if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<CONVERT_TYPE>) { |
| 204 | - uint32_t offset = row * alignSize_; | 363 | + if (medianMode_ != MEDIAN_MODE_STATIC) { |
| 205 | - if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<CONVERT_TYPE>) { | 364 | + for (uint32_t row = 0, offset = 0; row < rowNum; ++row, offset += alignSize_) { |
| 206 | - if (medianMode_ != MEDIAN_MODE_STATIC) { | ||
| 207 | CanonicalizeNanValues(xSortLocal[offset], alignSize_); | 365 | CanonicalizeNanValues(xSortLocal[offset], alignSize_); |
| 208 | } | 366 | } |
| 209 | } | 367 | } |
| 210 | - // MergeSort sorts sign-flipped ascending keys. Flip before Sort/Sort32 and flip back after Extract. | ||
| 211 | - FlipSignBit(xSortLocal, offset, alignSize_); | ||
| 212 | - if constexpr (isSort32SmallAxis == 1) { | ||
| 213 | - Sort32<CONVERT_TYPE>(sortedLocal, xSortLocal[offset], indexLocal_, 1); | ||
| 214 | - } else { | ||
| 215 | - LocalTensor<CONVERT_TYPE> concatLocal; | ||
| 216 | - Concat(concatLocal, xSortLocal[offset], concatTmp, concatRepeatTimes); | ||
| 217 | - AscendC::Sort<CONVERT_TYPE, true>(sortedLocal, concatLocal, indexLocal_, sortTmp, sortRepeatTimes); | ||
| 218 | - } | ||
| 219 | - Extract(sortedValueCast[offset], sortedIndexLocal[offset], sortedLocal, | ||
| 220 | - isSort32SmallAxis == 1 ? 1 : sortRepeatTimes); | ||
| 221 | - FlipSignBit(sortedValueCast, offset, alignSize_); | ||
| 222 | } | 368 | } |
| 369 | + // MergeSort sorts sign-flipped ascending keys. Flip before Sort/Sort32 and flip back after Extract. | ||
| 370 | + FlipSignBit(xSortLocal, 0, alignSize_ * rowNum); | ||
| 371 | + // Batched int16 selects k directly; the single-row fallback extracts its Sort32 block. | ||
| 372 | + constexpr bool compactExtract = !EnableMedian && IsSameType<T, int16_t>::value; | ||
| 373 | + uint32_t outputStride = compactExtract ? (mergeBatchRows_ > 1U ? 1U : DEALING_SORT_NUM_ONCE) : alignSize_; | ||
| 374 | + // Batch Sort32 and process rows at each merge level; each MergeStage remains confined to one row. | ||
| 375 | + if (compactExtract && mergeBatchRows_ > 1U) { | ||
| 376 | + SortCompactBatches(xSortLocal, sortedValueCast, sortedIndexLocal, rowNum, outputStride); | ||
| 377 | + } else { | ||
| 378 | + for (uint32_t row = 0, offset = 0; row < rowNum; ++row, offset += alignSize_) { | ||
| 379 | + bool resultInPing = PingPongMergeSortCommon::SortToProposal(sortedLocal, sortTmp, xSortLocal[offset], | ||
| 380 | + indexLocal_, sortRepeatTimes); | ||
| 381 | + // Packed results overwrite only rows whose input has already been consumed. | ||
| 382 | + LocalTensor<CONVERT_TYPE> proposal = resultInPing ? sortedLocal : sortTmp; | ||
| 383 | + uint32_t proposalOffset = compactExtract ? GetSortOffset<CONVERT_TYPE>(kthIndex_ / DEALING_SORT_NUM_ONCE * | ||
| 384 | + DEALING_SORT_NUM_ONCE) : | ||
| 385 | + 0U; | ||
| 386 | + Extract(sortedValueCast[row * outputStride], sortedIndexLocal[row * outputStride], proposal[proposalOffset], | ||
| 387 | + compactExtract || isSort32SmallAxis == 1 ? 1 : sortRepeatTimes); | ||
| 388 | + } | ||
| 389 | + } | ||
| 390 | + FlipSignBit(sortedValueCast, 0, outputStride * rowNum); | ||
| 223 | if constexpr (!IsSameType<T, CONVERT_TYPE>::value) { | 391 | if constexpr (!IsSameType<T, CONVERT_TYPE>::value) { |
| 224 | - Cast(sortedValueLocal, sortedValueCast, RoundMode::CAST_RINT, alignSize_ * rowNum); | 392 | + Cast(sortedValueLocal, sortedValueCast, RoundMode::CAST_RINT, outputStride * rowNum); |
| 393 | + } | ||
| 394 | +} | ||
| 395 | + | ||
| 396 | +template <typename T, typename CONVERT_TYPE, uint64_t isSort32SmallAxis, bool EnableMedian> | ||
| 397 | +__aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAxis, EnableMedian>::SelectScalarRows( | ||
| 398 | + LocalTensor<T> sortedValueLocal, LocalTensor<uint32_t> sortedIndexLocal, LocalTensor<T> compactValue, | ||
| 399 | + LocalTensor<int64_t> compactIndex, uint32_t rowNum, uint32_t outputStride, uint32_t outputK) | ||
| 400 | +{ | ||
| 401 | + for (uint32_t row = 0; row < rowNum; ++row) { | ||
| 402 | + uint32_t rowOffset = row * outputStride; | ||
| 403 | + uint32_t selectedK = outputK; | ||
| 404 | + if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) { | ||
| 405 | + if constexpr (KTH_VALUE_ENABLE_STATIC_MEDIAN_FAST_PATH) { | ||
| 406 | + if (medianMode_ != MEDIAN_MODE_STATIC) { | ||
| 407 | + uint32_t nonNanCount = CountNonNan(sortedValueLocal[rowOffset], numTileData_, | ||
| 408 | + sortTmpBuf_.Get<float>(), pipe_); | ||
| 409 | + selectedK = ResolveMedianK(kthIndex_, numTileData_, nonNanCount, medianMode_); | ||
| 410 | + } | ||
| 411 | + } else { | ||
| 412 | + uint32_t nonNanCount = CountNonNan(sortedValueLocal[rowOffset], numTileData_, sortTmpBuf_.Get<float>(), | ||
| 413 | + pipe_); | ||
| 414 | + selectedK = ResolveMedianK(kthIndex_, numTileData_, nonNanCount, medianMode_); | ||
| 415 | + } | ||
| 416 | + } | ||
| 417 | + uint32_t srcOffset = rowOffset + selectedK; | ||
| 418 | + compactValue.SetValue(row, sortedValueLocal.GetValue(srcOffset)); | ||
| 419 | + compactIndex.SetValue(row, static_cast<int64_t>(sortedIndexLocal.GetValue(srcOffset))); | ||
| 225 | } | 420 | } |
| 226 | } | 421 | } |
| 227 | 422 | ||
| @@ -232,31 +427,43 @@ __aicore__ inline void KthValueMergeSortOneCore<T, CONVERT_TYPE, isSort32SmallAx | |||
| 232 | LocalTensor<T> sortedValueLocal = outValueQueue_.DeQue<T>(); | 427 | LocalTensor<T> sortedValueLocal = outValueQueue_.DeQue<T>(); |
| 233 | LocalTensor<uint32_t> sortedIndexLocal = outIndexQueue_.DeQue<uint32_t>(); | 428 | LocalTensor<uint32_t> sortedIndexLocal = outIndexQueue_.DeQue<uint32_t>(); |
| 234 | 429 | ||
| 235 | - event_t eventIdVToS = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S)); | 430 | + constexpr bool compactExtract = !EnableMedian && IsSameType<T, int16_t>::value; |
| 236 | - SetFlag<HardEvent::V_S>(eventIdVToS); | 431 | + uint32_t outputStride = compactExtract ? (mergeBatchRows_ > 1U ? 1U : DEALING_SORT_NUM_ONCE) : alignSize_; |
| 237 | - WaitFlag<HardEvent::V_S>(eventIdVToS); | 432 | + uint32_t outputK = compactExtract ? (mergeBatchRows_ > 1U ? 0U : kthIndex_ % DEALING_SORT_NUM_ONCE) : kthIndex_; |
| 238 | - | 433 | + constexpr uint32_t gatherLanes = Ops::Base::GetVRegSize() / sizeof(uint32_t); |
| 434 | + constexpr uint32_t minGatherRows = 8U; | ||
| 435 | + // With k < stride, the largest B16 gather offset is 63*stride+k | ||
| 436 | + // <= 64*stride-1 <= 65471 for stride <= floor(65535/64) = 1023. | ||
| 437 | + constexpr uint32_t maxGatherStride = std::numeric_limits<uint16_t>::max() / gatherLanes; | ||
| 438 | + const bool useGather = !EnableMedian && IsSameType<T, int16_t>::value && rowNum >= minGatherRows && | ||
| 439 | + alignSize_ <= maxGatherStride; | ||
| 440 | + if (!useGather) { | ||
| 441 | + event_t eventIdVToS = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S)); | ||
| 442 | + SetFlag<HardEvent::V_S>(eventIdVToS); | ||
| 443 | + WaitFlag<HardEvent::V_S>(eventIdVToS); | ||
| 444 | + } | ||
| 239 | LocalTensor<T> compactValue = compactValueQueue_.AllocTensor<T>(); | 445 | LocalTensor<T> compactValue = compactValueQueue_.AllocTensor<T>(); |
| 240 | LocalTensor<int64_t> compactIndex = compactIndexQueue_.AllocTensor<int64_t>(); | 446 | LocalTensor<int64_t> compactIndex = compactIndexQueue_.AllocTensor<int64_t>(); |
| 241 | - for (uint32_t row = 0; row < rowNum; ++row) { | 447 | + if (useGather) { |
| 242 | - uint32_t rowOffset = row * alignSize_; | 448 | + LocalTensor<int16_t> sourceValues = sortedValueLocal.template ReinterpretCast<int16_t>(); |
| 243 | - uint32_t selectedK = kthIndex_; | 449 | + LocalTensor<int16_t> targetValues = compactValue.template ReinterpretCast<int16_t>(); |
| 244 | - if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) { | 450 | + GatherKthInt16(sourceValues, sortedIndexLocal, targetValues, compactIndex, rowNum, outputStride, outputK); |
| 245 | - uint32_t nonNanCount = CountNonNan(sortedValueLocal[rowOffset], numTileData_, concatTmpBuf_.Get<float>(), | 451 | + } else { |
| 246 | - pipe_); | 452 | + SelectScalarRows(sortedValueLocal, sortedIndexLocal, compactValue, compactIndex, rowNum, outputStride, outputK); |
| 247 | - selectedK = ResolveMedianK(kthIndex_, numTileData_, nonNanCount, medianMode_); | ||
| 248 | - } | ||
| 249 | - uint32_t srcOffset = rowOffset + selectedK; | ||
| 250 | - compactValue.SetValue(row, sortedValueLocal.GetValue(srcOffset)); | ||
| 251 | - compactIndex.SetValue(row, static_cast<int64_t>(sortedIndexLocal.GetValue(srcOffset))); | ||
| 252 | } | 453 | } |
| 253 | compactValueQueue_.EnQue<T>(compactValue); | 454 | compactValueQueue_.EnQue<T>(compactValue); |
| 254 | compactIndexQueue_.EnQue<int64_t>(compactIndex); | 455 | compactIndexQueue_.EnQue<int64_t>(compactIndex); |
| 255 | compactValue = compactValueQueue_.DeQue<T>(); | 456 | compactValue = compactValueQueue_.DeQue<T>(); |
| 256 | compactIndex = compactIndexQueue_.DeQue<int64_t>(); | 457 | compactIndex = compactIndexQueue_.DeQue<int64_t>(); |
| 257 | - event_t eventIdSToMte3 = static_cast<event_t>(pipe_->FetchEventID(HardEvent::S_MTE3)); | 458 | + if (useGather) { |
| 258 | - SetFlag<HardEvent::S_MTE3>(eventIdSToMte3); | 459 | + event_t ready = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_MTE3)); |
| 259 | - WaitFlag<HardEvent::S_MTE3>(eventIdSToMte3); | 460 | + SetFlag<HardEvent::V_MTE3>(ready); |
| 461 | + WaitFlag<HardEvent::V_MTE3>(ready); | ||
| 462 | + } else { | ||
| 463 | + event_t eventIdSToMte3 = static_cast<event_t>(pipe_->FetchEventID(HardEvent::S_MTE3)); | ||
| 464 | + SetFlag<HardEvent::S_MTE3>(eventIdSToMte3); | ||
| 465 | + WaitFlag<HardEvent::S_MTE3>(eventIdSToMte3); | ||
| 466 | + } | ||
| 260 | DataCopyExtParams valueCopyParam{1, static_cast<uint32_t>(rowNum * sizeof(T)), 0, 0, 0}; | 467 | DataCopyExtParams valueCopyParam{1, static_cast<uint32_t>(rowNum * sizeof(T)), 0, 0, 0}; |
| 261 | DataCopyPad(valuesGm_[outputOffset], compactValue, valueCopyParam); | 468 | DataCopyPad(valuesGm_[outputOffset], compactValue, valueCopyParam); |
| 262 | DataCopyExtParams indexCopyParam{1, static_cast<uint32_t>(rowNum * sizeof(int64_t)), 0, 0, 0}; | 469 | DataCopyExtParams indexCopyParam{1, static_cast<uint32_t>(rowNum * sizeof(int64_t)), 0, 0, 0}; |
| @@ -0,0 +1,173 @@ | |||
| 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 KTH_VALUE_NARROW_SELECT_H | ||
| 12 | +#define KTH_VALUE_NARROW_SELECT_H | ||
| 13 | + | ||
| 14 | +#include "kernel_operator.h" | ||
| 15 | +#include "simt_api/asc_simt.h" | ||
| 16 | +#include "common/small_axis_two_stage_base.h" | ||
| 17 | + | ||
| 18 | +namespace KthValue { | ||
| 19 | +constexpr uint32_t NARROW_SELECT_MODE = 1U; | ||
| 20 | +constexpr uint32_t NARROW_SELECT_WARP_SIZE = 32U; | ||
| 21 | +constexpr uint32_t NARROW_SELECT_CAPACITY_SMALL = 64U; | ||
| 22 | +constexpr uint32_t NARROW_SELECT_CAPACITY_MEDIUM = 128U; | ||
| 23 | +constexpr uint32_t NARROW_SELECT_CAPACITY_LARGE = 256U; | ||
| 24 | +constexpr uint32_t NARROW_SELECT_CAPACITY_MAX = 512U; | ||
| 25 | +constexpr uint32_t NARROW_SELECT_CAPACITY_BYTE = 384U; | ||
| 26 | +constexpr uint16_t NARROW_SELECT_ABS_MASK = 0x7FFFU; | ||
| 27 | +constexpr uint16_t NARROW_SELECT_SIGN_MASK = 0x8000U; | ||
| 28 | +constexpr uint16_t NARROW_SELECT_KEY_MASK = 0xFFFFU; | ||
| 29 | +constexpr uint16_t NARROW_SELECT_HALF_INF = 0x7C00U; | ||
| 30 | +constexpr uint16_t NARROW_SELECT_BFLOAT_INF = 0x7F80U; | ||
| 31 | +constexpr int32_t NARROW_SELECT_UPPER_SENTINEL = 65536; | ||
| 32 | +constexpr int32_t NARROW_SELECT_BYTE_SENTINEL = 256; | ||
| 33 | +constexpr int32_t NARROW_SELECT_LOWER_SENTINEL = -129; | ||
| 34 | +constexpr int32_t NARROW_SELECT_BINARY_RADIX = 2; | ||
| 35 | + | ||
| 36 | +// Ordering keys collapse both zeros and all NaN encodings. The final value is read | ||
| 37 | +// from the selected raw input, so the output preserves its zero sign or NaN payload. | ||
| 38 | +template <typename T> | ||
| 39 | +__simt_callee__ __aicore__ inline int32_t MakeNarrowFloatKey(uint16_t raw) | ||
| 40 | +{ | ||
| 41 | + uint16_t absBits = raw & NARROW_SELECT_ABS_MASK; | ||
| 42 | + constexpr uint16_t infBits = IsSameType<T, half>::value ? NARROW_SELECT_HALF_INF : NARROW_SELECT_BFLOAT_INF; | ||
| 43 | + if (absBits == 0U) { | ||
| 44 | + raw = 0U; | ||
| 45 | + } | ||
| 46 | + if (absBits > infBits) { | ||
| 47 | + return NARROW_SELECT_KEY_MASK; | ||
| 48 | + } | ||
| 49 | + if ((raw & NARROW_SELECT_SIGN_MASK) != 0U) { | ||
| 50 | + return (~raw) & NARROW_SELECT_KEY_MASK; | ||
| 51 | + } | ||
| 52 | + return raw ^ NARROW_SELECT_SIGN_MASK; | ||
| 53 | +} | ||
| 54 | + | ||
| 55 | +template <typename T> | ||
| 56 | +__simt_callee__ __aicore__ inline int32_t LoadNarrowKey(__ubuf__ T* input) | ||
| 57 | +{ | ||
| 58 | + if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) { | ||
| 59 | + return MakeNarrowFloatKey<T>(*reinterpret_cast<__ubuf__ uint16_t*>(input)); | ||
| 60 | + } else { | ||
| 61 | + return static_cast<int32_t>(*input); | ||
| 62 | + } | ||
| 63 | +} | ||
| 64 | + | ||
| 65 | +template <uint32_t Chunks> | ||
| 66 | +__simt_callee__ __aicore__ inline int32_t SelectNarrowKey(const int32_t (&keys)[Chunks], int32_t low, int32_t high, | ||
| 67 | + uint32_t kth) | ||
| 68 | +{ | ||
| 69 | + while (low < high) { | ||
| 70 | + int32_t mid = low + (high - low) / NARROW_SELECT_BINARY_RADIX; | ||
| 71 | + uint32_t count = 0; | ||
| 72 | +#pragma unroll | ||
| 73 | + for (uint32_t c = 0; c < Chunks; ++c) { | ||
| 74 | + count += keys[c] <= mid ? 1U : 0U; | ||
| 75 | + } | ||
| 76 | + count = asc_reduce_add(count); | ||
| 77 | + if (count > kth) { | ||
| 78 | + high = mid; | ||
| 79 | + } else { | ||
| 80 | + low = mid + 1; | ||
| 81 | + } | ||
| 82 | + } | ||
| 83 | + return low; | ||
| 84 | +} | ||
| 85 | + | ||
| 86 | +// Resolve ties in original column order: skip complete matching chunks, then select | ||
| 87 | +// the rank-th matching lane with an inclusive prefix. Every lane participates in collectives. | ||
| 88 | +template <uint32_t Chunks> | ||
| 89 | +__simt_callee__ __aicore__ inline uint32_t FindNarrowIndex(const int32_t (&keys)[Chunks], int32_t key, uint32_t kth, | ||
| 90 | + uint32_t lane, uint32_t cols) | ||
| 91 | +{ | ||
| 92 | + uint32_t less = 0; | ||
| 93 | +#pragma unroll | ||
| 94 | + for (uint32_t c = 0; c < Chunks; ++c) { | ||
| 95 | + less += keys[c] < key ? 1U : 0U; | ||
| 96 | + } | ||
| 97 | + less = asc_reduce_add(less); | ||
| 98 | + uint32_t rank = kth - less; | ||
| 99 | + uint32_t selected = cols; | ||
| 100 | +#pragma unroll | ||
| 101 | + for (uint32_t c = 0; c < Chunks; ++c) { | ||
| 102 | + uint32_t match = keys[c] == key ? 1U : 0U; | ||
| 103 | + uint32_t count = asc_reduce_add(match); | ||
| 104 | + if (rank < count) { | ||
| 105 | + uint32_t prefix = match; | ||
| 106 | +#pragma unroll | ||
| 107 | + for (uint32_t shift = 1; shift < NARROW_SELECT_WARP_SIZE; shift *= NARROW_SELECT_BINARY_RADIX) { | ||
| 108 | + uint32_t before = asc_shfl_up(prefix, shift); | ||
| 109 | + if (lane >= shift) { | ||
| 110 | + prefix += before; | ||
| 111 | + } | ||
| 112 | + } | ||
| 113 | + selected = match != 0 && prefix == rank + 1 ? c * NARROW_SELECT_WARP_SIZE + lane : cols; | ||
| 114 | + selected = asc_reduce_min(selected); | ||
| 115 | + break; | ||
| 116 | + } | ||
| 117 | + rank -= count; | ||
| 118 | + } | ||
| 119 | + return selected; | ||
| 120 | +} | ||
| 121 | + | ||
| 122 | +// One warp selects a narrow-key row without sorting; cached keys avoid repeated UB loads. | ||
| 123 | +template <typename T, uint32_t MaxAxis> | ||
| 124 | +__simt_vf__ LAUNCH_BOUND(SmallAxisCommon::TWO_STAGE_THREAD_NUM) __aicore__ | ||
| 125 | + void SimtSelectNarrowRows(uint32_t rows, uint32_t cols, uint32_t kth, uint64_t outputStart, __ubuf__ T* input, | ||
黄 | |||
| 126 | + __gm__ volatile T* values, __gm__ volatile int64_t* indices) | ||
| 127 | +{ | ||
| 128 | + constexpr bool isFloat = IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value; | ||
| 129 | + constexpr int32_t upperSentinel = isFloat ? NARROW_SELECT_UPPER_SENTINEL : NARROW_SELECT_BYTE_SENTINEL; | ||
| 130 | + constexpr uint32_t lanes = NARROW_SELECT_WARP_SIZE; | ||
| 131 | + constexpr uint32_t chunks = MaxAxis / lanes; | ||
| 132 | + // One warp owns one row; invalid tail lanes hold a key above every valid value. | ||
| 133 | + // Key selection finds the kth value first, then FindNarrowIndex recovers its stable index. | ||
| 134 | + uint32_t lane = static_cast<uint32_t>(threadIdx.x) % lanes; | ||
| 135 | + uint32_t warp = static_cast<uint32_t>(threadIdx.x) / lanes; | ||
| 136 | + for (uint32_t row = warp; row < rows; row += SmallAxisCommon::TWO_STAGE_THREAD_NUM / lanes) { | ||
| 137 | + int32_t keys[chunks]; | ||
| 138 | + int32_t low = upperSentinel; | ||
| 139 | + int32_t high = NARROW_SELECT_LOWER_SENTINEL; | ||
| 140 | +#pragma unroll | ||
| 141 | + for (uint32_t c = 0; c < chunks; ++c) { | ||
| 142 | + uint32_t col = c * lanes + lane; | ||
| 143 | + int32_t value = upperSentinel; | ||
| 144 | + if (col < cols) { | ||
| 145 | + value = LoadNarrowKey(input + row * cols + col); | ||
| 146 | + } | ||
| 147 | + keys[c] = value; | ||
| 148 | + if (col < cols) { | ||
| 149 | + low = value < low ? value : low; | ||
| 150 | + high = value > high ? value : high; | ||
| 151 | + } | ||
| 152 | + } | ||
| 153 | + low = asc_reduce_min(low); | ||
| 154 | + high = asc_reduce_max(high); | ||
| 155 | + uint32_t selected = kth; | ||
| 156 | + if (low != high) { | ||
| 157 | + low = SelectNarrowKey(keys, low, high, kth); | ||
| 158 | + selected = FindNarrowIndex(keys, low, kth, lane, cols); | ||
| 159 | + } | ||
| 160 | + if (lane == 0) { | ||
| 161 | + if constexpr (isFloat) { | ||
| 162 | + values[outputStart + row] = input[row * cols + selected]; | ||
| 163 | + } else { | ||
| 164 | + values[outputStart + row] = static_cast<T>(low); | ||
| 165 | + } | ||
| 166 | + indices[outputStart + row] = selected; | ||
| 167 | + } | ||
| 168 | + } | ||
| 169 | +} | ||
| 170 | + | ||
| 171 | +} // namespace KthValue | ||
| 172 | + | ||
| 173 | +#endif | ||
| @@ -194,6 +194,8 @@ __aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort, Enab | |||
| 194 | if (endTile > tileCount) { | 194 | if (endTile > tileCount) { |
| 195 | endTile = tileCount; | 195 | endTile = tileCount; |
| 196 | } | 196 | } |
| 197 | + // Merge Extract redirects sortedValue_ to sortInput_; preserve the proposal storage across tiles. | ||
| 198 | + LocalTensor<SortT> proposalPing = this->sortedValue_; | ||
| 197 | for (uint64_t tileId = startTile; tileId < endTile; ++tileId) { | 199 | for (uint64_t tileId = startTile; tileId < endTile; ++tileId) { |
| 198 | uint64_t outerId = tileId / this->innerLoopNum_; | 200 | uint64_t outerId = tileId / this->innerLoopNum_; |
| 199 | uint32_t innerTileId = static_cast<uint32_t>(tileId - outerId * static_cast<uint64_t>(this->innerLoopNum_)); | 201 | uint32_t innerTileId = static_cast<uint32_t>(tileId - outerId * static_cast<uint64_t>(this->innerLoopNum_)); |
| @@ -208,7 +210,7 @@ __aicore__ inline void KthValueNonLastSmallAxis<T, IsDescend, UseMergeSort, Enab | |||
| 208 | this->LoadTile(inputOffset, curInnerChunk); | 210 | this->LoadTile(inputOffset, curInnerChunk); |
| 209 | this->TransposeToSortMajor(curInnerChunk); | 211 | this->TransposeToSortMajor(curInnerChunk); |
| 210 | PrepareSortRows(curInnerChunk); | 212 | PrepareSortRows(curInnerChunk); |
| 211 | - this->SortRows(curInnerChunk); | 213 | + this->SortRows(curInnerChunk, proposalPing); |
| 212 | StoreTile(inputOffset, outputOffset, curInnerChunk); | 214 | StoreTile(inputOffset, outputOffset, curInnerChunk); |
| 213 | } | 215 | } |
| 214 | } | 216 | } |
| @@ -25,12 +25,13 @@ using namespace AscendC; | |||
| 25 | template <typename T, bool EnableMedian = false> | 25 | template <typename T, bool EnableMedian = false> |
| 26 | class KthValueRadixOneCore { | 26 | class KthValueRadixOneCore { |
| 27 | public: | 27 | public: |
| 28 | - __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, const KthValueTilingData* tiling, | 28 | + __aicore__ inline void Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, |
| 29 | - TPipe* pipe); | 29 | + const KthValueRadixOneCoreTilingData* tiling, TPipe* pipe); |
| 30 | __aicore__ inline void Process(); | 30 | __aicore__ inline void Process(); |
| 31 | 31 | ||
| 32 | private: | 32 | private: |
| 33 | __aicore__ inline void ParseTilingData(); | 33 | __aicore__ inline void ParseTilingData(); |
| 34 | + __aicore__ inline uint32_t PrepareMedianRank(LocalTensor<T> xLocal); | ||
| 34 | __aicore__ inline void CopyInputToUb(LocalTensor<T>& xLocal, int64_t row); | 35 | __aicore__ inline void CopyInputToUb(LocalTensor<T>& xLocal, int64_t row); |
| 35 | __aicore__ inline void ProcessOneRow(int64_t row, uint32_t localOffset, LocalTensor<T>& compactValue, | 36 | __aicore__ inline void ProcessOneRow(int64_t row, uint32_t localOffset, LocalTensor<T>& compactValue, |
| 36 | LocalTensor<int64_t>& compactIndex); | 37 | LocalTensor<int64_t>& compactIndex); |
| @@ -46,7 +47,7 @@ private: | |||
| 46 | GlobalTensor<int64_t> indicesGm_; | 47 | GlobalTensor<int64_t> indicesGm_; |
| 47 | 48 | ||
| 48 | TPipe* pipe_{nullptr}; | 49 | TPipe* pipe_{nullptr}; |
| 49 | - const KthValueTilingData* tiling_{nullptr}; | 50 | + const KthValueRadixOneCoreTilingData* tiling_{nullptr}; |
| 50 | TQue<QuePosition::VECIN, 1> inQueueX_; | 51 | TQue<QuePosition::VECIN, 1> inQueueX_; |
| 51 | TQue<QuePosition::VECOUT, 1> outValueQueue_; | 52 | TQue<QuePosition::VECOUT, 1> outValueQueue_; |
| 52 | TQue<QuePosition::VECOUT, 1> outIndexQueue_; | 53 | TQue<QuePosition::VECOUT, 1> outIndexQueue_; |
| @@ -73,7 +74,8 @@ private: | |||
| 73 | 74 | ||
| 74 | template <typename T, bool EnableMedian> | 75 | template <typename T, bool EnableMedian> |
| 75 | __aicore__ inline void KthValueRadixOneCore<T, EnableMedian>::Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, | 76 | __aicore__ inline void KthValueRadixOneCore<T, EnableMedian>::Init(GM_ADDR x, GM_ADDR values, GM_ADDR indices, |
| 76 | - const KthValueTilingData* tiling, TPipe* pipe) | 77 | + const KthValueRadixOneCoreTilingData* tiling, |
| 78 | + TPipe* pipe) | ||
| 77 | { | 79 | { |
| 78 | if (tiling == nullptr || pipe == nullptr) { | 80 | if (tiling == nullptr || pipe == nullptr) { |
| 79 | return; | 81 | return; |
| @@ -127,6 +129,29 @@ __aicore__ inline void KthValueRadixOneCore<T, EnableMedian>::CopyInputToUb(Loca | |||
| 127 | DataCopyPad(xLocal, xGm_[inputOffset], inputCopyParam, inputPadParam); | 129 | DataCopyPad(xLocal, xGm_[inputOffset], inputCopyParam, inputPadParam); |
| 128 | } | 130 | } |
| 129 | 131 | ||
| 132 | +template <typename T, bool EnableMedian> | ||
| 133 | +__aicore__ inline uint32_t KthValueRadixOneCore<T, EnableMedian>::PrepareMedianRank(LocalTensor<T> xLocal) | ||
| 134 | +{ | ||
| 135 | + uint32_t selectedK = kthIndex_; | ||
| 136 | + if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) { | ||
| 137 | + if constexpr (KTH_VALUE_ENABLE_STATIC_MEDIAN_FAST_PATH) { | ||
| 138 | + if (medianMode_ != MEDIAN_MODE_STATIC) { | ||
| 139 | + CanonicalizeNanValues(xLocal, numTileData_); | ||
| 140 | + uint32_t nonNanCount = CountNonNan(xLocal, numTileData_, tmpUb_.Get<float>(), pipe_); | ||
| 141 | + selectedK = ResolveMedianK(kthIndex_, numTileData_, nonNanCount, medianMode_); | ||
| 142 | + } | ||
| 143 | + } else { | ||
| 144 | + if (medianMode_ != MEDIAN_MODE_STATIC) { | ||
| 145 | + CanonicalizeNanValues(xLocal, numTileData_); | ||
| 146 | + } | ||
| 147 | + uint32_t nonNanCount = CountNonNan(xLocal, numTileData_, tmpUb_.Get<float>(), pipe_); | ||
| 148 | + selectedK = ResolveMedianK(kthIndex_, numTileData_, nonNanCount, medianMode_); | ||
| 149 | + } | ||
| 150 | + } | ||
| 151 | + | ||
| 152 | + return selectedK; | ||
| 153 | +} | ||
| 154 | + | ||
| 130 | template <typename T, bool EnableMedian> | 155 | template <typename T, bool EnableMedian> |
| 131 | __aicore__ inline void KthValueRadixOneCore<T, EnableMedian>::ProcessOneRow(int64_t row, uint32_t localOffset, | 156 | __aicore__ inline void KthValueRadixOneCore<T, EnableMedian>::ProcessOneRow(int64_t row, uint32_t localOffset, |
| 132 | LocalTensor<T>& compactValue, | 157 | LocalTensor<T>& compactValue, |
| @@ -139,14 +164,7 @@ __aicore__ inline void KthValueRadixOneCore<T, EnableMedian>::ProcessOneRow(int6 | |||
| 139 | LocalTensor<T> sortedValue = outValueQueue_.AllocTensor<T>(); | 164 | LocalTensor<T> sortedValue = outValueQueue_.AllocTensor<T>(); |
| 140 | LocalTensor<uint32_t> sortedIndex = outIndexQueue_.AllocTensor<uint32_t>(); | 165 | LocalTensor<uint32_t> sortedIndex = outIndexQueue_.AllocTensor<uint32_t>(); |
| 141 | LocalTensor<uint8_t> tmpUb = tmpUb_.Get<uint8_t>(); | 166 | LocalTensor<uint8_t> tmpUb = tmpUb_.Get<uint8_t>(); |
| 142 | - uint32_t selectedK = kthIndex_; | 167 | + uint32_t selectedK = PrepareMedianRank(xLocal); |
| 143 | - if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) { | ||
| 144 | - if (medianMode_ != MEDIAN_MODE_STATIC) { | ||
| 145 | - CanonicalizeNanValues(xLocal, numTileData_); | ||
| 146 | - } | ||
| 147 | - uint32_t nonNanCount = CountNonNan(xLocal, numTileData_, tmpUb_.Get<float>(), pipe_); | ||
| 148 | - selectedK = ResolveMedianK(kthIndex_, numTileData_, nonNanCount, medianMode_); | ||
| 149 | - } | ||
| 150 | 168 | ||
| 151 | bool hasNegativeZero = false; | 169 | bool hasNegativeZero = false; |
| 152 | if constexpr (SignedZeroSortCommon::IS_FLOATING_POINT_V<T>) { | 170 | if constexpr (SignedZeroSortCommon::IS_FLOATING_POINT_V<T>) { |
| @@ -18,6 +18,7 @@ | |||
| 18 | #include "simt_api/asc_simt.h" | 18 | #include "simt_api/asc_simt.h" |
| 19 | #include "kth_value_median_utils.h" | 19 | #include "kth_value_median_utils.h" |
| 20 | #include "kth_value_tiling_data.h" | 20 | #include "kth_value_tiling_data.h" |
| 21 | +#include "kth_value_narrow_select.h" | ||
| 21 | #include "common/small_axis_two_stage_base.h" | 22 | #include "common/small_axis_two_stage_base.h" |
| 22 | 23 | ||
| 23 | namespace KthValue { | 24 | namespace KthValue { |
| @@ -69,6 +70,9 @@ __simt_vf__ LAUNCH_BOUND(SmallAxisCommon::TWO_STAGE_THREAD_NUM) __aicore__ | |||
| 69 | } | 70 | } |
| 70 | } | 71 | } |
| 71 | 72 | ||
| 73 | +// The shared base sorts each short segment in two stages; this adapter emits only its kth pair. | ||
| 74 | +// Partial batches use validSegs throughout, and indices stay local to the selected axis. | ||
| 75 | +// Median hooks canonicalize ordering keys before sorting and resolve NaN policy before the final gather. | ||
| 72 | template <typename T, bool EnableMedian = false> | 76 | template <typename T, bool EnableMedian = false> |
| 73 | class KthValueSmallAxisTwoStage | 77 | class KthValueSmallAxisTwoStage |
| 74 | : public SmallAxisCommon::SmallAxisTwoStageBase<KthValueSmallAxisTwoStage<T, EnableMedian>, T, uint32_t, false> { | 78 | : public SmallAxisCommon::SmallAxisTwoStageBase<KthValueSmallAxisTwoStage<T, EnableMedian>, T, uint32_t, false> { |
| @@ -101,6 +105,7 @@ private: | |||
| 101 | __aicore__ inline void LoadBatch(uint32_t batchId, uint32_t validSegs, uint32_t totalElems); | 105 | __aicore__ inline void LoadBatch(uint32_t batchId, uint32_t validSegs, uint32_t totalElems); |
| 102 | __aicore__ inline void StoreKth(uint32_t batchId, int64_t segStart, uint32_t validSegs); | 106 | __aicore__ inline void StoreKth(uint32_t batchId, int64_t segStart, uint32_t validSegs); |
| 103 | 107 | ||
| 108 | + TBuf<TPosition::VECCALC> narrowInputBuf_; | ||
| 104 | GlobalTensor<T> inputGm_; | 109 | GlobalTensor<T> inputGm_; |
| 105 | GlobalTensor<T> valueGm_; | 110 | GlobalTensor<T> valueGm_; |
| 106 | GlobalTensor<int64_t> indexGm_; | 111 | GlobalTensor<int64_t> indexGm_; |
| @@ -143,6 +148,14 @@ __aicore__ inline void KthValueSmallAxisTwoStage<T, EnableMedian>::Init(GM_ADDR | |||
| 143 | return; | 148 | return; |
| 144 | } | 149 | } |
| 145 | 150 | ||
| 151 | + if constexpr (!EnableMedian && (IsSameType<T, int8_t>::value || IsSameType<T, uint8_t>::value || | ||
| 152 | + IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value)) { | ||
| 153 | + if (!IsNonLastMode() && tiling_->keyParams3 == NARROW_SELECT_MODE) { | ||
| 154 | + pipe_->InitBuffer(narrowInputBuf_, ROUND_UP_AGLIN(maxFlatElems_ * sizeof(T))); | ||
| 155 | + inputValues_ = narrowInputBuf_.template Get<T>(); | ||
| 156 | + return; | ||
| 157 | + } | ||
| 158 | + } | ||
| 146 | Base::InitSortBuffers(pipe, maxFlatElems_, tiling_->tmpUbSize, tiling_->keyParams2 != 0U, sizeof(uint32_t)); | 159 | Base::InitSortBuffers(pipe, maxFlatElems_, tiling_->tmpUbSize, tiling_->keyParams2 != 0U, sizeof(uint32_t)); |
| 147 | } | 160 | } |
| 148 | 161 | ||
| @@ -181,6 +194,44 @@ __aicore__ inline void KthValueSmallAxisTwoStage<T, EnableMedian>::ProcessBatch( | |||
| 181 | uint32_t totalElems = validSegs * segmentLen_; | 194 | uint32_t totalElems = validSegs * segmentLen_; |
| 182 | int64_t segStart = GetOutputStart(batchId); | 195 | int64_t segStart = GetOutputStart(batchId); |
| 183 | LoadBatch(batchId, validSegs, totalElems); | 196 | LoadBatch(batchId, validSegs, totalElems); |
| 197 | + if constexpr (!EnableMedian && (IsSameType<T, int8_t>::value || IsSameType<T, uint8_t>::value || | ||
| 198 | + IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value)) { | ||
| 199 | + if (!IsNonLastMode() && tiling_->keyParams3 == NARROW_SELECT_MODE) { | ||
| 200 | + bool useByteCapacity = (IsSameType<T, int8_t>::value || IsSameType<T, uint8_t>::value) && | ||
| 201 | + segmentLen_ > NARROW_SELECT_CAPACITY_MEDIUM && | ||
| 202 | + segmentLen_ <= NARROW_SELECT_CAPACITY_BYTE; | ||
| 203 | + if (useByteCapacity) { | ||
| 204 | + asc_vf_call<SimtSelectNarrowRows<T, NARROW_SELECT_CAPACITY_BYTE>>( | ||
| 205 | + dim3(SmallAxisCommon::TWO_STAGE_THREAD_NUM), validSegs, segmentLen_, kthIndex_, | ||
| 206 | + static_cast<uint64_t>(segStart), (__ubuf__ T*)inputValues_.GetPhyAddr(), | ||
| 207 | + (__gm__ volatile T*)valueGm_.GetPhyAddr(), (__gm__ volatile int64_t*)indexGm_.GetPhyAddr()); | ||
| 208 | + } else if (segmentLen_ <= NARROW_SELECT_CAPACITY_SMALL) { | ||
| 209 | + asc_vf_call<SimtSelectNarrowRows<T, NARROW_SELECT_CAPACITY_SMALL>>( | ||
| 210 | + dim3(SmallAxisCommon::TWO_STAGE_THREAD_NUM), validSegs, segmentLen_, kthIndex_, | ||
| 211 | + static_cast<uint64_t>(segStart), (__ubuf__ T*)inputValues_.GetPhyAddr(), | ||
| 212 | + (__gm__ volatile T*)valueGm_.GetPhyAddr(), (__gm__ volatile int64_t*)indexGm_.GetPhyAddr()); | ||
| 213 | + } else if (segmentLen_ <= NARROW_SELECT_CAPACITY_MEDIUM) { | ||
| 214 | + asc_vf_call<SimtSelectNarrowRows<T, NARROW_SELECT_CAPACITY_MEDIUM>>( | ||
| 215 | + dim3(SmallAxisCommon::TWO_STAGE_THREAD_NUM), validSegs, segmentLen_, kthIndex_, | ||
| 216 | + static_cast<uint64_t>(segStart), (__ubuf__ T*)inputValues_.GetPhyAddr(), | ||
| 217 | + (__gm__ volatile T*)valueGm_.GetPhyAddr(), (__gm__ volatile int64_t*)indexGm_.GetPhyAddr()); | ||
| 218 | + } else if (segmentLen_ <= NARROW_SELECT_CAPACITY_LARGE) { | ||
| 219 | + asc_vf_call<SimtSelectNarrowRows<T, NARROW_SELECT_CAPACITY_LARGE>>( | ||
| 220 | + dim3(SmallAxisCommon::TWO_STAGE_THREAD_NUM), validSegs, segmentLen_, kthIndex_, | ||
| 221 | + static_cast<uint64_t>(segStart), (__ubuf__ T*)inputValues_.GetPhyAddr(), | ||
| 222 | + (__gm__ volatile T*)valueGm_.GetPhyAddr(), (__gm__ volatile int64_t*)indexGm_.GetPhyAddr()); | ||
| 223 | + } else if (segmentLen_ <= NARROW_SELECT_CAPACITY_MAX) { | ||
| 224 | + asc_vf_call<SimtSelectNarrowRows<T, NARROW_SELECT_CAPACITY_MAX>>( | ||
| 225 | + dim3(SmallAxisCommon::TWO_STAGE_THREAD_NUM), validSegs, segmentLen_, kthIndex_, | ||
| 226 | + static_cast<uint64_t>(segStart), (__ubuf__ T*)inputValues_.GetPhyAddr(), | ||
| 227 | + (__gm__ volatile T*)valueGm_.GetPhyAddr(), (__gm__ volatile int64_t*)indexGm_.GetPhyAddr()); | ||
| 228 | + } | ||
| 229 | + event_t done = static_cast<event_t>(pipe_->FetchEventID(HardEvent::V_S)); | ||
| 230 | + SetFlag<HardEvent::V_S>(done); | ||
| 231 | + WaitFlag<HardEvent::V_S>(done); | ||
| 232 | + return; | ||
| 233 | + } | ||
| 234 | + } | ||
| 184 | if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) { | 235 | if constexpr (EnableMedian && IS_MEDIAN_FLOAT_TYPE<T>) { |
| 185 | asc_vf_call<SimtCanonicalizeMedianSortValues<T>>(dim3(SmallAxisCommon::TWO_STAGE_THREAD_NUM), totalElems, | 236 | asc_vf_call<SimtCanonicalizeMedianSortValues<T>>(dim3(SmallAxisCommon::TWO_STAGE_THREAD_NUM), totalElems, |
| 186 | reinterpret_cast<__ubuf__ T*>(inputValues_.GetPhyAddr())); | 237 | reinterpret_cast<__ubuf__ T*>(inputValues_.GetPhyAddr())); |
| @@ -17,6 +17,18 @@ constexpr uint32_t MEDIAN_MODE_STATIC = 0U; | |||
| 17 | constexpr uint32_t MEDIAN_MODE_PROPAGATE_NAN = 1U; | 17 | constexpr uint32_t MEDIAN_MODE_PROPAGATE_NAN = 1U; |
| 18 | constexpr uint32_t MEDIAN_MODE_IGNORE_NAN = 2U; | 18 | constexpr uint32_t MEDIAN_MODE_IGNORE_NAN = 2U; |
| 19 | 19 | ||
| 20 | +// Median opts into the legacy fixed-rank shortcut before including the shared kernels. | ||
| 21 | +// This translation-unit selector is independent of runtime medianMode; normal Median uses | ||
| 22 | +// PROPAGATE_NAN, so it still counts NaNs. Implementation branches below use if constexpr. | ||
| 23 | +#ifndef KTH_VALUE_STATIC_MEDIAN_FAST_PATH | ||
| 24 | +#define KTH_VALUE_STATIC_MEDIAN_FAST_PATH 0 | ||
| 25 | +#endif | ||
| 26 | +constexpr bool KTH_VALUE_ENABLE_STATIC_MEDIAN_FAST_PATH = KTH_VALUE_STATIC_MEDIAN_FAST_PATH != 0; | ||
| 27 | + | ||
| 28 | +constexpr uint32_t KTH_INT16_MERGE_BATCH_ROWS = 8U; | ||
| 29 | +constexpr uint32_t KTH_INT16_MERGE_BATCH_MIN_AXIS = 128U; | ||
| 30 | +constexpr uint32_t KTH_INT16_MERGE_BATCH_MAX_AXIS = 512U; | ||
| 31 | + | ||
| 20 | struct KthValueTilingData { | 32 | struct KthValueTilingData { |
| 21 | uint32_t numTileDataSize; // h轴ub一次处理个数 | 33 | uint32_t numTileDataSize; // h轴ub一次处理个数 |
| 22 | uint32_t unsortedDimParallel; // b轴使用的核数 | 34 | uint32_t unsortedDimParallel; // b轴使用的核数 |
| @@ -39,6 +51,7 @@ struct KthValueTilingData { | |||
| 39 | // small_axis_insertion/two_stage: batchNum_ | 51 | // small_axis_insertion/two_stage: batchNum_ |
| 40 | // non_last_small_axis: inputValueAxisBytes_ (bf16 merge) | 52 | // non_last_small_axis: inputValueAxisBytes_ (bf16 merge) |
| 41 | // axis_one_copy: loopTimes_ | 53 | // axis_one_copy: loopTimes_ |
| 54 | + // merge_more_core: sort_sync_merge_sort 逻辑块大小，0 表示普通 more-core | ||
| 42 | uint32_t keyParams1; | 55 | uint32_t keyParams1; |
| 43 | // radix: 清零的一次 ub 数据量 | 56 | // radix: 清零的一次 ub 数据量 |
| 44 | // merge_sort_one_core: y2OutQue 的 ub 大小 | 57 | // merge_sort_one_core: y2OutQue 的 ub 大小 |
| @@ -72,4 +85,33 @@ struct KthValueTilingData { | |||
| 72 | uint32_t medianMode; | 85 | uint32_t medianMode; |
| 73 | }; | 86 | }; |
| 74 | 87 | ||
| 88 | +// Single-core routes transfer only the fields consumed by their kernels. | ||
| 89 | +struct KthValueMergeOneCoreTilingData { | ||
| 90 | + int64_t kthIndex; | ||
| 91 | + int64_t unsortedDimNum; | ||
| 92 | + uint32_t numTileDataSize; | ||
| 93 | + uint32_t unsortedDimParallel; | ||
| 94 | + uint32_t sortLoopTimes; | ||
| 95 | + uint32_t keyParams0; // rows per UB batch (> 0) | ||
| 96 | + uint32_t keyParams1; // input queue capacity in bytes | ||
| 97 | + uint32_t keyParams2; // index queue capacity in bytes | ||
| 98 | + uint32_t keyParams3; // Sort32-aligned axis length in elements | ||
| 99 | + uint32_t keyParams4; // queue buffer count: 0/1 = single, 2 = double | ||
| 100 | + uint32_t medianMode; | ||
| 101 | +}; | ||
| 102 | + | ||
| 103 | +struct KthValueRadixOneCoreTilingData { | ||
| 104 | + int64_t kthIndex; | ||
| 105 | + int64_t lastAxisNum; | ||
| 106 | + int64_t unsortedDimNum; | ||
| 107 | + uint32_t numTileDataSize; | ||
| 108 | + uint32_t unsortedDimParallel; | ||
| 109 | + uint32_t keyParams0; // input UB capacity in bytes | ||
| 110 | + uint32_t keyParams1; // index UB capacity in bytes | ||
| 111 | + uint32_t keyParams3; // queue buffer count: 2 = double, otherwise single | ||
| 112 | + uint32_t keyParams4; // output rows per loop: 0 defaults to 1 | ||
| 113 | + uint32_t tmpUbSize; | ||
| 114 | + uint32_t medianMode; | ||
| 115 | +}; | ||
| 116 | + | ||
| 75 | #endif | 117 | #endif |
| @@ -26,15 +26,21 @@ | |||
| 26 | #define KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX 10 | 26 | #define KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX 10 |
| 27 | #define KTH_VALUE_SCHID_RADIX_SELECT 11 | 27 | #define KTH_VALUE_SCHID_RADIX_SELECT 11 |
| 28 | #define KTH_VALUE_SCHID_SMALL_AXIS_SHORT_RANK_SELECT 12 | 28 | #define KTH_VALUE_SCHID_SMALL_AXIS_SHORT_RANK_SELECT 12 |
| 29 | +#define KTH_VALUE_SCHID_RESIDENT_HISTOGRAM 13 | ||
| 29 | 30 | ||
| 30 | -#define KTH_VALUE_TPL_KEY_DECL() \ | 31 | +// Compact-layout registration uses string conditions consumed by the compiler. |
| 31 | - ASCENDC_TPL_UINT_DECL(schId, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, KTH_VALUE_SCHID_MERGE_SORT, \ | 32 | +// With isInt32 == 1 in the upper byte, schId 0/8 encode 256/264 and schId 1 encodes 257. |
| 32 | - KTH_VALUE_SCHID_RADIX_ONE_CORE, KTH_VALUE_SCHID_RADIX_MORE_CORE, \ | 33 | +#define KTH_VALUE_MERGE_ONE_CORE_TILING_CONDITION "TILING_KEY_VAR == 256 || TILING_KEY_VAR == 264" |
| 33 | - KTH_VALUE_SCHID_MERGE_MORE_CORE, KTH_VALUE_SCHID_MERGE_INTRA_CORE, \ | 34 | +#define KTH_VALUE_RADIX_ONE_CORE_TILING_CONDITION "TILING_KEY_VAR == 257" |
| 34 | - KTH_VALUE_SCHID_SMALL_AXIS_INSERTION, KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE, \ | 35 | + |
| 35 | - KTH_VALUE_SCHID_AXIS_ONE_COPY, KTH_VALUE_SCHID_SORT32_SMALL_AXIS, \ | 36 | +#define KTH_VALUE_TPL_KEY_DECL() \ |
| 36 | - KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS, KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX, \ | 37 | + ASCENDC_TPL_UINT_DECL( \ |
| 37 | - KTH_VALUE_SCHID_RADIX_SELECT, KTH_VALUE_SCHID_SMALL_AXIS_SHORT_RANK_SELECT), \ | 38 | + schId, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, KTH_VALUE_SCHID_MERGE_SORT, KTH_VALUE_SCHID_RADIX_ONE_CORE, \ |
| 39 | + KTH_VALUE_SCHID_RADIX_MORE_CORE, KTH_VALUE_SCHID_MERGE_MORE_CORE, KTH_VALUE_SCHID_MERGE_INTRA_CORE, \ | ||
| 40 | + KTH_VALUE_SCHID_SMALL_AXIS_INSERTION, KTH_VALUE_SCHID_SMALL_AXIS_TWO_STAGE, KTH_VALUE_SCHID_AXIS_ONE_COPY, \ | ||
| 41 | + KTH_VALUE_SCHID_SORT32_SMALL_AXIS, KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS, \ | ||
| 42 | + KTH_VALUE_SCHID_NON_LAST_SMALL_AXIS_RADIX, KTH_VALUE_SCHID_RADIX_SELECT, \ | ||
| 43 | + KTH_VALUE_SCHID_SMALL_AXIS_SHORT_RANK_SELECT, KTH_VALUE_SCHID_RESIDENT_HISTOGRAM), \ | ||
| 38 | ASCENDC_TPL_UINT_DECL(isInt32, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, 0, 1) | 44 | ASCENDC_TPL_UINT_DECL(isInt32, ASCENDC_TPL_8_BW, ASCENDC_TPL_UI_LIST, 0, 1) |
| 39 | 45 | ||
| 40 | #define KTH_VALUE_TPL_MERGE_SORT_KEY_SEL() \ | 46 | #define KTH_VALUE_TPL_MERGE_SORT_KEY_SEL() \ |
| @@ -99,8 +105,12 @@ | |||
| 99 | ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_RADIX_SELECT_KEY_SEL()), \ | 105 | ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_RADIX_SELECT_KEY_SEL()), \ |
| 100 | ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_SMALL_AXIS_SHORT_RANK_SELECT_KEY_SEL()) | 106 | ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_SMALL_AXIS_SHORT_RANK_SELECT_KEY_SEL()) |
| 101 | 107 | ||
| 108 | +#define KTH_VALUE_TPL_RESIDENT_HISTOGRAM_KEY_SEL() \ | ||
| 109 | + ASCENDC_TPL_UINT_SEL(schId, ASCENDC_TPL_UI_LIST, KTH_VALUE_SCHID_RESIDENT_HISTOGRAM), \ | ||
| 110 | + ASCENDC_TPL_UINT_SEL(isInt32, ASCENDC_TPL_UI_LIST, 1) | ||
| 111 | + | ||
| 102 | ASCENDC_TPL_ARGS_DECL(KthValue, KTH_VALUE_TPL_KEY_DECL()); | 112 | ASCENDC_TPL_ARGS_DECL(KthValue, KTH_VALUE_TPL_KEY_DECL()); |
| 103 | 113 | ||
| 104 | -ASCENDC_TPL_SEL(KTH_VALUE_TPL_ALL_KEY_SEL()); | 114 | +ASCENDC_TPL_SEL(KTH_VALUE_TPL_ALL_KEY_SEL(), ASCENDC_TPL_ARGS_SEL(KTH_VALUE_TPL_RESIDENT_HISTOGRAM_KEY_SEL())); |
| 105 | 115 | ||
| 106 | #endif | 116 | #endif |
| @@ -93,6 +93,22 @@ def _fill_pattern_along_axis(tensor, values, axis): | |||
| 93 | 93 | ||
| 94 | def kth_value_input(x, *, dim=-1, testcase_name="", **kwargs): | 94 | def kth_value_input(x, *, dim=-1, testcase_name="", **kwargs): |
| 95 | name = str(testcase_name) | 95 | name = str(testcase_name) |
| 96 | + if name.endswith("_mixed_nan"): | ||
| 97 | + # Preserve NaN signs and payloads rather than relying on float conversions. | ||
| 98 | + if str(x.dtype) == "float32": | ||
| 99 | + bits = x.view("uint32").reshape(-1) | ||
| 100 | + encodings = (0xFFC00000, 0x7FC00000, 0xFFC00001, 0x7FC00002) | ||
| 101 | + else: | ||
| 102 | + bits = x.view("uint16").reshape(-1) | ||
| 103 | + nan_bits = 0xFFC0 if "bfloat16" in str(x.dtype) else 0xFE00 | ||
| 104 | + encodings = ( | ||
| 105 | + nan_bits, | ||
| 106 | + nan_bits & 0x7FFF, | ||
| 107 | + nan_bits | 1, | ||
| 108 | + (nan_bits | 2) & 0x7FFF, | ||
| 109 | + ) | ||
| 110 | + bits[: len(encodings)] = encodings | ||
| 111 | + return (x,) | ||
| 96 | patterns = { | 112 | patterns = { |
| 97 | "duplicate": [2.0, 1.0, 1.0, 3.0, 2.0, 1.0, 3.0, 1.0], | 113 | "duplicate": [2.0, 1.0, 1.0, 3.0, 2.0, 1.0, 3.0, 1.0], |
| 98 | "signed_zero": [-0.0, 0.0, 0.0, -0.0, 1.0, -1.0, 0.0, -0.0], | 114 | "signed_zero": [-0.0, 0.0, 0.0, -0.0, 1.0, -1.0, 0.0, -0.0], |
| @@ -24,3 +24,8 @@ kth_value_duplicate,,kth_value,"((2, 8),)","('float32',)","('ND',)","((2, 1), (2 | |||
| 24 | kth_value_signed_zero,,kth_value,"((2, 8),)","('float32',)","('ND',)","((2, 1), (2, 1))","('float32', 'int64')","('ND', 'ND')",,,,,"{'k': 4, 'dim': -1}","((-4, 4),)",,,(),(),True,stable signed-zero index,,0,,(),() | 24 | kth_value_signed_zero,,kth_value,"((2, 8),)","('float32',)","('ND',)","((2, 1), (2, 1))","('float32', 'int64')","('ND', 'ND')",,,,,"{'k': 4, 'dim': -1}","((-4, 4),)",,,(),(),True,stable signed-zero index,,0,,(),() |
| 25 | kth_value_nan,,kth_value,"((2, 8),)","('float32',)","('ND',)","((2, 1), (2, 1))","('float32', 'int64')","('ND', 'ND')",,,,,"{'k': 7, 'dim': -1}","((-4, 4),)",,,(),(),True,stable first NaN rank,,0,,(),() | 25 | kth_value_nan,,kth_value,"((2, 8),)","('float32',)","('ND',)","((2, 1), (2, 1))","('float32', 'int64')","('ND', 'ND')",,,,,"{'k': 7, 'dim': -1}","((-4, 4),)",,,(),(),True,stable first NaN rank,,0,,(),() |
| 26 | kth_value_all_nan,,kth_value,"((2, 8),)","('float32',)","('ND',)","((2, 1), (2, 1))","('float32', 'int64')","('ND', 'ND')",,,,,"{'k': 4, 'dim': -1}","((-4, 4),)",,,(),(),True,stable all-NaN index,,0,,(),() | 26 | kth_value_all_nan,,kth_value,"((2, 8),)","('float32',)","('ND',)","((2, 1), (2, 1))","('float32', 'int64')","('ND', 'ND')",,,,,"{'k': 4, 'dim': -1}","((-4, 4),)",,,(),(),True,stable all-NaN index,,0,,(),() |
| 27 | +kth_value_int16_compact_301_k301,,kth_value,"((128, 301),)","('int16',)","('ND',)","((128, 1), (128, 1))","('int16', 'int64')","('ND', 'ND')",,,,,"{'k': 301, 'dim': -1}","((-32768,32767),)","((0,0),(0,0))","(0,0)",(),(),True,Compact merge extraction: original indices at rank-block and axis-tail boundaries,,0,,(),() | ||
| 28 | +kth_value_int16_compact_320_k33,,kth_value,"((128, 320),)","('int16',)","('ND',)","((128, 1), (128, 1))","('int16', 'int64')","('ND', 'ND')",,,,,"{'k': 33, 'dim': -1}","((-32768,32767),)","((0,0),(0,0))","(0,0)",(),(),True,Compact merge extraction: original indices at rank-block and axis-tail boundaries,,0,,(),() | ||
| 29 | +kth_value_int16_merge_batch_tail,,kth_value,"((1031, 301),)","('int16',)","('ND',)","((1031, 1), (1031, 1))","('int16', 'int64')","('ND','ND')","((1031, 301),)",,"((1031, 1), (1031, 1))",,"{'k': 301, 'dim': -1}","((-32768, 32767),)","((0,0), (0,0))",0,(),(),True,Batched merge with partial final groups and row-relative stable indices,,,,(),() | ||
| 30 | +kth_value_radix_fp32_k1_mixed_nan,,kth_value,"((1, 1024),)","('float32',)","('ND',)","((1, 1), (1, 1))","('float32', 'int64')","('ND', 'ND')",,,,,"{'k': 1, 'dim': -1, 'keepdim': True}","((-10, 10),)",,,(),(),True,Radix select: signed NaNs and distinct payloads,,0,,(),() | ||
| 31 | +kth_value_radix_fp32_k1024_mixed_nan,,kth_value,"((1, 1024),)","('float32',)","('ND',)","((1, 1), (1, 1))","('float32', 'int64')","('ND', 'ND')",,,,,"{'k': 1024, 'dim': -1, 'keepdim': True}","((-10, 10),)",,,(),(),True,Radix select: signed NaNs and distinct payloads,,0,,(),() | ||
| @@ -55,16 +55,16 @@ gert::TilingContextPara MakeKthValueTilingContext(const gert::StorageShape& xSha | |||
| 55 | } | 55 | } |
| 56 | } // namespace | 56 | } // namespace |
| 57 | 57 | ||
| 58 | -TEST_F(KthValueTilingTest, test_kthvalue_merge_sort_fp32_2x1024) | 58 | +TEST_F(KthValueTilingTest, test_kthvalue_radix_select_fp32_2x1024) |
| 59 | { | 59 | { |
| 60 | auto tilingContextPara = MakeKthValueTilingContext({{2, 1024}, {2, 1024}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}}, | 60 | auto tilingContextPara = MakeKthValueTilingContext({{2, 1024}, {2, 1024}}, {{2, 1}, {2, 1}}, {{2, 1}, {2, 1}}, |
| 61 | ge::DT_FLOAT, 5); | 61 | ge::DT_FLOAT, 5); |
| 62 | 62 | ||
| 63 | TilingInfo tilingInfo; | 63 | TilingInfo tilingInfo; |
| 64 | ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo)); | 64 | ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo)); |
| 65 | - EXPECT_EQ(tilingInfo.tilingKey, 256); | 65 | + EXPECT_EQ(tilingInfo.tilingKey, 267); |
| 66 | ASSERT_EQ(tilingInfo.workspaceSizes.size(), 1); | 66 | ASSERT_EQ(tilingInfo.workspaceSizes.size(), 1); |
| 67 | - EXPECT_EQ(tilingInfo.workspaceSizes[0], WORK_SPACE_SIZE); | 67 | + EXPECT_EQ(tilingInfo.workspaceSizes[0], WORK_SPACE_SIZE + 4224UL); |
| 68 | } | 68 | } |
| 69 | 69 | ||
| 70 | TEST_F(KthValueTilingTest, test_kthvalue_axis_one_copy_fp32_100x1) | 70 | TEST_F(KthValueTilingTest, test_kthvalue_axis_one_copy_fp32_100x1) |
| @@ -91,16 +91,16 @@ TEST_F(KthValueTilingTest, test_kthvalue_merge_sort_bf16_8x512) | |||
| 91 | EXPECT_EQ(tilingInfo.workspaceSizes[0], WORK_SPACE_SIZE); | 91 | EXPECT_EQ(tilingInfo.workspaceSizes[0], WORK_SPACE_SIZE); |
| 92 | } | 92 | } |
| 93 | 93 | ||
| 94 | -TEST_F(KthValueTilingTest, test_kthvalue_merge_sort_fp32_3d_8x16x1024) | 94 | +TEST_F(KthValueTilingTest, test_kthvalue_radix_select_fp32_3d_8x16x1024) |
| 95 | { | 95 | { |
| 96 | auto tilingContextPara = MakeKthValueTilingContext({{8, 16, 1024}, {8, 16, 1024}}, {{8, 16, 1}, {8, 16, 1}}, | 96 | auto tilingContextPara = MakeKthValueTilingContext({{8, 16, 1024}, {8, 16, 1024}}, {{8, 16, 1}, {8, 16, 1}}, |
| 97 | {{8, 16, 1}, {8, 16, 1}}, ge::DT_FLOAT, 500, 2); | 97 | {{8, 16, 1}, {8, 16, 1}}, ge::DT_FLOAT, 500, 2); |
| 98 | 98 | ||
| 99 | TilingInfo tilingInfo; | 99 | TilingInfo tilingInfo; |
| 100 | ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo)); | 100 | ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo)); |
| 101 | - EXPECT_EQ(tilingInfo.tilingKey, 256); | 101 | + EXPECT_EQ(tilingInfo.tilingKey, 267); |
| 102 | ASSERT_EQ(tilingInfo.workspaceSizes.size(), 1); | 102 | ASSERT_EQ(tilingInfo.workspaceSizes.size(), 1); |
| 103 | - EXPECT_EQ(tilingInfo.workspaceSizes[0], WORK_SPACE_SIZE); | 103 | + EXPECT_EQ(tilingInfo.workspaceSizes[0], WORK_SPACE_SIZE + 135168UL); |
| 104 | } | 104 | } |
| 105 | 105 | ||
| 106 | TEST_F(KthValueTilingTest, test_kthvalue_axis_one_copy_int32_64x1) | 106 | TEST_F(KthValueTilingTest, test_kthvalue_axis_one_copy_int32_64x1) |
| @@ -132,7 +132,7 @@ TEST_F(KthValueTilingTest, test_kthvalue_axis_one_copy_allows_unsorted_dim_over_ | |||
| 132 | EXPECT_EQ(tilingInfo.workspaceSizes[0], WORK_SPACE_SIZE); | 132 | EXPECT_EQ(tilingInfo.workspaceSizes[0], WORK_SPACE_SIZE); |
| 133 | } | 133 | } |
| 134 | 134 | ||
| 135 | -TEST_F(KthValueTilingTest, test_kthvalue_merge_intra_core_allows_unsorted_dim_over_uint32) | 135 | +TEST_F(KthValueTilingTest, test_kthvalue_radix_select_allows_unsorted_dim_over_uint32) |
| 136 | { | 136 | { |
| 137 | int64_t largeBatch = static_cast<int64_t>(std::numeric_limits<uint32_t>::max()) + 1; | 137 | int64_t largeBatch = static_cast<int64_t>(std::numeric_limits<uint32_t>::max()) + 1; |
| 138 | auto tilingContextPara = MakeKthValueTilingContext({{largeBatch, 4097}, {largeBatch, 4097}}, | 138 | auto tilingContextPara = MakeKthValueTilingContext({{largeBatch, 4097}, {largeBatch, 4097}}, |
| @@ -141,7 +141,7 @@ TEST_F(KthValueTilingTest, test_kthvalue_merge_intra_core_allows_unsorted_dim_ov | |||
| 141 | 141 | ||
| 142 | TilingInfo tilingInfo; | 142 | TilingInfo tilingInfo; |
| 143 | ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo)); | 143 | ASSERT_TRUE(ExecuteTiling(tilingContextPara, tilingInfo)); |
| 144 | - EXPECT_EQ(tilingInfo.tilingKey, 260); | 144 | + EXPECT_EQ(tilingInfo.tilingKey, 267); |
| 145 | ASSERT_GE(tilingInfo.tilingDataSize, sizeof(KthValueTilingData)); | 145 | ASSERT_GE(tilingInfo.tilingDataSize, sizeof(KthValueTilingData)); |
| 146 | const auto* tilingData = reinterpret_cast<const KthValueTilingData*>(tilingInfo.tilingData.get()); | 146 | const auto* tilingData = reinterpret_cast<const KthValueTilingData*>(tilingInfo.tilingData.get()); |
| 147 | EXPECT_EQ(tilingData->unsortedDimNum, largeBatch); | 147 | EXPECT_EQ(tilingData->unsortedDimNum, largeBatch); |
| @@ -11,6 +11,7 @@ | |||
| 11 | #include "kernel_operator.h" | 11 | #include "kernel_operator.h" |
| 12 | #include "kernel_tiling/kernel_tiling.h" | 12 | #include "kernel_tiling/kernel_tiling.h" |
| 13 | 13 | ||
| 14 | +#define KTH_VALUE_STATIC_MEDIAN_FAST_PATH 1 | ||
| 14 | #include "../../kth_value/arch35/kth_value_kernel_dispatch.h" | 15 | #include "../../kth_value/arch35/kth_value_kernel_dispatch.h" |
| 15 | 16 | ||
| 16 | using namespace AscendC; | 17 | using namespace AscendC; |
| @@ -31,7 +32,17 @@ __global__ __aicore__ void median(GM_ADDR x, GM_ADDR y, GM_ADDR indices, GM_ADDR | |||
| 31 | KERNEL_TASK_TYPE(MEDIAN_MERGE_MORE_CORE_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | 32 | KERNEL_TASK_TYPE(MEDIAN_MERGE_MORE_CORE_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); |
| 32 | KERNEL_TASK_TYPE(MEDIAN_RADIX_SELECT_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | 33 | KERNEL_TASK_TYPE(MEDIAN_RADIX_SELECT_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); |
| 33 | REGISTER_TILING_DEFAULT(KthValueTilingData); | 34 | REGISTER_TILING_DEFAULT(KthValueTilingData); |
| 34 | - GET_TILING_DATA_WITH_STRUCT(KthValueTilingData, tilingData, tiling); | ||
| 35 | TPipe pipe; | 35 | TPipe pipe; |
| 36 | - KthValue::Dispatch<true, schId, isInt32>(x, y, indices, workspace, &tilingData, &pipe); | 36 | + if constexpr (schId == KTH_VALUE_SCHID_MERGE_SORT || schId == KTH_VALUE_SCHID_SORT32_SMALL_AXIS) { |
| 37 | + REGISTER_TILING_FOR_TILINGKEY(KTH_VALUE_MERGE_ONE_CORE_TILING_CONDITION, KthValueMergeOneCoreTilingData); | ||
| 38 | + GET_TILING_DATA_WITH_STRUCT(KthValueMergeOneCoreTilingData, tilingData, tiling); | ||
| 39 | + KthValue::Dispatch<true, schId, isInt32>(x, y, indices, workspace, &tilingData, &pipe); | ||
| 40 | + } else if constexpr (schId == KTH_VALUE_SCHID_RADIX_ONE_CORE) { | ||
| 41 | + REGISTER_TILING_FOR_TILINGKEY(KTH_VALUE_RADIX_ONE_CORE_TILING_CONDITION, KthValueRadixOneCoreTilingData); | ||
| 42 | + GET_TILING_DATA_WITH_STRUCT(KthValueRadixOneCoreTilingData, tilingData, tiling); | ||
| 43 | + KthValue::Dispatch<true, schId, isInt32>(x, y, indices, workspace, &tilingData, &pipe); | ||
| 44 | + } else { | ||
| 45 | + GET_TILING_DATA_WITH_STRUCT(KthValueTilingData, tilingData, tiling); | ||
| 46 | + KthValue::Dispatch<true, schId, isInt32>(x, y, indices, workspace, &tilingData, &pipe); | ||
| 47 | + } | ||
| 37 | } | 48 | } |
| @@ -51,7 +51,8 @@ TEST(MedianTilingTest, FloatMedianUsesNanPropagatingMode) | |||
| 51 | TilingInfo tilingInfo; | 51 | TilingInfo tilingInfo; |
| 52 | ASSERT_TRUE(ExecuteTiling(context, tilingInfo)); | 52 | ASSERT_TRUE(ExecuteTiling(context, tilingInfo)); |
| 53 | EXPECT_EQ(tilingInfo.tilingKey, 264); | 53 | EXPECT_EQ(tilingInfo.tilingKey, 264); |
| 54 | - const auto* tilingData = GetTilingData(tilingInfo); | 54 | + ASSERT_EQ(tilingInfo.tilingDataSize, sizeof(KthValueMergeOneCoreTilingData)); |
| 55 | + const auto* tilingData = reinterpret_cast<const KthValueMergeOneCoreTilingData*>(tilingInfo.tilingData.get()); | ||
| 55 | EXPECT_EQ(tilingData->medianMode, 1U); | 56 | EXPECT_EQ(tilingData->medianMode, 1U); |
| 56 | EXPECT_EQ(tilingData->kthIndex, 2); | 57 | EXPECT_EQ(tilingData->kthIndex, 2); |
| 57 | ASSERT_EQ(tilingInfo.workspaceSizes.size(), 1); | 58 | ASSERT_EQ(tilingInfo.workspaceSizes.size(), 1); |
| @@ -31,7 +31,17 @@ __global__ __aicore__ void nan_median(GM_ADDR x, GM_ADDR y, GM_ADDR indices, GM_ | |||
| 31 | KERNEL_TASK_TYPE(NAN_MEDIAN_MERGE_MORE_CORE_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | 31 | KERNEL_TASK_TYPE(NAN_MEDIAN_MERGE_MORE_CORE_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); |
| 32 | KERNEL_TASK_TYPE(NAN_MEDIAN_RADIX_SELECT_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); | 32 | KERNEL_TASK_TYPE(NAN_MEDIAN_RADIX_SELECT_TILING_KEY, KERNEL_TYPE_MIX_AIV_1_0); |
| 33 | REGISTER_TILING_DEFAULT(KthValueTilingData); | 33 | REGISTER_TILING_DEFAULT(KthValueTilingData); |
| 34 | - GET_TILING_DATA_WITH_STRUCT(KthValueTilingData, tilingData, tiling); | ||
| 35 | TPipe pipe; | 34 | TPipe pipe; |
| 36 | - KthValue::Dispatch<true, schId, isInt32>(x, y, indices, workspace, &tilingData, &pipe); | 35 | + if constexpr (schId == KTH_VALUE_SCHID_MERGE_SORT || schId == KTH_VALUE_SCHID_SORT32_SMALL_AXIS) { |
| 36 | + REGISTER_TILING_FOR_TILINGKEY(KTH_VALUE_MERGE_ONE_CORE_TILING_CONDITION, KthValueMergeOneCoreTilingData); | ||
| 37 | + GET_TILING_DATA_WITH_STRUCT(KthValueMergeOneCoreTilingData, tilingData, tiling); | ||
| 38 | + KthValue::Dispatch<true, schId, isInt32>(x, y, indices, workspace, &tilingData, &pipe); | ||
| 39 | + } else if constexpr (schId == KTH_VALUE_SCHID_RADIX_ONE_CORE) { | ||
| 40 | + REGISTER_TILING_FOR_TILINGKEY(KTH_VALUE_RADIX_ONE_CORE_TILING_CONDITION, KthValueRadixOneCoreTilingData); | ||
| 41 | + GET_TILING_DATA_WITH_STRUCT(KthValueRadixOneCoreTilingData, tilingData, tiling); | ||
| 42 | + KthValue::Dispatch<true, schId, isInt32>(x, y, indices, workspace, &tilingData, &pipe); | ||
| 43 | + } else { | ||
| 44 | + GET_TILING_DATA_WITH_STRUCT(KthValueTilingData, tilingData, tiling); | ||
| 45 | + KthValue::Dispatch<true, schId, isInt32>(x, y, indices, workspace, &tilingData, &pipe); | ||
| 46 | + } | ||
| 37 | } | 47 | } |
| @@ -48,7 +48,8 @@ TEST(NanMedianTilingTest, UsesIgnoreNanMode) | |||
| 48 | TilingInfo tilingInfo; | 48 | TilingInfo tilingInfo; |
| 49 | ASSERT_TRUE(ExecuteTiling(context, tilingInfo)); | 49 | ASSERT_TRUE(ExecuteTiling(context, tilingInfo)); |
| 50 | EXPECT_EQ(tilingInfo.tilingKey, 264); | 50 | EXPECT_EQ(tilingInfo.tilingKey, 264); |
| 51 | - const auto* tilingData = GetTilingData(tilingInfo); | 51 | + ASSERT_EQ(tilingInfo.tilingDataSize, sizeof(KthValueMergeOneCoreTilingData)); |
| 52 | + const auto* tilingData = reinterpret_cast<const KthValueMergeOneCoreTilingData*>(tilingInfo.tilingData.get()); | ||
| 52 | EXPECT_EQ(tilingData->medianMode, 2U); | 53 | EXPECT_EQ(tilingData->medianMode, 2U); |
| 53 | EXPECT_EQ(tilingData->kthIndex, 3); | 54 | EXPECT_EQ(tilingData->kthIndex, 3); |
| 54 | } | 55 | } |
[P1] 扩大浮点 radix-select 路由前，需要为普通 KthValue 统一 NaN 排序键
新增 IsKthFloatSelect 将末轴长度 1024–4096 的 FP32 KthValue 优先送入 radix-select，但普通 KthValue 实例化 EnableMedian=false。kth_value_radix_select.h 的 MakeSortKeys 调用 TwiddleInFp32<T, UT, 0>，该 helper 不归一化 NaN；后续 MapNanKeysToMax 又仅在 EnableMedian=true 时执行。负 NaN 因此会被当作最小键。
具体反例：FP32 [1,1024]，k=1、dim=-1，输入 [-NaN,1,2,...,1023]，其中负 NaN 位模式为 0xffc00000。PR 的位变换将其映射为 0x003fffff，小于所有正数键，选择结果为 NaN/index 0；仓内 PyTorch golden 期望为 1/index 1。已实际运行 CPU torch.kthvalue 与复刻该位变换的最小模型，得到上述差异。新增 FP16/BF16 radix-select 范围也有同类问题。
建议浮点 KthValue 也将所有 NaN 映射到最大键，保留原始输入用于输出取值，并补充负 NaN、混合符号/不同 payload NaN、k=1/k=N 的回归用例。
验证范围：当前提交 0a31a902ccdbf2a9dce48681ea2c2bd311d00012 的静态分析及 CPU 键语义模型，未做 NPU 实跑。