Add TugOfWar implementation in PyTorch

benchmark/torchscripts/TugOfWar.py

@@ -0,0 +1,69 @@
+import torch
+import time
+import os
+import math
+
+def tug_of_war_mat(m: int, n: int) -> torch.Tensor:
+    e = 1/math.sqrt(m)
+    M = torch.randint(2, (m, n))
+    return e*(2*M - 1)
+
+
+@torch.jit.script
+def TugOfWar(A: torch.Tensor, B: torch.Tensor, l: int):
+    m, n = A.shape
+    n, p = B.shape
+
+    delta = 0.2
+
+    i_iters = int(-math.log(delta))
+    j_iters = int(2*(-math.log(delta) + math.log(-math.log(delta))))
+
+    z = torch.empty((i_iters,))
+    AS = []
+    SB = []
+
+    for i in range(i_iters):
+        S = tug_of_war_mat(l, n)
+        SB.append(S.matmul(B))
+        AS.append(A.matmul(S.T))
+
+        y = torch.empty((j_iters,))
+
+        for j in range(j_iters):
+            Q = tug_of_war_mat(16, p)
+            X = A.matmul(B.matmul(Q.T))
+            X_hat = AS[i].matmul(SB[i].matmul(Q.T))
+            y[j] = torch.norm(X - X_hat)**2
+        z[i] = torch.median(y)
+
+    i_star = torch.argmin(z)
+    return torch.matmul(AS[i_star], SB[i_star])
+
+
+def main():
+    width = 1000
+    A = torch.rand(10000, width)
+    B = torch.rand(width, 5000)
+
+    t = time.time()
+
+    aResult = TugOfWar(A, B, 500)
+    print("approximate: " + str(time.time() - t) + "s")
+
+    print(aResult)
+
+    # exact result
+    t = time.time()
+    eResult = torch.matmul(A, B)
+    print("\nExact: " + str(time.time() - t) + "s")
+
+    print(eResult)
+
+    print("\nerror: " + str(torch.norm(aResult - eResult, p='fro').item()))
+
+    TugOfWar_script = TugOfWar.save("TugOfWar.pt")
+
+
+if __name__ == '__main__':
+    main()

test/CMakeLists.txt

@@ -23,4 +23,5 @@ add_catch_test(sketch_test SystemTest/SketchTest.cpp IntelliStream)
 add_catch_test(crs_test SystemTest/CRSTest.cpp IntelliStream)
 add_catch_test(weighted_cr_test SystemTest/WeightedCRTest.cpp IntelliStream)
 add_catch_test(block_partition_test SystemTest/BlockPartitionTest.cpp IntelliStream)
+add_catch_test(tug_of_war_test SystemTest/TugOfWarTest.cpp IntelliStream)
 

test/SystemTest/TugOfWarTest.cpp

