cmd/compile: allow multiple induction variables in one block in prove

In this CL, the restriction that each block can only have one induction
variable has been removed. This reduces missed optimizations.

Fixes #76269

Change-Id: I14043182a40cc7887c5b6d9c1a5df8ea3a1bfedc
Reviewed-on: https://go-review.googlesource.com/c/go/+/719881
Reviewed-by: Keith Randall <khr@golang.org>
Auto-Submit: Keith Randall <khr@golang.org>
Reviewed-by: Carlos Amedee <carlos@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Keith Randall <khr@google.com>

3 files changed, 74 insertions, 37 deletions

diff --git a/src/cmd/compile/internal/ssa/downward_counting_loop.go b/src/cmd/compile/internal/ssa/downward_counting_loop.go
index ec32bcc391..ec0e5478a8 100644
--- a/src/cmd/compile/internal/ssa/downward_counting_loop.go
+++ b/src/cmd/compile/internal/ssa/downward_counting_loop.go
@@ -4,8 +4,6 @@
 
 package ssa
 
-import "fmt"
-
 // maybeRewriteLoopToDownwardCountingLoop tries to rewrite the loop to a
 // downward counting loop checking against start if the loop body does
 // not depend on ind or nxt and end is known before the loop.
@@ -76,32 +74,37 @@ func maybeRewriteLoopToDownwardCountingLoop(f *Func, v indVar) {
 		return
 	}
 
-	check := ind.Block.Controls[0]
-	// invert the check
-	check.Args[0], check.Args[1] = check.Args[1], check.Args[0]
+	check := v.entry.Preds[0].b.Controls[0]
 
