Update

benchmarks/README.md

@@ -0,0 +1,19 @@
+# Benchmark Tests
+
+This directory contains a series of benchmark tests designed to evaluate performance. The test files are located in the `src/basic` and `src/model` folders, specifically:
+
+- `basic` folder: Contains tests for basic matrix operations, event-driven CSR matrix operations and jitconn matrix operations.
+- `model` folder: Contains tests for complex models, such as `test_COBA.py` and `test_MultipleArea.py`.
+
+
+To run these tests, ensure all necessary dependencies are installed. You can install the dependencies using the following command:
+
+```bash
+pip install -r requirements.txt
+```
+
+or if you are using a CUDA-enabled GPU, you can install the dependencies using the following command:
+
+```bash
+pip install -r requirements_cuda.txt
+```

benchmarks/requirements.txt

@@ -4,4 +4,5 @@ taichi==1.7.0
 brainpylib==0.2.6
 pandas
 matplotlib
-seaborn
+seaborn
+numpy==1.26.4

benchmarks/requirements_cuda.txt

@@ -1,7 +1,9 @@
-jax[cuda12_pip]==0.4.14 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+jax[cuda12_pip]==0.4.14 
 brainpy==2.4.6.post5
 taichi==1.7.0
-brainpylib==0.2.6
+brainpylib-cu12x==0.2.6
 pandas
 matplotlib
-seaborn
+seaborn
+numpy==1.26.4

benchmarks/src/basic/csrmv_taichi_VS_csrmv.py

@@ -12,7 +12,9 @@
 import pandas as pd
 import taichi as ti
 
-bm.set_platform('cpu')
+
+bm.set_platform('gpu')
+
 
 s = [1000, 5000, 10000, 15000, 20000, 25000, 30000]
 p = [0.1, 0.2, 0.3, 0.4, 0.5]

benchmarks/src/model/linear.py

@@ -630,10 +630,11 @@ def __init__(
   def update(self, x):
     if x.ndim == 1:
       if self.method is None:
-        event_csrmv = bm.event.csrmv_taichi
+        return bm.event.csrmv_taichi(self.weight, self.indices, self.indptr, x,
+                            shape=(self.conn.pre_num, self.conn.post_num),
+                            transpose=self.transpose,)[0]
       else:
-        event_csrmv = bm.event.csrmv
-      return event_csrmv(self.weight, self.indices, self.indptr, x,
+        return bm.event.csrmv(self.weight, self.indices, self.indptr, x,
                             shape=(self.conn.pre_num, self.conn.post_num),
                             transpose=self.transpose,)
     elif x.ndim > 1:
@@ -646,10 +647,11 @@ def update(self, x):
 
   def _batch_csrmv(self, x):
     if self.method is None:
-        event_csrmv = bm.event.csrmv_taichi
+        return bm.event.csrmv_taichi(self.weight, self.indices, self.indptr, x,
+                          shape=(self.conn.pre_num, self.conn.post_num),
+                          transpose=self.transpose,)[0]
     else:
-        event_csrmv = bm.event.csrmv
-    return event_csrmv(self.weight, self.indices, self.indptr, x,
+        return bm.event.csrmv(self.weight, self.indices, self.indptr, x,
                           shape=(self.conn.pre_num, self.conn.post_num),
                           transpose=self.transpose,)
 

benchmarks/src/model/test_COBA.py

@@ -17,7 +17,7 @@
 EventJitFPNormalLinear_taichi = partial(EventJitFPNormalLinear, method=None)
 
 
-bm.set_platform('cpu')
+bm.set_platform('gpu')
 bm.disable_gpu_memory_preallocation()
 
 post_nums = [400, 4000, 40000, 400000, 4000000]

benchmarks/src/model/test_MultipleArea.py

@@ -15,7 +15,7 @@
 from multiple_area_customized_op import (multiple_area_customized_op_E as syn_fun_E,
                                           multiple_area_customized_op_I as syn_fun_I)
 
-bm.set_platform('cpu')
+bm.set_platform('gpu')
 bm.disable_gpu_memory_preallocation()
 
 PATH = os.path.dirname(os.path.abspath(__file__))