[wasm] Add support for atomic RMW operations

This introduces the WasmAtomicRMW operation to support various atomic Read-Modify-Write (RMW) operation instructions.

The supported operations now include atomic:
- Add
- Sub
- And
- Or
- Xor
- Exchange

These are available in various 8, 16, 32, and 64-bit versions.

Bug: 427134598
Change-Id: I5fb2f2e1e77b27a782a46c84c67254cf80a23a4b
Reviewed-on: https://chrome-internal-review.googlesource.com/c/v8/fuzzilli/+/8462196
Reviewed-by: Matthias Liedtke <[email protected]>
Commit-Queue: Dominik Klemba <[email protected]>

Author: Dominik Klemba

Sources/Fuzzilli/Base/ProgramBuilder.swift

@@ -3314,6 +3314,14 @@ public class ProgramBuilder {
             b.emit(op, withInputs: [memory, address, value])
         }
 
+        @discardableResult
+        func wasmAtomicRMW(memory: Variable, lhs: Variable, rhs: Variable, op: WasmAtomicRMWType, offset: Int64) -> Variable {
+            let op = WasmAtomicRMW(op: op, offset: offset)
+            let anyInt: ILType = .wasmi32 | .wasmi64
+            let valueType = op.op.type
+            return b.emit(op, withInputs: [memory, lhs, rhs], types: [.object(ofGroup: "WasmMemory"), anyInt, valueType]).output
+        }
+
         @discardableResult
         public func wasmMemorySize(memory: Variable) -> Variable {
             return b.emit(WasmMemorySize(), withInputs: [memory],

Sources/Fuzzilli/CodeGen/CodeGeneratorWeights.swift

@@ -220,6 +220,7 @@ public let codeGeneratorWeights = [
     "WasmMemoryStoreGenerator":                 10,
     "WasmAtomicLoadGenerator":                  10,
     "WasmAtomicStoreGenerator":                 10,
+    "WasmAtomicRMWGenerator":                   10,
     "WasmMemorySizeGenerator":                  5,
     "WasmMemoryGrowGenerator":                  1,
     "WasmMemoryFillGenerator":                  5,

Sources/Fuzzilli/CodeGen/WasmCodeGenerators.swift

@@ -371,6 +371,19 @@ public let WasmCodeGenerators: [CodeGenerator] = [
         function.wasmAtomicStore(memory: memory, address: address, value: value, storeType: storeType, offset: staticOffset)
     },
 
+    CodeGenerator("WasmAtomicRMWGenerator", inContext: .wasmFunction, inputs: .required(.object(ofGroup: "WasmMemory"))) { b, memory in
+        let function = b.currentWasmModule.currentWasmFunction
+        let op = chooseUniform(from: WasmAtomicRMWType.allCases)
+        let valueType = op.type
+        let alignment = op.naturalAlignment()
+
+        let rhs = function.findOrGenerateWasmVar(ofType: valueType)
+
+        let (lhs, staticOffset) = b.generateAlignedMemoryIndexes(forMemory: memory, alignment: alignment)
+
+        function.wasmAtomicRMW(memory: memory, lhs: lhs, rhs: rhs, op: op, offset: staticOffset)
+    },
+
     CodeGenerator("WasmMemorySizeGenerator", inContext: .wasmFunction, inputs: .required(.object(ofGroup: "WasmMemory"))) { b, memory in
         let function = b.currentWasmModule.currentWasmFunction
         function.wasmMemorySize(memory: memory)

Sources/Fuzzilli/FuzzIL/Instruction.swift

@@ -1305,6 +1305,11 @@ extension Instruction: ProtobufConvertible {
                     $0.storeType = convertEnum(op.storeType, WasmAtomicStoreType.allCases)
                     $0.offset = op.offset
                 }
+            case .wasmAtomicRMW(let op):
+                $0.wasmAtomicRmw = Fuzzilli_Protobuf_WasmAtomicRMW.with {
+                    $0.op = convertEnum(op.op, WasmAtomicRMWType.allCases)
+                    $0.offset = op.offset
+                }
             case .wasmMemorySize(_):
                 $0.wasmMemorySize = Fuzzilli_Protobuf_WasmMemorySize()
             case .wasmMemoryGrow(_):
@@ -2499,6 +2504,8 @@ extension Instruction: ProtobufConvertible {
             op = WasmAtomicLoad(loadType: try convertEnum(p.loadType, WasmAtomicLoadType.allCases), offset: p.offset)
         case .wasmAtomicStore(let p):
             op = WasmAtomicStore(storeType: try convertEnum(p.storeType, WasmAtomicStoreType.allCases), offset: p.offset)
+        case .wasmAtomicRmw(let p):
+            op = WasmAtomicRMW(op: try convertEnum(p.op, WasmAtomicRMWType.allCases), offset: p.offset)
         case .wasmAnyConvertExtern(_):
             op = WasmAnyConvertExtern()
         case .wasmExternConvertAny(_):

Sources/Fuzzilli/FuzzIL/JSTyper.swift

@@ -694,6 +694,10 @@ public struct JSTyper: Analyzer {
             case .wasmAtomicStore(_):
                 let definingInstruction = defUseAnalyzer.definition(of: instr.input(0))
                 dynamicObjectGroupManager.addWasmMemory(withType: type(of: instr.input(0)), forDefinition: definingInstruction, forVariable: instr.input(0))
+            case .wasmAtomicRMW(let op):
+                let definingInstruction = defUseAnalyzer.definition(of: instr.input(0))
+                dynamicObjectGroupManager.addWasmMemory(withType: type(of: instr.input(0)), forDefinition: definingInstruction, forVariable: instr.input(0))
+                setType(of: instr.output, to: op.op.type)
             case .wasmMemorySize(_),
                  .wasmMemoryGrow(_):
                 let isMemory64 = type(of: instr.input(0)).wasmMemoryType?.isMemory64 ?? false

Sources/Fuzzilli/FuzzIL/Opcodes.swift

@@ -285,6 +285,7 @@ enum Opcode {
     case wasmMemoryStore(WasmMemoryStore)
     case wasmAtomicLoad(WasmAtomicLoad)
     case wasmAtomicStore(WasmAtomicStore)
+    case wasmAtomicRMW(WasmAtomicRMW)
     case wasmMemorySize(WasmMemorySize)
     case wasmMemoryGrow(WasmMemoryGrow)
     case wasmMemoryFill(WasmMemoryFill)

Sources/Fuzzilli/FuzzIL/WasmOperations.swift

@@ -254,6 +254,105 @@ final class Wasmi64BinOp: WasmOperation {
     }
 }
 
+public enum WasmAtomicRMWType: UInt8, CaseIterable {
+    case i32Add = 0x1e
+    case i64Add = 0x1f
+    case i32Add8U = 0x20
+    case i32Add16U = 0x21
+    case i64Add8U = 0x22
+    case i64Add16U = 0x23
+    case i64Add32U = 0x24
+
+    case i32Sub = 0x25
+    case i64Sub = 0x26
+    case i32Sub8U = 0x27
+    case i32Sub16U = 0x28
+    case i64Sub8U = 0x29
+    case i64Sub16U = 0x2a
+    case i64Sub32U = 0x2b
+
+    case i32And = 0x2c
+    case i64And = 0x2d
+    case i32And8U = 0x2e
+    case i32And16U = 0x2f
+    case i64And8U = 0x30
+    case i64And16U = 0x31
+    case i64And32U = 0x32
+
+    case i32Or = 0x33
+    case i64Or = 0x34
+    case i32Or8U = 0x35
+    case i32Or16U = 0x36
+    case i64Or8U = 0x37
+    case i64Or16U = 0x38
+    case i64Or32U = 0x39
+
+    case i32Xor = 0x3a
+    case i64Xor = 0x3b
+    case i32Xor8U = 0x3c
+    case i32Xor16U = 0x3d
+    case i64Xor8U = 0x3e
+    case i64Xor16U = 0x3f
+    case i64Xor32U = 0x40
+
+    case i32Xchg = 0x41
+    case i64Xchg = 0x42
+    case i32Xchg8U = 0x43
+    case i32Xchg16U = 0x44
+    case i64Xchg8U = 0x45
+    case i64Xchg16U = 0x46
+    case i64Xchg32U = 0x47
+
+    var type: ILType {
+        switch self {
+        case .i32Add, .i32Add8U, .i32Add16U,
+             .i32Sub, .i32Sub8U, .i32Sub16U,
+             .i32And, .i32And8U, .i32And16U,
+             .i32Or, .i32Or8U, .i32Or16U,
+             .i32Xor, .i32Xor8U, .i32Xor16U,
+             .i32Xchg, .i32Xchg8U, .i32Xchg16U:
+            return .wasmi32
+        case .i64Add, .i64Add8U, .i64Add16U, .i64Add32U,
+             .i64Sub, .i64Sub8U, .i64Sub16U, .i64Sub32U,
+             .i64And, .i64And8U, .i64And16U, .i64And32U,
+             .i64Or, .i64Or8U, .i64Or16U, .i64Or32U,
+             .i64Xor, .i64Xor8U, .i64Xor16U, .i64Xor32U,
+             .i64Xchg, .i64Xchg8U, .i64Xchg16U, .i64Xchg32U:
+            return .wasmi64
+        }
+    }
+
+    func naturalAlignment() -> Int64 {
+        switch self {
+        case .i32Add8U, .i64Add8U,
+             .i32Sub8U, .i64Sub8U,
+             .i32And8U, .i64And8U,
+             .i32Or8U, .i64Or8U,
+             .i32Xor8U, .i64Xor8U,
+             .i32Xchg8U, .i64Xchg8U:
+            return 1
+        case .i32Add16U, .i64Add16U,
+             .i32Sub16U, .i64Sub16U,
+             .i32And16U, .i64And16U,
+             .i32Or16U, .i64Or16U,
+             .i32Xor16U, .i64Xor16U,
+             .i32Xchg16U, .i64Xchg16U:
+            return 2
+        case .i32Add, .i64Add32U,
+             .i32Sub, .i64Sub32U,
+             .i32And, .i64And32U,
+             .i32Or, .i64Or32U,
+             .i32Xor, .i64Xor32U,
+             .i32Xchg, .i64Xchg32U:
+            return 4
+        case .i64Add, .i64Sub,
+             .i64And, .i64Or,
+             .i64Xor, .i64Xchg:
+            return 8
+        }
+    }
+}
+
 final class Wasmi32UnOp: WasmOperation {
     override var opcode: Opcode { .wasmi32UnOp(self) }
     let unOpKind: WasmIntegerUnaryOpKind
@@ -2144,3 +2243,16 @@ final class WasmAtomicStore: WasmOperation {
         super.init(numInputs: 3, numOutputs: 0, attributes: [.isMutable], requiredContext: [.wasmFunction])
     }
 }
+
+final class WasmAtomicRMW: WasmOperation {
+    override var opcode: Opcode { .wasmAtomicRMW(self) }
+
+    let op: WasmAtomicRMWType
+    let offset: Int64
+
+    init(op: WasmAtomicRMWType, offset: Int64) {
+        self.op = op
+        self.offset = offset
+        super.init(numInputs: 3, numOutputs: 1, attributes: [.isMutable], requiredContext: [.wasmFunction])
+    }
+}

Sources/Fuzzilli/Lifting/FuzzILLifter.swift

@@ -853,6 +853,9 @@ public class FuzzILLifter: Lifter {
         case .wasmAtomicStore(let op):
             w.emit("WasmAtomicStore \(input(0))[\(input(1)) + \(op.offset)] <- \(input(2)) [\(op.storeType)]")
 
+        case .wasmAtomicRMW(let op):
+            w.emit("\(output()) <- WasmAtomicRMW \(input(0))[\(input(1)) + \(op.offset)] \(op.op) \(input(2))")
+
         case .wasmMemorySize(_):
             w.emit("\(output()) <- WasmMemorySize \(input(0))")
 

Sources/Fuzzilli/Lifting/JavaScriptLifter.swift

@@ -1637,6 +1637,7 @@ public class JavaScriptLifter: Lifter {
                  .wasmMemoryStore(_),
                  .wasmAtomicLoad(_),
                  .wasmAtomicStore(_),
+                 .wasmAtomicRMW(_),
                  .wasmMemorySize(_),
                  .wasmMemoryGrow(_),
                  .wasmMemoryFill(_),

Sources/Fuzzilli/Lifting/WasmLifter.swift

@@ -1754,6 +1754,10 @@ public class WasmLifter {
             let opcode = [Prefix.Atomic.rawValue, op.storeType.rawValue]
             let alignAndMemory = try alignmentAndMemoryBytes(wasmInstruction.input(0), alignment: op.storeType.naturalAlignment())
             return Data(opcode) + alignAndMemory + Leb128.signedEncode(Int(op.offset))
+        case .wasmAtomicRMW(let op):
+            let opcode = [Prefix.Atomic.rawValue, op.op.rawValue]
+            let alignAndMemory = try alignmentAndMemoryBytes(wasmInstruction.input(0), alignment: op.op.naturalAlignment())
+            return Data(opcode) + alignAndMemory + Leb128.signedEncode(Int(op.offset))
         case .wasmMemorySize(_):
             let memoryIdx = try resolveIdx(ofType: .memory, for: wasmInstruction.input(0))
             return Data([0x3F]) + Leb128.unsignedEncode(memoryIdx)