Repository: tetratelabs/wazero
Pull request #1635 (2023-08-17 14:48:40)

wazevo: adds bit count instructions clz, ctz
Continues #1496.

This introduces the front-end and back-end for clz, ctz (in terms of rbit+clz).

It originally also introduced popcnt for the front-end, but this requires 
VCNT in the backend, which is probably worth its own separate PR.

Contextually, this also introduces a partial implementation for bitRR.

First PR to wazevo, so I am not entirely sure that the generated
code for the SSA and the backend is correct, let me know if there are errors,
I am really figuring out a lot of stuff, including the framework :)

Signed-off-by: Edoardo Vacchi <evacchi@users.noreply.github.com>


evacchi: Although I am not sure about its correctness, this PR may be a good starting point if you haven't dug into wazevo yet since it is really tiny @ncruces 

Patch:
diff --git a/internal/emscripten/emscripten.go b/internal/emscripten/emscripten.go
index 3262e670..e6aa79f0 100644
--- a/internal/emscripten/emscripten.go
+++ b/internal/emscripten/emscripten.go
@@ -106,10 +106,7 @@ func (v *InvokeFunc) Call(ctx context.Context, mod api.Module, stack []uint64) {
 
 	// Lookup the table index we will call.
 	t := m.Tables[0] // Note: Emscripten doesn't use multiple tables
-	f, err := m.Engine.LookupFunction(t, typeID, tableOffset)
-	if err != nil {
-		panic(err)
-	}
+	f := m.LookupFunction(t, typeID, tableOffset)
 
 	// The Go implementation below mimics the Emscripten JS behaviour to support
 	// longjmps from indirect function calls. The implementation of these
diff --git a/internal/engine/compiler/engine.go b/internal/engine/compiler/engine.go
index 77df6994..f8b76fa4 100644
--- a/internal/engine/compiler/engine.go
+++ b/internal/engine/compiler/engine.go
@@ -681,24 +681,20 @@ func (e *moduleEngine) newFunction(f *function) api.Function {
 }
 
 // LookupFunction implements the same method as documented on wasm.ModuleEngine.
-func (e *moduleEngine) LookupFunction(t *wasm.TableInstance, typeId wasm.FunctionTypeID, tableOffset wasm.Index) (f api.Function, err error) {
+func (e *moduleEngine) LookupFunction(t *wasm.TableInstance, typeId wasm.FunctionTypeID, tableOffset wasm.Index) (*wasm.ModuleInstance, wasm.Index) {
 	if tableOffset >= uint32(len(t.References)) || t.Type != wasm.RefTypeFuncref {
-		err = wasmruntime.ErrRuntimeInvalidTableAccess
-		return
+		panic(wasmruntime.ErrRuntimeInvalidTableAccess)
 	}
 	rawPtr := t.References[tableOffset]
 	if rawPtr == 0 {
-		err = wasmruntime.ErrRuntimeInvalidTableAccess
-		return
+		panic(wasmruntime.ErrRuntimeInvalidTableAccess)
 	}
 
 	tf := functionFromUintptr(rawPtr)
 	if tf.typeID != typeId {
-		err = wasmruntime.ErrRuntimeIndirectCallTypeMismatch
-		return
+		panic(wasmruntime.ErrRuntimeIndirectCallTypeMismatch)
 	}
-	f = e.newFunction(tf)
-	return
+	return tf.moduleInstance, tf.parent.index
 }
 
 // functionFromUintptr resurrects the original *function from the given uintptr
diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go
index f4a79b71..8f0db3c7 100644
--- a/internal/engine/interpreter/interpreter.go
+++ b/internal/engine/interpreter/interpreter.go
@@ -462,25 +462,20 @@ func (e *moduleEngine) NewFunction(index wasm.Index) (ce api.Function) {
 }
 
 // LookupFunction implements the same method as documented on wasm.ModuleEngine.
-func (e *moduleEngine) LookupFunction(t *wasm.TableInstance, typeId wasm.FunctionTypeID, tableOffset wasm.Index) (f api.Function, err error) {
+func (e *moduleEngine) LookupFunction(t *wasm.TableInstance, typeId wasm.FunctionTypeID, tableOffset wasm.Index) (*wasm.ModuleInstance, wasm.Index) {
 	if tableOffset >= uint32(len(t.References)) {
-		err = wasmruntime.ErrRuntimeInvalidTableAccess
-		return
+		panic(wasmruntime.ErrRuntimeInvalidTableAccess)
 	}
 	rawPtr := t.References[tableOffset]
 	if rawPtr == 0 {
-		err = wasmruntime.ErrRuntimeInvalidTableAccess
-		return
+		panic(wasmruntime.ErrRuntimeInvalidTableAccess)
 	}
 
 	tf := functionFromUintptr(rawPtr)
 	if tf.typeID != typeId {
-		err = wasmruntime.ErrRuntimeIndirectCallTypeMismatch
-		return
+		panic(wasmruntime.ErrRuntimeIndirectCallTypeMismatch)
 	}
-
-	f = e.newCallEngine(tf)
-	return
+	return tf.moduleInstance, tf.parent.index
 }
 
 // Definition implements the same method as documented on api.Function.
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr.go b/internal/engine/wazevo/backend/isa/arm64/instr.go
index cf601b47..4137313d 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr.go
@@ -50,6 +50,7 @@ var defKinds = [numInstructionKinds]defKind{
 	aluRRRShift:     defKindRD,
 	aluRRImmShift:   defKindRD,
 	aluRRRExtend:    defKindRD,
+	bitRR:           defKindRD,
 	movZ:            defKindRD,
 	movN:            defKindRD,
 	mov32:           defKindRD,
@@ -143,6 +144,7 @@ var useKinds = [numInstructionKinds]useKind{
 	aluRRRShift:     useKindRNRM,
 	aluRRImmShift:   useKindRN,
 	aluRRRExtend:    useKindRNRM,
+	bitRR:           useKindRN,
 	movZ:            useKindNone,
 	movN:            useKindNone,
 	mov32:           useKindRN,
@@ -612,6 +614,15 @@ func (i *instruction) asALUBitmaskImm(aluOp aluOp, rn, rd regalloc.VReg, imm uin
 	}
 }
 
+func (i *instruction) asBitRR(bitOp bitOp, rd, rn regalloc.VReg, is64bit bool) {
+	i.kind = bitRR
+	i.rn, i.rd = operandNR(rn), operandNR(rd)
+	i.u1 = uint64(bitOp)
+	if is64bit {
+		i.u2 = 1
+	}
+}
+
 func (i *instruction) asFpuRRR(op fpuBinOp, rd, rn, rm operand, dst64bit bool) {
 	i.kind = fpuRRR
 	i.u1 = uint64(op)
@@ -722,7 +733,12 @@ func (i *instruction) String() (str string) {
 			e,
 		)
 	case bitRR:
-		panic("TODO")
+		size := is64SizeBitToSize(i.u2)
+		str = fmt.Sprintf("%s %s, %s",
+			bitOp(i.u1),
+			formatVRegSized(i.rd.nr(), size),
+			formatVRegSized(i.rn.nr(), size),
+		)
 	case uLoad8:
 		str = fmt.Sprintf("ldrb %s, %s", formatVRegSized(i.rd.nr(), 32), i.amode.format(32))
 	case sLoad8:
@@ -1235,6 +1251,28 @@ const (
 	aluOpMSub
 )
 
+// bitOp determines the type of bitwise operation. Instructions whose kind is one of
+// bitOpRbit and bitOpClz would use this type.
+type bitOp int
+
+// String implements fmt.Stringer.
+func (b bitOp) String() string {
+	switch b {
+	case bitOpRbit:
+		return "rbit"
+	case bitOpClz:
+		return "clz"
+	}
+	panic(int(b))
+}
+
+const (
+	// 32/64-bit Rbit.
+	bitOpRbit bitOp = iota
+	// 32/64-bit Clz.
+	bitOpClz
+)
+
 // fpuBinOp represents a binary floating-point unit (FPU) operation.
 type fpuBinOp byte
 
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go b/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
index 52baefed..13a9cfe2 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr_encoding.go
@@ -157,6 +157,13 @@ func (i *instruction) encode(c backend.Compiler) {
 			i.u2,
 			i.u3 == 1,
 		))
+	case bitRR:
+		c.Emit4Bytes(encodeBitRR(
+			bitOp(i.u1),
+			regNumberInEncoding[i.rd.realReg()],
+			regNumberInEncoding[i.rn.realReg()],
+			uint32(i.u2)),
+		)
 	case aluRRImm12:
 		imm12, shift := i.rm.imm12()
 		c.Emit4Bytes(encodeAluRRImm12(
@@ -309,6 +316,21 @@ func encodeAluRRRR(op aluOp, rd, rn, rm, ra, _64bit uint32) uint32 {
 	return _64bit<<31 | 0b11011<<24 | op31<<21 | rm<<16 | oO<<15 | ra<<10 | rn<<5 | rd
 }
 
+// encodeBitRR encodes as Data-processing (1 source) in
+// https://developer.arm.com/documentation/ddi0596/2020-12/Index-by-Encoding/Data-Processing----Register?lang=en
+func encodeBitRR(op bitOp, rd, rn, _64bit uint32) uint32 {
+	var opcode2, opcode uint32
+	switch op {
+	case bitOpRbit:
+		opcode2, opcode = 0b00000, 0b000000
+	case bitOpClz:
+		opcode2, opcode = 0b00000, 0b000100
+	default:
+		panic("TODO/BUG")
+	}
+	return _64bit<<31 | 0b1_0_11010110<<21 | opcode2<<15 | opcode<<10 | rn<<5 | rd
+}
+
 func encodeAsMov32(rn, rd uint32) uint32 {
 	// This is an alias of ORR (shifted register):
 	// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/MOV--register---Move--register---an-alias-of-ORR--shifted-register--
diff --git a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
index f9ccd7ca..789b37bb 100644
--- a/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
+++ b/internal/engine/wazevo/backend/isa/arm64/lower_instr.go
@@ -147,6 +147,14 @@ func (m *machine) LowerInstr(instr *ssa.Instruction) {
 	case ssa.OpcodeSelect:
 		c, x, y := instr.SelectData()
 		m.lowerSelect(c, x, y, instr.Return())
+	case ssa.OpcodeClz:
+		x := instr.UnaryData()
+		result := instr.Return()
+		m.lowerClz(x, result)
+	case ssa.OpcodeCtz:
+		x := instr.UnaryData()
+		result := instr.Return()
+		m.lowerCtz(x, result)
 	default:
 		panic("TODO: lowering " + instr.Opcode().String())
 	}
@@ -295,6 +303,26 @@ func (m *machine) lowerImul(x, y, result ssa.Value) {
 	m.insert(mul)
 }
 
+func (m *machine) lowerClz(x, result ssa.Value) {
+	rd := m.compiler.VRegOf(result)
+	rn := m.getOperand_NR(m.compiler.ValueDefinition(x), extModeNone)
+	clz := m.allocateInstr()
+	clz.asBitRR(bitOpClz, rd, rn.nr(), x.Type().Bits() == 64)
+	m.insert(clz)
+}
+
+func (m *machine) lowerCtz(x, result ssa.Value) {
+	rd := m.compiler.VRegOf(result)
+	rn := m.getOperand_NR(m.compiler.ValueDefinition(x), extModeNone)
+	rbit := m.allocateInstr()
+	rbit.asBitRR(bitOpRbit, rd, rn.nr(), x.Type().Bits() == 64)
+	m.insert(rbit)
+
+	clz := m.allocateInstr()
+	clz.asBitRR(bitOpClz, rd, rd, x.Type().Bits() == 64)
+	m.insert(clz)
+}
+
 const exitWithCodeEncodingSize = exitSequenceSize + 8
 
 // lowerExitWithCode lowers the lowerExitWithCode takes a context pointer as argument.
diff --git a/internal/engine/wazevo/frontend/lower.go b/internal/engine/wazevo/frontend/lower.go
index 9445b65b..08d60abe 100644
--- a/internal/engine/wazevo/frontend/lower.go
+++ b/internal/engine/wazevo/frontend/lower.go
@@ -427,6 +427,26 @@ func (c *Compiler) lowerOpcode(op wasm.Opcode) {
 		builder.InsertInstruction(ishl)
 		value := ishl.Return()
 		state.push(value)
+	case wasm.OpcodeI32Clz, wasm.OpcodeI64Clz:
+		if state.unreachable {
+			return
+		}
+		x := state.pop()
+		clz := builder.AllocateInstruction()
+		clz.AsClz(x)
+		builder.InsertInstruction(clz)
+		value := clz.Return()
+		state.push(value)
+	case wasm.OpcodeI32Ctz, wasm.OpcodeI64Ctz:
+		if state.unreachable {
+			return
+		}
+		x := state.pop()
+		ctz := builder.AllocateInstruction()
+		ctz.AsCtz(x)
+		builder.InsertInstruction(ctz)
+		value := ctz.Return()
+		state.push(value)
 	case wasm.OpcodeGlobalGet:
 		index := c.readI32u()
 		if state.unreachable {
diff --git a/internal/engine/wazevo/module_engine.go b/internal/engine/wazevo/module_engine.go
index 0bf801f1..6eaa560d 100644
--- a/internal/engine/wazevo/module_engine.go
+++ b/internal/engine/wazevo/module_engine.go
@@ -206,6 +206,6 @@ func (m *moduleEngine) FunctionInstanceReference(funcIndex wasm.Index) wasm.Refe
 }
 
 // LookupFunction implements wasm.ModuleEngine.
-func (m *moduleEngine) LookupFunction(t *wasm.TableInstance, typeId wasm.FunctionTypeID, tableOffset wasm.Index) (api.Function, error) {
+func (m *moduleEngine) LookupFunction(t *wasm.TableInstance, typeId wasm.FunctionTypeID, tableOffset wasm.Index) (*wasm.ModuleInstance, wasm.Index) {
 	panic("TODO")
 }
diff --git a/internal/engine/wazevo/ssa/instructions.go b/internal/engine/wazevo/ssa/instructions.go
index b7dfb217..130ad435 100644
--- a/internal/engine/wazevo/ssa/instructions.go
+++ b/internal/engine/wazevo/ssa/instructions.go
@@ -836,6 +836,8 @@ var instructionSideEffects = [opcodeEnd]sideEffect{
 	OpcodeIcmp:               sideEffectFalse,
 	OpcodeFcmp:               sideEffectFalse,
 	OpcodeFadd:               sideEffectFalse,
+	OpcodeClz:                sideEffectFalse,
+	OpcodeCtz:                sideEffectFalse,
 	OpcodeLoad:               sideEffectFalse,
 	OpcodeUload8:             sideEffectFalse,
 	OpcodeUload16:            sideEffectFalse,
@@ -933,6 +935,8 @@ var instructionReturnTypes = [opcodeEnd]returnTypesFn{
 	OpcodeFmin:               returnTypesFnSingle,
 	OpcodeF32const:           returnTypesFnF32,
 	OpcodeF64const:           returnTypesFnF64,
+	OpcodeClz:                returnTypesFnSingle,
+	OpcodeCtz:                returnTypesFnSingle,
 	OpcodeStore:              returnTypesFnNoReturns,
 	OpcodeIstore8:            returnTypesFnNoReturns,
 	OpcodeIstore16:           returnTypesFnNoReturns,
@@ -1302,6 +1306,25 @@ func (i *Instruction) CallIndirectData() (funcPtr Value, sigID SignatureID, args
 	return
 }
 
+// AsClz initializes this instruction as a Count Leading Zeroes instruction with OpcodeClz.
+func (i *Instruction) AsClz(x Value) {
+	i.opcode = OpcodeClz
+	i.v = x
+	i.typ = x.Type()
+}
+
+// AsCtz initializes this instruction as a Count Trailing Zeroes instruction with OpcodeCtz.
+func (i *Instruction) AsCtz(x Value) {
+	i.opcode = OpcodeCtz
+	i.v = x
+	i.typ = x.Type()
+}
+
+// UnaryData return the operand for a unary instruction.
+func (i *Instruction) UnaryData() Value {
+	return i.v
+}
+
 // AsSExtend initializes this instruction as a sign extension instruction with OpcodeSExtend.
 func (i *Instruction) AsSExtend(v Value, from, to byte) {
 	i.opcode = OpcodeSExtend
@@ -1438,6 +1461,8 @@ func (i *Instruction) Format(b Builder) string {
 	case OpcodeIshl, OpcodeSshr, OpcodeUshr:
 		instSuffix = fmt.Sprintf(" %s, %s", i.v.Format(b), i.v2.Format(b))
 	case OpcodeUndefined:
+	case OpcodeClz, OpcodeCtz:
+		instSuffix = " " + i.v.Format(b)
 	default:
 		panic(fmt.Sprintf("TODO: format for %s", i.opcode))
 	}
diff --git a/internal/wasm/engine.go b/internal/wasm/engine.go
index 26d196b0..3d483807 100644
--- a/internal/wasm/engine.go
+++ b/internal/wasm/engine.go
@@ -51,8 +51,8 @@ type ModuleEngine interface {
 	// ResolveImportedMemory is called when this module imports a memory from another module.
 	ResolveImportedMemory(importedModuleEngine ModuleEngine)
 
-	// LookupFunction returns the api.Function created from the function in the function table at the given offset.
-	LookupFunction(t *TableInstance, typeId FunctionTypeID, tableOffset Index) (api.Function, error)
+	// LookupFunction returns the FunctionModule and the Index of the function in the returned ModuleInstance at the given offset in the table.
+	LookupFunction(t *TableInstance, typeId FunctionTypeID, tableOffset Index) (*ModuleInstance, Index)
 
 	// FunctionInstanceReference returns Reference for the given Index for a FunctionInstance. The returned values are used by
 	// the initialization via ElementSegment.
diff --git a/internal/wasm/module_instance_lookup.go b/internal/wasm/module_instance_lookup.go
new file mode 100644
index 00000000..442d26a2
--- /dev/null
+++ b/internal/wasm/module_instance_lookup.go
@@ -0,0 +1,73 @@
+package wasm
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/internal/internalapi"
+)
+
+// LookupFunction looks up the table by the given index, and returns the api.Function implementation if found,
+// otherwise this panics according to the same semantics as call_indirect instruction.
+// Currently, this is only used by emscripten which needs to do call_indirect-like operation in the host function.
+func (m *ModuleInstance) LookupFunction(t *TableInstance, typeId FunctionTypeID, tableOffset Index) api.Function {
+	fm, index := m.Engine.LookupFunction(t, typeId, tableOffset)
+	if source := fm.Source; source.IsHostModule {
+		// This case, the found function is a host function stored in the table. Generally, Engine.NewFunction are only
+		// responsible for calling Wasm-defined functions (not designed for calling Go functions!). Hence we need to wrap
+		// the host function as a special case.
+		def := &source.FunctionDefinitionSection[index]
+		goF := source.CodeSection[index].GoFunc
+		switch typed := goF.(type) {
+		case api.GoFunction:
+			// GoFunction doesn't need looked up module.
+			return &lookedUpGoFunction{def: def, g: goFunctionAsGoModuleFunction(typed)}
+		case api.GoModuleFunction:
+			return &lookedUpGoFunction{def: def, lookedUpModule: m, g: typed}
+		default:
+			panic(fmt.Sprintf("unexpected GoFunc type: %T", goF))
+		}
+	} else {
+		return fm.Engine.NewFunction(index)
+	}
+}
+
+// lookedUpGoFunction implements lookedUpGoModuleFunction.
+type lookedUpGoFunction struct {
+	internalapi.WazeroOnly
+	def *FunctionDefinition
+	// lookedUpModule is the *ModuleInstance from which this Go function is looked up, i.e. owner of the table.
+	lookedUpModule *ModuleInstance
+	g              api.GoModuleFunction
+}
+
+// goFunctionAsGoModuleFunction converts api.GoFunction to api.GoModuleFunction which ignores the api.Module argument.
+func goFunctionAsGoModuleFunction(g api.GoFunction) api.GoModuleFunction {
+	return api.GoModuleFunc(func(ctx context.Context, _ api.Module, stack []uint64) {
+		g.Call(ctx, stack)
+	})
+}
+
+// Definition implements api.Function.
+func (l *lookedUpGoFunction) Definition() api.FunctionDefinition { return l.def }
+
+// Call implements api.Function.
+func (l *lookedUpGoFunction) Call(ctx context.Context, params ...uint64) ([]uint64, error) {
+	typ := l.def.Functype
+	stackSize := typ.ParamNumInUint64
+	rn := typ.ResultNumInUint64
+	if rn > stackSize {
+		stackSize = rn
+	}
+	stack := make([]uint64, stackSize)
+	copy(stack, params)
+	return stack[:rn], l.CallWithStack(ctx, stack)
+}
+
+// CallWithStack implements api.Function.
+func (l *lookedUpGoFunction) CallWithStack(ctx context.Context, stack []uint64) error {
+	// The Go host function always needs to access caller's module, in this case the one holding the table.
+	l.g.Call(ctx, l.lookedUpModule, stack)
+	return nil
+}


Test patch:
diff --git a/internal/engine/wazevo/backend/backend_test.go b/internal/engine/wazevo/backend/backend_test.go
index e36ee2a1..2df19b39 100644
--- a/internal/engine/wazevo/backend/backend_test.go
+++ b/internal/engine/wazevo/backend/backend_test.go
@@ -1415,6 +1415,37 @@ L1 (SSA Block: blk0):
 	sxth w6, w8
 	ldr x30, [sp], #0x10
 	ret
+`,
+		},
+		{
+			name: "integer bit counts", m: testcases.IntegerBitCounts.Module,
+			afterLoweringARM64: `
+L1 (SSA Block: blk0):
+	mov x2?, x2
+	mov x3?, x3
+	clz w4?, w2?
+	rbit w5?, w2?
+	clz w5?, w5?
+	clz x6?, x3?
+	rbit x7?, x3?
+	clz x7?, x7?
+	mov x3, x7?
+	mov x2, x6?
+	mov x1, x5?
+	mov x0, x4?
+	ret
+`,
+			afterFinalizeARM64: `
+L1 (SSA Block: blk0):
+	str x30, [sp, #-0x10]!
+	clz w0, w2
+	rbit w1, w2
+	clz w1, w1
+	clz x2, x3
+	rbit x3, x3
+	clz x3, x3
+	ldr x30, [sp], #0x10
+	ret
 `,
 		},
 		{
diff --git a/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go b/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go
index 98774814..e57b7568 100644
--- a/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go
+++ b/internal/engine/wazevo/backend/isa/arm64/instr_encoding_test.go
@@ -434,6 +434,10 @@ func TestInstruction_encode(t *testing.T) {
 		{want: "400041d3", setup: func(i *instruction) {
 			i.asALUShift(aluOpLsl, operandNR(x0VReg), operandNR(x2VReg), operandShiftImm(63), true)
 		}},
+		{want: "4000c05a", setup: func(i *instruction) { i.asBitRR(bitOpRbit, x0VReg, x2VReg, false) }},
+		{want: "4000c0da", setup: func(i *instruction) { i.asBitRR(bitOpRbit, x0VReg, x2VReg, true) }},
+		{want: "4010c05a", setup: func(i *instruction) { i.asBitRR(bitOpClz, x0VReg, x2VReg, false) }},
+		{want: "4010c0da", setup: func(i *instruction) { i.asBitRR(bitOpClz, x0VReg, x2VReg, true) }},
 	} {
 		tc := tc
 		t.Run(tc.want, func(t *testing.T) {
diff --git a/internal/engine/wazevo/frontend/frontend_test.go b/internal/engine/wazevo/frontend/frontend_test.go
index 6b6014a0..9bf4aed1 100644
--- a/internal/engine/wazevo/frontend/frontend_test.go
+++ b/internal/engine/wazevo/frontend/frontend_test.go
@@ -609,6 +609,17 @@ blk0: (exec_ctx:i64, module_ctx:i64, v2:i32, v3:i64)
 	v9:i32 = SExtend v2, 8->32
 	v10:i32 = SExtend v2, 16->32
 	Jump blk_ret, v4, v5, v6, v7, v8, v9, v10
+`,
+		},
+		{
+			name: "integer bit counts", m: testcases.IntegerBitCounts.Module,
+			exp: `
+blk0: (exec_ctx:i64, module_ctx:i64, v2:i32, v3:i64)
+	v4:i32 = Clz v2
+	v5:i32 = Ctz v2
+	v6:i64 = Clz v3
+	v7:i64 = Ctz v3
+	Jump blk_ret, v4, v5, v6, v7
 `,
 		},
 		{
diff --git a/internal/engine/wazevo/testcases/testcases.go b/internal/engine/wazevo/testcases/testcases.go
index 28666ca0..6487aeb0 100644
--- a/internal/engine/wazevo/testcases/testcases.go
+++ b/internal/engine/wazevo/testcases/testcases.go
@@ -731,6 +731,27 @@ var (
 			wasm.OpcodeEnd,
 		}, []wasm.ValueType{}),
 	}
+	IntegerBitCounts = TestCase{
+		Name: "integer_bit_counts",
+		Module: SingleFunctionModule(wasm.FunctionType{
+			Params:  []wasm.ValueType{i32, i64},
+			Results: []wasm.ValueType{i32, i32, i64, i64},
+		}, []byte{
+			wasm.OpcodeLocalGet, 0,
+			wasm.OpcodeI32Clz,
+
+			wasm.OpcodeLocalGet, 0,
+			wasm.OpcodeI32Ctz,
+
+			wasm.OpcodeLocalGet, 1,
+			wasm.OpcodeI64Clz,
+
+			wasm.OpcodeLocalGet, 1,
+			wasm.OpcodeI64Ctz,
+
+			wasm.OpcodeEnd,
+		}, []wasm.ValueType{}),
+	}
 	FloatComparisons = TestCase{
 		Name: "float_comparisons",
 		Module: SingleFunctionModule(wasm.FunctionType{
diff --git a/internal/testing/enginetest/enginetest.go b/internal/testing/enginetest/enginetest.go
index b9dbbb77..edd80527 100644
--- a/internal/testing/enginetest/enginetest.go
+++ b/internal/testing/enginetest/enginetest.go
@@ -252,6 +252,7 @@ func RunTestModuleEngineLookupFunction(t *testing.T, et EngineTester) {
 	require.NoError(t, err)
 	m := &wasm.ModuleInstance{
 		TypeIDs: []wasm.FunctionTypeID{0, 1},
+		Source:  mod,
 	}
 	m.Tables = []*wasm.TableInstance{
 		{Min: 2, References: make([]wasm.Reference, 2), Type: wasm.RefTypeFuncref},
@@ -264,9 +265,13 @@ func RunTestModuleEngineLookupFunction(t *testing.T, et EngineTester) {
 	linkModuleToEngine(m, me)
 
 	t.Run("null reference", func(t *testing.T) {
-		_, err := me.LookupFunction(m.Tables[0], m.TypeIDs[0], 0) // offset 0 is not initialized yet.
+		err = require.CapturePanic(func() {
+			m.LookupFunction(m.Tables[0], m.TypeIDs[0], 0) // offset 0 is not initialized yet.
+		})
 		require.Equal(t, wasmruntime.ErrRuntimeInvalidTableAccess, err)
-		_, err = me.LookupFunction(m.Tables[0], m.TypeIDs[0], 1) // offset 1 is not initialized yet.
+		err = require.CapturePanic(func() {
+			m.LookupFunction(m.Tables[0], m.TypeIDs[0], 1) // offset 1 is not initialized yet.
+		})
 		require.Equal(t, wasmruntime.ErrRuntimeInvalidTableAccess, err)
 	})
 
@@ -274,39 +279,41 @@ func RunTestModuleEngineLookupFunction(t *testing.T, et EngineTester) {
 	m.Tables[0].References[1] = me.FunctionInstanceReference(0)
 
 	t.Run("initialized", func(t *testing.T) {
-		f1, err := me.LookupFunction(m.Tables[0], m.TypeIDs[0], 0) // offset 0 is now initialized.
-		require.NoError(t, err)
+		f1 := m.LookupFunction(m.Tables[0], m.TypeIDs[0], 0) // offset 0 is now initialized.
 		require.Equal(t, wasm.Index(2), f1.Definition().Index())
-		f2, err := me.LookupFunction(m.Tables[0], m.TypeIDs[0], 1) // offset 1 is now initialized.
-		require.NoError(t, err)
+		f2 := m.LookupFunction(m.Tables[0], m.TypeIDs[0], 1) // offset 1 is now initialized.
 		require.Equal(t, wasm.Index(0), f2.Definition().Index())
 	})
 
 	t.Run("out of range", func(t *testing.T) {
-		_, err := me.LookupFunction(m.Tables[0], m.TypeIDs[0], 100 /* out of range */)
+		err = require.CapturePanic(func() {
+			me.LookupFunction(m.Tables[0], m.TypeIDs[0], 100 /* out of range */)
+		})
 		require.Equal(t, wasmruntime.ErrRuntimeInvalidTableAccess, err)
 	})
 
 	t.Run("access to externref table", func(t *testing.T) {
-		_, err := me.LookupFunction(m.Tables[1], /* table[1] has externref type. */
-			m.TypeIDs[0], 0)
+		err := require.CapturePanic(func() {
+			m.LookupFunction(m.Tables[1], /* table[1] has externref type. */
+				m.TypeIDs[0], 0)
+		})
 		require.Equal(t, wasmruntime.ErrRuntimeInvalidTableAccess, err)
 	})
 
 	t.Run("access to externref table", func(t *testing.T) {
-		_, err := me.LookupFunction(m.Tables[0], /* type mismatch */
-			m.TypeIDs[1], 0)
+		err = require.CapturePanic(func() {
+			me.LookupFunction(m.Tables[0], /* type mismatch */
+				m.TypeIDs[1], 0)
+		})
 		require.Equal(t, wasmruntime.ErrRuntimeIndirectCallTypeMismatch, err)
 	})
 
 	m.Tables[2].References[0] = me.FunctionInstanceReference(1)
 	m.Tables[2].References[5] = me.FunctionInstanceReference(2)
 	t.Run("initialized - tables[2]", func(t *testing.T) {
-		f1, err := me.LookupFunction(m.Tables[2], m.TypeIDs[0], 0)
-		require.NoError(t, err)
+		f1 := m.LookupFunction(m.Tables[2], m.TypeIDs[0], 0)
 		require.Equal(t, wasm.Index(1), f1.Definition().Index())
-		f2, err := me.LookupFunction(m.Tables[2], m.TypeIDs[0], 5)
-		require.NoError(t, err)
+		f2 := m.LookupFunction(m.Tables[2], m.TypeIDs[0], 5)
 		require.Equal(t, wasm.Index(2), f2.Definition().Index())
 	})
 }
diff --git a/internal/wasm/module_instance_lookup_test.go b/internal/wasm/module_instance_lookup_test.go
new file mode 100644
index 00000000..3e8c839e
--- /dev/null
+++ b/internal/wasm/module_instance_lookup_test.go
@@ -0,0 +1,96 @@
+package wasm
+
+import (
+	"context"
+	"math"
+	"testing"
+
+	"github.com/tetratelabs/wazero/api"
+	"github.com/tetratelabs/wazero/internal/testing/require"
+)
+
+func TestModuleInstance_LookupFunction(t *testing.T) {
+	var called int
+	hostModule := &Module{
+		IsHostModule: true,
+		CodeSection: []Code{
+			{GoFunc: api.GoFunc(func(context.Context, []uint64) {
+				called++
+			})},
+			{GoFunc: api.GoModuleFunc(func(context.Context, api.Module, []uint64) {
+				called++
+			})},
+		},
+		FunctionDefinitionSection: []FunctionDefinition{{}, {}},
+	}
+
+	me := &mockModuleEngine{
+		lookupEntries: map[Index]mockModuleEngineLookupEntry{
+			0: {m: &ModuleInstance{Source: hostModule}, index: 0},
+			1: {m: &ModuleInstance{Source: hostModule}, index: 1},
+		},
+	}
+	m := &ModuleInstance{Engine: me}
+
+	t.Run("host", func(t *testing.T) {
+		gf, ok := m.LookupFunction(nil, 0, 0).(*lookedUpGoFunction)
+		require.True(t, ok)
+		require.Nil(t, gf.lookedUpModule) // GoFunction doesn't need looked up module.
+		require.Equal(t, &hostModule.FunctionDefinitionSection[0], gf.def)
+		err := gf.CallWithStack(context.Background(), nil)
+		require.NoError(t, err)
+
+		gmf, ok := m.LookupFunction(nil, 0, 1).(*lookedUpGoFunction)
+		require.True(t, ok)
+		require.Equal(t, m, gmf.lookedUpModule)
+		require.Equal(t, hostModule.CodeSection[1].GoFunc, gmf.g)
+		require.Equal(t, &hostModule.FunctionDefinitionSection[1], gmf.def)
+		err = gmf.CallWithStack(context.Background(), nil)
+		require.NoError(t, err)
+
+		require.Equal(t, 2, called)
+	})
+
+	t.Run("wasm", func(t *testing.T) {
+		me.lookupEntries[2] = mockModuleEngineLookupEntry{
+			m: &ModuleInstance{
+				Source: &Module{},
+				Engine: me,
+			},
+			index: 100,
+		}
+		wf, ok := m.LookupFunction(nil, 0, 2).(*mockCallEngine)
+		require.True(t, ok)
+		require.Equal(t, Index(100), wf.index)
+	})
+}
+
+func Test_lookedUpGoModuleFunction(t *testing.T) {
+	def := &FunctionDefinition{
+		Functype: &FunctionType{
+			ParamNumInUint64: 2, ResultNumInUint64: 1,
+			Params: []ValueType{ValueTypeI64, ValueTypeF64}, Results: []ValueType{ValueTypeI64},
+		},
+	}
+
+	expMod := &ModuleInstance{}
+	var called bool
+	l := &lookedUpGoFunction{
+		lookedUpModule: expMod,
+		def:            def,
+		g: api.GoModuleFunc(
+			func(ctx context.Context, mod api.Module, stack []uint64) {
+				require.Equal(t, expMod, mod)
+				require.Equal(t, []uint64{math.MaxUint64, math.Float64bits(math.Pi)}, stack)
+				stack[0] = 1
+				called = true
+			},
+		),
+	}
+	require.Equal(t, l.Definition().(*FunctionDefinition), l.def)
+	result, err := l.Call(context.Background(), math.MaxUint64, math.Float64bits(math.Pi))
+	require.NoError(t, err)
+	require.True(t, called)
+	require.Equal(t, []uint64{1}, result)
+	require.Equal(t, 1, len(result))
+}
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index 195095fb..00b418f0 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -413,6 +413,12 @@ type mockModuleEngine struct {
 	functionRefs         map[Index]Reference
 	resolveImportsCalled map[Index]Index
 	importedMemModEngine ModuleEngine
+	lookupEntries        map[Index]mockModuleEngineLookupEntry
+}
+
+type mockModuleEngineLookupEntry struct {
+	m     *ModuleInstance
+	index Index
 }
 
 type mockCallEngine struct {
@@ -436,8 +442,11 @@ func (e *mockEngine) CompileModule(context.Context, *Module, []experimental.Func
 }
 
 // LookupFunction implements the same method as documented on wasm.Engine.
-func (e *mockModuleEngine) LookupFunction(*TableInstance, FunctionTypeID, Index) (api.Function, error) {
-	return nil, nil
+func (e *mockModuleEngine) LookupFunction(_ *TableInstance, _ FunctionTypeID, offset Index) (*ModuleInstance, Index) {
+	if entry, ok := e.lookupEntries[offset]; ok {
+		return entry.m, entry.index
+	}
+	return nil, 0
 }
 
 // CompiledModuleCount implements the same method as documented on wasm.Engine.