Skip to content

Fix quantizer for UInt8 type #418

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

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Fix quantizer for UInt8 type #418

Show file tree
Hide file tree
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
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down Expand Up @@ -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>());
}

Expand Down
5 changes: 3 additions & 2 deletions AnnService/inc/Core/Common/OPQQuantizer.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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();
}
Expand Down
56 changes: 44 additions & 12 deletions AnnService/inc/Core/Common/PQQuantizer.h
View file
Open in desktop
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_
Expand All @@ -7,6 +7,7 @@
#include "CommonUtils.h"
#include "DistanceUtils.h"
#include "IQuantizer.h"
#include "inc/Helper/StringConvert.h"
#include <iostream>
#include <fstream>
#include <limits>
Expand All @@ -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();

Expand Down Expand Up @@ -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();
Expand All @@ -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();
}
Expand Down Expand Up @@ -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;
Expand All @@ -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++) {
Expand Down Expand Up @@ -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>
Expand All @@ -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;
}
Expand All @@ -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();
Expand All @@ -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);
Expand Down Expand Up @@ -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);
}
}
}
Expand Down
20 changes: 11 additions & 9 deletions AnnService/inc/Quantizer/Training.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand All @@ -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.");
Expand All @@ -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() {}
Expand All @@ -53,6 +55,8 @@ class QuantizerOptions : public Helper::ReaderOptions

SPTAG::QuantizerType m_quantizerType;

SPTAG::DistCalcMethod m_distMethod;

bool m_debug;

float m_KmeansLambda;
Expand All @@ -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;
Expand Down Expand Up @@ -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
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand All @@ -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;
}
Expand Down Expand Up @@ -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;
}
Expand Down
Loading
Loading