JIT compile polynomial subtraction routine

src/galois/_polys/_dense.py

@@ -74,6 +74,44 @@ def subtract(a: Array, b: Array) -> Array:
  return c
 
 
+class subtract_jit(Function):
+ """
+ Computes polynomial subtraction of two polynomials.
+
+ Algorithm:
+ c(x) = a(x) - b(x)
+ """
+ def __call__(self, a: Array, b: Array) -> Array:
+ verify_isinstance(a, self.field)
+ verify_isinstance(b, self.field)
+ assert a.ndim == 1 and b.ndim == 1
+ dtype = a.dtype
+
+ if self.field.ufunc_mode != "python-calculate":
+ r = self.jit(a.astype(np.int64), b.astype(np.int64))
+ r = r.astype(dtype)
+ else:
+ r = self.python(a.view(np.ndarray), b.view(np.ndarray))
+ r = self.field._view(r)
+
+ return r
+
+ def set_globals(self):
+ # pylint: disable=global-variable-undefined
+ global SUBTRACT
+ SUBTRACT = self.field._subtract.ufunc
+
+ _SIGNATURE = numba.types.FunctionType(int64[:](int64[:], int64[:]))
+
+ @staticmethod
+ def implementation(a, b):
+ dtype = a.dtype
+ c = np.zeros(max(a.size, b.size), dtype=dtype)
+ c[-a.size:] = a
+ c[-b.size:] = SUBTRACT(c[-b.size:], b)
+ return c
+
+
 def multiply(a: Array, b: Array) -> Array:
  """
  c(x) = a(x) * b(x)

src/galois/_polys/_poly.py

@@ -1800,7 +1800,7 @@ def __sub__(self, other: Poly | Array) -> Poly:
  else:
  a = _convert_to_coeffs(self, self.field)
  b = _convert_to_coeffs(other, self.field)
- c = _dense.subtract(a, b)
+ c = _dense.subtract_jit(self.field)(a, b)
  return Poly(c, field=self.field)
 
  def __rsub__(self, other: Poly | Array) -> Poly:
@@ -1820,7 +1820,7 @@ def __rsub__(self, other: Poly | Array) -> Poly:
  else:
  a = _convert_to_coeffs(other, self.field)
  b = _convert_to_coeffs(self, self.field)
- c = _dense.subtract(a, b)
+ c = _dense.subtract_jit(self.field)(a, b)
  return Poly(c, field=self.field)
 
  def __mul__(self, other: Poly | Array | int) -> Poly: