further optimizations -- specializing and inlining

Author: martyall

Diff for: arithmetic-circuits.cabal

@@ -242,7 +242,7 @@ benchmark circuit-benchmarks
 
   hs-source-dirs:     bench
 
-  ghc-options:        -freverse-errors -O2 -Wall -main-is Main
+  ghc-options:        -freverse-errors -O4 -Wall -main-is Main
   build-depends:
       arithmetic-circuits
     , arithmetic-circuits:language

Diff for: bench/Circuit/Bench.hs

@@ -9,7 +9,6 @@ import Data.Field.Galois (Prime)
 import Data.IntMap qualified as IntMap
 import Data.Map qualified as Map
 import Data.Vector (generateM)
-import GHC.TypeNats (Natural, SNat, withKnownNat, withSomeSNat)
 import Protolude
 
 type Fr = Prime 21888242871839275222246405745257275088548364400416034343698204186575808495617
@@ -24,24 +23,22 @@ benchmarks =
       bench "1_000_000" $ whnf largeMult 1_000_000
     ]
 
-largeMult :: Natural -> Fr
+largeMult :: Int -> Fr
 largeMult n =
-  withSomeSNat n $ \(sn :: SNat n) ->
-    withKnownNat sn $
-      let BuilderState {bsVars, bsCircuit} = snd $ runCircuitBuilder (program (Proxy @n))
-          inputs =
-            assignInputs bsVars $
-              Map.fromList $
-                map (\i -> (nameInput i, fromIntegral i + 1)) [0 .. n - 1]
-          w = altSolve bsCircuit inputs
-       in fromMaybe (panic "output not found") $ lookupVar bsVars "out" w
+  let BuilderState {bsVars, bsCircuit} = snd $ runCircuitBuilder (program n)
+      inputs =
+        assignInputs bsVars $
+          Map.fromList $
+            map (\i -> (nameInput i, fromIntegral i + 1)) [0 .. n - 1]
+      w = altSolve bsCircuit inputs
+   in fromMaybe (panic "output not found") $ lookupVar bsVars "out" w
 
 nameInput :: (Integral a) => a -> Text
 nameInput i = "x" <> show (toInteger i)
 
-program :: forall n. (KnownNat n) => Proxy n -> ExprM Fr (Var Wire Fr Fr)
+program :: Int -> ExprM Fr (Var Wire Fr 'TField)
 program p = do
-  xs <- generateM (fromIntegral $ natVal p) $ \i ->
+  xs <- generateM p $ \i ->
     var_ <$> fieldInput Public (nameInput i)
   fieldOutput "out" $ product xs
 

Diff for: circom-compat/app/Main.hs

@@ -26,7 +26,7 @@ main = do
 
 type Fr = Prime 21888242871839275222246405745257275088548364400416034343698204186575808495617
 
-program :: ExprM Fr (Var Wire Fr Bool)
+program :: ExprM Fr (Var Wire Fr 'TBool)
 program = do
   n <- var_ <$> fieldInput Public "n"
   a <- var_ <$> fieldInput Private "a"

Diff for: circom-compat/src/R1CS/Circom.hs

@@ -17,6 +17,7 @@ module R1CS.Circom
   )
 where
 
+import Circuit (CircuitVars (..))
 import Data.Binary (Binary (..), Get, Put)
 import Data.Binary.Get (getInt32le, getInt64le, getWord32le, getWord64le, lookAhead, skip)
 import Data.Binary.Put (putInt32le, putLazyByteString, putWord32le, putWord64le, runPut)
@@ -27,7 +28,6 @@ import Data.IntSet qualified as IntSet
 import Data.Vector (Vector)
 import Data.Vector qualified as V
 import Protolude
-import Circuit (CircuitVars(..))
 import R1CS (LinearPoly (..), R1C (..), R1CS (..), Witness (..))
 import Prelude (fail)
 

Diff for: circuit/src/Circuit/Arithmetic.hs

@@ -22,6 +22,7 @@ module Circuit.Arithmetic
     InputBidings (..),
     insertInputBinding,
     Reindexable (..),
+    restrictVars,
   )
 where
 
@@ -316,6 +317,7 @@ instance (Ord label) => Monoid (CircuitVars label) where
         cvOutputs = mempty,
         cvInputsLabels = mempty
       }
+
 instance Reindexable (CircuitVars label) where
   reindex f CircuitVars {..} =
     CircuitVars
@@ -355,6 +357,16 @@ collectCircuitVars (ArithCircuit gates) =
           cvInputsLabels = inputBindingsFromList ls
         }
 
+restrictVars :: CircuitVars label -> IntSet -> CircuitVars label
+restrictVars CircuitVars {..} vars =
+  CircuitVars
+    { cvVars = IntSet.intersection cvVars vars,
+      cvPrivateInputs = IntSet.intersection cvPrivateInputs vars,
+      cvPublicInputs = IntSet.intersection cvPublicInputs vars,
+      cvOutputs = IntSet.intersection cvOutputs vars,
+      cvInputsLabels = cvInputsLabels
+    }
+
 assignInputs :: (Ord label) => CircuitVars label -> Map label f -> IntMap f
 assignInputs CircuitVars {..} inputs =
   IntMap.mapMaybe (\label -> Map.lookup label inputs) (varToLabel cvInputsLabels)
