cmd/compile: rewrite proved multiplies by 0 or 1 into CondSelect

Updates #76056

Change-Id: I64fe631ab381c74f902f877392530d7cc91860ab
Reviewed-on: https://go-review.googlesource.com/c/go/+/715044
Reviewed-by: Michael Knyszek <mknyszek@google.com>
Auto-Submit: Jorropo <jorropo.pgm@gmail.com>
Reviewed-by: Keith Randall <khr@google.com>
Reviewed-by: Keith Randall <khr@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

src/cmd/compile/internal/ssa/branchelim.go

@@ -70,7 +70,8 @@ func branchelim(f *Func) {
 }
 
 func canCondSelect(v *Value, arch string, loadAddr *sparseSet) bool {
-	if loadAddr.contains(v.ID) {
+	if loadAddr != nil && // prove calls this on some multiplies and doesn't take care of loadAddrs
+		loadAddr.contains(v.ID) {
 		// The result of the soon-to-be conditional move is used to compute a load address.
 		// We want to avoid generating a conditional move in this case
 		// because the load address would now be data-dependent on the condition.

src/cmd/compile/internal/ssa/prove.go

@@ -2646,11 +2646,23 @@ var bytesizeToConst = [...]Op{
 	32 / 8: OpConst32,
 	64 / 8: OpConst64,
 }
+var bytesizeToNeq = [...]Op{
+	8 / 8:  OpNeq8,
+	16 / 8: OpNeq16,
+	32 / 8: OpNeq32,
+	64 / 8: OpNeq64,
+}
+var bytesizeToAnd = [...]Op{
+	8 / 8:  OpAnd8,
+	16 / 8: OpAnd16,
+	32 / 8: OpAnd32,
+	64 / 8: OpAnd64,
+}
 
 // simplifyBlock simplifies some constant values in b and evaluates
 // branches to non-uniquely dominated successors of b.
 func simplifyBlock(sdom SparseTree, ft *factsTable, b *Block) {
-	for _, v := range b.Values {
+	for iv, v := range b.Values {
 		switch v.Op {
 		case OpSlicemask:
 			// Replace OpSlicemask operations in b with constants where possible.
@@ -2757,6 +2769,43 @@ func simplifyBlock(sdom SparseTree, ft *factsTable, b *Block) {
 					b.Func.Warnl(v.Pos, "Proved %v does not need fix-up", v.Op)
 				}
 			}
+		case OpMul64, OpMul32, OpMul16, OpMul8:
+			x := v.Args[0]
+			xl := ft.limits[x.ID]
+			y := v.Args[1]
+			yl := ft.limits[y.ID]
+			switch xOne, yOne := xl.umax <= 1, yl.umax <= 1; {
+			case xOne && yOne:
+				v.Op = bytesizeToAnd[v.Type.Size()]
+				if b.Func.pass.debug > 0 {
+					b.Func.Warnl(v.Pos, "Rewrote Mul %v into And", v)
+				}
+			case yOne && b.Func.Config.haveCondSelect:
+				x, y = y, x
+				fallthrough
+			case xOne && b.Func.Config.haveCondSelect:
+				if !canCondSelect(v, b.Func.Config.arch, nil) {
+					break
+				}
+				zero := b.Func.constVal(bytesizeToConst[v.Type.Size()], v.Type, 0, true)
+				ft.initLimitForNewValue(zero)
+				check := b.NewValue2(v.Pos, bytesizeToNeq[v.Type.Size()], types.Types[types.TBOOL], zero, x)
+				ft.initLimitForNewValue(check)
+				v.reset(OpCondSelect)
+				v.AddArg3(y, zero, check)
+
+				// FIXME: workaround for go.dev/issues/76060
+				// we need to schedule the Neq before the CondSelect even tho
+				// scheduling is meaningless until we reach the schedule pass.
+				if b.Values[len(b.Values)-1] != check {
+					panic("unreachable; failed sanity check, new value isn't at the end of the block")
+				}
+				b.Values[iv], b.Values[len(b.Values)-1] = b.Values[len(b.Values)-1], b.Values[iv]
+
+				if b.Func.pass.debug > 0 {
+					b.Func.Warnl(v.Pos, "Rewrote Mul %v into CondSelect; %v is bool", v, x)
+				}
+			}
 		}
 		// Fold provable constant results.
 		// Helps in cases where we reuse a value after branching on its equality.

test/prove.go

@@ -2501,6 +2501,19 @@ func issue75144ifNot(a, b []uint64) bool {
 	return false
 }
 
+func mulIntoAnd(a, b uint) uint {
+	if a > 1 || b > 1 {
+		return 0
+	}
+	return a * b // ERROR "Rewrote Mul v[0-9]+ into And$"
+}
+func mulIntoCondSelect(a, b uint) uint {
+	if a > 1 {
+		return 0
+	}
+	return a * b // ERROR "Rewrote Mul v[0-9]+ into CondSelect"
+}
+
 //go:noinline
 func useInt(a int) {
 }