Star1413
2156
Code
Code
Issues923
Pull Requests749
Actions
Discussion
Wiki
Analysis
Star1413
2156
Merged
feat: aclnnTopk 支持 sortLen 超 int32 上限的 RadixTopK int64 索引 #10291
Nice tryAdd on Sep 11
feat: aclnnTopk 支持 sortLen 超 int32 上限的 RadixTopK int64 索引 #10291
Merged
Nice tryAdd on Sep 11
Merged fromNice_try/ops-nn-workRadixTopKInt64intocann/ops-nnmaster
Total: 1 files changed+58-32
Mindex/gather_elements/op_api/aclnn_topk.cpp+58-32
@@ -136,8 +136,8 @@ static bool CheckParamValid(const aclTensor* self, int64_t k, int64_t dim)
136 return false;136 return false;
137 } else if (tmpDim > 0 && (dim < -tmpDim || dim >= tmpDim)) {137 } else if (tmpDim > 0 && (dim < -tmpDim || dim >= tmpDim)) {
138 OP_LOGE(ACLNN_ERR_PARAM_INVALID,138 OP_LOGE(ACLNN_ERR_PARAM_INVALID,
139- "Dimension out of range (expected to be in range of [-%ld, %ld], "139+ "Dimension out of range (expected to be in range of [-%ld, %ld],"
140- "but got %ld)",140+ " but got %ld)",
141 tmpDim, tmpDim - 1, dim);141 tmpDim, tmpDim - 1, dim);
142 return false;142 return false;
143 }143 }
@@ -147,8 +147,8 @@ static bool CheckParamValid(const aclTensor* self, int64_t k, int64_t dim)
147 int64_t tmpK = (tmpDim > 0) ? inputShape.GetDim(positiveDim) : 1;147 int64_t tmpK = (tmpDim > 0) ? inputShape.GetDim(positiveDim) : 1;
148 if (k < 0 || k > tmpK) {148 if (k < 0 || k > tmpK) {
149 OP_LOGE(ACLNN_ERR_PARAM_INVALID,149 OP_LOGE(ACLNN_ERR_PARAM_INVALID,
150- "Selected index k out of range (max num of self.size(%ld) is %ld, "150+ "Selected index k out of range (max num of self.size(%ld) is %ld,"
151- "but k is %ld)",151+ " but k is %ld)",
152 tmpDim, tmpK, k);152 tmpDim, tmpK, k);
153 return false;153 return false;
154 }154 }
@@ -578,7 +578,9 @@ static bool IsDataTypeDouble(op::DataType xDataType)
578// 2. 输入数据类型为 FP16 或 BF16（Radix TopK 仅支持这两种低精度类型）578// 2. 输入数据类型为 FP16 或 BF16（Radix TopK 仅支持这两种低精度类型）
579// 3. 最后一维长度 sortLen 与 k 的比例满足阈值要求（见下方 shapeCheck 分支）579// 3. 最后一维长度 sortLen 与 k 的比例满足阈值要求（见下方 shapeCheck 分支）
580// 4. k > 1000（Radix TopK 仅在大 k 场景下相比排序方案有收益）580// 4. k > 1000（Radix TopK 仅在大 k 场景下相比排序方案有收益）
581-static bool IsRadixTopKSupported(const aclTensor* self, int64_t k)581+// 例外：sortLen > int32 上限时走 RadixTopK int64 模板，此时 k 只要求 > 0，
582+// 与 ops-math math/topk/op_api/topk.cpp 的同名门控保持一致。
583+static bool IsRadixTopKSupported(const aclTensor* self, int64_t k, bool sorted, op::DataType indicesDType)
582{584{
583 SocVersion version = GetCurrentPlatformInfo().GetSocVersion();585 SocVersion version = GetCurrentPlatformInfo().GetSocVersion();
584 auto inputShape = self->GetViewShape();586 auto inputShape = self->GetViewShape();
@@ -588,7 +590,15 @@ static bool IsRadixTopKSupported(const aclTensor* self, int64_t k)
588 bool dtypeCheck = self->GetDataType() == op::DataType::DT_FLOAT16 || self->GetDataType() == op::DataType::DT_BF16;590 bool dtypeCheck = self->GetDataType() == op::DataType::DT_FLOAT16 || self->GetDataType() == op::DataType::DT_BF16;
589 bool shapeCheck = false;591 bool shapeCheck = false;
590 if (sortLen > MAX_INT32_INPUTSIZE) {592 if (sortLen > MAX_INT32_INPUTSIZE) {
591- shapeCheck = false;593+ // 超过 int32 上限走 RadixTopK int64 模板（TopKV2 在该量级会 59.5GB workspace OOM，
594+ // 比例门控失去意义），要求 indices 输出为 int64，k 比例不设限
595+ // 注：indicesDType 在 aclnn 入口已被 CheckDtypeValid 强制为 int64，此处判断恒为真；
596+ // 保留它是为了在将来放宽 indices 输出类型时，该分支仍能拦住非 int64 组合（int64 模板
597+ // 无法把 >2^31 的索引写进 int32），并与 ops-math 侧 l0op::Topk 的入参语义对齐。
598+ // k 门控同步放宽为 > 0：该量级下替代路径是 TopKV2 的 workspace OOM，而不是更慢的方案，
599+ // 因此"仅大 k 才有收益"的收益门控在此不适用（与 ops-math 的改动一致）。
600+ // 且此种情况，无法进行sort后gather_elements_v2，故不支持sorted为true的情况。
601+ return socCheck && dtypeCheck && (indicesDType == op::DataType::DT_INT64) && k > 0 && !sorted;
592 } else if (sortLen >= RADIX_TOP_K_S_THRESHOLD_2) {602 } else if (sortLen >= RADIX_TOP_K_S_THRESHOLD_2) {
593 shapeCheck = sortLen > RADIX_TOP_K_S_K_RATIO_2 * k;603 shapeCheck = sortLen > RADIX_TOP_K_S_K_RATIO_2 * k;
594 } else if (sortLen >= RADIX_TOP_K_S_THRESHOLD_1) {604 } else if (sortLen >= RADIX_TOP_K_S_THRESHOLD_1) {
@@ -731,7 +741,8 @@ aclnnStatus aclnnTopkGetWorkspaceSizeCommon(const aclTensor* self, int64_t k, in
731 auto selfReshape = TopkAdaptInputZeroDimTensor(selfContiguous, dimNum, uniqueExecutor.get());741 auto selfReshape = TopkAdaptInputZeroDimTensor(selfContiguous, dimNum, uniqueExecutor.get());
732 CHECK_RET(selfReshape != nullptr, ACLNN_ERR_INNER_NULLPTR);742 CHECK_RET(selfReshape != nullptr, ACLNN_ERR_INNER_NULLPTR);
733 743 
734- const aclTensor* indicesCastInt32 = nullptr;744+ // RadixTopK int64 大 shape 路径下这里可能持有 int64 索引，故用中性名字
745+ const aclTensor* indicesTopkOut = nullptr;
735 const aclTensor* valuesTopkOut = nullptr;746 const aclTensor* valuesTopkOut = nullptr;
736 747 
737 // 在910上,当输入fp32时,路径3的ge侧会插入cast,转换到fp16.此处修改与路径3保持一致。748 // 在910上,当输入fp32时,路径3的ge侧会插入cast,转换到fp16.此处修改与路径3保持一致。
@@ -774,7 +785,7 @@ aclnnStatus aclnnTopkGetWorkspaceSizeCommon(const aclTensor* self, int64_t k, in
774 topkOut = TopkWithPolicyCheck(selfCast, k, positiveDim, largest, sorted, indicesDType, sortPolicy,785 topkOut = TopkWithPolicyCheck(selfCast, k, positiveDim, largest, sorted, indicesDType, sortPolicy,
775 uniqueExecutor.get());786 uniqueExecutor.get());
776 valuesTopkOut = std::get<0>(topkOut);787 valuesTopkOut = std::get<0>(topkOut);
777- indicesCastInt32 = std::get<1>(topkOut);788+ indicesTopkOut = std::get<1>(topkOut);
778 } else if (positiveDim != lastDim) {789 } else if (positiveDim != lastDim) {
779 aclIntArray* axes = GetDimTransposeArray(dimNum, lastDim, positiveDim, uniqueExecutor.get());790 aclIntArray* axes = GetDimTransposeArray(dimNum, lastDim, positiveDim, uniqueExecutor.get());
780 CHECK_RET(axes != nullptr, ACLNN_ERR_INNER_NULLPTR);791 CHECK_RET(axes != nullptr, ACLNN_ERR_INNER_NULLPTR);
@@ -796,18 +807,26 @@ aclnnStatus aclnnTopkGetWorkspaceSizeCommon(const aclTensor* self, int64_t k, in
796 OP_LOGD("sort and topk, positiveDim not equal lastDim.");807 OP_LOGD("sort and topk, positiveDim not equal lastDim.");
797 topkOut = SortAndTopK(selfTranspose, k, lastDim, sortDimValue, largest, indicesDType, uniqueExecutor.get());808 topkOut = SortAndTopK(selfTranspose, k, lastDim, sortDimValue, largest, indicesDType, uniqueExecutor.get());
798 isHasCasted = true;809 isHasCasted = true;
799- } else if (IsRadixTopKSupported(selfTranspose, k)) {810+ } else if (IsRadixTopKSupported(selfTranspose, k, sorted, indicesDType)) {
800 OP_LOGD("radix topk supported, positiveDim not equal lastDim.");811 OP_LOGD("radix topk supported, positiveDim not equal lastDim.");
801 auto topkOutFirst = l0op::Topk(selfTranspose, k, lastDim, largest, sorted, indicesDType,812 auto topkOutFirst = l0op::Topk(selfTranspose, k, lastDim, largest, sorted, indicesDType,
802 uniqueExecutor.get());813 uniqueExecutor.get());
803 valuesTopkOut = std::get<0>(topkOutFirst);814 valuesTopkOut = std::get<0>(topkOutFirst);
804 aclTensor* indicesCast = std::get<1>(topkOutFirst);815 aclTensor* indicesCast = std::get<1>(topkOutFirst);
805- 816+ if (indicesCast->GetDataType() == op::DataType::DT_INT64) {
yuning_chen
yuning_chenyuning_chen28 days ago

sorted 为 true 的时候，int64 这条大 shape 路径把后面的排序丢掉了。

  • 现在代码看来 indices 是 int64 时直接把 RadixTopK 的索引透传出去，原来的 Sort 加 GatherElementsV2 没有走，lastDim 那个分支也是一样
  • 可能存在有结果没按值排好的问题，注释里写了结果无序，但接口在 sorted 为 true 时应该把值和对应下标一起排好
  • 建议可以在 sorted 为 true 时补一条不依赖 int32 GatherElements 的排序，或者这条路径直接拒绝 sorted 为 true
likedislike
Nice try
Nice try
28 days ago Comments：
806- auto sortOut = l0op::Sort(valuesTopkOut, -1, largest, true, indicesDType, uniqueExecutor.get());817+ // 与 positiveDim == lastDim 分支保持一致：GatherElementsV2 只接受 int32 索引、
807- valuesTopkOut = std::get<0>(sortOut);818+ // self 也不支持 int64（见 index/gather_elements_v2/op_host/op_api/gather_elements_v2.cpp），
808- aclTensor* indicesSort = std::get<1>(sortOut);819+ // 且 A2/A3 上 l0op::Sort 的索引输出恒为 int32（math/sort/op_api/sort.cpp），
809- indicesCastInt32 = l0op::GatherElementsV2(indicesCast, indicesSort, lastDim, uniqueExecutor.get());820+ // IsRadixTopKSupported 已拦截sorted为true的情况
810- topkOut = std::tie(valuesTopkOut, indicesCastInt32);821+ OP_LOGD("radix topk int64 indices passthrough, positiveDim not equal lastDim.");
822+ indicesTopkOut = indicesCast;
823+ } else {
824+ auto sortOut = l0op::Sort(valuesTopkOut, -1, largest, true, indicesDType, uniqueExecutor.get());
825+ valuesTopkOut = std::get<0>(sortOut);
826+ aclTensor* indicesSort = std::get<1>(sortOut);
827+ indicesTopkOut = l0op::GatherElementsV2(indicesCast, indicesSort, lastDim, uniqueExecutor.get());
828+ }
829+ topkOut = std::tie(valuesTopkOut, indicesTopkOut);
811 } else {830 } else {
812 OP_LOGD("topk, positiveDim not equal lastDim.");831 OP_LOGD("topk, positiveDim not equal lastDim.");
813 topkOut = TopkWithPolicyCheck(selfTranspose, k, lastDim, largest, sorted, indicesDType, sortPolicy,832 topkOut = TopkWithPolicyCheck(selfTranspose, k, lastDim, largest, sorted, indicesDType, sortPolicy,
@@ -820,7 +839,7 @@ aclnnStatus aclnnTopkGetWorkspaceSizeCommon(const aclTensor* self, int64_t k, in
820 valuesTopkOut = l0op::Transpose(std::get<0>(topkOut), axes, uniqueExecutor.get());839 valuesTopkOut = l0op::Transpose(std::get<0>(topkOut), axes, uniqueExecutor.get());
821 840 
822 // 将结果indices_transpose进行transpose，转换成正确的shape841 // 将结果indices_transpose进行transpose，转换成正确的shape
823- indicesCastInt32 = l0op::Transpose(std::get<1>(topkOut), axes, uniqueExecutor.get());842+ indicesTopkOut = l0op::Transpose(std::get<1>(topkOut), axes, uniqueExecutor.get());
824 } else {843 } else {
825 if (k > 0 && k < MAX_AICORE_CALC_DIM && !Ops::NN::AclnnUtil::IsRegbase()) {844 if (k > 0 && k < MAX_AICORE_CALC_DIM && !Ops::NN::AclnnUtil::IsRegbase()) {
826 int64_t kFirst = std::min(PARALLEL_K, selfContiguous->GetViewShape().GetDim(positiveDim));845 int64_t kFirst = std::min(PARALLEL_K, selfContiguous->GetViewShape().GetDim(positiveDim));
@@ -835,7 +854,7 @@ aclnnStatus aclnnTopkGetWorkspaceSizeCommon(const aclTensor* self, int64_t k, in
835 valuesTopkOut = std::get<0>(topkOut);854 valuesTopkOut = std::get<0>(topkOut);
836 aclTensor* indicesCast = std::get<1>(topkOut);855 aclTensor* indicesCast = std::get<1>(topkOut);
837 856 
838- indicesCastInt32 = l0op::GatherElements(indicesCastFirst, positiveDim, indicesCast, uniqueExecutor.get());857+ indicesTopkOut = l0op::GatherElements(indicesCastFirst, positiveDim, indicesCast, uniqueExecutor.get());
839 } else {858 } else {
840 std::tuple<const aclTensor*, const aclTensor*> topkOut(nullptr, nullptr);859 std::tuple<const aclTensor*, const aclTensor*> topkOut(nullptr, nullptr);
841 if (IsSort(sortDimValue, k, selfCast) && !IsDataTypeDouble(xDType) &&860 if (IsSort(sortDimValue, k, selfCast) && !IsDataTypeDouble(xDType) &&
@@ -849,31 +868,37 @@ aclnnStatus aclnnTopkGetWorkspaceSizeCommon(const aclTensor* self, int64_t k, in
849 topkOut = SortAndTopK(selfCast, k, positiveDim, sortDimValue, largest, indicesDType,868 topkOut = SortAndTopK(selfCast, k, positiveDim, sortDimValue, largest, indicesDType,
850 uniqueExecutor.get());869 uniqueExecutor.get());
851 isHasCasted = true;870 isHasCasted = true;
852- } else if (IsRadixTopKSupported(selfCast, k)) {871+ } else if (IsRadixTopKSupported(selfCast, k, sorted, indicesDType)) {
853 OP_LOGD("radix topk supported, positiveDim equal lastDim.");872 OP_LOGD("radix topk supported, positiveDim equal lastDim.");
854 auto topkOutFirst = l0op::Topk(selfCast, k, positiveDim, largest, sorted, indicesDType,873 auto topkOutFirst = l0op::Topk(selfCast, k, positiveDim, largest, sorted, indicesDType,
855 uniqueExecutor.get());874 uniqueExecutor.get());
856 valuesTopkOut = std::get<0>(topkOutFirst);875 valuesTopkOut = std::get<0>(topkOutFirst);
857 aclTensor* indicesCast = std::get<1>(topkOutFirst);876 aclTensor* indicesCast = std::get<1>(topkOutFirst);
858- 877+ if (indicesCast->GetDataType() == op::DataType::DT_INT64) {
859- auto sortOut = l0op::Sort(valuesTopkOut, -1, largest, true, indicesDType, uniqueExecutor.get());878+ // int64 大 shape 路径直接透传 RadixTopK 的 int64 索引，Sort/GatherElementsV2 组合暂不接管
860- valuesTopkOut = std::get<0>(sortOut);879+ // IsRadixTopKSupported 已拦截sorted为true的情况
861- aclTensor* indicesSort = std::get<1>(sortOut);880+ indicesTopkOut = indicesCast;
862- indicesCastInt32 = l0op::GatherElementsV2(indicesCast, indicesSort, positiveDim, uniqueExecutor.get());881+ } else {
863- topkOut = std::tie(valuesTopkOut, indicesCastInt32);882+ auto sortOut = l0op::Sort(valuesTopkOut, -1, largest, true, indicesDType, uniqueExecutor.get());
883+ valuesTopkOut = std::get<0>(sortOut);
884+ aclTensor* indicesSort = std::get<1>(sortOut);
885+ indicesTopkOut = l0op::GatherElementsV2(indicesCast, indicesSort, positiveDim,
886+ uniqueExecutor.get());
887+ }
888+ topkOut = std::tie(valuesTopkOut, indicesTopkOut);
864 } else {889 } else {
865 OP_LOGD("topk, positiveDim equal lastDim.");890 OP_LOGD("topk, positiveDim equal lastDim.");
866 topkOut = TopkWithPolicyCheck(selfCast, k, positiveDim, largest, sorted, indicesDType, sortPolicy,891 topkOut = TopkWithPolicyCheck(selfCast, k, positiveDim, largest, sorted, indicesDType, sortPolicy,
867 uniqueExecutor.get());892 uniqueExecutor.get());
868 }893 }
869 valuesTopkOut = std::get<0>(topkOut);894 valuesTopkOut = std::get<0>(topkOut);
870- indicesCastInt32 = std::get<1>(topkOut);895+ indicesTopkOut = std::get<1>(topkOut);
871 }896 }
872 }897 }
873- CHECK_RET(valuesTopkOut != nullptr && indicesCastInt32 != nullptr, ACLNN_ERR_INNER_NULLPTR);898+ CHECK_RET(valuesTopkOut != nullptr && indicesTopkOut != nullptr, ACLNN_ERR_INNER_NULLPTR);
874 899 
875 CHECK_RET(CheckReduceOutShape(valuesOut, valuesTopkOut), ACLNN_ERR_PARAM_INVALID);900 CHECK_RET(CheckReduceOutShape(valuesOut, valuesTopkOut), ACLNN_ERR_PARAM_INVALID);
876- CHECK_RET(CheckReduceOutShape(indicesOut, indicesCastInt32), ACLNN_ERR_PARAM_INVALID);901+ CHECK_RET(CheckReduceOutShape(indicesOut, indicesTopkOut), ACLNN_ERR_PARAM_INVALID);
877 902 
878 // 在910上，输入fp32转换到fp16计算完后需要重新转换到fp32。903 // 在910上，输入fp32转换到fp16计算完后需要重新转换到fp32。
879 auto valuesCast = TopkAdaptGeCastTensor(selfReshape, valuesTopkOut, k, op::DataType::DT_FLOAT,904 auto valuesCast = TopkAdaptGeCastTensor(selfReshape, valuesTopkOut, k, op::DataType::DT_FLOAT,
@@ -884,15 +909,16 @@ aclnnStatus aclnnTopkGetWorkspaceSizeCommon(const aclTensor* self, int64_t k, in
884 auto viewCopyValuesResult = l0op::ViewCopy(valuesCast, valuesOut, uniqueExecutor.get());909 auto viewCopyValuesResult = l0op::ViewCopy(valuesCast, valuesOut, uniqueExecutor.get());
885 CHECK_RET(viewCopyValuesResult != nullptr, ACLNN_ERR_INNER_NULLPTR);910 CHECK_RET(viewCopyValuesResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
886 911 
887- if (!indicesOutNeedsCast(k, sorted, isHasCasted)) {912+ // int64 索引（大 shape 透传路径）本身已与 indicesOut 同类型，无需再做 int64→int64 的 Cast
888- auto viewCopyIndicesResult = l0op::ViewCopy(indicesCastInt32, indicesOut, uniqueExecutor.get());913+ if (!indicesOutNeedsCast(k, sorted, isHasCasted) || indicesTopkOut->GetDataType() == op::DataType::DT_INT64) {
914+ auto viewCopyIndicesResult = l0op::ViewCopy(indicesTopkOut, indicesOut, uniqueExecutor.get());
889 CHECK_RET(viewCopyIndicesResult != nullptr, ACLNN_ERR_INNER_NULLPTR);915 CHECK_RET(viewCopyIndicesResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
890 } else {916 } else {
891- // 将结果indices_cast_int32进行cast，转换成int64类型917+ // 将结果索引 cast 成 int64 类型
892- auto indicesCastInt64 = l0op::Cast(indicesCastInt32, op::DataType::DT_INT64, uniqueExecutor.get());918+ auto indicesCastInt64 = l0op::Cast(indicesTopkOut, op::DataType::DT_INT64, uniqueExecutor.get());
893 CHECK_RET(indicesCastInt64 != nullptr, ACLNN_ERR_INNER_NULLPTR);919 CHECK_RET(indicesCastInt64 != nullptr, ACLNN_ERR_INNER_NULLPTR);
894 920 
895- // 将indices_cast_int64结果拷贝到values上921+ // 将 indices_cast_int64 结果拷贝到 indicesOut 上
896 auto viewCopyIndicesResult = l0op::ViewCopy(indicesCastInt64, indicesOut, uniqueExecutor.get());922 auto viewCopyIndicesResult = l0op::ViewCopy(indicesCastInt64, indicesOut, uniqueExecutor.get());
897 CHECK_RET(viewCopyIndicesResult != nullptr, ACLNN_ERR_INNER_NULLPTR);923 CHECK_RET(viewCopyIndicesResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
898 }924 }