Repository: tetratelabs/wazero
Pull request #71 (2021-12-10 05:43:18)

jit: support select instruction
part of #65 

mathetake: done writing test!
mathetake: > Also, not to address in this PR, but I think we'll need a less hairy way to setup tests at some point as the logic to test is a little hairy. It looks very similar to writing main code and less like assertions. It would be cool to have a test utility at some point to setup the preconditions and assertions in an easier to read way. granted that is easier said than done!

Totally agree the test is getting hairy.. at some point before I finish all the impl of instructions, I would do some refactoring and add helpers to clean up! 

Patch:
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index c1abf746..287665bb 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -78,12 +78,15 @@ func (e *engine) compileWasmFunction(f *wasm.FunctionInstance) (*compiledWasmFun
 				return nil, fmt.Errorf("error handling call operation: %w", err)
 			}
 		case *wazeroir.OperationCallIndirect:
+			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
 		case *wazeroir.OperationDrop:
 			if err := builder.handleDrop(o); err != nil {
 				return nil, fmt.Errorf("error handling drop operation: %w", err)
 			}
 		case *wazeroir.OperationSelect:
-			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
+			if err := builder.handleSelect(); err != nil {
+				return nil, fmt.Errorf("error handling select operation: %w", err)
+			}
 		case *wazeroir.OperationPick:
 			if err := builder.handlePick(o); err != nil {
 				return nil, fmt.Errorf("error handling pick operation: %w", err)
@@ -374,9 +377,9 @@ func (b *amd64Builder) handleBrIf(o *wazeroir.OperationBrIf) error {
 		// as the operand for br_if instruction.
 		if cond.onStack() {
 			// This case even worse, the operand is not on a allocated register, but
-			// actually in the stack memory, so we have to assig a register to it
+			// actually in the stack memory, so we have to assign a register to it
 			// before we judge if we should jump to the Then branch or Else.
-			if err := b.moveStackToRegister(cond.registerType(), cond); err != nil {
+			if err := b.moveStackToRegisterWithAllocation(cond.registerType(), cond); err != nil {
 				return err
 			}
 		}
@@ -552,7 +555,7 @@ func (b *amd64Builder) emitDropRange(r *wazeroir.InclusiveRange) error {
 				topIsConditional = false
 			}
 			// Write the value in the old stack location to a register
-			if err := b.moveStackToRegister(live.registerType(), live); err != nil {
+			if err := b.moveStackToRegisterWithAllocation(live.registerType(), live); err != nil {
 				return err
 			}
 			// Modify the location in the stack with new stack pointer.
@@ -564,6 +567,84 @@ func (b *amd64Builder) emitDropRange(r *wazeroir.InclusiveRange) error {
 	return nil
 }
 
+// handleSelect uses top three values on the stack:
+// Assume we have stack as [..., x1, x2, c], if the value of c
+// equals zero, then the stack results in [..., x1]
+// otherwise, [..., x2].
+// The emitted native code depends on whether the values are on
+// the physical registers or memory stack, or maybe conditional register.
+func (b *amd64Builder) handleSelect() error {
+	c := b.locationStack.pop()
+	x2 := b.locationStack.pop()
+	// We do not consume x1 here, but modify the value according to
+	// the conditional value "c" above.
+	peekedX1 := b.locationStack.peek()
+
+	// Ensure the conditional value lives in a gp register.
+	if c.onConditionalRegister() {
+		if err := b.moveConditionalToGPRegister(c); err != nil {
+			return err
+		}
+	} else if c.onStack() {
+		if err := b.moveStackToRegisterWithAllocation(c.registerType(), c); err != nil {
+			return err
+		}
+	}
+
+	// Compare the conditional value with zero.
+	cmpZero := b.newProg()
+	cmpZero.As = x86.ACMPQ
+	cmpZero.From.Type = obj.TYPE_REG
+	cmpZero.From.Reg = c.register
+	cmpZero.To.Type = obj.TYPE_CONST
+	cmpZero.To.Offset = 0
+	b.addInstruction(cmpZero)
+
+	// Now we can use c.register as temporary location.
+	// We alias it here for readability.
+	tmpRegister := c.register
+
+	// Set the jump if the top value is not zero.
+	jmpIfNotZero := b.newProg()
+	jmpIfNotZero.As = x86.AJNE
+	jmpIfNotZero.To.Type = obj.TYPE_BRANCH
+	b.addInstruction(jmpIfNotZero)
+
+	// If the value is zero, we must place the value of x2 onto the stack position of x1.
+
+	// First we copy the value of x2 to the temporary register if x2 is not currently on a register.
+	if x2.onStack() {
+		x2.register = tmpRegister
+		b.moveStackToRegister(x2)
+	}
+
+	//
+	// At this point x2's value is always on a register.
+	//
+
+	// Then release the value in the x2's register to the x1's stack position.
+	if peekedX1.onRegister() {
+		movX2ToX1 := b.newProg()
+		movX2ToX1.As = x86.AMOVQ
+		movX2ToX1.From.Type = obj.TYPE_REG
+		movX2ToX1.From.Reg = x2.register
+		movX2ToX1.To.Type = obj.TYPE_REG
+		movX2ToX1.To.Reg = peekedX1.register
+		b.addInstruction(movX2ToX1)
+	} else {
+		peekedX1.register = x2.register
+		b.releaseRegisterToStack(peekedX1) // Note inside we mark the register unused!
+	}
+
+	// Else, we don't need to adjust value, just need to jump to the next instruction.
+	b.setJmpOrigin = jmpIfNotZero
+
+	// In any case, we don't need x2 and c anymore!
+	b.locationStack.releaseRegister(x2)
+	b.locationStack.releaseRegister(c)
+	return nil
+}
+
 func (b *amd64Builder) handlePick(o *wazeroir.OperationPick) error {
 	// TODO: if we track the type of values on the stack,
 	// we could optimize the instruction according to the bit size of the value.
@@ -628,7 +709,7 @@ func (b *amd64Builder) handleAdd(o *wazeroir.OperationAdd) error {
 
 	x2 := b.locationStack.pop()
 	if x2.onStack() {
-		if err := b.moveStackToRegister(tp, x2); err != nil {
+		if err := b.moveStackToRegisterWithAllocation(tp, x2); err != nil {
 			return err
 		}
 	} else if x2.onConditionalRegister() {
@@ -639,7 +720,7 @@ func (b *amd64Builder) handleAdd(o *wazeroir.OperationAdd) error {
 
 	x1 := b.locationStack.peek() // Note this is peek, pop!
 	if x1.onStack() {
-		if err := b.moveStackToRegister(tp, x1); err != nil {
+		if err := b.moveStackToRegisterWithAllocation(tp, x1); err != nil {
 			return err
 		}
 	} else if x1.onConditionalRegister() {
@@ -689,7 +770,7 @@ func (b *amd64Builder) handleSub(o *wazeroir.OperationSub) error {
 
 	x2 := b.locationStack.pop()
 	if x2.onStack() {
-		if err := b.moveStackToRegister(tp, x2); err != nil {
+		if err := b.moveStackToRegisterWithAllocation(tp, x2); err != nil {
 			return err
 		}
 	} else if x2.onConditionalRegister() {
@@ -700,7 +781,7 @@ func (b *amd64Builder) handleSub(o *wazeroir.OperationSub) error {
 
 	x1 := b.locationStack.peek() // Note this is peek, pop!
 	if x1.onStack() {
-		if err := b.moveStackToRegister(tp, x1); err != nil {
+		if err := b.moveStackToRegisterWithAllocation(tp, x1); err != nil {
 			return err
 		}
 	} else if x1.onConditionalRegister() {
@@ -755,7 +836,7 @@ func (b *amd64Builder) handleLe(o *wazeroir.OperationLe) error {
 
 	x2 := b.locationStack.pop()
 	if x2.onStack() {
-		if err := b.moveStackToRegister(tp, x2); err != nil {
+		if err := b.moveStackToRegisterWithAllocation(tp, x2); err != nil {
 			return err
 		}
 	} else if x2.onConditionalRegister() {
@@ -766,7 +847,7 @@ func (b *amd64Builder) handleLe(o *wazeroir.OperationLe) error {
 
 	x1 := b.locationStack.pop()
 	if x1.onStack() {
-		if err := b.moveStackToRegister(tp, x1); err != nil {
+		if err := b.moveStackToRegisterWithAllocation(tp, x1); err != nil {
 			return err
 		}
 	} else if x1.onConditionalRegister() {
@@ -807,27 +888,33 @@ func (b *amd64Builder) handleConstI64(o *wazeroir.OperationConstI64) error {
 	return nil
 }
 
-func (b *amd64Builder) moveStackToRegister(tp generalPurposeRegisterType, loc *valueLocation) error {
+// TODO: maybe split this function as this is doing too much at once to say at once.
+func (b *amd64Builder) moveStackToRegisterWithAllocation(tp generalPurposeRegisterType, loc *valueLocation) error {
 	// Allocate the register.
 	reg, err := b.allocateRegister(tp)
 	if err != nil {
 		return err
 	}
 
-	// Then copy the value from the stack.
+	// Mark it uses the register.
+	loc.setRegister(reg)
+	b.locationStack.markRegisterUsed(reg)
+
+	// Now ready to move value.
+	b.moveStackToRegister(loc)
+	return nil
+}
+
+func (b *amd64Builder) moveStackToRegister(loc *valueLocation) {
+	// Copy the value from the stack.
 	prog := b.newProg()
 	prog.As = x86.AMOVQ
 	prog.From.Type = obj.TYPE_MEM
 	prog.From.Reg = cachedStackBasePointerReg
 	prog.From.Offset = int64(loc.stackPointer) * 8
 	prog.To.Type = obj.TYPE_REG
-	prog.To.Reg = reg
+	prog.To.Reg = loc.register
 	b.addInstruction(prog)
-
-	// Mark it uses the register.
-	loc.setRegister(reg)
-	b.locationStack.markRegisterUsed(reg)
-	return nil
 }
 
 func (b *amd64Builder) moveConditionalToGPRegister(loc *valueLocation) error {
@@ -908,7 +995,7 @@ func (b *amd64Builder) allocateRegister(t generalPurposeRegisterType) (reg int16
 
 	// Release the steal target register value onto stack location.
 	reg = stealTarget.register
-	b.releaseRegister(stealTarget)
+	b.releaseRegisterToStack(stealTarget)
 	return
 }
 
@@ -985,7 +1072,7 @@ func (b *amd64Builder) releaseAllRegistersToStack() {
 	used := len(b.locationStack.usedRegisters)
 	for i := len(b.locationStack.stack) - 1; i >= 0 && used > 0; i-- {
 		if loc := b.locationStack.stack[i]; loc.onRegister() {
-			b.releaseRegister(loc)
+			b.releaseRegisterToStack(loc)
 			used--
 		}
 	}
@@ -1063,7 +1150,7 @@ func (b *amd64Builder) movConstToRegister(val int64, targetRegister int16) *obj.
 	return prog
 }
 
-func (b *amd64Builder) releaseRegister(loc *valueLocation) {
+func (b *amd64Builder) releaseRegisterToStack(loc *valueLocation) {
 	// Push value.
 	prog := b.newProg()
 	prog.As = x86.AMOVQ


Test patch:
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index 5c7656dd..9b833e1e 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -84,7 +84,7 @@ func TestRecursiveFunctionCalls(t *testing.T) {
 	builder.addInstruction(jmp)
 	// If not zero, we call push back the value to the stack
 	// and call itself recursively.
-	builder.releaseRegister(loc)
+	builder.releaseRegisterToStack(loc)
 	require.NotContains(t, builder.locationStack.usedRegisters, loc.register)
 	builder.callFunctionFromConstIndex(0)
 	// ::End
@@ -92,7 +92,7 @@ func TestRecursiveFunctionCalls(t *testing.T) {
 	builder.assignRegisterToValue(loc, tmpReg)
 	prog = builder.movConstToRegister(5, loc.register)
 	jmp.To.SetTarget(prog) // the above mov instruction is the jump target of the JEQ.
-	builder.releaseRegister(loc)
+	builder.releaseRegisterToStack(loc)
 	builder.setJITStatus(jitCallStatusCodeReturned)
 	builder.returnFunction()
 	// Compile.
@@ -157,7 +157,7 @@ func TestPushValueWithGoroutines(t *testing.T) {
 			loc := builder.locationStack.pushValueOnRegister(pushTargetRegister)
 			builder.movConstToRegister(int64(targetValue), pushTargetRegister)
 			// Push pushTargetRegister into the engine.stack[engine.sp].
-			builder.releaseRegister(loc)
+			builder.releaseRegisterToStack(loc)
 			// Finally increment the stack pointer and write it back to the eng.sp
 			builder.returnFunction()
 
@@ -282,7 +282,7 @@ func Test_setContinuationAtNextInstruction(t *testing.T) {
 	_ = builder.locationStack.pushValueOnStack() // dummy value, not actually used!
 	loc := builder.locationStack.pushValueOnRegister(tmpReg)
 	builder.movConstToRegister(int64(50), tmpReg)
-	builder.releaseRegister(loc)
+	builder.releaseRegisterToStack(loc)
 	require.NotContains(t, builder.locationStack.usedRegisters, tmpReg)
 	builder.setJITStatus(jitCallStatusCodeCallWasmFunction)
 	builder.returnFunction()
@@ -326,7 +326,7 @@ func Test_callFunction(t *testing.T) {
 		_ = builder.locationStack.pushValueOnStack() // dummy value, not actually used!
 		loc := builder.locationStack.pushValueOnRegister(tmpReg)
 		builder.movConstToRegister(int64(50), tmpReg)
-		builder.releaseRegister(loc)
+		builder.releaseRegisterToStack(loc)
 		require.NotContains(t, builder.locationStack.usedRegisters, tmpReg)
 		builder.returnFunction()
 		// Compile.
@@ -366,7 +366,7 @@ func Test_callFunction(t *testing.T) {
 		_ = builder.locationStack.pushValueOnStack() // dummy value, not actually used!
 		loc := builder.locationStack.pushValueOnRegister(tmpReg)
 		builder.movConstToRegister(int64(50), tmpReg)
-		builder.releaseRegister(loc)
+		builder.releaseRegisterToStack(loc)
 		require.NotContains(t, builder.locationStack.usedRegisters, tmpReg)
 		builder.returnFunction()
 		// Compile.
@@ -409,7 +409,7 @@ func TestEngine_exec_callHostFunction(t *testing.T) {
 		_ = builder.locationStack.pushValueOnStack() // dummy value, not actually used!
 		loc := builder.locationStack.pushValueOnRegister(tmpReg)
 		builder.movConstToRegister(int64(50), tmpReg)
-		builder.releaseRegister(loc)
+		builder.releaseRegisterToStack(loc)
 		require.NotContains(t, builder.locationStack.usedRegisters, tmpReg)
 		builder.callHostFunctionFromConstIndex(functionIndex)
 		// On the continuation after function call,
@@ -449,7 +449,7 @@ func TestEngine_exec_callHostFunction(t *testing.T) {
 		_ = builder.locationStack.pushValueOnStack() // dummy value, not actually used!
 		loc := builder.locationStack.pushValueOnRegister(x86.REG_AX)
 		builder.movConstToRegister(int64(50), tmpReg)
-		builder.releaseRegister(loc)
+		builder.releaseRegisterToStack(loc)
 		require.NotContains(t, builder.locationStack.usedRegisters, tmpReg)
 		// Set the function index
 		builder.movConstToRegister(int64(1), tmpReg)
@@ -518,7 +518,7 @@ func Test_popFromStackToRegister(t *testing.T) {
 	result := builder.locationStack.pushValueOnRegister(targetRegister1)
 	builder.addInstruction(prog)
 	// Push it back to the stack.
-	builder.releaseRegister(result)
+	builder.releaseRegisterToStack(result)
 	builder.returnFunction()
 	// Compile.
 	code, err := builder.assemble()
@@ -590,7 +590,7 @@ func TestAmd64Builder_allocateRegister(t *testing.T) {
 		// Create new value using the stolen register.
 		loc := builder.locationStack.pushValueOnRegister(reg)
 		builder.movConstToRegister(int64(2000), loc.register)
-		builder.releaseRegister(loc)
+		builder.releaseRegisterToStack(loc)
 		builder.returnFunction()
 
 		// Assemble.
@@ -668,9 +668,9 @@ func TestAmd64Builder_handlePick(t *testing.T) {
 		builder.addInstruction(prog)
 		// To verify the behavior, we push the incremented picked value
 		// to the stack.
-		builder.releaseRegister(pickedLocation)
+		builder.releaseRegisterToStack(pickedLocation)
 		// Also write the original location back to the stack.
-		builder.releaseRegister(pickTargetLocation)
+		builder.releaseRegisterToStack(pickTargetLocation)
 		builder.returnFunction()
 
 		// Assemble.
@@ -715,7 +715,7 @@ func TestAmd64Builder_handlePick(t *testing.T) {
 
 		// To verify the behavior, we push the incremented picked value
 		// to the stack.
-		builder.releaseRegister(pickedLocation)
+		builder.releaseRegisterToStack(pickedLocation)
 		builder.returnFunction()
 
 		// Assemble.
@@ -757,7 +757,7 @@ func TestAmd64Builder_handleConstI64(t *testing.T) {
 			prog.To.Type = obj.TYPE_REG
 			prog.To.Reg = loc.register
 			builder.addInstruction(prog)
-			builder.releaseRegister(loc)
+			builder.releaseRegisterToStack(loc)
 			builder.returnFunction()
 
 			// Assemble.
@@ -798,7 +798,7 @@ func TestAmd64Builder_handleAdd(t *testing.T) {
 
 			// To verify the behavior, we push the value
 			// to the stack.
-			builder.releaseRegister(x1Location)
+			builder.releaseRegisterToStack(x1Location)
 			builder.returnFunction()
 
 			// Assemble.
@@ -830,7 +830,7 @@ func TestAmd64Builder_handleAdd(t *testing.T) {
 
 			// To verify the behavior, we push the value
 			// to the stack.
-			builder.releaseRegister(builder.locationStack.peek())
+			builder.releaseRegisterToStack(builder.locationStack.peek())
 			builder.returnFunction()
 
 			// Assemble.
@@ -864,7 +864,7 @@ func TestAmd64Builder_handleAdd(t *testing.T) {
 
 			// To verify the behavior, we push the value
 			// to the stack.
-			builder.releaseRegister(builder.locationStack.peek())
+			builder.releaseRegisterToStack(builder.locationStack.peek())
 			builder.returnFunction()
 
 			// Assemble.
@@ -898,7 +898,7 @@ func TestAmd64Builder_handleAdd(t *testing.T) {
 
 			// To verify the behavior, we push the value
 			// to the stack.
-			builder.releaseRegister(builder.locationStack.peek())
+			builder.releaseRegisterToStack(builder.locationStack.peek())
 			builder.returnFunction()
 
 			// Assemble.
@@ -957,7 +957,7 @@ func TestAmd64Builder_handleLe(t *testing.T) {
 			err = builder.moveConditionalToGPRegister(top)
 			require.NoError(t, err)
 			require.True(t, !top.onConditionalRegister() && top.onRegister())
-			builder.releaseRegister(top)
+			builder.releaseRegisterToStack(top)
 			builder.returnFunction()
 
 			// Assemble.
@@ -1015,7 +1015,7 @@ func TestAmd64Builder_handleLe(t *testing.T) {
 			err = builder.moveConditionalToGPRegister(top)
 			require.NoError(t, err)
 			require.True(t, !top.onConditionalRegister() && top.onRegister())
-			builder.releaseRegister(top)
+			builder.releaseRegisterToStack(top)
 			builder.returnFunction()
 
 			// Assemble.
@@ -1059,7 +1059,7 @@ func TestAmd64Builder_handleSub(t *testing.T) {
 
 			// To verify the behavior, we push the value
 			// to the stack.
-			builder.releaseRegister(x1Location)
+			builder.releaseRegisterToStack(x1Location)
 			builder.returnFunction()
 
 			// Assemble.
@@ -1091,7 +1091,7 @@ func TestAmd64Builder_handleSub(t *testing.T) {
 
 			// To verify the behavior, we push the value
 			// to the stack.
-			builder.releaseRegister(builder.locationStack.peek())
+			builder.releaseRegisterToStack(builder.locationStack.peek())
 			builder.returnFunction()
 
 			// Assemble.
@@ -1125,7 +1125,7 @@ func TestAmd64Builder_handleSub(t *testing.T) {
 
 			// To verify the behavior, we push the value
 			// to the stack.
-			builder.releaseRegister(builder.locationStack.peek())
+			builder.releaseRegisterToStack(builder.locationStack.peek())
 			builder.returnFunction()
 
 			// Assemble.
@@ -1159,7 +1159,7 @@ func TestAmd64Builder_handleSub(t *testing.T) {
 
 			// To verify the behavior, we push the value
 			// to the stack.
-			builder.releaseRegister(builder.locationStack.peek())
+			builder.releaseRegisterToStack(builder.locationStack.peek())
 			builder.returnFunction()
 
 			// Assemble.
@@ -1375,8 +1375,8 @@ func TestAmd64Builder_handleDrop(t *testing.T) {
 				Range: &wazeroir.InclusiveRange{Start: 1, End: 1},
 			})
 			require.NoError(t, err)
-			builder.releaseRegister(bottom)
-			builder.releaseRegister(top)
+			builder.releaseRegisterToStack(bottom)
+			builder.releaseRegisterToStack(top)
 			builder.returnFunction()
 			// Assemble.
 			code, err := builder.assemble()
@@ -1481,3 +1481,134 @@ func TestAmd64Builder_handleUnreachable(t *testing.T) {
 	require.Equal(t, uint64(300), eng.stack[0])
 	require.Equal(t, uint64(51), eng.stack[1])
 }
+
+func TestAmd64Builder_handleSelect(t *testing.T) {
+	// There are mainly 8 cases we have to test:
+	// - [x1 = reg, x2 = reg] select x1
+	// - [x1 = reg, x2 = reg] select x2
+	// - [x1 = reg, x2 = stack] select x1
+	// - [x1 = reg, x2 = stack] select x2
+	// - [x1 = stack, x2 = reg] select x1
+	// - [x1 = stack, x2 = reg] select x2
+	// - [x1 = stack, x2 = stack] select x1
+	// - [x1 = stack, x2 = stack] select x2
+	// And for each case, we have to test with
+	// three conditional value location: stack, gp register, conditional register.
+	// So in total we have 24 cases.
+	const x1Value, x2Value = 100, 200
+	for i, tc := range []struct {
+		x1OnRegister, x2OnRegister                                        bool
+		selectX1                                                          bool
+		condlValueOnStack, condValueOnGPRegister, condValueOnCondRegister bool
+	}{
+		// Conditional value on stack.
+		{x1OnRegister: true, x2OnRegister: true, selectX1: true, condlValueOnStack: true},
+		{x1OnRegister: true, x2OnRegister: true, selectX1: false, condlValueOnStack: true},
+		{x1OnRegister: true, x2OnRegister: false, selectX1: true, condlValueOnStack: true},
+		{x1OnRegister: true, x2OnRegister: false, selectX1: false, condlValueOnStack: true},
+		{x1OnRegister: false, x2OnRegister: true, selectX1: true, condlValueOnStack: true},
+		{x1OnRegister: false, x2OnRegister: true, selectX1: false, condlValueOnStack: true},
+		{x1OnRegister: false, x2OnRegister: false, selectX1: true, condlValueOnStack: true},
+		{x1OnRegister: false, x2OnRegister: false, selectX1: false, condlValueOnStack: true},
+		// Conditional value on register.
+		{x1OnRegister: true, x2OnRegister: true, selectX1: true, condValueOnGPRegister: true},
+		{x1OnRegister: true, x2OnRegister: true, selectX1: false, condValueOnGPRegister: true},
+		{x1OnRegister: true, x2OnRegister: false, selectX1: true, condValueOnGPRegister: true},
+		{x1OnRegister: true, x2OnRegister: false, selectX1: false, condValueOnGPRegister: true},
+		{x1OnRegister: false, x2OnRegister: true, selectX1: true, condValueOnGPRegister: true},
+		{x1OnRegister: false, x2OnRegister: true, selectX1: false, condValueOnGPRegister: true},
+		{x1OnRegister: false, x2OnRegister: false, selectX1: true, condValueOnGPRegister: true},
+		{x1OnRegister: false, x2OnRegister: false, selectX1: false, condValueOnGPRegister: true},
+		// Conditional value on conditional register.
+		{x1OnRegister: true, x2OnRegister: true, selectX1: true, condValueOnCondRegister: true},
+		{x1OnRegister: true, x2OnRegister: true, selectX1: false, condValueOnCondRegister: true},
+		{x1OnRegister: true, x2OnRegister: false, selectX1: true, condValueOnCondRegister: true},
+		{x1OnRegister: true, x2OnRegister: false, selectX1: false, condValueOnCondRegister: true},
+		{x1OnRegister: false, x2OnRegister: true, selectX1: true, condValueOnCondRegister: true},
+		{x1OnRegister: false, x2OnRegister: true, selectX1: false, condValueOnCondRegister: true},
+		{x1OnRegister: false, x2OnRegister: false, selectX1: true, condValueOnCondRegister: true},
+		{x1OnRegister: false, x2OnRegister: false, selectX1: false, condValueOnCondRegister: true},
+	} {
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			builder := requireNewBuilder(t)
+			eng := newEngine()
+			builder.initializeReservedRegisters()
+			var x1, x2, c *valueLocation
+			if tc.x1OnRegister {
+				x1 = builder.locationStack.pushValueOnRegister(x86.REG_AX)
+				builder.movConstToRegister(x1Value, x1.register)
+			} else {
+				x1 = builder.locationStack.pushValueOnStack()
+				eng.stack[x1.stackPointer] = x1Value
+			}
+			if tc.x2OnRegister {
+				x2 = builder.locationStack.pushValueOnRegister(x86.REG_R10)
+				builder.movConstToRegister(x2Value, x2.register)
+			} else {
+				x2 = builder.locationStack.pushValueOnStack()
+				eng.stack[x2.stackPointer] = x2Value
+			}
+			if tc.condlValueOnStack {
+				c = builder.locationStack.pushValueOnStack()
+				if tc.selectX1 {
+					eng.stack[c.stackPointer] = 1
+				} else {
+					eng.stack[c.stackPointer] = 0
+				}
+			} else if tc.condValueOnGPRegister {
+				c = builder.locationStack.pushValueOnRegister(x86.REG_R9)
+				if tc.selectX1 {
+					builder.movConstToRegister(1, c.register)
+				} else {
+					builder.movConstToRegister(0, c.register)
+				}
+			} else if tc.condValueOnCondRegister {
+				builder.movConstToRegister(0, x86.REG_CX)
+				cmp := builder.newProg()
+				cmp.As = x86.ACMPQ
+				cmp.From.Type = obj.TYPE_REG
+				cmp.From.Reg = x86.REG_CX
+				cmp.To.Type = obj.TYPE_CONST
+				if tc.selectX1 {
+					cmp.To.Offset = 0
+				} else {
+					cmp.To.Offset = 1
+				}
+				builder.addInstruction(cmp)
+				builder.locationStack.pushValueOnConditionalRegister(conditionalRegisterStateE)
+			}
+
+			// Now emit code for select.
+			err := builder.handleSelect()
+			require.NoError(t, err)
+			// The code generation should not affect the x1's placement in any case.
+			require.Equal(t, tc.x1OnRegister, x1.onRegister())
+			// Plus x1 is top of the stack.
+			require.Equal(t, x1, builder.locationStack.peek())
+
+			// Now write back the x1 to the memory if it is on a register.
+			if tc.x1OnRegister {
+				builder.releaseRegisterToStack(x1)
+			}
+			builder.returnFunction()
+
+			// Run code.
+			code, err := builder.assemble()
+			require.NoError(t, err)
+			mem := newMemoryInst()
+			jitcall(
+				uintptr(unsafe.Pointer(&code[0])),
+				uintptr(unsafe.Pointer(eng)),
+				uintptr(unsafe.Pointer(&mem.Buffer[0])),
+			)
+
+			// Check the selected value.
+			require.Equal(t, uint64(1), eng.currentStackPointer)
+			if tc.selectX1 {
+				require.Equal(t, eng.stack[x1.stackPointer], uint64(x1Value))
+			} else {
+				require.Equal(t, eng.stack[x1.stackPointer], uint64(x2Value))
+			}
+		})
+	}
+}