Skip to content

Commit

Permalink
optimise memory usage during replay generation
Browse files Browse the repository at this point in the history 
fix runtime errors due to memory table or global wrongly imported
  • Loading branch information
@jakobgetz
jakobgetz committed Nov 24, 2023
1 parent 4fd69b0 commit 5a935e0
Show file tree
Hide file tree
Showing 9 changed files with 129 additions and 100 deletions.
2 changes: 1 addition & 1 deletion src/benchmark.cts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ export default class Benchmark {
// console.log('wrote temp trace to disk and start stream code generation')
const code = await new Generator().generateReplayFromStream(fss.createReadStream(diskSave))
// console.log('stream code generation finished. Now stream js string to file')
code.toWriteStream(fss.createWriteStream(path.join(binPath, 'replay.js')))
await code.toWriteStream(fss.createWriteStream(path.join(binPath, 'replay.js')))
// const jsString = new Generator().generateReplay(trace).toString()
// console.log('js code generation finished. Now dump wasm to file')
// await fs.writeFile(path.join(binPath, 'replay.js'), jsString)
Expand Down
114 changes: 60 additions & 54 deletions src/replay-generator.cts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,11 @@ type TableGrow = { type: "TableGrow", idx: number, amount: number } & ImpExp
type GlobalSet = { type: "GlobalSet", value: number, bigInt: boolean } & ImpExp
type Event = Call | Store | MemGrow | TableSet | TableGrow | GlobalSet
type Import = { module: string, name: string }
type Function = Import & { body: { results: number[], events: Event[] }[] }
type Memory = Import & { pages: number, maxPages: number }
type BodyPart = { results: number[], events: Event[] }
type Function = Import & { body: BodyPart[] }
type Memory = Import & WebAssembly.MemoryDescriptor
type Table = Import & WebAssembly.TableDescriptor
type Global = Import & { valtype: string, value: number }
type Global = Import & WebAssembly.GlobalDescriptor & { initial: number }
type State = { callStack: Function[], lastFunc?: Function }

export default class Generator {
Expand Down Expand Up @@ -64,17 +65,11 @@ export default class Generator {
this.pushEvent({ type: 'Call', name: event.name, params: event.params })
break
case "ImportCall":
if (this.code.funcImports[event.idx] === undefined) {
console.log('body')
}
this.code.funcImports[event.idx].body.push({ results: [], events: [] })
this.state.callStack.push(this.code.funcImports[event.idx])
break
case "ImportReturn":
this.state.lastFunc = this.state.callStack.pop()
if (this.state.lastFunc === undefined) {
console.log('yo')
}
this.state.lastFunc.body.slice(-1)[0].results = event.results
break
case "Load":
Expand Down Expand Up @@ -118,8 +113,8 @@ export default class Generator {
this.code.memImports[event.idx] = {
module: event.module,
name: event.name,
pages: event.pages,
maxPages: event.maxPages
initial: event.initial,
maximum: event.maximum
}
break
case 'GlobalGet':
Expand All @@ -146,8 +141,9 @@ export default class Generator {
this.code.globalImports[event.idx] = {
module: event.module,
name: event.name,
valtype: event.valtype,
value: event.value,
initial: event.initial,
mutable: event.mutable
}
break
case 'ImportFunc':
Expand Down Expand Up @@ -220,20 +216,20 @@ class Code {
// Init memories
for (let memidx in this.memImports) {
let mem = this.memImports[memidx]
jsString += `const ${mem.name} = new WebAssembly.Memory({ initial: ${mem.pages}, maximum: ${mem.maxPages} })\n`
jsString += `const ${mem.name} = new WebAssembly.Memory({ initial: ${mem.initial}, maximum: ${mem.maximum} })\n`
jsString += `${writeImport(mem.module, mem.name)}${mem.name}\n`
}
// Init globals
for (let globalIdx in this.globalImports) {
let global = this.globalImports[globalIdx]
jsString += `const ${global.name} = new WebAssembly.Global({ value: '${global.valtype}', mutable: true}, ${global.value})\n`
jsString += `${global.name}.value = ${global.value}\n`
jsString += `const ${global.name} = new WebAssembly.Global({ value: '${global.value}', mutable: ${global.mutable}}, ${global.initial})\n`
// jsString += `${global.name}.value = ${global.initial}\n`
jsString += `${writeImport(global.module, global.name)}${global.name}\n`
}
// Init tables
for (let tableidx in this.tableImports) {
let table = this.tableImports[tableidx]
jsString += `const ${table.name} = new WebAssembly.Table({ initial: ${table.initial}, element: '${table.element}'})\n`
jsString += `const ${table.name} = new WebAssembly.Table({ initial: ${table.initial}, maximum: ${table.maximum}, element: '${table.element}'})\n`
jsString += `${writeImport(table.module, table.name)}${table.name}\n`
}
// Init entity states
Expand Down Expand Up @@ -304,7 +300,7 @@ class Code {
return jsString
}

toWriteStream(stream: WriteStream) {
async toWriteStream(stream: WriteStream) {
stream.write(`import fs from 'fs'\n`)
stream.write(`import path from 'path'\n`)
stream.write(`export default async function replay(wasmBinary) {\n`)
Expand All @@ -318,20 +314,20 @@ class Code {
// Init memories
for (let memidx in this.memImports) {
let mem = this.memImports[memidx]
stream.write(`const ${mem.name} = new WebAssembly.Memory({ initial: ${mem.pages}, maximum: ${mem.maxPages} })\n`)
stream.write(`const ${mem.name} = new WebAssembly.Memory({ initial: ${mem.initial}, maximum: ${mem.maximum} })\n`)
stream.write(`${writeImport(mem.module, mem.name)}${mem.name}\n`)
}
// Init globals
for (let globalIdx in this.globalImports) {
let global = this.globalImports[globalIdx]
stream.write(`const ${global.name} = new WebAssembly.Global({ value: '${global.valtype}', mutable: true}, ${global.value})\n`)
stream.write(`${global.name}.value = ${global.value}\n`)
stream.write(`const ${global.name} = new WebAssembly.Global({ value: '${global.value}', mutable: ${global.mutable}}, ${global.initial})\n`)
// stream.write(`${global.name}.value = ${global.initial}\n`)
stream.write(`${writeImport(global.module, global.name)}${global.name}\n`)
}
// Init tables
for (let tableidx in this.tableImports) {
let table = this.tableImports[tableidx]
stream.write(`const ${table.name} = new WebAssembly.Table({ initial: ${table.initial}, element: '${table.element}'})\n`)
stream.write(`const ${table.name} = new WebAssembly.Table({ initial: ${table.initial}, maximum: ${table.maximum}, element: '${table.element}'})\n`)
stream.write(`${writeImport(table.module, table.name)}${table.name}\n`)
}
// Init entity states
Expand All @@ -358,39 +354,9 @@ class Code {
stream.write(`${writeImport(func.module, func.name)}() => {\n`)
stream.write(`${writeFuncGlobal(funcidx)}++\n`)
stream.write(`switch (${writeFuncGlobal(funcidx)}) {\n`)
func.body.forEach((b, i) => {
if (b.events.length !== 0 || b.results.length !== 0) {
stream.write(`case ${i}:\n`)
}
if (b.events.length !== 0) {
for (let event of b.events) {
switch (event.type) {
case 'Call':
stream.write(`instance.exports.${event.name}(${writeParamsString(event.params)})\n`)
break
case 'Store':
stream.write(this.storeEvent(event))
break
case 'MemGrow':
stream.write(this.memGrowEvent(event))
break
case 'TableSet':
stream.write(this.tableSetEvent(event))
break
case 'TableGrow':
stream.write(this.tableGrowEvent(event))
break
case 'GlobalSet':
stream.write(this.globalSet(event))
break
default: unreachable(event)
}
}
}
if (b.results.length !== 0) {
stream.write(`return ${b.results[0]} \n`)
}
})
for (let i = 0; i < func.body.length; i++) {
await this.writeBody(stream, func.body[i], i)
}
stream.write('}\n')
stream.write('}\n')
}
Expand All @@ -408,6 +374,46 @@ class Code {
stream.close()
}

private async writeBody(stream: WriteStream, b: BodyPart, i: number) {
if (b.events.length !== 0 || b.results.length !== 0) {
await this.write(stream, `case ${i}:\n`)
}
if (b.events.length !== 0) {
for (let event of b.events) {
switch (event.type) {
case 'Call':
stream.write(`instance.exports.${event.name}(${writeParamsString(event.params)})\n`)
break
case 'Store':
stream.write(this.storeEvent(event))
break
case 'MemGrow':
stream.write(this.memGrowEvent(event))
break
case 'TableSet':
stream.write(this.tableSetEvent(event))
break
case 'TableGrow':
stream.write(this.tableGrowEvent(event))
break
case 'GlobalSet':
stream.write(this.globalSet(event))
break
default: unreachable(event)
}
}
}
if (b.results.length !== 0) {
await this.write(stream, `return ${b.results[0]} \n`)
}
}

private async write(stream: WriteStream, s: string) {
if (stream.write(s) === false) {
await new Promise((resolve) => stream.once('drain', resolve))
}
}

private storeEvent(event: Store) {
let jsString = ''
event.data.forEach((byte, j) => {
Expand Down
51 changes: 27 additions & 24 deletions src/tracer.cts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,13 +71,13 @@ export class Trace {
trace.push({type: 'ImportReturn', idx: e[1], name: e[2], results: e[3]})
break
case 'IM':
trace.push({type: 'ImportMemory', idx: e[1], module: e[2], name: e[3], pages: e[4], maxPages: e[5]})
trace.push({type: 'ImportMemory', idx: e[1], module: e[2], name: e[3], initial: e[4], maximum: e[5]})
break
case 'IT':
trace.push({type: 'ImportTable', idx: e[1], module: e[2], name: e[3], initial: e[4], element: e[5]})
trace.push({type: 'ImportTable', idx: e[1], module: e[2], name: e[3], initial: e[4], maximum: e[5], element: e[6]})
break
case 'IG':
trace.push({type: 'ImportGlobal', idx: e[1], module: e[2], name: e[3], valtype: e[4], value: e[5]})
trace.push({type: 'ImportGlobal', idx: e[1], module: e[2], name: e[3], value: e[4], mutable: e[5] === 1, initial: e[6]})
break
case 'IF':
trace.push({type: 'ImportFunc', idx: e[1], module: e[2], name: e[3]})
Expand All @@ -101,7 +101,9 @@ export class Trace {
}
})
} else {
eventString += value
if (value !== undefined && value !== null) {
eventString += value
}
}
if (i < event.length - 1) {
eventString += ';'
Expand Down Expand Up @@ -137,7 +139,7 @@ export class Trace {
components[2],
components[3],
parseInt(components[4]),
parseInt(components[5])
components[5] === '' ? undefined : parseInt(components[5])
]
case "EC":
return [
Expand Down Expand Up @@ -199,12 +201,13 @@ export class Trace {
]
case 'IG':
return [
components[0],
parseInt(components[1]),
components[2],
components[3],
components[0],
parseInt(components[1]),
components[2],
components[3],
components[4] as ValType,
parseInt(components[5]),
parseInt(components[5]) as 0 | 1,
parseInt(components[6]),
]
case 'IF':
return [
Expand All @@ -215,12 +218,13 @@ export class Trace {
]
case 'IT':
return [
components[0],
components[0],
parseInt(components[1]),
components[2],
components[3],
parseInt(components[4]),
components[5] as 'anyfunc'
components[5] === '' ? undefined : parseInt(components[5]),
components[6] as 'anyfunc'
]
default:
throw new Error(`${components[0]}: Not a valid trace event. The whole event: ${event}.`)
Expand Down Expand Up @@ -253,8 +257,8 @@ export class Trace {
idx: parseInt(components[1]),
module: components[2],
name: components[3],
pages: parseInt(components[4]),
maxPages: parseInt(components[5])
initial: parseInt(components[4]),
maximum: components[5] === '' ? undefined : parseInt(components[5])
}
case "EC":
return {
Expand Down Expand Up @@ -320,8 +324,9 @@ export class Trace {
idx: parseInt(components[1]),
module: components[2],
name: components[3],
valtype: components[4] as ValType,
value: parseInt(components[5]),
initial: parseInt(components[6]),
value: components[4] as ValType,
mutable: parseInt(components[5]) === 1
}
case 'IF':
return {
Expand All @@ -337,7 +342,8 @@ export class Trace {
module: components[2],
name: components[3],
initial: parseInt(components[4]),
element: components[5] as 'anyfunc'
maximum: components[5] === '' ? undefined : parseInt(components[5]),
element: components[6] as 'anyfunc'
}
default:
throw new Error(`${components[0]}: Not a valid trace event. The whole event: ${event}.`)
Expand Down Expand Up @@ -725,29 +731,26 @@ export default class Analysis implements AnalysisI<Trace> {
})
this.Wasabi.module.info.memories.forEach((m, idx) => {
if (m.import !== null) {
const memory = this.Wasabi.module.memories[idx]
const pages = memory.buffer.byteLength / this.MEM_PAGE_SIZE
const maxPages = memory.buffer.byteLength / (64 * 1024)
this.trace.push(['IM', idx, m.import[0], m.import[1], pages, maxPages ])
this.trace.push(['IM', idx, m.import[0], m.import[1], m.initial, m.maximum ])
}
})
// Init Tables
this.Wasabi.module.tables.forEach((t, i) => {
this.shadowTables.push(new WebAssembly.Table({ initial: this.Wasabi.module.tables[i].length, element: 'anyfunc' }))
this.shadowTables.push(new WebAssembly.Table({ initial: this.Wasabi.module.tables[i].length, element: 'anyfunc' })) // want to replace anyfunc through t.refType but it holds the wrong string ('funcref')
for (let y = 0; y < this.Wasabi.module.tables[i].length; y++) {
this.shadowTables[i].set(y, t.get(y))
}
})
this.Wasabi.module.info.tables.forEach((t, idx) => {
if (t.import !== null) {
this.trace.push(['IT', idx, t.import![0], t.import![1], this.Wasabi.module.tables[idx].length , 'anyfunc'])
this.trace.push(['IT', idx, t.import![0], t.import![1], t.initial, t.maximum, 'anyfunc']) // want to replace anyfunc through t.refType but it holds the wrong string ('funcref')
}
})
// Init Globals
this.shadowGlobals = this.Wasabi.module.globals.map(g => g.value)
this.Wasabi.module.info.globals.forEach((g, idx) => {
if (g.import !== null) {
this.trace.push([ 'IG', idx, g.import[0], g.import[1], g.valType, this.Wasabi.module.globals[idx].value])
this.trace.push([ 'IG', idx, g.import[0], g.import[1], g.valType, g.mutability === 'Mut' ? 1 : 0, this.Wasabi.module.globals[idx].value])
}
})
// Init Functions
Expand Down
7 changes: 7 additions & 0 deletions tests/node/glob-imp-const/index.wat
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
(module
(import "env" "global" (global $global i32))
(func $main (export "main")
global.get $global
drop
)
)
12 changes: 12 additions & 0 deletions tests/node/glob-imp-const/test.js
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
export default async function test(wasmBinary) {
let instance
const global = new WebAssembly.Global({ value: "i32", mutable: false }, 4);
let imports = {
env: {
global: global
}
}
let wasm = await WebAssembly.instantiate(wasmBinary, imports)
instance = wasm.instance
instance.exports.main()
}
Loading

0 comments on commit 5a935e0

Please sign in to comment.