optimizes the connect time when using gpu (#293)

optimizes the connect time when using gpu (#293)

Author: ztqakita

brainpy/connect/base.py

@@ -5,9 +5,10 @@
 
 import jax.numpy as jnp
 import numpy as onp
-
+from jax import config
 from brainpy import tools, math as bm
 from brainpy.errors import ConnectorError
+from brainpy.tools.others import numba_jit, numba_range
 
 __all__ = [
   # the connection types
@@ -24,7 +25,7 @@
   'Connector', 'TwoEndConnector', 'OneEndConnector',
 
   # methods
-  'csr2csc', 'csr2mat', 'mat2csr', 'ij2csr'
+  'csr2csc', 'csr2mat', 'mat2csr', 'ij2csr', 'ij2csr2'
 ]
 
 CONN_MAT = 'conn_mat'
@@ -43,9 +44,8 @@
                            PRE2SYN, POST2SYN,
                            PRE_SLICE, POST_SLICE]
 
-MAT_DTYPE = onp.bool_
-IDX_DTYPE = onp.uint32
-
+MAT_DTYPE = jnp.bool_
+IDX_DTYPE = jnp.uint32
 
 def set_default_dtype(mat_dtype=None, idx_dtype=None):
   """Set the default dtype.
@@ -203,15 +203,15 @@ def _return_by_mat(self, structures, mat, all_data: dict):
 
     require_other_structs = len([s for s in structures if s != CONN_MAT]) > 0
     if require_other_structs:
-      np = onp if isinstance(mat, onp.ndarray) else bm
+      np = jnp if isinstance(mat, jnp.ndarray) else bm
       pre_ids, post_ids = np.where(mat > 0)
       pre_ids = np.asarray(pre_ids, dtype=IDX_DTYPE)
       post_ids = np.asarray(post_ids, dtype=IDX_DTYPE)
       self._return_by_ij(structures, ij=(pre_ids, post_ids), all_data=all_data)
 
   def _return_by_csr(self, structures, csr: tuple, all_data: dict):
     indices, indptr = csr
-    np = onp if isinstance(indices, onp.ndarray) else bm
+    np = jnp if isinstance(indices, jnp.ndarray) else bm
     assert self.pre_num == indptr.size - 1
 
     if (CONN_MAT in structures) and (CONN_MAT not in all_data):
@@ -260,7 +260,10 @@ def _return_by_ij(self, structures, ij: tuple, all_data: dict):
     require_other_structs = len([s for s in structures
                                  if s not in [CONN_MAT, PRE_IDS, POST_IDS]]) > 0
     if require_other_structs:
-      csr = ij2csr(pre_ids, post_ids, self.pre_num)
+      if config.read('jax_platform_name') == "gpu":
+        csr = ij2csr(pre_ids, post_ids, self.pre_num)
+      else:
+        csr = ij2csr2(pre_ids, post_ids, self.pre_num)   
       self._return_by_csr(structures, csr=csr, all_data=all_data)
 
   def make_returns(self, structures, conn_data, csr=None, mat=None, ij=None):
@@ -421,19 +424,19 @@ def _reset_conn(self, pre_size, post_size=None):
 def csr2csc(csr, post_num, data=None):
   """Convert csr to csc."""
   indices, indptr = csr
-  np = onp if isinstance(indices, onp.ndarray) else bm
-  kind = 'quicksort' if isinstance(indices, onp.ndarray) else 'stable'
+  np = jnp if isinstance(indices, jnp.ndarray) else bm
+  # kind = 'quicksort' if isinstance(indices, jnp.ndarray) else 'stable'
 
   pre_ids = np.repeat(np.arange(indptr.size - 1), np.diff(indptr))
 
-  sort_ids = np.argsort(indices, kind=kind)  # to maintain the original order of the elements with the same value
+  sort_ids = np.argsort(indices)  # to maintain the original order of the elements with the same value
   if isinstance(sort_ids, bm.JaxArray):
     sort_ids = sort_ids.value
   pre_ids_new = np.asarray(pre_ids[sort_ids], dtype=IDX_DTYPE)
 
   unique_post_ids, count = np.unique(indices, return_counts=True)
   post_count = np.zeros(post_num, dtype=IDX_DTYPE)
-  post_count[unique_post_ids] = count
+  post_count = post_count.at[unique_post_ids.value if isinstance(unique_post_ids, bm.JaxArray) else unique_post_ids].set(count.value if isinstance(count, bm.JaxArray) else count)
 
   indptr_new = post_count.cumsum()
   indptr_new = np.insert(indptr_new, 0, 0)
@@ -448,14 +451,14 @@ def csr2csc(csr, post_num, data=None):
 
 def mat2csr(dense):
   """convert a dense matrix to (indices, indptr)."""
-  np = onp if isinstance(dense, onp.ndarray) else bm
+  np = jnp if isinstance(dense, jnp.ndarray) else bm
 
   pre_ids, post_ids = np.where(dense > 0)
   pre_num = dense.shape[0]
 
   uni_idx, count = np.unique(pre_ids, return_counts=True)
   pre_count = np.zeros(pre_num, dtype=IDX_DTYPE)
-  pre_count[uni_idx] = count
+  pre_count = pre_count.at[uni_idx.value if isinstance(uni_idx, bm.JaxArray) else uni_idx].set(count.value if isinstance(count, bm.JaxArray) else count)
   indptr = count.cumsum()
   indptr = np.insert(indptr, 0, 0)
 
@@ -465,38 +468,54 @@ def mat2csr(dense):
 def csr2mat(csr, num_pre, num_post):
   """convert (indices, indptr) to a dense matrix."""
   indices, indptr = csr
-  np = onp if isinstance(indices, onp.ndarray) else bm
+  np = jnp if isinstance(indices, jnp.ndarray) else bm
 
   d = np.zeros((num_pre, num_post), dtype=MAT_DTYPE)  # num_pre, num_post
   pre_ids = np.repeat(np.arange(indptr.size - 1), np.diff(indptr))
-  d[pre_ids, indices] = True
+  d = d.at[pre_ids.value if isinstance(pre_ids, bm.JaxArray) else pre_ids, indices.value if isinstance(indices, bm.JaxArray) else indices].set(True)
   return d
 
 
 def ij2mat(ij, num_pre, num_post):
   """convert (indices, indptr) to a dense matrix."""
   pre_ids, post_ids = ij
-  np = onp if isinstance(pre_ids, onp.ndarray) else bm
+  np = jnp if isinstance(pre_ids, jnp.ndarray) else bm
 
   d = np.zeros((num_pre, num_post), dtype=MAT_DTYPE)  # num_pre, num_post
-  d[pre_ids, post_ids] = True
+  d = d.at[pre_ids.value if isinstance(pre_ids, bm.JaxArray) else pre_ids, post_ids.value if isinstance(post_ids, bm.JaxArray) else post_ids].set(True)
   return d
 
 
 def ij2csr(pre_ids, post_ids, num_pre):
-  """convert pre_ids, post_ids to (indices, indptr)."""
-  np = onp if isinstance(pre_ids, onp.ndarray) else bm
-  kind = 'quicksort' if isinstance(pre_ids, onp.ndarray) else 'stable'
-
-  # sorting
-  sort_ids = np.argsort(pre_ids, kind=kind)
+  """convert pre_ids, post_ids to (indices, indptr) when'jax_platform_name' = 'gpu'"""
+  np = jnp if isinstance(pre_ids, jnp.ndarray) else bm
+  sort_ids = np.argsort(pre_ids)
   post_ids = post_ids[sort_ids.value if isinstance(sort_ids, bm.JaxArray) else sort_ids]
-
   indices = post_ids
   unique_pre_ids, pre_count = np.unique(pre_ids, return_counts=True)
-  final_pre_count = np.zeros(num_pre, dtype=IDX_DTYPE)
-  final_pre_count[unique_pre_ids] = pre_count
+  final_pre_count = np.zeros(num_pre, dtype=jnp.uint32)
+  final_pre_count = final_pre_count.at[unique_pre_ids.value if isinstance(unique_pre_ids, bm.JaxArray) else unique_pre_ids].set(pre_count.value if isinstance(pre_count, bm.JaxArray) else pre_count)
   indptr = final_pre_count.cumsum()
   indptr = np.insert(indptr, 0, 0)
 
   return np.asarray(indices, dtype=IDX_DTYPE), np.asarray(indptr, dtype=IDX_DTYPE)
+
+def ij2csr2(pre_ids, post_ids, num_pre):
+  """convert pre_ids, post_ids to (indices, indptr). and use numba for sort function when'jax_platform_name' = 'cpu'"""
+  np = jnp if isinstance(pre_ids, jnp.ndarray) else bm
+  post_ids = onp.asarray(post_ids)
+  unique_pre_ids, pre_count = onp.unique(pre_ids, return_counts=True)
+  final_pre_count = onp.zeros(num_pre, dtype=onp.uint32)
+  final_pre_count[unique_pre_ids] = pre_count
+  indptr = final_pre_count.cumsum()
+  indptr = onp.insert(indptr, 0, 0)
+  @numba_jit (parallel=True, nogil=True)
+  def single_sort(pre_ids,post_ids,indptr):
+    pre_tmp = indptr.copy()
+    indices= onp.zeros((indptr[-1],))
+    for i in numba_range(indptr[-1]):
+      indices[pre_tmp[pre_ids[i]]]=post_ids[i]
+      pre_tmp[pre_ids[i]]+=1
+    return indices
+  indices = single_sort(onp.asarray(pre_ids),onp.asarray(post_ids),indptr)
+  return np.asarray(indices, dtype=IDX_DTYPE), np.asarray(indptr, dtype=IDX_DTYPE)