Improved SetNewCenters [skip ci]

This commit is contained in:
Andrew Kane
2024-04-24 15:56:50 -07:00
parent 558953ca6b
commit fb6683ade7

View File

@@ -308,40 +308,48 @@ SumCenters(VectorArray samples, VectorArray aggCenters, int *closestCenters, Ivf
} }
} }
static void
HalfvecSetNewCenter(const void *v, Vector * aggCenter)
{
HalfVector *newCenter = (HalfVector *) v;
for (int k = 0; k < aggCenter->dim; k++)
newCenter->x[k] = Float4ToHalfUnchecked(aggCenter->x[k]);
}
static void
BitSetNewCenter(const void *v, Vector * aggCenter)
{
VarBit *newCenter = (VarBit *) v;
unsigned char *nx = VARBITS(newCenter);
for (uint32 k = 0; k < VARBITBYTES(newCenter); k++)
nx[k] = 0;
for (int k = 0; k < aggCenter->dim; k++)
nx[k / 8] |= (aggCenter->x[k] > 0.5) << (7 - (k % 8));
}
/* /*
* Set new centers * Set new centers
*/ */
static void static void
SetNewCenters(VectorArray aggCenters, VectorArray newCenters, IvfflatType type) SetNewCenters(VectorArray aggCenters, VectorArray newCenters, IvfflatType type)
{ {
int dimensions = aggCenters->dim; void (*setNewCenter) (const void *v, Vector * aggCenter);
int numCenters = aggCenters->maxlen;
if (type == IVFFLAT_TYPE_HALFVEC) if (type == IVFFLAT_TYPE_HALFVEC)
{ setNewCenter = HalfvecSetNewCenter;
for (int j = 0; j < numCenters; j++)
{
Vector *aggCenter = (Vector *) VectorArrayGet(aggCenters, j);
HalfVector *newCenter = (HalfVector *) VectorArrayGet(newCenters, j);
for (int k = 0; k < dimensions; k++)
newCenter->x[k] = Float4ToHalfUnchecked(aggCenter->x[k]);
}
}
else if (type == IVFFLAT_TYPE_BIT) else if (type == IVFFLAT_TYPE_BIT)
setNewCenter = BitSetNewCenter;
else
return;
for (int j = 0; j < aggCenters->length; j++)
{ {
for (int j = 0; j < numCenters; j++) Vector *aggCenter = (Vector *) VectorArrayGet(aggCenters, j);
{
Vector *aggCenter = (Vector *) VectorArrayGet(aggCenters, j);
VarBit *newCenter = (VarBit *) VectorArrayGet(newCenters, j);
unsigned char *nx = VARBITS(newCenter);
for (uint32 k = 0; k < VARBITBYTES(newCenter); k++) setNewCenter(VectorArrayGet(newCenters, j), aggCenter);
nx[k] = 0;
for (int k = 0; k < dimensions; k++)
nx[k / 8] |= (aggCenter->x[k] > 0.5) << (7 - (k % 8));
}
} }
} }