cmd/compile: extend condselect into math code to handle other constants than 1

On amd64 along:
  if b { x += 1 } => x += b

We can also implement constants 2 4 and 8:
  if b { x += 2 } => x += b * 2

This compiles to a displacement LEA.

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

4 files changed, 252 insertions, 131 deletions

diff --git a/src/cmd/compile/internal/ssa/_gen/generic.rules b/src/cmd/compile/internal/ssa/_gen/generic.rules
index e75f371790..5ceadd79a1 100644
--- a/src/cmd/compile/internal/ssa/_gen/generic.rules
+++ b/src/cmd/compile/internal/ssa/_gen/generic.rules
@@ -2297,12 +2297,24 @@
 => (StaticLECall {f} [argsize] dict_ (StringMake <typ.String> ptr len) mem)
 
 // Transform some CondSelect into math operations.
-// if b { x++ } => x += b // but not on arm64 because it has CSINC
-(CondSelect (Add8 <t> x (Const8 [1])) x bool) && config.arch != "arm64" => (Add8 x (CvtBoolToUint8 <t> bool))
-(CondSelect (Add(64|32|16) <t> x (Const(64|32|16) [1])) x bool) && config.arch != "arm64" => (Add(64|32|16) x (ZeroExt8to(64|32|16) <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> bool)))
-// if !b { x++ } => x += !b // but not on arm64 because it has CSINC
-(CondSelect x (Add8 <t> x (Const8 [1])) bool) && config.arch != "arm64" => (Add8 x (CvtBoolToUint8 <t> (Not <bool.Type> bool)))
-(CondSelect x (Add(64|32|16) <t> x (Const(64|32|16) [1])) bool) && config.arch != "arm64" => (Add(64|32|16) x (ZeroExt8to(64|32|16) <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> (Not <bool.Type> bool))))
+// if b { x += c } => x += b * c // but not on arm64 because it has CSINC
+(CondSelect op:(Add8 <t> x c:(Const8)) x bool) &&
+	canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) &&
+	config.arch != "arm64" =>
+	(Add8 x (Mul8 <t> c (CvtBoolToUint8 <t> bool)))
+(CondSelect op:(Add(64|32|16) <t> x c:(Const(64|32|16))) x bool) &&
+	canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) &&
+	config.arch != "arm64" =>
+	(Add(64|32|16) x (Mul(64|32|16) <t> c (ZeroExt8to(64|32|16) <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> bool))))
+// if !b { x += c } => x += !b * c // but not on arm64 because it has CSINC
+(CondSelect x op:(Add8 <t> x c:(Const8)) bool) &&
+	canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) &&
+	config.arch != "arm64" =>
+	(Add8 x (Mul8 <t> c (CvtBoolToUint8 <t> (Not <bool.Type> bool))))
+(CondSelect x op:(Add(64|32|16) <t> x c:(Const(64|32|16))) bool) &&
+	canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) &&
+	config.arch != "arm64" =>
+	(Add(64|32|16) x (Mul(64|32|16) <t> c (ZeroExt8to(64|32|16) <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> (Not <bool.Type> bool)))))
 
 // if b { x-- } => x -= b
 (CondSelect (Add8 <t> x (Const8 [-1])) x bool) => (Sub8 x (CvtBoolToUint8 <t> bool))
diff --git a/src/cmd/compile/internal/ssa/rewrite.go b/src/cmd/compile/internal/ssa/rewrite.go
index 7358d575ad..3bc872bd76 100644
--- a/src/cmd/compile/internal/ssa/rewrite.go
+++ b/src/cmd/compile/internal/ssa/rewrite.go
@@ -2799,3 +2799,23 @@ func bool2int(x bool) int {
 	}
 	return b
 }
