JIT compile polynomial addition routine

src/galois/_polys/_dense.py

@@ -14,17 +14,42 @@
 # pylint: disable=unidiomatic-typecheck
 
 
-def add(a: Array, b: Array) -> Array:
+class add_jit(Function):
     """
-    c(x) = a(x) + b(x)
+    Computes polynomial addition of two polynomials.
+
+    Algorithm:
+        c(x) = a(x) + b(x)
     """
-    field = type(a)
+    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
 
-    c = field.Zeros(max(a.size, b.size))
-    c[-a.size:] = a
-    c[-b.size:] += b
+        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 c
+        return r
+
+    def set_globals(self):
+        # pylint: disable=global-variable-undefined
+        global ADD
+        ADD = self.field._add.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:] = ADD(c[-b.size:], b)
+        return c
 
 
 def negative(a: Array) -> Array:

src/galois/_polys/_poly.py

@@ -1746,7 +1746,7 @@ def __add__(self, other: Poly | Array) -> Poly:
         else:
             a = _convert_to_coeffs(self, self.field)
             b = _convert_to_coeffs(other, self.field)
-            c = _dense.add(a, b)
+            c = _dense.add_jit(self.field)(a, b)
             return Poly(c, field=self.field)
 
     def __radd__(self, other: Poly | Array) -> Poly:
@@ -1766,7 +1766,7 @@ def __radd__(self, other: Poly | Array) -> Poly:
         else:
             a = _convert_to_coeffs(other, self.field)
             b = _convert_to_coeffs(self, self.field)
-            c = _dense.add(a, b)
+            c = _dense.add_jit(self.field)(a, b)
             return Poly(c, field=self.field)
 
     def __neg__(self):