liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-06-30 01:31:15 +08:00
Code Issues Packages Projects Releases Wiki Activity

Calculate all allocation sizes

Browse Source
This commit is contained in:
Andrew Kane
2022-02-11 03:33:43 -08:00
parent cab103a2e7
commit f59b1184fc
1 changed files with 17 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
src/ivfkmeans.c
View File

@@ -190,10 +190,18 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
bool rjreset;
double dxcx;
double dxc;
Size lowerBoundSize = sizeof(float) * numSamples * numCenters;
/* TODO Add other allocations */
Size totalSize = lowerBoundSize;
/* Calculate allocation sizes */
Size centerCountsSize = sizeof(int) * numCenters;
Size closestCentersSize = sizeof(int) * numSamples;
Size lowerBoundSize = sizeof(float) * numSamples * numCenters;
Size upperBoundSize = sizeof(float) * numSamples;
Size sSize = sizeof(float) * numCenters;
Size halfcdistSize = sizeof(float) * numCenters * numCenters;
Size newcdistSize = sizeof(float) * numCenters;
/* Calculate total size */
Size totalSize = centerCountsSize + closestCentersSize + lowerBoundSize + upperBoundSize + sSize + halfcdistSize + newcdistSize;
/* Check memory requirements */
if (totalSize / 1024 > maintenance_work_mem)
@@ -209,13 +217,13 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
/* Allocate space */
/* Use float instead of double to save memory */
centerCounts = palloc(sizeof(int) * numCenters);
closestCenters = palloc(sizeof(int) * numSamples);
centerCounts = palloc(centerCountsSize);
closestCenters = palloc(closestCentersSize);
lowerBound = palloc_extended(lowerBoundSize, MCXT_ALLOC_HUGE);
upperBound = palloc(sizeof(float) * numSamples);
s = palloc(sizeof(float) * numCenters);
halfcdist = palloc(sizeof(float) * numCenters * numCenters);
newcdist = palloc(sizeof(float) * numCenters);
upperBound = palloc(upperBoundSize);
s = palloc(sSize);
halfcdist = palloc(halfcdistSize);
newcdist = palloc(newcdistSize);
newCenters = VectorArrayInit(numCenters, dimensions);
for (j = 0; j < numCenters; j++)