Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix quantizer for UInt8 type #418

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
+156 −47
Open

Fix quantizer for UInt8 type #418

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
08fa091
fix uint8 quantizer training and avx2 distance calculation
Aug 27, 2024
d78dde7
fix UInt8 distance
Sep 4, 2024
9dd823a
add innerproduct PQ/OPQ support
Jan 3, 2025
c7ed3b2
fix include file and tests
Jan 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion AnnService/inc/Core/Common/BKTree.h
View file
Original file line number Diff line number Diff line change
@@ -54,10 +54,12 @@ namespace SPTAG

KmeansArgs(int k, DimensionType dim, SizeType datasize, int threadnum, DistCalcMethod distMethod, const std::shared_ptr<IQuantizer>& quantizer = nullptr) : _K(k), _DK(k), _D(dim), _RD(dim), _T(threadnum), _M(distMethod), m_pQuantizer(quantizer){
if (m_pQuantizer) {
SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "KmeansArgs: Using quantizer!\n");
_RD = m_pQuantizer->ReconstructDim();
fComputeDistance = m_pQuantizer->DistanceCalcSelector<T>(distMethod);
}
else {
SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "KmeansArgs: Using none quantizer!\n");
fComputeDistance = COMMON::DistanceCalcSelector<T>(distMethod);
}

@@ -179,7 +181,7 @@ namespace SPTAG
currCenters[j] /= args.counts[k];
}

if (args._M == DistCalcMethod::Cosine) {
if (args._M != DistCalcMethod::L2) {
COMMON::Utils::Normalize(currCenters, args._RD, COMMON::Utils::GetBase<T>());
}

5 changes: 3 additions & 2 deletions AnnService/inc/Core/Common/OPQQuantizer.h
View file
Original file line number Diff line number Diff line change
@@ -24,7 +24,7 @@ namespace SPTAG

OPQQuantizer();

OPQQuantizer(DimensionType NumSubvectors, SizeType KsPerSubvector, DimensionType DimPerSubvector, bool EnableADC, std::unique_ptr<T[]>&& Codebooks, std::unique_ptr<OPQMatrixType[]>&& OPQMatrix);
OPQQuantizer(DimensionType NumSubvectors, SizeType KsPerSubvector, DimensionType DimPerSubvector, bool EnableADC, std::unique_ptr<T[]>&& Codebooks, std::unique_ptr<OPQMatrixType[]>&& OPQMatrix, DistCalcMethod distMethod);

virtual void QuantizeVector(const void* vec, std::uint8_t* vecout, bool ADC = true) const;

@@ -78,7 +78,8 @@ namespace SPTAG
}

template <typename T>
OPQQuantizer<T>::OPQQuantizer(DimensionType NumSubvectors, SizeType KsPerSubvector, DimensionType DimPerSubvector, bool EnableADC, std::unique_ptr<T[]>&& Codebooks, std::unique_ptr<OPQMatrixType[]>&& OPQMatrix) : m_OPQMatrix(std::move(OPQMatrix)), PQQuantizer<T>::PQQuantizer(NumSubvectors, KsPerSubvector, DimPerSubvector, EnableADC, std::move(Codebooks)), m_matrixDim(NumSubvectors * DimPerSubvector)
OPQQuantizer<T>::OPQQuantizer(DimensionType NumSubvectors, SizeType KsPerSubvector, DimensionType DimPerSubvector, bool EnableADC, std::unique_ptr<T[]>&& Codebooks, std::unique_ptr<OPQMatrixType[]>&& OPQMatrix, DistCalcMethod distMethod) :
m_OPQMatrix(std::move(OPQMatrix)), PQQuantizer<T>::PQQuantizer(NumSubvectors, KsPerSubvector, DimPerSubvector, EnableADC, std::move(Codebooks), distMethod), m_matrixDim(NumSubvectors * DimPerSubvector)
{
m_InitMatrixTranspose();
}
56 changes: 44 additions & 12 deletions AnnService/inc/Core/Common/PQQuantizer.h
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#ifndef _SPTAG_COMMON_PQQUANTIZER_H_
@@ -7,6 +7,7 @@
#include "CommonUtils.h"
#include "DistanceUtils.h"
#include "IQuantizer.h"
#include "inc/Helper/StringConvert.h"
#include <iostream>
#include <fstream>
#include <limits>
@@ -25,7 +26,7 @@ namespace SPTAG
public:
PQQuantizer();

