Eliminate runtime shift operations for register access

Author: kateinoigakukun

Sources/WasmKit/Execution/Execution.swift

@@ -130,14 +130,50 @@ extension Sp {
             return self[Int(index)] = newValue.storage
         }
     }
+
+    subscript<R: ShiftedVReg>(_ index: R) -> UntypedValue {
+        get {
+            return UntypedValue(storage: read(shifted: index))
+        }
+        nonmutating set {
+            write(shifted: index, newValue)
+        }
+    }
+
+    private func read<T: FixedWidthInteger, R: ShiftedVReg>(shifted index: R) -> T {
+        return UnsafeRawPointer(self).advanced(by: Int(index.value)).withMemoryRebound(to: T.self, capacity: 1) {
+            $0.pointee
+        }
+    }
     private func read<T: FixedWidthInteger, R: FixedWidthInteger>(_ index: R) -> T {
         return self.advanced(by: Int(index)).withMemoryRebound(to: T.self, capacity: 1) {
             $0.pointee
         }
     }
+    private func write<R: ShiftedVReg>(shifted index: R, _ value: UntypedValue) {
+        UnsafeMutableRawPointer(self).advanced(by: Int(index.value)).storeBytes(of: value.storage, as: UInt64.self)
+    }
     private func write<R: FixedWidthInteger>(_ index: R, _ value: UntypedValue) {
         self[Int(index)] = value
     }
+
+    subscript<R: ShiftedVReg>(i32 index: R) -> UInt32 {
+        get { return read(shifted: index) }
+        nonmutating set { write(shifted: index, .i32(newValue)) }
+    }
+    subscript<R: ShiftedVReg>(i64 index: R) -> UInt64 {
+        get { return read(shifted: index) }
+        nonmutating set { write(shifted: index, .i64(newValue)) }
+    }
+    subscript<R: ShiftedVReg>(f32 index: R) -> Float32 {
+        get { return Float32(bitPattern: read(shifted: index)) }
+        nonmutating set { write(shifted: index, .f32(newValue)) }
+    }
+    subscript<R: ShiftedVReg>(f64 index: R) -> Float64 {
+        get { return Float64(bitPattern: read(shifted: index)) }
+        nonmutating set { write(shifted: index, .f64(newValue)) }
+    }
+
     subscript<R: FixedWidthInteger>(i32 index: R) -> UInt32 {
         get { return read(index) }
         nonmutating set { write(index, .i32(newValue)) }

Sources/WasmKit/Execution/Instructions/Instruction.swift

@@ -431,7 +431,7 @@ extension Instruction {
             return Self(reg: reg, rawGlobal: rawGlobal)
         }
         @inline(__always) static func emit(to emitSlot: ((Self) -> CodeSlot) -> Void) {
-            emitSlot { CodeSlot(bitPattern: $0.reg) }
+            emitSlot { unsafeBitCast(($0.reg) as (LLVReg), to: CodeSlot.self) }
             emitSlot { $0.rawGlobal }
         }
     }
@@ -623,7 +623,7 @@ extension Instruction {
         }
         @inline(__always) static func emit(to emitSlot: ((Self) -> CodeSlot) -> Void) {
             emitSlot { unsafeBitCast(($0.value) as (UntypedValue), to: CodeSlot.self) }
-            emitSlot { CodeSlot(bitPattern: $0.result) }
+            emitSlot { unsafeBitCast(($0.result) as (LLVReg), to: CodeSlot.self) }
         }
     }
 

Sources/WasmKit/Execution/Instructions/InstructionSupport.swift

@@ -1,8 +1,48 @@
 import WasmParser
 
+/// A register that is used to store a value in the stack.
 typealias VReg = Int16
-typealias LVReg = Int32
-typealias LLVReg = Int64
+
+/// A register value that is pre-shifted to avoid runtime shift operation.
+protocol ShiftedVReg {
+    associatedtype Storage: FixedWidthInteger
+
+    /// The value of the shifted register.
+    /// Must be a multiple of `MemoryLayout<StackSlot>.size`.
+    var value: Storage { get }
+}
+
+/// A larger (32-bit) version of `VReg`
+/// Used to utilize halfword loads instructions.
+struct LVReg: Equatable, ShiftedVReg {
+    let value: Int32
+
+    init(_ value: VReg) {
+        // Pre-shift to avoid runtime shift operation by using
+        // unused high bits.
+        self.value = Int32(value) * Int32(MemoryLayout<StackSlot>.size)
+    }
+
+    init(storage: Int32) {
+        self.value = storage
+    }
+}
+
+/// A larger (64-bit) version of `VReg`
+/// Used to utilize word loads instructions.
+struct LLVReg: Equatable, ShiftedVReg {
+    let value: Int64
+
+    init(_ value: VReg) {
+        // Pre-shift to avoid runtime shift operation by using
+        // unused high bits.
+        self.value = Int64(value) * Int64(MemoryLayout<StackSlot>.size)
+    }
+
+    init(storage: Int64) {
+        self.value = storage
+    }
+}
 
 extension RawUnsignedInteger {
     init(_ slot: CodeSlot, shiftWidth: Int) {
@@ -85,10 +125,10 @@ extension VReg: InstructionImmediate {
 
 extension LLVReg: InstructionImmediate {
     static func load(from pc: inout Pc) -> Self {
-        Self(bitPattern: pc.read())
+        Self(storage: Int64(bitPattern: pc.read()))
     }
     static func emit(to emitSlot: ((Self) -> CodeSlot) -> Void) {
-        emitSlot { UInt64(bitPattern: $0) }
+        emitSlot { UInt64(bitPattern: $0.value) }
     }
 }
 
@@ -184,6 +224,8 @@ struct InstructionPrintingContext {
             return "reg:\(reg)"
         }
     }
+    func reg(_ x: LVReg) -> String { reg(x.value) }
+    func reg(_ x: LLVReg) -> String { reg(x.value) }
 
     func offset(_ offset: UInt64) -> String {
         "offset: \(offset)"

Sources/WasmKit/Translator.swift

@@ -853,7 +853,7 @@ struct InstructionTranslator<Context: TranslatorContext>: InstructionVisitor {
     @discardableResult
     private mutating func emitCopyStack(from source: VReg, to dest: VReg) -> Bool {
         guard source != dest else { return false }
-        emit(.copyStack(Instruction.CopyStackOperand(source: Int32(source), dest: Int32(dest))))
+        emit(.copyStack(Instruction.CopyStackOperand(source: LVReg(source), dest: LVReg(dest))))
         return true
     }
 

Utilities/Sources/VMSpec/Immediate.swift

@@ -125,12 +125,6 @@ extension VMGen {
                         output += """
                                 emitSlot { $0.\(slot[0].name) }
 
-                        """
-                        continue
-                    case "LLVReg":
-                        output += """
-                                emitSlot { CodeSlot(bitPattern: $0.\(slot[0].name)) }
-
                         """
                         continue
                     default: break