@@ -0,0 +1,56 @@
+#include <vector>
+
+#define CATCH_CONFIG_MAIN
+#include "catch.hpp"
+#include <AMMBench.h>
+using namespace std;
+using namespace INTELLI;
+using namespace torch;
+void runSingleThreadTest(std::string configName) {
+  ConfigMapPtr cfg = newConfigMap();
+  cfg->fromFile(configName);
+  AMMBench::MatrixLoaderTable mLoaderTable;
+  uint64_t sketchDimension;
+  sketchDimension = cfg->tryU64("sketchDimension", 50, true);
+  uint64_t coreBind = cfg->tryU64("coreBind", 0, true);
+  UtilityFunctions::bind2Core((int) coreBind);
+  torch::set_num_threads(1);
+  std::string ptFile = cfg->tryString("ptFile", "torchscripts/FDAMM.pt", true);
+
+  //uint64_t customResultName = cfg->tryU64("customResultName", 0, true);
+  INTELLI_INFO("Place me at core" + to_string(coreBind));
+  INTELLI_INFO(
+      "with sketch" + to_string(sketchDimension));
+  torch::jit::script::Module module;
+  INTELLI_INFO("Try pt file " + ptFile);
+  module = torch::jit::load(ptFile);
+  std::string matrixLoaderTag = cfg->tryString("matrixLoaderTag", "random", true);
+  auto matLoaderPtr = mLoaderTable.findMatrixLoader(matrixLoaderTag);
+  assert(matLoaderPtr);
+  matLoaderPtr->setConfig(cfg);
+  auto A = matLoaderPtr->getA();
+  auto B = matLoaderPtr->getB();
+  /*torch::manual_seed(114514);
+//555
+auto A = torch::rand({(long) aRow, (long) aCol});
+auto B = torch::rand({(long) aCol, (long) bCol});*/
+  INTELLI_INFO("Generation done, conducting...");
+  ThreadPerf pef((int) coreBind);
+  pef.setPerfList();
+  pef.start();
+  auto C =module.forward({A, B, (long) sketchDimension}).toTensor();
+  pef.end();
+  std::string ruName = "default";
+
+  auto resultCsv = pef.resultToConfigMap();
+  resultCsv->toFile(ruName + ".csv");
+  INTELLI_INFO("Done. here is result");
+  std::cout << resultCsv->toString() << endl;
+}
+TEST_CASE("Test the Tug of War", "[short]")
+{
+  int a = 0;
+  runSingleThreadTest("scripts/config_tugOfWar.csv");
+  // place your test here
+  REQUIRE(a == 0);
+}

test/scripts/config_tugOfWar.csv

@@ -0,0 +1,6 @@
+key,value,type
+aRow,100,U64
+aCol,1000,U64
+bCol,500,U64
+sketchDimension,25,U64
+ptFile,torchscripts/CRSV2.pt,String

test/torchscripts/TugOfWar.py

@@ -0,0 +1,69 @@
+import torch
+import time
+import os
+import math
+
+def tug_of_war_mat(m: int, n: int) -> torch.Tensor:
+    e = 1/math.sqrt(m)
+    M = torch.randint(2, (m, n))
+    return e*(2*M - 1)
+
+
+@torch.jit.script
+def TugOfWar(A: torch.Tensor, B: torch.Tensor, l: int):
+    m, n = A.shape
+    n, p = B.shape
+
+    delta = 0.2
+
+    i_iters = int(-math.log(delta))
+    j_iters = int(2*(-math.log(delta) + math.log(-math.log(delta))))
+
+    z = torch.empty((i_iters,))
+    AS = []
+    SB = []
+
+    for i in range(i_iters):
+        S = tug_of_war_mat(l, n)
+        SB.append(S.matmul(B))
+        AS.append(A.matmul(S.T))
+
+        y = torch.empty((j_iters,))
+
+        for j in range(j_iters):
+            Q = tug_of_war_mat(16, p)
+            X = A.matmul(B.matmul(Q.T))
+            X_hat = AS[i].matmul(SB[i].matmul(Q.T))
+            y[j] = torch.norm(X - X_hat)**2
+        z[i] = torch.median(y)
+
+    i_star = torch.argmin(z)
+    return torch.matmul(AS[i_star], SB[i_star])
+
+
+def main():
+    width = 1000
+    A = torch.rand(10000, width)
+    B = torch.rand(width, 5000)
+
+    t = time.time()
+
+    aResult = TugOfWar(A, B, 500)
+    print("approximate: " + str(time.time() - t) + "s")
+
+    print(aResult)
+
+    # exact result
+    t = time.time()
+    eResult = torch.matmul(A, B)
+    print("\nExact: " + str(time.time() - t) + "s")
+
+    print(eResult)
+
+    print("\nerror: " + str(torch.norm(aResult - eResult, p='fro').item()))
+
+    TugOfWar_script = TugOfWar.save("TugOfWar.pt")
+
+
+if __name__ == '__main__':
+    main()