Merge pull request #652 from pydata/finch-benchmark-test

mtsokol · web-flow · commit 18c1596f9cd2 · 2024-04-03T13:46:53.000+02:00
Finch+PyData benchmarks
diff --git a/asv.conf.json b/asv.conf.json
@@ -3,20 +3,54 @@
     "HEAD"
   ],
   "build_command": [],
-  "environment_type": "conda",
+  "environment_type": "virtualenv",
+  "exclude": [
+    {
+      "env": {
+        "SPARSE_BACKEND": "Finch"
+      },
+      "req": {
+        "finch-tensor": null
+      }
+    },
+    {
+      "env": {
+        "SPARSE_BACKEND": "PyData"
+      },
+      "req": {
+        "finch-tensor": ""
+      }
+    }
+  ],
   "install_command": [
     "pip install --no-deps ."
   ],
   "matrix": {
-    "env": {},
+    "env": {
+      "SPARSE_BACKEND": [
+        "Finch",
+        "PyData"
+      ]
+    },
     "env_nobuild": {},
     "req": {
-      "numba": [],
-      "numpy": []
+      "finch-tensor": [
+        "",
+        null
+      ],
+      "numba": [
+        ""
+      ],
+      "numpy": [
+        ""
+      ]
     }
   },
   "project": "sparse",
   "project_url": "https://sparse.pydata.org/",
