liuyu/pgvector
1
0
Fork 0
mirror of https://github.com/pgvector/pgvector.git synced 2026-07-12 07:36:55 +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; bool rjreset;
double dxcx; double dxcx;
double dxc; double dxc;
Size lowerBoundSize = sizeof(float) * numSamples * numCenters;
/* TODO Add other allocations */ /* Calculate allocation sizes */
Size totalSize = lowerBoundSize; 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 */ /* Check memory requirements */
if (totalSize / 1024 > maintenance_work_mem) if (totalSize / 1024 > maintenance_work_mem)
@@ -209,13 +217,13 @@ ElkanKmeans(Relation index, VectorArray samples, VectorArray centers)
/* Allocate space */ /* Allocate space */
/* Use float instead of double to save memory */ /* Use float instead of double to save memory */
centerCounts = palloc(sizeof(int) * numCenters); centerCounts = palloc(centerCountsSize);
closestCenters = palloc(sizeof(int) * numSamples); closestCenters = palloc(closestCentersSize);
lowerBound = palloc_extended(lowerBoundSize, MCXT_ALLOC_HUGE); lowerBound = palloc_extended(lowerBoundSize, MCXT_ALLOC_HUGE);
upperBound = palloc(sizeof(float) * numSamples); upperBound = palloc(upperBoundSize);
s = palloc(sizeof(float) * numCenters); s = palloc(sSize);
halfcdist = palloc(sizeof(float) * numCenters * numCenters); halfcdist = palloc(halfcdistSize);
newcdist = palloc(sizeof(float) * numCenters); newcdist = palloc(newcdistSize);
newCenters = VectorArrayInit(numCenters, dimensions); newCenters = VectorArrayInit(numCenters, dimensions);
for (j = 0; j < numCenters; j++) for (j = 0; j < numCenters; j++)