Repository: tetratelabs/wazero
Pull request #89 (2021-12-24 00:36:24)

jit: support for load{8,16,32}
- [x] load8
- [x] load16
- [x] load32

mathetake: https://github.com/tetratelabs/wazero/pull/89/commits/8d08a5611a1d4f4d9199ae546608ae53c1850461 removed the setValueType for valueLocation, and instead just have setRegisterType to distinguish whether it should be placed on Int register or float!

Patch:
diff --git a/wasm/jit/jit_amd64.go b/wasm/jit/jit_amd64.go
index 5e69a2dd..79dedc4f 100644
--- a/wasm/jit/jit_amd64.go
+++ b/wasm/jit/jit_amd64.go
@@ -102,11 +102,17 @@ func (e *engine) compileWasmFunction(f *wasm.FunctionInstance) (*compiledWasmFun
 				return nil, fmt.Errorf("error handling load operation: %w", err)
 			}
 		case *wazeroir.OperationLoad8:
-			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
+			if err := builder.handleLoad8(o); err != nil {
+				return nil, fmt.Errorf("error handling load8 operation: %w", err)
+			}
 		case *wazeroir.OperationLoad16:
-			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
+			if err := builder.handleLoad16(o); err != nil {
+				return nil, fmt.Errorf("error handling load16 operation: %w", err)
+			}
 		case *wazeroir.OperationLoad32:
-			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
+			if err := builder.handleLoad32(o); err != nil {
+				return nil, fmt.Errorf("error handling load16 operation: %w", err)
+			}
 		case *wazeroir.OperationStore:
 			return nil, fmt.Errorf("unsupported operation in JIT compiler: %v", o)
 		case *wazeroir.OperationStore8:
@@ -255,7 +261,12 @@ func (b *amd64Builder) newCompiledWasmFunction(code []byte) *compiledWasmFunctio
 func (b *amd64Builder) pushFunctionInputs() {
 	for _, t := range b.f.Signature.InputTypes {
 		loc := b.locationStack.pushValueOnStack()
-		loc.setValueType(wazeroir.WasmValueTypeToUnsignedType(t))
+		switch t {
+		case wasm.ValueTypeI32, wasm.ValueTypeI64:
+			loc.setRegisterType(generalPurposeRegisterTypeInt)
+		case wasm.ValueTypeF32, wasm.ValueTypeF64:
+			loc.setRegisterType(generalPurposeRegisterTypeFloat)
+		}
 	}
 }
 
