Repository: tetratelabs/wazero
Pull request #113 (2022-01-04 01:11:00)

jit: support for clz, ctz and popcnt
- [x] clz
- [x]  ctz
- [x] popcnt

mathetake: lol CLZ is only failing on macos 😄 
mathetake: Looks like we have to implement something like this only for mac (from https://developer.apple.com/documentation/apple-silicon/addressing-architectural-differences-in-your-macos-code)
```
int cnt(int x){
    int count;
    if(x == 0){
        return (sizeof(x) * CHAR_BIT);
    }
    __asm__ (
        "bsrl %1, %0\n\t" // undefined when operating on 0
        "xorl $0x1f, %1"
        : "=r" (count)
        : "r" (x)
    );
    return count;
}
```
mathetake: I hope https://github.com/tetratelabs/wazero/pull/113/commits/898a8b639a608363b9f9e90916a99d00b80e1172 fixed for mac.
mathetake: now only failing for the case where the target value is zero
mathetake: ok realized that now CLZ succeeded on mac but the only failing case is for CTZ
mathetake: ready for review! 

Patch:
diff --git a/wasm/jit/compiler.go b/wasm/jit/compiler.go
index df03d693..d13cc9b5 100644
--- a/wasm/jit/compiler.go
+++ b/wasm/jit/compiler.go
@@ -26,6 +26,9 @@ type compiler interface {
 	compileAdd(o *wazeroir.OperationAdd) error
 	compileSub(o *wazeroir.OperationSub) error
 	compileMul(o *wazeroir.OperationMul) error
+	compileClz(o *wazeroir.OperationClz) error
+	compileCtz(o *wazeroir.OperationCtz) error
+	compilePopcnt(o *wazeroir.OperationPopcnt) error
 	compileEq(o *wazeroir.OperationEq) error
 	compileNe(o *wazeroir.OperationNe) error
 	compileEqz(o *wazeroir.OperationEqz) error
diff --git a/wasm/jit/engine.go b/wasm/jit/engine.go
index 8277135a..3d870a87 100644
--- a/wasm/jit/engine.go
+++ b/wasm/jit/engine.go
@@ -539,11 +539,11 @@ func (e *engine) compileWasmFunction(f *wasm.FunctionInstance) (*compiledWasmFun
 		case *wazeroir.OperationMul:
 			err = compiler.compileMul(o)
 		case *wazeroir.OperationClz:
-			err = fmt.Errorf("unsupported operation")
+			err = compiler.compileClz(o)
 		case *wazeroir.OperationCtz:
-			err = fmt.Errorf("unsupported operation")
+			err = compiler.compileCtz(o)
 		case *wazeroir.OperationPopcnt:
-			err = fmt.Errorf("unsupported operation")
+			err = compiler.compilePopcnt(o)
 		case *wazeroir.OperationDiv:
 			err = fmt.Errorf("unsupported operation")
 		case *wazeroir.OperationRem:
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index 0667f1de..3b9e7f15 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -13,6 +13,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"math"
+	"runtime"
 
 	asm "github.com/twitchyliquid64/golang-asm"
 	"github.com/twitchyliquid64/golang-asm/obj"
@@ -886,7 +887,8 @@ func (c *amd64Compiler) compileMulForInts(is32Bit bool, mulInstruction obj.As) e
 
 	// Now we have the result in the AX register,
 	// so we record it.
-	c.locationStack.pushValueOnRegister(resultRegister)
+	result := c.locationStack.pushValueOnRegister(resultRegister)
+	result.setRegisterType(generalPurposeRegisterTypeInt)
 	return nil
 }
 
@@ -916,6 +918,210 @@ func (c *amd64Compiler) compileMulForFloats(instruction obj.As) error {
 	return nil
 }
 
+// compileClz emits instructions to count up the leading zeros in the
+// current top of the stack, and push the count result.
+// For example, stack of [..., 0x00_ff_ff_ff] results in [..., 8].
+func (c *amd64Compiler) compileClz(o *wazeroir.OperationClz) error {
+	target := c.locationStack.pop()
+	if err := c.ensureOnGeneralPurposeRegister(target); err != nil {
+		return err
+	}
+
+	if runtime.GOOS != "darwin" {
+		countZeros := c.newProg()
+		countZeros.From.Type = obj.TYPE_REG
+		countZeros.From.Reg = target.register
+		countZeros.To.Type = obj.TYPE_REG
+		countZeros.To.Reg = target.register
+		if o.Type == wazeroir.UnsignedInt32 {
+			countZeros.As = x86.ALZCNTL
+		} else {
+			countZeros.As = x86.ALZCNTQ
+		}
+		c.addInstruction(countZeros)
+	} else {
+		// On x86 mac, we cannot use LZCNT as it always results in zero.
+		// Instead we combine BSR (calculating most significant set bit)
+		// with XOR. This logic is described in
+		// "Replace Raw Assembly Code with Builtin Intrinsics" section in:
+		// https://developer.apple.com/documentation/apple-silicon/addressing-architectural-differences-in-your-macos-code.
+
+		// First, we have to check if the target is non-zero as BSR is undefined
+		// on zero. See https://www.felixcloutier.com/x86/bsr.
+		cmpZero := c.newProg()
+		cmpZero.As = x86.ACMPQ
+		cmpZero.From.Type = obj.TYPE_REG
+		cmpZero.From.Reg = target.register
+		cmpZero.To.Type = obj.TYPE_CONST
+		cmpZero.To.Offset = 0
+		c.addInstruction(cmpZero)
+
+		jmpIfNonZero := c.newProg()
+		jmpIfNonZero.As = x86.AJNE
+		jmpIfNonZero.To.Type = obj.TYPE_BRANCH
+		c.addInstruction(jmpIfNonZero)
+
+		// If the value is zero, we just push the const value.
+		if o.Type == wazeroir.UnsignedInt32 {
+			c.emitConstI32(32, target.register)
+		} else {
+			c.emitConstI64(64, target.register)
+		}
+
+		// Emit the jmp instruction to jump to the position right after
+		// the non-zero case.
+		jmpAtEndOfZero := c.newProg()
+		jmpAtEndOfZero.As = obj.AJMP
+		jmpAtEndOfZero.To.Type = obj.TYPE_BRANCH
+		c.addInstruction(jmpAtEndOfZero)
+
+		// Start emitting non-zero case.
+		// First, we calculate the most significant set bit.
+		mostSignificantSetBit := c.newProg()
+		// Set the jump target of the first non-zero check.
+		jmpIfNonZero.To.SetTarget(mostSignificantSetBit)
+		if o.Type == wazeroir.UnsignedInt32 {
+			mostSignificantSetBit.As = x86.ABSRL
+		} else {
+			mostSignificantSetBit.As = x86.ABSRQ
+		}
+		mostSignificantSetBit.From.Type = obj.TYPE_REG
+		mostSignificantSetBit.From.Reg = target.register
+		mostSignificantSetBit.To.Type = obj.TYPE_REG
+		mostSignificantSetBit.To.Reg = target.register
+		c.addInstruction(mostSignificantSetBit)
+
+		// Now we XOR the value with the bit length minus one.
+		xorWithBitLength := c.newProg()
+		xorWithBitLength.To.Type = obj.TYPE_REG
+		xorWithBitLength.To.Reg = target.register
+		xorWithBitLength.From.Type = obj.TYPE_CONST
+		if o.Type == wazeroir.UnsignedInt32 {
+			xorWithBitLength.As = x86.AXORL
+			xorWithBitLength.From.Offset = 31
+		} else {
+			xorWithBitLength.As = x86.AXORQ
+			xorWithBitLength.From.Offset = 63
+		}
+		c.addInstruction(xorWithBitLength)
+
+		// Finally the end jump instruction of zero case must target towards
+		// the next instruction.
+		c.setJmpOrigin = jmpAtEndOfZero
+	}
+
+	// We reused the same register of target for the result.
+	result := c.locationStack.pushValueOnRegister(target.register)
+	result.setRegisterType(generalPurposeRegisterTypeInt)
+	return nil
+}
+
+// compileCtz emits instructions to count up the trailing zeros in the
+// current top of the stack, and push the count result.
+// For example, stack of [..., 0xff_ff_ff_00] results in [..., 8].
+func (c *amd64Compiler) compileCtz(o *wazeroir.OperationCtz) error {
+	target := c.locationStack.pop()
+	if err := c.ensureOnGeneralPurposeRegister(target); err != nil {
+		return err
+	}
+
+	if runtime.GOOS != "darwin" {
+		countZeros := c.newProg()
+		countZeros.From.Type = obj.TYPE_REG
+		countZeros.From.Reg = target.register
+		countZeros.To.Type = obj.TYPE_REG
+		countZeros.To.Reg = target.register
+		if o.Type == wazeroir.UnsignedInt32 {
+			countZeros.As = x86.ATZCNTL
+		} else {
+			countZeros.As = x86.ATZCNTQ
+		}
+		c.addInstruction(countZeros)
+	} else {
+		// Somehow, if the target value is zero, TZCNT always returns zero: this is wrong.
+		// Meanwhile, we need branches for non-zero and zero cases on macos.
+		// TODO: find the reference to this behavior and put the link here.
+
+		// First we compare the target with zero.
+		cmpZero := c.newProg()
+		cmpZero.As = x86.ACMPQ
+		cmpZero.From.Type = obj.TYPE_REG
+		cmpZero.From.Reg = target.register
+		cmpZero.To.Type = obj.TYPE_CONST
+		cmpZero.To.Offset = 0
+		c.addInstruction(cmpZero)
+
+		jmpIfNonZero := c.newProg()
+		jmpIfNonZero.As = x86.AJNE
+		jmpIfNonZero.To.Type = obj.TYPE_BRANCH
+		c.addInstruction(jmpIfNonZero)
+
+		// If the value is zero, we just push the const value.
+		if o.Type == wazeroir.UnsignedInt32 {
+			c.emitConstI32(32, target.register)
+		} else {
+			c.emitConstI64(64, target.register)
+		}
+
+		// Emit the jmp instruction to jump to the position right after
+		// the non-zero case.
+		jmpAtEndOfZero := c.newProg()
+		jmpAtEndOfZero.As = obj.AJMP
+		jmpAtEndOfZero.To.Type = obj.TYPE_BRANCH
+		c.addInstruction(jmpAtEndOfZero)
+
+		// Otherwise, emit the TZCNT.
+		countZeros := c.newProg()
+		jmpIfNonZero.To.SetTarget(countZeros)
+		countZeros.From.Type = obj.TYPE_REG
+		countZeros.From.Reg = target.register
+		countZeros.To.Type = obj.TYPE_REG
+		countZeros.To.Reg = target.register
+		if o.Type == wazeroir.UnsignedInt32 {
+			countZeros.As = x86.ATZCNTL
+		} else {
+			countZeros.As = x86.ATZCNTQ
+		}
+		c.addInstruction(countZeros)
+
+		// Finally the end jump instruction of zero case must target towards
+		// the next instruction.
+		c.setJmpOrigin = jmpAtEndOfZero
+	}
+
+	// We reused the same register of target for the result.
+	result := c.locationStack.pushValueOnRegister(target.register)
+	result.setRegisterType(generalPurposeRegisterTypeInt)
+	return nil
+}
+
+// compilePopcnt emits instructions to count up the number of set bits in the
+// current top of the stack, and push the count result.
+// For example, stack of [..., 0b00_00_00_11] results in [..., 2].
+func (c *amd64Compiler) compilePopcnt(o *wazeroir.OperationPopcnt) error {
+	target := c.locationStack.pop()
+	if err := c.ensureOnGeneralPurposeRegister(target); err != nil {
+		return err
+	}
+
+	countBits := c.newProg()
+	countBits.From.Type = obj.TYPE_REG
+	countBits.From.Reg = target.register
+	countBits.To.Type = obj.TYPE_REG
+	countBits.To.Reg = target.register
+	if o.Type == wazeroir.UnsignedInt32 {
+		countBits.As = x86.APOPCNTL
+	} else {
+		countBits.As = x86.APOPCNTQ
+	}
+	c.addInstruction(countBits)
+
+	// We reused the same register of target for the result.
+	result := c.locationStack.pushValueOnRegister(target.register)
+	result.setRegisterType(generalPurposeRegisterTypeInt)
+	return nil
+}
+
 func (c *amd64Compiler) compileEq(o *wazeroir.OperationEq) error {
 	return c.emitEqOrNe(o.Type, true)
 }
@@ -1542,14 +1748,18 @@ func (c *amd64Compiler) compileConstI32(o *wazeroir.OperationConstI32) error {
 	loc := c.locationStack.pushValueOnRegister(reg)
 	loc.setRegisterType(generalPurposeRegisterTypeInt)
 
+	c.emitConstI32(o.Value, reg)
+	return nil
+}
+
+func (c *amd64Compiler) emitConstI32(val uint32, register int16) {
 	prog := c.newProg()
 	prog.As = x86.AMOVL // Note 32-bit move!
 	prog.From.Type = obj.TYPE_CONST
-	prog.From.Offset = int64(o.Value)
+	prog.From.Offset = int64(val)
 	prog.To.Type = obj.TYPE_REG
-	prog.To.Reg = reg
+	prog.To.Reg = register
 	c.addInstruction(prog)
-	return nil
 }
 
 func (c *amd64Compiler) compileConstI64(o *wazeroir.OperationConstI64) error {
@@ -1560,14 +1770,18 @@ func (c *amd64Compiler) compileConstI64(o *wazeroir.OperationConstI64) error {
 	loc := c.locationStack.pushValueOnRegister(reg)
 	loc.setRegisterType(generalPurposeRegisterTypeInt)
 
+	c.emitConstI64(o.Value, reg)
+	return nil
+}
+
+func (c *amd64Compiler) emitConstI64(val uint64, register int16) {
 	prog := c.newProg()
 	prog.As = x86.AMOVQ
 	prog.From.Type = obj.TYPE_CONST
-	prog.From.Offset = int64(o.Value)
+	prog.From.Offset = int64(val)
 	prog.To.Type = obj.TYPE_REG
-	prog.To.Reg = reg
+	prog.To.Reg = register
 	c.addInstruction(prog)
-	return nil
 }
 
 func (c *amd64Compiler) compileConstF32(o *wazeroir.OperationConstF32) error {


Test patch:
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index 8a67b920..6cf6d5b9 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -9,6 +9,7 @@ import (
 	"io"
 	"math"
 	"reflect"
+	"runtime"
 	"sync"
 	"testing"
 	"unsafe"
@@ -2384,6 +2385,7 @@ func TestAmd64Compiler_compileMul(t *testing.T) {
 				err := compiler.compileMul(&wazeroir.OperationMul{Type: wazeroir.UnsignedTypeI32})
 				require.NoError(t, err)
 				require.Equal(t, int16(x86.REG_AX), compiler.locationStack.peek().register)
+				require.Equal(t, generalPurposeRegisterTypeInt, compiler.locationStack.peek().regType)
 				require.Equal(t, uint64(2), compiler.locationStack.sp)
 				require.Len(t, compiler.locationStack.usedRegisters, 1)
 				// At this point, the previous value on the DX register is saved to the stack.
@@ -2498,6 +2500,7 @@ func TestAmd64Compiler_compileMul(t *testing.T) {
 				err := compiler.compileMul(&wazeroir.OperationMul{Type: wazeroir.UnsignedTypeI64})
 				require.NoError(t, err)
 				require.Equal(t, int16(x86.REG_AX), compiler.locationStack.peek().register)
+				require.Equal(t, generalPurposeRegisterTypeInt, compiler.locationStack.peek().regType)
 				require.Equal(t, uint64(2), compiler.locationStack.sp)
 				require.Len(t, compiler.locationStack.usedRegisters, 1)
 				// At this point, the previous value on the DX register is saved to the stack.
@@ -2636,6 +2639,315 @@ func TestAmd64Compiler_compileMul(t *testing.T) {
 	})
 }
 
+func TestAmd64Compiler_compilClz(t *testing.T) {
+	t.Run("32bit", func(t *testing.T) {
+		for _, tc := range []struct{ input, expectedLeadingZeros uint32 }{
+			{input: 0xff_ff_ff_ff, expectedLeadingZeros: 0},
+			{input: 0xf0_00_00_00, expectedLeadingZeros: 0},
+			{input: 0x00_ff_ff_ff, expectedLeadingZeros: 8},
+			{input: 0, expectedLeadingZeros: 32},
+		} {
+			tc := tc
+			t.Run(fmt.Sprintf("%032b", tc.input), func(t *testing.T) {
+				compiler := requireNewCompiler(t)
+				compiler.initializeReservedRegisters()
+				// Setup the target value.
+				err := compiler.compileConstI32(&wazeroir.OperationConstI32{Value: tc.input})
+				require.NoError(t, err)
+
+				// Emit the clz instruction.
+				err = compiler.compileClz(&wazeroir.OperationClz{Type: wazeroir.UnsignedInt32})
+				require.NoError(t, err)
+				// Verify that the result is pushed, meaning that
+				// stack pointer must not be changed.
+				require.Equal(t, uint64(1), compiler.locationStack.sp)
+				// Also the location must be register.
+				require.True(t, compiler.locationStack.peek().onRegister())
+				// On darwin, we have two branches and one must jump to the next
+				// instruction after compileClz.
+				require.True(t, runtime.GOOS != "darwin" || compiler.setJmpOrigin != nil)
+
+				// To verify the behavior, we release the value
+				// to the stack.
+				compiler.releaseAllRegistersToStack()
+				compiler.returnFunction()
+
+				// Generate and run the code under test.
+				code, _, err := compiler.generate()
+				require.NoError(t, err)
+				eng := newEngine()
+				mem := newMemoryInst()
+				jitcall(
+					uintptr(unsafe.Pointer(&code[0])),
+					uintptr(unsafe.Pointer(eng)),
+					uintptr(unsafe.Pointer(&mem.Buffer[0])),
+				)
+
+				// Check the stack.
+				require.Equal(t, uint64(1), eng.stackPointer)
+				require.Equal(t, tc.expectedLeadingZeros, uint32(eng.stack[eng.stackPointer-1]))
+			})
+		}
+	})
+	t.Run("64bit", func(t *testing.T) {
+		for _, tc := range []struct{ input, expectedLeadingZeros uint64 }{
+			{input: 0xf0_00_00_00_00_00_00_00, expectedLeadingZeros: 0},
+			{input: 0xff_ff_ff_ff_ff_ff_ff_ff, expectedLeadingZeros: 0},
+			{input: 0x00_ff_ff_ff_ff_ff_ff_ff, expectedLeadingZeros: 8},
+			{input: 0x00_00_00_00_ff_ff_ff_ff, expectedLeadingZeros: 32},
+			{input: 0x00_00_00_00_00_ff_ff_ff, expectedLeadingZeros: 40},
+			{input: 0, expectedLeadingZeros: 64},
+		} {
+			tc := tc
+			t.Run(fmt.Sprintf("%064b", tc.expectedLeadingZeros), func(t *testing.T) {
+				compiler := requireNewCompiler(t)
+				compiler.initializeReservedRegisters()
+				// Setup the target value.
+				err := compiler.compileConstI64(&wazeroir.OperationConstI64{Value: tc.input})
+				require.NoError(t, err)
+
+				// Emit the clz instruction.
+				err = compiler.compileClz(&wazeroir.OperationClz{Type: wazeroir.UnsignedInt64})
+				require.NoError(t, err)
+				// Verify that the result is pushed, meaning that
+				// stack pointer must not be changed.
+				require.Equal(t, uint64(1), compiler.locationStack.sp)
+				// Also the location must be register.
+				require.True(t, compiler.locationStack.peek().onRegister())
+				// On darwin, we have two branches and one must jump to the next
+				// instruction after compileClz.
+				require.True(t, runtime.GOOS != "darwin" || compiler.setJmpOrigin != nil)
+
+				// To verify the behavior, we release the value
+				// to the stack.
+				compiler.releaseAllRegistersToStack()
+				compiler.returnFunction()
+
+				// Generate and run the code under test.
+				code, _, err := compiler.generate()
+				require.NoError(t, err)
+				eng := newEngine()
+				mem := newMemoryInst()
+				jitcall(
+					uintptr(unsafe.Pointer(&code[0])),
+					uintptr(unsafe.Pointer(eng)),
+					uintptr(unsafe.Pointer(&mem.Buffer[0])),
+				)
+
+				// Check the stack.
+				require.Equal(t, uint64(1), eng.stackPointer)
+				require.Equal(t, tc.expectedLeadingZeros, eng.stack[eng.stackPointer-1])
+			})
+		}
+	})
+}
+
+func TestAmd64Compiler_compilCtz(t *testing.T) {
+	t.Run("32bit", func(t *testing.T) {
+		for _, tc := range []struct{ input, expectedTrailingZeros uint32 }{
+			{input: 0xff_ff_ff_ff, expectedTrailingZeros: 0},
+			{input: 0x00_00_00_01, expectedTrailingZeros: 0},
+			{input: 0xff_ff_ff_00, expectedTrailingZeros: 8},
+			{input: 0, expectedTrailingZeros: 32},
+		} {
+			tc := tc
+			t.Run(fmt.Sprintf("%032b", tc.input), func(t *testing.T) {
+				compiler := requireNewCompiler(t)
+				compiler.initializeReservedRegisters()
+				// Setup the target value.
+				err := compiler.compileConstI32(&wazeroir.OperationConstI32{Value: tc.input})
+				require.NoError(t, err)
+
+				// Emit the clz instruction.
+				err = compiler.compileCtz(&wazeroir.OperationCtz{Type: wazeroir.UnsignedInt32})
+				require.NoError(t, err)
+				// Verify that the result is pushed, meaning that
+				// stack pointer must not be changed.
+				require.Equal(t, uint64(1), compiler.locationStack.sp)
+				// Also the location must be register.
+				require.True(t, compiler.locationStack.peek().onRegister())
+				// On darwin, we have two branches and one must jump to the next
+				// instruction after compileCtz.
+				require.True(t, runtime.GOOS != "darwin" || compiler.setJmpOrigin != nil)
+
+				// To verify the behavior, we release the value
+				// to the stack.
+				compiler.releaseAllRegistersToStack()
+				compiler.returnFunction()
+
+				// Generate and run the code under test.
+				code, _, err := compiler.generate()
+				require.NoError(t, err)
+				eng := newEngine()
+				mem := newMemoryInst()
+				jitcall(
+					uintptr(unsafe.Pointer(&code[0])),
+					uintptr(unsafe.Pointer(eng)),
+					uintptr(unsafe.Pointer(&mem.Buffer[0])),
+				)
+
+				// Check the stack.
+				require.Equal(t, uint64(1), eng.stackPointer)
+				require.Equal(t, tc.expectedTrailingZeros, uint32(eng.stack[eng.stackPointer-1]))
+			})
+		}
+	})
+	t.Run("64bit", func(t *testing.T) {
+		for _, tc := range []struct{ input, expectedTrailingZeros uint64 }{
+			{input: 0xff_ff_ff_ff_ff_ff_ff_ff, expectedTrailingZeros: 0},
+			{input: 0x00_00_00_00_00_00_00_01, expectedTrailingZeros: 0},
+			{input: 0xff_ff_ff_ff_ff_ff_ff_00, expectedTrailingZeros: 8},
+			{input: 0xff_ff_ff_ff_00_00_00_00, expectedTrailingZeros: 32},
+			{input: 0xff_ff_ff_00_00_00_00_00, expectedTrailingZeros: 40},
+			{input: 0, expectedTrailingZeros: 64},
+		} {
+			tc := tc
+			t.Run(fmt.Sprintf("%064b", tc.input), func(t *testing.T) {
+				compiler := requireNewCompiler(t)
+				compiler.initializeReservedRegisters()
+				// Setup the target value.
+				err := compiler.compileConstI64(&wazeroir.OperationConstI64{Value: tc.input})
+				require.NoError(t, err)
+
+				// Emit the clz instruction.
+				err = compiler.compileCtz(&wazeroir.OperationCtz{Type: wazeroir.UnsignedInt64})
+				require.NoError(t, err)
+				// Verify that the result is pushed, meaning that
+				// stack pointer must not be changed.
+				require.Equal(t, uint64(1), compiler.locationStack.sp)
+				// Also the location must be register.
+				require.True(t, compiler.locationStack.peek().onRegister())
+				// On darwin, we have two branches and one must jump to the next
+				// instruction after compileCtz.
+				require.True(t, runtime.GOOS != "darwin" || compiler.setJmpOrigin != nil)
+
+				// To verify the behavior, we release the value
+				// to the stack.
+				compiler.releaseAllRegistersToStack()
+				compiler.returnFunction()
+
+				// Generate and run the code under test.
+				code, _, err := compiler.generate()
+				require.NoError(t, err)
+				eng := newEngine()
+				mem := newMemoryInst()
+				jitcall(
+					uintptr(unsafe.Pointer(&code[0])),
+					uintptr(unsafe.Pointer(eng)),
+					uintptr(unsafe.Pointer(&mem.Buffer[0])),
+				)
+
+				// Check the stack.
+				require.Equal(t, uint64(1), eng.stackPointer)
+				require.Equal(t, tc.expectedTrailingZeros, eng.stack[eng.stackPointer-1])
+			})
+		}
+	})
+}
+func TestAmd64Compiler_compilPopcnt(t *testing.T) {
+	t.Run("32bit", func(t *testing.T) {
+		for _, tc := range []struct{ input, expectedSetBits uint32 }{
+			{input: 0xff_ff_ff_ff, expectedSetBits: 32},
+			{input: 0x00_00_00_01, expectedSetBits: 1},
+			{input: 0x10_00_00_00, expectedSetBits: 1},
+			{input: 0x00_00_10_00, expectedSetBits: 1},
+			{input: 0x00_01_00_01, expectedSetBits: 2},
+			{input: 0xff_ff_00_ff, expectedSetBits: 24},
+			{input: 0, expectedSetBits: 0},
+		} {
+			tc := tc
+			t.Run(fmt.Sprintf("%032b", tc.input), func(t *testing.T) {
+				compiler := requireNewCompiler(t)
+				compiler.initializeReservedRegisters()
+				// Setup the target value.
+				err := compiler.compileConstI32(&wazeroir.OperationConstI32{Value: tc.input})
+				require.NoError(t, err)
+
+				// Emit the clz instruction.
+				err = compiler.compilePopcnt(&wazeroir.OperationPopcnt{Type: wazeroir.UnsignedInt32})
+				require.NoError(t, err)
+				// Verify that the result is pushed, meaning that
+				// stack pointer must not be changed.
+				require.Equal(t, uint64(1), compiler.locationStack.sp)
+				// Also the location must be register.
+				require.True(t, compiler.locationStack.peek().onRegister())
+
+				// To verify the behavior, we release the value
+				// to the stack.
+				compiler.releaseAllRegistersToStack()
+				compiler.returnFunction()
+
+				// Generate and run the code under test.
+				code, _, err := compiler.generate()
+				require.NoError(t, err)
+				eng := newEngine()
+				mem := newMemoryInst()
+				jitcall(
+					uintptr(unsafe.Pointer(&code[0])),
+					uintptr(unsafe.Pointer(eng)),
+					uintptr(unsafe.Pointer(&mem.Buffer[0])),
+				)
+
+				// Check the stack.
+				require.Equal(t, uint64(1), eng.stackPointer)
+				require.Equal(t, tc.expectedSetBits, uint32(eng.stack[eng.stackPointer-1]))
+			})
+		}
+	})
+	t.Run("64bit", func(t *testing.T) {
+		for _, tc := range []struct{ in, exp uint64 }{
+			{in: 0xff_ff_ff_ff_ff_ff_ff_ff, exp: 64},
+			{in: 0x00_00_00_00_00_00_00_01, exp: 1},
+			{in: 0x00_00_00_01_00_00_00_00, exp: 1},
+			{in: 0x10_00_00_00_00_00_00_00, exp: 1},
+			{in: 0xf0_00_00_00_00_00_01_00, exp: 5},
+			{in: 0xff_ff_ff_ff_ff_ff_ff_00, exp: 56},
+			{in: 0xff_ff_ff_00_ff_ff_ff_ff, exp: 56},
+			{in: 0xff_ff_ff_ff_00_00_00_00, exp: 32},
+			{in: 0xff_ff_ff_00_00_00_00_00, exp: 24},
+			{in: 0, exp: 0},
+		} {
+			tc := tc
+			t.Run(fmt.Sprintf("%064b", tc.in), func(t *testing.T) {
+				compiler := requireNewCompiler(t)
+				compiler.initializeReservedRegisters()
+				// Setup the target value.
+				err := compiler.compileConstI64(&wazeroir.OperationConstI64{Value: tc.in})
+				require.NoError(t, err)
+
+				// Emit the clz instruction.
+				err = compiler.compilePopcnt(&wazeroir.OperationPopcnt{Type: wazeroir.UnsignedInt64})
+				require.NoError(t, err)
+				// Verify that the result is pushed, meaning that
+				// stack pointer must not be changed.
+				require.Equal(t, uint64(1), compiler.locationStack.sp)
+				// Also the location must be register.
+				require.True(t, compiler.locationStack.peek().onRegister())
+
+				// To verify the behavior, we release the value
+				// to the stack.
+				compiler.releaseAllRegistersToStack()
+				compiler.returnFunction()
+
+				// Generate and run the code under test.
+				code, _, err := compiler.generate()
+				require.NoError(t, err)
+				eng := newEngine()
+				mem := newMemoryInst()
+				jitcall(
+					uintptr(unsafe.Pointer(&code[0])),
+					uintptr(unsafe.Pointer(eng)),
+					uintptr(unsafe.Pointer(&mem.Buffer[0])),
+				)
+
+				// Check the stack.
+				require.Equal(t, uint64(1), eng.stackPointer)
+				require.Equal(t, tc.exp, eng.stack[eng.stackPointer-1])
+			})
+		}
+	})
+}
+
 func TestAmd64Compiler_compileCall(t *testing.T) {
 	t.Run("host function", func(t *testing.T) {
 		const functionIndex = 5