-	// swap start and end in the loop
+	idxEnd, idxStart := -1, -1
 	for i, v := range check.Args {
-		if v != end {
-			continue
+		if v == end {
+			idxEnd = i
+			break
 		}
-
-		check.SetArg(i, start)
-		goto replacedEnd
 	}
-	panic(fmt.Sprintf("unreachable, ind: %v, start: %v, end: %v", ind, start, end))
-replacedEnd:
-
 	for i, v := range ind.Args {
-		if v != start {
-			continue
+		if v == start {
+			idxStart = i
+			break
 		}
+	}
+	if idxEnd < 0 || idxStart < 0 {
+		return
+	}
 
-		ind.SetArg(i, end)
-		goto replacedStart
+	sdom := f.Sdom()
+	// the end may not dominate the ind after rewrite, check it first
+	if !sdom.IsAncestorEq(end.Block, ind.Block) {
+		return
 	}
-	panic(fmt.Sprintf("unreachable, ind: %v, start: %v, end: %v", ind, start, end))
-replacedStart:
+
+	// swap start and end in the loop
+	check.SetArg(idxEnd, start)
+	ind.SetArg(idxStart, end)
+
+	// invert the check
+	check.Args[0], check.Args[1] = check.Args[1], check.Args[0]
 
 	if nxt.Args[0] != ind {
 		// unlike additions subtractions are not commutative so be sure we get it right
diff --git a/src/cmd/compile/internal/ssa/loopbce.go b/src/cmd/compile/internal/ssa/loopbce.go
index aa6cc48cac..76fc703fc2 100644
--- a/src/cmd/compile/internal/ssa/loopbce.go
+++ b/src/cmd/compile/internal/ssa/loopbce.go
@@ -23,9 +23,9 @@ type indVar struct {
 	nxt   *Value // the incremented variable
 	min   *Value // minimum value, inclusive/exclusive depends on flags
 	max   *Value // maximum value, inclusive/exclusive depends on flags
-	entry *Block // entry block in the loop.
+	entry *Block // the block where the edge from the succeeded comparison of the induction variable goes to, means when the bound check has passed.
 	flags indVarFlags
-	// Invariant: for all blocks strictly dominated by entry:
+	// Invariant: for all blocks dominated by entry:
 	//	min <= ind <  max    [if flags == 0]
 	//	min <  ind <  max    [if flags == indVarMinExc]
 	//	min <= ind <= max    [if flags == indVarMaxInc]
@@ -83,12 +83,43 @@ func parseIndVar(ind *Value) (min, inc, nxt *Value, loopReturn Edge) {
 //		goto loop
 //
 //	 exit_loop:
+//
+// We may have more than one induction variables, the loop in the go
+// source code may looks like this:
+//
+//	for i >= 0 && j >= 0 {
+//		// use i and j
+//		i--
+//		j--
+//	}
+//
+// So, also look for variables and blocks that satisfy the following
+//
+//	loop:
+//	  i = (Phi maxi nxti)
+//	  j = (Phi maxj nxtj)
+//	  if i >= mini
+//	    then goto check_j
+//	    else goto exit_loop
+//
+//	check_j:
+//	  if j >= minj
+//	    then goto enter_loop
+//	    else goto exit_loop
+//
+//	enter_loop:
+//	  do something
+//	  nxti = i - di
+//	  nxtj = j - dj
+//	  goto loop
+//
+//	exit_loop:
 func findIndVar(f *Func) []indVar {
 	var iv []indVar
 	sdom := f.Sdom()
 
 	for _, b := range f.Blocks {
-		if b.Kind != BlockIf || len(b.Preds) != 2 {
+		if b.Kind != BlockIf {
 			continue
 		}
 
@@ -130,10 +161,9 @@ func findIndVar(f *Func) []indVar {
 			less = false
 		}
 
-		if ind.Block != b {
-			// TODO: Could be extended to include disjointed loop headers.
-			// I don't think this is causing missed optimizations in real world code often.
-			// See https://go.dev/issue/63955
+		// This is ind.Block.Preds, not b.Preds. That's a restriction on the loop header,
+		// not the comparison block.
+		if len(ind.Block.Preds) != 2 {
 			continue
 		}
 
@@ -184,9 +214,10 @@ func findIndVar(f *Func) []indVar {
 		// Two conditions must happen listed below to accept ind
 		// as an induction variable.
 
-		// First condition: loop entry has a single predecessor, which
-		// is the header block.  This implies that b.Succs[0] is
-		// reached iff ind < limit.
+		// First condition: the entry block has a single predecessor.
+		// The entry now means the in-loop edge where the induction variable
+		// comparison succeeded. Its predecessor is not necessarily the header
+		// block. This implies that b.Succs[0] is reached iff ind < limit.
 		if len(startBody.b.Preds) != 1 {
 			// the other successor must exit the loop.
 			continue
@@ -195,7 +226,8 @@ func findIndVar(f *Func) []indVar {
 		// Second condition: startBody.b dominates nxt so that
 		// nxt is computed when inc < limit.
 		if !sdom.IsAncestorEq(startBody.b, nxt.Block) {
-			// inc+ind can only be reached through the branch that enters the loop.
+			// inc+ind can only be reached through the branch that confirmed the
+			// induction variable is in bounds.
 			continue
 		}
 
@@ -309,10 +341,13 @@ func findIndVar(f *Func) []indVar {
 			}
 
 			iv = append(iv, indVar{
-				ind:   ind,
-				nxt:   nxt,
-				min:   min,
-				max:   max,
+				ind: ind,
+				nxt: nxt,
+				min: min,
+				max: max,
+				// This is startBody.b, where startBody is the edge from the comparison for the
+				// induction variable, not necessarily the in-loop edge from the loop header.
+				// Induction variable bounds are not valid in the loop before this edge.
 				entry: startBody.b,
 				flags: flags,
 			})
diff --git a/src/cmd/compile/internal/ssa/prove.go b/src/cmd/compile/internal/ssa/prove.go
index c0f6b967da..abe43fd1ae 100644
--- a/src/cmd/compile/internal/ssa/prove.go
+++ b/src/cmd/compile/internal/ssa/prove.go
@@ -1562,8 +1562,7 @@ func getSliceInfo(vp *Value) (inf sliceInfo) {
 // its negation. If either leads to a contradiction, it can trim that
 // successor.
 func prove(f *Func) {
-	// Find induction variables. Currently, findIndVars
-	// is limited to one induction variable per block.
+	// Find induction variables.
 	var indVars map[*Block]indVar
 	for _, v := range findIndVar(f) {
 		ind := v.ind