Skip to content

Commit 58e5fe6

Browse files
committed
[WEB] Initial support for asyncify
This PR enables asyncify support for web runtime. Asyncify is a feature to allow C++ to call async function in javascript. The emcc compiler will unwind and store the stack, returning control to JS runtime. The JS runtime needs to be able to await the promise and then call rewind to get to the original suspended point. This feature can be potentially useful when we would like to call WebGPU sync in C++ runtime. As on web platform everything have to be non-blocking. Because asyncify can increase the wams size by 2x, we don't enable it by default in emcc.py and still would need to pass in options. We will confirm potential benefit tradeoffs before turning it on by default. Another catch is that as of now asyncify is not compatible with wasm exception, so we temporary turn wasm-exception it off for now. This is an item that is being worked on by emscripten so we might be able to turn it back on later. The testcases are added. reference: https://emscripten.org/docs/porting/asyncify.html
1 parent 40dd376 commit 58e5fe6

File tree

12 files changed

+396
-26
lines changed

12 files changed

+396
-26
lines changed

python/tvm/contrib/emcc.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,14 @@ def create_tvmjs_wasm(output, objects, options=None, cc="emcc"):
4242
cmd += ["-O3"]
4343
cmd += ["-std=c++17"]
4444
cmd += ["--no-entry"]
45-
cmd += ["-fwasm-exceptions"]
45+
# NOTE: asynctify conflicts with wasm-exception
46+
# so we temp disable exception handling for now
47+
#
48+
# We also expect user to explicitly pass in
49+
# -s ASYNCIFY=1 as it can increase wasm size by 2xq
50+
#
51+
# cmd += ["-s", "ASYNCIFY=1"]
52+
# cmd += ["-fwasm-exceptions"]
4653
cmd += ["-s", "WASM_BIGINT=1"]
4754
cmd += ["-s", "ERROR_ON_UNDEFINED_SYMBOLS=0"]
4855
cmd += ["-s", "STANDALONE_WASM=1"]

