Minor improvements to GF arithmetic bloqs (#1705)

Some minor fixes:
- Use `GF2MulViaKaratsuba` instead of `GF2Multiplication` in
`GF2Inverse` to reduce the cost of inversion
- Add support for symbolic gate counts to `GF2MulViaKaratsuba` 
- Reduce the number of GF2 multiplications in `GF2Inverse` by 1 since
the last multiplication was not really needed.
- Reduce the number of junk registers in `GF2Inverse` by 1 since
`junk[0]` was same as the input `x`.

Author: tanujkhattar

qualtran/bloqs/gf_arithmetic/gf2_inverse.py

@@ -28,7 +28,7 @@
     Signature,
 )
 from qualtran.bloqs.gf_arithmetic.gf2_addition import GF2Addition
-from qualtran.bloqs.gf_arithmetic.gf2_multiplication import GF2Multiplication
+from qualtran.bloqs.gf_arithmetic.gf2_multiplication import GF2MulViaKaratsuba
 from qualtran.bloqs.gf_arithmetic.gf2_square import GF2Square
 from qualtran.resource_counting.generalizers import ignore_alloc_free, ignore_split_join
 from qualtran.symbolics import bit_length, ceil, is_symbolic, log2, SymbolicInt
@@ -109,7 +109,7 @@ def qgf(self) -> QGF:
 
     @cached_property
     def n_junk_regs(self) -> SymbolicInt:
-        return 2 * bit_length(self.bitsize - 1) + self.bitsize_hamming_weight
+        return 2 * bit_length(self.bitsize - 1) + self.bitsize_hamming_weight - 3
 
     @cached_property
     def bitsize_hamming_weight(self) -> SymbolicInt:
@@ -139,11 +139,11 @@ def build_composite_bloq(self, bb: 'BloqBuilder', *, x: 'Soquet') -> Dict[str, '
             return {'x': x, 'result': result}
 
         junk = []
-        beta = bb.allocate(dtype=self.qgf)
-        x, beta = bb.add(GF2Addition(self.bitsize), x=x, y=beta)
+        beta = x
         is_first = True
         bitsize_minus_one = int(self.bitsize - 1)
-        for i in range(bitsize_minus_one.bit_length()):
+        n_iters = bitsize_minus_one.bit_length()
+        for i in range(n_iters):
             if (1 << i) & bitsize_minus_one:
                 if is_first:
                     beta, result = bb.add(GF2Addition(self.bitsize), x=beta, y=result)
@@ -152,21 +152,23 @@ def build_composite_bloq(self, bb: 'BloqBuilder', *, x: 'Soquet') -> Dict[str, '
                     for j in range(2**i):
                         result = bb.add(GF2Square(self.bitsize), x=result)
                     beta, result, new_result = bb.add(
-                        GF2Multiplication(self.bitsize), x=beta, y=result
+                        GF2MulViaKaratsuba(self.bitsize), x=beta, y=result
                     )
                     junk.append(result)
                     result = new_result
-            beta_squared = bb.allocate(dtype=self.qgf)
-            beta, beta_squared = bb.add(GF2Addition(self.bitsize), x=beta, y=beta_squared)
-            for j in range(2**i):
-                beta_squared = bb.add(GF2Square(self.bitsize), x=beta_squared)
-            beta, beta_squared, beta_new = bb.add(
-                GF2Multiplication(self.bitsize), x=beta, y=beta_squared
-            )
-            junk.extend([beta, beta_squared])
-            beta = beta_new
+            if i != n_iters - 1:
+                beta_squared = bb.allocate(dtype=self.qgf)
+                beta, beta_squared = bb.add(GF2Addition(self.bitsize), x=beta, y=beta_squared)
+                for j in range(2**i):
+                    beta_squared = bb.add(GF2Square(self.bitsize), x=beta_squared)
+                beta, beta_squared, beta_new = bb.add(
+                    GF2MulViaKaratsuba(self.bitsize), x=beta, y=beta_squared
+                )
+                junk.extend([beta, beta_squared])
+                beta = beta_new
         junk.append(beta)
         result = bb.add(GF2Square(self.bitsize), x=result)
+        x = junk.pop(0)
         assert len(junk) == self.n_junk_regs, f'{len(junk)=}, {self.n_junk_regs=}'
         return {'x': x, 'result': result, 'junk': np.array(junk)}
 
@@ -179,13 +181,12 @@ def build_call_graph(
         if not is_symbolic(self.bitsize):
             n = self.bitsize - 1
             square_count -= n & (-n)
+            square_count -= 1 << (n.bit_length() - 1)
+        mul_count = ceil(log2(self.bitsize)) + self.bitsize_hamming_weight - 2
         return {
-            GF2Addition(self.bitsize): 2 + ceil(log2(self.bitsize)),
+            GF2Addition(self.bitsize): ceil(log2(self.bitsize)),
             GF2Square(self.bitsize): square_count,
-            GF2Multiplication(self.bitsize): ceil(log2(self.bitsize))
-            + self.bitsize_hamming_weight
-            - 1,
-        }
+        } | ({GF2MulViaKaratsuba(self.bitsize): mul_count} if mul_count else {})
 
     def on_classical_vals(self, *, x) -> Dict[str, 'ClassicalValT']:
         assert isinstance(x, self.qgf.gf_type)
@@ -204,10 +205,13 @@ def on_classical_vals(self, *, x) -> Dict[str, 'ClassicalValT']:
                         result = result**2
                     junk.append(result)
                     result = result * beta
-            beta_squared = beta ** (2 ** (2**i))
-            junk.extend([beta, beta_squared])
-            beta = beta * beta_squared
+            if i != bitsize_minus_one.bit_length() - 1:
+                beta_squared = beta ** (2 ** (2**i))
+                junk.extend([beta, beta_squared])
+                beta = beta * beta_squared
         junk.append(beta)
+        assert x == junk[0]
+        junk = junk[1:]
         return {'x': x, 'result': x ** (-1) if x else self.qgf.gf_type(0), 'junk': np.array(junk)}
 
 

qualtran/bloqs/gf_arithmetic/gf2_inverse_test.py

@@ -38,9 +38,9 @@ def test_gf2_inverse_symbolic(bloq_autotester):
 def test_gf2_inverse_symbolic_toffoli_complexity():
     bloq = _gf2_inverse_symbolic.make()
     m = bloq.bitsize
-    expected_expr = m**2 * (2 * ceil(log2(m)) - 1)
+    expected_expr = m ** log2(3) * (2 * ceil(log2(m)) - 2)
     assert get_cost_value(bloq, QECGatesCost()).total_toffoli_only() - expected_expr == 0
-    expected_expr = m * (3 * ceil(log2(m)) + 2)
+    expected_expr = m * (3 * ceil(log2(m)) - 1)
     assert isinstance(expected_expr, sympy.Expr)
     assert sympy.simplify(get_cost_value(bloq, QubitCount()) - expected_expr) == 0
 
@@ -60,7 +60,7 @@ def test_gf2_inverse_classical_sim(m):
     assert_consistent_classical_action(bloq, x=GFM.elements)
 
 
-@pytest.mark.parametrize('m', [*range(1, 12)])
+@pytest.mark.parametrize('m', [*range(1, 17)])
 def test_gf2_equivalent_bloq_counts(m):
     bloq = GF2Inverse(m)
     assert_equivalent_bloq_counts(bloq, generalizer=[ignore_split_join, ignore_alloc_free])

qualtran/bloqs/gf_arithmetic/gf2_multiplication.py

@@ -94,6 +94,9 @@ def build_call_graph(
     def adjoint(self) -> 'SynthesizeLRCircuit':
         return attrs.evolve(self, is_adjoint=not self.is_adjoint)
 
+    def __str__(self):
+        return f'{self.__class__.__name__}†' if self.is_adjoint else f'{self.__class__.__name__}'
+
 
 def _qgf_converter(x: Union[QGF, int, Poly, SymbolicInt, Sequence[int]]) -> QGF:
     if isinstance(x, QGF):
@@ -289,8 +292,8 @@ def m_x(self) -> Poly:
         return self.dtype.gf_type.irreducible_poly
 
     @cached_property
-    def n(self) -> int:
-        return self.m_x.degree
+    def n(self) -> SymbolicInt:
+        return self.dtype.bitsize
 
     @cached_property
     def qgf(self) -> QGF:
@@ -315,7 +318,7 @@ def lup(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
         by a full rank matrix that can be decomposed into PLU where L and U are lower
         and upper traingular matricies and P is a permutation matrix.
         """
-        n = self.n
+        n = int(self.n)
         matrix = np.zeros((n, n), dtype=int)
         for i in range(n):
             p = self._const * self.galois_field(2**i)
@@ -970,12 +973,12 @@ def m_x(self):
         return self.dtype.gf_type.irreducible_poly
 
     def __attrs_post_init__(self):
-        if self.m_x.degree < 2:
+        if not is_symbolic(self.dtype) and self.m_x.degree < 2:
             raise ValueError(f'GF2MulViaKaratsuba is not supported for {self.m_x}')
 
     @cached_property
-    def n(self):
-        return int(self.m_x.degrees.max())
+    def n(self) -> SymbolicInt:
+        return self.dtype.bitsize
 
     @cached_property
     def gf(self):
@@ -988,6 +991,9 @@ def qgf(self):
     def adjoint(self) -> 'GF2MulViaKaratsuba':
         return attrs.evolve(self, uncompute=not self.uncompute)
 
+    def __str__(self):
+        return f'{self.__class__.__name__}†' if self.uncompute else f'{self.__class__.__name__}'
+
     @cached_property
     def signature(self) -> 'Signature':
         # C is directional
@@ -1036,9 +1042,12 @@ def build_composite_bloq(
     def build_call_graph(
         self, ssa: 'SympySymbolAllocator'
     ) -> Union['BloqCountDictT', Set['BloqCountT']]:
+        if is_symbolic(self.n):
+            return {Toffoli(): self.n ** (log2(3)), CNOT(): self.n**2}
+
         if self.n == 1:
             return {Toffoli(): 1}
-        if not is_symbolic(self.n) and 2 * self.k == self.n:
+        if 2 * self.k == self.n:
             return {
                 CNOT(): 4 * (self.n - self.k),
                 BinaryPolynomialMultiplication(self.k): 3,