@@ -394,11 +406,10 @@ instance Reindexable (InputBidings label) where
       { labelToVar = Map.mapMaybe (flip IntMap.lookup f) labelToVar,
         varToLabel = IntMap.compose varToLabel (reverseMap f)
       }
-      where
+    where
       reverseMap :: IntMap Int -> IntMap Int
       reverseMap = IntMap.foldlWithKey' (\acc k v -> IntMap.insert v k acc) mempty
 
-
 instance (Ord label) => Semigroup (InputBidings label) where
   a <> b =
     InputBidings

Diff for: language/src/Circuit/Language/Compile.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE DataKinds #-}
 {-# LANGUAGE PatternSynonyms #-}
 
 module Circuit.Language.Compile
@@ -174,17 +175,20 @@ mulToImm l r = do
   o <- imm
   emit $ Mul (addVar l) (addVar r) o
   pure o
+{-# INLINE mulToImm #-}
 
 -- | Add a Mul and its output to the ArithCircuit
 emit :: (MonadState (BuilderState f) m) => Gate f Wire -> m ()
 emit c = modify $ \s@(BuilderState {bsCircuit = ArithCircuit cs}) ->
   s {bsCircuit = ArithCircuit (c : cs)}
+{-# INLINE emit #-}
 
 -- | Turn a wire into an affine circuit, or leave it be
 addVar :: SignalSource f -> AffineCircuit f Wire
 addVar s = case s of
   WireSource w -> Var w
   AffineSource c -> c
+{-# INLINE addVar #-}
 
 -- | Turn an affine circuit into a wire, or leave it be
 addWire :: (MonadState (BuilderState f) m, Num f) => SignalSource f -> m Wire
@@ -194,29 +198,26 @@ addWire x = case x of
     mulOut <- imm
     emit $ Mul (ConstGate 1) c mulOut
     pure mulOut
+{-# INLINE addWire #-}
 
 --------------------------------------------------------------------------------
 
 compileWithWire ::
   (Hashable f) =>
   (GaloisField f) =>
-  (MonadState (BuilderState f) m) =>
-  (MonadError (CircuitBuilderError f) m) =>
-  Var Wire f f ->
-  Expr Wire f f ->
-  m (Var Wire f f)
+  Var Wire f 'TField ->
+  Expr Wire f 'TField ->
+  ExprM f (Var Wire f 'TField)
 compileWithWire freshWire e = do
   res <- compileWithWires (V.singleton freshWire) e
   pure . V.head $ res
 
 compileWithWires ::
   (Hashable f) =>
   (GaloisField f) =>
-  (MonadState (BuilderState f) m) =>
-  (MonadError (CircuitBuilderError f) m) =>
-  V.Vector (Var Wire f f) ->
+  V.Vector (Var Wire f 'TField) ->
   Expr Wire f ty ->
-  m (V.Vector (Var Wire f f))
+  ExprM f (V.Vector (Var Wire f 'TField))
 compileWithWires ws expr = do
   compileOut <- compile expr
   for (V.zip compileOut ws) $ \(o, freshWire) -> do
@@ -233,25 +234,20 @@ compileWithWires ws expr = do
 {-# SCC compile #-}
 compile ::
   (Hashable f, GaloisField f) =>
-  (MonadState (BuilderState f) m) =>
-  (MonadError (CircuitBuilderError f) m) =>
   Expr Wire f ty ->
-  m (V.Vector (SignalSource f))
+  ExprM f (V.Vector (SignalSource f))
 compile e = do
-  let g = reifyGraph e
-  res <- traverse _compile g
-  case res of
-    (_ :|> x) -> pure x
+  case reifyGraph e of
+    (xs :|> x) ->
+      traverse_ _compile xs >> _compile x
     _ -> panic "empty graph"
 
 {-# SCC _compile #-}
 _compile ::
-  forall f m.
+  forall f.
   (Hashable f, GaloisField f) =>
-  (MonadState (BuilderState f) m) =>
-  (MonadError (CircuitBuilderError f) m) =>
   (Hash, Node Wire f) ->
-  m (V.Vector (SignalSource f))
+  ExprM f (V.Vector (SignalSource f))
 _compile (h, expr) = case expr of
   NVal f -> do
     let source = V.singleton $ AffineSource $ ConstGate f
@@ -398,8 +394,8 @@ exprToArithCircuit expr output = do
 
 fieldToBool ::
   (Hashable f, GaloisField f) =>
-  Expr Wire f f ->
-  ExprM f (Expr Wire f Bool)
+  Expr Wire f 'TField ->
+  ExprM f (Expr Wire f 'TBool)
 fieldToBool e = do
   -- let eOut = unType e
   a <- compile e >>= assertSingleSource >>= addWire
@@ -411,8 +407,8 @@ _unBundle ::
   (KnownNat n) =>
   (GaloisField f) =>
   (Hashable f) =>
-  Expr Wire f (SV.Vector n ty) ->
-  ExprM f (SV.Vector n (Expr Wire f f))
+  Expr Wire f (TVec n ty) ->
+  ExprM f (SV.Vector n (Expr Wire f 'TField))
 _unBundle b = do
   bis <- compile b
   ws <- traverse addWire bis