src/runtime/c_runtime_api.cc

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -569,7 +569,6 @@ int TVMByteArrayFree(TVMByteArray* arr) {
569569
int TVMFuncCall(TVMFunctionHandle func, TVMValue* args, int* arg_type_codes, int num_args,
570570
TVMValue* ret_val, int* ret_type_code) {
571571
API_BEGIN();
572-
573572
TVMRetValue rv;
574573
(static_cast<const PackedFuncObj*>(func))
575574
->CallPacked(TVMArgs(args, arg_type_codes, num_args), &rv);

web/Makefile

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,10 +27,11 @@ all: dist/wasm/tvmjs_runtime.wasm dist/wasm/tvmjs_runtime.wasi.js src/tvmjs_runt
2727

2828
EMCC = emcc
2929

30-
EMCC_CFLAGS = $(INCLUDE_FLAGS) -O3 -std=c++17 -Wno-ignored-attributes -fwasm-exceptions
30+
EMCC_CFLAGS = $(INCLUDE_FLAGS) -O3 -std=c++17 -Wno-ignored-attributes
3131

3232
EMCC_LDFLAGS = --no-entry -s WASM_BIGINT=1 -s ALLOW_MEMORY_GROWTH=1 -s STANDALONE_WASM=1\
33-
-s ERROR_ON_UNDEFINED_SYMBOLS=0 --pre-js emcc/preload.js
33+
-s ERROR_ON_UNDEFINED_SYMBOLS=0 --pre-js emcc/preload.js\
34+
-s ASYNCIFY=1
3435

3536
dist/wasm/%.bc: emcc/%.cc
3637
@mkdir -p $(@D)

web/apps/node/example.js

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
*/
2222
const path = require("path");
2323
const fs = require("fs");
24-
const tvmjs = require("../../lib");
24+
const tvmjs = require("../../dist/tvmjs.bundle");
2525

2626
const wasmPath = tvmjs.wasmPath();
2727
const wasmSource = fs.readFileSync(path.join(wasmPath, "tvmjs_runtime.wasm"));

web/emcc/decorate_as_wasi.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
template_head = """
2222
function EmccWASI() {
23+
var asyncifyStubs = {};
2324
"""
2425

2526
template_tail = """

web/emcc/wasm_runtime.cc

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,12 @@ TVM_REGISTER_GLOBAL("testing.echo").set_body([](TVMArgs args, TVMRetValue* ret)
100100
*ret = args[0];
101101
});
102102

103+
TVM_REGISTER_GLOBAL("testing.call").set_body([](TVMArgs args, TVMRetValue* ret) {
104+
(args[0].operator PackedFunc()).CallPacked(
105+
TVMArgs(args.values + 1, args.type_codes + 1, args.num_args - 1), ret
106+
);
107+
});
108+
103109
TVM_REGISTER_GLOBAL("testing.ret_string").set_body([](TVMArgs args, TVMRetValue* ret) {
104110
*ret = args[0].operator String();
105111
});

web/emcc/webgpu_runtime.cc

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,11 @@ class WebGPUDeviceAPI : public DeviceAPI {
112112
LOG(FATAL) << "Not implemented";
113113
}
114114

115-
void StreamSync(Device dev, TVMStreamHandle stream) final { LOG(FATAL) << "Not implemented"; }
115+
void StreamSync(Device dev, TVMStreamHandle stream) final {
116+
static const PackedFunc* func = runtime::Registry::Get("__asyncify.WebGPUWaitForTasks");
117+
ICHECK(func != nullptr) << "Stream sync inside c++ only supported in asyncify mode";
118+
(*func)();
119+
}
116120

117121
void SetStream(Device dev, TVMStreamHandle stream) final { LOG(FATAL) << "Not implemented"; }
118122

web/src/artifact_cache.ts

Lines changed: 32 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,37 @@
11
/*
2-
Common Interface for the artifact cache
3-
*/
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
/**
20+
* Common Interface for the artifact cache
21+
*/
422
export interface ArtifactCacheTemplate {
5-
/**
6-
* fetch key url from cache
7-
*/
8-
fetchWithCache(url: string);
23+
/**
24+
* fetch key url from cache
25+
*/
26+
fetchWithCache(url: string);
927

10-
/**
11-
* check if cache has all keys in Cache
12-
*/
13-
hasAllKeys(keys: string[]);
28+
/**
29+
* check if cache has all keys in Cache
30+
*/
31+
hasAllKeys(keys: string[]);
1432

15-
/**
16-
* Delete url in cache if url exists
17-
*/
18-
deleteInCache(url: string);
33+
/**
34+
* Delete url in cache if url exists
35+
*/
36+
deleteInCache(url: string);
1937
}

web/src/asyncify.ts

Lines changed: 227 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,227 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
// Helper tools to enable asynctify handling
20+
// Thie following code is used to support wrapping of
21+
// functins that can have async await calls in the backend runtime
22+
// reference
23+
// - https://kripken.github.io/blog/wasm/2019/07/16/asyncify.html
24+
// - https://github.com/GoogleChromeLabs/asyncify
25+
import { assert, isPromise } from "./support";
26+
27+
/**
28+
* enums to check the current state of asynctify
29+
*/
30+
const enum AsyncifyStateKind {
31+
None = 0,
32+
Unwinding = 1,
33+
Rewinding = 2
34+
}
35+
36+
/** The start location of asynctify stack data */
37+
const ASYNCIFY_DATA_ADDR = 16;
38+
/** The data start of stack rewind/unwind */
39+
const ASYNCIFY_DATA_START = ASYNCIFY_DATA_ADDR + 8;
40+
/** The data end of stack rewind/unwind */
41+
const ASYNCIFY_DATA_END = 1024;
42+
43+
/** Hold asynctify handler instance that runtime can use */
44+
export class AsyncifyHandler {
45+
/** exports from wasm */
46+
private exports: Record<string, Function>;
47+
/** current state kind */
48+
private state: AsyncifyStateKind = AsyncifyStateKind.None;
49+
/** The stored value before unwind */
50+
private storedPromiseBeforeUnwind : Promise<any> = null;
51+
// NOTE: asynctify do not work with exceptions
52+
// this implementation here is mainly for possible future compact
53+
/** The stored value that is resolved */
54+
private storedValueBeforeRewind: any = null;
55+
/** The stored exception */
56+
private storedExceptionBeforeRewind: any = null;
57+
58+
constructor(exports: Record<string, Function>, memory: WebAssembly.Memory) {
59+
this.exports = exports;
60+
this.initMemory(memory);
61+
}
62+
63+
// NOTE: wrapImport and wrapExport are closely related to each other
64+
// We mark the logical jump pt in comments to increase the readability
65+
/**
66+
* Whether the wasm enables asynctify
67+
* @returns Whether the wasm enables asynctify
68+
*/
69+
enabled(): boolean {
70+
return this.exports.asyncify_stop_rewind !== undefined;
71+
}
72+
73+
/**
74+
* Get the current asynctify state
75+
*
76+
* @returns The current asynctify state
77+
*/
78+
getState(): AsyncifyStateKind {
79+
return this.state;
80+
}
81+
82+
/**
83+
* Wrap a function that can be used as import of the wasm asynctify layer
84+
*
85+
* @param func The input import function
86+
* @returns The wrapped function that can be registered to the system
87+
*/
88+
wrapImport(func: (...args: Array<any>) => any): (...args: Array<any>) => any {
89+
return (...args: any) => {
90+
// this is being called second time
91+
// where we are rewinding the stack
92+
if (this.getState() == AsyncifyStateKind.Rewinding) {
93+
// JUMP-PT-REWIND: rewind will jump to this pt
94+
// while rewinding the stack
95+
this.stopRewind();
96+
// the value has been resolved
97+
if (this.storedValueBeforeRewind !== null) {
98+
assert(this.storedExceptionBeforeRewind === null);
99+
const result = this.storedValueBeforeRewind;
100+
this.storedValueBeforeRewind = null;
101+
return result;
102+
} else {
103+
assert(this.storedValueBeforeRewind === null);
104+
const error = this.storedExceptionBeforeRewind;
105+
this.storedExceptionBeforeRewind = null;
106+
throw error;
107+
}
108+
}
109+
// this function is being called for the first time
110+
assert(this.getState() == AsyncifyStateKind.None);
111+
112+
// call the function
113+
const value = func(...args);
114+
// if the value is promise
115+
// we need to unwind the stack
116+
// so the caller will be able to evaluate the promise
117+
if (isPromise(value)) {
118+
// The next code step is JUMP-PT-UNWIND in wrapExport
119+
// The value will be passed to that pt through storedPromiseBeforeUnwind
120+
// getState() == Unwinding and we will enter the while loop in wrapExport
121+
this.startUnwind();
122+
assert(this.storedPromiseBeforeUnwind == null);
123+
this.storedPromiseBeforeUnwind = value;
124+
return undefined;
125+
} else {
126+
// The next code step is JUMP-PT-UNWIND in wrapExport
127+
// normal value, we don't have to do anything
128+
// getState() == None and we will exit while loop there
129+
return value;
130+
}
131+
};
132+
}
133+
134+
/**
135+
* Warp an exported asynctify function so it can return promise
136+
*
137+
* @param func The input function
138+
* @returns The wrapped async function
139+
*/
140+
wrapExport(func: (...args: Array<any>) => any): (...args: Array<any>) => Promise<any> {
141+
return async (...args: Array<any>) => {
142+
assert(this.getState() == AsyncifyStateKind.None);
143+
144+
// call the original function
145+
let result = func(...args);
146+
147+
// JUMP-PT-UNWIND
148+
// after calling the function
149+
// the caller may hit a unwinding point depending on
150+
// the if (isPromise(value)) condition in wrapImport
151+
while (this.getState() == AsyncifyStateKind.Unwinding) {
152+
this.stopUnwind();
153+
// try to resolve the promise that the internal requested
154+
// we then store it into the temp value in storedValueBeforeRewind
155+
// which then get passed onto the function(see wrapImport)
156+
// that can return the value
157+
const storedPromiseBeforeUnwind = this.storedPromiseBeforeUnwind;
158+
this.storedPromiseBeforeUnwind = null;
159+
assert(this.storedExceptionBeforeRewind === null);
160+
assert(this.storedValueBeforeRewind == null);
161+
162+
try {
163+
this.storedValueBeforeRewind = await storedPromiseBeforeUnwind;
164+
} catch (error) {
165+
// the store exception
166+
this.storedExceptionBeforeRewind = error;
167+
}
168+
assert(!isPromise(this.storedValueBeforeRewind));
169+
// because we called asynctify_stop_unwind,the state is now none
170+
assert(this.getState() == AsyncifyStateKind.None);
171+
172+
// re-enter the function, jump to JUMP-PT-REWIND in wrapImport
173+
// the value will be passed to that point via storedValueBeforeRewind
174+
//
175+
// NOTE: we guarantee that if exception is throw the asynctify state
176+
// will already be at None, this is because we will goto JUMP-PT-REWIND
177+
// which will call aynctify_stop_rewind
178+
this.startRewind();
179+
result = func(...args);
180+
}
181+
return result;
182+
};
183+
}
184+
185+
private startRewind() : void {
186+
if (this.exports.asyncify_start_rewind === undefined) {
187+
throw Error("Asynctify is not enabled, please compile with -s ASYNCIFY=1 in emcc");
188+
}
189+
this.exports.asyncify_start_rewind(ASYNCIFY_DATA_ADDR);
190+
this.state = AsyncifyStateKind.Rewinding;
191+
}
192+
193+
private stopRewind() : void {
194+
if (this.exports.asyncify_stop_rewind === undefined) {
195+
throw Error("Asynctify is not enabled, please compile with -s ASYNCIFY=1 in emcc");
196+
}
197+
this.exports.asyncify_stop_rewind();
198+
this.state = AsyncifyStateKind.None;
199+
}
200+
201+
private startUnwind() : void {
202+
if (this.exports.asyncify_start_unwind === undefined) {
203+
throw Error("Asynctify is not enabled, please compile with -s ASYNCIFY=1 in emcc");
204+
}
205+
this.exports.asyncify_start_unwind(ASYNCIFY_DATA_ADDR);
206+
this.state = AsyncifyStateKind.Unwinding;
207+
}
208+
209+
private stopUnwind() : void {
210+
if (this.exports.asyncify_stop_unwind === undefined) {
211+
throw Error("Asynctify is not enabled, please compile with -s ASYNCIFY=1 in emcc");
212+
}
213+
this.exports.asyncify_stop_unwind();
214+
this.state = AsyncifyStateKind.None;
215+
}
216+
/**
217+
* Initialize the wasm memory to setup necessary meta-data
218+
* for asynctify handling
219+
* @param memory The memory ti
220+
*/
221+
private initMemory(memory: WebAssembly.Memory): void {
222+
// Set the meta-data at address ASYNCTIFY_DATA_ADDR
223+
new Int32Array(memory.buffer, ASYNCIFY_DATA_ADDR, 2).set(
224+
[ASYNCIFY_DATA_START, ASYNCIFY_DATA_END]
225+
);
226+
}
227+
}

0 commit comments

Comments
 (0)