PQQuantizer(DimensionType NumSubvectors, SizeType KsPerSubvector, DimensionType DimPerSubvector, bool EnableADC, std::unique_ptr<T[]>&& Codebooks);
PQQuantizer(DimensionType NumSubvectors, SizeType KsPerSubvector, DimensionType DimPerSubvector, bool EnableADC, std::unique_ptr<T[]>&& Codebooks, DistCalcMethod distMethod);

~PQQuantizer();

@@ -84,6 +85,7 @@ namespace SPTAG
DimensionType m_DimPerSubvector;
SizeType m_BlockSize;
bool m_EnableADC;
DistCalcMethod m_distMethod;

inline SizeType m_DistIndexCalc(SizeType i, SizeType j, SizeType k) const;
void InitializeDistanceTables();
@@ -93,12 +95,13 @@ namespace SPTAG
};

template <typename T>
PQQuantizer<T>::PQQuantizer() : m_NumSubvectors(0), m_KsPerSubvector(0), m_DimPerSubvector(0), m_BlockSize(0), m_EnableADC(false)
PQQuantizer<T>::PQQuantizer() : m_NumSubvectors(0), m_KsPerSubvector(0), m_DimPerSubvector(0), m_BlockSize(0), m_EnableADC(false), m_distMethod(DistCalcMethod::L2)
{
}

template <typename T>
PQQuantizer<T>::PQQuantizer(DimensionType NumSubvectors, SizeType KsPerSubvector, DimensionType DimPerSubvector, bool EnableADC, std::unique_ptr<T[]>&& Codebooks) : m_NumSubvectors(NumSubvectors), m_KsPerSubvector(KsPerSubvector), m_DimPerSubvector(DimPerSubvector), m_BlockSize(KsPerSubvector* KsPerSubvector), m_codebooks(std::move(Codebooks)), m_EnableADC(EnableADC)
PQQuantizer<T>::PQQuantizer(DimensionType NumSubvectors, SizeType KsPerSubvector, DimensionType DimPerSubvector, bool EnableADC, std::unique_ptr<T[]>&& Codebooks, DistCalcMethod distMethod) :
m_NumSubvectors(NumSubvectors), m_KsPerSubvector(KsPerSubvector), m_DimPerSubvector(DimPerSubvector), m_BlockSize(KsPerSubvector* KsPerSubvector), m_codebooks(std::move(Codebooks)), m_EnableADC(EnableADC), m_distMethod(distMethod)
{
InitializeDistanceTables();
}
@@ -131,16 +134,28 @@ namespace SPTAG
float PQQuantizer<T>::CosineDistance(const std::uint8_t* pX, const std::uint8_t* pY) const
// pX must be query distance table for ADC
{
SPTAGLIB_LOG(Helper::LogLevel::LL_Error, "Quantizer does not support CosineDistance!\n");
return 0;
float out = 0;
if (GetEnableADC()) {
float* ptr = (float*)pX;
for (int i = 0; i < m_NumSubvectors; i++) {
out += ptr[pY[i]];
ptr += m_KsPerSubvector;
}
}
else {
for (int i = 0; i < m_NumSubvectors; i++) {
out += m_L2DistanceTables[m_DistIndexCalc(i, pX[i], pY[i])];
}
}
return out;
}

