perf: GKR Levels (#1735)

Author: Tabaie

constraint/gkr.go

@@ -0,0 +1,80 @@
+package constraint
+
+type (
+	// GkrClaimSource identifies an incoming evaluation claim for a wire.
+	// Level is the level that produced the claim.
+	// OutgoingClaimIndex selects which of that level's outgoing evaluation points is referenced;
+	// always 0 for SumcheckLevels, 0..M-1 for SkipLevels with M inherited evaluation points.
+	// The initial verifier challenge is represented as {Level: len(schedule), OutgoingClaimIndex: 0}.
+	GkrClaimSource struct {
+		Level              int `json:"level"`
+		OutgoingClaimIndex int `json:"outgoingClaimIndex"`
+	}
+
+	// GkrClaimGroup represents a set of wires sharing identical claim sources.
+	// finalEvalProof index = pos(wire, srcLevel) * NbOutgoingEvalPoints(srcLevel) + ClaimSources[claimI].OutgoingClaimIndex,
+	// where pos(wire, srcLevel) is the wire's position in srcLevel's UniqueGateInputs list.
+	GkrClaimGroup struct {
+		Wires        []int            `json:"wires"`
+		ClaimSources []GkrClaimSource `json:"claimSources"`
+	}
+
+	// GkrProvingLevel is a single level in the proving schedule.
+	GkrProvingLevel interface {
+		NbOutgoingEvalPoints() int
+		// NbClaims returns the total number of claims at this level.
+		NbClaims() int
+		ClaimGroups() []GkrClaimGroup
+		// FinalEvalProofIndex returns where to find the evaluationPointI'th evaluation claim for the wireI'th input wire to the layer,
+		// in the layer's final evaluation proof.
+		FinalEvalProofIndex(wireI, evaluationPointI int) int
+	}
+
+	// GkrSkipLevel represents a level where zerocheck is skipped.
+	// Claims propagate through at their existing evaluation points.
+	GkrSkipLevel GkrClaimGroup
+
+	// GkrSumcheckLevel represents a level where one or more zerochecks are batched
+	// together in a single sumcheck. Each GkrClaimGroup within may have different
+	// claim sources (sumcheck-level batching), or the same source (enabling
+	// zerocheck-level batching with shared eq tables).
+	GkrSumcheckLevel []GkrClaimGroup
+
+	// GkrProvingSchedule is a sequence of levels defining how to prove a GKR circuit.
+	GkrProvingSchedule []GkrProvingLevel
+)
+
+func (g GkrClaimGroup) NbClaims() int { return len(g.Wires) * len(g.ClaimSources) }
+
+func (l GkrSumcheckLevel) NbOutgoingEvalPoints() int { return 1 }
+func (l GkrSumcheckLevel) NbClaims() int {
+	n := 0
+	for _, g := range l {
+		n += len(g.Wires) * len(g.ClaimSources)
+	}
+	return n
+}
+func (l GkrSumcheckLevel) ClaimGroups() []GkrClaimGroup         { return l }
+func (l GkrSumcheckLevel) FinalEvalProofIndex(wireI, _ int) int { return wireI }
+
+func (l GkrSkipLevel) NbOutgoingEvalPoints() int { return len(l.ClaimSources) }
+func (l GkrSkipLevel) NbClaims() int {
+	return GkrClaimGroup(l).NbClaims()
+}
+func (l GkrSkipLevel) ClaimGroups() []GkrClaimGroup { return []GkrClaimGroup{GkrClaimGroup(l)} }
+func (l GkrSkipLevel) FinalEvalProofIndex(wireI, evaluationPointI int) int {
+	return wireI*l.NbOutgoingEvalPoints() + evaluationPointI
+}
+
+// BindGkrFinalEvalProof binds the non-input-wire entries of finalEvalProof into the transcript.
+// Input-wire evaluations are fully determined by the public assignment (and by evaluation points
+// already committed to the transcript), so hashing them contributes nothing to Fiat-Shamir security.
+// uniqueGateInputs is the deduplicated list of gate-input wire indices for the level in the same
+// order as the finalEvalProof entries (i.e. the order returned by UniqueGateInputs).
+func BindGkrFinalEvalProof[F any](transcript interface{ Bind(...F) }, finalEvalProof []F, uniqueGateInputs []int, isInput func(wireI int) bool, level GkrProvingLevel) {
+	for i, inputWireI := range uniqueGateInputs {
+		if !isInput(inputWireI) {
+			transcript.Bind(finalEvalProof[level.FinalEvalProofIndex(i, 0):level.FinalEvalProofIndex(i+1, 0)]...)
+		}
+	}
+}

constraint/marshal.go

@@ -398,6 +398,9 @@ func getTagSet() cbor.TagSet {
 	addType(reflect.TypeOf(BlueprintBatchInverse[U32]{}))
 	addType(reflect.TypeOf(BlueprintBatchInverse[U64]{}))
 
+	addType(reflect.TypeOf(GkrSkipLevel{}))
+	addType(reflect.TypeOf(GkrSumcheckLevel{}))
+
 	// Add types registered by external packages (e.g., GKR blueprints)
 	// These use explicit tag numbers to ensure stability regardless of init() order
 	for _, rt := range registeredBlueprintTypes {

constraint/solver/gkrgates/registry.go

@@ -10,7 +10,7 @@ import (
 	"runtime"
 
 	"github.com/consensys/gnark-crypto/ecc"
-	"github.com/consensys/gnark/internal/gkr/gkrtypes"
+	"github.com/consensys/gnark/internal/gkr/gkrcore"
 	"github.com/consensys/gnark/std/gkrapi/gkr"
 )
 
@@ -119,7 +119,7 @@ func Register(f gkr.GateFunction, nbIn int, options ...RegisterOption) error {
 	}
 
 	for _, curve := range s.curves {
-		compiled, err := gkrtypes.CompileGateFunction(f, nbIn, curve.ScalarField())
+		compiled, err := gkrcore.CompileGateFunction(f, nbIn, curve.ScalarField())
 		if err != nil {
 			return err
 		}

internal/generator/backend/template/gkr/blueprint.go.tmpl

@@ -8,10 +8,9 @@ import (
 	"github.com/consensys/gnark-crypto/ecc"
 	"{{ .FieldPackagePath }}"
 	"{{ .FieldPackagePath }}/polynomial"
-	fiatshamir "github.com/consensys/gnark-crypto/fiat-shamir"
 	"github.com/consensys/gnark-crypto/hash"
 	"github.com/consensys/gnark/constraint"
-	"github.com/consensys/gnark/internal/gkr/gkrtypes"
+	"github.com/consensys/gnark/internal/gkr/gkrcore"
 )
 
 func init() {
@@ -27,7 +26,7 @@ type circuitEvaluator struct {
 // BlueprintSolve is a {{.FieldID}}-specific blueprint for solving GKR circuit instances.
 type BlueprintSolve struct {
 	// Circuit structure (serialized)
-	Circuit     gkrtypes.SerializableCircuit
+	Circuit     gkrcore.SerializableCircuit
 	NbInstances uint32
 
 	// Not serialized - recreated lazily at solve time
@@ -41,10 +40,13 @@ type BlueprintSolve struct {
 
 // Ensures BlueprintSolve implements BlueprintStateful
 var _ constraint.BlueprintStateful[constraint.U64] = (*BlueprintSolve)(nil)
+
 // Equal returns true if the serialized fields of two BlueprintSolve are equal.
 // Used for testing serialization round-trips.
 func (b *BlueprintSolve) Equal(other constraint.BlueprintComparable) bool {
-	if other == nil { return false }
+	if other == nil {
+		return false
+	}
 	o, ok := other.(*BlueprintSolve)
 	if !ok {
 		return false
@@ -107,7 +109,7 @@ func (b *BlueprintSolve) Solve(s constraint.Solver[constraint.U64], inst constra
 		if w.IsInput() {
 			val, delta := s.Read(calldata)
 			calldata = calldata[delta:]
-			// Copy directly from constraint.U64 to fr.Element (both in Montgomery form)
+			// Copy directly from constraint.U64 to {{ .ElementType }} (both in Montgomery form)
 			copy(b.assignments[wI][instanceI][:], val[:])
 		} else {
 			// Get evaluator for this wire from the circuit evaluator
@@ -123,7 +125,7 @@ func (b *BlueprintSolve) Solve(s constraint.Solver[constraint.U64], inst constra
 		}
 	}
 
-	// Set output wires (copy fr.Element to U64 in Montgomery form)
+	// Set output wires (copy {{ .ElementType }} to U64 in Montgomery form)
 	for outI, outWI := range b.outputWires {
 		var val constraint.U64
 		copy(val[:], b.assignments[outWI][instanceI][:])
@@ -150,9 +152,9 @@ func (b *BlueprintSolve) NbConstraints() int {
 
 // NbOutputs implements Blueprint
 func (b *BlueprintSolve) NbOutputs(inst constraint.Instruction) int {
-    if b.outputWires == nil {
-        b.outputWires = b.Circuit.Outputs()
-    }
+	if b.outputWires == nil {
+		b.outputWires = b.Circuit.Outputs()
+	}
 	return len(b.outputWires)
 }
 
@@ -194,21 +196,25 @@ func (b *BlueprintSolve) UpdateInstructionTree(inst constraint.Instruction, tree
 type BlueprintProve struct {
 	SolveBlueprintID constraint.BlueprintID
 	SolveBlueprint   *BlueprintSolve `cbor:"-"` // not serialized, set at compile time
+	Schedule         constraint.GkrProvingSchedule
 	HashName         string
 
 	lock sync.Mutex
 }
 
 // Ensures BlueprintProve implements BlueprintSolvable
 var _ constraint.BlueprintSolvable[constraint.U64] = (*BlueprintProve)(nil)
+
 // Equal returns true if the serialized fields of two BlueprintProve are equal.
 func (b *BlueprintProve) Equal(other constraint.BlueprintComparable) bool {
-	if other == nil { return false }
+	if other == nil {
+		return false
+	}
 	o, ok := other.(*BlueprintProve)
 	if !ok {
 		return false
 	}
-	return b.SolveBlueprintID == o.SolveBlueprintID && b.HashName == o.HashName
+	return b.SolveBlueprintID == o.SolveBlueprintID && b.HashName == o.HashName && reflect.DeepEqual(b.Schedule, o.Schedule)
 }
 
 // Solve implements the BlueprintSolvable interface for proving.
@@ -243,28 +249,27 @@ func (b *BlueprintProve) Solve(s constraint.Solver[constraint.U64], inst constra
 		}
 	}
 
+	// Create hasher and write base challenges
+	hsh := hash.NewHash(b.HashName + "_{{.FieldID}}")
+
 	// Read initial challenges from instruction calldata (parse dynamically, no metadata)
 	// Format: [0]=totalSize, [1...]=challenge linear expressions
-	insBytes := make([][]byte, 0) // first challenges
 	calldata := inst.Calldata[1:] // skip size prefix
 	for len(calldata) != 0 {
 		val, delta := s.Read(calldata)
 		calldata = calldata[delta:]
 
-		// Copy directly from constraint.U64 to fr.Element (both in Montgomery form)
+		// Copy directly from constraint.U64 to {{ .ElementType }} (both in Montgomery form)
 		var challenge {{ .ElementType }}
 		copy(challenge[:], val[:])
-		insBytes = append(insBytes, challenge.Marshal())
+		challengeBytes := challenge.Bytes()
+		hsh.Write(challengeBytes[:])
 	}
 
-	// Create Fiat-Shamir settings
-	hsh := hash.NewHash(b.HashName + "_{{.FieldID}}")
-	fsSettings := fiatshamir.WithHash(hsh, insBytes...)
-
 	// Call the {{.FieldID}}-specific Prove function (assignments already WireAssignment type)
-	proof, err := Prove(solveBlueprint.Circuit, assignments, fsSettings)
+	proof, err := Prove(solveBlueprint.Circuit, b.Schedule, assignments, hsh)
 	if err != nil {
-		return fmt.Errorf("{{toLower .FieldID}} prove failed: %w", err)
+		return fmt.Errorf("{{.FieldID}} prove failed: %w", err)
 	}
 
 	for i, elem := range proof.flatten() {
@@ -292,7 +297,7 @@ func (b *BlueprintProve) proofSize() int {
 	}
 	nbPaddedInstances := ecc.NextPowerOfTwo(uint64(b.SolveBlueprint.NbInstances))
 	logNbInstances := bits.TrailingZeros64(nbPaddedInstances)
-	return b.SolveBlueprint.Circuit.ProofSize(logNbInstances)
+	return b.SolveBlueprint.Circuit.ProofSize(b.Schedule, logNbInstances)
 }
 
 // NbOutputs implements Blueprint
@@ -344,9 +349,12 @@ type BlueprintGetAssignment struct {
 
 // Ensures BlueprintGetAssignment implements BlueprintSolvable
 var _ constraint.BlueprintSolvable[constraint.U64] = (*BlueprintGetAssignment)(nil)
+
 // Equal returns true if the serialized fields of two BlueprintGetAssignment are equal.
 func (b *BlueprintGetAssignment) Equal(other constraint.BlueprintComparable) bool {
-	if other == nil { return false }
+	if other == nil {
+		return false
+	}
 	o, ok := other.(*BlueprintGetAssignment)
 	if !ok {
 		return false
@@ -418,7 +426,7 @@ func (b *BlueprintGetAssignment) UpdateInstructionTree(inst constraint.Instructi
 }
 
 // NewBlueprints creates and registers all GKR blueprints for {{.FieldID}}
-func NewBlueprints(circuit gkrtypes.SerializableCircuit, hashName string, compiler constraint.CustomizableSystem) gkrtypes.Blueprints {
+func NewBlueprints(circuit gkrcore.SerializableCircuit, schedule constraint.GkrProvingSchedule, hashName string, compiler constraint.CustomizableSystem) gkrcore.Blueprints {
 	// Create and register solve blueprint
 	solve := &BlueprintSolve{Circuit: circuit}
 	solveID := compiler.AddBlueprint(solve)
@@ -427,6 +435,7 @@ func NewBlueprints(circuit gkrtypes.SerializableCircuit, hashName string, compil
 	prove := &BlueprintProve{
 		SolveBlueprintID: solveID,
 		SolveBlueprint:   solve,
+		Schedule:         schedule,
 		HashName:         hashName,
 	}
 	proveID := compiler.AddBlueprint(prove)
@@ -437,7 +446,7 @@ func NewBlueprints(circuit gkrtypes.SerializableCircuit, hashName string, compil
 	}
 	getAssignmentID := compiler.AddBlueprint(getAssignment)
 
-	return gkrtypes.Blueprints{
+	return gkrcore.Blueprints{
 		SolveID:         solveID,
 		Solve:           solve,
 		ProveID:         proveID,