Support Sequential ParExec

This puts in place support for using a sequential ParExec implementation. This is a first step towards parallelising trace expansion. Specifically, there are two outstanding issues to resolve: firstly, we must be able to add expanded columns in arbitrary order; secondly, we need to actually make `ParExec` operate in parallel (which make require making trace thread-safe or something). The issue around the order in which expanded columns are added to the trace is problematic as the current design essentially expects them to be added in a specific order. Instead, we need assignments and input columns to know what their target column index is. Then, we expanding the trace, they can specify this.
Consensys · Jul 29, 2024 · b9c2786 · b9c2786
1 parent 5876153
commit b9c2786
Show file tree

Hide file tree

Showing 19 changed files with 665 additions and 19 deletions.
diff --git a/pkg/air/expr.go b/pkg/air/expr.go
@@ -46,6 +46,15 @@ func (p *Add) Context(schema sc.Schema) trace.Context {
 	return sc.JoinContexts[Expr](p.Args, schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Add) RequiredColumns() *util.SortedSet[uint] {
+	return util.UnionSortedSets(p.Args, func(e Expr) *util.SortedSet[uint] {
+		return e.RequiredColumns()
+	})
+}
+
 // Add two expressions together, producing a third.
 func (p *Add) Add(other Expr) Expr { return &Add{Args: []Expr{p, other}} }
 
@@ -75,6 +84,15 @@ func (p *Sub) Context(schema sc.Schema) trace.Context {
 	return sc.JoinContexts[Expr](p.Args, schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Sub) RequiredColumns() *util.SortedSet[uint] {
+	return util.UnionSortedSets(p.Args, func(e Expr) *util.SortedSet[uint] {
+		return e.RequiredColumns()
+	})
+}
+
 // Add two expressions together, producing a third.
 func (p *Sub) Add(other Expr) Expr { return &Add{Args: []Expr{p, other}} }
 
@@ -104,6 +122,15 @@ func (p *Mul) Context(schema sc.Schema) trace.Context {
 	return sc.JoinContexts[Expr](p.Args, schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Mul) RequiredColumns() *util.SortedSet[uint] {
+	return util.UnionSortedSets(p.Args, func(e Expr) *util.SortedSet[uint] {
+		return e.RequiredColumns()
+	})
+}
+
 // Add two expressions together, producing a third.
 func (p *Mul) Add(other Expr) Expr { return &Add{Args: []Expr{p, other}} }
 
@@ -156,6 +183,13 @@ func (p *Constant) Context(schema sc.Schema) trace.Context {
 	return trace.VoidContext()
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Constant) RequiredColumns() *util.SortedSet[uint] {
+	return util.NewSortedSet[uint]()
+}
+
 // Add two expressions together, producing a third.
 func (p *Constant) Add(other Expr) Expr { return &Add{Args: []Expr{p, other}} }
 
@@ -196,6 +230,16 @@ func (p *ColumnAccess) Context(schema sc.Schema) trace.Context {
 	return col.Context()
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *ColumnAccess) RequiredColumns() *util.SortedSet[uint] {
+	r := util.NewSortedSet[uint]()
+	r.Insert(p.Column)
+	// Done
+	return r
+}
+
 // Add two expressions together, producing a third.
 func (p *ColumnAccess) Add(other Expr) Expr { return &Add{Args: []Expr{p, other}} }
 

diff --git a/pkg/air/gadgets/normalisation.go b/pkg/air/gadgets/normalisation.go
@@ -89,6 +89,13 @@ func (e *Inverse) Context(schema sc.Schema) tr.Context {
 	return e.Expr.Context(schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (e *Inverse) RequiredColumns() *util.SortedSet[uint] {
+	return e.Expr.RequiredColumns()
+}
+
 func (e *Inverse) String() string {
 	return fmt.Sprintf("(inv %s)", e.Expr)
 }
diff --git a/pkg/hir/expr.go b/pkg/hir/expr.go
@@ -52,6 +52,15 @@ func (p *Add) Context(schema sc.Schema) trace.Context {
 	return sc.JoinContexts[Expr](p.Args, schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Add) RequiredColumns() *util.SortedSet[uint] {
+	return util.UnionSortedSets(p.Args, func(e Expr) *util.SortedSet[uint] {
+		return e.RequiredColumns()
+	})
+}
+
 // ============================================================================
 // Subtraction
 // ============================================================================
@@ -69,6 +78,15 @@ func (p *Sub) Context(schema sc.Schema) trace.Context {
 	return sc.JoinContexts[Expr](p.Args, schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Sub) RequiredColumns() *util.SortedSet[uint] {
+	return util.UnionSortedSets(p.Args, func(e Expr) *util.SortedSet[uint] {
+		return e.RequiredColumns()
+	})
+}
+
 // ============================================================================
 // Multiplication
 // ============================================================================
@@ -86,6 +104,15 @@ func (p *Mul) Context(schema sc.Schema) trace.Context {
 	return sc.JoinContexts[Expr](p.Args, schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Mul) RequiredColumns() *util.SortedSet[uint] {
+	return util.UnionSortedSets(p.Args, func(e Expr) *util.SortedSet[uint] {
+		return e.RequiredColumns()
+	})
+}
+
 // ============================================================================
 // Exponentiation
 // ============================================================================
@@ -106,6 +133,13 @@ func (p *Exp) Context(schema sc.Schema) trace.Context {
 	return p.Arg.Context(schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Exp) RequiredColumns() *util.SortedSet[uint] {
+	return p.Arg.RequiredColumns()
+}
+
 // ============================================================================
 // List
 // ============================================================================
@@ -123,6 +157,15 @@ func (p *List) Context(schema sc.Schema) trace.Context {
 	return sc.JoinContexts[Expr](p.Args, schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *List) RequiredColumns() *util.SortedSet[uint] {
+	return util.UnionSortedSets(p.Args, func(e Expr) *util.SortedSet[uint] {
+		return e.RequiredColumns()
+	})
+}
+
 // ============================================================================
 // Constant
 // ============================================================================
@@ -140,6 +183,13 @@ func (p *Constant) Context(schema sc.Schema) trace.Context {
 	return trace.VoidContext()
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Constant) RequiredColumns() *util.SortedSet[uint] {
+	return util.NewSortedSet[uint]()
+}
+
 // ============================================================================
 // IfZero
 // ============================================================================
@@ -190,6 +240,23 @@ func (p *IfZero) Context(schema sc.Schema) trace.Context {
 	return sc.JoinContexts[Expr](args, schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *IfZero) RequiredColumns() *util.SortedSet[uint] {
+	set := p.Condition.RequiredColumns()
+	// Include true branch (if applicable)
+	if p.TrueBranch != nil {
+		set.InsertSorted(p.TrueBranch.RequiredColumns())
+	}
+	// Include false branch (if applicable)
+	if p.FalseBranch != nil {
+		set.InsertSorted(p.FalseBranch.RequiredColumns())
+	}
+	// Done
+	return set
+}
+
 // ============================================================================
 // Normalise
 // ============================================================================
@@ -210,6 +277,13 @@ func (p *Normalise) Context(schema sc.Schema) trace.Context {
 	return p.Arg.Context(schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Normalise) RequiredColumns() *util.SortedSet[uint] {
+	return p.Arg.RequiredColumns()
+}
+
 // ============================================================================
 // ColumnAccess
 // ============================================================================
@@ -242,3 +316,13 @@ func (p *ColumnAccess) Context(schema sc.Schema) trace.Context {
 	col := schema.Columns().Nth(p.Column)
 	return col.Context()
 }
+
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *ColumnAccess) RequiredColumns() *util.SortedSet[uint] {
+	r := util.NewSortedSet[uint]()
+	r.Insert(p.Column)
+	// Done
+	return r
+}
diff --git a/pkg/hir/util.go b/pkg/hir/util.go
@@ -50,6 +50,13 @@ func (p ZeroArrayTest) Context(schema sc.Schema) trace.Context {
 	return p.Expr.Context(schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p ZeroArrayTest) RequiredColumns() *util.SortedSet[uint] {
+	return p.Expr.RequiredColumns()
+}
+
 // ============================================================================
 // UnitExpr
 // ============================================================================
@@ -98,3 +105,10 @@ func (e UnitExpr) Bounds() util.Bounds {
 func (e UnitExpr) Context(schema sc.Schema) trace.Context {
 	return e.expr.Context(schema)
 }
+
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (e UnitExpr) RequiredColumns() *util.SortedSet[uint] {
+	return e.expr.RequiredColumns()
+}
diff --git a/pkg/mir/expr.go b/pkg/mir/expr.go
@@ -36,6 +36,15 @@ func (p *Add) Context(schema sc.Schema) trace.Context {
 	return sc.JoinContexts[Expr](p.Args, schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Add) RequiredColumns() *util.SortedSet[uint] {
+	return util.UnionSortedSets(p.Args, func(e Expr) *util.SortedSet[uint] {
+		return e.RequiredColumns()
+	})
+}
+
 // ============================================================================
 // Subtraction
 // ============================================================================
@@ -53,6 +62,15 @@ func (p *Sub) Context(schema sc.Schema) trace.Context {
 	return sc.JoinContexts[Expr](p.Args, schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Sub) RequiredColumns() *util.SortedSet[uint] {
+	return util.UnionSortedSets(p.Args, func(e Expr) *util.SortedSet[uint] {
+		return e.RequiredColumns()
+	})
+}
+
 // ============================================================================
 // Multiplication
 // ============================================================================
@@ -70,6 +88,15 @@ func (p *Mul) Context(schema sc.Schema) trace.Context {
 	return sc.JoinContexts[Expr](p.Args, schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Mul) RequiredColumns() *util.SortedSet[uint] {
+	return util.UnionSortedSets(p.Args, func(e Expr) *util.SortedSet[uint] {
+		return e.RequiredColumns()
+	})
+}
+
 // ============================================================================
 // Exponentiation
 // ============================================================================
@@ -90,6 +117,13 @@ func (p *Exp) Context(schema sc.Schema) trace.Context {
 	return p.Arg.Context(schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Exp) RequiredColumns() *util.SortedSet[uint] {
+	return p.Arg.RequiredColumns()
+}
+
 // ============================================================================
 // Constant
 // ============================================================================
@@ -107,6 +141,13 @@ func (p *Constant) Context(schema sc.Schema) trace.Context {
 	return trace.VoidContext()
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Constant) RequiredColumns() *util.SortedSet[uint] {
+	return util.NewSortedSet[uint]()
+}
+
 // ============================================================================
 // Normalise
 // ============================================================================
@@ -125,6 +166,13 @@ func (p *Normalise) Context(schema sc.Schema) trace.Context {
 	return p.Arg.Context(schema)
 }
 
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *Normalise) RequiredColumns() *util.SortedSet[uint] {
+	return p.Arg.RequiredColumns()
+}
+
 // ============================================================================
 // ColumnAccess
 // ============================================================================
@@ -157,3 +205,13 @@ func (p *ColumnAccess) Context(schema sc.Schema) trace.Context {
 	col := schema.Columns().Nth(p.Column)
 	return col.Context()
 }
+
+// RequiredColumns returns the set of columns on which this term depends.
+// That is, columns whose values may be accessed when evaluating this term
+// on a given trace.
+func (p *ColumnAccess) RequiredColumns() *util.SortedSet[uint] {
+	r := util.NewSortedSet[uint]()
+	r.Insert(p.Column)
+	// Done
+	return r
+}
diff --git a/pkg/schema/assignment/byte_decomposition.go b/pkg/schema/assignment/byte_decomposition.go
@@ -99,6 +99,12 @@ func (p *ByteDecomposition) RequiredSpillage() uint {
 	return uint(0)
 }
 
+// Dependencies returns the set of columns that this assignment depends upon.
+// That can include both input columns, as well as other computed columns.
+func (p *ByteDecomposition) Dependencies() []uint {
+	return []uint{p.source}
+}
+
 // Decompose a given element into n bytes in little endian form.  For example,
 // decomposing 41b into 2 bytes gives [0x1b,0x04].
 func decomposeIntoBytes(val *fr.Element, n int) []*fr.Element {

diff --git a/pkg/schema/assignment/computed_column.go b/pkg/schema/assignment/computed_column.go
@@ -107,3 +107,9 @@ func (p *ComputedColumn[E]) ExpandTrace(tr trace.Trace) error {
 	// Done
 	return nil
 }
+
+// Dependencies returns the set of columns that this assignment depends upon.
+// That can include both input columns, as well as other computed columns.
+func (p *ComputedColumn[E]) Dependencies() []uint {
+	return *p.expr.RequiredColumns()
+}
diff --git a/pkg/schema/assignment/interleave.go b/pkg/schema/assignment/interleave.go
@@ -114,3 +114,9 @@ func (p *Interleaving) ExpandTrace(tr tr.Trace) error {
 	//
 	return nil
 }
+
+// Dependencies returns the set of columns that this assignment depends upon.
+// That can include both input columns, as well as other computed columns.
+func (p *Interleaving) Dependencies() []uint {
+	return p.sources
+}