cmd/compile: fix phi-function updates for preemptible loops

Previous code failed to account for particular control flow
involving nested loops when updating phi function inputs.
Fix involves:
1) remove incorrect shortcut
2) generate a "better" order for children in dominator tree
3) note inner-loop updates and check before applying
   outer-loop updates.

Fixes #20675.

Change-Id: I2fe21470604b5c259e777ad8b15de95f7706894d
Reviewed-on: https://go-review.googlesource.com/45791
Run-TryBot: David Chase <drchase@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Cherry Zhang <cherryyz@google.com>

2 files changed, 69 insertions, 21 deletions

diff --git a/src/cmd/compile/internal/ssa/loopreschedchecks.go b/src/cmd/compile/internal/ssa/loopreschedchecks.go
index 98b6e92e93..4222bf81c5 100644
--- a/src/cmd/compile/internal/ssa/loopreschedchecks.go
+++ b/src/cmd/compile/internal/ssa/loopreschedchecks.go
@@ -73,9 +73,13 @@ func insertLoopReschedChecks(f *Func) {
 	lastMems := findLastMems(f)
 
 	idom := f.Idom()
-	sdom := f.sdom()
+	po := f.postorder()
+	// The ordering in the dominator tree matters; it's important that
+	// the walk of the dominator tree also be a preorder (i.e., a node is
+	// visited only after all its non-backedge predecessors have been visited).
+	sdom := newSparseOrderedTree(f, idom, po)
 
-	if f.pass.debug > 2 {
+	if f.pass.debug > 1 {
 		fmt.Printf("before %s = %s\n", f.Name, sdom.treestructure(f.Entry))
 	}
 
@@ -93,7 +97,6 @@ func insertLoopReschedChecks(f *Func) {
 	memDefsAtBlockEnds := make([]*Value, f.NumBlocks()) // For each block, the mem def seen at its bottom. Could be from earlier block.
 
 	// Propagate last mem definitions forward through successor blocks.
-	po := f.postorder()
 	for i := len(po) - 1; i >= 0; i-- {
 		b := po[i]
 		mem := lastMems[b.ID]
@@ -102,6 +105,9 @@ func insertLoopReschedChecks(f *Func) {
 			mem = memDefsAtBlockEnds[b.Preds[j].b.ID]
 		}
 		memDefsAtBlockEnds[b.ID] = mem
+		if f.pass.debug > 2 {
+			fmt.Printf("memDefsAtBlockEnds[%s] = %s\n", b, mem)
+		}
 	}
 
 	// Maps from block to newly-inserted phi function in block.
@@ -126,18 +132,27 @@ func insertLoopReschedChecks(f *Func) {
 			mem0 := memDefsAtBlockEnds[idom[h.ID].ID]
 			headerMemPhi = newPhiFor(h, mem0)
 			newmemphis[h] = rewrite{before: mem0, after: headerMemPhi}
-			addDFphis(mem0, h, h, f, memDefsAtBlockEnds, newmemphis)
+			addDFphis(mem0, h, h, f, memDefsAtBlockEnds, newmemphis, sdom)
 
 		}
 		tofixBackedges[i].m = headerMemPhi
 
 	}
+	if f.pass.debug > 0 {
+		for b, r := range newmemphis {
+			fmt.Printf("before b=%s, rewrite=%s\n", b, r.String())
+		}
+	}
+
+	// dfPhiTargets notes inputs to phis in dominance frontiers that should not
+	// be rewritten as part of the dominated children of some outer rewrite.
+	dfPhiTargets := make(map[rewriteTarget]bool)
 
-	rewriteNewPhis(f.Entry, f.Entry, f, memDefsAtBlockEnds, newmemphis)
+	rewriteNewPhis(f.Entry, f.Entry, f, memDefsAtBlockEnds, newmemphis, dfPhiTargets, sdom)
 
 	if f.pass.debug > 0 {
 		for b, r := range newmemphis {
-			fmt.Printf("b=%s, rewrite=%s\n", b, r.String())
+			fmt.Printf("after b=%s, rewrite=%s\n", b, r.String())
 		}
 	}
 
@@ -248,7 +263,7 @@ func insertLoopReschedChecks(f *Func) {
 
 	f.invalidateCFG()
 
-	if f.pass.debug > 2 {
+	if f.pass.debug > 1 {
 		sdom = newSparseTree(f, f.Idom())
 		fmt.Printf("after %s = %s\n", f.Name, sdom.treestructure(f.Entry))
 	}
@@ -272,7 +287,10 @@ func newPhiFor(b *Block, v *Value) *Value {
 // if b has its own phi definition then it takes the place of h.
 // defsForUses provides information about other definitions of the variable that are present
 // (and if nil, indicates that the variable is no longer live)
-func rewriteNewPhis(h, b *Block, f *Func, defsForUses []*Value, newphis map[*Block]rewrite) {
+// sdom must yield a preorder of the flow graph if recursively walked, root-to-children.
+// The result of newSparseOrderedTree with order supplied by a dfs-postorder satisfies this
+// requirement.
+func rewriteNewPhis(h, b *Block, f *Func, defsForUses []*Value, newphis map[*Block]rewrite, dfPhiTargets map[rewriteTarget]bool, sdom SparseTree) {
 	// If b is a block with a new phi, then a new rewrite applies below it in the dominator tree.
 	if _, ok := newphis[b]; ok {
 		h = b
@@ -292,7 +310,19 @@ func rewriteNewPhis(h, b *Block, f *Func, defsForUses []*Value, newphis map[*Blo
 				if w != x {
 					continue
 				}
-				*p = append(*p, rewriteTarget{v, i})
+				tgt := rewriteTarget{v, i}
+
+				// It's possible dominated control flow will rewrite this instead.
+				// Visiting in preorder (a property of how sdom was constructed)
+				// ensures that these are seen in the proper order.
+				if dfPhiTargets[tgt] {
+					continue
+				}
+				*p = append(*p, tgt)
+				if f.pass.debug > 1 {
+					fmt.Printf("added block target for h=%v, b=%v, x=%v, y=%v, tgt.v=%s, tgt.i=%d\n",
+						h, b, x, y, v, i)
+				}
 			}
 		}
 
@@ -304,13 +334,16 @@ func rewriteNewPhis(h, b *Block, f *Func, defsForUses []*Value, newphis map[*Blo
 		if dfu := defsForUses[b.ID]; dfu != nil && dfu.Block != b {
 			for _, e := range b.Succs {
 				s := e.b
-				if sphi, ok := newphis[s]; ok { // saves time to find the phi this way.
-					*p = append(*p, rewriteTarget{sphi.after, e.i})
-					continue
-				}
+
 				for _, v := range s.Values {
 					if v.Op == OpPhi && v.Args[e.i] == x {
-						*p = append(*p, rewriteTarget{v, e.i})
+						tgt := rewriteTarget{v, e.i}
+						*p = append(*p, tgt)
+						dfPhiTargets[tgt] = true
+						if f.pass.debug > 1 {
+							fmt.Printf("added phi target for h=%v, b=%v, s=%v, x=%v, y=%v, tgt.v=%s, tgt.i=%d\n",
+								h, b, s, x, y, v.LongString(), e.i)
+						}
 						break
 					}
 				}
@@ -319,10 +352,8 @@ func rewriteNewPhis(h, b *Block, f *Func, defsForUses []*Value, newphis map[*Blo
 		newphis[h] = change
 	}
 
-	sdom := f.sdom()
-
 	for c := sdom[b.ID].child; c != nil; c = sdom[c.ID].sibling {
-		rewriteNewPhis(h, c, f, defsForUses, newphis) // TODO: convert to explicit stack from recursion.
+		rewriteNewPhis(h, c, f, defsForUses, newphis, dfPhiTargets, sdom) // TODO: convert to explicit stack from recursion.
 	}
 }
 
@@ -333,12 +364,11 @@ func rewriteNewPhis(h, b *Block, f *Func, defsForUses []*Value, newphis map[*Blo
 // either b = h or h strictly dominates b.
 // These newly created phis are themselves new definitions that may require addition of their
 // own trivial phi functions in their own dominance frontier, and this is handled recursively.
-func addDFphis(x *Value, h, b *Block, f *Func, defForUses []*Value, newphis map[*Block]rewrite) {
+func addDFphis(x *Value, h, b *Block, f *Func, defForUses []*Value, newphis map[*Block]rewrite, sdom SparseTree) {
 	oldv := defForUses[b.ID]
 	if oldv != x { // either a new definition replacing x, or nil if it is proven that there are no uses reachable from b
 		return
 	}
-	sdom := f.sdom()
 	idom := f.Idom()
 outer:
 	for _, e := range b.Succs {
@@ -362,10 +392,10 @@ outer:
 		headerPhi := newPhiFor(s, old)
 		// the new phi will replace "old" in block s and all blocks dominated by s.
 		newphis[s] = rewrite{before: old, after: headerPhi} // record new phi, to have inputs labeled "old" rewritten to "headerPhi"
-		addDFphis(old, s, s, f, defForUses, newphis)        // the new definition may also create new phi functions.
+		addDFphis(old, s, s, f, defForUses, newphis, sdom)  // the new definition may also create new phi functions.
 	}
 	for c := sdom[b.ID].child; c != nil; c = sdom[c.ID].sibling {
-		addDFphis(x, h, c, f, defForUses, newphis) // TODO: convert to explicit stack from recursion.
+		addDFphis(x, h, c, f, defForUses, newphis, sdom) // TODO: convert to explicit stack from recursion.
 	}
 }
 
diff --git a/src/cmd/compile/internal/ssa/sparsetree.go b/src/cmd/compile/internal/ssa/sparsetree.go
index 8e5b9f3e5b..f7af85446b 100644
--- a/src/cmd/compile/internal/ssa/sparsetree.go
+++ b/src/cmd/compile/internal/ssa/sparsetree.go
@@ -70,6 +70,24 @@ func newSparseTree(f *Func, parentOf []*Block) SparseTree {
 	return t
 }
 
+// newSparseOrderedTree creates a SparseTree from a block-to-parent map (array indexed by Block.ID)
+// children will appear in the reverse of their order in reverseOrder
+// in particular, if reverseOrder is a dfs-reversePostOrder, then the root-to-children
+// walk of the tree will yield a pre-order.
+func newSparseOrderedTree(f *Func, parentOf, reverseOrder []*Block) SparseTree {
+	t := make(SparseTree, f.NumBlocks())
+	for _, b := range reverseOrder {
+		n := &t[b.ID]
+		if p := parentOf[b.ID]; p != nil {
+			n.parent = p
+			n.sibling = t[p.ID].child
+			t[p.ID].child = b
+		}
+	}
+	t.numberBlock(f.Entry, 1)
+	return t
+}
+
 // treestructure provides a string description of the dominator
 // tree and flow structure of block b and all blocks that it
 // dominates.