+
+// canImplementXOpBooleanLikeTimesAConstantCheaply reports whether we can cheaply implement:
+// x OP (y * constant) when we know y is either 0 or 1.
+func canImplementXOpBooleanLikeTimesAConstantCheaply(config *Config, op Op, constant int64) bool {
+	if constant == 1 {
+		return true
+	}
+	switch config.arch {
+	case "amd64":
+		switch op {
+		case OpAdd64, OpAdd32, OpAdd16, OpAdd8:
+			switch constant {
+			case 2, 4, 8:
+				// Implemented with LEA a + b * displacement form
+				return true
+			}
+		}
+	}
+	return false
+}
diff --git a/src/cmd/compile/internal/ssa/rewritegeneric.go b/src/cmd/compile/internal/ssa/rewritegeneric.go
index 0386763639..688dbcbe9f 100644
--- a/src/cmd/compile/internal/ssa/rewritegeneric.go
+++ b/src/cmd/compile/internal/ssa/rewritegeneric.go
@@ -5912,245 +5912,289 @@ func rewriteValuegeneric_OpCondSelect(v *Value) bool {
 		v.copyOf(x)
 		return true
 	}
-	// match: (CondSelect (Add8 <t> x (Const8 [1])) x bool)
-	// cond: config.arch != "arm64"
-	// result: (Add8 x (CvtBoolToUint8 <t> bool))
+	// match: (CondSelect op:(Add8 <t> x c:(Const8)) x bool)
+	// cond: canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64"
+	// result: (Add8 x (Mul8 <t> c (CvtBoolToUint8 <t> bool)))
 	for {
-		if v_0.Op != OpAdd8 {
+		op := v_0
+		if op.Op != OpAdd8 {
 			break
 		}
-		t := v_0.Type
-		_ = v_0.Args[1]
-		v_0_0 := v_0.Args[0]
-		v_0_1 := v_0.Args[1]
-		for _i0 := 0; _i0 <= 1; _i0, v_0_0, v_0_1 = _i0+1, v_0_1, v_0_0 {
-			x := v_0_0
-			if v_0_1.Op != OpConst8 || auxIntToInt8(v_0_1.AuxInt) != 1 || x != v_1 {
+		t := op.Type
+		_ = op.Args[1]
+		op_0 := op.Args[0]
+		op_1 := op.Args[1]
+		for _i0 := 0; _i0 <= 1; _i0, op_0, op_1 = _i0+1, op_1, op_0 {
+			x := op_0
+			c := op_1
+			if c.Op != OpConst8 || x != v_1 {
 				continue
 			}
 			bool := v_2
-			if !(config.arch != "arm64") {
+			if !(canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64") {
 				continue
 			}
 			v.reset(OpAdd8)
-			v0 := b.NewValue0(v.Pos, OpCvtBoolToUint8, t)
-			v0.AddArg(bool)
+			v0 := b.NewValue0(v.Pos, OpMul8, t)
+			v1 := b.NewValue0(v.Pos, OpCvtBoolToUint8, t)
+			v1.AddArg(bool)
+			v0.AddArg2(c, v1)
 			v.AddArg2(x, v0)
 			return true
 		}
 		break
 	}
-	// match: (CondSelect (Add64 <t> x (Const64 [1])) x bool)
-	// cond: config.arch != "arm64"
-	// result: (Add64 x (ZeroExt8to64 <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> bool)))
+	// match: (CondSelect op:(Add64 <t> x c:(Const64)) x bool)
+	// cond: canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64"
+	// result: (Add64 x (Mul64 <t> c (ZeroExt8to64 <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> bool))))
 	for {
-		if v_0.Op != OpAdd64 {
+		op := v_0
+		if op.Op != OpAdd64 {
 			break
 		}
-		t := v_0.Type
-		_ = v_0.Args[1]
-		v_0_0 := v_0.Args[0]
-		v_0_1 := v_0.Args[1]
-		for _i0 := 0; _i0 <= 1; _i0, v_0_0, v_0_1 = _i0+1, v_0_1, v_0_0 {
-			x := v_0_0
-			if v_0_1.Op != OpConst64 || auxIntToInt64(v_0_1.AuxInt) != 1 || x != v_1 {
+		t := op.Type
+		_ = op.Args[1]
+		op_0 := op.Args[0]
+		op_1 := op.Args[1]
+		for _i0 := 0; _i0 <= 1; _i0, op_0, op_1 = _i0+1, op_1, op_0 {
+			x := op_0
+			c := op_1
+			if c.Op != OpConst64 || x != v_1 {
 				continue
 			}
 			bool := v_2
-			if !(config.arch != "arm64") {
+			if !(canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64") {
 				continue
 			}
 			v.reset(OpAdd64)
-			v0 := b.NewValue0(v.Pos, OpZeroExt8to64, t)
-			v1 := b.NewValue0(v.Pos, OpCvtBoolToUint8, types.Types[types.TUINT8])
-			v1.AddArg(bool)
-			v0.AddArg(v1)
+			v0 := b.NewValue0(v.Pos, OpMul64, t)
+			v1 := b.NewValue0(v.Pos, OpZeroExt8to64, t)
+			v2 := b.NewValue0(v.Pos, OpCvtBoolToUint8, types.Types[types.TUINT8])
+			v2.AddArg(bool)
+			v1.AddArg(v2)
+			v0.AddArg2(c, v1)
 			v.AddArg2(x, v0)
 			return true
 		}
 		break
 	}
-	// match: (CondSelect (Add32 <t> x (Const32 [1])) x bool)
-	// cond: config.arch != "arm64"
-	// result: (Add32 x (ZeroExt8to32 <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> bool)))
+	// match: (CondSelect op:(Add32 <t> x c:(Const32)) x bool)
+	// cond: canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64"
+	// result: (Add32 x (Mul32 <t> c (ZeroExt8to32 <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> bool))))
 	for {
-		if v_0.Op != OpAdd32 {
+		op := v_0
+		if op.Op != OpAdd32 {
 			break
 		}
-		t := v_0.Type
-		_ = v_0.Args[1]
-		v_0_0 := v_0.Args[0]
-		v_0_1 := v_0.Args[1]
-		for _i0 := 0; _i0 <= 1; _i0, v_0_0, v_0_1 = _i0+1, v_0_1, v_0_0 {
-			x := v_0_0
-			if v_0_1.Op != OpConst32 || auxIntToInt32(v_0_1.AuxInt) != 1 || x != v_1 {
+		t := op.Type
+		_ = op.Args[1]
+		op_0 := op.Args[0]
+		op_1 := op.Args[1]
+		for _i0 := 0; _i0 <= 1; _i0, op_0, op_1 = _i0+1, op_1, op_0 {
+			x := op_0
+			c := op_1
+			if c.Op != OpConst32 || x != v_1 {
 				continue
 			}
 			bool := v_2
-			if !(config.arch != "arm64") {
+			if !(canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64") {
 				continue
 			}
 			v.reset(OpAdd32)
-			v0 := b.NewValue0(v.Pos, OpZeroExt8to32, t)
-			v1 := b.NewValue0(v.Pos, OpCvtBoolToUint8, types.Types[types.TUINT8])
-			v1.AddArg(bool)
-			v0.AddArg(v1)
+			v0 := b.NewValue0(v.Pos, OpMul32, t)
+			v1 := b.NewValue0(v.Pos, OpZeroExt8to32, t)
+			v2 := b.NewValue0(v.Pos, OpCvtBoolToUint8, types.Types[types.TUINT8])
+			v2.AddArg(bool)
+			v1.AddArg(v2)
+			v0.AddArg2(c, v1)
 			v.AddArg2(x, v0)
 			return true
 		}
 		break
 	}
-	// match: (CondSelect (Add16 <t> x (Const16 [1])) x bool)
-	// cond: config.arch != "arm64"
-	// result: (Add16 x (ZeroExt8to16 <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> bool)))
+	// match: (CondSelect op:(Add16 <t> x c:(Const16)) x bool)
+	// cond: canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64"
+	// result: (Add16 x (Mul16 <t> c (ZeroExt8to16 <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> bool))))
 	for {
-		if v_0.Op != OpAdd16 {
+		op := v_0
+		if op.Op != OpAdd16 {
 			break
 		}
-		t := v_0.Type
-		_ = v_0.Args[1]
-		v_0_0 := v_0.Args[0]
-		v_0_1 := v_0.Args[1]
-		for _i0 := 0; _i0 <= 1; _i0, v_0_0, v_0_1 = _i0+1, v_0_1, v_0_0 {
-			x := v_0_0
-			if v_0_1.Op != OpConst16 || auxIntToInt16(v_0_1.AuxInt) != 1 || x != v_1 {
+		t := op.Type
+		_ = op.Args[1]
+		op_0 := op.Args[0]
+		op_1 := op.Args[1]
+		for _i0 := 0; _i0 <= 1; _i0, op_0, op_1 = _i0+1, op_1, op_0 {
+			x := op_0
+			c := op_1
+			if c.Op != OpConst16 || x != v_1 {
 				continue
 			}
 			bool := v_2
-			if !(config.arch != "arm64") {
+			if !(canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64") {
 				continue
 			}
 			v.reset(OpAdd16)
-			v0 := b.NewValue0(v.Pos, OpZeroExt8to16, t)
-			v1 := b.NewValue0(v.Pos, OpCvtBoolToUint8, types.Types[types.TUINT8])
-			v1.AddArg(bool)
-			v0.AddArg(v1)
+			v0 := b.NewValue0(v.Pos, OpMul16, t)
+			v1 := b.NewValue0(v.Pos, OpZeroExt8to16, t)
+			v2 := b.NewValue0(v.Pos, OpCvtBoolToUint8, types.Types[types.TUINT8])
+			v2.AddArg(bool)
+			v1.AddArg(v2)
+			v0.AddArg2(c, v1)
 			v.AddArg2(x, v0)
 			return true
 		}
 		break
 	}
-	// match: (CondSelect x (Add8 <t> x (Const8 [1])) bool)
-	// cond: config.arch != "arm64"
-	// result: (Add8 x (CvtBoolToUint8 <t> (Not <bool.Type> bool)))
+	// match: (CondSelect x op:(Add8 <t> x c:(Const8)) bool)
+	// cond: canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64"
+	// result: (Add8 x (Mul8 <t> c (CvtBoolToUint8 <t> (Not <bool.Type> bool))))
 	for {
 		x := v_0
-		if v_1.Op != OpAdd8 {
+		op := v_1
+		if op.Op != OpAdd8 {
 			break
 		}
-		t := v_1.Type
-		_ = v_1.Args[1]
-		v_1_0 := v_1.Args[0]
-		v_1_1 := v_1.Args[1]
-		for _i0 := 0; _i0 <= 1; _i0, v_1_0, v_1_1 = _i0+1, v_1_1, v_1_0 {
-			if x != v_1_0 || v_1_1.Op != OpConst8 || auxIntToInt8(v_1_1.AuxInt) != 1 {
+		t := op.Type
+		_ = op.Args[1]
+		op_0 := op.Args[0]
+		op_1 := op.Args[1]
+		for _i0 := 0; _i0 <= 1; _i0, op_0, op_1 = _i0+1, op_1, op_0 {
+			if x != op_0 {
+				continue
+			}
+			c := op_1
+			if c.Op != OpConst8 {
 				continue
 			}
 			bool := v_2
-			if !(config.arch != "arm64") {
+			if !(canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64") {
 				continue
 			}
 			v.reset(OpAdd8)
-			v0 := b.NewValue0(v.Pos, OpCvtBoolToUint8, t)
-			v1 := b.NewValue0(v.Pos, OpNot, bool.Type)
-			v1.AddArg(bool)
-			v0.AddArg(v1)
+			v0 := b.NewValue0(v.Pos, OpMul8, t)
+			v1 := b.NewValue0(v.Pos, OpCvtBoolToUint8, t)
+			v2 := b.NewValue0(v.Pos, OpNot, bool.Type)
+			v2.AddArg(bool)
+			v1.AddArg(v2)
+			v0.AddArg2(c, v1)
 			v.AddArg2(x, v0)
 			return true
 		}
 		break
 	}
-	// match: (CondSelect x (Add64 <t> x (Const64 [1])) bool)
-	// cond: config.arch != "arm64"
-	// result: (Add64 x (ZeroExt8to64 <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> (Not <bool.Type> bool))))
+	// match: (CondSelect x op:(Add64 <t> x c:(Const64)) bool)
+	// cond: canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64"
+	// result: (Add64 x (Mul64 <t> c (ZeroExt8to64 <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> (Not <bool.Type> bool)))))
 	for {
 		x := v_0
-		if v_1.Op != OpAdd64 {
+		op := v_1
+		if op.Op != OpAdd64 {
 			break
 		}
-		t := v_1.Type
-		_ = v_1.Args[1]
-		v_1_0 := v_1.Args[0]
-		v_1_1 := v_1.Args[1]
-		for _i0 := 0; _i0 <= 1; _i0, v_1_0, v_1_1 = _i0+1, v_1_1, v_1_0 {
-			if x != v_1_0 || v_1_1.Op != OpConst64 || auxIntToInt64(v_1_1.AuxInt) != 1 {
+		t := op.Type
+		_ = op.Args[1]
+		op_0 := op.Args[0]
+		op_1 := op.Args[1]
+		for _i0 := 0; _i0 <= 1; _i0, op_0, op_1 = _i0+1, op_1, op_0 {
+			if x != op_0 {
+				continue
+			}
+			c := op_1
+			if c.Op != OpConst64 {
 				continue
 			}
 			bool := v_2
-			if !(config.arch != "arm64") {
+			if !(canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64") {
 				continue
 			}
 			v.reset(OpAdd64)
-			v0 := b.NewValue0(v.Pos, OpZeroExt8to64, t)
-			v1 := b.NewValue0(v.Pos, OpCvtBoolToUint8, types.Types[types.TUINT8])
-			v2 := b.NewValue0(v.Pos, OpNot, bool.Type)
-			v2.AddArg(bool)
+			v0 := b.NewValue0(v.Pos, OpMul64, t)
+			v1 := b.NewValue0(v.Pos, OpZeroExt8to64, t)
+			v2 := b.NewValue0(v.Pos, OpCvtBoolToUint8, types.Types[types.TUINT8])
+			v3 := b.NewValue0(v.Pos, OpNot, bool.Type)
+			v3.AddArg(bool)
+			v2.AddArg(v3)
 			v1.AddArg(v2)
-			v0.AddArg(v1)
+			v0.AddArg2(c, v1)
 			v.AddArg2(x, v0)
 			return true
 		}
 		break
 	}
-	// match: (CondSelect x (Add32 <t> x (Const32 [1])) bool)
-	// cond: config.arch != "arm64"
-	// result: (Add32 x (ZeroExt8to32 <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> (Not <bool.Type> bool))))
+	// match: (CondSelect x op:(Add32 <t> x c:(Const32)) bool)
+	// cond: canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64"
+	// result: (Add32 x (Mul32 <t> c (ZeroExt8to32 <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> (Not <bool.Type> bool)))))
 	for {
 		x := v_0
-		if v_1.Op != OpAdd32 {
+		op := v_1
+		if op.Op != OpAdd32 {
 			break
 		}
-		t := v_1.Type
-		_ = v_1.Args[1]
-		v_1_0 := v_1.Args[0]
-		v_1_1 := v_1.Args[1]
-		for _i0 := 0; _i0 <= 1; _i0, v_1_0, v_1_1 = _i0+1, v_1_1, v_1_0 {
-			if x != v_1_0 || v_1_1.Op != OpConst32 || auxIntToInt32(v_1_1.AuxInt) != 1 {
+		t := op.Type
+		_ = op.Args[1]
+		op_0 := op.Args[0]
+		op_1 := op.Args[1]
+		for _i0 := 0; _i0 <= 1; _i0, op_0, op_1 = _i0+1, op_1, op_0 {
+			if x != op_0 {
+				continue
+			}
+			c := op_1
+			if c.Op != OpConst32 {
 				continue
 			}
 			bool := v_2
-			if !(config.arch != "arm64") {
+			if !(canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64") {
 				continue
 			}
 			v.reset(OpAdd32)
-			v0 := b.NewValue0(v.Pos, OpZeroExt8to32, t)
-			v1 := b.NewValue0(v.Pos, OpCvtBoolToUint8, types.Types[types.TUINT8])
-			v2 := b.NewValue0(v.Pos, OpNot, bool.Type)
-			v2.AddArg(bool)
+			v0 := b.NewValue0(v.Pos, OpMul32, t)
+			v1 := b.NewValue0(v.Pos, OpZeroExt8to32, t)
+			v2 := b.NewValue0(v.Pos, OpCvtBoolToUint8, types.Types[types.TUINT8])
+			v3 := b.NewValue0(v.Pos, OpNot, bool.Type)
+			v3.AddArg(bool)
+			v2.AddArg(v3)
 			v1.AddArg(v2)
-			v0.AddArg(v1)
+			v0.AddArg2(c, v1)
 			v.AddArg2(x, v0)
 			return true
 		}
 		break
 	}
-	// match: (CondSelect x (Add16 <t> x (Const16 [1])) bool)
-	// cond: config.arch != "arm64"
-	// result: (Add16 x (ZeroExt8to16 <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> (Not <bool.Type> bool))))
+	// match: (CondSelect x op:(Add16 <t> x c:(Const16)) bool)
+	// cond: canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64"
+	// result: (Add16 x (Mul16 <t> c (ZeroExt8to16 <t> (CvtBoolToUint8 <types.Types[types.TUINT8]> (Not <bool.Type> bool)))))
 	for {
 		x := v_0
-		if v_1.Op != OpAdd16 {
+		op := v_1
+		if op.Op != OpAdd16 {
 			break
 		}
-		t := v_1.Type
-		_ = v_1.Args[1]
-		v_1_0 := v_1.Args[0]
-		v_1_1 := v_1.Args[1]
-		for _i0 := 0; _i0 <= 1; _i0, v_1_0, v_1_1 = _i0+1, v_1_1, v_1_0 {
-			if x != v_1_0 || v_1_1.Op != OpConst16 || auxIntToInt16(v_1_1.AuxInt) != 1 {
+		t := op.Type
+		_ = op.Args[1]
+		op_0 := op.Args[0]
+		op_1 := op.Args[1]
+		for _i0 := 0; _i0 <= 1; _i0, op_0, op_1 = _i0+1, op_1, op_0 {
+			if x != op_0 {
+				continue
+			}
+			c := op_1
+			if c.Op != OpConst16 {
 				continue
 			}
 			bool := v_2
-			if !(config.arch != "arm64") {
+			if !(canImplementXOpBooleanLikeTimesAConstantCheaply(config, op.Op, c.AuxInt) && config.arch != "arm64") {
 				continue
 			}
 			v.reset(OpAdd16)
-			v0 := b.NewValue0(v.Pos, OpZeroExt8to16, t)
-			v1 := b.NewValue0(v.Pos, OpCvtBoolToUint8, types.Types[types.TUINT8])
-			v2 := b.NewValue0(v.Pos, OpNot, bool.Type)
-			v2.AddArg(bool)
+			v0 := b.NewValue0(v.Pos, OpMul16, t)
+			v1 := b.NewValue0(v.Pos, OpZeroExt8to16, t)
+			v2 := b.NewValue0(v.Pos, OpCvtBoolToUint8, types.Types[types.TUINT8])
+			v3 := b.NewValue0(v.Pos, OpNot, bool.Type)
+			v3.AddArg(bool)
+			v2.AddArg(v3)
 			v1.AddArg(v2)
-			v0.AddArg(v1)
+			v0.AddArg2(c, v1)
 			v.AddArg2(x, v0)
 			return true
 		}
diff --git a/test/codegen/condmove.go b/test/codegen/condmove.go
index ec1b971ed7..2b28077758 100644
--- a/test/codegen/condmove.go
+++ b/test/codegen/condmove.go
@@ -475,6 +475,51 @@ func cmovmathaddelse(a uint, b bool) uint {
 	return a
 }
 
+func cmovmathadd2(a uint, b bool) uint {
+	if b {
+		a += 2
+	}
+	// amd64:"LEAQ" -"CMOV" -"MUL"
+	return a
+}
+func cmovmathadd2else(a uint, b bool) uint {
+	if !b {
+		a += 2
+	}
+	// amd64:"LEAQ" -"CMOV" -"MUL"
+	return a
+}
+
+func cmovmathadd4(a uint, b bool) uint {
+	if b {
+		a += 4
+	}
+	// amd64:"LEAQ" -"CMOV" -"MUL"
+	return a
+}
+func cmovmathadd4else(a uint, b bool) uint {
+	if !b {
+		a += 4
+	}
+	// amd64:"LEAQ" -"CMOV" -"MUL"
+	return a
+}
+
+func cmovmathadd8(a uint, b bool) uint {
+	if b {
+		a += 8
+	}
+	// amd64:"LEAQ" -"CMOV" -"MUL"
+	return a
+}
+func cmovmathadd8else(a uint, b bool) uint {
+	if !b {
+		a += 8
+	}
+	// amd64:"LEAQ" -"CMOV" -"MUL"
+	return a
+}
+
 func cmovmathsub(a uint, b bool) uint {
 	if b {
 		a--