solver.go

     1  // Copyright 2026 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package specexpr manipulates symbolic constraints on vector shapes.
     6  //
     7  // # Shapes
     8  //
     9  // A vector shape consists of a base type B (e.g., int, float), element width N
    10  // (8, 16, 32, or 64), and a vector width W (128, 256, 512, or scalable). It
    11  // also has a lane count L, which is the vector width / element width. These are
    12  // written like "Int32x4" or "Int32w128" or, for a scalable vector, "Int32s"
    13  // Masks are represented as vectors with a base type of "Mask", e.g.,
    14  // "Mask32x4".
    15  //
    16  // Scalar shapes consist only of a base type and an element width, e.g.,
    17  // "uint32".
    18  //
    19  // # Expressions
    20  //
    21  // Constraints are written as boolean expressions over shapes and a few basic
    22  // types.
    23  //
    24  // The primary expressions are:
    25  //
    26  //   - Variable ([a-zA-Z]+), such as x or xL
    27  //
    28  //   - Integers ([0-9]+)
    29  //
    30  //   - Shapes, written in the form given above, but where each component can be
    31  //     written as a bracketed expression, such as "Int32x{z/2}".
    32  //
    33  // These can be combined with operators * or /, or comparison operators =, >, <,
    34  // >=, <=. Arithmetic operators bind more tightly than comparison operators.
    35  //
    36  // # Example
    37  //
    38  // Consider a DotProductPairs function that takes two vectors x and y that have
    39  // the same shape and produces a vector z that has the same base type as x and
    40  // y, but has half as many elements, each of double the width. This can be
    41  // expressed as:
    42  //
    43  //	y=x
    44  //	z={xB}{xN*2}x{xL/2}
    45  //
    46  // # Width rounding
    47  //
    48  // The minimum vector width is 128 bits. Sometimes, operations would naturally
    49  // produce a width smaller than this, so hardware simply pads the vector out to
    50  // 128 bits. Shapes implement this behavior. For example, consider a "convert to
    51  // float32 operation" with constraints
    52  //
    53  //	z=Float32x{xL}
    54  //
    55  // If x is Float64x2, then z would naturally be Float32x2, but since this is
    56  // only 64 bits, the shape is "rounded" up to Float32x4.
    57  //
    58  // # Limitations
    59  //
    60  // The solver is intentionally simple. See [Solver] for a description of its
    61  // limitations. If you run up against its limitations, you're probably being too
    62  // clever.
    63  package specexpr
    64  
    65  import (
    66  	"cmp"
    67  	"container/heap"
    68  	"fmt"
    69  	"io"
    70  	"iter"
    71  	"log"
    72  	"maps"
    73  	"slices"
    74  	"strings"
    75  )
    76  
    77  // A Solver solves a set of constraints.
    78  //
    79  // This is a simple monotonic solver. It looks for a single order in which it
    80  // can resolve all constraints, assertions of the form "var=expr" are treated as
    81  // candidates for resolving the value of "var", and anything else is treated as
    82  // a boolean check. Variables that appear only on the right hand side of
    83  // assignments are "independent" and it will enumerate all possible values of
    84  // these variables. It never tries to invert any formulas and refuses to solve a
    85  // system with cycles. This is intentional to keep this solver fast: if your
    86  // formulas are cyclic, you're doing something too complicated.
    87  type Solver struct {
    88  	vars    map[Variable][]any
    89  	asserts []Expr
    90  	tracer  *tracer
    91  }
    92  
    93  // SetTrace enables emitting a solver trace to w.
    94  func (s *Solver) SetTrace(w io.Writer) {
    95  	if w == nil {
    96  		s.tracer = nil
    97  	} else {
    98  		s.tracer = &tracer{w: w}
    99  	}
   100  }
   101  
   102  // Declare declares a variable and its domain.
   103  //
   104  // Any "int" values in domain will be converted to [Int].
   105  func (s *Solver) Declare(v Variable, domain []any) {
   106  	if s.vars == nil {
   107  		s.vars = make(map[Variable][]any)
   108  	}
   109  	if _, ok := s.vars[v]; ok {
   110  		panic(v + " redeclared")
   111  	}
   112  	for i, d := range domain {
   113  		if d, ok := d.(int); ok {
   114  			domain[i] = Int(d)
   115  		}
   116  	}
   117  	s.vars[v] = domain
   118  }
   119  
   120  // Assign is a convenience for asserting that v=val.
   121  func (s *Solver) Assign(v Variable, val Expr) Variable {
   122  	s.Assert(&BinExpr{Op: OpEqual, X: v, Y: val})
   123  	return v
   124  }
   125  
   126  // Assert asserts a boolean condition must be true.
   127  func (s *Solver) Assert(cond Expr) {
   128  	s.asserts = append(s.asserts, cond)
   129  }
   130  
   131  func (s *Solver) Fprint(w io.Writer) {
   132  	for _, v := range slices.Sorted(maps.Keys(s.vars)) {
   133  		fmt.Fprintf(w, "%s in %v\n", v, s.vars[v])
   134  	}
   135  	for _, expr := range s.asserts {
   136  		fmt.Fprintf(w, "%s\n", expr)
   137  	}
   138  }
   139  
   140  // Bindings is a set of variable values.
   141  type Bindings struct {
   142  	varNames map[Variable]int // Shared between all solutions
   143  	vals     []any
   144  }
   145  
   146  // Get returns the value of v if resolved, or nil.
   147  func (b *Bindings) Get(v Variable) any {
   148  	vid, ok := b.varNames[v]
   149  	if !ok || vid >= len(b.vals) {
   150  		return nil
   151  	}
   152  	return b.vals[vid]
   153  }
   154  
   155  // All yields all variable bindings.
   156  func (b *Bindings) All() iter.Seq2[Variable, any] {
   157  	return func(yield func(Variable, any) bool) {
   158  		for _, varName := range slices.Sorted(maps.Keys(b.varNames)) {
   159  			vid := b.varNames[varName]
   160  			if vid < len(b.vals) && !yield(varName, b.vals[vid]) {
   161  				return
   162  			}
   163  		}
   164  	}
   165  }
   166  
   167  func (b *Bindings) String() string {
   168  	var buf strings.Builder
   169  	buf.WriteByte('{')
   170  	for v, val := range b.All() {
   171  		if buf.Len() > 1 {
   172  			buf.WriteByte(' ')
   173  		}
   174  		fmt.Fprintf(&buf, "%s=%v", v, val)
   175  	}
   176  	buf.WriteByte('}')
   177  	return buf.String()
   178  }
   179  
   180  // Solve yields all satisfying assignments of the variables in s.
   181  func (s *Solver) Solve() iter.Seq2[*Bindings, error] {
   182  	steps, err := s.topoSort()
   183  	if err != nil {
   184  		return func(yield func(*Bindings, error) bool) {
   185  			yield(nil, err)
   186  		}
   187  	}
   188  
   189  	// Assign variable indexes
   190  	varIDs := make(map[Variable]int)
   191  	for _, step := range steps {
   192  		if step.kind == solverStepCheck {
   193  			continue
   194  		}
   195  		if _, ok := varIDs[step.bind]; ok {
   196  			panic(fmt.Sprintf("variable %s resolved multiple times by solver sequence", step.bind))
   197  		}
   198  		varIDs[step.bind] = len(varIDs)
   199  	}
   200  
   201  	// If there are no solutions, then we report any evaluation errors. As soon
   202  	// as we yield anything, we set this to nil to indicate that.
   203  	errors := make(map[string]bool)
   204  	addErr := func(err error) {
   205  		if errors != nil {
   206  			errors[err.Error()] = true
   207  		}
   208  	}
   209  
   210  	// Walk solver steps
   211  	b := Bindings{
   212  		varNames: varIDs,
   213  		vals:     make([]any, 0, len(varIDs)),
   214  	}
   215  	pop := func() {
   216  		b.vals = b.vals[:len(b.vals)-1]
   217  	}
   218  	var visit func(steps []*solverStep, yield func(*Bindings, error) bool) bool
   219  	visit = func(steps []*solverStep, yield func(*Bindings, error) bool) bool {
   220  		if len(steps) == 0 {
   221  			errors = nil // Discard any errors
   222  			// Snapshot Bindings.
   223  			s.tracer.sat()
   224  			return yield(&Bindings{varNames: b.varNames, vals: slices.Clone(b.vals)}, nil)
   225  		}
   226  		step := steps[0]
   227  		steps = steps[1:]
   228  		switch step.kind {
   229  		case solverStepAssign:
   230  			val, err := step.expr.(*BinExpr).Y.eval(&b)
   231  			if err == nil {
   232  				if domain, ok := s.vars[step.bind]; ok && !slices.Contains(domain, val) {
   233  					err = fmt.Errorf("cannot assign %s=%v: not in domain", step.bind, val)
   234  				}
   235  			}
   236  			s.tracer.assign(step.expr, step.bind, val, err)
   237  			if err != nil {
   238  				addErr(err)
   239  				return true
   240  			}
   241  			b.vals = append(b.vals, val)
   242  			defer pop()
   243  			return visit(steps, yield)
   244  
   245  		case solverStepCheck:
   246  			val, err := step.expr.eval(&b)
   247  			s.tracer.check(step.expr, val, err)
   248  			if err != nil {
   249  				addErr(err)
   250  				return true
   251  			}
   252  			vBool, ok := val.(bool)
   253  			if !ok {
   254  				panic(fmt.Errorf("%s has type %T, expected bool", step.expr, val))
   255  			}
   256  			if vBool {
   257  				return visit(steps, yield)
   258  			}
   259  			return true
   260  
   261  		case solverStepIndep:
   262  			i := len(b.vals)
   263  			b.vals = append(b.vals, nil)
   264  			defer pop()
   265  			for _, val := range s.vars[step.bind] {
   266  				b.vals[i] = val
   267  				s.tracer.enter(step.bind, val)
   268  				if !visit(steps, yield) {
   269  					return false
   270  				}
   271  				s.tracer.exit()
   272  			}
   273  			return true
   274  		}
   275  		panic("bad step kind")
   276  	}
   277  	return func(yield func(*Bindings, error) bool) {
   278  		if visit(steps, yield) {
   279  			if len(errors) > 0 {
   280  				err := fmt.Errorf("%s", strings.Join(slices.Sorted(maps.Keys(errors)), "\n"))
   281  				yield(nil, err)
   282  			}
   283  		}
   284  	}
   285  }
   286  
   287  type solverStep struct {
   288  	kind solverStepKind
   289  	id   int
   290  	expr Expr
   291  	bind Variable // Variable to bind for solverStepAssign or solverStepIndep
   292  	hid  int      // heap index
   293  }
   294  
   295  type solverStepKind int
   296  
   297  const (
   298  	// solverStepIndep is a solverStep that simultaneously binds
   299  	// [solverStep.bind] to every possible value in bind's domain.
   300  	solverStepIndep solverStepKind = iota
   301  	// solverStepAssign is a solverStep where expr is an [OpEqual] [BinExpr]
   302  	// where the LHS is a Variable. It evaluates the RHS and assigns it to the
   303  	// variable. Variable must not have been bound by an earlier step (any
   304  	// subsequent OpEqual expressions for this variable should instead be a
   305  	// solverStepCheck).
   306  	solverStepAssign
   307  	// solverStepCheck is a solverStep that checks that expr is true and
   308  	// otherwise terminates the current solver branch.
   309  	solverStepCheck
   310  )
   311  
   312  func (s *solverStep) Compare(t *solverStep) int {
   313  	// Put assertions before independent variables because they may cut off paths
   314  	// before we have to enumerate values.
   315  	if s.kind != t.kind {
   316  		return cmp.Compare(s.kind, t.kind)
   317  	}
   318  	switch s.kind {
   319  	case solverStepCheck, solverStepAssign:
   320  		return cmp.Compare(s.id, t.id)
   321  	case solverStepIndep:
   322  		return cmp.Compare(s.bind, t.bind)
   323  	}
   324  	panic("bad solverStep kind")
   325  }
   326  
   327  type solverHeap []*solverStep
   328  
   329  func (sh solverHeap) Len() int { return len(sh) }
   330  
   331  func (sh solverHeap) Less(i, j int) bool {
   332  	return sh[i].Compare(sh[j]) < 0
   333  }
   334  
   335  func (sh solverHeap) Swap(i, j int) {
   336  	sh[i], sh[j] = sh[j], sh[i]
   337  	sh[i].hid, sh[j].hid = i, j
   338  }
   339  
   340  func (sh *solverHeap) Push(x any) {
   341  	item := x.(*solverStep)
   342  	item.hid = len(*sh)
   343  	*sh = append(*sh, item)
   344  }
   345  
   346  func (sh *solverHeap) Pop() any {
   347  	old := *sh
   348  	n := len(old)
   349  	item := old[n-1]
   350  	old[n-1] = nil
   351  	item.hid = -1
   352  	*sh = old[0 : n-1]
   353  	return item
   354  }
   355  
   356  func (s *Solver) topoSort() (order []*solverStep, err error) {
   357  	// This is a topo-sort with a few tricks: an assertion can be evaluated once
   358  	// all of its input variables are available, BUT a variable value can be
   359  	// resolved by potentially more than one assertion. Hence, we have a mix of
   360  	// "AND" and "OR" dependencies. For example, if we have:
   361  	//
   362  	//   x=Basic{xB, xN}
   363  	//   x=y
   364  	//
   365  	// Then x depends on "Basic{xB, xN}" OR "y" because we can assign x's value
   366  	// as soon as we resolve either of these. But resolving the first of these
   367  	// two expressions depends on xB AND xN.
   368  	//
   369  	// To handle this mix of "AND" and "OR" dependencies, we use a
   370  	// wavefront-style topo sort where we track the number of unresolved input
   371  	// variables to each assertion and whenever we resolve one of these inputs
   372  	// for the first time, we decrement that count. Once it reaches zero, that
   373  	// assertion is let out of the gate.
   374  	//
   375  	// The second trick is that in the set of possible next steps, we bias
   376  	// toward taking steps that are more likely to cut off a path and less
   377  	// likely to cause more fan-out.
   378  
   379  	var queue solverHeap
   380  	defs := make(map[Variable][]*solverStep)
   381  	uses := make(map[Variable]map[*solverStep]bool)
   382  	remaining := make(map[*solverStep]int)
   383  
   384  	isVarAssign := func(e Expr) (Variable, Expr, bool) {
   385  		switch e := e.(type) {
   386  		case *BinExpr:
   387  			if e.Op == OpEqual {
   388  				switch x := e.X.(type) {
   389  				case Variable:
   390  					return x, e.Y, true
   391  				}
   392  			}
   393  		}
   394  		return "", nil, false
   395  	}
   396  
   397  	for i, assert := range s.asserts {
   398  		var rhs Expr
   399  		// Wrap the assertion in a step, assigning an ID.
   400  		var step *solverStep
   401  		if def, val, ok := isVarAssign(assert); ok {
   402  			if def == val {
   403  				// x=x assertion. It's always true and will muck up the sorting,
   404  				// so throw it out.
   405  				continue
   406  			}
   407  			// Variable assignment
   408  			step = &solverStep{kind: solverStepAssign, id: i, expr: assert, hid: -1, bind: def}
   409  			// Record variable definition.
   410  			defs[def] = append(defs[def], step)
   411  			rhs = val
   412  		} else {
   413  			// Boolean check
   414  			step = &solverStep{kind: solverStepCheck, id: i, expr: assert, hid: -1}
   415  			rhs = assert
   416  		}
   417  
   418  		// Record variables this step depends on.
   419  		deps := 0
   420  		for use := range exprVars(rhs) {
   421  			if uses[use] == nil {
   422  				uses[use] = make(map[*solverStep]bool)
   423  			}
   424  			if !uses[use][step] {
   425  				uses[use][step] = true
   426  				deps++
   427  			}
   428  		}
   429  		if deps == 0 {
   430  			// Already solvable, enqueue it.
   431  			step.hid = len(queue)
   432  			queue = append(queue, step)
   433  		} else {
   434  			remaining[step] = deps
   435  		}
   436  	}
   437  
   438  	// Find the independent variables and also seed the frontier with them.
   439  	for v := range uses {
   440  		if defs[v] == nil {
   441  			if _, ok := s.vars[v]; !ok {
   442  				return nil, fmt.Errorf("no domain for independent variable %q", v)
   443  			}
   444  			queue = append(queue, &solverStep{kind: solverStepIndep, bind: v, hid: -1})
   445  		}
   446  	}
   447  	// Any variables that are declared but not referenced are also independent.
   448  	for v := range s.vars {
   449  		if uses[v] == nil && defs[v] == nil {
   450  			queue = append(queue, &solverStep{kind: solverStepIndep, bind: v, hid: -1})
   451  		}
   452  	}
   453  
   454  	// Drive the frontier. queue is the frontier of variables whose dependencies
   455  	// are all resolved, maintained in a heuristic order.
   456  	heap.Init(&queue)
   457  	for len(queue) > 0 {
   458  		step := queue[0]
   459  		heap.Pop(&queue)
   460  
   461  		order = append(order, step)
   462  
   463  		if step.kind == solverStepCheck {
   464  			continue
   465  		}
   466  
   467  		// This step resolved variable v.
   468  		v := step.bind
   469  
   470  		// Demote any other assignments of this variable to checks.
   471  		for _, def := range defs[v] {
   472  			if def != step {
   473  				def.kind = solverStepCheck
   474  				// Adjust heap
   475  				if def.hid != -1 {
   476  					heap.Fix(&queue, def.hid)
   477  				}
   478  			}
   479  		}
   480  		delete(defs, v)
   481  
   482  		// Check steps that depend on v.
   483  		for use := range uses[v] {
   484  			remaining[use]--
   485  			if remaining[use] == 0 {
   486  				// All variables used by this step are now resolved. Add it to the
   487  				// queue.
   488  				heap.Push(&queue, use)
   489  				delete(remaining, use)
   490  			}
   491  		}
   492  	}
   493  
   494  	if len(remaining) > 0 {
   495  		return nil, reportCycle(defs, remaining)
   496  	}
   497  	return order, nil
   498  }
   499  
   500  func reportCycle(defs map[Variable][]*solverStep, remaining map[*solverStep]int) error {
   501  	// There was a cycle. There could be more than one cycle, but report one.
   502  	// Start with the "minimum" remaining step for stability.
   503  	var step *solverStep
   504  	for rem := range remaining {
   505  		if rem.kind == solverStepAssign && (step == nil || rem.Compare(step) < 0) {
   506  			step = rem
   507  		}
   508  	}
   509  	// Walk forward through the graph, filtered to unresolved nodes.
   510  	var cycle []Expr
   511  	have := make(map[*solverStep]int)
   512  	var visit func(step *solverStep) error
   513  	visit = func(step *solverStep) error {
   514  		if step.kind != solverStepAssign || remaining[step] == 0 {
   515  			return nil
   516  		}
   517  		if i, ok := have[step]; ok {
   518  			// Found the cycle
   519  			return fmt.Errorf("cyclic requirements: %v", cycle[i:])
   520  		}
   521  		have[step] = len(cycle)
   522  		cycle = append(cycle, step.expr)
   523  		for v := range exprVars(step.expr.(*BinExpr).Y) {
   524  			for _, def := range defs[v] {
   525  				if err := visit(def); err != nil {
   526  					return err
   527  				}
   528  			}
   529  		}
   530  		delete(have, step)
   531  		cycle = cycle[:len(cycle)-1]
   532  		return nil
   533  	}
   534  	err := visit(step)
   535  	if err == nil {
   536  		log.Printf("remaining:")
   537  		for rem, count := range remaining {
   538  			log.Printf("  %v (%d)", rem, count)
   539  		}
   540  		log.Fatal("unresolved assertions, but failed to find a cycle")
   541  	}
   542  	return err
   543  }
   544
View as plain text