perf: reduce mem allocs in scs frontend (#654)

* perf: avoid []uint32 alloc when compressing constraints * perf: reduce allocs in scs frontend * fix: remove System.lbHints from go-cmp test
Consensys · Apr 25, 2023 · 94f1ddc · 94f1ddc
1 parent 96aad98
commit 94f1ddc
Show file tree

Hide file tree

Showing 30 changed files with 165 additions and 111 deletions.
diff --git a/constraint/bls12-377/coeff.go b/constraint/bls12-377/coeff.go
diff --git a/constraint/bls12-377/r1cs_test.go b/constraint/bls12-377/r1cs_test.go
diff --git a/constraint/bls12-381/coeff.go b/constraint/bls12-381/coeff.go
diff --git a/constraint/bls12-381/r1cs_test.go b/constraint/bls12-381/r1cs_test.go
diff --git a/constraint/bls24-315/coeff.go b/constraint/bls24-315/coeff.go
diff --git a/constraint/bls24-315/r1cs_test.go b/constraint/bls24-315/r1cs_test.go
diff --git a/constraint/bls24-317/coeff.go b/constraint/bls24-317/coeff.go
diff --git a/constraint/bls24-317/r1cs_test.go b/constraint/bls24-317/r1cs_test.go
diff --git a/constraint/blueprint_hint.go b/constraint/blueprint_hint.go
@@ -43,7 +43,7 @@ func (b *BlueprintGenericHint) CompressHint(h HintMapping) []uint32 {
 
 	nbInputs += 2 // output range start / end
 
-	r := make([]uint32, 0, nbInputs)
+	r := getBuffer(nbInputs)
 	r = append(r, uint32(nbInputs))
 	r = append(r, uint32(h.HintID))
 	r = append(r, uint32(len(h.Inputs)))

diff --git a/constraint/blueprint_r1cs.go b/constraint/blueprint_r1cs.go
@@ -17,7 +17,7 @@ func (b *BlueprintGenericR1C) NbConstraints() int {
 func (b *BlueprintGenericR1C) CompressR1C(c *R1C) []uint32 {
 	// we store total nb inputs, len L, len R, len O, and then the "flatten" linear expressions
 	nbInputs := 4 + 2*(len(c.L)+len(c.R)+len(c.O))
-	r := make([]uint32, 0, nbInputs)
+	r := getBuffer(nbInputs)
 	r = append(r, uint32(nbInputs))
 	r = append(r, uint32(len(c.L)), uint32(len(c.R)), uint32(len(c.O)))
 	for _, t := range c.L {
@@ -56,3 +56,17 @@ func (b *BlueprintGenericR1C) DecompressR1C(c *R1C, calldata []uint32) {
 	copySlice(&c.R, lenR, offset+2*lenL)
 	copySlice(&c.O, lenO, offset+2*(lenL+lenR))
 }
+
+// since frontend is single threaded, to avoid allocating slices at each compress call
+// we transit the compressed output through here
+var bufCalldata []uint32
+
+// getBuffer return a slice with at least the given capacity to use in Compress methods
+// this is obviously not thread safe, but the frontend is single threaded anyway.
+func getBuffer(size int) []uint32 {
+	if cap(bufCalldata) < size {
+		bufCalldata = make([]uint32, 0, size*2)
+	}
+	bufCalldata = bufCalldata[:0]
+	return bufCalldata
+}
diff --git a/constraint/blueprint_scs.go b/constraint/blueprint_scs.go
@@ -14,7 +14,8 @@ var (
 // Encodes
 //
 //	qL⋅xa + qR⋅xb + qO⋅xc + qM⋅(xaxb) + qC == 0
-type BlueprintGenericSparseR1C struct{}
+type BlueprintGenericSparseR1C struct {
+}
 
 func (b *BlueprintGenericSparseR1C) NbInputs() int {
 	return 9 // number of fields in SparseR1C
@@ -24,17 +25,16 @@ func (b *BlueprintGenericSparseR1C) NbConstraints() int {
 }
 
 func (b *BlueprintGenericSparseR1C) CompressSparseR1C(c *SparseR1C) []uint32 {
-	return []uint32{
-		c.XA,
-		c.XB,
-		c.XC,
-		c.QL,
-		c.QR,
-		c.QO,
-		c.QM,
-		c.QC,
-		uint32(c.Commitment),
-	}
+	bufSCS[0] = c.XA
+	bufSCS[1] = c.XB
+	bufSCS[2] = c.XC
+	bufSCS[3] = c.QL
+	bufSCS[4] = c.QR
+	bufSCS[5] = c.QO
+	bufSCS[6] = c.QM
+	bufSCS[7] = c.QC
+	bufSCS[8] = uint32(c.Commitment)
+	return bufSCS[:]
 }
 
 func (b *BlueprintGenericSparseR1C) DecompressSparseR1C(c *SparseR1C, calldata []uint32) {
@@ -167,12 +167,11 @@ func (b *BlueprintSparseR1CMul) NbConstraints() int {
 }
 
 func (b *BlueprintSparseR1CMul) CompressSparseR1C(c *SparseR1C) []uint32 {
-	return []uint32{
-		c.XA,
-		c.XB,
-		c.XC,
-		c.QM,
-	}
+	bufSCS[0] = c.XA
+	bufSCS[1] = c.XB
+	bufSCS[2] = c.XC
+	bufSCS[3] = c.QM
+	return bufSCS[:4]
 }
 
 func (b *BlueprintSparseR1CMul) Solve(s Solver, calldata []uint32) error {
@@ -209,14 +208,13 @@ func (b *BlueprintSparseR1CAdd) NbConstraints() int {
 }
 
 func (b *BlueprintSparseR1CAdd) CompressSparseR1C(c *SparseR1C) []uint32 {
-	return []uint32{
-		c.XA,
-		c.XB,
-		c.XC,
-		c.QL,
-		c.QR,
-		c.QC,
-	}
+	bufSCS[0] = c.XA
+	bufSCS[1] = c.XB
+	bufSCS[2] = c.XC
+	bufSCS[3] = c.QL
+	bufSCS[4] = c.QR
+	bufSCS[5] = c.QC
+	return bufSCS[:6]
 }
 
 func (blueprint *BlueprintSparseR1CAdd) Solve(s Solver, calldata []uint32) error {
@@ -258,11 +256,10 @@ func (b *BlueprintSparseR1CBool) NbConstraints() int {
 }
 
 func (b *BlueprintSparseR1CBool) CompressSparseR1C(c *SparseR1C) []uint32 {
-	return []uint32{
-		c.XA,
-		c.QL,
-		c.QM,
-	}
+	bufSCS[0] = c.XA
+	bufSCS[1] = c.QL
+	bufSCS[2] = c.QM
+	return bufSCS[:3]
 }
 
 func (blueprint *BlueprintSparseR1CBool) Solve(s Solver, calldata []uint32) error {
@@ -285,3 +282,7 @@ func (b *BlueprintSparseR1CBool) DecompressSparseR1C(c *SparseR1C, calldata []ui
 	c.QL = calldata[1]
 	c.QM = calldata[2]
 }
+
+// since frontend is single threaded, to avoid allocating slices at each compress call
+// we transit the compressed output through here
+var bufSCS [9]uint32
diff --git a/constraint/bn254/coeff.go b/constraint/bn254/coeff.go
diff --git a/constraint/bn254/r1cs_test.go b/constraint/bn254/r1cs_test.go
diff --git a/constraint/bw6-633/coeff.go b/constraint/bw6-633/coeff.go
diff --git a/constraint/bw6-633/r1cs_test.go b/constraint/bw6-633/r1cs_test.go
diff --git a/constraint/bw6-761/coeff.go b/constraint/bw6-761/coeff.go
diff --git a/constraint/bw6-761/r1cs_test.go b/constraint/bw6-761/r1cs_test.go
diff --git a/constraint/core.go b/constraint/core.go
@@ -88,9 +88,8 @@ type System struct {
 	bitLen int      `cbor:"-"`
 
 	// level builder
-	lbWireLevel []int            `cbor:"-"` // at which level we solve a wire. init at -1.
-	lbOutputs   []uint32         `cbor:"-"` // wire outputs for current constraint.
-	lbHints     map[int]struct{} `cbor:"-"` // hints we processed in current round
+	lbWireLevel []int    `cbor:"-"` // at which level we solve a wire. init at -1.
+	lbOutputs   []uint32 `cbor:"-"` // wire outputs for current constraint.
 
 	CommitmentInfo Commitment
 
@@ -108,9 +107,11 @@ func NewSystem(scalarField *big.Int, capacity int, t SystemType) System {
 		MHintsDependencies: make(map[solver.HintID]string),
 		q:                  new(big.Int).Set(scalarField),
 		bitLen:             scalarField.BitLen(),
-		lbHints:            map[int]struct{}{},
 		Instructions:       make([]Instruction, 0, capacity),
-		CallData:           make([]uint32, 0, capacity*2),
+		CallData:           make([]uint32, 0, capacity*8),
+		lbOutputs:          make([]uint32, 0, 256),
+		lbWireLevel:        make([]int, 0, capacity),
+		Levels:             make([][]int, 0, capacity/2),
 	}
 	system.genericHint = system.AddBlueprint(&BlueprintGenericHint{})
 	return system

diff --git a/constraint/hint.go b/constraint/hint.go
@@ -21,14 +21,14 @@ func (h *HintMapping) WireIterator() func() int {
 	for i := 0; i < len(h.Inputs); i++ {
 		n += len(h.Inputs[i])
 	}
-	inputs := make([]int, 0, n)
+	inputs := getBuffer(n)
 	for i := 0; i < len(h.Inputs); i++ {
 		for j := 0; j < len(h.Inputs[i]); j++ {
 			term := h.Inputs[i][j]
 			if term.IsConstant() {
 				continue
 			}
-			inputs = append(inputs, int(term.VID))
+			inputs = append(inputs, term.VID)
 		}
 	}
 	lenOutputs := int(h.OutputRange.End - h.OutputRange.Start)
@@ -40,7 +40,7 @@ func (h *HintMapping) WireIterator() func() int {
 		}
 		if curr < lenOutputs+len(inputs) {
 			curr++
-			return inputs[curr-1-lenOutputs]
+			return int(inputs[curr-1-lenOutputs])
 		}
 		return -1
 	}

diff --git a/constraint/level_builder.go b/constraint/level_builder.go
@@ -34,7 +34,6 @@ func (system *System) updateLevel(iID int, c Iterable) {
 	// clean the table. NB! Do not remove or move, this is required to make the
 	// compilation deterministic.
 	system.lbOutputs = system.lbOutputs[:0]
-	system.lbHints = map[int]struct{}{}
 }
 
 func (system *System) processWire(wireID uint32, maxLevel *int) {

diff --git a/constraint/r1cs_test.go b/constraint/r1cs_test.go
@@ -31,26 +31,26 @@ func ExampleR1CS_GetR1Cs() {
 
 	// X² == X * X
 	r1cs.AddR1C(constraint.R1C{
-		L: constraint.LinearExpression{r1cs.MakeTerm(&cOne, X)},
-		R: constraint.LinearExpression{r1cs.MakeTerm(&cOne, X)},
-		O: constraint.LinearExpression{r1cs.MakeTerm(&cOne, v0)},
+		L: constraint.LinearExpression{r1cs.MakeTerm(cOne, X)},
+		R: constraint.LinearExpression{r1cs.MakeTerm(cOne, X)},
+		O: constraint.LinearExpression{r1cs.MakeTerm(cOne, v0)},
 	}, blueprint)
 
 	// X³ == X² * X
 	r1cs.AddR1C(constraint.R1C{
-		L: constraint.LinearExpression{r1cs.MakeTerm(&cOne, v0)},
-		R: constraint.LinearExpression{r1cs.MakeTerm(&cOne, X)},
-		O: constraint.LinearExpression{r1cs.MakeTerm(&cOne, v1)},
+		L: constraint.LinearExpression{r1cs.MakeTerm(cOne, v0)},
+		R: constraint.LinearExpression{r1cs.MakeTerm(cOne, X)},
+		O: constraint.LinearExpression{r1cs.MakeTerm(cOne, v1)},
 	}, blueprint)
 
 	// Y == X³ + X + 5
 	r1cs.AddR1C(constraint.R1C{
-		R: constraint.LinearExpression{r1cs.MakeTerm(&cOne, ONE)},
-		L: constraint.LinearExpression{r1cs.MakeTerm(&cOne, Y)},
+		R: constraint.LinearExpression{r1cs.MakeTerm(cOne, ONE)},
+		L: constraint.LinearExpression{r1cs.MakeTerm(cOne, Y)},
 		O: constraint.LinearExpression{
-			r1cs.MakeTerm(&cFive, ONE),
-			r1cs.MakeTerm(&cOne, X),
-			r1cs.MakeTerm(&cOne, v1),
+			r1cs.MakeTerm(cFive, ONE),
+			r1cs.MakeTerm(cOne, X),
+			r1cs.MakeTerm(cOne, v1),
 		},
 	}, blueprint)
 

diff --git a/constraint/system.go b/constraint/system.go
@@ -53,7 +53,7 @@ type ConstraintSystem interface {
 
 	// MakeTerm returns a new Term. The constraint system may store coefficients in a map, so
 	// calls to this function will grow the memory usage of the constraint system.
-	MakeTerm(coeff *Element, variableID int) Term
+	MakeTerm(coeff Element, variableID int) Term
 
 	// AddCoeff adds a coefficient to the underlying constraint system. The system will not store duplicate,
 	// but is not purging for unused coeff either, so this grows memory usage.

diff --git a/constraint/tinyfield/coeff.go b/constraint/tinyfield/coeff.go
diff --git a/constraint/tinyfield/r1cs_test.go b/constraint/tinyfield/r1cs_test.go