lang: Add unification optimizations

This adds a unification optimizations API, and uses it to optimize the embedded provisioner. With these turned on, type unification drops from around 1m45s to 2.5s which is a 40x speedup.
xlai89 · Mar 30, 2024 · 1b00af6 · 1b00af6
1 parent ddf1be6
commit 1b00af6
Show file tree

Hide file tree

Showing 6 changed files with 91 additions and 4 deletions.
diff --git a/cli/util/args.go b/cli/util/args.go
@@ -84,9 +84,10 @@ type LangArgs struct {
  OnlyDownload bool `arg:"--only-download" help:"stop after downloading any missing imports"`
  Update bool `arg:"--update" help:"update all dependencies to the latest versions"`
 
- OnlyUnify bool `arg:"--only-unify" help:"stop after type unification"`
- SkipUnify bool `arg:"--skip-unify" help:"skip type unification"`
- UnifySolver *string `arg:"--unify-name" help:"pick a specific unification solver"`
+ OnlyUnify bool `arg:"--only-unify" help:"stop after type unification"`
+ SkipUnify bool `arg:"--skip-unify" help:"skip type unification"`
+ UnifySolver *string `arg:"--unify-name" help:"pick a specific unification solver"`
+ UnifyOptimizations []string `arg:"--unify-optimizations" help:"list of unification optimizations to request (experts only)"`
 
  Depth int `arg:"--depth" default:"-1" help:"max recursion depth limit (-1 is unlimited)"`
 

diff --git a/lang/core/embedded/provisioner/provisioner.go b/lang/core/embedded/provisioner/provisioner.go
@@ -46,6 +46,7 @@ import (
  "github.com/purpleidea/mgmt/lang/embedded"
  "github.com/purpleidea/mgmt/lang/funcs/simple"
  "github.com/purpleidea/mgmt/lang/types"
+ "github.com/purpleidea/mgmt/lang/unification/simplesolver" // TODO: remove me!
  "github.com/purpleidea/mgmt/util"
  "github.com/purpleidea/mgmt/util/errwrap"
  "github.com/purpleidea/mgmt/util/password"
@@ -383,6 +384,12 @@ func (obj *provisioner) Customize(a interface{}) (*cli.RunArgs, error) {
 
  // Make any changes here that we want to...
  runArgs.RunLang.SkipUnify = true // speed things up for known good code
+ name := simplesolver.Name
+ // TODO: Remove these optimizations when the solver is faster overall.
+ runArgs.RunLang.UnifySolver = &name
+ runArgs.RunLang.UnifyOptimizations = []string{
+ simplesolver.OptimizationSkipFuncCmp,
+ }
  libConfig.TmpPrefix = true
  libConfig.NoPgp = true
 

diff --git a/lang/gapi/gapi.go b/lang/gapi/gapi.go
@@ -268,6 +268,10 @@ func (obj *GAPI) Cli(info *gapi.Info) (*gapi.Deploy, error) {
  if name := args.UnifySolver; name != nil && *name != "" {
  unificationStrategy[unification.StrategyNameKey] = *name
  }
+ if len(args.UnifyOptimizations) > 0 {
+ // TODO: use a query string parser instead?
+ unificationStrategy[unification.StrategyOptimizationsKey] = strings.Join(args.UnifyOptimizations, ",")
+ }
 
  if !args.SkipUnify {
  // apply type unification

diff --git a/lang/lang.go b/lang/lang.go
@@ -242,6 +242,7 @@ func (obj *Lang) Init(ctx context.Context) error {
 
  // apply type unification
  logf := func(format string, v ...interface{}) {
+ // TODO: Remove the masked logger here when unification is clean!
  if obj.Debug { // unification only has debug messages...
  obj.Logf("unification: "+format, v...)
  }

diff --git a/lang/unification/interfaces.go b/lang/unification/interfaces.go
@@ -45,6 +45,11 @@ const (
 
  // StrategyNameKey is the string key used when choosing a solver name.
  StrategyNameKey = "name"
+
+ // StrategyOptimizationsKey is the string key used to tell the solver
+ // about the specific optimizations you'd like to request. The format
+ // can be specific to each solver.
+ StrategyOptimizationsKey = "optimizations"
 )
 
 // Init contains some handles that are used to initialize every solver. Each

diff --git a/lang/unification/simplesolver/simplesolver.go b/lang/unification/simplesolver/simplesolver.go
@@ -32,6 +32,8 @@ package simplesolver
 import (
  "context"
  "fmt"
+ "sort"
+ "strings"
 
  "github.com/purpleidea/mgmt/lang/interfaces"
  "github.com/purpleidea/mgmt/lang/types"
@@ -43,6 +45,17 @@ const (
  // Name is the prefix for our solver log messages.
  Name = "simple"
 
+ // OptimizationSkipFuncCmp is the magic flag name to include the skip
+ // func cmp optimization which can speed up some simple programs. If
+ // this is specified, then OptimizationHeuristicalDrop is redundant.
+ OptimizationSkipFuncCmp = "skip-func-cmp"
+
+ // OptimizationHeuristicalDrop is the magic flag name to include to tell
+ // the solver to drop some func compares. This is a less aggressive form
+ // of OptimizationSkipFuncCmp. This is redundant if
+ // OptimizationSkipFuncCmp is true.
+ OptimizationHeuristicalDrop = "heuristical-drop"
+
  // AllowRecursion specifies whether we're allowed to use the recursive
  // solver or not. It uses an absurd amount of memory, and might hang
  // your system if a simple solution doesn't exist.
@@ -71,6 +84,19 @@ type SimpleInvariantSolver struct {
 
  Debug bool
  Logf func(format string, v ...interface{})
+
+ // skipFuncCmp tells the solver to skip the slow loop entirely. This may
+ // prevent some correct programs from completing type unification.
+ skipFuncCmp bool
+
+ // heuristicalDrop tells the solver to drop some func compares. This was
+ // determined heuristically and needs checking to see if it's even a
+ // sensible algorithmic approach. This is a less aggressive form of
+ // skipFuncCmp. This is redundant if skipFuncCmp is true.
+ heuristicalDrop bool
+
+ // zTotal is a heuristic counter to measure the size of the slow loop.
+ zTotal int
 }
 
 // Init contains some handles that are used to initialize the solver.
@@ -80,6 +106,22 @@ func (obj *SimpleInvariantSolver) Init(init *unification.Init) error {
  obj.Debug = init.Debug
  obj.Logf = init.Logf
 
+ optimizations, exists := init.Strategy[unification.StrategyOptimizationsKey]
+ if !exists {
+ return nil
+ }
+ // TODO: use a query string parser instead?
+ for _, x := range strings.Split(optimizations, ",") {
+ if x == OptimizationSkipFuncCmp {
+ obj.skipFuncCmp = true
+ continue
+ }
+ if x == OptimizationHeuristicalDrop {
+ obj.heuristicalDrop = true
+ continue
+ }
+ }
+
  return nil
 }
 
@@ -646,7 +688,12 @@ Loop:
  return false
  }
  // is there another EqualityWrapFuncInvariant with the same Expr1 pointer?
- for _, fn := range fnInvariants {
+ fnDone := make(map[int]struct{})
+ for z, fn := range fnInvariants {
+ if obj.skipFuncCmp {
+ break
+ }
+ obj.zTotal++
  // XXX: I think we're busy in this loop a lot.
  select {
  case <-ctx.Done():
@@ -693,6 +740,7 @@ Loop:
  eqInvariants = append(eqInvariants, newEq)
  // TODO: add it as a generator instead?
  equalities = append(equalities, newEq)
+ fnDone[z] = struct{}{} // XXX: heuristical drop
  }
 
  // both solved or both unsolved we skip
@@ -706,6 +754,7 @@ Loop:
  solved[rhsExpr] = lhsTyp // yay, we learned something!
  //used = append(used, i) // mark equality as used up when complete!
  obj.Logf("%s: solved partial rhs func arg equality", Name)
+ fnDone[z] = struct{}{} // XXX: heuristical drop
  } else if err := newTyp.Cmp(lhsTyp); err != nil {
  return nil, errwrap.Wrapf(err, "can't unify, invariant illogicality with partial rhs func arg equality")
  }
@@ -726,6 +775,7 @@ Loop:
  solved[lhsExpr] = rhsTyp // yay, we learned something!
  //used = append(used, i) // mark equality as used up when complete!
  obj.Logf("%s: solved partial lhs func arg equality", Name)
+ fnDone[z] = struct{}{} // XXX: heuristical drop
  } else if err := newTyp.Cmp(rhsTyp); err != nil {
  return nil, errwrap.Wrapf(err, "can't unify, invariant illogicality with partial lhs func arg equality")
  }
@@ -756,6 +806,7 @@ Loop:
  eqInvariants = append(eqInvariants, newEq)
  // TODO: add it as a generator instead?
  equalities = append(equalities, newEq)
+ fnDone[z] = struct{}{} // XXX: heuristical drop
  }
 
  // both solved or both unsolved we skip
@@ -769,6 +820,7 @@ Loop:
  solved[rhsExpr] = lhsTyp // yay, we learned something!
  //used = append(used, i) // mark equality as used up when complete!
  obj.Logf("%s: solved partial rhs func return equality", Name)
+ fnDone[z] = struct{}{} // XXX: heuristical drop
  } else if err := newTyp.Cmp(lhsTyp); err != nil {
  return nil, errwrap.Wrapf(err, "can't unify, invariant illogicality with partial rhs func return equality")
  }
@@ -789,6 +841,7 @@ Loop:
  solved[lhsExpr] = rhsTyp // yay, we learned something!
  //used = append(used, i) // mark equality as used up when complete!
  obj.Logf("%s: solved partial lhs func return equality", Name)
+ fnDone[z] = struct{}{} // XXX: heuristical drop
  } else if err := newTyp.Cmp(rhsTyp); err != nil {
  return nil, errwrap.Wrapf(err, "can't unify, invariant illogicality with partial lhs func return equality")
  }
@@ -800,6 +853,21 @@ Loop:
  }
  }
 
+ } // end big slow loop
+ if obj.heuristicalDrop {
+ fnDoneList := []int{}
+ for k := range fnDone {
+ fnDoneList = append(fnDoneList, k)
+ }
+ sort.Ints(fnDoneList)
+
+ for z := len(fnDoneList) - 1; z >= 0; z-- {
+ zx := fnDoneList[z] // delete index that was marked as done!
+ fnInvariants = append(fnInvariants[:zx], fnInvariants[zx+1:]...)
+ if obj.Debug {
+ obj.Logf("zTotal: %d, had: %d, removing: %d", obj.zTotal, len(fnInvariants), len(fnDoneList))
+ }
+ }
  }
 
  // can we solve anything?
@@ -1279,6 +1347,7 @@ Loop:
  }
  solutions = append(solutions, invar)
  }
+ obj.Logf("zTotal: %d", obj.zTotal)
  return &unification.InvariantSolution{
  Solutions: solutions,
  }, nil