template <typename T>
void PQQuantizer<T>::QuantizeVector(const void* vec, std::uint8_t* vecout, bool ADC) const
{
if (ADC && GetEnableADC())
{
auto distCalc = DistanceCalcSelector<T>(DistCalcMethod::L2);
auto distCalc = COMMON::DistanceCalcSelector<T>(m_distMethod);
float* ADCtable = (float*) vecout;
T* subcodebooks = m_codebooks.get();
T* subvec = (T*)vec;
@@ -157,7 +172,7 @@ namespace SPTAG
}
else
{
auto distCalc = DistanceCalcSelector<T>(DistCalcMethod::L2);
auto distCalc = COMMON::DistanceCalcSelector<T>(DistCalcMethod::L2);
T* subvec = (T*)vec;
T* subcodebooks = m_codebooks.get();
for (int i = 0; i < m_NumSubvectors; i++) {
@@ -220,7 +235,7 @@ namespace SPTAG
std::uint64_t PQQuantizer<T>::BufferSize() const
{
return sizeof(T) * m_NumSubvectors * m_KsPerSubvector * m_DimPerSubvector +
sizeof(DimensionType) + sizeof(SizeType) + sizeof(DimensionType) + sizeof(VectorValueType) + sizeof(QuantizerType);
sizeof(DimensionType) + sizeof(SizeType) + sizeof(DimensionType) + sizeof(VectorValueType) + sizeof(QuantizerType) + sizeof(DistCalcMethod);
}

template <typename T>
@@ -234,6 +249,7 @@ namespace SPTAG
IOBINARY(p_out, WriteBinary, sizeof(SizeType), (char*)&m_KsPerSubvector);
IOBINARY(p_out, WriteBinary, sizeof(DimensionType), (char*)&m_DimPerSubvector);
IOBINARY(p_out, WriteBinary, sizeof(T) * m_NumSubvectors * m_KsPerSubvector * m_DimPerSubvector, (char*)m_codebooks.get());
IOBINARY(p_out, WriteBinary, sizeof(DistCalcMethod), (char*)&m_distMethod);
SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Saving quantizer: Subvectors:%d KsPerSubvector:%d DimPerSubvector:%d\n", m_NumSubvectors, m_KsPerSubvector, m_DimPerSubvector);
return ErrorCode::Success;
}
@@ -252,6 +268,13 @@ namespace SPTAG
SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "sizeof(T): %s.\n", std::to_string(sizeof(T)).c_str());
IOBINARY(p_in, ReadBinary, sizeof(T) * m_NumSubvectors * m_KsPerSubvector * m_DimPerSubvector, (char*)m_codebooks.get());
SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "After read codebooks.\n");

if (p_in->ReadBinary(sizeof(DistCalcMethod), (char*)&m_distMethod) != sizeof(DistCalcMethod)) {
m_distMethod = DistCalcMethod::L2;
}

SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "After read dist: %s.\n", Helper::Convert::ConvertToString(m_distMethod).c_str());


m_BlockSize = m_KsPerSubvector * m_KsPerSubvector;
InitializeDistanceTables();
@@ -275,8 +298,17 @@ namespace SPTAG
m_codebooks = std::make_unique<T[]>(m_NumSubvectors * m_KsPerSubvector * m_DimPerSubvector);
SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "sizeof(T): %s.\n", std::to_string(sizeof(T)).c_str());
std::memcpy(m_codebooks.get(), raw_bytes, sizeof(T) * m_NumSubvectors * m_KsPerSubvector * m_DimPerSubvector);
raw_bytes += sizeof(T) * m_NumSubvectors * m_KsPerSubvector * m_DimPerSubvector;
SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "After read codebooks.\n");


m_distMethod = *(DistCalcMethod*)raw_bytes;
raw_bytes += sizeof(DistCalcMethod);
if (m_distMethod >= DistCalcMethod::Undefined) {
m_distMethod = DistCalcMethod::L2;
raw_bytes -= sizeof(DistCalcMethod);
}
SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "After read dist: %s.\n", Helper::Convert::ConvertToString(m_distMethod).c_str());

