Parallelize JIT functions

src/galois/_domains/_function.py

@@ -45,6 +45,9 @@ def set_globals(self):
  _SIGNATURE: numba.types.FunctionType
  """The function's Numba signature."""
 
+ _PARALLEL = False
+ """Indicates if parallel processing should be performed."""
+
  implementation: Callable
  """The function's implementation in pure Python."""
 
@@ -78,7 +81,7 @@ def jit(self) -> numba.types.FunctionType:
 
  if key_2 not in self._CACHE[key_1]:
  self.set_globals() # Set the globals once before JIT compiling the function
- self._CACHE[key_1][key_2] = numba.jit(self._SIGNATURE.signature, nopython=True)(self.implementation)
+ self._CACHE[key_1][key_2] = numba.jit(self._SIGNATURE.signature, parallel=self._PARALLEL, nopython=True)(self.implementation)
 
  return self._CACHE[key_1][key_2]
 

src/galois/_domains/_linalg.py

@@ -219,6 +219,7 @@ def set_globals(self):
  MULTIPLY = self.field._multiply.ufunc_call_only
 
  _SIGNATURE = numba.types.FunctionType(int64[:,:](int64[:,:], int64[:,:]))
+ _PARALLEL = True
 
  @staticmethod
  def implementation(A, B):
@@ -228,8 +229,8 @@ def implementation(A, B):
  M, K = A.shape
  K, N = B.shape
  C = np.zeros((M, N), dtype=A.dtype)
- for i in range(M):
- for j in range(N):
+ for i in numba.prange(M): # pylint: disable=not-an-iterable
+ for j in numba.prange(N): # pylint: disable=not-an-iterable
  for k in range(K):
  C[i,j] = ADD(C[i,j], MULTIPLY(A[i,k], B[k,j]))
 

src/galois/_polys/_dense.py

@@ -428,11 +428,12 @@ def set_globals(self):
  MULTIPLY = self.field._multiply.ufunc_call_only
 
  _SIGNATURE = numba.types.FunctionType(int64[:](int64[:], int64[:]))
+ _PARALLEL = True
 
  @staticmethod
  def implementation(coeffs, values):
  y = np.zeros(values.size, dtype=values.dtype)
- for i in range(values.size):
+ for i in numba.prange(values.size): # pylint: disable=not-an-iterable
  y[i] = coeffs[0]
  for j in range(1, coeffs.size):
  y[i] = ADD(coeffs[j], MULTIPLY(y[i], values[i]))