add Production Quantization

include/AMMBench.h

@@ -142,7 +142,8 @@
 #include <CPPAlgos/EWSCPPAlgo.h>
 #include <CPPAlgos/CoOccurringFDCPPAlgo.h>
 #include <CPPAlgos/BetaCoOFDCPPAlgo.h>
-
+#include <CPPAlgos/ProductQuantizationRaw.h>
+#include <CPPAlgos/ProductQuantizationHash.h>
 #include <CPPAlgos/INT8CPPAlgo.h>
 
 #include <CPPAlgos/TugOfWarCPPAlgo.h>

include/CPPAlgos/ProductQuantizationHash.h

@@ -0,0 +1,57 @@
+//
+// Created by haolan on 25/6/23.
+//
+
+#ifndef INTELLISTREAM_PRODUCTQUANTIZATIONHASH_H
+#define INTELLISTREAM_PRODUCTQUANTIZATIONHASH_H
+#include <CPPAlgos/AbstractCPPAlgo.h>
+
+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<class AMMBench::ProductQuantizationHash> ProductQuantizationHashPtr;
+/**
+ * @ingroup AMMBENCH_CppAlgos
+ * @def newProductQuantizationHashAlgo
+ * @brief (Macro) To creat a new @ref  ProductQuantizationHashAlgounder shared pointer.
+ */
+#define newProductQuantizationHashAlgo std::make_shared<AMMBench::ProductQuantizationHash>
+}
+/**
+ * @}
+ */
+#endif //INTELLISTREAM_PRODUCTQUANTIZATIONHASH_H

include/CPPAlgos/ProductQuantizationRaw.h

@@ -0,0 +1,57 @@
+//
+// Created by haolan on 22/6/23.
+//
+
+#ifndef AMMBENCH_PRODUCTQUANTIZATIONRAW_H
+#define AMMBENCH_PRODUCTQUANTIZATIONRAW_H
+#include <CPPAlgos/AbstractCPPAlgo.h>
+
+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<class AMMBench::ProductQuantizationRaw> ProductQuantizationRawPtr;
+/**
+ * @ingroup AMMBENCH_CppAlgos
+ * @def newProductQuantizationRawAlgo
+ * @brief (Macro) To creat a new @ref  ProductQuantizationRawAlgounder shared pointer.
+ */
+#define newProductQuantizationRawAlgo std::make_shared<AMMBench::ProductQuantizationRaw>
+}
+/**
+ * @}
+ */
+#endif //AMMBENCH_PRODUCTQUANTIZATIONRAW_H

src/CPPAlgos/CMakeLists.txt

@@ -12,5 +12,7 @@ add_sources(
         INT8CPPAlgo.cpp
         TugOfWarCPPAlgo.cpp
         SMPPCACPPAlgo.cpp
+        ProductQuantizationRaw.cpp
+        ProductQuantizationHash.cpp
 )
 

src/CPPAlgos/CPPAlgoTable.cpp

@@ -7,7 +7,8 @@
 #include <CPPAlgos/CRSV2CPPAlgo.h>
 
 #include <CPPAlgos/CountSketchCPPAlgo.h>
-
+#include <CPPAlgos/ProductQuantizationRaw.h>
+#include <CPPAlgos/ProductQuantizationHash.h>
 #include <CPPAlgos/BCRSCPPAlgo.h>
 #include <CPPAlgos/EWSCPPAlgo.h>
 #include <CPPAlgos/CoOccurringFDCPPAlgo.h>
@@ -18,19 +19,20 @@
 #include <CPPAlgos/TugOfWarCPPAlgo.h>
 
 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

src/CPPAlgos/ProductQuantizationHash.cpp

@@ -0,0 +1,92 @@
+//
+// Created by haolan on 25/6/23.
+//
+#include <CPPAlgos/ProductQuantizationHash.h>
+
+int compute_hash_bucket(const std::vector<int>& split_indices, const std::vector<torch::Tensor>& 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<torch::Tensor> A_encoded;
+
+    std::vector<int> 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<torch::Tensor> split_thresholds = {v1, v2, v3, v4};
+
+
+    for (int i = 0; i < A.size(0); ++i) {
+        torch::Tensor a = A[i];
+        std::vector<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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<int>();
+            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<int>& split_indices, const std::vector<torch::Tensor>& 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<float>();
+        int b = (x[j_t].item<float>() >= v) ? 1 : 0;
+        i = 2 * i - 1 + b;
+    }
+    return i;
+}

src/CPPAlgos/ProductQuantizationRaw.cpp

@@ -0,0 +1,65 @@
+//
+// Created by haolan on 22/6/23.
+//
+#include <CPPAlgos/ProductQuantizationRaw.h>
+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<torch::Tensor> A_encoded;
+
+    for (int i = 0; i < A.size(0); ++i) {
+        torch::Tensor a = A[i];
+        std::vector<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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<int>();
+            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);
+}

src/MatrixLoader/BinomialMatrixLoader.cpp

@@ -21,15 +21,18 @@ 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});
 
         // Add the results of the Bernoulli trial to the binomial tensor
         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});
 

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)
-