From c69b20dae97ffe346891ce86deda0242f3785ae2 Mon Sep 17 00:00:00 2001 From: Ryan Berger Date: Wed, 19 Apr 2023 22:13:04 -0600 Subject: [PATCH 1/5] cmd/compile: add reassociate ssa pass to rebalance commutative operations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Currently the compiler groups expressions with commutative operations such as a + b + c + d as so: (a + (b + (c + d))) which is suboptimal for CPU instruction pipelining. This pass balances commutative expressions as shown above to (a + b) + (c + d) to optimally pipeline them. It also attempts to reassociate constants to as far right of the commutative expression as possible for better constant folding opportunities. Below is a benchmark from crypto/md5 on an MacBook Pro M2: trunk reassociate Hash1K-8 433.7Mi ± 0% 499.4Mi ± 4% +15.17% (p=0.000 n=10) Hash8K-8 454.3Mi ± 1% 524.9Mi ± 1% +15.53% (p=0.000 n=10) .... geomean 284.4Mi 327.5Mi +15.15% Other CPU architectures tried showed very little change (+/-1%) on this particular benchmark but tight mathematical code stands to gain greatly from this optimization Fixes #49331 --- src/cmd/compile/internal/ssa/compile.go | 5 +- src/cmd/compile/internal/ssa/reassociate.go | 190 +++++++++++++++++++ src/cmd/compile/internal/ssa/shortcircuit.go | 2 +- 3 files changed, 194 insertions(+), 3 deletions(-) create mode 100644 src/cmd/compile/internal/ssa/reassociate.go diff --git a/src/cmd/compile/internal/ssa/compile.go b/src/cmd/compile/internal/ssa/compile.go index 8618cf34cd..d353535625 100644 --- a/src/cmd/compile/internal/ssa/compile.go +++ b/src/cmd/compile/internal/ssa/compile.go @@ -463,8 +463,9 @@ var passes = [...]pass{ {name: "short circuit", fn: shortcircuit}, {name: "decompose user", fn: decomposeUser, required: true}, {name: "pre-opt deadcode", fn: deadcode}, - {name: "opt", fn: opt, required: true}, // NB: some generic rules know the name of the opt pass. TODO: split required rules and optimizing rules - {name: "zero arg cse", fn: zcse, required: true}, // required to merge OpSB values + {name: "opt", fn: opt, required: true}, // NB: some generic rules know the name of the opt pass. TODO: split required rules and optimizing rules + {name: "zero arg cse", fn: zcse, required: true}, // required to merge OpSB values + {name: "reassociate", fn: reassociate}, {name: "opt deadcode", fn: deadcode, required: true}, // remove any blocks orphaned during opt {name: "generic cse", fn: cse}, {name: "phiopt", fn: phiopt}, diff --git a/src/cmd/compile/internal/ssa/reassociate.go b/src/cmd/compile/internal/ssa/reassociate.go new file mode 100644 index 0000000000..42a5e4f135 --- /dev/null +++ b/src/cmd/compile/internal/ssa/reassociate.go @@ -0,0 +1,190 @@ +// Copyright 2023 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package ssa + +import ( + "fmt" + "sort" +) + +// balanceExprTree repurposes all nodes and leafs into a +// balanced expression tree +func balanceExprTree(v *Value, visited map[*Value]bool, nodes, leafs []*Value) { + // reset all arguments of nodes to help rebalancing + for i, n := range nodes { + n.reset(n.Op) + + // sometimes nodes in the tree are in different blocks + // so pull them in into a common block (v's block) + // to make sure nodes don't end up dominating their leaves + if v.Block != n.Block { + copied := n.copyInto(v.Block) + n.Op = OpInvalid + visited[n] = true // "revisit" the copied node + nodes[i] = copied + } + } + + // we bfs'ed through the nodes in reverse topological order + // (expression dominated by all others to expression dominated by none of the others), + // we want to rebuild the tree reverse topological order + for i, j := 0, len(nodes)-1; i <= j; i, j = i+1, j-1 { + nodes[i], nodes[j] = nodes[j], nodes[i] + } + + // push all leafs which are constants as far off to the + // right as possible to give the constant folder more opportunities + sort.Slice(leafs, func(i, j int) bool { + switch leafs[j].Op { + case OpConst8, OpConst16, OpConst32, OpConst64: + return false + default: + return true + } + }) + + // build tree in reverse topological order + for i := 0; i < len(nodes); i++ { + if len(leafs) < 2 { // we need at least two leafs per node, balance went very wrong + panic(fmt.Sprint("leafs needs to be >= 2, got", len(leafs))) + } + + // Take two leaves out and attach them to a node, + // use the node as a new leaf in the "next layer" of the tree + nodes[i].AddArg2(leafs[0], leafs[1]) + leafs = append(leafs[2:], nodes[i]) + } +} + +func isOr(op Op) bool { + switch op { + case OpOr8, OpOr16, OpOr32, OpOr64: + return true + default: + return false + } +} + +// probablyMemcombine helps find a pattern of leaves that form +// a load that can be widened which looks like: +// +// (l | l << 8 | l << 18 | l << 24) +// +// which cannot be rebalanced or else it won't fire rewrite rules +func probablyMemcombine(op Op, leafs []*Value) bool { + if !isOr(op) { + return false + } + + lshCount := 0 + for _, l := range leafs { + switch l.Op { + case OpLsh8x8, OpLsh8x16, OpLsh8x32, OpLsh8x64, + OpLsh16x8, OpLsh16x16, OpLsh16x32, OpLsh16x64, + OpLsh32x8, OpLsh32x16, OpLsh32x32, OpLsh32x64, + OpLsh64x8, OpLsh64x16, OpLsh64x32, OpLsh64x64: + lshCount++ + } + } + + return lshCount == len(leafs)-1 +} + +// rebalance balances associative computation to better help CPU instruction pipelining (#49331) +// and groups constants together catch more constant folding opportunities. +// +// a + b + c + d compiles to to (a + (b + (c + d)) which is an unbalanced expression tree +// that looks like +// +// - (v1) +// / \ +// a + (v2) +// / \ +// b + (v3) +// / \ +// c d +// +// Which is suboptimal since it requires the CPU to compute v3 before fetching it use its result in +// v2, and v2 before its use in v1 +// This optimization rebalances the expression tree to: +// +// - (v1) +// / \ +// +// (v2) + + (v3) +// +// / \ / \ +// a b c d +// +// Which removes such dependencies and frees up the CPU pipeline. +// +// The above optimization is a good starting point for other sorts of operations such as +// turning a + a + a => 3*a, cancelling pairs a + (-a), collecting up common factors TODO(ryan-berger) +func rebalance(v *Value, visited map[*Value]bool) { + // We cannot apply this optimization to non-commutative operations, + // values that have more than one use, or non-binary ops (would need more log math). + // Try and save time by not revisiting nodes + if visited[v] || !(v.Uses == 1 && opcodeTable[v.Op].commutative) || len(v.Args) > 2 { + return + } + + // The smallest possible rebalanceable expression has 3 nodes and 4 leafs, + // so preallocate the lists to save time if it is not rebalanceable + leafs := make([]*Value, 0, 4) + nodes := make([]*Value, 0, 3) + + // Do a bfs on v to keep a nice reverse topological order + haystack := []*Value{v} + for len(haystack) != 0 { + nextHaystack := make([]*Value, 0, len(v.Args)*len(haystack)) + for _, needle := range haystack { + // if we are searching a value, it must be a node so add it to our node list + nodes = append(nodes, needle) + + // Only visit nodes. Leafs may be rebalancable for a different op type + visited[needle] = true + + for _, a := range needle.Args { + // If the ops aren't the same or have more than one use it must be a leaf. + if a.Op != v.Op || a.Uses != 1 { + leafs = append(leafs, a) + continue + } + + // nodes in the tree now hold the invariants that: + // - they are of a common associative operation as the rest of the tree + // - they have only a single use (this invariant could be removed with further analysis TODO(ryan-berger) + nextHaystack = append(nextHaystack, a) + } + } + haystack = nextHaystack + } + + // we need at least 4 leafs for this expression to be rebalanceable, + // and we can't balance a potential load widening (memcombine) + if len(leafs) < 4 || probablyMemcombine(v.Op, leafs) { + return + } + + balanceExprTree(v, visited, nodes, leafs) +} + +// reassociate balances trees of commutative computation +// to better group expressions for better constant folding, +// cse, etc. +func reassociate(f *Func) { + visited := make(map[*Value]bool) + + for _, b := range f.Postorder() { + for i := len(b.Values) - 1; i >= 0; i-- { + val := b.Values[i] + rebalance(val, visited) + } + } + + for k := range visited { + delete(visited, k) + } +} diff --git a/src/cmd/compile/internal/ssa/shortcircuit.go b/src/cmd/compile/internal/ssa/shortcircuit.go index d7d0b6fe33..9ee55941da 100644 --- a/src/cmd/compile/internal/ssa/shortcircuit.go +++ b/src/cmd/compile/internal/ssa/shortcircuit.go @@ -499,7 +499,7 @@ func (v *Value) moveTo(dst *Block, i int) { } src := v.Block if src.Values[i] != v { - v.Fatalf("moveTo bad index %d", v, i) + v.Fatalf("moveTo bad index %d", i) } if src == dst { return From c7ac24f04d09e70a83e480aa163f0641571aeae7 Mon Sep 17 00:00:00 2001 From: Ryan Berger Date: Fri, 5 May 2023 16:38:50 -0600 Subject: [PATCH 2/5] remove badly fmt'd ASCII art from comments --- src/cmd/compile/internal/ssa/reassociate.go | 22 ++------------------- 1 file changed, 2 insertions(+), 20 deletions(-) diff --git a/src/cmd/compile/internal/ssa/reassociate.go b/src/cmd/compile/internal/ssa/reassociate.go index 42a5e4f135..56d3c31132 100644 --- a/src/cmd/compile/internal/ssa/reassociate.go +++ b/src/cmd/compile/internal/ssa/reassociate.go @@ -96,29 +96,11 @@ func probablyMemcombine(op Op, leafs []*Value) bool { // and groups constants together catch more constant folding opportunities. // // a + b + c + d compiles to to (a + (b + (c + d)) which is an unbalanced expression tree -// that looks like -// -// - (v1) -// / \ -// a + (v2) -// / \ -// b + (v3) -// / \ -// c d -// // Which is suboptimal since it requires the CPU to compute v3 before fetching it use its result in // v2, and v2 before its use in v1 -// This optimization rebalances the expression tree to: // -// - (v1) -// / \ -// -// (v2) + + (v3) -// -// / \ / \ -// a b c d -// -// Which removes such dependencies and frees up the CPU pipeline. +// This optimization rebalances this expression tree to look like (a + b) + (c + d) , +// which removes such dependencies and frees up the CPU pipeline. // // The above optimization is a good starting point for other sorts of operations such as // turning a + a + a => 3*a, cancelling pairs a + (-a), collecting up common factors TODO(ryan-berger) From 5ad3fc9629c25100f5036be1b1205a90a000fbbd Mon Sep 17 00:00:00 2001 From: Ryan Berger Date: Sun, 7 May 2023 21:56:06 -0600 Subject: [PATCH 3/5] add basic codegen test, fix binary tree rebalancing --- src/cmd/compile/internal/ssa/reassociate.go | 63 +++++++++++---------- test/codegen/reassociate.go | 18 ++++++ 2 files changed, 50 insertions(+), 31 deletions(-) create mode 100644 test/codegen/reassociate.go diff --git a/src/cmd/compile/internal/ssa/reassociate.go b/src/cmd/compile/internal/ssa/reassociate.go index 56d3c31132..e11f5b7bb1 100644 --- a/src/cmd/compile/internal/ssa/reassociate.go +++ b/src/cmd/compile/internal/ssa/reassociate.go @@ -4,13 +4,10 @@ package ssa -import ( - "fmt" - "sort" -) - -// balanceExprTree repurposes all nodes and leafs into a -// balanced expression tree +// balanceExprTree repurposes all nodes and leafs into a well-balanced expression tree. +// It doesn't truly balance the tree in the sense of a BST, rather it +// prioritizes pairing up innermost (rightmost) expressions and their results and only +// pairing results of outermost (leftmost) expressions up with them when no other nice pairing exists func balanceExprTree(v *Value, visited map[*Value]bool, nodes, leafs []*Value) { // reset all arguments of nodes to help rebalancing for i, n := range nodes { @@ -34,27 +31,27 @@ func balanceExprTree(v *Value, visited map[*Value]bool, nodes, leafs []*Value) { nodes[i], nodes[j] = nodes[j], nodes[i] } - // push all leafs which are constants as far off to the - // right as possible to give the constant folder more opportunities - sort.Slice(leafs, func(i, j int) bool { - switch leafs[j].Op { - case OpConst8, OpConst16, OpConst32, OpConst64: - return false - default: - return true + // rebuild expression trees from the bottom up, prioritizing + // right grouping. + // if the number of leaves is not even, skip the first leaf + // and add it to be paired up later + i := 0 + subTrees := leafs + for len(subTrees) != 1 { + nextSubTrees := make([]*Value, 0, (len(subTrees)+1)/2) + + start := len(subTrees)%2 + if start != 0 { + nextSubTrees = append(nextSubTrees, subTrees[0]) } - }) - - // build tree in reverse topological order - for i := 0; i < len(nodes); i++ { - if len(leafs) < 2 { // we need at least two leafs per node, balance went very wrong - panic(fmt.Sprint("leafs needs to be >= 2, got", len(leafs))) + + for j := start; j < len(subTrees)-1; j+=2 { + nodes[i].AddArg2(subTrees[j], subTrees[j+1]) + nextSubTrees = append(nextSubTrees, nodes[i]) + i++ } - - // Take two leaves out and attach them to a node, - // use the node as a new leaf in the "next layer" of the tree - nodes[i].AddArg2(leafs[0], leafs[1]) - leafs = append(leafs[2:], nodes[i]) + + subTrees = nextSubTrees } } @@ -72,7 +69,7 @@ func isOr(op Op) bool { // // (l | l << 8 | l << 18 | l << 24) // -// which cannot be rebalanced or else it won't fire rewrite rules +// which cannot be rebalanced or else it won't fire load widening rewrite rules func probablyMemcombine(op Op, leafs []*Value) bool { if !isOr(op) { return false @@ -89,7 +86,11 @@ func probablyMemcombine(op Op, leafs []*Value) bool { } } - return lshCount == len(leafs)-1 + // there are a few algorithms in the std lib expressed as two 32 bit loads + // which can get turned into a 64 bit load + // conservatively estimate that if there are more shifts than not then it is + // some sort of load waiting to be widened + return lshCount > len(leafs)/2 } // rebalance balances associative computation to better help CPU instruction pipelining (#49331) @@ -145,7 +146,7 @@ func rebalance(v *Value, visited map[*Value]bool) { } // we need at least 4 leafs for this expression to be rebalanceable, - // and we can't balance a potential load widening (memcombine) + // and we can't balance a potential load widening (see memcombine) if len(leafs) < 4 || probablyMemcombine(v.Op, leafs) { return } @@ -154,8 +155,8 @@ func rebalance(v *Value, visited map[*Value]bool) { } // reassociate balances trees of commutative computation -// to better group expressions for better constant folding, -// cse, etc. +// to better group expressions to expose easy optimizations in +// cse, cancelling/counting/factoring expressions, etc. func reassociate(f *Func) { visited := make(map[*Value]bool) diff --git a/test/codegen/reassociate.go b/test/codegen/reassociate.go new file mode 100644 index 0000000000..1309d9dc0a --- /dev/null +++ b/test/codegen/reassociate.go @@ -0,0 +1,18 @@ +// asmcheck + +package codegen + +// reassociateAddition expects very specific sequence of registers +// of the form: +// R2 += R3 +// R1 += R0 +// R1 += R2 +func reassociateAddition(a, b, c, d int) int { + // arm64:`ADD\tR2,\sR3,\sR2` + x := b + a + // arm64:`ADD\tR0,\sR1,\sR1` + y := x + c + // arm64:`ADD\tR1,\sR2,\sR0` + z := y + d + return z +} \ No newline at end of file From 1c2d14836e9b89ec42784cf2854b0ef19118511c Mon Sep 17 00:00:00 2001 From: Ryan Berger Date: Mon, 8 May 2023 08:14:03 -0600 Subject: [PATCH 4/5] use correct plural of leaf, fix license file --- src/cmd/compile/internal/ssa/reassociate.go | 24 ++++++++++----------- test/codegen/reassociate.go | 4 ++++ 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/src/cmd/compile/internal/ssa/reassociate.go b/src/cmd/compile/internal/ssa/reassociate.go index e11f5b7bb1..3363d778df 100644 --- a/src/cmd/compile/internal/ssa/reassociate.go +++ b/src/cmd/compile/internal/ssa/reassociate.go @@ -4,11 +4,11 @@ package ssa -// balanceExprTree repurposes all nodes and leafs into a well-balanced expression tree. +// balanceExprTree repurposes all nodes and leaves into a well-balanced expression tree. // It doesn't truly balance the tree in the sense of a BST, rather it // prioritizes pairing up innermost (rightmost) expressions and their results and only // pairing results of outermost (leftmost) expressions up with them when no other nice pairing exists -func balanceExprTree(v *Value, visited map[*Value]bool, nodes, leafs []*Value) { +func balanceExprTree(v *Value, visited map[*Value]bool, nodes, leaves []*Value) { // reset all arguments of nodes to help rebalancing for i, n := range nodes { n.reset(n.Op) @@ -36,7 +36,7 @@ func balanceExprTree(v *Value, visited map[*Value]bool, nodes, leafs []*Value) { // if the number of leaves is not even, skip the first leaf // and add it to be paired up later i := 0 - subTrees := leafs + subTrees := leaves for len(subTrees) != 1 { nextSubTrees := make([]*Value, 0, (len(subTrees)+1)/2) @@ -70,13 +70,13 @@ func isOr(op Op) bool { // (l | l << 8 | l << 18 | l << 24) // // which cannot be rebalanced or else it won't fire load widening rewrite rules -func probablyMemcombine(op Op, leafs []*Value) bool { +func probablyMemcombine(op Op, leaves []*Value) bool { if !isOr(op) { return false } lshCount := 0 - for _, l := range leafs { + for _, l := range leaves { switch l.Op { case OpLsh8x8, OpLsh8x16, OpLsh8x32, OpLsh8x64, OpLsh16x8, OpLsh16x16, OpLsh16x32, OpLsh16x64, @@ -90,7 +90,7 @@ func probablyMemcombine(op Op, leafs []*Value) bool { // which can get turned into a 64 bit load // conservatively estimate that if there are more shifts than not then it is // some sort of load waiting to be widened - return lshCount > len(leafs)/2 + return lshCount > len(leaves)/2 } // rebalance balances associative computation to better help CPU instruction pipelining (#49331) @@ -113,9 +113,9 @@ func rebalance(v *Value, visited map[*Value]bool) { return } - // The smallest possible rebalanceable expression has 3 nodes and 4 leafs, + // The smallest possible rebalanceable expression has 3 nodes and 4 leaves, // so preallocate the lists to save time if it is not rebalanceable - leafs := make([]*Value, 0, 4) + leaves := make([]*Value, 0, 4) nodes := make([]*Value, 0, 3) // Do a bfs on v to keep a nice reverse topological order @@ -132,7 +132,7 @@ func rebalance(v *Value, visited map[*Value]bool) { for _, a := range needle.Args { // If the ops aren't the same or have more than one use it must be a leaf. if a.Op != v.Op || a.Uses != 1 { - leafs = append(leafs, a) + leaves = append(leaves, a) continue } @@ -145,13 +145,13 @@ func rebalance(v *Value, visited map[*Value]bool) { haystack = nextHaystack } - // we need at least 4 leafs for this expression to be rebalanceable, + // we need at least 4 leaves for this expression to be rebalanceable, // and we can't balance a potential load widening (see memcombine) - if len(leafs) < 4 || probablyMemcombine(v.Op, leafs) { + if len(leaves) < 4 || probablyMemcombine(v.Op, leaves) { return } - balanceExprTree(v, visited, nodes, leafs) + balanceExprTree(v, visited, nodes, leaves) } // reassociate balances trees of commutative computation diff --git a/test/codegen/reassociate.go b/test/codegen/reassociate.go index 1309d9dc0a..9fa795b530 100644 --- a/test/codegen/reassociate.go +++ b/test/codegen/reassociate.go @@ -1,5 +1,9 @@ // asmcheck +// Copyright 2023 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + package codegen // reassociateAddition expects very specific sequence of registers From 4846c63a974a931498197cee71e89fd90717a54a Mon Sep 17 00:00:00 2001 From: Ryan Berger Date: Mon, 8 May 2023 15:16:56 -0600 Subject: [PATCH 5/5] simplify bfs, other code review comments --- src/cmd/compile/internal/ssa/reassociate.go | 106 +++++++++----------- 1 file changed, 50 insertions(+), 56 deletions(-) diff --git a/src/cmd/compile/internal/ssa/reassociate.go b/src/cmd/compile/internal/ssa/reassociate.go index 3363d778df..2a6a23776e 100644 --- a/src/cmd/compile/internal/ssa/reassociate.go +++ b/src/cmd/compile/internal/ssa/reassociate.go @@ -4,22 +4,36 @@ package ssa +// reassociate balances trees of commutative computation +// to better group expressions to expose easy optimizations in +// cse, cancelling/counting/factoring expressions, etc. +func reassociate(f *Func) { + visited := f.newSparseSet(f.NumValues()) + + for _, b := range f.Postorder() { + for i := len(b.Values) - 1; i >= 0; i-- { + val := b.Values[i] + rebalance(val, visited) + } + } +} + // balanceExprTree repurposes all nodes and leaves into a well-balanced expression tree. -// It doesn't truly balance the tree in the sense of a BST, rather it +// It doesn't truly balance the tree in the sense of a BST, rather it // prioritizes pairing up innermost (rightmost) expressions and their results and only -// pairing results of outermost (leftmost) expressions up with them when no other nice pairing exists -func balanceExprTree(v *Value, visited map[*Value]bool, nodes, leaves []*Value) { +// pairing results of outermost (leftmost) expressions up with them when no other nice pairing exists +func balanceExprTree(v *Value, visited *sparseSet, nodes, leaves []*Value) { // reset all arguments of nodes to help rebalancing for i, n := range nodes { n.reset(n.Op) // sometimes nodes in the tree are in different blocks // so pull them in into a common block (v's block) - // to make sure nodes don't end up dominating their leaves + // to make sure nodes don't end up dominating their leaves TODO(ryan-berger), not necessary if v.Block != n.Block { copied := n.copyInto(v.Block) n.Op = OpInvalid - visited[n] = true // "revisit" the copied node + visited.add(copied.ID) // "revisit" the copied node nodes[i] = copied } } @@ -33,24 +47,24 @@ func balanceExprTree(v *Value, visited map[*Value]bool, nodes, leaves []*Value) // rebuild expression trees from the bottom up, prioritizing // right grouping. - // if the number of leaves is not even, skip the first leaf + // if the number of leaves is not even, skip the first leaf // and add it to be paired up later i := 0 subTrees := leaves for len(subTrees) != 1 { nextSubTrees := make([]*Value, 0, (len(subTrees)+1)/2) - - start := len(subTrees)%2 + + start := len(subTrees) % 2 if start != 0 { nextSubTrees = append(nextSubTrees, subTrees[0]) } - - for j := start; j < len(subTrees)-1; j+=2 { + + for j := start; j < len(subTrees)-1; j += 2 { nodes[i].AddArg2(subTrees[j], subTrees[j+1]) nextSubTrees = append(nextSubTrees, nodes[i]) i++ } - + subTrees = nextSubTrees } } @@ -96,78 +110,58 @@ func probablyMemcombine(op Op, leaves []*Value) bool { // rebalance balances associative computation to better help CPU instruction pipelining (#49331) // and groups constants together catch more constant folding opportunities. // -// a + b + c + d compiles to to (a + (b + (c + d)) which is an unbalanced expression tree +// a + b + c + d compiles to to v1:(a + v2:(b + v3:(c + d)) which is an unbalanced expression tree // Which is suboptimal since it requires the CPU to compute v3 before fetching it use its result in // v2, and v2 before its use in v1 // // This optimization rebalances this expression tree to look like (a + b) + (c + d) , // which removes such dependencies and frees up the CPU pipeline. // -// The above optimization is a good starting point for other sorts of operations such as +// The above optimization is also a good starting point for other sorts of operations such as // turning a + a + a => 3*a, cancelling pairs a + (-a), collecting up common factors TODO(ryan-berger) -func rebalance(v *Value, visited map[*Value]bool) { +func rebalance(v *Value, visited *sparseSet) { // We cannot apply this optimization to non-commutative operations, - // values that have more than one use, or non-binary ops (would need more log math). // Try and save time by not revisiting nodes - if visited[v] || !(v.Uses == 1 && opcodeTable[v.Op].commutative) || len(v.Args) > 2 { + if visited.contains(v.ID) || !opcodeTable[v.Op].commutative { return } - // The smallest possible rebalanceable expression has 3 nodes and 4 leaves, + // The smallest possible rebalanceable binary expression has 3 nodes and 4 leaves, // so preallocate the lists to save time if it is not rebalanceable leaves := make([]*Value, 0, 4) nodes := make([]*Value, 0, 3) // Do a bfs on v to keep a nice reverse topological order haystack := []*Value{v} - for len(haystack) != 0 { - nextHaystack := make([]*Value, 0, len(v.Args)*len(haystack)) - for _, needle := range haystack { - // if we are searching a value, it must be a node so add it to our node list - nodes = append(nodes, needle) + for i := 0; i < len(haystack); i++ { + needle := haystack[i] + // if we are searching a value, it must be a node so add it to our node list + nodes = append(nodes, needle) - // Only visit nodes. Leafs may be rebalancable for a different op type - visited[needle] = true + // Only visit nodes. Leafs may be rebalancable for a different op type + visited.add(v.ID) - for _, a := range needle.Args { - // If the ops aren't the same or have more than one use it must be a leaf. - if a.Op != v.Op || a.Uses != 1 { - leaves = append(leaves, a) - continue - } - - // nodes in the tree now hold the invariants that: - // - they are of a common associative operation as the rest of the tree - // - they have only a single use (this invariant could be removed with further analysis TODO(ryan-berger) - nextHaystack = append(nextHaystack, a) + for _, a := range needle.Args { + // If the ops aren't the same or have more than one use it must be a leaf. + if a.Op != v.Op || a.Uses != 1 { + leaves = append(leaves, a) + continue } + + // nodes in the tree now hold the invariants that: + // - they are of a common associative operation as the rest of the tree + // - they have only a single use (this invariant could be removed with further analysis TODO(ryan-berger)) + haystack = append(haystack, a) } - haystack = nextHaystack } - // we need at least 4 leaves for this expression to be rebalanceable, + minLeaves := len(v.Args) * len(v.Args) + + // we need at least args^2 leaves for this expression to be rebalanceable, // and we can't balance a potential load widening (see memcombine) - if len(leaves) < 4 || probablyMemcombine(v.Op, leaves) { + if len(leaves) < minLeaves || probablyMemcombine(v.Op, leaves) { return } balanceExprTree(v, visited, nodes, leaves) } - -// reassociate balances trees of commutative computation -// to better group expressions to expose easy optimizations in -// cse, cancelling/counting/factoring expressions, etc. -func reassociate(f *Func) { - visited := make(map[*Value]bool) - - for _, b := range f.Postorder() { - for i := len(b.Values) - 1; i >= 0; i-- { - val := b.Values[i] - rebalance(val, visited) - } - } - - for k := range visited { - delete(visited, k) - } -}