Replace ModInv shim with KaliskiModInverse in ECAdd bloq (#1485)

* Swap modInv bloq pending classical simulation for the adjoint

* Fix gate counts test for ec add using kaliski mod inverse

* remove unused mod inverse shims

* Remove assertions in ModInverse for cases where we don't care

* Fix lint errors

* Fix bug with mod inversion of 0

* Fix bug in step 5

* staging for mod inverse fix

* Adjust toffoli count for adjusted mod inverse circuit

* remove alternative solution bloqs from call graph

* remove unused import

* Resolve nit

* Fix bloq counts

qualtran/bloqs/factoring/ecc/ec_add.py

@@ -41,12 +41,12 @@
     CModNeg,
     CModSub,
     DirtyOutOfPlaceMontgomeryModMul,
+    KaliskiModInverse,
     ModAdd,
     ModDbl,
     ModNeg,
     ModSub,
 )
-from qualtran.bloqs.mod_arithmetic._shims import ModInv
 from qualtran.resource_counting import BloqCountDictT, SympySymbolAllocator
 from qualtran.simulation.classical_sim import ClassicalValT
 from qualtran.symbolics.types import HasLength, is_symbolic
@@ -285,7 +285,7 @@ def build_composite_bloq(
         ctrl, b, y = bb.add(CModSub(QMontgomeryUInt(self.n), mod=self.mod), ctrl=ctrl, x=b, y=y)
 
         # Perform modular inversion s.t. x = (x - a)^-1 % p.
-        x, z1, z2 = bb.add(ModInv(n=self.n, mod=self.mod), x=x)
+        x, junk = bb.add(KaliskiModInverse(bitsize=self.n, mod=self.mod), x=x)
 
         # Perform modular multiplication z4 = (y / x) % p.
         x, y, z4, z3, reduced = bb.add(
@@ -336,7 +336,7 @@ def build_composite_bloq(
             qrom_indices=z3,
             reduced=reduced,
         )
-        x = bb.add(ModInv(n=self.n, mod=self.mod).adjoint(), x=x, garbage1=z1, garbage2=z2)
+        x = bb.add(KaliskiModInverse(bitsize=self.n, mod=self.mod).adjoint(), x=x, junk=junk)
 
         # Return the output registers.
         return {'f1': f1, 'ctrl': ctrl, 'a': a, 'b': b, 'x': x, 'y': y, 'lam': lam, 'lam_r': lam_r}
@@ -346,7 +346,7 @@ def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT:
             Equals(QMontgomeryUInt(self.n)): 1,
             ModSub(QMontgomeryUInt(self.n), mod=self.mod): 1,
             CModSub(QMontgomeryUInt(self.n), mod=self.mod): 1,
-            ModInv(n=self.n, mod=self.mod): 1,
+            KaliskiModInverse(bitsize=self.n, mod=self.mod): 1,
             DirtyOutOfPlaceMontgomeryModMul(
                 bitsize=self.n, window_size=self.window_size, mod=self.mod
             ): 1,
@@ -355,7 +355,7 @@ def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT:
             DirtyOutOfPlaceMontgomeryModMul(
                 bitsize=self.n, window_size=self.window_size, mod=self.mod
             ).adjoint(): 1,
-            ModInv(n=self.n, mod=self.mod).adjoint(): 1,
+            KaliskiModInverse(bitsize=self.n, mod=self.mod).adjoint(): 1,
         }
 
 
@@ -706,7 +706,7 @@ def build_composite_bloq(
             raise DecomposeTypeError(f"Cannot decompose {self} with symbolic `n`.")
 
         # x = x ^ -1 % p.
-        x, z1, z2 = bb.add(ModInv(n=self.n, mod=self.mod), x=x)
+        x, junk = bb.add(KaliskiModInverse(bitsize=self.n, mod=self.mod), x=x)
 
         # z4 = x * y % p.
         x, y, z4, z3, reduced = bb.add(
@@ -744,7 +744,7 @@ def build_composite_bloq(
             qrom_indices=z3,
             reduced=reduced,
         )
-        x = bb.add(ModInv(n=self.n, mod=self.mod).adjoint(), x=x, garbage1=z1, garbage2=z2)
+        x = bb.add(KaliskiModInverse(bitsize=self.n, mod=self.mod).adjoint(), x=x, junk=junk)
 
         # If ctrl: x = x_r - a % p.
         ctrl, x = bb.add(CModNeg(QMontgomeryUInt(self.n), mod=self.mod), ctrl=ctrl, x=x)
@@ -761,14 +761,14 @@ def build_composite_bloq(
     def build_call_graph(self, ssa: SympySymbolAllocator) -> BloqCountDictT:
         return {
             CModSub(QMontgomeryUInt(self.n), mod=self.mod): 1,
-            ModInv(n=self.n, mod=self.mod): 1,
+            KaliskiModInverse(bitsize=self.n, mod=self.mod): 1,
             DirtyOutOfPlaceMontgomeryModMul(
                 bitsize=self.n, window_size=self.window_size, mod=self.mod
             ): 1,
             DirtyOutOfPlaceMontgomeryModMul(
                 bitsize=self.n, window_size=self.window_size, mod=self.mod
             ).adjoint(): 1,
-            ModInv(n=self.n, mod=self.mod).adjoint(): 1,
+            KaliskiModInverse(bitsize=self.n, mod=self.mod).adjoint(): 1,
             ModAdd(self.n, mod=self.mod): 1,
             MultiControlX(cvs=[1, 1]): self.n,
             CModNeg(QMontgomeryUInt(self.n), mod=self.mod): 1,

qualtran/bloqs/factoring/ecc/ec_add_test.py

@@ -407,19 +407,18 @@ def test_ec_add_symbolic_cost():
     # This is why instead of using bitsize=n directly, we use bitsize=4*m=n.
     b = ECAdd(n=4 * m, window_size=4, mod=p)
     cost = get_cost_value(b, QECGatesCost()).total_t_and_ccz_count()
-    assert cost['n_t'] == 0
+    # We have some T gates since we use CSwapApprox instead of n CSWAPs in KaliskiModInverse.
+    total_toff = (cost['n_t'] / 4 + cost['n_ccz']) * sympy.Integer(1)
+    total_toff = total_toff.subs(m, n / 4).expand()
 
     # Litinski 2023 https://arxiv.org/abs/2306.08585
     # Based on the counts from Figures 3, 5, and 8 the toffoli count for ECAdd is 126.5n^2 + 189n.
-    # The following formula is 126.5n^2 + 175.5n - 35. We account for the discrepancy in the
-    # coefficient of n by a reduction in the toffoli cost of Montgomery ModMult, n extra toffolis
-    # in ModNeg, and 2n extra toffolis to do n 3-controlled toffolis in step 2. The expression is
-    # written with rationals because sympy comparison fails with floats.
-    assert isinstance(cost['n_ccz'], sympy.Expr)
-    assert (
-        cost['n_ccz'].subs(m, n / 4).expand()
-        == sympy.Rational(253, 2) * n**2 + sympy.Rational(351, 2) * n - 35
-    )
+    # The following formula is 126.5n^2 + 195.5n - 31. We account for the discrepancy in the
+    # coefficient of n by a reduction in the toffoli cost of Montgomery ModMult, an increase in the
+    # toffoli cost for Kaliski Mod Inverse, n extra toffolis in ModNeg, 2n extra toffolis to do n
+    # 3-controlled toffolis in step 2. The expression is written with rationals because sympy
+    # comparison fails with floats.
+    assert total_toff == sympy.Rational(253, 2) * n**2 + sympy.Rational(391, 2) * n - 31
 
 
 def test_ec_add(bloq_autotester):

qualtran/bloqs/mod_arithmetic/__init__.py

@@ -12,7 +12,6 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from ._shims import ModInv
 from .mod_addition import CModAdd, CModAddK, CtrlScaleModAdd, ModAdd, ModAddK
 from .mod_division import KaliskiModInverse
 from .mod_multiplication import CModMulK, DirtyOutOfPlaceMontgomeryModMul, ModDbl