已合并
提升topk int64索引的部分shape的性能 #2505
JimmyLam2创建于 4月27日
提升topk int64索引的部分shape的性能 #2505
已合并
共 1 个文件变更+11-1
| @@ -38,6 +38,8 @@ const int64_t MIN_AICORE_CALC_REG_BASE_INT64_INPUTSIZE = 65000; | |||
| 38 | const int64_t MAX_K = 16; | 38 | const int64_t MAX_K = 16; |
| 39 | 39 | ||
| 40 | constexpr int64_t TWO_THOUSAND = 2000; | 40 | constexpr int64_t TWO_THOUSAND = 2000; |
| 41 | +// 排序轴大于该阈值时，走SortAndTopK模板，SortWithIndex场景不涉及 | ||
| 42 | +constexpr int64_t SORT_AND_TOP_K_THRESHOLD = 10000000; | ||
| 41 | 43 | ||
| 42 | static const std::initializer_list<op::DataType> ANCIENT_DTYPE_SUPPORT_LIST = { | 44 | static const std::initializer_list<op::DataType> ANCIENT_DTYPE_SUPPORT_LIST = { |
| 43 | op::DataType::DT_FLOAT16, op::DataType::DT_FLOAT}; | 45 | op::DataType::DT_FLOAT16, op::DataType::DT_FLOAT}; |
| @@ -150,6 +152,14 @@ std::tuple<aclTensor*, aclTensor*> TopkV3( | |||
| 150 | return std::tuple<aclTensor*, aclTensor*>(values, indices); | 152 | return std::tuple<aclTensor*, aclTensor*>(values, indices); |
| 151 | } | 153 | } |
| 152 | 154 | ||
| 155 | +// SortWithIndex排序时，将输出索引类型int64转为int32，计算结束后再将结果转为int64，能获取更好的性能 | ||
| 156 | +static bool IsSortWithIndex(const aclTensor* self, int64_t k, bool sorted) { | ||
| 157 | + auto inputShape = self->GetViewShape(); | ||
| 158 | + int64_t dimNum = static_cast<int64_t>(inputShape.GetDimNum()); | ||
| 159 | + | ||
| 160 | + return (k > TWO_THOUSAND) && (inputShape.GetDim(dimNum - 1) <= SORT_AND_TOP_K_THRESHOLD) && sorted; | ||
| 161 | +} | ||
| 162 | + | ||
| 153 | // AICPU算子kernel | 163 | // AICPU算子kernel |
| 154 | std::tuple<aclTensor*, aclTensor*> TopkAiCpu( | 164 | std::tuple<aclTensor*, aclTensor*> TopkAiCpu( |
| 155 | const aclTensor* self, const aclTensor* k, int64_t dim, bool largest, bool sorted, aclTensor* values, | 165 | const aclTensor* self, const aclTensor* k, int64_t dim, bool largest, bool sorted, aclTensor* values, |
| @@ -181,7 +191,7 @@ std::tuple<aclTensor*, aclTensor*> Topk( | |||
| 181 | const aclTensor* kTensor = executor->ConvertToTensor(kScalar, op::ToOpDataType(ACL_INT32)); | 191 | const aclTensor* kTensor = executor->ConvertToTensor(kScalar, op::ToOpDataType(ACL_INT32)); |
| 182 | auto valuesOut = executor->AllocTensor(outShape, self->GetDataType(), self->GetStorageFormat()); | 192 | auto valuesOut = executor->AllocTensor(outShape, self->GetDataType(), self->GetStorageFormat()); |
| 183 | aclTensor* indicesOut = nullptr; | 193 | aclTensor* indicesOut = nullptr; |
| 184 | - if (IsRegBase()) { | 194 | + if (IsRegBase() && !IsSortWithIndex(self, k, sorted)) { |
| 185 | indicesOut = executor->AllocTensor(outShape, indicesDType, self->GetStorageFormat()); | 195 | indicesOut = executor->AllocTensor(outShape, indicesDType, self->GetStorageFormat()); |
| 186 | } else { | 196 | } else { |
| 187 | indicesOut = executor->AllocTensor(outShape, op::DataType::DT_INT32, self->GetStorageFormat()); | 197 | indicesOut = executor->AllocTensor(outShape, op::DataType::DT_INT32, self->GetStorageFormat()); |