| @@ -751,8 +751,7 @@ void NpuIndexIVFPQ::trainPQCodeBook_( | |||
| 751 | std::vector<float> sampled; | 751 | std::vector<float> sampled; |
| 752 | const float* trainData = x; | 752 | const float* trainData = x; |
| 753 | if (trainCount < n) { | 753 | if (trainCount < n) { |
| 754 | - sampled = sampleTrainData( | 754 | + sampled = sampleTrainData(x, n, this->d, trainCount, pq.cp.seed); |
| 755 | - x, n, this->d, trainCount, ivfpqConfig_.cp.seed); | ||
| 756 | trainData = sampled.data(); | 755 | trainData = sampled.data(); |
| 757 | } | 756 | } |
| 758 | 757 | ||
| @@ -771,7 +770,7 @@ void NpuIndexIVFPQ::trainPQCodeBook_( | |||
| 771 | } | 770 | } |
| 772 | 771 | ||
| 773 | std::vector<std::vector<float>> trainedCentroids(pq.M); | 772 | std::vector<std::vector<float>> trainedCentroids(pq.M); |
| 774 | - const int pqNiter = ivfpqConfig_.cp.niter; | 773 | + const int pqNiter = pq.cp.niter; |
| 775 | const size_t workerCount = | 774 | const size_t workerCount = |
| 776 | std::min(devices.size(), static_cast<size_t>(pq.M)); | 775 | std::min(devices.size(), static_cast<size_t>(pq.M)); |
| 777 | std::vector<std::future<void>> workers; | 776 | std::vector<std::future<void>> workers; |
| @@ -783,6 +782,14 @@ void NpuIndexIVFPQ::trainPQCodeBook_( | |||
| 783 | try { | 782 | try { |
| 784 | for (size_t m = worker; m < pq.M; m += workerCount) { | 783 | for (size_t m = worker; m < pq.M; m += workerCount) { |
| 785 | try { | 784 | try { |
| 785 | + if (verbose) { | ||
| 786 | + printf(" NpuIndexIVFPQ::train: [pq %zu/%zu] " | ||
| 787 | + "begin device=%d n=%ld\n", | ||
| 788 | + m + 1, | ||
| 789 | + static_cast<size_t>(pq.M), | ||
| 790 | + devices[worker], | ||
| 791 | + (long)trainCount); | ||
| 792 | + } | ||
| 786 | indexTrainImpl_( | 793 | indexTrainImpl_( |
| 787 | trainCount, | 794 | trainCount, |
| 788 | subspaceData[m].data(), | 795 | subspaceData[m].data(), |
| @@ -791,8 +798,20 @@ void NpuIndexIVFPQ::trainPQCodeBook_( | |||
| 791 | devices[worker], | 798 | devices[worker], |
| 792 | trainingResources[worker].get(), | 799 | trainingResources[worker].get(), |
| 793 | pqNiter, | 800 | pqNiter, |
Y | |||
| 794 | - ivfpqConfig_.cp.seed, | 801 | + pq.cp.nredo, |
| 802 | + pq.cp.seed, | ||
| 803 | + pq.cp.spherical, | ||
| 804 | + // PQ codebook K-means uses L2, matching | ||
| 805 | + // ProductQuantizer::train's default | ||
| 806 | + // IndexFlatL2. | ||
| 807 | + faiss::METRIC_L2, | ||
L 严重程度: 提示
问题: PQ 子量化器训练
![]() ![]() | |||
| 795 | trainedCentroids[m]); | 808 | trainedCentroids[m]); |
| 809 | + if (verbose) { | ||
| 810 | + printf(" NpuIndexIVFPQ::train: [pq %zu/%zu] " | ||
| 811 | + "done\n", | ||
| 812 | + m + 1, | ||
| 813 | + static_cast<size_t>(pq.M)); | ||
| 814 | + } | ||
| 796 | } catch (std::exception& e) { | 815 | } catch (std::exception& e) { |
| 797 | FAISS_THROW_FMT( | 816 | FAISS_THROW_FMT( |
| 798 | "NPU training failed for sub-quantizer %zu: %s", | 817 | "NPU training failed for sub-quantizer %zu: %s", |
| @@ -831,8 +850,7 @@ void NpuIndexIVFPQ::trainSubQuantizer_(int m, idx_t n, const float* x) { | |||
| 831 | "Insufficient data for sub-quantizer clustering"); | 850 | "Insufficient data for sub-quantizer clustering"); |
| 832 | 851 | ||
| 833 | // CPU path: faiss::Clustering with configurable parameters | 852 | // CPU path: faiss::Clustering with configurable parameters |
| 834 | - faiss::ClusteringParameters pqCp = ivfpqConfig_.cp; | 853 | + faiss::ClusteringParameters pqCp = pq.cp; |
| 835 | - pqCp.niter = ivfpqConfig_.cp.niter; | ||
| 836 | faiss::Clustering clus((int)pq.dsub, (int)pq.ksub, pqCp); | 854 | faiss::Clustering clus((int)pq.dsub, (int)pq.ksub, pqCp); |
| 837 | faiss::IndexFlatL2 index((int)pq.dsub); | 855 | faiss::IndexFlatL2 index((int)pq.dsub); |
| 838 | clus.train(n_data, subspace_data.data(), index); | 856 | clus.train(n_data, subspace_data.data(), index); |
| @@ -873,7 +891,10 @@ void NpuIndexIVFPQ::indexTrainImpl_( | |||
| 873 | int device, | 891 | int device, |
| 874 | NpuResources* resources, | 892 | NpuResources* resources, |
| 875 | int niter, | 893 | int niter, |
| 894 | + int nredo, | ||
| 876 | int64_t seed, | 895 | int64_t seed, |
| 896 | + bool spherical, | ||
| 897 | + int64_t assignmentMetric, | ||
| 877 | std::vector<float>& centroidsOut) { | 898 | std::vector<float>& centroidsOut) { |
| 878 | if (n <= 0 || x == nullptr) { | 899 | if (n <= 0 || x == nullptr) { |
| 879 | FAISS_THROW_MSG("train data invalid"); | 900 | FAISS_THROW_MSG("train data invalid"); |
| @@ -882,7 +903,7 @@ void NpuIndexIVFPQ::indexTrainImpl_( | |||
| 882 | FAISS_THROW_IF_NOT_MSG( | 903 | FAISS_THROW_IF_NOT_MSG( |
| 883 | niter > 0, "training iteration count must be positive"); | 904 | niter > 0, "training iteration count must be positive"); |
| 884 | if (verbose) { | 905 | if (verbose) { |
| 885 | - printf("AscendIndex::train: training %ld vectors of dims %d into %ld clusters\n", | 906 | + printf(" NpuIVFPQ::train: KMeans n=%ld dim=%d nlist=%ld\n", |
| 886 | (long)n, | 907 | (long)n, |
| 887 | dim, | 908 | dim, |
| 888 | (long)nlistTrain); | 909 | (long)nlistTrain); |
| @@ -897,8 +918,11 @@ void NpuIndexIVFPQ::indexTrainImpl_( | |||
| 897 | nlistTrain, | 918 | nlistTrain, |
| 898 | centroidsOut, | 919 | centroidsOut, |
| 899 | niter, | 920 | niter, |
| 921 | + nredo, | ||
| 900 | seed, | 922 | seed, |
| 901 | - ivfpqConfig_.cp.spherical); | 923 | + spherical, |
| 924 | + assignmentMetric, | ||
| 925 | + verbose); | ||
| 902 | } | 926 | } |
| 903 | 927 | ||
| 904 | std::vector<std::shared_ptr<NpuResources>> NpuIndexIVFPQ:: | 928 | std::vector<std::shared_ptr<NpuResources>> NpuIndexIVFPQ:: |
| @@ -1073,8 +1097,24 @@ void NpuIndexIVFPQ::train(idx_t n, const float* x) { | |||
| 1073 | ivfpqConfig_.cp.spherical = true; | 1097 | ivfpqConfig_.cp.spherical = true; |
| 1074 | } | 1098 | } |
| 1075 | 1099 | ||
| 1100 | + const bool npuTrain = ivfpqConfig_.useNpuTrain; | ||
| 1101 | + const double trainStart = verbose ? getmillisecs() : 0.0; | ||
| 1102 | + if (verbose) { | ||
| 1103 | + printf("NpuIndexIVFPQ::train: n=%ld d=%d nlist=%ld M=%d nbits=%d " | ||
| 1104 | + "metric=%s mode=%s niter=%d nredo=%d devices=%zu\n", | ||
| 1105 | + (long)n, | ||
| 1106 | + this->d, | ||
| 1107 | + (long)this->nlist, | ||
| 1108 | + subQuantizers_, | ||
| 1109 | + bitsPerCode_, | ||
| 1110 | + metric_type == faiss::METRIC_INNER_PRODUCT ? "IP" : "L2", | ||
| 1111 | + npuTrain ? "NPU" : "CPU", | ||
| 1112 | + ivfpqConfig_.cp.niter, | ||
| 1113 | + ivfpqConfig_.cp.nredo, | ||
| 1114 | + deviceIds_.size()); | ||
| 1115 | + } | ||
| 1076 | std::vector<std::shared_ptr<NpuResources>> trainingResources; | 1116 | std::vector<std::shared_ptr<NpuResources>> trainingResources; |
| 1077 | - if (ivfpqConfig_.useNpuTrain) { | 1117 | + if (npuTrain) { |
| 1078 | trainingResources = createTrainingResources_(deviceIds_); | 1118 | trainingResources = createTrainingResources_(deviceIds_); |
| 1079 | } | 1119 | } |
| 1080 | 1120 | ||
| @@ -1084,7 +1124,12 @@ void NpuIndexIVFPQ::train(idx_t n, const float* x) { | |||
| 1084 | idx_t pqSamplePrefixCount = 0; | 1124 | idx_t pqSamplePrefixCount = 0; |
| 1085 | 1125 | ||
| 1086 | // ---- Train coarse quantizer (IVF centroids) ---- | 1126 | // ---- Train coarse quantizer (IVF centroids) ---- |
| 1087 | - if (ivfpqConfig_.useNpuTrain) { | 1127 | + if (verbose) { |
| 1128 | + printf(" NpuIndexIVFPQ::train: [coarse] begin niter=%d nredo=%d\n", | ||
| 1129 | + ivfpqConfig_.cp.niter, | ||
| 1130 | + ivfpqConfig_.cp.nredo); | ||
| 1131 | + } | ||
| 1132 | + if (npuTrain) { | ||
| 1088 | const idx_t trainCount = computeTrainSize( | 1133 | const idx_t trainCount = computeTrainSize( |
| 1089 | this->nlist, n, ivfpqConfig_.cp.max_points_per_centroid); | 1134 | this->nlist, n, ivfpqConfig_.cp.max_points_per_centroid); |
| 1090 | std::vector<float> sampled = sampleTrainData( | 1135 | std::vector<float> sampled = sampleTrainData( |
| @@ -1117,8 +1162,11 @@ void NpuIndexIVFPQ::train(idx_t n, const float* x) { | |||
| 1117 | static_cast<int>(this->nlist), | 1162 | static_cast<int>(this->nlist), |
| 1118 | centroidsData_, | 1163 | centroidsData_, |
| 1119 | ivfpqConfig_.cp.niter, | 1164 | ivfpqConfig_.cp.niter, |
| 1165 | + ivfpqConfig_.cp.nredo, | ||
| 1120 | ivfpqConfig_.cp.seed, | 1166 | ivfpqConfig_.cp.seed, |
| 1121 | - ivfpqConfig_.cp.spherical); | 1167 | + ivfpqConfig_.cp.spherical, |
| 1168 | + metric_type, | ||
| 1169 | + verbose); | ||
| 1122 | } else { | 1170 | } else { |
| 1123 | indexTrainImpl_( | 1171 | indexTrainImpl_( |
| 1124 | trainCount, | 1172 | trainCount, |
| @@ -1128,7 +1176,10 @@ void NpuIndexIVFPQ::train(idx_t n, const float* x) { | |||
| 1128 | deviceIds_.front(), | 1176 | deviceIds_.front(), |
| 1129 | trainingResources.front().get(), | 1177 | trainingResources.front().get(), |
| 1130 | ivfpqConfig_.cp.niter, | 1178 | ivfpqConfig_.cp.niter, |
| 1179 | + ivfpqConfig_.cp.nredo, | ||
| 1131 | ivfpqConfig_.cp.seed, | 1180 | ivfpqConfig_.cp.seed, |
| 1181 | + ivfpqConfig_.cp.spherical, | ||
| 1182 | + metric_type, | ||
| 1132 | centroidsData_); | 1183 | centroidsData_); |
| 1133 | } | 1184 | } |
| 1134 | // stash centroidsOnHost. | 1185 | // stash centroidsOnHost. |
| @@ -1144,9 +1195,13 @@ void NpuIndexIVFPQ::train(idx_t n, const float* x) { | |||
| 1144 | this->cp = ivfpqConfig_.cp; | 1195 | this->cp = ivfpqConfig_.cp; |
| 1145 | trainQuantizer_(n, x); | 1196 | trainQuantizer_(n, x); |
| 1146 | } | 1197 | } |
| 1198 | + if (verbose) { | ||
| 1199 | + printf(" NpuIndexIVFPQ::train: [coarse] done elapsed=%.3f s\n", | ||
| 1200 | + (getmillisecs() - trainStart) / 1000.0); | ||
| 1201 | + } | ||
| 1147 | 1202 | ||
| 1148 | // Residual PQ training needs coarse assignments. | 1203 | // Residual PQ training needs coarse assignments. |
| 1149 | - if (byResidual_ && ivfpqConfig_.useNpuTrain) { | 1204 | + if (byResidual_ && npuTrain) { |
| 1150 | DeviceScope scope(config_.device); | 1205 | DeviceScope scope(config_.device); |
| 1151 | setIndex_( | 1206 | setIndex_( |
| 1152 | resources_.get(), | 1207 | resources_.get(), |
| @@ -1176,7 +1231,7 @@ void NpuIndexIVFPQ::train(idx_t n, const float* x) { | |||
| 1176 | pqTrainCount = pqSamplePrefixCount; | 1231 | pqTrainCount = pqSamplePrefixCount; |
| 1177 | } | 1232 | } |
| 1178 | if (byResidual_) { | 1233 | if (byResidual_) { |
| 1179 | - if (ivfpqConfig_.useNpuTrain && pqSamplePrefixCount == 0) { | 1234 | + if (npuTrain && pqSamplePrefixCount == 0) { |
| 1180 | const idx_t sampleCount = computeTrainSize( | 1235 | const idx_t sampleCount = computeTrainSize( |
| 1181 | static_cast<idx_t>(pq.ksub), | 1236 | static_cast<idx_t>(pq.ksub), |
| 1182 | n, | 1237 | n, |
| @@ -1195,7 +1250,21 @@ void NpuIndexIVFPQ::train(idx_t n, const float* x) { | |||
| 1195 | pqTrainCount, pqTrainData, residuals.data(), assign.data()); | 1250 | pqTrainCount, pqTrainData, residuals.data(), assign.data()); |
| 1196 | pqTrainData = residuals.data(); | 1251 | pqTrainData = residuals.data(); |
| 1197 | } | 1252 | } |
| 1253 | + if (verbose) { | ||
| 1254 | + printf(" NpuIndexIVFPQ::train: [pq] begin n=%ld M=%d dsub=%d " | ||
| 1255 | + "niter=%d nredo=%d source=%s\n", | ||
| 1256 | + (long)pqTrainCount, | ||
| 1257 | + subQuantizers_, | ||
| 1258 | + static_cast<int>(pq.dsub), | ||
| 1259 | + pq.cp.niter, | ||
| 1260 | + pq.cp.nredo, | ||
| 1261 | + byResidual_ ? "residual" : "raw"); | ||
| 1262 | + } | ||
| 1198 | trainPQCodeBook_(pqTrainCount, pqTrainData, deviceIds_, trainingResources); | 1263 | trainPQCodeBook_(pqTrainCount, pqTrainData, deviceIds_, trainingResources); |
| 1264 | + if (verbose) { | ||
| 1265 | + printf(" NpuIndexIVFPQ::train: [pq] done elapsed=%.3f s\n", | ||
| 1266 | + (getmillisecs() - trainStart) / 1000.0); | ||
| 1267 | + } | ||
| 1199 | 1268 | ||
| 1200 | // ---- Initialize the impl-layer index with trained PQ centroids ---- | 1269 | // ---- Initialize the impl-layer index with trained PQ centroids ---- |
| 1201 | { | 1270 | { |
| @@ -1239,6 +1308,10 @@ void NpuIndexIVFPQ::train(idx_t n, const float* x) { | |||
| 1239 | 1308 | ||
| 1240 | FAISS_ASSERT(index_); | 1309 | FAISS_ASSERT(index_); |
| 1241 | this->is_trained = true; | 1310 | this->is_trained = true; |
| 1311 | + if (verbose) { | ||
| 1312 | + printf("NpuIndexIVFPQ::train done: %.3f s\n", | ||
| 1313 | + (getmillisecs() - trainStart) / 1000.0); | ||
| 1314 | + } | ||
| 1242 | } | 1315 | } |
| 1243 | 1316 | ||
| 1244 | // Initialize the underlying IVFPQ impl index with the given parameters. | 1317 | // Initialize the underlying IVFPQ impl index with the given parameters. |
| @@ -1777,9 +1850,8 @@ void NpuIndexIVFPQ::addImplCore_( | |||
| 1777 | 1850 | ||
| 1778 | // Assign each of the n vectors to its best IVF centroid. | 1851 | // Assign each of the n vectors to its best IVF centroid. |
| 1779 | // NPU path (useNpuTrain && centroidsOnHost_ trained): reuses the K-means | 1852 | // NPU path (useNpuTrain && centroidsOnHost_ trained): reuses the K-means |
| 1780 | -// centroids with the fused exact coarse-assignment operator. The impl maps | 1853 | +// centroids with the fused exact coarse-assignment operator. CPU path uses |
| 1781 | -// inner product to its argmin formulation without materializing distances. | 1854 | +// quantizer->assign. |
| 1782 | -// CPU path: quantizer->assign. | ||
| 1783 | void NpuIndexIVFPQ::addL1_( | 1855 | void NpuIndexIVFPQ::addL1_( |
| 1784 | idx_t n, | 1856 | idx_t n, |
| 1785 | const float* x, | 1857 | const float* x, |
| @@ -261,7 +261,10 @@ class NpuIndexIVFPQ : public NpuIndexIVF { | |||
| 261 | int device, | 261 | int device, |
| 262 | NpuResources* resources, | 262 | NpuResources* resources, |
| 263 | int niter, | 263 | int niter, |
| 264 | + int nredo, | ||
| 264 | int64_t seed, | 265 | int64_t seed, |
| 266 | + bool spherical, | ||
| 267 | + int64_t assignmentMetric, | ||
| 265 | std::vector<float>& centroidsOut); | 268 | std::vector<float>& centroidsOut); |
| 266 | 269 | ||
| 267 | /// Train each sub-quantizer independently on the configured encoding | 270 | /// Train each sub-quantizer independently on the configured encoding |
| @@ -13,6 +13,8 @@ | |||
| 13 | 13 | ||
| 14 | 14 | ||
| 15 | 15 | ||
| 16 | + | ||
| 17 | + | ||
| 16 | 18 | ||
| 17 | 19 | ||
| 18 | 20 | ||
| @@ -44,8 +46,12 @@ constexpr int IVF_PQ_BURST_LEN = 64; | |||
| 44 | constexpr int IVF_PQ_TRAIN_MAX_BATCH = 1024; | 46 | constexpr int IVF_PQ_TRAIN_MAX_BATCH = 1024; |
| 45 | const int IVF_PQ_CORE_NUM = faiss::npu::NpuSocInfo::getInstance().getCoreNum(); | 47 | const int IVF_PQ_CORE_NUM = faiss::npu::NpuSocInfo::getInstance().getCoreNum(); |
| 46 | constexpr int IVF_PQ_SIZE_ALIGN = 8; | 48 | constexpr int IVF_PQ_SIZE_ALIGN = 8; |
| 49 | +// Match faiss::Clustering's redo seed schedule so a fixed seed produces the | ||
| 50 | +// same redo initialization sequence on the CPU and NPU paths. | ||
| 51 | +constexpr int64_t IVF_PQ_REDO_SEED_STRIDE = 15486557LL; | ||
| 47 | constexpr int IVF_PQ_FLAG_SIZE = 16; | 52 | constexpr int IVF_PQ_FLAG_SIZE = 16; |
| 48 | constexpr int TOPK_FLAT_ATTR_IDX_COUNT = 9; | 53 | constexpr int TOPK_FLAT_ATTR_IDX_COUNT = 9; |
| 54 | +constexpr double IVF_PQ_VERBOSE_CONVERGENCE_TOL = 1e-6; | ||
| 49 | // L3 topk merge attrs indices (mirror of TopkIvfpqL3AttrIdx) | 55 | // L3 topk merge attrs indices (mirror of TopkIvfpqL3AttrIdx) |
| 50 | constexpr int TOPK_IVFPQ_L3_ATTR_ASC_IDX = 0; | 56 | constexpr int TOPK_IVFPQ_L3_ATTR_ASC_IDX = 0; |
| 51 | constexpr int TOPK_IVFPQ_L3_ATTR_K_IDX = 1; | 57 | constexpr int TOPK_IVFPQ_L3_ATTR_K_IDX = 1; |
| @@ -162,6 +168,54 @@ static int getFusedAssignDim(int dim) { | |||
| 162 | return ((dim + alignment - 1) / alignment) * alignment; | 168 | return ((dim + alignment - 1) / alignment) * alignment; |
| 163 | } | 169 | } |
| 164 | 170 | ||
| 171 | +static double computeKMeansObjective( | ||
| 172 | + int dim, | ||
| 173 | + int totalSize, | ||
| 174 | + const std::vector<int32_t>& assignments, | ||
| 175 | + const std::vector<float>& data, | ||
| 176 | + const std::vector<float>& centroids, | ||
| 177 | + int64_t metricType) { | ||
| 178 | + double objective = 0.0; | ||
| 179 | + for (int i = 0; i < totalSize; ++i) { | ||
| 180 | + const float* x = data.data() + static_cast<size_t>(i) * dim; | ||
| 181 | + const float* c = centroids.data() + | ||
| 182 | + static_cast<size_t>(assignments[static_cast<size_t>(i)]) * dim; | ||
| 183 | + double value = 0.0; | ||
| 184 | + for (int j = 0; j < dim; ++j) { | ||
| 185 | + if (metricType == faiss::METRIC_INNER_PRODUCT) { | ||
| 186 | + value += static_cast<double>(x[j]) * c[j]; | ||
| 187 | + } else { | ||
| 188 | + const double diff = static_cast<double>(x[j]) - c[j]; | ||
| 189 | + value += diff * diff; | ||
| 190 | + } | ||
| 191 | + } | ||
| 192 | + objective += value; | ||
| 193 | + } | ||
Y 严重程度: 建议 问题: computeKMeansObjective 内部使用 double 累积目标值,但返回类型为 float。对于大规模训练集(如 PR 描述中提到的 200M 向量场景,即使经过子采样),目标值累加可能超出 float 精度范围,导致不同 redo 之间的比较不可靠。 原因: K-means 的 redo 选择依赖于 objective 的相对大小比较。当训练数据量较大或维度较高时,sum over all vectors 的 objective 值可能很大,float 的 23 位尾数(约 7 位有效数字)可能无法区分相近的 objective 值,使得 redo 选择退化为随机。内部已经用 double 累积说明作者意识到了精度需求,但返回时的 float 截断浪费了这份精度。 怎么改: 将返回类型改为 double,同时修改调用方 trainImpl_ 和 trainKMeansOnNpuDistributed 中 bestObjective 的类型为 double:
![]() ![]() | |||
| 194 | + return objective; | ||
| 195 | +} | ||
| 196 | + | ||
| 197 | +static const char* getKMeansProgressStatus( | ||
| 198 | + double objective, | ||
| 199 | + double previousObjective, | ||
| 200 | + bool hasPreviousObjective, | ||
| 201 | + int64_t metricType) { | ||
| 202 | + if (!hasPreviousObjective) { | ||
| 203 | + return "initial"; | ||
| 204 | + } | ||
| 205 | + const double delta = objective - previousObjective; | ||
| 206 | + const double relativeChange = | ||
| 207 | + std::abs(delta) / std::max(std::abs(previousObjective), 1.0); | ||
| 208 | + if (relativeChange <= IVF_PQ_VERBOSE_CONVERGENCE_TOL) { | ||
| 209 | + return "converged"; | ||
| 210 | + } | ||
| 211 | + const bool improved = metricType == faiss::METRIC_INNER_PRODUCT | ||
| 212 | + ? delta > 0.0 | ||
| 213 | + : delta < 0.0; | ||
| 214 | + return improved ? "improving" : "regressed"; | ||
| 215 | +} | ||
| 216 | + | ||
| 217 | +// Verify that the device has the vector-core count required by fused IVFPQ | ||
| 218 | +// training and assignment operators. | ||
| 165 | static void verifyTrainingCoreCount(int device) { | 219 | static void verifyTrainingCoreCount(int device) { |
| 166 | int64_t vectorCoreNum = 0; | 220 | int64_t vectorCoreNum = 0; |
| 167 | ACL_VERIFY(aclrtGetDeviceInfo( | 221 | ACL_VERIFY(aclrtGetDeviceInfo( |
| @@ -237,6 +291,7 @@ using aclnnAscendcIvfpqCoarseAssignGetWorkspaceSizeFuncType = aclnnStatus (*)( | |||
| 237 | const aclTensor* query, | 291 | const aclTensor* query, |
| 238 | const aclTensor* centroids, | 292 | const aclTensor* centroids, |
| 239 | const aclTensor* centroidsSqr, | 293 | const aclTensor* centroidsSqr, |
| 294 | + const int64_t metricType, | ||
| 240 | aclTensor* labels, | 295 | aclTensor* labels, |
| 241 | uint64_t* workspaceSize, | 296 | uint64_t* workspaceSize, |
| 242 | aclOpExecutor** executor); | 297 | aclOpExecutor** executor); |
| @@ -478,14 +533,18 @@ void IVFPQ::trainKMeansOnNpu( | |||
| 478 | int nlist, | 533 | int nlist, |
| 479 | std::vector<float>& centroidsOut, | 534 | std::vector<float>& centroidsOut, |
| 480 | int niter, | 535 | int niter, |
| 536 | + int nredo, | ||
| 481 | int64_t seed, | 537 | int64_t seed, |
| 482 | - bool spherical) { | 538 | + bool spherical, |
| 539 | + int64_t assignmentMetric, | ||
| 540 | + bool verbose) { | ||
| 483 | FAISS_ASSERT(resources != nullptr); | 541 | FAISS_ASSERT(resources != nullptr); |
| 484 | FAISS_ASSERT(n > 0); | 542 | FAISS_ASSERT(n > 0); |
| 485 | FAISS_ASSERT(x != nullptr); | 543 | FAISS_ASSERT(x != nullptr); |
| 486 | FAISS_ASSERT(dim > 0); | 544 | FAISS_ASSERT(dim > 0); |
| 487 | FAISS_THROW_IF_NOT_MSG( | 545 | FAISS_THROW_IF_NOT_MSG( |
| 488 | niter > 0, "K-means iteration count must be positive"); | 546 | niter > 0, "K-means iteration count must be positive"); |
| 547 | + FAISS_THROW_IF_NOT_MSG(nredo > 0, "K-means redo count must be positive"); | ||
| 489 | FAISS_ASSERT(n >= nlist); | 548 | FAISS_ASSERT(n >= nlist); |
| 490 | FAISS_ASSERT(nlist >= IVF_PQ_BURST_LEN); | 549 | FAISS_ASSERT(nlist >= IVF_PQ_BURST_LEN); |
| 491 | FAISS_THROW_IF_NOT_MSG( | 550 | FAISS_THROW_IF_NOT_MSG( |
| @@ -495,6 +554,11 @@ void IVFPQ::trainKMeansOnNpu( | |||
| 495 | 554 | ||
| 496 | const int totalSize = static_cast<int>(n); | 555 | const int totalSize = static_cast<int>(n); |
| 497 | const int deviceDim = getFusedAssignDim(dim); | 556 | const int deviceDim = getFusedAssignDim(dim); |
| 557 | + const int64_t actualSeed = seed >= 0 | ||
| 558 | + ? seed | ||
| 559 | + : static_cast<int64_t>(std::chrono::high_resolution_clock::now() | ||
| 560 | + .time_since_epoch() | ||
| 561 | + .count()); | ||
| 498 | 562 | ||
| 499 | aclrtStream stream = resources->getDefaultStreamCurrentDevice(); | 563 | aclrtStream stream = resources->getDefaultStreamCurrentDevice(); |
| 500 | 564 | ||
| @@ -518,7 +582,7 @@ void IVFPQ::trainKMeansOnNpu( | |||
| 518 | x, | 582 | x, |
| 519 | trainDataHost, | 583 | trainDataHost, |
| 520 | centroidsHost, | 584 | centroidsHost, |
| 521 | - seed, | 585 | + actualSeed, |
| 522 | spherical); | 586 | spherical); |
| 523 | 587 | ||
| 524 | resetTrainOp_(); | 588 | resetTrainOp_(); |
| @@ -546,24 +610,65 @@ void IVFPQ::trainKMeansOnNpu( | |||
| 546 | trainDataDeviceBytes, | 610 | trainDataDeviceBytes, |
| 547 | ACL_MEMCPY_HOST_TO_DEVICE)); | 611 | ACL_MEMCPY_HOST_TO_DEVICE)); |
| 548 | 612 | ||
| 549 | - trainImpl_( | 613 | + std::vector<float> bestCentroids; |
| 550 | - resources, | 614 | + double bestObjective = assignmentMetric == faiss::METRIC_INNER_PRODUCT |
| 551 | - (int)n, | 615 | + ? -HUGE_VAL |
| 552 | - x, | 616 | + : HUGE_VAL; |
| 553 | - dim, | 617 | + const auto trainStart = verbose ? std::chrono::steady_clock::now() |
| 554 | - deviceDim, | 618 | + : std::chrono::steady_clock::time_point{}; |
| 555 | - nlist, | 619 | + for (int redo = 0; redo < nredo; ++redo) { |
| 556 | - stream, | 620 | + if (redo != 0) { |
| 557 | - dataDev.data(), | 621 | + initTraining_( |
| 558 | - centroidsDev.data(), | 622 | + totalSize, |
| 559 | - centroidsSqrDev.data(), | 623 | + dim, |
| 560 | - assignmentLabelsDev.data(), | 624 | + nlist, |
| 561 | - centroidsHost, | 625 | + x, |
| 562 | - trainDataHost, | 626 | + trainDataHost, |
| 563 | - niter, | 627 | + centroidsHost, |
| 564 | - spherical); | 628 | + actualSeed + redo * IVF_PQ_REDO_SEED_STRIDE, |
| 565 | - | 629 | + spherical); |
| 566 | - centroidsOut.assign(centroidsHost.begin(), centroidsHost.end()); | 630 | + } |
| 631 | + const double objective = trainImpl_( | ||
| 632 | + resources, | ||
| 633 | + (int)n, | ||
| 634 | + x, | ||
| 635 | + dim, | ||
| 636 | + deviceDim, | ||
| 637 | + nlist, | ||
| 638 | + stream, | ||
| 639 | + dataDev.data(), | ||
| 640 | + centroidsDev.data(), | ||
| 641 | + centroidsSqrDev.data(), | ||
| 642 | + assignmentLabelsDev.data(), | ||
| 643 | + centroidsHost, | ||
| 644 | + trainDataHost, | ||
| 645 | + niter, | ||
| 646 | + spherical, | ||
| 647 | + assignmentMetric, | ||
| 648 | + verbose, | ||
| 649 | + redo, | ||
| 650 | + nredo); | ||
| 651 | + const bool isBetter = bestCentroids.empty() || | ||
| 652 | + (assignmentMetric == faiss::METRIC_INNER_PRODUCT | ||
| 653 | + ? objective > bestObjective | ||
| 654 | + : objective < bestObjective); | ||
| 655 | + if (verbose) { | ||
| 656 | + printf(" NpuIVFPQ::train: redo %d/%d final objective=%.6e " | ||
| 657 | + "selected=%s elapsed=%.3f s\n", | ||
| 658 | + redo + 1, | ||
| 659 | + nredo, | ||
| 660 | + objective, | ||
| 661 | + isBetter ? "yes" : "no", | ||
| 662 | + std::chrono::duration<double>( | ||
| 663 | + std::chrono::steady_clock::now() - trainStart) | ||
| 664 | + .count()); | ||
| 665 | + } | ||
| 666 | + if (isBetter) { | ||
Y 严重程度: 建议 问题: redo 循环结束后直接 std::move(bestCentroids) 到 centroidsOut,未处理 bestCentroids 可能为空的情况 原因: 当 nredo=1 时,第一次 redo 的 objective 会与初始值(HUGE_VALF 或 -HUGE_VALF)比较。正常情况下第一次结果总是 'better',bestCentroids 会被赋值。但如果 objective 因浮点溢出变为 +inf(L2)或 -inf(IP),比较结果为 false,bestCentroids 保持为空,std::move 后 centroidsOut 为空向量,后续代码访问空向量的数据会导致未定义行为或崩溃。trainKMeansOnNpuDistributed 中存在同样的问题。 怎么改: 在 redo 循环后添加防护:if (bestCentroids.empty()) { bestCentroids = centroidsHost; } 确保即使比较逻辑异常,也有兜底结果。或在循环中保证至少第一次 redo 的结果总是被采纳。 ![]() ![]() | |||
| 667 | + bestObjective = objective; | ||
| 668 | + bestCentroids = centroidsHost; | ||
| 669 | + } | ||
| 670 | + } | ||
| 671 | + centroidsOut = std::move(bestCentroids); | ||
D 严重程度: 建议 问题: 原因: 当 怎么改:
或保证第一次 redo 无条件采纳;同时建议 objective 全程用 ![]() ![]() | |||
| 567 | } | 672 | } |
| 568 | 673 | ||
| 569 | void IVFPQ::trainKMeansOnNpuDistributed( | 674 | void IVFPQ::trainKMeansOnNpuDistributed( |
| @@ -575,8 +680,11 @@ void IVFPQ::trainKMeansOnNpuDistributed( | |||
| 575 | int nlist, | 680 | int nlist, |
| 576 | std::vector<float>& centroidsOut, | 681 | std::vector<float>& centroidsOut, |
| 577 | int niter, | 682 | int niter, |
| 683 | + int nredo, | ||
| 578 | int64_t seed, | 684 | int64_t seed, |
| 579 | - bool spherical) { | 685 | + bool spherical, |
| 686 | + int64_t assignmentMetric, | ||
| 687 | + bool verbose) { | ||
| 580 | FAISS_THROW_IF_NOT_MSG( | 688 | FAISS_THROW_IF_NOT_MSG( |
| 581 | !resources.empty() && resources.size() == devices.size(), | 689 | !resources.empty() && resources.size() == devices.size(), |
| 582 | "distributed K-means requires one resource manager per device"); | 690 | "distributed K-means requires one resource manager per device"); |
| @@ -587,6 +695,7 @@ void IVFPQ::trainKMeansOnNpuDistributed( | |||
| 587 | FAISS_THROW_IF_NOT_MSG(dim > 0, "training dimension must be positive"); | 695 | FAISS_THROW_IF_NOT_MSG(dim > 0, "training dimension must be positive"); |
| 588 | FAISS_THROW_IF_NOT_MSG( | 696 | FAISS_THROW_IF_NOT_MSG( |
| 589 | niter > 0, "K-means iteration count must be positive"); | 697 | niter > 0, "K-means iteration count must be positive"); |
| 698 | + FAISS_THROW_IF_NOT_MSG(nredo > 0, "K-means redo count must be positive"); | ||
| 590 | 699 | ||
| 591 | struct TrainingState { | 700 | struct TrainingState { |
| 592 | int device = -1; | 701 | int device = -1; |
| @@ -605,6 +714,11 @@ void IVFPQ::trainKMeansOnNpuDistributed( | |||
| 605 | const int deviceDim = getFusedAssignDim(dim); | 714 | const int deviceDim = getFusedAssignDim(dim); |
| 606 | const int workerCount = static_cast<int>(resources.size()); | 715 | const int workerCount = static_cast<int>(resources.size()); |
| 607 | const int fragmentSize = (totalSize + workerCount - 1) / workerCount; | 716 | const int fragmentSize = (totalSize + workerCount - 1) / workerCount; |
| 717 | + const int64_t actualSeed = seed >= 0 | ||
| 718 | + ? seed | ||
| 719 | + : static_cast<int64_t>(std::chrono::high_resolution_clock::now() | ||
| 720 | + .time_since_epoch() | ||
| 721 | + .count()); | ||
| 608 | std::vector<float> trainDataHost(static_cast<size_t>(totalSize) * dim); | 722 | std::vector<float> trainDataHost(static_cast<size_t>(totalSize) * dim); |
| 609 | std::vector<float> trainDataDeviceHost; | 723 | std::vector<float> trainDataDeviceHost; |
| 610 | std::vector<float> centroidsHost(static_cast<size_t>(nlist) * dim); | 724 | std::vector<float> centroidsHost(static_cast<size_t>(nlist) * dim); |
| @@ -615,7 +729,7 @@ void IVFPQ::trainKMeansOnNpuDistributed( | |||
| 615 | x, | 729 | x, |
| 616 | trainDataHost, | 730 | trainDataHost, |
| 617 | centroidsHost, | 731 | centroidsHost, |
| 618 | - seed, | 732 | + actualSeed, |
| 619 | spherical); | 733 | spherical); |
| 620 | const float* trainDataDevice = trainDataHost.data(); | 734 | const float* trainDataDevice = trainDataHost.data(); |
| 621 | if (deviceDim != dim) { | 735 | if (deviceDim != dim) { |
| @@ -701,62 +815,151 @@ void IVFPQ::trainKMeansOnNpuDistributed( | |||
| 701 | } | 815 | } |
| 702 | 816 | ||
| 703 | std::vector<int32_t> assignments(static_cast<size_t>(totalSize)); | 817 | std::vector<int32_t> assignments(static_cast<size_t>(totalSize)); |
| 704 | - CentroidUpdateWorkspace_ centroidWorkspace; | 818 | + std::vector<float> bestCentroids; |
| 705 | - for (int iter = 0; iter < niter; ++iter) { | 819 | + double bestObjective = assignmentMetric == faiss::METRIC_INNER_PRODUCT |
| 706 | - std::vector<std::future<void>> workers; | 820 | + ? -HUGE_VAL |
| 707 | - workers.reserve(states.size()); | 821 | + : HUGE_VAL; |
| 708 | - for (const auto& statePtr : states) { | 822 | + const auto trainStart = verbose ? std::chrono::steady_clock::now() |
| 709 | - TrainingState* state = statePtr.get(); | 823 | + : std::chrono::steady_clock::time_point{}; |
| 710 | - workers.emplace_back(std::async( | 824 | + for (int redo = 0; redo < nredo; ++redo) { |
| 711 | - std::launch::async, | 825 | + if (redo != 0) { |
| 712 | - [state, ¢roidsHost, nlist, dim, deviceDim, iter]() { | 826 | + initTraining_( |
| 713 | - DeviceScope scope(state->device); | 827 | + totalSize, |
| 714 | - updateCentroidsToDevice_( | 828 | + dim, |
| 715 | - nlist, | 829 | + nlist, |
| 716 | - dim, | 830 | + x, |
| 717 | - deviceDim, | 831 | + trainDataHost, |
| 718 | - centroidsHost, | 832 | + centroidsHost, |
| 719 | - state->centroidsDev->data()); | 833 | + actualSeed + redo * IVF_PQ_REDO_SEED_STRIDE, |
| 720 | - runKMeans_( | 834 | + spherical); |
| 721 | - state->resources, | ||
| 722 | - nlist, | ||
| 723 | - dim, | ||
| 724 | - deviceDim, | ||
| 725 | - state->size, | ||
| 726 | - iter, | ||
| 727 | - centroidsHost, | ||
| 728 | - state->dataDev->data(), | ||
| 729 | - state->assignments, | ||
| 730 | - state->stream, | ||
| 731 | - state->centroidsDev->data(), | ||
| 732 | - state->centroidsSqrDev->data(), | ||
| 733 | - state->assignmentLabelsDev->data(), | ||
| 734 | - true); | ||
| 735 | - })); | ||
| 736 | } | 835 | } |
| 737 | - for (auto& worker : workers) { | 836 | + CentroidUpdateWorkspace_ centroidWorkspace; |
| 738 | - worker.get(); | 837 | + double objective = 0.0; |
| 838 | + double previousObjective = 0.0; | ||
| 839 | + bool hasPreviousObjective = false; | ||
| 840 | + for (int iter = 0; iter < niter; ++iter) { | ||
| 841 | + std::vector<std::future<void>> workers; | ||
| 842 | + workers.reserve(states.size()); | ||
| 843 | + for (const auto& statePtr : states) { | ||
| 844 | + TrainingState* state = statePtr.get(); | ||
| 845 | + workers.emplace_back(std::async( | ||
| 846 | + std::launch::async, | ||
| 847 | + [state, | ||
| 848 | + ¢roidsHost, | ||
| 849 | + nlist, | ||
| 850 | + dim, | ||
| 851 | + deviceDim, | ||
| 852 | + iter, | ||
| 853 | + assignmentMetric]() { | ||
| 854 | + DeviceScope scope(state->device); | ||
| 855 | + updateCentroidsToDevice_( | ||
| 856 | + nlist, | ||
| 857 | + dim, | ||
| 858 | + deviceDim, | ||
| 859 | + centroidsHost, | ||
| 860 | + state->centroidsDev->data()); | ||
| 861 | + runKMeans_( | ||
| 862 | + state->resources, | ||
| 863 | + nlist, | ||
| 864 | + dim, | ||
| 865 | + deviceDim, | ||
| 866 | + state->size, | ||
| 867 | + iter, | ||
| 868 | + centroidsHost, | ||
| 869 | + state->dataDev->data(), | ||
| 870 | + state->assignments, | ||
| 871 | + state->stream, | ||
| 872 | + state->centroidsDev->data(), | ||
| 873 | + state->centroidsSqrDev->data(), | ||
| 874 | + state->assignmentLabelsDev->data(), | ||
| 875 | + true, | ||
| 876 | + assignmentMetric); | ||
| 877 | + })); | ||
| 878 | + } | ||
| 879 | + for (auto& worker : workers) { | ||
| 880 | + worker.get(); | ||
| 881 | + } | ||
| 882 | + for (const auto& state : states) { | ||
| 883 | + std::copy( | ||
| 884 | + state->assignments.begin(), | ||
| 885 | + state->assignments.end(), | ||
| 886 | + assignments.begin() + state->offset); | ||
| 887 | + } | ||
| 888 | + // Verbose mode reports convergence every iteration; otherwise the | ||
| 889 | + // final objective alone is sufficient to select the best redo. | ||
| 890 | + if (verbose || iter == niter - 1) { | ||
| 891 | + objective = computeKMeansObjective( | ||
| 892 | + dim, | ||
| 893 | + totalSize, | ||
| 894 | + assignments, | ||
| 895 | + trainDataHost, | ||
| 896 | + centroidsHost, | ||
| 897 | + assignmentMetric); | ||
| 898 | + if (verbose) { | ||
| 899 | + const double delta = hasPreviousObjective | ||
| 900 | + ? objective - previousObjective | ||
| 901 | + : 0.0; | ||
| 902 | + const double relativeChange = hasPreviousObjective | ||
| 903 | + ? std::abs(delta) / | ||
| 904 | + std::max(std::abs(previousObjective), 1.0) | ||
| 905 | + : 0.0; | ||
| 906 | + printf(" NpuIVFPQ::train: redo %d/%d iter %d/%d " | ||
| 907 | + "objective=%.6e delta=%.6e rel_change=%.3e " | ||
| 908 | + "status=%s elapsed=%.3f s\n", | ||
| 909 | + redo + 1, | ||
| 910 | + nredo, | ||
| 911 | + iter + 1, | ||
| 912 | + niter, | ||
| 913 | + objective, | ||
| 914 | + delta, | ||
| 915 | + relativeChange, | ||
| 916 | + getKMeansProgressStatus( | ||
| 917 | + objective, | ||
| 918 | + previousObjective, | ||
| 919 | + hasPreviousObjective, | ||
| 920 | + assignmentMetric), | ||
| 921 | + std::chrono::duration<double>( | ||
| 922 | + std::chrono::steady_clock::now() - | ||
| 923 | + trainStart) | ||
| 924 | + .count()); | ||
| 925 | + previousObjective = objective; | ||
| 926 | + hasPreviousObjective = true; | ||
| 927 | + } | ||
| 928 | + } | ||
| 929 | + updateCentroids_( | ||
| 930 | + nlist, | ||
| 931 | + dim, | ||
| 932 | + totalSize, | ||
| 933 | + assignments, | ||
| 934 | + trainDataHost, | ||
| 935 | + centroidsHost, | ||
| 936 | + centroidWorkspace); | ||
| 937 | + if (spherical) { | ||
| 938 | + normL2_(dim, nlist, centroidsHost.data()); | ||
| 939 | + } | ||
| 739 | } | 940 | } |
| 740 | - for (const auto& state : states) { | 941 | + const bool isBetter = bestCentroids.empty() || |
| 741 | - std::copy( | 942 | + (assignmentMetric == faiss::METRIC_INNER_PRODUCT |
| 742 | - state->assignments.begin(), | 943 | + ? objective > bestObjective |
| 743 | - state->assignments.end(), | 944 | + : objective < bestObjective); |
| 744 | - assignments.begin() + state->offset); | 945 | + if (verbose) { |
| 946 | + printf(" NpuIVFPQ::train: redo %d/%d final objective=%.6e " | ||
| 947 | + "selected=%s elapsed=%.3f s\n", | ||
| 948 | + redo + 1, | ||
| 949 | + nredo, | ||
| 950 | + objective, | ||
| 951 | + isBetter ? "yes" : "no", | ||
| 952 | + std::chrono::duration<double>( | ||
| 953 | + std::chrono::steady_clock::now() - trainStart) | ||
| 954 | + .count()); | ||
| 745 | } | 955 | } |
| 746 | - updateCentroids_( | 956 | + if (isBetter) { |
| 747 | - nlist, | 957 | + bestObjective = objective; |
| 748 | - dim, | 958 | + bestCentroids = centroidsHost; |
| 749 | - totalSize, | ||
| 750 | - assignments, | ||
| 751 | - trainDataHost, | ||
| 752 | - centroidsHost, | ||
| 753 | - centroidWorkspace); | ||
| 754 | - if (spherical) { | ||
| 755 | - normL2_(dim, nlist, centroidsHost.data()); | ||
| 756 | } | 959 | } |
| 757 | } | 960 | } |
| 758 | 961 | ||
| 759 | - centroidsOut = std::move(centroidsHost); | 962 | + centroidsOut = std::move(bestCentroids); |
| 760 | } | 963 | } |
| 761 | 964 | ||
| 762 | void IVFPQ::assignCentroidOnDevice( | 965 | void IVFPQ::assignCentroidOnDevice( |
| @@ -771,8 +974,8 @@ void IVFPQ::assignCentroidOnDevice( | |||
| 771 | "add assignment size exceeds supported range"); | 974 | "add assignment size exceeds supported range"); |
| 772 | verifyTrainingCoreCount(getCurrentDevice()); | 975 | verifyTrainingCoreCount(getCurrentDevice()); |
| 773 | 976 | ||
| 774 | - // The add path needs only argmin(k=1). Use the dedicated fused operator | 977 | + // The add path needs only top-1 coarse assignment. The fused operator |
| 775 | - // instead of materializing batch x nlist distances for AICPU TopK. | 978 | + // uses direct IP argmax for IP indices and L2 argmin otherwise. |
| 776 | FAISS_THROW_IF_NOT_MSG( | 979 | FAISS_THROW_IF_NOT_MSG( |
| 777 | aclnnAscendcIvfpqCoarseAssignGetWorkspaceSize, | 980 | aclnnAscendcIvfpqCoarseAssignGetWorkspaceSize, |
| 778 | "aclnnAscendcIvfpqCoarseAssignGetWorkspaceSize not loaded; " | 981 | "aclnnAscendcIvfpqCoarseAssignGetWorkspaceSize not loaded; " |
| @@ -792,17 +995,6 @@ void IVFPQ::assignCentroidOnDevice( | |||
| 792 | DeviceVector<uint8_t>& workspaceDev = | 995 | DeviceVector<uint8_t>& workspaceDev = |
| 793 | ensureDeviceVectorCache(assignWorkspaceDev_, resources_, stream); | 996 | ensureDeviceVectorCache(assignWorkspaceDev_, resources_, stream); |
| 794 | DeviceVector<float>* centroidNormsDev = ¢roidsSqrSum_; | 997 | DeviceVector<float>* centroidNormsDev = ¢roidsSqrSum_; |
| 795 | - if (metric_ == faiss::METRIC_INNER_PRODUCT) { | ||
| 796 | - DeviceVector<float>& zeroCentroidNormsDev = ensureDeviceVectorCache( | ||
| 797 | - assignZeroCentroidNormsDev_, resources_, stream); | ||
| 798 | - const bool resized = | ||
| 799 | - zeroCentroidNormsDev.size() != static_cast<size_t>(numLists_); | ||
| 800 | - zeroCentroidNormsDev.resize(static_cast<size_t>(numLists_), stream); | ||
| 801 | - if (resized) { | ||
| 802 | - zeroCentroidNormsDev.setAll(0.0f, stream); | ||
| 803 | - } | ||
| 804 | - centroidNormsDev = &zeroCentroidNormsDev; | ||
| 805 | - } | ||
| 806 | dataDev.resize(static_cast<size_t>(totalSize) * dim_, stream); | 998 | dataDev.resize(static_cast<size_t>(totalSize) * dim_, stream); |
| 807 | labelsDev.resize(totalSize, stream); | 999 | labelsDev.resize(totalSize, stream); |
| 808 | 1000 | ||
| @@ -846,6 +1038,9 @@ void IVFPQ::assignCentroidOnDevice( | |||
| 846 | queryTensor.get(), | 1038 | queryTensor.get(), |
| 847 | centroidsTensor.get(), | 1039 | centroidsTensor.get(), |
| 848 | centroidsSqrTensor.get(), | 1040 | centroidsSqrTensor.get(), |
| 1041 | + metric_ == faiss::METRIC_INNER_PRODUCT | ||
| 1042 | + ? faiss::METRIC_INNER_PRODUCT | ||
| 1043 | + : faiss::METRIC_L2, | ||
| 849 | labelsTensor.get(), | 1044 | labelsTensor.get(), |
| 850 | &workspaceSize, | 1045 | &workspaceSize, |
| 851 | &executor)); | 1046 | &executor)); |
| @@ -893,6 +1088,8 @@ bool IVFPQ::encodeResidualsOnDevice( | |||
| 893 | FAISS_ASSERT(x != nullptr); | 1088 | FAISS_ASSERT(x != nullptr); |
| 894 | FAISS_ASSERT(assignments != nullptr); | 1089 | FAISS_ASSERT(assignments != nullptr); |
| 895 | FAISS_ASSERT(codes != nullptr); | 1090 | FAISS_ASSERT(codes != nullptr); |
| 1091 | + // The residual encoder requires each sub-quantizer dimension to be | ||
| 1092 | + // aligned to 8 elements; use the CPU fallback for unsupported sizes. | ||
| 896 | if (dimPerSubQuantizer_ % 8 != 0) { | 1093 | if (dimPerSubQuantizer_ % 8 != 0) { |
D 严重程度: 严重 问题: 原因: 典型配置 怎么改:
或保留算子路径,但删除/改写误导性注释,并补全末尾残缺句子。 ![]() ![]() | |||
| 897 | return false; | 1094 | return false; |
| 898 | } | 1095 | } |
| @@ -1028,7 +1225,7 @@ bool IVFPQ::encodeResidualsOnDevice( | |||
| 1028 | } | 1225 | } |
| 1029 | 1226 | ||
| 1030 | // Training loop: runKMeans_ + updateCentroids_ + normL2_ per iteration. | 1227 | // Training loop: runKMeans_ + updateCentroids_ + normL2_ per iteration. |
| 1031 | -void IVFPQ::trainImpl_( | 1228 | +double IVFPQ::trainImpl_( |
| 1032 | NpuResources* resources, | 1229 | NpuResources* resources, |
| 1033 | int n, | 1230 | int n, |
| 1034 | const float* x, | 1231 | const float* x, |
| @@ -1043,13 +1240,22 @@ void IVFPQ::trainImpl_( | |||
| 1043 | std::vector<float>& centroidsHost, | 1240 | std::vector<float>& centroidsHost, |
| 1044 | std::vector<float>& trainDataHost, | 1241 | std::vector<float>& trainDataHost, |
| 1045 | int niter, | 1242 | int niter, |
| 1046 | - bool spherical) { | 1243 | + bool spherical, |
| 1244 | + int64_t assignmentMetric, | ||
| 1245 | + bool verbose, | ||
| 1246 | + int redo, | ||
| 1247 | + int nredo) { | ||
| 1047 | (void)n; | 1248 | (void)n; |
| 1048 | (void)x; | 1249 | (void)x; |
| 1049 | 1250 | ||
| 1050 | int totalSize = (int)(trainDataHost.size() / dim); | 1251 | int totalSize = (int)(trainDataHost.size() / dim); |
| 1051 | std::vector<int32_t> totalAssigns(totalSize); | 1252 | std::vector<int32_t> totalAssigns(totalSize); |
| 1052 | CentroidUpdateWorkspace_ centroidWorkspace; | 1253 | CentroidUpdateWorkspace_ centroidWorkspace; |
| 1254 | + double objective = 0.0; | ||
| 1255 | + double previousObjective = 0.0; | ||
| 1256 | + bool hasPreviousObjective = false; | ||
| 1257 | + const auto trainStart = verbose ? std::chrono::steady_clock::now() | ||
| 1258 | + : std::chrono::steady_clock::time_point{}; | ||
| 1053 | for (int iter = 0; iter < niter; iter++) { | 1259 | for (int iter = 0; iter < niter; iter++) { |
| 1054 | updateCentroidsToDevice_( | 1260 | updateCentroidsToDevice_( |
| 1055 | nlist, dim, deviceDim, centroidsHost, centroidsDev); | 1261 | nlist, dim, deviceDim, centroidsHost, centroidsDev); |
| @@ -1067,7 +1273,49 @@ void IVFPQ::trainImpl_( | |||
| 1067 | centroidsDev, | 1273 | centroidsDev, |
Y 严重程度: 建议 问题: computeKMeansObjective 在 trainImpl_ 的每次 K-means 迭代中都被调用,但只有最后一次迭代的返回值被使用(用于 redo 间的最优选择)。同样的模式也出现在 trainKMeansOnNpuDistributed 中。 原因: computeKMeansObjective 遍历全部训练向量 × 全部维度,复杂度为 O(totalSize × dim)。在典型配置下(如 niter=25, nredo=3),该函数会被调用 75 次,但仅最后一次的结果有效。前 74 次计算完全浪费 CPU 时间,且在 NPU 训练流程中这段 CPU 计算会阻塞流水线,降低整体训练吞吐。 怎么改: 将 objective 计算移到迭代循环之后,或仅在最后一次迭代时计算:
同理修改 trainKMeansOnNpuDistributed 中的对应逻辑。 ![]() ![]() | |||
| 1068 | centroidsSqrDev, | 1274 | centroidsSqrDev, |
| 1069 | assignmentLabelsDev, | 1275 | assignmentLabelsDev, |
| 1070 | - true); | 1276 | + true, |
| 1277 | + assignmentMetric); | ||
| 1278 | + | ||
| 1279 | + // Verbose mode reports convergence every iteration; otherwise the | ||
| 1280 | + // final objective alone is sufficient to select the best redo. | ||
| 1281 | + if (verbose || iter == niter - 1) { | ||
| 1282 | + objective = computeKMeansObjective( | ||
| 1283 | + dim, | ||
| 1284 | + totalSize, | ||
| 1285 | + totalAssigns, | ||
| 1286 | + trainDataHost, | ||
| 1287 | + centroidsHost, | ||
| 1288 | + assignmentMetric); | ||
| 1289 | + if (verbose) { | ||
| 1290 | + const double delta = hasPreviousObjective | ||
| 1291 | + ? objective - previousObjective | ||
| 1292 | + : 0.0; | ||
| 1293 | + const double relativeChange = hasPreviousObjective | ||
| 1294 | + ? std::abs(delta) / | ||
| 1295 | + std::max(std::abs(previousObjective), 1.0) | ||
| 1296 | + : 0.0; | ||
| 1297 | + printf(" NpuIVFPQ::train: redo %d/%d iter %d/%d " | ||
| 1298 | + "objective=%.6e delta=%.6e rel_change=%.3e " | ||
| 1299 | + "status=%s elapsed=%.3f s\n", | ||
| 1300 | + redo + 1, | ||
| 1301 | + nredo, | ||
| 1302 | + iter + 1, | ||
| 1303 | + niter, | ||
| 1304 | + objective, | ||
| 1305 | + delta, | ||
| 1306 | + relativeChange, | ||
| 1307 | + getKMeansProgressStatus( | ||
| 1308 | + objective, | ||
| 1309 | + previousObjective, | ||
| 1310 | + hasPreviousObjective, | ||
| 1311 | + assignmentMetric), | ||
| 1312 | + std::chrono::duration<double>( | ||
| 1313 | + std::chrono::steady_clock::now() - trainStart) | ||
| 1314 | + .count()); | ||
| 1315 | + previousObjective = objective; | ||
| 1316 | + hasPreviousObjective = true; | ||
| 1317 | + } | ||
| 1318 | + } | ||
| 1071 | 1319 | ||
| 1072 | updateCentroids_( | 1320 | updateCentroids_( |
| 1073 | nlist, | 1321 | nlist, |
| @@ -1083,6 +1331,7 @@ void IVFPQ::trainImpl_( | |||
| 1083 | } | 1331 | } |
| 1084 | updateCentroidsToDevice_( | 1332 | updateCentroidsToDevice_( |
| 1085 | nlist, dim, deviceDim, centroidsHost, centroidsDev); | 1333 | nlist, dim, deviceDim, centroidsHost, centroidsDev); |
| 1334 | + return objective; | ||
| 1086 | } | 1335 | } |
| 1087 | 1336 | ||
| 1088 | // Initialize training data and centroids from input vectors. | 1337 | // Initialize training data and centroids from input vectors. |
| @@ -1147,19 +1396,24 @@ void IVFPQ::runKMeans_( | |||
| 1147 | float* centroidsDev, | 1396 | float* centroidsDev, |
| 1148 | float* centroidsSqrDev, | 1397 | float* centroidsSqrDev, |
| 1149 | int32_t* assignmentLabelsDev, | 1398 | int32_t* assignmentLabelsDev, |
| 1150 | - bool centroidsResident) { | 1399 | + bool centroidsResident, |
| 1400 | + int64_t assignmentMetric) { | ||
| 1151 | (void)iter; | 1401 | (void)iter; |
| 1152 | 1402 | ||
| 1153 | - // Compute centroid squared norms (centroidsDoubleHost). | 1403 | + std::vector<float> centroidsDoubleHost; |
| 1154 | - std::vector<float> centroidsDoubleHost(nlist); | 1404 | + if (assignmentMetric == faiss::METRIC_L2) { |
| 1405 | + centroidsDoubleHost.resize(nlist); | ||
| 1155 | 1406 | ||
| 1156 | - for (int i = 0; i < nlist; i++) { | 1407 | + for (int i = 0; i < nlist; i++) { |
| 1157 | - float normSq = 0.0f; | 1408 | + float normSq = 0.0f; |
| 1158 | - for (int j = 0; j < dim; j++) { | 1409 | + for (int j = 0; j < dim; j++) { |
| 1159 | - float val = centroidsHost[static_cast<size_t>(i) * dim + j]; | 1410 | + float val = centroidsHost[static_cast<size_t>(i) * dim + j]; |
| 1160 | - normSq += val * val; | 1411 | + normSq += val * val; |
| 1412 | + } | ||
| 1413 | + centroidsDoubleHost[i] = normSq; | ||
| 1161 | } | 1414 | } |
| 1162 | - centroidsDoubleHost[i] = normSq; | 1415 | + } else { |
| 1416 | + centroidsDoubleHost.assign(nlist, 0.0f); | ||
| 1163 | } | 1417 | } |
| 1164 | 1418 | ||
| 1165 | if (!centroidsResident) { | 1419 | if (!centroidsResident) { |
| @@ -1192,6 +1446,7 @@ void IVFPQ::runKMeans_( | |||
| 1192 | assignmentLabelsDev + processed, | 1446 | assignmentLabelsDev + processed, |
| 1193 | nlist, | 1447 | nlist, |
| 1194 | deviceDim, | 1448 | deviceDim, |
| 1449 | + assignmentMetric, | ||
| 1195 | stream, | 1450 | stream, |
| 1196 | workspaceDev); | 1451 | workspaceDev); |
| 1197 | } | 1452 | } |
| @@ -1225,6 +1480,7 @@ void IVFPQ::runTrainAssignOp_( | |||
| 1225 | int32_t* labelsDev, | 1480 | int32_t* labelsDev, |
| 1226 | int nlist, | 1481 | int nlist, |
| 1227 | int dim, | 1482 | int dim, |
| 1483 | + int64_t assignmentMetric, | ||
| 1228 | aclrtStream stream, | 1484 | aclrtStream stream, |
| 1229 | DeviceVector<uint8_t>& workspaceDev) { | 1485 | DeviceVector<uint8_t>& workspaceDev) { |
| 1230 | (void)resources; | 1486 | (void)resources; |
| @@ -1247,6 +1503,7 @@ void IVFPQ::runTrainAssignOp_( | |||
| 1247 | queryTensor.get(), | 1503 | queryTensor.get(), |
| 1248 | centroidsTensor.get(), | 1504 | centroidsTensor.get(), |
| 1249 | centroidsSqrTensor.get(), | 1505 | centroidsSqrTensor.get(), |
| 1506 | + assignmentMetric, | ||
| 1250 | labelsTensor.get(), | 1507 | labelsTensor.get(), |
| 1251 | &workspaceSize, | 1508 | &workspaceSize, |
| 1252 | &executor)); | 1509 | &executor)); |
| @@ -1280,6 +1537,7 @@ void IVFPQ::updateCentroids_( | |||
| 1280 | static_cast<size_t>(maxThreads) * nlist); | 1537 | static_cast<size_t>(maxThreads) * nlist); |
| 1281 | workspace.clusterOffsets.resize(static_cast<size_t>(nlist) + 1); | 1538 | workspace.clusterOffsets.resize(static_cast<size_t>(nlist) + 1); |
| 1282 | workspace.clusterCounts.resize(nlist); | 1539 | workspace.clusterCounts.resize(nlist); |
| 1540 | + workspace.centroidCounts.resize(nlist); | ||
| 1283 | workspace.bucketedTrainData.resize( | 1541 | workspace.bucketedTrainData.resize( |
| 1284 | static_cast<size_t>(totalSize) * dim); | 1542 | static_cast<size_t>(totalSize) * dim); |
| 1285 | } catch (const std::bad_alloc&) { | 1543 | } catch (const std::bad_alloc&) { |
| @@ -1288,9 +1546,9 @@ void IVFPQ::updateCentroids_( | |||
| 1288 | } | 1546 | } |
| 1289 | } | 1547 | } |
| 1290 | 1548 | ||
Y 严重程度: 提示 问题: updateCentroids_ 中 hassign(聚类计数)的类型从 std::vector 原因: 聚类计数本质上是整数值,使用 float 存储虽然对于典型 nlist 范围内精确无误,但语义上不够清晰。同时 clusterCounts(int)和 centroidCounts(float)并存容易让后续维护者困惑两者的关系和用途。 怎么改: 如果 float 的引入是为了后续支持加权聚类(非整数权重),建议添加注释说明设计意图。如果仅是为了统一类型,考虑直接将 clusterCounts 也改为 float 并合并为同一字段,减少冗余。 ![]() ![]() | |||
| 1291 | - std::vector<int> fallbackCounts; | 1549 | + std::vector<float> fallbackCounts; |
| 1292 | - std::vector<int>& hassign = | 1550 | + std::vector<float>& hassign = |
| 1293 | - useBucketedUpdate ? workspace.clusterCounts : fallbackCounts; | 1551 | + useBucketedUpdate ? workspace.centroidCounts : fallbackCounts; |
| 1294 | if (useBucketedUpdate) { | 1552 | if (useBucketedUpdate) { |
| 1295 | // Use fixed logical chunks instead of OpenMP thread ids so the bucket | 1553 | // Use fixed logical chunks instead of OpenMP thread ids so the bucket |
| 1296 | // layout remains valid even when OpenMP changes the active team size. | 1554 | // layout remains valid even when OpenMP changes the active team size. |
| @@ -1314,13 +1572,14 @@ void IVFPQ::updateCentroids_( | |||
| 1314 | count += workspace.threadOffsets | 1572 | count += workspace.threadOffsets |
| 1315 | [static_cast<size_t>(rank) * nlist + ci]; | 1573 | [static_cast<size_t>(rank) * nlist + ci]; |
| 1316 | } | 1574 | } |
| 1317 | - hassign[ci] = count; | 1575 | + workspace.clusterCounts[ci] = count; |
| 1576 | + hassign[ci] = static_cast<float>(count); | ||
| 1318 | } | 1577 | } |
| 1319 | 1578 | ||
| 1320 | workspace.clusterOffsets[0] = 0; | 1579 | workspace.clusterOffsets[0] = 0; |
| 1321 | for (int ci = 0; ci < nlist; ++ci) { | 1580 | for (int ci = 0; ci < nlist; ++ci) { |
| 1322 | workspace.clusterOffsets[ci + 1] = | 1581 | workspace.clusterOffsets[ci + 1] = |
| 1323 | - workspace.clusterOffsets[ci] + hassign[ci]; | 1582 | + workspace.clusterOffsets[ci] + workspace.clusterCounts[ci]; |
| 1324 | } | 1583 | } |
| 1325 | 1584 | ||
| 1326 | 1585 | ||
| @@ -1380,7 +1639,7 @@ void IVFPQ::updateCentroids_( | |||
| 1380 | 0, | 1639 | 0, |
| 1381 | static_cast<size_t>(nlist) * static_cast<size_t>(dim) * | 1640 | static_cast<size_t>(nlist) * static_cast<size_t>(dim) * |
| 1382 | sizeof(float)); | 1641 | sizeof(float)); |
| 1383 | - hassign.assign(nlist, 0); | 1642 | + hassign.assign(nlist, 0.0f); |
| 1384 | 1643 | ||
| 1385 | { | 1644 | { |
| 1386 | const int nt = omp_get_num_threads(); | 1645 | const int nt = omp_get_num_threads(); |
| @@ -1405,10 +1664,10 @@ void IVFPQ::updateCentroids_( | |||
| 1405 | } | 1664 | } |
| 1406 | 1665 | ||
| 1407 | for (int ci = 0; ci < nlist; ci++) { | 1666 | for (int ci = 0; ci < nlist; ci++) { |
| 1408 | - if (hassign[ci] == 0) { | 1667 | + if (hassign[ci] <= 0.0f) { |
| 1409 | continue; | 1668 | continue; |
| 1410 | } | 1669 | } |
| 1411 | - float norm = 1.0f / static_cast<float>(hassign[ci]); | 1670 | + float norm = 1.0f / hassign[ci]; |
| 1412 | float* c = centroidsHost.data() + static_cast<size_t>(ci) * dim; | 1671 | float* c = centroidsHost.data() + static_cast<size_t>(ci) * dim; |
| 1413 | for (int j = 0; j < dim; j++) { | 1672 | for (int j = 0; j < dim; j++) { |
| 1414 | c[j] *= norm; | 1673 | c[j] *= norm; |
| @@ -1426,7 +1685,7 @@ void IVFPQ::updateCentroids_( | |||
| 1426 | } | 1685 | } |
| 1427 | size_t cj; | 1686 | size_t cj; |
| 1428 | for (cj = 0; true; cj = (cj + 1) % nlist) { | 1687 | for (cj = 0; true; cj = (cj + 1) % nlist) { |
| 1429 | - const float p = (static_cast<float>(hassign[cj]) - 1.0f) / | 1688 | + const float p = (hassign[cj] - 1.0f) / |
| 1430 | static_cast<float>(totalSize - nlist); | 1689 | static_cast<float>(totalSize - nlist); |
| 1431 | const float r = rng.rand_float(); | 1690 | const float r = rng.rand_float(); |
| 1432 | if (r < p) { | 1691 | if (r < p) { |
| @@ -1459,20 +1718,7 @@ void IVFPQ::updateCentroids_( | |||
| 1459 | 1718 | ||
| 1460 | // L2-normalize each centroid. | 1719 | // L2-normalize each centroid. |
| 1461 | void IVFPQ::normL2_(int dim, int nlist, float* data) { | 1720 | void IVFPQ::normL2_(int dim, int nlist, float* data) { |
| 1462 | -#pragma omp parallel for | 1721 | + faiss::fvec_renorm_L2(dim, nlist, data); |
| 1463 | - for (int i = 0; i < nlist; i++) { | ||
| 1464 | - float* vec = data + static_cast<size_t>(i) * dim; | ||
| 1465 | - float norm2 = 0.0f; | ||
| 1466 | - for (int d = 0; d < dim; d++) { | ||
| 1467 | - norm2 += vec[d] * vec[d]; | ||
| 1468 | - } | ||
| 1469 | - if (norm2 > 1e-30f) { | ||
| 1470 | - float inv = 1.0f / sqrtf(norm2); | ||
| 1471 | - for (int d = 0; d < dim; d++) { | ||
| 1472 | - vec[d] *= inv; | ||
| 1473 | - } | ||
| 1474 | - } | ||
| 1475 | - } | ||
| 1476 | } | 1722 | } |
| 1477 | 1723 | ||
| 1478 | // Upload final centroids to the device buffer. | 1724 | // Upload final centroids to the device buffer. |
| @@ -131,8 +131,11 @@ class IVFPQ : public IVFBase { | |||
| 131 | int nlist, | 131 | int nlist, |
| 132 | std::vector<float>& centroidsOut, | 132 | std::vector<float>& centroidsOut, |
| 133 | int niter, | 133 | int niter, |
| 134 | + int nredo, | ||
| 134 | int64_t seed, | 135 | int64_t seed, |
| 135 | - bool spherical); | 136 | + bool spherical, |
| 137 | + int64_t assignmentMetric, | ||
| 138 | + bool verbose); | ||
| 136 | 139 | ||
| 137 | static void trainKMeansOnNpuDistributed( | 140 | static void trainKMeansOnNpuDistributed( |
| 138 | const std::vector<NpuResources*>& resources, | 141 | const std::vector<NpuResources*>& resources, |
| @@ -143,8 +146,11 @@ class IVFPQ : public IVFBase { | |||
| 143 | int nlist, | 146 | int nlist, |
| 144 | std::vector<float>& centroidsOut, | 147 | std::vector<float>& centroidsOut, |
| 145 | int niter, | 148 | int niter, |
| 149 | + int nredo, | ||
| 146 | int64_t seed, | 150 | int64_t seed, |
| 147 | - bool spherical); | 151 | + bool spherical, |
| 152 | + int64_t assignmentMetric, | ||
| 153 | + bool verbose); | ||
| 148 | 154 | ||
| 149 | /// Assign vectors to resident IVF centroids using cached instance | 155 | /// Assign vectors to resident IVF centroids using cached instance |
| 150 | /// workspaces. Intended for addL1 fast path. | 156 | /// workspaces. Intended for addL1 fast path. |
| @@ -170,12 +176,15 @@ class IVFPQ : public IVFBase { | |||
| 170 | struct CentroidUpdateWorkspace_ { | 176 | struct CentroidUpdateWorkspace_ { |
| 171 | std::vector<int> threadOffsets; | 177 | std::vector<int> threadOffsets; |
| 172 | std::vector<int> clusterOffsets; | 178 | std::vector<int> clusterOffsets; |
| 179 | + // Integer counts are required to build bucket offsets. | ||
| 173 | std::vector<int> clusterCounts; | 180 | std::vector<int> clusterCounts; |
| 181 | + // Counts become fractional after empty-cluster splitting. | ||
| 182 | + std::vector<float> centroidCounts; | ||
| 174 | std::vector<float> bucketedTrainData; | 183 | std::vector<float> bucketedTrainData; |
| 175 | }; | 184 | }; |
| 176 | 185 | ||
| 177 | /// Runs the full K-means training pipeline on NPU. | 186 | /// Runs the full K-means training pipeline on NPU. |
| 178 | - static void trainImpl_( | 187 | + static double trainImpl_( |
| 179 | NpuResources* resources, | 188 | NpuResources* resources, |
| 180 | int n, | 189 | int n, |
| 181 | const float* x, | 190 | const float* x, |
| @@ -190,7 +199,11 @@ class IVFPQ : public IVFBase { | |||
| 190 | std::vector<float>& centroidsHost, | 199 | std::vector<float>& centroidsHost, |
| 191 | std::vector<float>& trainDataHost, | 200 | std::vector<float>& trainDataHost, |
| 192 | int niter, | 201 | int niter, |
| 193 | - bool spherical); | 202 | + bool spherical, |
| 203 | + int64_t assignmentMetric, | ||
| 204 | + bool verbose, | ||
| 205 | + int redo, | ||
| 206 | + int nredo); | ||
| 194 | 207 | ||
| 195 | /// Subsample `totalSize` vectors from `x` into `trainDataHost`, pick | 208 | /// Subsample `totalSize` vectors from `x` into `trainDataHost`, pick |
| 196 | /// `nlist` initial centroids via faiss::rand_perm (seed + 1, matching | 209 | /// `nlist` initial centroids via faiss::rand_perm (seed + 1, matching |
| @@ -226,10 +239,11 @@ class IVFPQ : public IVFBase { | |||
| 226 | float* centroidsDev, | 239 | float* centroidsDev, |
| 227 | float* centroidsSqrDev, | 240 | float* centroidsSqrDev, |
| 228 | int32_t* assignmentLabelsDev, | 241 | int32_t* assignmentLabelsDev, |
| 229 | - bool centroidsResident); | 242 | + bool centroidsResident, |
| 243 | + int64_t assignmentMetric); | ||
| 230 | 244 | ||
| 231 | - /// Executes the fused coarse-assignment operator for one L2 K-means | 245 | + /// Executes one fused coarse-assignment batch, writing one int32 centroid |
| 232 | - /// batch, writing one int32 centroid label per training vector. | 246 | + /// label per training vector using the requested IP or L2 metric. |
| 233 | static void runTrainAssignOp_( | 247 | static void runTrainAssignOp_( |
| 234 | NpuResources* resources, | 248 | NpuResources* resources, |
| 235 | int batch, | 249 | int batch, |
| @@ -239,6 +253,7 @@ class IVFPQ : public IVFBase { | |||
| 239 | int32_t* labelsDev, | 253 | int32_t* labelsDev, |
| 240 | int nlist, | 254 | int nlist, |
| 241 | int dim, | 255 | int dim, |
| 256 | + int64_t assignmentMetric, | ||
| 242 | aclrtStream stream, | 257 | aclrtStream stream, |
| 243 | DeviceVector<uint8_t>& workspaceDev); | 258 | DeviceVector<uint8_t>& workspaceDev); |
| 244 | 259 | ||
| @@ -251,8 +266,7 @@ class IVFPQ : public IVFBase { | |||
| 251 | std::vector<float>& centroidsHost, | 266 | std::vector<float>& centroidsHost, |
| 252 | CentroidUpdateWorkspace_& workspace); | 267 | CentroidUpdateWorkspace_& workspace); |
| 253 | 268 | ||
| 254 | - /// L2-normalize each centroid in-place (threshold 1e-30 on ||c||^2). | 269 | + /// L2-normalize each centroid in-place with FAISS's CPU implementation. |
| 255 | - /// `#pragma omp parallel for` is used for parallelism. | ||
| 256 | static void normL2_(int dim, int nlist, float* data); | 270 | static void normL2_(int dim, int nlist, float* data); |
| 257 | 271 | ||
| 258 | /// Upload centroids to a device buffer with row stride `deviceDim`. | 272 | /// Upload centroids to a device buffer with row stride `deviceDim`. |
| @@ -75,6 +75,25 @@ constexpr int OPQ_PQ_FLAG_SIZE = 16; | |||
| 75 | constexpr int OPQ_TOPK_ATTR_IDX_COUNT = 9; | 75 | constexpr int OPQ_TOPK_ATTR_IDX_COUNT = 9; |
| 76 | // 与 IVFPQ 分块策略一致, 按 512MB 自适应每批行数, 避免大底库 OOM。 | 76 | // 与 IVFPQ 分块策略一致, 按 512MB 自适应每批行数, 避免大底库 OOM。 |
| 77 | constexpr size_t OPQ_BATCH_BUDGET_BYTES = 512ULL * 1024 * 1024; | 77 | constexpr size_t OPQ_BATCH_BUDGET_BYTES = 512ULL * 1024 * 1024; |
| 78 | +constexpr double OPQ_VERBOSE_CONVERGENCE_TOL = 1e-6; | ||
| 79 | + | ||
| 80 | +static double computeReconstructionError( | ||
| 81 | + idx_t n, | ||
| 82 | + int dim, | ||
| 83 | + const float* source, | ||
| 84 | + const float* reconstruction) { | ||
| 85 | + double error = 0.0; | ||
| 86 | + for (idx_t i = 0; i < n; ++i) { | ||
| 87 | + const float* sourceRow = source + i * dim; | ||
| 88 | + const float* reconstructionRow = reconstruction + i * dim; | ||
| 89 | + for (int j = 0; j < dim; ++j) { | ||
| 90 | + const double diff = | ||
| 91 | + static_cast<double>(sourceRow[j]) - reconstructionRow[j]; | ||
| 92 | + error += diff * diff; | ||
| 93 | + } | ||
| 94 | + } | ||
| 95 | + return error / static_cast<double>(n); | ||
| 96 | +} | ||
| 78 | 97 | ||
| 79 | // 距离/topk 算子函数指针 | 98 | // 距离/topk 算子函数指针 |
| 80 | const aclnnDistanceFlatExtremaAtFp32GetWorkspaceSizeFuncType& | 99 | const aclnnDistanceFlatExtremaAtFp32GetWorkspaceSizeFuncType& |
| @@ -926,11 +945,16 @@ void OPQ::train( | |||
| 926 | 945 | ||
| 927 | if (verbose) { | 946 | if (verbose) { |
| 928 | printf("NpuOPQ::train: training OPQ rotation matrix " | 947 | printf("NpuOPQ::train: training OPQ rotation matrix " |
| 929 | - "for M=%d from %" PRId64 " vectors in %dD -> %dD\n", | 948 | + "for M=%d from %" PRId64 |
| 949 | + " vectors in %dD -> %dD " | ||
| 950 | + "niter=%d niter_pq0=%d niter_pq=%d\n", | ||
| 930 | M_, | 951 | M_, |
| 931 | n, | 952 | n, |
| 932 | d_, | 953 | d_, |
| 933 | - d2_); | 954 | + d2_, |
| 955 | + niter, | ||
| 956 | + niterPq0, | ||
| 957 | + niterPq); | ||
| 934 | } | 958 | } |
| 935 | 959 | ||
| 936 | // 与 CPU OPQMatrix::train 对齐: 训练维度 d = max(d_in, d_out), | 960 | // 与 CPU OPQMatrix::train 对齐: 训练维度 d = max(d_in, d_out), |
| @@ -1016,7 +1040,9 @@ void OPQ::train( | |||
| 1016 | std::vector<float> Umat(d2 * d2); | 1040 | std::vector<float> Umat(d2 * d2); |
| 1017 | std::vector<float> VTmat(d2 * d); | 1041 | std::vector<float> VTmat(d2 * d); |
| 1018 | 1042 | ||
| 1019 | - double t0 = getmillisecs(); | 1043 | + const double t0 = verbose ? getmillisecs() : 0.0; |
| 1044 | + double previousObjective = 0.0; | ||
| 1045 | + bool hasPreviousObjective = false; | ||
| 1020 | for (int iter = 0; iter < niter; iter++) { | 1046 | for (int iter = 0; iter < niter; iter++) { |
| 1021 | // ---- a. 投影: xproj = xtrain @ rotation^T ---- | 1047 | // ---- a. 投影: xproj = xtrain @ rotation^T ---- |
| 1022 | for (int k = 0; k < d; k++) { | 1048 | for (int k = 0; k < d; k++) { |
| @@ -1073,6 +1099,39 @@ void OPQ::train( | |||
| 1073 | } | 1099 | } |
| 1074 | } | 1100 | } |
| 1075 | 1101 | ||
| 1102 | + if (verbose) { | ||
| 1103 | + // This is the same PQ reconstruction error reported by the CPU | ||
| 1104 | + // OPQ implementation, evaluated on the projected vectors. | ||
| 1105 | + const double objective = computeReconstructionError( | ||
| 1106 | + n, d2, xproj.get(), pq_recons.get()); | ||
| 1107 | + const double delta = | ||
| 1108 | + hasPreviousObjective ? objective - previousObjective : 0.0; | ||
| 1109 | + const double relativeChange = hasPreviousObjective | ||
| 1110 | + ? std::abs(delta) / | ||
| 1111 | + std::max(std::abs(previousObjective), 1.0) | ||
| 1112 | + : 0.0; | ||
| 1113 | + const char* status = "initial"; | ||
| 1114 | + if (hasPreviousObjective) { | ||
| 1115 | + status = relativeChange <= OPQ_VERBOSE_CONVERGENCE_TOL | ||
| 1116 | + ? "converged" | ||
| 1117 | + : delta < 0.0 ? "improving" | ||
| 1118 | + : "regressed"; | ||
| 1119 | + } | ||
| 1120 | + previousObjective = objective; | ||
| 1121 | + hasPreviousObjective = true; | ||
| 1122 | + printf(" NpuOPQ::train: iter %d/%d pq_iter=%d " | ||
| 1123 | + "objective=%.6e delta=%.6e rel_change=%.3e " | ||
| 1124 | + "status=%s elapsed=%.3f s\n", | ||
| 1125 | + iter + 1, | ||
| 1126 | + niter, | ||
| 1127 | + iter == 0 ? niterPq0 : niterPq, | ||
| 1128 | + objective, | ||
| 1129 | + delta, | ||
| 1130 | + relativeChange, | ||
| 1131 | + status, | ||
| 1132 | + (getmillisecs() - t0) / 1000.0); | ||
| 1133 | + } | ||
| 1134 | + | ||
| 1076 | // ---- d. xxr = pq_recons^T @ xtrain (d2 * d) ---- | 1135 | // ---- d. xxr = pq_recons^T @ xtrain (d2 * d) ---- |
| 1077 | // 按内积维 n 分批累加: 每次 reconsT_chunk(d2×rows) @ | 1136 | // 按内积维 n 分批累加: 每次 reconsT_chunk(d2×rows) @ |
| 1078 | // xtrain_chunk(rows×d) 得到部分 (d2×d), host 端用 double 累加避免大 n | 1137 | // xtrain_chunk(rows×d) 得到部分 (d2×d), host 端用 double 累加避免大 n |
| @@ -1114,15 +1173,6 @@ void OPQ::train( | |||
| 1114 | } | 1173 | } |
| 1115 | } | 1174 | } |
| 1116 | matmul_(d2, d2, d, Umat.data(), VTmat.data(), A.data()); | 1175 | matmul_(d2, d2, d, Umat.data(), VTmat.data(), A.data()); |
| 1117 | - | ||
| 1118 | - if (verbose) { | ||
| 1119 | - printf(" NpuOPQ::train: iter %d/%d (%d PQ iterations): " | ||
| 1120 | - "%.3f s\n", | ||
| 1121 | - iter, | ||
| 1122 | - niter, | ||
| 1123 | - iter == 0 ? niterPq0 : niterPq, | ||
| 1124 | - (getmillisecs() - t0) / 1000.0); | ||
| 1125 | - } | ||
| 1126 | } | 1176 | } |
| 1127 | 1177 | ||
| 1128 | if (verbose) { | 1178 | if (verbose) { |
| @@ -31,6 +31,7 @@ class AscendcIvfpqCoarseAssign : public OpDef { | |||
| 31 | .DataType({ge::DT_INT32}) | 31 | .DataType({ge::DT_INT32}) |
| 32 | .Format({ge::FORMAT_ND}) | 32 | .Format({ge::FORMAT_ND}) |
| 33 | .UnknownShapeFormat({ge::FORMAT_ND}); | 33 | .UnknownShapeFormat({ge::FORMAT_ND}); |
| 34 | + this->Attr("metric_type").AttrType(REQUIRED).Int(); | ||
| 34 | 35 | ||
| 35 | this->AICore() | 36 | this->AICore() |
| 36 | .AddConfig("ascend910b") | 37 | .AddConfig("ascend910b") |
| @@ -19,6 +19,8 @@ constexpr int32_t kBaseN = 128; | |||
| 19 | constexpr int32_t kBaseCode = 4096; | 19 | constexpr int32_t kBaseCode = 4096; |
| 20 | constexpr int32_t kReduceSize = 64; | 20 | constexpr int32_t kReduceSize = 64; |
| 21 | constexpr int32_t kFloatBytes = 4; | 21 | constexpr int32_t kFloatBytes = 4; |
| 22 | +constexpr int32_t kMetricInnerProduct = 0; | ||
| 23 | +constexpr int32_t kMetricL2 = 1; | ||
| 22 | } // namespace | 24 | } // namespace |
| 23 | 25 | ||
| 24 | namespace optiling { | 26 | namespace optiling { |
| @@ -116,13 +118,19 @@ static ge::graphStatus TilingFunc(gert::TilingContext* context) { | |||
| 116 | if (context == nullptr || context->GetInputShape(0) == nullptr || | 118 | if (context == nullptr || context->GetInputShape(0) == nullptr || |
| 117 | context->GetInputShape(1) == nullptr || | 119 | context->GetInputShape(1) == nullptr || |
| 118 | context->GetInputShape(2) == nullptr || | 120 | context->GetInputShape(2) == nullptr || |
| 119 | - context->GetOutputShape(0) == nullptr) { | 121 | + context->GetOutputShape(0) == nullptr || |
| 122 | + context->GetAttrs() == nullptr) { | ||
| 120 | return ge::GRAPH_FAILED; | 123 | return ge::GRAPH_FAILED; |
| 121 | } | 124 | } |
| 122 | auto* tiling = context->GetTilingData<AscendcIvfpqCoarseAssignTilingData>(); | 125 | auto* tiling = context->GetTilingData<AscendcIvfpqCoarseAssignTilingData>(); |
| 123 | if (tiling == nullptr) { | 126 | if (tiling == nullptr) { |
| 124 | return ge::GRAPH_FAILED; | 127 | return ge::GRAPH_FAILED; |
| 125 | } | 128 | } |
| 129 | + const auto metricType = context->GetAttrs()->GetAttrPointer<int64_t>(0); | ||
| 130 | + if (metricType == nullptr || | ||
| 131 | + (*metricType != kMetricInnerProduct && *metricType != kMetricL2)) { | ||
| 132 | + return ge::GRAPH_FAILED; | ||
| 133 | + } | ||
| 126 | 134 | ||
| 127 | const auto queryShape = context->GetInputShape(0)->GetStorageShape(); | 135 | const auto queryShape = context->GetInputShape(0)->GetStorageShape(); |
| 128 | const auto centroidShape = context->GetInputShape(1)->GetStorageShape(); | 136 | const auto centroidShape = context->GetInputShape(1)->GetStorageShape(); |
| @@ -175,6 +183,7 @@ static ge::graphStatus TilingFunc(gert::TilingContext* context) { | |||
| 175 | tiling->codesTileLength = tiling->baseCodesNumLength; | 183 | tiling->codesTileLength = tiling->baseCodesNumLength; |
| 176 | tiling->codesLastTileLength = tiling->codesNumLength - | 184 | tiling->codesLastTileLength = tiling->codesNumLength - |
| 177 | (tiling->codesTileNum - 1) * tiling->codesTileLength; | 185 | (tiling->codesTileNum - 1) * tiling->codesTileLength; |
| 186 | + tiling->metricType = static_cast<int32_t>(*metricType); | ||
| 178 | SetQuerySplit(*tiling, aivNum, tiling->queryNumLength, queryTile); | 187 | SetQuerySplit(*tiling, aivNum, tiling->queryNumLength, queryTile); |
| 179 | if (SetMatmulTiling(*tiling, l1Size, l0cSize) != ge::GRAPH_SUCCESS) { | 188 | if (SetMatmulTiling(*tiling, l1Size, l0cSize) != ge::GRAPH_SUCCESS) { |
| 180 | return ge::GRAPH_FAILED; | 189 | return ge::GRAPH_FAILED; |
| @@ -19,11 +19,16 @@ using namespace AscendC; | |||
| 19 | 19 | ||
| 20 | namespace AscendC { | 20 | namespace AscendC { |
| 21 | 21 | ||
| 22 | -// Fused exact coarse assignment (argmin over all centroids) for the IVFPQ add | 22 | +// Metric type values of the op attr `metric_type`, mirroring |
| 23 | -// path. The distance matrix D[q][c] = |q - c|^2 is computed on the Cube and | 23 | +// faiss::METRIC_INNER_PRODUCT(0) / faiss::METRIC_L2(1). The kernel cannot |
| 24 | -// reduced in AIV, but never materialized in GM: per 64-centroid segment a | 24 | +// include faiss headers, so keep the mapping explicit here. |
| 25 | -// (min, index) pair is produced by WholeReduceMin, then a per-query running | 25 | +constexpr int32_t kMetricInnerProduct = 0; |
| 26 | -// minimum with the global centroid id is maintained across all code tiles. | 26 | +constexpr int32_t kMetricL2 = 1; |
| 27 | + | ||
| 28 | +// Fused exact coarse assignment for the IVFPQ add path. L2 uses | ||
| 29 | +// |q - c|^2; IP uses -q.c so the same argmin reduction selects the largest | ||
| 30 | +// inner product without norm accumulation or cancellation. The score matrix | ||
| 31 | +// is computed on the Cube and reduced in AIV, but never materialized in GM. | ||
| 27 | // At the end only one int32 centroid label per query is written back. | 32 | // At the end only one int32 centroid label per query is written back. |
| 28 | // | 33 | // |
| 29 | // Tie-break: strictly-less comparison while scanning centroids in increasing | 34 | // Tie-break: strictly-less comparison while scanning centroids in increasing |
| @@ -54,6 +59,7 @@ class AscendcIvfpqCoarseAssign { | |||
| 54 | this->codesNumLength = tiling_data.codesNumLength; | 59 | this->codesNumLength = tiling_data.codesNumLength; |
| 55 | this->dimLength = tiling_data.dimLength; | 60 | this->dimLength = tiling_data.dimLength; |
| 56 | this->bufferSize = tiling_data.bufferSize; | 61 | this->bufferSize = tiling_data.bufferSize; |
| 62 | + this->metricType = tiling_data.metricType; | ||
| 57 | } | 63 | } |
| 58 | 64 | ||
| 59 | __aicore__ inline void InitBuffer( | 65 | __aicore__ inline void InitBuffer( |
| @@ -260,29 +266,42 @@ class AscendcIvfpqCoarseAssign { | |||
| 260 | LocalTensor<T> workspaceLocal = inQueueWorkspace.DeQue<T>(); | 266 | LocalTensor<T> workspaceLocal = inQueueWorkspace.DeQue<T>(); |
| 261 | LocalTensor<T> distResultLocal = distResultQueue.AllocTensor<T>(); | 267 | LocalTensor<T> distResultLocal = distResultQueue.AllocTensor<T>(); |
| 262 | LocalTensor<T> minResultLocal = minResultQueue.AllocTensor<T>(); | 268 | LocalTensor<T> minResultLocal = minResultQueue.AllocTensor<T>(); |
| 263 | - uint32_t nShape[2] = {static_cast<uint32_t>(this->queryLength), 1}; | 269 | + if (this->metricType == kMetricInnerProduct) { |
| 264 | - uint32_t nBroadShape[2] = { | 270 | + Duplicate( |
| 265 | - static_cast<uint32_t>(this->queryLength), | 271 | + distResultLocal, |
| 266 | - static_cast<uint32_t>(this->codesLength)}; | 272 | + this->zero_float, |
| 267 | - BroadCast<T, 2, 1>(distResultLocal, processLocal, nBroadShape, nShape); | 273 | + this->queryLength * this->codesLength); |
| 268 | - PipeBarrier<PIPE_V>(); | 274 | + PipeBarrier<PIPE_V>(); |
| 269 | - for (int32_t i = 0; i < this->queryLength; i++) { | 275 | + Sub(distResultLocal, |
| 270 | - Add(distResultLocal[i * this->codesLength], | 276 | + distResultLocal, |
| 271 | - codesDoubleLocal, | 277 | + workspaceLocal, |
| 272 | - distResultLocal[i * this->codesLength], | 278 | + this->queryLength * this->codesLength); |
| 273 | - this->codesLength); | 279 | + } else { |
| 280 | + uint32_t nShape[2] = {static_cast<uint32_t>(this->queryLength), 1}; | ||
| 281 | + uint32_t nBroadShape[2] = { | ||
| 282 | + static_cast<uint32_t>(this->queryLength), | ||
| 283 | + static_cast<uint32_t>(this->codesLength)}; | ||
| 284 | + BroadCast<T, 2, 1>( | ||
| 285 | + distResultLocal, processLocal, nBroadShape, nShape); | ||
| 286 | + PipeBarrier<PIPE_V>(); | ||
| 287 | + for (int32_t i = 0; i < this->queryLength; i++) { | ||
| 288 | + Add(distResultLocal[i * this->codesLength], | ||
| 289 | + codesDoubleLocal, | ||
| 290 | + distResultLocal[i * this->codesLength], | ||
| 291 | + this->codesLength); | ||
| 292 | + } | ||
| 293 | + PipeBarrier<PIPE_V>(); | ||
| 294 | + Sub(distResultLocal, | ||
| 295 | + distResultLocal, | ||
| 296 | + workspaceLocal, | ||
| 297 | + this->queryLength * this->codesLength); | ||
| 298 | + PipeBarrier<PIPE_V>(); | ||
| 299 | + Sub(distResultLocal, | ||
| 300 | + distResultLocal, | ||
| 301 | + workspaceLocal, | ||
| 302 | + this->queryLength * this->codesLength); | ||
| 274 | } | 303 | } |
| 275 | PipeBarrier<PIPE_V>(); | 304 | PipeBarrier<PIPE_V>(); |
| 276 | - Sub(distResultLocal, | ||
| 277 | - distResultLocal, | ||
| 278 | - workspaceLocal, | ||
| 279 | - this->queryLength * this->codesLength); | ||
| 280 | - PipeBarrier<PIPE_V>(); | ||
| 281 | - Sub(distResultLocal, | ||
| 282 | - distResultLocal, | ||
| 283 | - workspaceLocal, | ||
| 284 | - this->queryLength * this->codesLength); | ||
| 285 | - PipeBarrier<PIPE_V>(); | ||
| 286 | // First stage: one (min, segment-local index) pair per 64 centroids. | 305 | // First stage: one (min, segment-local index) pair per 64 centroids. |
| 287 | const uint32_t minColsPerRow = | 306 | const uint32_t minColsPerRow = |
| 288 | (this->codesLength + this->mask - 1) / this->mask * 2; | 307 | (this->codesLength + this->mask - 1) / this->mask * 2; |
| @@ -362,16 +381,23 @@ class AscendcIvfpqCoarseAssign { | |||
| 362 | 381 | ||
| 363 | __aicore__ inline void CopyLabels(int32_t queryProcess) { | 382 | __aicore__ inline void CopyLabels(int32_t queryProcess) { |
| 364 | LocalTensor<int32_t> runningIdxLocal = runningIdxBuf.Get<int32_t>(); | 383 | LocalTensor<int32_t> runningIdxLocal = runningIdxBuf.Get<int32_t>(); |
| 384 | + // ReduceTile writes runningIdxLocal with scalar SetValue. Synchronize | ||
| 385 | + // that producer with MTE3 before the tail tile is copied to GM. | ||
| 386 | + set_flag(PIPE_S, PIPE_MTE3, EVENT_ID0); | ||
| 387 | + wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID0); | ||
| 365 | DataCopyExtParams labelCopyParams; | 388 | DataCopyExtParams labelCopyParams; |
| 366 | labelCopyParams.blockCount = 1; | 389 | labelCopyParams.blockCount = 1; |
| 367 | labelCopyParams.blockLen = this->queryLength * sizeof(int32_t); | 390 | labelCopyParams.blockLen = this->queryLength * sizeof(int32_t); |
| 368 | labelCopyParams.srcStride = 0; | 391 | labelCopyParams.srcStride = 0; |
| 369 | labelCopyParams.dstStride = 0; | 392 | labelCopyParams.dstStride = 0; |
| 370 | labelCopyParams.rsv = 0; | 393 | labelCopyParams.rsv = 0; |
| 394 | + // queryProcess is indexed by full tiles. The last tile may be shorter, | ||
| 395 | + // but its output starts after all preceding formerTileLength-sized | ||
| 396 | + // tiles. | ||
| 371 | DataCopyPad( | 397 | DataCopyPad( |
| 372 | labelsGlobal | 398 | labelsGlobal |
| 373 | [this->queryCoreOffset + | 399 | [this->queryCoreOffset + |
| 374 | - queryProcess * this->tileLength], | 400 | + queryProcess * this->formerTileLength], |
| 375 | runningIdxLocal, | 401 | runningIdxLocal, |
| 376 | labelCopyParams); | 402 | labelCopyParams); |
| 377 | } | 403 | } |
| @@ -441,6 +467,7 @@ class AscendcIvfpqCoarseAssign { | |||
| 441 | int32_t codesNumLength; | 467 | int32_t codesNumLength; |
| 442 | int32_t dimLength; | 468 | int32_t dimLength; |
| 443 | int32_t bufferSize; | 469 | int32_t bufferSize; |
| 470 | + int32_t metricType; | ||
| 444 | int32_t tileNum; | 471 | int32_t tileNum; |
| 445 | int32_t tileLength; | 472 | int32_t tileLength; |
| 446 | int32_t lastTileLength; | 473 | int32_t lastTileLength; |
| @@ -28,6 +28,7 @@ struct AscendcIvfpqCoarseAssignTilingData { | |||
| 28 | int32_t codesTileLength; | 28 | int32_t codesTileLength; |
| 29 | int32_t codesLastTileLength; | 29 | int32_t codesLastTileLength; |
| 30 | int32_t bufferSize; | 30 | int32_t bufferSize; |
| 31 | + int32_t metricType; | ||
| 31 | }; | 32 | }; |
| 32 | 33 | ||
| 33 | 34 | ||
| @@ -17,7 +17,7 @@ | |||
| 17 | namespace { | 17 | namespace { |
| 18 | constexpr int64_t kQueryNum = 41; | 18 | constexpr int64_t kQueryNum = 41; |
| 19 | constexpr int64_t kNlist = 4160; | 19 | constexpr int64_t kNlist = 4160; |
| 20 | -constexpr int64_t kDim = 128; | 20 | +constexpr int64_t kDim = 1024; |
| 21 | 21 | ||
| 22 | std::vector<int64_t> Strides(const std::vector<int64_t>& shape) { | 22 | std::vector<int64_t> Strides(const std::vector<int64_t>& shape) { |
| 23 | std::vector<int64_t> strides(shape.size(), 1); | 23 | std::vector<int64_t> strides(shape.size(), 1); |
| @@ -160,6 +160,7 @@ int main() { | |||
| 160 | queryTensor.tensor, | 160 | queryTensor.tensor, |
| 161 | centroidTensor.tensor, | 161 | centroidTensor.tensor, |
| 162 | normTensor.tensor, | 162 | normTensor.tensor, |
| 163 | + 1, | ||
| 163 | labelTensor.tensor, | 164 | labelTensor.tensor, |
| 164 | &workspaceSize, | 165 | &workspaceSize, |
| 165 | &executor); | 166 | &executor); |
| @@ -200,6 +201,7 @@ int main() { | |||
| 200 | queryTensor.tensor, | 201 | queryTensor.tensor, |
| 201 | centroidTensor.tensor, | 202 | centroidTensor.tensor, |
| 202 | zeroNormTensor.tensor, | 203 | zeroNormTensor.tensor, |
| 204 | + 0, | ||
| 203 | ipLabelTensor.tensor, | 205 | ipLabelTensor.tensor, |
| 204 | &workspaceSize, | 206 | &workspaceSize, |
| 205 | &executor); | 207 | &executor); |
| @@ -18,6 +18,7 @@ | |||
| 18 | 18 | ||
| 19 | 19 | ||
| 20 | 20 | ||
| 21 | + | ||
| 21 | 22 | ||
| 22 | 23 | ||
| 23 | 24 | ||
| @@ -1251,6 +1252,8 @@ TEST_F(TestNpuIVFPQExtended, TestEndToEndPrecision) { | |||
| 1251 | setTestSeed(12345); | 1252 | setTestSeed(12345); |
| 1252 | std::vector<float> trainVectors = randVecs((size_t)nb_, (size_t)dim_); | 1253 | std::vector<float> trainVectors = randVecs((size_t)nb_, (size_t)dim_); |
| 1253 | std::vector<float> queries = randVecs((size_t)nq_, (size_t)dim_); | 1254 | std::vector<float> queries = randVecs((size_t)nq_, (size_t)dim_); |
| 1255 | + faiss::fvec_renorm_L2(dim_, nb_, trainVectors.data()); | ||
| 1256 | + faiss::fvec_renorm_L2(dim_, nq_, queries.data()); | ||
| 1254 | std::vector<idx_t> ids(nb_); | 1257 | std::vector<idx_t> ids(nb_); |
| 1255 | std::iota(ids.begin(), ids.end(), 0); | 1258 | std::iota(ids.begin(), ids.end(), 0); |
| 1256 | 1259 | ||
| @@ -1268,6 +1271,7 @@ TEST_F(TestNpuIVFPQExtended, TestEndToEndPrecision) { | |||
| 1268 | cpuIndex->cp.seed = 1234; | 1271 | cpuIndex->cp.seed = 1234; |
| 1269 | cpuIndex->cp.nredo = 1; | 1272 | cpuIndex->cp.nredo = 1; |
| 1270 | cpuIndex->pq.cp = cpuIndex->cp; | 1273 | cpuIndex->pq.cp = cpuIndex->cp; |
| 1274 | + cpuIndex->pq.cp.spherical = false; | ||
| 1271 | cpuIndex->by_residual = false; | 1275 | cpuIndex->by_residual = false; |
| 1272 | cpuIndex->train(nb_, trainVectors.data()); | 1276 | cpuIndex->train(nb_, trainVectors.data()); |
| 1273 | cpuIndex->add_with_ids(nb_, trainVectors.data(), ids.data()); | 1277 | cpuIndex->add_with_ids(nb_, trainVectors.data(), ids.data()); |
| @@ -1293,6 +1297,7 @@ TEST_F(TestNpuIVFPQExtended, TestEndToEndPrecision) { | |||
| 1293 | faiss::METRIC_INNER_PRODUCT, | 1297 | faiss::METRIC_INNER_PRODUCT, |
| 1294 | false, | 1298 | false, |
| 1295 | config); | 1299 | config); |
| 1300 | + npuIndex.pq.cp = cpuIndex->pq.cp; | ||
| 1296 | npuIndex.train(nb_, trainVectors.data()); | 1301 | npuIndex.train(nb_, trainVectors.data()); |
| 1297 | ASSERT_TRUE(npuIndex.is_trained); | 1302 | ASSERT_TRUE(npuIndex.is_trained); |
| 1298 | 1303 | ||
| @@ -1361,16 +1366,23 @@ TEST_F(TestNpuIVFPQExtended, TestEndToEndPrecision) { | |||
| 1361 | } | 1366 | } |
| 1362 | } | 1367 | } |
| 1363 | 1368 | ||
| 1364 | - // Search recall | 1369 | + // Search top-k label set overlap |
| 1365 | - int totalMatch = 0; | 1370 | + int overlapTotal = 0; |
| 1366 | for (idx_t i = 0; i < nq_; i++) { | 1371 | for (idx_t i = 0; i < nq_; i++) { |
| 1367 | for (int j = 0; j < k_; j++) { | 1372 | for (int j = 0; j < k_; j++) { |
| 1368 | - if (cpuLabels[i * k_ + j] == npuLabels[i * k_ + j]) { | 1373 | + const idx_t cpuLabel = cpuLabels[i * k_ + j]; |
| 1369 | - totalMatch++; | 1374 | + if (cpuLabel < 0) { |
| 1375 | + continue; | ||
| 1376 | + } | ||
| 1377 | + for (int p = 0; p < k_; p++) { | ||
| 1378 | + if (cpuLabel == npuLabels[i * k_ + p]) { | ||
| 1379 | + overlapTotal++; | ||
| 1380 | + break; | ||
| 1381 | + } | ||
| 1370 | } | 1382 | } |
| 1371 | } | 1383 | } |
| 1372 | } | 1384 | } |
| 1373 | - double recall = 100.0 * totalMatch / (nq_ * k_); | 1385 | + double overlap = 100.0 * overlapTotal / (nq_ * k_); |
| 1374 | 1386 | ||
| 1375 | // Assertions with loose thresholds | 1387 | // Assertions with loose thresholds |
| 1376 | EXPECT_EQ(npuIndex.ntotal, (idx_t)nb_); | 1388 | EXPECT_EQ(npuIndex.ntotal, (idx_t)nb_); |
| @@ -1379,7 +1391,7 @@ TEST_F(TestNpuIVFPQExtended, TestEndToEndPrecision) { | |||
| 1379 | EXPECT_GE((double)listCodesMatch / nonEmptyLists, 0.8); | 1391 | EXPECT_GE((double)listCodesMatch / nonEmptyLists, 0.8); |
| 1380 | EXPECT_GE((double)listIdsMatch / nonEmptyLists, 0.8); | 1392 | EXPECT_GE((double)listIdsMatch / nonEmptyLists, 0.8); |
| 1381 | EXPECT_GE((double)l1Match / nb_, 0.7); | 1393 | EXPECT_GE((double)l1Match / nb_, 0.7); |
| 1382 | - EXPECT_GE(recall, 50.0); | 1394 | + EXPECT_GE(overlap, 50.0); |
| 1383 | } | 1395 | } |
| 1384 | 1396 | ||
| 1385 | // ============================================================ | 1397 | // ============================================================ |
| @@ -15,6 +15,7 @@ | |||
| 15 | 15 | ||
| 16 | 16 | ||
| 17 | 17 | ||
| 18 | + | ||
| 18 | 19 | ||
| 19 | 20 | ||
| 20 | 21 | ||
| @@ -139,6 +140,8 @@ TEST_F(TestNpuOpq, TestEndToEndOpqIvfpqVsCpu) { | |||
| 139 | setTestSeed(12345); | 140 | setTestSeed(12345); |
| 140 | std::vector<float> xb = randVecs((size_t)nb_, (size_t)d_); | 141 | std::vector<float> xb = randVecs((size_t)nb_, (size_t)d_); |
| 141 | std::vector<float> xq = randVecs((size_t)nq_, (size_t)d_); | 142 | std::vector<float> xq = randVecs((size_t)nq_, (size_t)d_); |
| 143 | + faiss::fvec_renorm_L2(d_, nb_, xb.data()); | ||
| 144 | + faiss::fvec_renorm_L2(d_, nq_, xq.data()); | ||
| 142 | std::vector<idx_t> ids(nb_); | 145 | std::vector<idx_t> ids(nb_); |
| 143 | std::iota(ids.begin(), ids.end(), 0); | 146 | std::iota(ids.begin(), ids.end(), 0); |
| 144 | 147 | ||
| @@ -198,6 +201,7 @@ TEST_F(TestNpuOpq, TestEndToEndOpqIvfpqVsCpu) { | |||
| 198 | cpuIndex->cp.seed = 1234; | 201 | cpuIndex->cp.seed = 1234; |
| 199 | cpuIndex->cp.nredo = 1; | 202 | cpuIndex->cp.nredo = 1; |
| 200 | cpuIndex->pq.cp = cpuIndex->cp; | 203 | cpuIndex->pq.cp = cpuIndex->cp; |
| 204 | + cpuIndex->pq.cp.spherical = false; | ||
| 201 | cpuIndex->by_residual = false; | 205 | cpuIndex->by_residual = false; |
| 202 | cpuIndex->train(nb_, xbUsed); | 206 | cpuIndex->train(nb_, xbUsed); |
| 203 | ASSERT_TRUE(cpuIndex->is_trained); | 207 | ASSERT_TRUE(cpuIndex->is_trained); |
| @@ -222,6 +226,7 @@ TEST_F(TestNpuOpq, TestEndToEndOpqIvfpqVsCpu) { | |||
| 222 | faiss::METRIC_INNER_PRODUCT, | 226 | faiss::METRIC_INNER_PRODUCT, |
| 223 | false, | 227 | false, |
| 224 | config); | 228 | config); |
| 229 | + npuIndex.pq.cp = cpuIndex->pq.cp; | ||
| 225 | npuIndex.train(nb_, xbUsed); | 230 | npuIndex.train(nb_, xbUsed); |
| 226 | ASSERT_TRUE(npuIndex.is_trained); | 231 | ASSERT_TRUE(npuIndex.is_trained); |
| 227 | npuIndex.add_with_ids(nb_, xbUsed, ids.data()); | 232 | npuIndex.add_with_ids(nb_, xbUsed, ids.data()); |
| @@ -259,22 +264,29 @@ TEST_F(TestNpuOpq, TestEndToEndOpqIvfpqVsCpu) { | |||
| 259 | (int)nb_); | 264 | (int)nb_); |
| 260 | EXPECT_GE((double)l1Match / nb_, 0.7); | 265 | EXPECT_GE((double)l1Match / nb_, 0.7); |
| 261 | 266 | ||
| 262 | - // ---- 对比: 检索 top-k 标签重合率 ---- | 267 | + // ---- 对比: 检索 top-k 标签集合重合率 ---- |
| 263 | - int matchTotal = 0; | 268 | + int overlapTotal = 0; |
| 264 | for (idx_t i = 0; i < nq_; i++) { | 269 | for (idx_t i = 0; i < nq_; i++) { |
| 265 | for (int j = 0; j < k; j++) { | 270 | for (int j = 0; j < k; j++) { |
| 266 | - if (cpuLabels[i * k + j] == npuLabels[i * k + j]) { | 271 | + const idx_t cpuLabel = cpuLabels[i * k + j]; |
| 267 | - matchTotal++; | 272 | + if (cpuLabel < 0) { |
| 273 | + continue; | ||
| 274 | + } | ||
| 275 | + for (int p = 0; p < k; p++) { | ||
| 276 | + if (cpuLabel == npuLabels[i * k + p]) { | ||
| 277 | + overlapTotal++; | ||
| 278 | + break; | ||
| 279 | + } | ||
| 268 | } | 280 | } |
| 269 | } | 281 | } |
| 270 | } | 282 | } |
| 271 | - double recall = 100.0 * matchTotal / (nq_ * k); | 283 | + double overlap = 100.0 * overlapTotal / (nq_ * k); |
| 272 | - printf("[COMPARE] OPQ+IVFPQ end-to-end: recall@%d = %.1f%% (%d/%d)\n", | 284 | + printf("[COMPARE] OPQ+IVFPQ end-to-end: overlap@%d = %.1f%% (%d/%d)\n", |
| 273 | k, | 285 | k, |
| 274 | - recall, | 286 | + overlap, |
| 275 | - matchTotal, | 287 | + overlapTotal, |
| 276 | (int)(nq_ * k)); | 288 | (int)(nq_ * k)); |
| 277 | - EXPECT_GE(recall, 70.0); | 289 | + EXPECT_GE(overlap, 70.0); |
| 278 | } | 290 | } |
| 279 | 291 | ||
| 280 | // ============================================================ | 292 | // ============================================================ |
| @@ -359,8 +359,8 @@ def prepare_data(d=128, nb=50000, nq=4, seed=12345, use_float16=True): | |||
| 359 | use_float16: True 返回 float16, False 返回 float32 | 359 | use_float16: True 返回 float16, False 返回 float32 |
| 360 | """ | 360 | """ |
| 361 | np.random.seed(seed) | 361 | np.random.seed(seed) |
| 362 | - xb = np.random.random((nb, d)).astype('float32') | 362 | + xb = np.random.standard_normal((nb, d)).astype('float32') |
| 363 | - xq = np.random.random((nq, d)).astype('float32') | 363 | + xq = np.random.standard_normal((nq, d)).astype('float32') |
| 364 | 364 | ||
| 365 | out_dtype = np.float16 if use_float16 else np.float32 | 365 | out_dtype = np.float16 if use_float16 else np.float32 |
| 366 | 366 | ||
| @@ -471,7 +471,7 @@ def configure_ivfpq_training(index, nb, niter, train_seed, nredo, by_residual=Tr | |||
| 471 | index.cp.nredo = nredo | 471 | index.cp.nredo = nredo |
| 472 | index.cp.max_points_per_centroid = nb | 472 | index.cp.max_points_per_centroid = nb |
| 473 | 473 | ||
| 474 | - index.pq.cp.spherical = True | 474 | + index.pq.cp.spherical = False |
| 475 | index.pq.cp.niter = niter | 475 | index.pq.cp.niter = niter |
| 476 | index.pq.cp.seed = train_seed | 476 | index.pq.cp.seed = train_seed |
| 477 | index.pq.cp.nredo = nredo | 477 | index.pq.cp.nredo = nredo |


严重程度: 建议
问题: 在 trainPQCodeBook_ 中调用 indexTrainImpl_ 时,assignmentMetric 参数被硬编码为 faiss::METRIC_L2,而粗量化器训练(train 方法中)使用的是 metric_type。PQ 子量化器训练始终使用 L2 距离进行分配,即使整体索引使用内积度量(metric_type == METRIC_INNER_PRODUCT)。
原因: 如果索引配置为内积度量(如 spherical=true 时),粗量化器训练使用 metric_type(可能是 INNER_PRODUCT)进行分配,但 PQ 子量化器训练却固定使用 L2。这种不一致可能是设计意图(PQ 编码残差时通常使用 L2),但也可能是遗漏。在原始代码中,trainPQCodeBook_ 使用 ivfpqConfig_.cp.spherical 控制 NPU 训练行为,但没有传递 metric type。新代码引入了 assignmentMetric 参数但只在 PQ 训练中固定为 L2,这与粗量化器的 metric_type 传递方式不一致,可能导致内积索引的 PQ 编码质量次优。
怎么改: 如果 PQ 训练确实应始终使用 L2(这是 faiss CPU 实现中的常见做法,因为 PQ 编码的是残差向量),建议在代码中添加注释说明:
// PQ 子量化器始终使用 L2 距离进行聚类分配,因为 PQ 编码的是残差向量, // 残差空间中的距离度量与原始空间的度量类型无关。 faiss::METRIC_L2,如果 PQ 训练应跟随索引的 metric_type,则应将 metric_type 传递给 trainPQCodeBook_ 而非硬编码 L2。