JIT parallelize convolution

mhostetter · Nov 30, 2022 · 498c3cb · 498c3cb
1 parent 1e1fdd8
commit 498c3cb
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/galois/_domains/_function.py b/src/galois/_domains/_function.py
@@ -126,6 +126,7 @@ def set_globals(self):
  MULTIPLY = self.field._multiply.ufunc
 
  _SIGNATURE = numba.types.FunctionType(int64[:](int64[:], int64[:]))
+ _PARALLEL = True
 
  @staticmethod
  def implementation(a, b):
@@ -149,8 +150,9 @@ def implementation(a, b):
 
  # Fall-back brute force method
  c = np.zeros(a.size + b.size - 1, dtype=dtype)
- for i in range(a.size):
- for j in range(b.size - 1, -1, -1):
+ for i in numba.prange(a.size):
+ for jj in numba.prange(1, b.size):
+ j = b.size - jj
  c[i + j] = ADD(c[i + j], MULTIPLY(a[i], b[j]))
 
  return c