@@ -463,21 +474,20 @@ func (b *amd64Builder) handleGlobalGet(o *wazeroir.OperationGlobalGet) error {
 
 	// Record that the retrieved global value on the top of the stack is now in a register.
 	loc := b.locationStack.pushValueOnRegister(valueReg)
-	loc.setValueType(wazeroir.WasmValueTypeToUnsignedType(wasmType))
+	switch wasmType {
+	case wasm.ValueTypeI32, wasm.ValueTypeI64:
+		loc.setRegisterType(generalPurposeRegisterTypeInt)
+	case wasm.ValueTypeF32, wasm.ValueTypeF64:
+		loc.setRegisterType(generalPurposeRegisterTypeFloat)
+	}
 	return nil
 }
 
 func (b *amd64Builder) handleGlobalSet(o *wazeroir.OperationGlobalSet) error {
 	// First, move the value to set into a temporary register.
 	val := b.locationStack.pop()
-	if val.onStack() {
-		if err := b.moveStackToRegisterWithAllocation(val.registerType(), val); err != nil {
-			return err
-		}
-	} else if val.onConditionalRegister() {
-		if err := b.moveConditionalToGeneralPurposeRegister(val); err != nil {
-			return err
-		}
+	if err := b.ensureOnGeneralPurposeRegister(val); err != nil {
+		return err
 	}
 
 	// Allocate a register to hold the memory location of the target global instance.
@@ -796,22 +806,15 @@ func (b *amd64Builder) emitDropRange(r *wazeroir.InclusiveRange) error {
 // the physical registers or memory stack, or maybe conditional register.
 func (b *amd64Builder) handleSelect() error {
 	c := b.locationStack.pop()
+	if err := b.ensureOnGeneralPurposeRegister(c); err != nil {
+		return err
+	}
+
 	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.moveConditionalToGeneralPurposeRegister(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
@@ -900,7 +903,7 @@ func (b *amd64Builder) handlePick(o *wazeroir.OperationPick) error {
 	// Now we already placed the picked value on the register,
 	// so push the location onto the stack.
 	loc := b.locationStack.pushValueOnRegister(reg)
-	loc.setValueType(pickTarget.valueType)
+	loc.setRegisterType(pickTarget.registerType())
 	return nil
 }
 
@@ -1094,20 +1097,14 @@ func (b *amd64Builder) handleLe(o *wazeroir.OperationLe) error {
 	// Finally we have the result on the conditional register,
 	// so record it.
 	loc := b.locationStack.pushValueOnConditionalRegister(resultConditionState)
-	loc.setValueType(wazeroir.UnsignedTypeI32)
+	loc.setRegisterType(generalPurposeRegisterTypeInt)
 	return nil
 }
 
 func (b *amd64Builder) handleLoad(o *wazeroir.OperationLoad) error {
 	base := b.locationStack.pop()
-	if base.onStack() {
-		if err := b.moveStackToRegisterWithAllocation(generalPurposeRegisterTypeInt, base); err != nil {
-			return err
-		}
-	} else if base.onConditionalRegister() {
-		if err := b.moveConditionalToGeneralPurposeRegister(base); err != nil {
-			return err
-		}
+	if err := b.ensureOnGeneralPurposeRegister(base); err != nil {
+		return err
 	}
 
 	// At this point, base's value is on the integer general purpose reg.
@@ -1158,7 +1155,7 @@ func (b *amd64Builder) handleLoad(o *wazeroir.OperationLoad) error {
 		moveFromMemory.From.Scale = 1
 		b.addInstruction(moveFromMemory)
 		top := b.locationStack.pushValueOnRegister(reg)
-		top.setValueType(o.Type)
+		top.setRegisterType(generalPurposeRegisterTypeInt)
 	} else {
 		// For float types, we read the value to the float register.
 		floatReg, err := b.allocateRegister(generalPurposeRegisterTypeFloat)
@@ -1175,13 +1172,123 @@ func (b *amd64Builder) handleLoad(o *wazeroir.OperationLoad) error {
 		moveFromMemory.From.Scale = 1
 		b.addInstruction(moveFromMemory)
 		top := b.locationStack.pushValueOnRegister(floatReg)
-		top.setValueType(o.Type)
+		top.setRegisterType(generalPurposeRegisterTypeFloat)
 		// We no longer need the int register so mark it unused.
 		b.locationStack.markRegisterUnused(reg)
 	}
 	return nil
 }
 
+func (b *amd64Builder) handleLoad8(o *wazeroir.OperationLoad8) error {
+	base := b.locationStack.pop()
+	if err := b.ensureOnGeneralPurposeRegister(base); err != nil {
+		return err
+	}
+
+	// At this point, base's value is on the integer general purpose reg.
+	// We reuse the register below, so we alias it here for readability.
+	reg := base.register
+
+	// We have to calculate the offset on the memory region.
+	addOffsetToBase := b.newProg()
+	addOffsetToBase.As = x86.AADDL // 32-bit!
+	addOffsetToBase.To.Type = obj.TYPE_REG
+	addOffsetToBase.To.Reg = reg
+	addOffsetToBase.From.Type = obj.TYPE_CONST
+	addOffsetToBase.From.Offset = int64(o.Arg.Offest)
+	b.addInstruction(addOffsetToBase)
+
+	// Then move a byte at the offset to the register.
+	// Note that Load8 is only for integer types.
+	moveFromMemory := b.newProg()
+	moveFromMemory.As = x86.AMOVB // 1 byte!
+	moveFromMemory.To.Type = obj.TYPE_REG
+	moveFromMemory.To.Reg = reg
+	moveFromMemory.From.Type = obj.TYPE_MEM
+	moveFromMemory.From.Reg = reservedRegisterForMemory
+	moveFromMemory.From.Index = reg
+	moveFromMemory.From.Scale = 1
+	b.addInstruction(moveFromMemory)
+	top := b.locationStack.pushValueOnRegister(reg)
+
+	// The result of load8 is always int type.
+	top.setRegisterType(generalPurposeRegisterTypeInt)
+	return nil
+}
+
+func (b *amd64Builder) handleLoad16(o *wazeroir.OperationLoad16) error {
+	base := b.locationStack.pop()
+	if err := b.ensureOnGeneralPurposeRegister(base); err != nil {
+		return err
+	}
+
+	// At this point, base's value is on the integer general purpose reg.
+	// We reuse the register below, so we alias it here for readability.
+	reg := base.register
+
+	// We have to calculate the offset on the memory region.
+	addOffsetToBase := b.newProg()
+	addOffsetToBase.As = x86.AADDL // 32-bit!
+	addOffsetToBase.To.Type = obj.TYPE_REG
+	addOffsetToBase.To.Reg = reg
+	addOffsetToBase.From.Type = obj.TYPE_CONST
+	addOffsetToBase.From.Offset = int64(o.Arg.Offest)
+	b.addInstruction(addOffsetToBase)
+
+	// Then move a byte at the offset to the register.
+	// Note that Load16 is only for integer types.
+	moveFromMemory := b.newProg()
+	moveFromMemory.As = x86.AMOVW // 2 byte!
+	moveFromMemory.To.Type = obj.TYPE_REG
+	moveFromMemory.To.Reg = reg
+	moveFromMemory.From.Type = obj.TYPE_MEM
+	moveFromMemory.From.Reg = reservedRegisterForMemory
+	moveFromMemory.From.Index = reg
+	moveFromMemory.From.Scale = 1
+	b.addInstruction(moveFromMemory)
+	top := b.locationStack.pushValueOnRegister(reg)
+
+	// The result of load16 is always int type.
+	top.setRegisterType(generalPurposeRegisterTypeInt)
+	return nil
+}
+
+func (b *amd64Builder) handleLoad32(o *wazeroir.OperationLoad32) error {
+	base := b.locationStack.pop()
+	if err := b.ensureOnGeneralPurposeRegister(base); err != nil {
+		return err
+	}
+
+	// At this point, base's value is on the integer general purpose reg.
+	// We reuse the register below, so we alias it here for readability.
+	reg := base.register
+
+	// We have to calculate the offset on the memory region.
+	addOffsetToBase := b.newProg()
+	addOffsetToBase.As = x86.AADDL // 32-bit!
+	addOffsetToBase.To.Type = obj.TYPE_REG
+	addOffsetToBase.To.Reg = reg
+	addOffsetToBase.From.Type = obj.TYPE_CONST
+	addOffsetToBase.From.Offset = int64(o.Arg.Offest)
+	b.addInstruction(addOffsetToBase)
+
+	// Then move a byte at the offset to the register.
+	moveFromMemory := b.newProg()
+	moveFromMemory.As = x86.AMOVL // 3 byte!
+	moveFromMemory.To.Type = obj.TYPE_REG
+	moveFromMemory.To.Reg = reg
+	moveFromMemory.From.Type = obj.TYPE_MEM
+	moveFromMemory.From.Reg = reservedRegisterForMemory
+	moveFromMemory.From.Index = reg
+	moveFromMemory.From.Scale = 1
+	b.addInstruction(moveFromMemory)
+	top := b.locationStack.pushValueOnRegister(reg)
+
+	// The result of load32 is always int type.
+	top.setRegisterType(generalPurposeRegisterTypeInt)
+	return nil
+}
+
 func (b *amd64Builder) handleMemoryGrow() {
 	b.callBuiltinFunctionFromConstIndex(builtinFunctionIndexMemoryGrow)
 }
@@ -1189,7 +1296,7 @@ func (b *amd64Builder) handleMemoryGrow() {
 func (b *amd64Builder) handleMemorySize() {
 	b.callBuiltinFunctionFromConstIndex(builtinFunctionIndexMemorySize)
 	loc := b.locationStack.pushValueOnStack() // The size is pushed on the top.
-	loc.setValueType(wazeroir.UnsignedTypeI32)
+	loc.setRegisterType(generalPurposeRegisterTypeInt)
 }
 
 func (b *amd64Builder) callBuiltinFunctionFromConstIndex(index int64) {
@@ -1209,7 +1316,7 @@ func (b *amd64Builder) handleConstI32(o *wazeroir.OperationConstI32) error {
 		return err
 	}
 	loc := b.locationStack.pushValueOnRegister(reg)
-	loc.setValueType(wazeroir.UnsignedTypeI32)
+	loc.setRegisterType(generalPurposeRegisterTypeInt)
 
 	prog := b.newProg()
 	prog.As = x86.AMOVL // Note 32-bit move!
@@ -1227,7 +1334,7 @@ func (b *amd64Builder) handleConstI64(o *wazeroir.OperationConstI64) error {
 		return err
 	}
 	loc := b.locationStack.pushValueOnRegister(reg)
-	loc.setValueType(wazeroir.UnsignedTypeI64)
+	loc.setRegisterType(generalPurposeRegisterTypeInt)
 
 	prog := b.newProg()
 	prog.As = x86.AMOVQ
@@ -1245,7 +1352,7 @@ func (b *amd64Builder) handleConstF32(o *wazeroir.OperationConstF32) error {
 		return err
 	}
 	loc := b.locationStack.pushValueOnRegister(reg)
-	loc.setValueType(wazeroir.UnsignedTypeF32)
+	loc.setRegisterType(generalPurposeRegisterTypeFloat)
 
 	// We cannot directly load the value from memory to float regs,
 	// so we move it to int reg temporarily.
@@ -1279,7 +1386,7 @@ func (b *amd64Builder) handleConstF64(o *wazeroir.OperationConstF64) error {
 		return err
 	}
 	loc := b.locationStack.pushValueOnRegister(reg)
-	loc.setValueType(wazeroir.UnsignedTypeF64)
+	loc.setRegisterType(generalPurposeRegisterTypeFloat)
 
 	// We cannot directly load the value from memory to float regs,
 	// so we move it to int reg temporarily.
@@ -1660,3 +1767,18 @@ func (b *amd64Builder) initializeReservedRegisters() {
 	prog.From.Reg = reg
 	b.addInstruction(prog)
 }
+
+// ensureOnGeneralPurposeRegister ensures that the given value is located on a
+// general purpose register of an appropriate type.
+func (b *amd64Builder) ensureOnGeneralPurposeRegister(loc *valueLocation) error {
+	if loc.onStack() {
+		if err := b.moveStackToRegisterWithAllocation(loc.registerType(), loc); err != nil {
+			return err
+		}
+	} else if loc.onConditionalRegister() {
+		if err := b.moveConditionalToGeneralPurposeRegister(loc); err != nil {
+			return err
+		}
+	}
+	return nil
+}
diff --git a/wasm/jit/jit_value_location_amd64.go b/wasm/jit/jit_value_location_amd64.go
index 00b7cb09..d9140721 100644
--- a/wasm/jit/jit_value_location_amd64.go
+++ b/wasm/jit/jit_value_location_amd64.go
@@ -5,8 +5,6 @@ package jit
 
 import (
 	"github.com/twitchyliquid64/golang-asm/obj/x86"
-
-	"github.com/tetratelabs/wazero/wasm/wazeroir"
 )
 
 // Reserved registers.
@@ -68,7 +66,7 @@ const (
 // and it has the information about where it exists in the physical machine.
 // It might exist in registers, or maybe on in the non-virtual physical stack allocated in memory.
 type valueLocation struct {
-	valueType wazeroir.UnsignedType
+	regType generalPurposeRegisterType
 	// Set to -1 if the value is stored in the memory stack.
 	register int16
 	// Set to conditionalRegisterStateUnset if the value is not on the conditional register.
@@ -79,19 +77,11 @@ type valueLocation struct {
 }
 
 func (v *valueLocation) registerType() (t generalPurposeRegisterType) {
-	switch v.valueType {
-	case wazeroir.UnsignedTypeI32, wazeroir.UnsignedTypeI64:
-		t = generalPurposeRegisterTypeInt
-	case wazeroir.UnsignedTypeF32, wazeroir.UnsignedTypeF64:
-		t = generalPurposeRegisterTypeFloat
-	default:
-		panic("unreachable")
-	}
-	return
+	return v.regType
 }
 
-func (v *valueLocation) setValueType(t wazeroir.UnsignedType) {
-	v.valueType = t
+func (v *valueLocation) setRegisterType(t generalPurposeRegisterType) {
+	v.regType = t
 }
 
 func (v *valueLocation) setRegister(reg int16) {
@@ -147,7 +137,7 @@ func (s *valueLocationStack) clone() *valueLocationStack {
 	ret.stack = make([]*valueLocation, len(s.stack))
 	for i, v := range s.stack {
 		ret.stack[i] = &valueLocation{
-			valueType:           v.valueType,
+			regType:             v.regType,
 			conditionalRegister: v.conditionalRegister,
 			stackPointer:        v.stackPointer,
 			register:            v.register,
diff --git a/wasm/wazeroir/compiler.go b/wasm/wazeroir/compiler.go
index a66c21a9..6c694784 100644
--- a/wasm/wazeroir/compiler.go
+++ b/wasm/wazeroir/compiler.go
@@ -149,7 +149,7 @@ func Compile(f *wasm.FunctionInstance) (*CompilationResult, error) {
 
 	// Push function arguments.
 	for _, t := range f.Signature.InputTypes {
-		c.stackPush(WasmValueTypeToUnsignedType(t))
+		c.stackPush(wasmValueTypeToUnsignedType(t))
 	}
 	// Emit const expressions for locals.
 	// Note that here we don't take function arguments
@@ -162,7 +162,7 @@ func Compile(f *wasm.FunctionInstance) (*CompilationResult, error) {
 	// Insert the function control frame.
 	returns := make([]UnsignedType, 0, len(f.Signature.ReturnTypes))
 	for _, t := range f.Signature.ReturnTypes {
-		returns = append(returns, WasmValueTypeToUnsignedType(t))
+		returns = append(returns, wasmValueTypeToUnsignedType(t))
 	}
 	c.controlFrames.push(&controlFrame{
 		frameID:          c.nextID(),
@@ -232,7 +232,7 @@ operatorSwitch:
 			kind:             controlFrameKindBlockWithoutContinuationLabel,
 		}
 		for _, t := range bt.ReturnTypes {
-			frame.returns = append(frame.returns, WasmValueTypeToUnsignedType(t))
+			frame.returns = append(frame.returns, wasmValueTypeToUnsignedType(t))
 		}
 		c.controlFrames.push(frame)
 
@@ -258,7 +258,7 @@ operatorSwitch:
 			kind:             controlFrameKindLoop,
 		}
 		for _, t := range bt.ReturnTypes {
-			frame.returns = append(frame.returns, WasmValueTypeToUnsignedType(t))
+			frame.returns = append(frame.returns, wasmValueTypeToUnsignedType(t))
 		}
 		c.controlFrames.push(frame)
 
@@ -298,7 +298,7 @@ operatorSwitch:
 			kind: controlFrameKindIfWithoutElse,
 		}
 		for _, t := range bt.ReturnTypes {
-			frame.returns = append(frame.returns, WasmValueTypeToUnsignedType(t))
+			frame.returns = append(frame.returns, wasmValueTypeToUnsignedType(t))
 		}
 		c.controlFrames.push(frame)
 
diff --git a/wasm/wazeroir/signature.go b/wasm/wazeroir/signature.go
index 2745531f..09aa66e5 100644
--- a/wasm/wazeroir/signature.go
+++ b/wasm/wazeroir/signature.go
@@ -176,9 +176,9 @@ func wasmOpcodeSignature(f *wasm.FunctionInstance, op wasm.Opcode, index uint32)
 		}
 		var t UnsignedType
 		if index < inputLen {
-			t = WasmValueTypeToUnsignedType(f.Signature.InputTypes[index])
+			t = wasmValueTypeToUnsignedType(f.Signature.InputTypes[index])
 		} else {
-			t = WasmValueTypeToUnsignedType(f.LocalTypes[index-inputLen])
+			t = wasmValueTypeToUnsignedType(f.LocalTypes[index-inputLen])
 		}
 		return &signature{out: []UnsignedType{t}}, nil
 	case wasm.OpcodeLocalSet:
@@ -188,9 +188,9 @@ func wasmOpcodeSignature(f *wasm.FunctionInstance, op wasm.Opcode, index uint32)
 		}
 		var t UnsignedType
 		if index < inputLen {
-			t = WasmValueTypeToUnsignedType(f.Signature.InputTypes[index])
+			t = wasmValueTypeToUnsignedType(f.Signature.InputTypes[index])
 		} else {
-			t = WasmValueTypeToUnsignedType(f.LocalTypes[index-inputLen])
+			t = wasmValueTypeToUnsignedType(f.LocalTypes[index-inputLen])
 		}
 		return &signature{in: []UnsignedType{t}}, nil
 	case wasm.OpcodeLocalTee:
@@ -200,9 +200,9 @@ func wasmOpcodeSignature(f *wasm.FunctionInstance, op wasm.Opcode, index uint32)
 		}
 		var t UnsignedType
 		if index < inputLen {
-			t = WasmValueTypeToUnsignedType(f.Signature.InputTypes[index])
+			t = wasmValueTypeToUnsignedType(f.Signature.InputTypes[index])
 		} else {
-			t = WasmValueTypeToUnsignedType(f.LocalTypes[index-inputLen])
+			t = wasmValueTypeToUnsignedType(f.LocalTypes[index-inputLen])
 		}
 		return &signature{in: []UnsignedType{t}, out: []UnsignedType{t}}, nil
 	case wasm.OpcodeGlobalGet:
@@ -210,14 +210,14 @@ func wasmOpcodeSignature(f *wasm.FunctionInstance, op wasm.Opcode, index uint32)
 			return nil, fmt.Errorf("invalid global index for global.get %d >= %d", index, len(f.ModuleInstance.Globals))
 		}
 		return &signature{
-			out: []UnsignedType{WasmValueTypeToUnsignedType(f.ModuleInstance.Globals[index].Type.ValType)},
+			out: []UnsignedType{wasmValueTypeToUnsignedType(f.ModuleInstance.Globals[index].Type.ValType)},
 		}, nil
 	case wasm.OpcodeGlobalSet:
 		if len(f.ModuleInstance.Globals) <= int(index) {
 			return nil, fmt.Errorf("invalid global index for global.get %d >= %d", index, len(f.ModuleInstance.Globals))
 		}
 		return &signature{
-			in: []UnsignedType{WasmValueTypeToUnsignedType(f.ModuleInstance.Globals[index].Type.ValType)},
+			in: []UnsignedType{wasmValueTypeToUnsignedType(f.ModuleInstance.Globals[index].Type.ValType)},
 		}, nil
 	case wasm.OpcodeI32Load:
 		return signature_I32_I32, nil
@@ -354,15 +354,15 @@ func wasmOpcodeSignature(f *wasm.FunctionInstance, op wasm.Opcode, index uint32)
 func funcTypeToSignature(tps *wasm.FunctionType) *signature {
 	ret := &signature{}
 	for _, vt := range tps.InputTypes {
-		ret.in = append(ret.in, WasmValueTypeToUnsignedType(vt))
+		ret.in = append(ret.in, wasmValueTypeToUnsignedType(vt))
 	}
 	for _, vt := range tps.ReturnTypes {
-		ret.out = append(ret.out, WasmValueTypeToUnsignedType(vt))
+		ret.out = append(ret.out, wasmValueTypeToUnsignedType(vt))
 	}
 	return ret
 }
 
-func WasmValueTypeToUnsignedType(vt wasm.ValueType) UnsignedType {
+func wasmValueTypeToUnsignedType(vt wasm.ValueType) UnsignedType {
 	switch vt {
 	case wasm.ValueTypeI32:
 		return UnsignedTypeI32


Test patch:
diff --git a/wasm/jit/jit_amd64_test.go b/wasm/jit/jit_amd64_test.go
index a1fb7889..e66a958e 100644
--- a/wasm/jit/jit_amd64_test.go
+++ b/wasm/jit/jit_amd64_test.go
@@ -661,7 +661,7 @@ func TestAmd64Builder_handlePick(t *testing.T) {
 	t.Run("on reg", func(t *testing.T) {
 		// Set up the pick target original value.
 		pickTargetLocation := builder.locationStack.pushValueOnRegister(int16(x86.REG_R10))
-		pickTargetLocation.setValueType(wazeroir.UnsignedTypeI32)
+		pickTargetLocation.setRegisterType(generalPurposeRegisterTypeInt)
 		builder.locationStack.pushValueOnStack() // Dummy value!
 		builder.movIntConstToRegister(100, pickTargetLocation.register)
 		// Now insert pick code.
@@ -759,7 +759,7 @@ func TestAmd64Builder_handleConstI32(t *testing.T) {
 			// To verify the behavior, we increment and push the const value
 			// to the stack.
 			loc := builder.locationStack.peek()
-			require.Equal(t, wazeroir.UnsignedTypeI32, loc.valueType)
+			require.Equal(t, generalPurposeRegisterTypeInt, loc.registerType())
 			prog := builder.newProg()
 			prog.As = x86.AINCQ
 			prog.To.Type = obj.TYPE_REG
@@ -801,7 +801,7 @@ func TestAmd64Builder_handleConstI64(t *testing.T) {
 			// To verify the behavior, we increment and push the const value
 			// to the stack.
 			loc := builder.locationStack.peek()
-			require.Equal(t, wazeroir.UnsignedTypeI64, loc.valueType)
+			require.Equal(t, generalPurposeRegisterTypeInt, loc.registerType())
 			prog := builder.newProg()
 			prog.As = x86.AINCQ
 			prog.To.Type = obj.TYPE_REG
@@ -843,7 +843,7 @@ func TestAmd64Builder_handleConstF32(t *testing.T) {
 			// To verify the behavior, we double and push the const value
 			// to the stack.
 			loc := builder.locationStack.peek()
-			require.Equal(t, wazeroir.UnsignedTypeF32, loc.valueType)
+			require.Equal(t, generalPurposeRegisterTypeFloat, loc.registerType())
 			prog := builder.newProg()
 			prog.As = x86.AADDSS
 			prog.To.Type = obj.TYPE_REG
@@ -887,7 +887,7 @@ func TestAmd64Builder_handleConstF64(t *testing.T) {
 			// To verify the behavior, we double and push the const value
 			// to the stack.
 			loc := builder.locationStack.peek()
-			require.Equal(t, wazeroir.UnsignedTypeF64, loc.valueType)
+			require.Equal(t, generalPurposeRegisterTypeFloat, loc.registerType())
 			prog := builder.newProg()
 			prog.As = x86.AADDSD
 			prog.To.Type = obj.TYPE_REG
@@ -1426,22 +1426,25 @@ func TestAmd64Builder_handleLoad(t *testing.T) {
 
 			// At this point, the loaded value must be on top of the stack, and placed on a register.
 			loadedValue := builder.locationStack.peek()
-			require.Equal(t, o.Type, loadedValue.valueType)
 			require.True(t, loadedValue.onRegister())
 
 			// Double the loaded value in order to verify the behavior.
 			var addInst obj.As
 			switch tp {
 			case wazeroir.UnsignedTypeI32:
+				require.Equal(t, generalPurposeRegisterTypeInt, loadedValue.registerType())
 				require.True(t, isIntRegister(loadedValue.register))
 				addInst = x86.AADDL
 			case wazeroir.UnsignedTypeI64:
+				require.Equal(t, generalPurposeRegisterTypeInt, loadedValue.registerType())
 				require.True(t, isIntRegister(loadedValue.register))
 				addInst = x86.AADDQ
 			case wazeroir.UnsignedTypeF32:
+				require.Equal(t, generalPurposeRegisterTypeFloat, loadedValue.registerType())
 				require.True(t, isFloatRegister(loadedValue.register))
 				addInst = x86.AADDSS
 			case wazeroir.UnsignedTypeF64:
+				require.Equal(t, generalPurposeRegisterTypeFloat, loadedValue.registerType())
 				require.True(t, isFloatRegister(loadedValue.register))
 				addInst = x86.AADDSD
 			}
@@ -1499,6 +1502,188 @@ func TestAmd64Builder_handleLoad(t *testing.T) {
 	}
 }
 
+func TestAmd64Builder_handleLoad8(t *testing.T) {
+	for i, tp := range []wazeroir.SignedInt{
+		wazeroir.SignedInt32,
+		wazeroir.SignedInt64,
+		wazeroir.SignedUint32,
+		wazeroir.SignedUint64,
+	} {
+		tp := tp
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			eng := newEngine()
+			builder := requireNewBuilder(t)
+			builder.initializeReservedRegisters()
+			// Before load operations, we must push the base offset value.
+			const baseOffset = 100 // For testing. Arbitrary number is fine.
+			base := builder.locationStack.pushValueOnStack()
+			eng.stack[base.stackPointer] = baseOffset
+
+			// Emit the memory load instructions.
+			o := &wazeroir.OperationLoad8{Type: tp, Arg: &wazeroir.MemoryImmediate{Offest: 361}}
+			err := builder.handleLoad8(o)
+			require.NoError(t, err)
+
+			// At this point, the loaded value must be on top of the stack, and placed on a register.
+			loadedValue := builder.locationStack.peek()
+			require.Equal(t, generalPurposeRegisterTypeInt, loadedValue.registerType())
+			require.True(t, loadedValue.onRegister())
+
+			// Increment the loaded value in order to verify the behavior.
+			doubleLoadedValue := builder.newProg()
+			doubleLoadedValue.As = x86.AINCB
+			doubleLoadedValue.To.Type = obj.TYPE_REG
+			doubleLoadedValue.To.Reg = loadedValue.register
+			builder.addInstruction(doubleLoadedValue)
+
+			// We need to write the result back to the memory stack.
+			builder.releaseRegisterToStack(loadedValue)
+
+			// Compile.
+			builder.returnFunction()
+			code, err := builder.compile()
+			require.NoError(t, err)
+
+			// Place the load target value to the memory.
+			mem := newMemoryInst()
+			// For testing, arbitrary byte is be fine.
+			original := byte(0x10)
+			mem.Buffer[baseOffset+o.Arg.Offest] = byte(original)
+
+			// Run code.
+			jitcall(
+				uintptr(unsafe.Pointer(&code[0])),
+				uintptr(unsafe.Pointer(eng)),
+				uintptr(unsafe.Pointer(&mem.Buffer[0])),
+			)
+
+			// Load instruction must push the loaded value to the top of the stack,
+			// so the stack pointer must be incremented.
+			require.Equal(t, uint64(1), eng.stackPointer)
+			// The loaded value must be incremented via x86.AINCB.
+			require.Equal(t, original+1, byte(eng.stack[eng.stackPointer-1]))
+		})
+	}
+}
+
+func TestAmd64Builder_handleLoad16(t *testing.T) {
+	for i, tp := range []wazeroir.SignedInt{
+		wazeroir.SignedInt32,
+		wazeroir.SignedInt64,
+		wazeroir.SignedUint32,
+		wazeroir.SignedUint64,
+	} {
+		tp := tp
+		t.Run(fmt.Sprintf("%d", i), func(t *testing.T) {
+			eng := newEngine()
+			builder := requireNewBuilder(t)
+			builder.initializeReservedRegisters()
+			// Before load operations, we must push the base offset value.
+			const baseOffset = 100 // For testing. Arbitrary number is fine.
+			base := builder.locationStack.pushValueOnStack()
+			eng.stack[base.stackPointer] = baseOffset
+
+			// Emit the memory load instructions.
+			o := &wazeroir.OperationLoad16{Type: tp, Arg: &wazeroir.MemoryImmediate{Offest: 361}}
+			err := builder.handleLoad16(o)
+			require.NoError(t, err)
+
+			// At this point, the loaded value must be on top of the stack, and placed on a register.
+			loadedValue := builder.locationStack.peek()
+			require.Equal(t, generalPurposeRegisterTypeInt, loadedValue.registerType())
+			require.True(t, loadedValue.onRegister())
+
+			// Increment the loaded value in order to verify the behavior.
+			doubleLoadedValue := builder.newProg()
+			doubleLoadedValue.As = x86.AINCW
+			doubleLoadedValue.To.Type = obj.TYPE_REG
+			doubleLoadedValue.To.Reg = loadedValue.register
+			builder.addInstruction(doubleLoadedValue)
+
+			// We need to write the result back to the memory stack.
+			builder.releaseRegisterToStack(loadedValue)
+
+			// Compile.
+			builder.returnFunction()
+			code, err := builder.compile()
+			require.NoError(t, err)
+
+			// Place the load target value to the memory.
+			mem := newMemoryInst()
+			// For testing, arbitrary uint16 is be fine.
+			original := uint16(0xff_fe)
+			binary.LittleEndian.PutUint16(mem.Buffer[baseOffset+o.Arg.Offest:], original)
+
+			// Run code.
+			jitcall(
+				uintptr(unsafe.Pointer(&code[0])),
+				uintptr(unsafe.Pointer(eng)),
+				uintptr(unsafe.Pointer(&mem.Buffer[0])),
+			)
+
+			// Load instruction must push the loaded value to the top of the stack,
+			// so the stack pointer must be incremented.
+			require.Equal(t, uint64(1), eng.stackPointer)
+			// The loaded value must be incremented via x86.AINCW.
+			require.Equal(t, original+1, uint16(eng.stack[eng.stackPointer-1]))
+		})
+	}
+}
+
+func TestAmd64Builder_handleLoad32(t *testing.T) {
+	eng := newEngine()
+	builder := requireNewBuilder(t)
+	builder.initializeReservedRegisters()
+	// Before load operations, we must push the base offset value.
+	const baseOffset = 100 // For testing. Arbitrary number is fine.
+	base := builder.locationStack.pushValueOnStack()
+	eng.stack[base.stackPointer] = baseOffset
+
+	// Emit the memory load instructions.
+	o := &wazeroir.OperationLoad32{Arg: &wazeroir.MemoryImmediate{Offest: 361}}
+	err := builder.handleLoad32(o)
+	require.NoError(t, err)
+
+	// At this point, the loaded value must be on top of the stack, and placed on a register.
+	loadedValue := builder.locationStack.peek()
+	require.Equal(t, generalPurposeRegisterTypeInt, loadedValue.registerType())
+	require.True(t, loadedValue.onRegister())
+
+	// Increment the loaded value in order to verify the behavior.
+	doubleLoadedValue := builder.newProg()
+	doubleLoadedValue.As = x86.AINCL
+	doubleLoadedValue.To.Type = obj.TYPE_REG
+	doubleLoadedValue.To.Reg = loadedValue.register
+	builder.addInstruction(doubleLoadedValue)
+
+	// We need to write the result back to the memory stack.
+	builder.releaseRegisterToStack(loadedValue)
+
+	// Compile.
+	builder.returnFunction()
+	code, err := builder.compile()
+	require.NoError(t, err)
+
+	// Place the load target value to the memory.
+	mem := newMemoryInst()
+	// For testing, arbitrary uint32 is be fine.
+	original := uint32(0xff_ff_fe)
+	binary.LittleEndian.PutUint32(mem.Buffer[baseOffset+o.Arg.Offest:], original)
+
+	// Run code.
+	jitcall(
+		uintptr(unsafe.Pointer(&code[0])),
+		uintptr(unsafe.Pointer(eng)),
+		uintptr(unsafe.Pointer(&mem.Buffer[0])),
+	)
+
+	// Load instruction must push the loaded value to the top of the stack,
+	// so the stack pointer must be incremented.
+	require.Equal(t, uint64(1), eng.stackPointer)
+	// The loaded value must be incremented via x86.AINCL.
+	require.Equal(t, original+1, uint32(eng.stack[eng.stackPointer-1]))
+}
+
 func TestAmd64Builder_handleMemoryGrow(t *testing.T) {
 	builder := requireNewBuilder(t)
 
@@ -1529,7 +1714,7 @@ func TestAmd64Builder_handleMemorySize(t *testing.T) {
 	builder.handleMemorySize()
 	// At this point, the size of memory should be pushed onto the stack.
 	require.Equal(t, uint64(1), builder.locationStack.sp)
-	require.Equal(t, wazeroir.UnsignedTypeI32, builder.locationStack.peek().valueType)
+	require.Equal(t, generalPurposeRegisterTypeInt, builder.locationStack.peek().registerType())
 
 	// Compile.
 	code, err := builder.compile()