Merge branch 'clean-up' into optimize-sap

python/mops/benchmarks/benchmark.py

@@ -6,7 +6,7 @@
 import numpy as np
 
 
-def benchmark(function, repeats=1000, warmup=10, plot=True):
+def benchmark(function, repeats=100, warmup=10, plot=True):
     gc.disable()
 
     for _ in range(warmup):
@@ -29,6 +29,7 @@ def benchmark(function, repeats=1000, warmup=10, plot=True):
         plt.plot(np.arange(times_array.shape[0]), times_array, ".")
         plt.savefig("benchmark_plot.pdf")
 
+    gc.enable()
     return mean, std
 
 

python/mops/benchmarks/hpe.py

@@ -7,8 +7,8 @@
 np.random.seed(0xDEADBEEF)
 
 A = np.random.rand(1000, 300)
-C = np.random.rand(2000)
-P = np.random.randint(300, size=(2000, 4))
+C = np.random.rand(10000)
+P = np.random.randint(300, size=(10000, 4))
 
 ref_mean, ref_std = benchmark(lambda: ref_hpe(A, C, P))
 mean, std = benchmark(lambda: hpe(A, C, P))

python/mops/benchmarks/opsa.py

@@ -6,10 +6,10 @@
 
 np.random.seed(0xDEADBEEF)
 
-A = np.random.rand(1000, 20)
-B = np.random.rand(1000, 5)
+A = np.random.rand(60000, 13)
+B = np.random.rand(60000, 20)
 
-indices = np.sort(np.random.randint(10, size=(1000,)))
+indices = np.sort(np.random.randint(1000, size=(60000,)))
 
 ref_mean, ref_std = benchmark(lambda: ref_opsa(A, B, indices, np.max(indices) + 1))
 mean, std = benchmark(lambda: opsa(A, B, indices, np.max(indices) + 1))

python/mops/benchmarks/opsax.py

@@ -9,12 +9,12 @@
 np.random.seed(0xDEADBEEF)
 
 
-A = np.random.rand(100, 10)
-R = np.random.rand(100, 5)
-X = np.random.rand(20, 5)
+A = np.random.rand(60000, 13)
+R = np.random.rand(60000, 32)
+X = np.random.rand(1000, 32)
 
-I = np.random.randint(20, size=(100,))
-J = np.random.randint(20, size=(100,))
+I = np.random.randint(1000, size=(60000,))
+J = np.random.randint(1000, size=(60000,))
 
 ref_mean, ref_std = benchmark(lambda: ref_opsax(A, R, X, I, J))
 mean, std = benchmark(lambda: opsax(A, R, X, I, J))

python/mops/benchmarks/sap.py

@@ -6,14 +6,14 @@
 
 np.random.seed(0xDEADBEEF)
 
-A = np.random.rand(1000, 20)
-B = np.random.rand(1000, 6)
-C = np.random.rand(100)
+A = np.random.rand(32000, 13)
+B = np.random.rand(32000, 7)
+C = np.random.rand(900)
 
-P_A = np.random.randint(20, size=(100,))
-P_B = np.random.randint(6, size=(100,))
-n_O = 50
-P_O = np.random.randint(n_O, size=(100,))
+P_A = np.random.randint(13, size=(900,))
+P_B = np.random.randint(7, size=(900,))
+n_O = 100
+P_O = np.random.randint(n_O, size=(900,))
 
 ref_mean, ref_std = benchmark(lambda: ref_sap(A, B, C, P_A, P_B, P_O, n_O))
 mean, std = benchmark(lambda: sap(A, B, C, P_A, P_B, P_O, n_O))

python/mops/benchmarks/sasax.py

@@ -1,28 +1,31 @@
 import numpy as np
 from benchmark import benchmark, format_mean_std
-from mops.reference_implementations import (
-    sparse_accumulation_scatter_add_with_weights as ref_sasax,
-)
 
 from mops import sparse_accumulation_scatter_add_with_weights as sasax
 
+# from mops.reference_implementations import (
+#     sparse_accumulation_scatter_add_with_weights as ref_sasax,
+# )
+
 np.random.seed(0xDEADBEEF)
 
-A = np.random.rand(100, 20)
-R = np.random.rand(100, 200)
-X = np.random.rand(25, 13, 200)
-C = np.random.rand(50)
-I = np.random.randint(25, size=(100,))
-J = np.random.randint(25, size=(100,))
-n_O = 15
-M_1 = np.random.randint(20, size=(50,))
-M_2 = np.random.randint(13, size=(50,))
-M_3 = np.random.randint(n_O, size=(50,))
+A = np.random.rand(60000, 13)
+R = np.random.rand(60000, 32)
+X = np.random.rand(1000, 7, 32)
+C = np.random.rand(900)
+I = np.random.randint(1000, size=(60000,))
+J = np.random.randint(1000, size=(60000,))
+n_O = 100
+M_1 = np.random.randint(13, size=(900,))
+M_2 = np.random.randint(7, size=(900,))
+M_3 = np.random.randint(n_O, size=(900,))
 
-ref_mean, ref_std = benchmark(lambda: ref_sasax(A, R, X, C, I, J, M_1, M_2, M_3, n_O))
+# ref_mean, ref_std = benchmark(
+#     lambda: ref_sasax(A, R, X, C, I, J, M_1, M_2, M_3, n_O), repeats=3, warmup=1
+# )  # very slow
 mean, std = benchmark(lambda: sasax(A, R, X, C, I, J, M_1, M_2, M_3, n_O))
 
-print("Reference implementation:", format_mean_std(ref_mean, ref_std))
+# print("Reference implementation:", format_mean_std(ref_mean, ref_std))
 print("Optimized implementation:", format_mean_std(mean, std))
 
-print("Speed-up:", ref_mean / mean)
+# print("Speed-up:", ref_mean / mean)

python/mops/tests/sasax.py

@@ -23,8 +23,6 @@ def test_sasax():
 
     reference = ref_sasax(A, R, X, C, I, J, M_1, M_2, M_3, n_O)
     actual = sasax(A, R, X, C, I, J, M_1, M_2, M_3, n_O)
-    print(reference)
-    print(actual)
     assert np.allclose(reference, actual)
 
 

setup.py

@@ -20,7 +20,7 @@ def run(self):
         os.makedirs(build_dir, exist_ok=True)
 
         cmake_options = [
-            "-DCMAKE_BUILD_TYPE=release",
+            "-DCMAKE_BUILD_TYPE=Release",
             "-DBUILD_SHARED_LIBS=ON",
             f"-DCMAKE_INSTALL_PREFIX={install_dir}",
         ]