已合并
feat: 对齐CPU训练一致性 #168
lan_xin创建于 18 天前
feat: 对齐CPU训练一致性 #168
已合并
lan_xin创建于 18 天前
13 个文件变更+653-204
@@ -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
Yyihao123415 天前

严重程度: 建议

问题: 在 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。

likedislike
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
Llan_xin18 天前

严重程度: 提示 问题: PQ 子量化器训练 indexTrainImpl_(..., faiss::METRIC_L2, ...) 硬编码 L2 度量,且无注释说明,与 coarse quantizer 使用 metric_type 的做法不一致。 原因: PQ 残差编码始终使用 L2 是 faiss 的标准做法(残差范数最小化),但缺少注释使该设计意图不明显,后续维护者可能误改为跟随索引 metric_type。 怎么改: 补充注释说明设计依据:

// PQ 残差量化始终使用 L2:残差编码以 ||x - c||² 最小化为目标,
// 与索引的 coarse metric 无关(faiss CPU 路径同此约定)。
indexTrainImpl_(..., pq.cp.spherical, faiss::METRIC_L2, trainedCentroids[m]);
likedislike
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 parameters852 // 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 
904std::vector<std::shared_ptr<NpuResources>> NpuIndexIVFPQ::928std::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-means1852// NPU path (useNpuTrain && centroidsOnHost_ trained): reuses the K-means
1780-// centroids with the fused exact coarse-assignment operator. The impl maps1853+// 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.
1783void NpuIndexIVFPQ::addL1_(1855void 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 encoding270 /// Train each sub-quantizer independently on the configured encoding
@@ -13,6 +13,8 @@
13#include <cfloat>13#include <cfloat>
14#include <chrono>14#include <chrono>
15#include <cmath>15#include <cmath>
16+#include <cstdio>
17+#include <cstdlib>
16#include <cstring>18#include <cstring>
17#include <future>19#include <future>
18#include <limits>20#include <limits>
@@ -44,8 +46,12 @@ constexpr int IVF_PQ_BURST_LEN = 64;
44constexpr int IVF_PQ_TRAIN_MAX_BATCH = 1024;46constexpr int IVF_PQ_TRAIN_MAX_BATCH = 1024;
45const int IVF_PQ_CORE_NUM = faiss::npu::NpuSocInfo::getInstance().getCoreNum();47const int IVF_PQ_CORE_NUM = faiss::npu::NpuSocInfo::getInstance().getCoreNum();
46constexpr int IVF_PQ_SIZE_ALIGN = 8;48constexpr 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;
47constexpr int IVF_PQ_FLAG_SIZE = 16;52constexpr int IVF_PQ_FLAG_SIZE = 16;
48constexpr int TOPK_FLAT_ATTR_IDX_COUNT = 9;53constexpr 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)
50constexpr int TOPK_IVFPQ_L3_ATTR_ASC_IDX = 0;56constexpr int TOPK_IVFPQ_L3_ATTR_ASC_IDX = 0;
51constexpr int TOPK_IVFPQ_L3_ATTR_K_IDX = 1;57constexpr 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
Yyihao123415 天前

严重程度: 建议

问题: 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:

static double computeKMeansObjective(...)
{
    ...
    return objective;  // 直接返回 double,无需 cast
}
likedislike
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.
165static void verifyTrainingCoreCount(int device) {219static 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
Yyihao123415 天前

严重程度: 建议

问题: 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 的结果总是被采纳。

likedislike
lan_xin
7 天前 评论:
667+ bestObjective = objective;
668+ bestCentroids = centroidsHost;
669+ }
670+ }
671+ centroidsOut = std::move(bestCentroids);
D
Dduliqiang15 天前

严重程度: 建议

问题: nredo 循环结束后直接 centroidsOut = std::move(bestCentroids);,未处理 bestCentroids 仍为空的情况(单卡与 distributed 两处同构)。

原因:objective 因 float 截断/溢出变成 ±inf,或出现 NaN 时,isBetter 比较失败,bestCentroids 保持默认空向量;后续访问会 UB/崩溃。与 computeKMeansObjective 返回 float 的风险叠加。

怎么改:

if (bestCentroids.empty()) {
    bestCentroids = centroidsHost;  // 至少兜底最后一轮
}
centroidsOut = std::move(bestCentroids);

或保证第一次 redo 无条件采纳;同时建议 objective 全程用 doubletrainKMeansOnNpuDistributed 中对称处理。

likedislike
567}672}
568 673 
569void IVFPQ::trainKMeansOnNpuDistributed(674void 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, &centroidsHost, 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+ &centroidsHost,
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 
762void IVFPQ::assignCentroidOnDevice(965void 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 operator977+ // 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 = &centroidsSqrSum_;997 DeviceVector<float>* centroidNormsDev = &centroidsSqrSum_;
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
Dduliqiang15 天前

严重程度: 严重

问题: encodeResidualsOnDevice 注释声称「强制走 CPU fallback 以规避 dim=1024 task-step bug」,但实际仍只在 dimPerSubQuantizer_ % 8 != 0return false;注释还以残缺的 The 结尾。

原因: 典型配置 dim=1024, M=32dsub=32% 8 == 0,会继续进入下方 AscendcIvfpqResidualEncode 算子路径。注释描述的防护并未生效,且与「对齐 CPU 训练一致性 / 修复残差召回约 2%」的目标相矛盾。

怎么改:

// AscendcIvfpqResidualEncode 存在 task-step bug(vectorCoreNum_ = usedCoreNum_ * 2),
// 会导致 dim=1024 残差召回率降至约 2%。bug 修复前一律走 CPU fallback。
// TODO: 算子修复后恢复 NPU 路径与 dsub%8 判定。
return false;

或保留算子路径,但删除/改写误导性注释,并补全末尾残缺句子。

likedislike
lan_xin
7 天前 评论:
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
Yyihao123415 天前

严重程度: 建议

问题: computeKMeansObjective 在 trainImpl_ 的每次 K-means 迭代中都被调用,但只有最后一次迭代的返回值被使用(用于 redo 间的最优选择)。同样的模式也出现在 trainKMeansOnNpuDistributed 中。

原因: computeKMeansObjective 遍历全部训练向量 × 全部维度,复杂度为 O(totalSize × dim)。在典型配置下(如 niter=25, nredo=3),该函数会被调用 75 次,但仅最后一次的结果有效。前 74 次计算完全浪费 CPU 时间,且在 NPU 训练流程中这段 CPU 计算会阻塞流水线,降低整体训练吞吐。

怎么改: 将 objective 计算移到迭代循环之后,或仅在最后一次迭代时计算:

for (int iter = 0; iter < niter; iter++) {
    // ... runKMeans_, updateCentroids_ ...
    if (iter == niter - 1) {
        objective = computeKMeansObjective(
                dim, totalSize, totalAssigns,
                trainDataHost, centroidsHost, assignmentMetric);
    }
}

同理修改 trainKMeansOnNpuDistributed 中的对应逻辑。

likedislike
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#pragma omp parallel for if (nlist > 8192)1406#pragma omp parallel for if (nlist > 8192)
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
Yyihao123415 天前

严重程度: 提示

问题: updateCentroids_ 中 hassign(聚类计数)的类型从 std::vector 改为 std::vector,同时新增了 workspace.centroidCounts(float)与原有 workspace.clusterCounts(int)并存。

原因: 聚类计数本质上是整数值,使用 float 存储虽然对于典型 nlist 范围内精确无误,但语义上不够清晰。同时 clusterCounts(int)和 centroidCounts(float)并存容易让后续维护者困惑两者的关系和用途。

怎么改: 如果 float 的引入是为了后续支持加权聚类(非整数权重),建议添加注释说明设计意图。如果仅是为了统一类型,考虑直接将 clusterCounts 也改为 float 并合并为同一字段,减少冗余。

likedislike
lan_xin
7 天前 评论:
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 bucket1553 // 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.threadOffsets1572 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#pragma omp parallel for1585#pragma omp parallel for
@@ -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#pragma omp parallel1643#pragma omp parallel
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#pragma omp parallel for1665#pragma omp parallel for
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.
1461void IVFPQ::normL2_(int dim, int nlist, float* data) {1720void IVFPQ::normL2_(int dim, int nlist, float* data) {
1462-#pragma omp parallel for1721+ 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 instance155 /// 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`, pick208 /// Subsample `totalSize` vectors from `x` into `trainDataHost`, pick
196 /// `nlist` initial centroids via faiss::rand_perm (seed + 1, matching209 /// `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-means245+ /// 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;
75constexpr int OPQ_TOPK_ATTR_IDX_COUNT = 9;75constexpr int OPQ_TOPK_ATTR_IDX_COUNT = 9;
76// 与 IVFPQ 分块策略一致, 按 512MB 自适应每批行数, 避免大底库 OOM。76// 与 IVFPQ 分块策略一致, 按 512MB 自适应每批行数, 避免大底库 OOM。
77constexpr size_t OPQ_BATCH_BUDGET_BYTES = 512ULL * 1024 * 1024;77constexpr 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 算子函数指针
80const aclnnDistanceFlatExtremaAtFp32GetWorkspaceSizeFuncType&99const 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 累加避免大 n1137 // 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;
19constexpr int32_t kBaseCode = 4096;19constexpr int32_t kBaseCode = 4096;
20constexpr int32_t kReduceSize = 64;20constexpr int32_t kReduceSize = 64;
21constexpr int32_t kFloatBytes = 4;21constexpr int32_t kFloatBytes = 4;
22+constexpr int32_t kMetricInnerProduct = 0;
23+constexpr int32_t kMetricL2 = 1;
22} // namespace24} // namespace
23 25 
24namespace optiling {26namespace 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 
20namespace AscendC {20namespace AscendC {
21 21 
22-// Fused exact coarse assignment (argmin over all centroids) for the IVFPQ add22+// 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 and23+// 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 a24+// include faiss headers, so keep the mapping explicit here.
25-// (min, index) pair is produced by WholeReduceMin, then a per-query running25+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 increasing34// 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 labelsGlobal398 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#endif // ASCENDC_IVFPQ_COARSE_ASSIGN_TILING_DATA_H34#endif // ASCENDC_IVFPQ_COARSE_ASSIGN_TILING_DATA_H
@@ -17,7 +17,7 @@
17namespace {17namespace {
18constexpr int64_t kQueryNum = 41;18constexpr int64_t kQueryNum = 41;
19constexpr int64_t kNlist = 4160;19constexpr int64_t kNlist = 4160;
20-constexpr int64_t kDim = 128;20+constexpr int64_t kDim = 1024;
21 21 
22std::vector<int64_t> Strides(const std::vector<int64_t>& shape) {22std::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#include <faiss/npu/utils/CopyUtils.h>18#include <faiss/npu/utils/CopyUtils.h>
19#include <faiss/npu/utils/DeviceUtils.h>19#include <faiss/npu/utils/DeviceUtils.h>
20#include <faiss/npu/utils/Float16.h>20#include <faiss/npu/utils/Float16.h>
21+#include <faiss/utils/distances.h>
21 22 
22#include <gtest/gtest.h>23#include <gtest/gtest.h>
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 recall1369+ // 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 thresholds1387 // 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#include <faiss/npu/StandardNpuResources.h>15#include <faiss/npu/StandardNpuResources.h>
16#include <faiss/npu/test/TestUtils.h>16#include <faiss/npu/test/TestUtils.h>
17#include <faiss/npu/utils/DeviceUtils.h>17#include <faiss/npu/utils/DeviceUtils.h>
18+#include <faiss/utils/distances.h>
18 19 
19#include <gtest/gtest.h>20#include <gtest/gtest.h>
20#include <algorithm>21#include <algorithm>
@@ -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 返回 float32359 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.float32365 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 = nredo471 index.cp.nredo = nredo
472 index.cp.max_points_per_centroid = nb472 index.cp.max_points_per_centroid = nb
473 473 
474- index.pq.cp.spherical = True474+ index.pq.cp.spherical = False
475 index.pq.cp.niter = niter475 index.pq.cp.niter = niter
476 index.pq.cp.seed = train_seed476 index.pq.cp.seed = train_seed
477 index.pq.cp.nredo = nredo477 index.pq.cp.nredo = nredo