+  "pythons": [
+    "3.10"
+  ],
   "repo": ".",
   "uninstall_command": [
     "pip uninstall sparse --yes"
diff --git a/benchmarks/benchmark_backends.py b/benchmarks/benchmark_backends.py
@@ -0,0 +1,142 @@
+import sparse
+
+import numpy as np
+
+from .utils import SkipNotImplemented
+
+TIMEOUT: float = 200.0
+BACKEND: sparse.BackendType = sparse.backend_var.get()
+
+
+class Tensordot:
+    timeout = TIMEOUT
+
+    def setup(self):
+        rng = np.random.default_rng(0)
+
+        random_kwargs = {"density": 0.01, "random_state": rng}
+        if sparse.BackendType.PyData == BACKEND:
+            random_kwargs["format"] = "gcxs"
+
+        self.s1 = sparse.random((100, 10), **random_kwargs)
+        self.s2 = sparse.random((100, 100, 10), **random_kwargs)
+
+        if sparse.BackendType.Finch == BACKEND:
+            import finch
+
+            self.s1 = self.s1.to_device(
+                finch.Storage(finch.Dense(finch.SparseList(finch.Element(0.0))), order=self.s1.get_order())
+            )
+            self.s2 = self.s2.to_device(
+                finch.Storage(
+                    finch.Dense(finch.SparseList(finch.SparseList(finch.Element(0.0)))),
+                    order=self.s2.get_order(),
+                )
+            )
+
+        sparse.tensordot(self.s1, self.s2, axes=([0, 1], [0, 2]))  # compilation
+
+    def time_tensordot(self):
+        sparse.tensordot(self.s1, self.s2, axes=([0, 1], [0, 2]))
+
+
+class SpMv:
+    timeout = TIMEOUT
+    # NOTE: https://github.com/willow-ahrens/Finch.jl/issues/488
+    params = [[True, False], [(10, 0.01)]]  # (1000, 0.01), (1_000_000, 1e-05)
+    param_names = ["lazy_mode", "size_and_density"]
+
+    def setup(self, lazy_mode, size_and_density):
+        rng = np.random.default_rng(0)
+        size, density = size_and_density
+
+        random_kwargs = {"density": density, "random_state": rng}
+        if sparse.BackendType.PyData == BACKEND:
+            random_kwargs["format"] = "gcxs"
+
+        self.M = sparse.random((size, size), **random_kwargs)
+        # NOTE: Once https://github.com/willow-ahrens/Finch.jl/issues/487 is fixed change to (size, 1).
+        self.v1 = rng.normal(size=(size, 2))
+        self.v2 = rng.normal(size=(size, 2))
+
+        if sparse.BackendType.Finch == BACKEND:
+            import finch
+
+            self.M = self.M.to_device(
+                finch.Storage(finch.Dense(finch.SparseList(finch.Element(0.0))), order=self.M.get_order())
+            )
+            self.v1 = finch.Tensor(self.v1)
+            self.v2 = finch.Tensor(self.v2)
+            if lazy_mode:
+
+                @sparse.compiled
+                def fn(tns1, tns2, tns3):
+                    return tns1 @ tns2 + tns3
+            else:
+
+                def fn(tns1, tns2, tns3):
+                    return tns1 @ tns2 + tns3
+
+        elif sparse.BackendType.PyData == BACKEND:
+            if lazy_mode:
+                raise SkipNotImplemented("PyData doesn't have lazy mode")
+
+            def fn(tns1, tns2, tns3):
+                return tns1 @ tns2 + tns3
+
+        else:
+            raise Exception(f"Invalid backend: {BACKEND}")
+
+        self.fn = fn
+        self.fn(self.M, self.v1, self.v2)
+
+    def time_spmv(self, lazy_mode, size_and_density):
+        self.fn(self.M, self.v1, self.v2)
+
+
+class Elemwise:
+    timeout = TIMEOUT
+
+    def setup(self):
+        rng = np.random.default_rng(0)
+
+        random_kwargs = {"density": 0.01, "random_state": rng}
+        if sparse.BackendType.PyData == BACKEND:
+            random_kwargs["format"] = "gcxs"
+
+        self.s1 = sparse.random((100, 10), **random_kwargs)
+        self.s2 = sparse.random((100, 100, 10), **random_kwargs)
+
+        if sparse.BackendType.Finch == BACKEND:
+            import finch
+
+            self.s1 = self.s1.to_device(
+                finch.Storage(finch.Dense(finch.SparseList(finch.Element(0.0))), order=self.s1.get_order())
+            )
+            self.s2 = self.s2.to_device(
+                finch.Storage(
+                    finch.Dense(finch.SparseList(finch.SparseList(finch.Element(0.0)))), order=self.s2.get_order()
+                )
+            )
+
+        self.s1 + self.s2
+        self.s1 * self.s2
+
+    def time_add(self):
+        self.s1 + self.s2
+
+    def time_mul(self):
+        self.s1 * self.s2
+
+
+# class SDDMM:
+#     timeout = TIMEOUT
+
+#     def setup():
+#         pass
+
+# class Reductions:
+#     timeout = TIMEOUT
+
+#     def setup():
+#         pass
diff --git a/benchmarks/benchmark_coo.py b/benchmarks/benchmark_coo.py
@@ -2,9 +2,13 @@
 
 import numpy as np
 
+from .utils import skip_if_finch
+
 
 class MatrixMultiplySuite:
     def setup(self):
+        skip_if_finch()
+
         rng = np.random.default_rng(0)
         self.x = sparse.random((100, 100), density=0.01, random_state=rng)
         self.y = sparse.random((100, 100), density=0.01, random_state=rng)
@@ -17,11 +21,14 @@ def time_matmul(self):
 
 class ElemwiseSuite:
     def setup(self):
+        skip_if_finch()
+
         rng = np.random.default_rng(0)
         self.x = sparse.random((100, 100, 100), density=0.01, random_state=rng)
         self.y = sparse.random((100, 100, 100), density=0.01, random_state=rng)
 
         self.x + self.y  # Numba compilation
+        self.x * self.y  # Numba compilation
 
     def time_add(self):
         self.x + self.y
@@ -32,10 +39,15 @@ def time_mul(self):
 
 class ElemwiseBroadcastingSuite:
     def setup(self):
+        skip_if_finch()
+
         rng = np.random.default_rng(0)
         self.x = sparse.random((100, 1, 100), density=0.01, random_state=rng)
         self.y = sparse.random((100, 100), density=0.01, random_state=rng)
 
+        self.x + self.y  # Numba compilation
+        self.x * self.y  # Numba compilation
+
     def time_add(self):
         self.x + self.y
 
@@ -45,6 +57,8 @@ def time_mul(self):
 
 class IndexingSuite:
     def setup(self):
+        skip_if_finch()
+
         rng = np.random.default_rng(0)
         self.index = rng.integers(0, 100, 50)
         self.x = sparse.random((100, 100, 100), density=0.01, random_state=rng)
diff --git a/benchmarks/benchmark_gcxs.py b/benchmarks/benchmark_gcxs.py
@@ -2,9 +2,13 @@
 
 import numpy as np
 
+from .utils import skip_if_finch
+
 
 class MatrixMultiplySuite:
     def setup(self):
+        skip_if_finch()
+
         rng = np.random.default_rng(0)
         self.x = sparse.random((100, 100), density=0.01, format="gcxs", random_state=rng)
         self.y = sparse.random((100, 100), density=0.01, format="gcxs", random_state=rng)
@@ -17,6 +21,8 @@ def time_matmul(self):
 
 class ElemwiseSuite:
     def setup(self):
+        skip_if_finch()
+
         rng = np.random.default_rng(0)
         self.x = sparse.random((100, 100, 100), density=0.01, format="gcxs", random_state=rng)
         self.y = sparse.random((100, 100, 100), density=0.01, format="gcxs", random_state=rng)
@@ -32,6 +38,8 @@ def time_mul(self):
 
 class ElemwiseBroadcastingSuite:
     def setup(self):
+        skip_if_finch()
+
         rng = np.random.default_rng(0)
         self.x = sparse.random((100, 1, 100), density=0.01, format="gcxs", random_state=rng)
         self.y = sparse.random((100, 100), density=0.01, format="gcxs", random_state=rng)
@@ -45,6 +53,8 @@ def time_mul(self):
 
 class IndexingSuite:
     def setup(self):
+        skip_if_finch()
+
         rng = np.random.default_rng(0)
         self.index = rng.integers(0, 100, 50)
         self.x = sparse.random((100, 100, 100), density=0.01, format="gcxs", random_state=rng)
@@ -74,6 +84,8 @@ class DenseMultiplySuite:
     param_names = ["compressed axis", "n_vectors"]
 
     def setup(self, compressed_axis, n_vecs):
+        skip_if_finch()
+
         rng = np.random.default_rng(1337)
         n = 10000
         x = sparse.random((n, n), density=0.001, format="gcxs", random_state=rng).change_compressed_axes(
diff --git a/benchmarks/benchmark_matmul.py b/benchmarks/benchmark_matmul.py
@@ -2,11 +2,15 @@
 
 import numpy as np
 
+from .utils import skip_if_finch
+
 
 class Matmul_Sparse:
     params = (["coo", "gcxs"], [0, 1, None])
 
     def setup(self, p, dens_arg):
+        skip_if_finch()
+
         rng = np.random.default_rng(0)
         self.x = sparse.random((100, 100), density=0.01, format=p, random_state=rng)
         self.y = sparse.random((100, 100), density=0.01, format=p, random_state=rng)
diff --git a/benchmarks/benchmark_tensordot.py b/benchmarks/benchmark_tensordot.py
@@ -2,6 +2,8 @@
 
 import numpy as np
 
+from .utils import skip_if_finch
+
 
 class TensordotSuiteDenseSparse:
     """
@@ -10,6 +12,8 @@ class TensordotSuiteDenseSparse:
     """
 
     def setup(self):
+        skip_if_finch()
+
         rng = np.random.default_rng(0)
         self.n = rng.random((100, 100))
         self.s = sparse.random((100, 100, 100, 100), density=0.01, random_state=rng)
@@ -28,6 +32,8 @@ class TensordotSuiteSparseSparse:
     """
 
     def setup(self):
+        skip_if_finch()
+
         rng = np.random.default_rng(0)
         self.s1 = sparse.random((100, 100), density=0.01, random_state=rng)
         self.s2 = sparse.random((100, 100, 100, 100), density=0.01, random_state=rng)
@@ -46,6 +52,8 @@ class TensordotSuiteSparseDense:
     """
 
     def setup(self):
+        skip_if_finch()
+
         rng = np.random.default_rng(0)
         self.s = sparse.random((100, 100, 100, 100), density=0.01, random_state=rng)
         self.n = rng.random((100, 100))
diff --git a/benchmarks/utils.py b/benchmarks/utils.py
@@ -0,0 +1,8 @@
+import sparse
+
+from asv_runner.benchmarks.mark import SkipNotImplemented
+
+
+def skip_if_finch():
+    if sparse.backend_var.get() == sparse.BackendType.Finch:
+        raise SkipNotImplemented("Finch backend is skipped.")
diff --git a/sparse/finch_backend/__init__.py b/sparse/finch_backend/__init__.py
@@ -22,6 +22,7 @@
     int64,
     int_,
     lazy,
+    matmul,
     multiply,
     negative,
     permute_dims,
@@ -30,6 +31,7 @@
     random,
     subtract,
     sum,
+    tensordot,
     uint,
     uint8,
     uint16,
@@ -56,6 +58,7 @@
     "int64",
     "int_",
     "lazy",
+    "matmul",
     "multiply",
     "negative",
     "permute_dims",
@@ -64,6 +67,7 @@
     "random",
     "subtract",
     "sum",
+    "tensordot",
     "uint",
     "uint8",
     "uint16",
