diff --git a/benchmark/torchscripts/prototypes.pt b/benchmark/torchscripts/prototypes.pt new file mode 100644 index 00000000..f9c12cb6 Binary files /dev/null and b/benchmark/torchscripts/prototypes.pt differ diff --git a/include/AMMBench.h b/include/AMMBench.h index ce7085fe..56ed9706 100755 --- a/include/AMMBench.h +++ b/include/AMMBench.h @@ -142,7 +142,8 @@ #include #include #include - +#include +#include #include #include diff --git a/include/CPPAlgos/ProductQuantizationHash.h b/include/CPPAlgos/ProductQuantizationHash.h new file mode 100644 index 00000000..cc121dc6 --- /dev/null +++ b/include/CPPAlgos/ProductQuantizationHash.h @@ -0,0 +1,57 @@ +// +// Created by haolan on 25/6/23. +// + +#ifndef INTELLISTREAM_PRODUCTQUANTIZATIONHASH_H +#define INTELLISTREAM_PRODUCTQUANTIZATIONHASH_H +#include + +namespace AMMBench { +/** + * @ingroup AMMBENCH_CppAlgos The algorithms written in c++ + * @{ + */ +/** + * @class ProductQuantizationRawAlgo CPPAlgos/ProductQuantizationRaw.h + * @brief The Product Quantization AMM class of c++ algos, using hash function to find matching prototypes + * + */ + class ProductQuantizationHash : public AMMBench::AbstractCPPAlgo { + public: + ProductQuantizationHash() { + + } + + ~ProductQuantizationHash() { + + } + + /** + * @brief the virtual function provided for outside callers, rewrite in children classes + * @param A the A matrix + * @param B the B matrix + * @param sketchSize the size of sketc or sampling + * @return the output c matrix + */ + virtual torch::Tensor amm(torch::Tensor A, torch::Tensor B, uint64_t sketchSize); + + }; + +/** + * @ingroup AMMBENCH_CppAlgos + * @typedef AbstractMatrixCppAlgoPtr + * @brief The class to describe a shared pointer to @ref ProductQuantizationHashAlgo + + */ + typedef std::shared_ptr ProductQuantizationHashPtr; +/** + * @ingroup AMMBENCH_CppAlgos + * @def newProductQuantizationHashAlgo + * @brief (Macro) To creat a new @ref ProductQuantizationHashAlgounder shared pointer. + */ +#define newProductQuantizationHashAlgo std::make_shared +} +/** + * @} + */ +#endif //INTELLISTREAM_PRODUCTQUANTIZATIONHASH_H diff --git a/include/CPPAlgos/ProductQuantizationRaw.h b/include/CPPAlgos/ProductQuantizationRaw.h new file mode 100644 index 00000000..23342529 --- /dev/null +++ b/include/CPPAlgos/ProductQuantizationRaw.h @@ -0,0 +1,57 @@ +// +// Created by haolan on 22/6/23. +// + +#ifndef AMMBENCH_PRODUCTQUANTIZATIONRAW_H +#define AMMBENCH_PRODUCTQUANTIZATIONRAW_H +#include + +namespace AMMBench { +/** + * @ingroup AMMBENCH_CppAlgos The algorithms written in c++ + * @{ + */ +/** + * @class ProductQuantizationRawAlgo CPPAlgos/ProductQuantizationRaw.h + * @brief The Product Quantization AMM class of c++ algos, using Euclidean distance + * + */ + class ProductQuantizationRaw : public AMMBench::AbstractCPPAlgo { + public: + ProductQuantizationRaw() { + + } + + ~ProductQuantizationRaw() { + + } + + /** + * @brief the virtual function provided for outside callers, rewrite in children classes + * @param A the A matrix + * @param B the B matrix + * @param sketchSize the size of sketc or sampling + * @return the output c matrix + */ + virtual torch::Tensor amm(torch::Tensor A, torch::Tensor B, uint64_t sketchSize); + + }; + +/** + * @ingroup AMMBENCH_CppAlgos + * @typedef AbstractMatrixCppAlgoPtr + * @brief The class to describe a shared pointer to @ref ProductQuantizationRawAlgo + + */ + typedef std::shared_ptr ProductQuantizationRawPtr; +/** + * @ingroup AMMBENCH_CppAlgos + * @def newProductQuantizationRawAlgo + * @brief (Macro) To creat a new @ref ProductQuantizationRawAlgounder shared pointer. + */ +#define newProductQuantizationRawAlgo std::make_shared +} +/** + * @} + */ +#endif //AMMBENCH_PRODUCTQUANTIZATIONRAW_H diff --git a/src/CPPAlgos/CMakeLists.txt b/src/CPPAlgos/CMakeLists.txt index 139ce09d..b31e0a0f 100644 --- a/src/CPPAlgos/CMakeLists.txt +++ b/src/CPPAlgos/CMakeLists.txt @@ -12,5 +12,7 @@ add_sources( INT8CPPAlgo.cpp TugOfWarCPPAlgo.cpp SMPPCACPPAlgo.cpp + ProductQuantizationRaw.cpp + ProductQuantizationHash.cpp ) diff --git a/src/CPPAlgos/CPPAlgoTable.cpp b/src/CPPAlgos/CPPAlgoTable.cpp index 24cb46eb..138bdea8 100644 --- a/src/CPPAlgos/CPPAlgoTable.cpp +++ b/src/CPPAlgos/CPPAlgoTable.cpp @@ -7,7 +7,8 @@ #include #include - +#include +#include #include #include #include @@ -18,19 +19,20 @@ #include namespace AMMBench { - AMMBench::CPPAlgoTable::CPPAlgoTable() { - algoMap["mm"] = newAbstractCPPAlgo(); - algoMap["crs"] = newCRSCPPAlgo(); - algoMap["crsV2"] = newCRSV2CPPAlgo(); - algoMap["countSketch"] = newCountSketchCPPAlgo(); - algoMap["bcrs"] = newBCRSCPPAlgo(); - algoMap["ews"] = newEWSCPPAlgo(); - algoMap["cooFD"] = newCoOccurringFDCPPAlgo(); - algoMap["bcooFD"] = newBetaCoOFDCPPAlgo(); - algoMap["int8"] = newINT8CPPAlgo(); - algoMap["tugOfWar"] = newTugOfWarCPPAlgo(); - algoMap["weighted-cr"] = newWeightedCRCPPAlgo(); - algoMap["smp-pca"] = newSMPPCACPPAlgo(); - } - +AMMBench::CPPAlgoTable::CPPAlgoTable() { + algoMap["mm"] = newAbstractCPPAlgo(); + algoMap["crs"] = newCRSCPPAlgo(); + algoMap["crsV2"] = newCRSV2CPPAlgo(); + algoMap["countSketch"] = newCountSketchCPPAlgo(); + algoMap["bcrs"] = newBCRSCPPAlgo(); + algoMap["ews"] = newEWSCPPAlgo(); + algoMap["cooFD"] = newCoOccurringFDCPPAlgo(); + algoMap["bcooFD"] = newBetaCoOFDCPPAlgo(); + algoMap["int8"] = newINT8CPPAlgo(); + algoMap["tugOfWar"] = newTugOfWarCPPAlgo(); + algoMap["weighted-cr"] = newWeightedCRCPPAlgo(); + algoMap["smp-pca"] = newSMPPCACPPAlgo(); + algoMap["pq-raw"] = newProductQuantizationRawAlgo(); + algoMap["pq-hash"] = newProductQuantizationHashAlgo(); +} } // AMMBench diff --git a/src/CPPAlgos/ProductQuantizationHash.cpp b/src/CPPAlgos/ProductQuantizationHash.cpp new file mode 100644 index 00000000..222989f6 --- /dev/null +++ b/src/CPPAlgos/ProductQuantizationHash.cpp @@ -0,0 +1,92 @@ +// +// Created by haolan on 25/6/23. +// +#include + +int compute_hash_bucket(const std::vector& split_indices, const std::vector& split_thresholds, const torch::Tensor& x); +torch::Tensor AMMBench::ProductQuantizationHash::amm(torch::Tensor A, torch::Tensor B, uint64_t sketchSize) { + const int D = A.size(1); + int C; + if (sketchSize < 50) C = (int) sketchSize; + C = 10; + const int D_c = D / C; + + torch::Tensor prototypes; + + torch::serialize::InputArchive archive; + archive.load_from("torchscripts/prototypes.pt"); + archive.read("prototypes", prototypes); + + std::vector A_encoded; + + std::vector split_indices = {59, 54, 24, 79}; + torch::Tensor v1; + torch::Tensor v2; + torch::Tensor v3; + torch::Tensor v4; + + v1 = torch::tensor({{2.4377}}); + v2 = torch::tensor({{2.0435, 0.0047}}); + v3 = torch::tensor({{2.5292, -0.2861, 0.5683, 0.5683}}); + v4 = torch::tensor({{2.3362, -0.2747, -1.0972, -1.0972, 0.4769, 0.4769, 0.4769, 0.4769}}); + std::vector split_thresholds = {v1, v2, v3, v4}; + + + for (int i = 0; i < A.size(0); ++i) { + torch::Tensor a = A[i]; + std::vector a_encoded; + for (int c = 0; c < C; ++c) { + auto prototypes_c = prototypes[c]; + auto a_subvector = a.slice(0, c * D_c, (c + 1) * D_c); + + auto closest_prototype_index = compute_hash_bucket(split_indices, split_thresholds, a_subvector); + + a_encoded.push_back(torch::tensor(closest_prototype_index).unsqueeze(0)); + } + A_encoded.push_back(torch::cat(a_encoded)); + } + torch::Tensor A_encoded_tensor = torch::stack(A_encoded); + + std::vector tables; + + for (int c = 0; c < C; ++c) { + auto prototypes_c = prototypes[c]; + auto B_subspace = B.slice(0, c * D_c, (c + 1) * D_c); + + std::vector table_c; + for (int i = 0; i < prototypes_c.size(0); ++i) { + auto prototype = prototypes_c[i]; + auto dot_products = prototype.matmul(B_subspace); + table_c.push_back(dot_products); + } + tables.push_back(torch::stack(table_c)); + } + + std::vector result; + + for (int i = 0; i < A_encoded_tensor.size(0); ++i) { + auto a_encoded = A_encoded_tensor[i]; + auto row_sum = torch::zeros({B.size(1)}); + for (int c = 0; c < C; ++c) { + int prototype_index = a_encoded[c].item(); + auto table_c = tables[c]; + auto dot_products = table_c[prototype_index]; + row_sum += dot_products; + } + result.push_back(row_sum); + } + + return torch::stack(result); +} + +int compute_hash_bucket(const std::vector& split_indices, const std::vector& split_thresholds, const torch::Tensor& x) { + int i = 0; + for (int t = 0; t < 4; t++) { + int j_t = split_indices[t]; + torch::Tensor v_t = split_thresholds[t]; + float v = v_t[0][i].item(); + int b = (x[j_t].item() >= v) ? 1 : 0; + i = 2 * i - 1 + b; + } + return i; +} \ No newline at end of file diff --git a/src/CPPAlgos/ProductQuantizationRaw.cpp b/src/CPPAlgos/ProductQuantizationRaw.cpp new file mode 100644 index 00000000..56680fc1 --- /dev/null +++ b/src/CPPAlgos/ProductQuantizationRaw.cpp @@ -0,0 +1,65 @@ +// +// Created by haolan on 22/6/23. +// +#include +torch::Tensor AMMBench::ProductQuantizationRaw::amm(torch::Tensor A, torch::Tensor B, uint64_t sketchSize) { + const int D = A.size(1); + int C; + if (sketchSize < 50) C = (int) sketchSize; + C = 10; + const int D_c = D / C; + + torch::Tensor prototypes; + + torch::serialize::InputArchive archive; + archive.load_from("torchscripts/prototypes.pt"); + archive.read("prototypes", prototypes); + + std::vector A_encoded; + + for (int i = 0; i < A.size(0); ++i) { + torch::Tensor a = A[i]; + std::vector a_encoded; + for (int c = 0; c < C; ++c) { + auto prototypes_c = prototypes[c]; + auto a_subvector = a.slice(0, c * D_c, (c + 1) * D_c); + + auto distances = torch::norm(prototypes_c - a_subvector.expand_as(prototypes_c), 1); + auto closest_prototype_index = torch::argmin(distances); + a_encoded.push_back(closest_prototype_index.unsqueeze(0)); + } + A_encoded.push_back(torch::cat(a_encoded)); + } + torch::Tensor A_encoded_tensor = torch::stack(A_encoded); + + std::vector tables; + + for (int c = 0; c < C; ++c) { + auto prototypes_c = prototypes[c]; + auto B_subspace = B.slice(0, c * D_c, (c + 1) * D_c); + + std::vector table_c; + for (int i = 0; i < prototypes_c.size(0); ++i) { + auto prototype = prototypes_c[i]; + auto dot_products = prototype.matmul(B_subspace); + table_c.push_back(dot_products); + } + tables.push_back(torch::stack(table_c)); + } + + std::vector result; + + for (int i = 0; i < A_encoded_tensor.size(0); ++i) { + auto a_encoded = A_encoded_tensor[i]; + auto row_sum = torch::zeros({B.size(1)}); + for (int c = 0; c < C; ++c) { + int prototype_index = a_encoded[c].item(); + auto table_c = tables[c]; + auto dot_products = table_c[prototype_index]; + row_sum += dot_products; + } + result.push_back(row_sum); + } + + return torch::stack(result); +} \ No newline at end of file diff --git a/src/MatrixLoader/BinomialMatrixLoader.cpp b/src/MatrixLoader/BinomialMatrixLoader.cpp index b6f9d744..9b363e4b 100644 --- a/src/MatrixLoader/BinomialMatrixLoader.cpp +++ b/src/MatrixLoader/BinomialMatrixLoader.cpp @@ -21,7 +21,9 @@ void AMMBench::BinomialMatrixLoader::generateAB() { A = torch::zeros({(long) aRow, (long) aCol}); B = torch::zeros({(long) aCol, (long) bCol}); - for (int i = 0; i < trials; i++) { + + for(uint64_t i = 0; i < trials; i++) { + // Create a tensor filled with random numbers between 0 and 1 torch::Tensor rand_tensor = torch::rand({(long) aRow, (long) aCol}); @@ -29,7 +31,8 @@ void AMMBench::BinomialMatrixLoader::generateAB() { A += (rand_tensor < probability).to(torch::kInt); } - for (int i = 0; i < trials; i++) { + for(uint64_t i = 0; i < trials; i++) { + // Create a tensor filled with random numbers between 0 and 1 torch::Tensor rand_tensor = torch::rand({(long) aCol, (long) bCol}); diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index db9e353f..f7f3f20a 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -29,4 +29,3 @@ add_catch_test(tug_of_war_test SystemTest/TugOfWarTest.cpp IntelliStream) add_catch_test(int8_test SystemTest/INT8Test.cpp IntelliStream) add_catch_test(pq_test SystemTest/PQTest.cpp IntelliStream) add_catch_test(streaming_test SystemTest/StreamingTest.cpp IntelliStream) - diff --git a/test/SystemTest/PQTest.cpp b/test/SystemTest/PQTest.cpp index f008f040..d234908a 100644 --- a/test/SystemTest/PQTest.cpp +++ b/test/SystemTest/PQTest.cpp @@ -1,25 +1,40 @@ + +// +// Created by haolan on 25/6/23. +// #include #define CATCH_CONFIG_MAIN - #include "catch.hpp" #include - -using namespace std; -using namespace INTELLI; -using namespace torch; - -TEST_CASE("Test Load PQ", "[short]") +#include +TEST_CASE("Test PQ", "[short]") +{ + torch::manual_seed(114514); + AMMBench::ProductQuantizationRaw pqRaw; + auto A = torch::rand({1000, 1000}); + auto B = torch::rand({1000, 1000}); + auto realC = torch::matmul(A, B); + auto ammC = pqRaw.amm(A, B, 20); + std::cout << "PQ:" << std::endl; + std::cout << ammC << std::endl; + std::cout << "exact:" << std::endl; + std::cout << realC << std::endl; + double froError = INTELLI::UtilityFunctions::relativeFrobeniusNorm(realC, ammC); + REQUIRE(froError < 0.5); +} +TEST_CASE("Test PQ Hash", "[short]") { - torch::serialize::InputArchive archive; - archive.load_from("torchscripts/PQ/prototypes.pt"); - torch::Tensor prototypes; - archive.read("prototypes", prototypes); - auto pt_size = prototypes.sizes(); - std::cout << "prototype size:" + to_string(pt_size[0]) + "x" + to_string(pt_size[1]) + "x" + to_string(pt_size[2]) - << std::endl; - std::cout << prototypes << endl; - std::cout << "print first one" << endl; - std::cout << prototypes[0] << endl; - std::cout << prototypes[0][0][0] << endl; + torch::manual_seed(114514); + AMMBench::ProductQuantizationHash pqHash; + auto A = torch::rand({1000, 1000}); + auto B = torch::rand({1000, 1000}); + auto realC = torch::matmul(A, B); + auto ammC = pqHash.amm(A, B, 20); + std::cout << "PQ:" << std::endl; + std::cout << ammC << std::endl; + std::cout << "exact:" << std::endl; + std::cout << realC << std::endl; + double froError = INTELLI::UtilityFunctions::relativeFrobeniusNorm(realC, ammC); + REQUIRE(froError < 0.5); } \ No newline at end of file diff --git a/test/torchscripts/prototypes.pt b/test/torchscripts/prototypes.pt new file mode 100644 index 00000000..f9c12cb6 Binary files /dev/null and b/test/torchscripts/prototypes.pt differ