m_BlockSize = m_KsPerSubvector * m_KsPerSubvector;
InitializeDistanceTables();
SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Loaded quantizer: Subvectors:%d KsPerSubvector:%d DimPerSubvector:%d\n", m_NumSubvectors, m_KsPerSubvector, m_DimPerSubvector);
@@ -334,13 +366,13 @@ namespace SPTAG
void PQQuantizer<T>::InitializeDistanceTables()
{
auto temp_m_L2DistanceTables = std::make_unique<float[]>(m_BlockSize * m_NumSubvectors);
auto L2Dist = DistanceCalcSelector<T>(DistCalcMethod::L2);
auto distFunc = COMMON::DistanceCalcSelector<T>(m_distMethod);

for (int i = 0; i < m_NumSubvectors; i++) {
SizeType baseIdx = i * m_KsPerSubvector * m_DimPerSubvector;
for (int j = 0; j < m_KsPerSubvector; j++) {
for (int k = 0; k < m_KsPerSubvector; k++) {
temp_m_L2DistanceTables[m_DistIndexCalc(i, j, k)] = L2Dist(&m_codebooks[baseIdx + j * m_DimPerSubvector], &m_codebooks[baseIdx + k * m_DimPerSubvector], m_DimPerSubvector);
temp_m_L2DistanceTables[m_DistIndexCalc(i, j, k)] = distFunc(&m_codebooks[baseIdx + j * m_DimPerSubvector], &m_codebooks[baseIdx + k * m_DimPerSubvector], m_DimPerSubvector);
}
}
}
20 changes: 11 additions & 9 deletions AnnService/inc/Quantizer/Training.h
View file
Original file line number Diff line number Diff line change
@@ -2,6 +2,7 @@
#include <inc/Core/Common/DistanceUtils.h>
#include <inc/Core/Common/IQuantizer.h>
#include <inc/Core/Common/PQQuantizer.h>
#include <inc/Helper/VectorSetReader.h>

#include <memory>
#include <inc/Core/VectorSet.h>
@@ -13,7 +14,7 @@ using namespace SPTAG;
class QuantizerOptions : public Helper::ReaderOptions
{
public:
QuantizerOptions(SizeType trainingSamples, bool debug, float lambda, SPTAG::QuantizerType qtype, std::string qfile, DimensionType qdim, std::string fullvecs, std::string recvecs) : Helper::ReaderOptions(VectorValueType::Float, 0, VectorFileType::TXT, "|", 32), m_trainingSamples(trainingSamples), m_debug(debug), m_KmeansLambda(lambda), m_quantizerType(qtype), m_outputQuantizerFile(qfile), m_quantizedDim(qdim), m_outputFullVecFile(fullvecs), m_outputReconstructVecFile(recvecs)
QuantizerOptions(SizeType trainingSamples, bool debug, float lambda, SPTAG::QuantizerType qtype, SPTAG::DistCalcMethod dm, std::string qfile, DimensionType qdim, std::string fullvecs, std::string recvecs) : Helper::ReaderOptions(VectorValueType::Float, 0, VectorFileType::TXT, "|", 32), m_trainingSamples(trainingSamples), m_debug(debug), m_KmeansLambda(lambda), m_quantizerType(qtype), m_distMethod(dm), m_outputQuantizerFile(qfile), m_quantizedDim(qdim), m_outputFullVecFile(fullvecs), m_outputReconstructVecFile(recvecs)
{
AddRequiredOption(m_inputFiles, "-i", "--input", "Input raw data.");
AddRequiredOption(m_outputFile, "-o", "--output", "Output quantized vectors.");
@@ -29,6 +30,7 @@ class QuantizerOptions : public Helper::ReaderOptions
AddOptionalOption(m_KmeansLambda, "-kml", "--lambda", "Kmeans lambda parameter.");
AddOptionalOption(m_outputFullVecFile, "-ofv", "--output_full", "Output Uncompressed vectors.");
AddOptionalOption(m_outputFullVecFile, "-orv", "--output_reconstruct", "Output reconstructed vectors.");
AddOptionalOption(m_distMethod, "-dist", "--distance_function", "The distance calculation method.");
}

~QuantizerOptions() {}
@@ -53,6 +55,8 @@ class QuantizerOptions : public Helper::ReaderOptions

SPTAG::QuantizerType m_quantizerType;

SPTAG::DistCalcMethod m_distMethod;

bool m_debug;

float m_KmeansLambda;
@@ -73,21 +77,19 @@ std::unique_ptr<T[]> TrainPQQuantizer(std::shared_ptr<QuantizerOptions> options,
#pragma omp parallel for
for (int codebookIdx = 0; codebookIdx < options->m_quantizedDim; codebookIdx++) {
SPTAGLIB_LOG(Helper::LogLevel::LL_Info, "Training Codebook %d.\n", codebookIdx);
auto kargs = COMMON::KmeansArgs<T>(numCentroids, subdim, raw_vectors->Count(), options->m_threadNum, DistCalcMethod::L2, nullptr);
auto dset = COMMON::Dataset<T>(raw_vectors->Count(), subdim, blockRows, raw_vectors->Count());
COMMON::Dataset<T> dset(raw_vectors->Count(), subdim, blockRows, raw_vectors->Count());

for (int vectorIdx = 0; vectorIdx < raw_vectors->Count(); vectorIdx++) {
auto raw_addr = reinterpret_cast<T*>(raw_vectors->GetVector(vectorIdx)) + (codebookIdx * subdim);
auto dset_addr = dset[vectorIdx];
for (int k = 0; k < subdim; k++) {
dset_addr[k] = raw_addr[k];
}
T* raw_addr = reinterpret_cast<T*>(raw_vectors->GetVector(vectorIdx)) + (codebookIdx * subdim);
T* dset_addr = dset[vectorIdx];
std::memcpy(dset_addr, raw_addr, sizeof(T)*subdim);
}

std::vector<SizeType> localindices;
localindices.resize(dset.R());
for (SizeType il = 0; il < localindices.size(); il++) localindices[il] = il;

auto kargs = COMMON::KmeansArgs<T>(numCentroids, subdim, raw_vectors->Count(), options->m_threadNum, options->m_distMethod, nullptr);
auto nclusters = COMMON::KmeansClustering<T>(dset, localindices, 0, dset.R(), kargs, options->m_trainingSamples, options->m_KmeansLambda, options->m_debug, nullptr);

std::vector<SizeType> reverselocalindex;
@@ -120,4 +122,4 @@ std::unique_ptr<T[]> TrainPQQuantizer(std::shared_ptr<QuantizerOptions> options,
}

return codebooks;
}
}
90 changes: 75 additions & 15 deletions AnnService/src/Core/Common/DistanceUtils.cpp
View file
Original file line number Diff line number Diff line change
@@ -416,14 +416,34 @@ float DistanceUtils::ComputeL2Distance_SSE(const std::uint8_t* pX, const std::ui
}
float diff = DIFF128[0] + DIFF128[1] + DIFF128[2] + DIFF128[3];

float c1;
uint8_t a, b;
while (pX < pEnd4) {
float c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
Copy link
Contributor

@PhilipBAdams PhilipBAdams Sep 4, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you share a bit more detail what's the difference in behavior here (floating point quirks?) to me the logic looks the same

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is the AMD cpu compiling issue. When using the original two lines it will cause illegal instruction segmentation fault.

pX++; pY++;
diff += c1 * c1;

a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
pX++; pY++;
diff += c1 * c1;

a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
pX++; pY++;
diff += c1 * c1;

a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
pX++; pY++;
diff += c1 * c1;
}
while (pX < pEnd1) {
float c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
pX++; pY++;
diff += c1 * c1;
}
return diff;
}
@@ -445,14 +465,34 @@ float DistanceUtils::ComputeL2Distance_AVX(const std::uint8_t* pX, const std::ui
}
float diff = DIFF128[0] + DIFF128[1] + DIFF128[2] + DIFF128[3];

float c1;
uint8_t a, b;
while (pX < pEnd4) {
float c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
pX++; pY++;
diff += c1 * c1;

a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
pX++; pY++;
diff += c1 * c1;

a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
pX++; pY++;
diff += c1 * c1;

a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
pX++; pY++;
diff += c1 * c1;
}
while (pX < pEnd1) {
float c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
pX++; pY++;
diff += c1 * c1;
}
return diff;
}
@@ -484,14 +524,34 @@ float DistanceUtils::ComputeL2Distance_AVX512(const std::uint8_t* pX, const std:
}
float diff = DIFF128[0] + DIFF128[1] + DIFF128[2] + DIFF128[3];

float c1;
uint8_t a, b;
while (pX < pEnd4) {
float c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
pX++; pY++;
diff += c1 * c1;

a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
pX++; pY++;
diff += c1 * c1;

a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
pX++; pY++;
diff += c1 * c1;

a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
pX++; pY++;
diff += c1 * c1;
}
while (pX < pEnd1) {
float c1 = ((float)(*pX++) - (float)(*pY++)); diff += c1 * c1;
a = (*pX); b = (*pY);
c1 = a + 0.0f - b;
pX++; pY++;
diff += c1 * c1;
}
return diff;
}
Loading