Skip to content

[wasm] Fix wasm backend not starting any threads #4957

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mattsoulanille merged 12 commits into from
Apr 22, 2021
Merged

[wasm] Fix wasm backend not starting any threads #4957

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
ef322ce
Fix threaded wasm not creating any workers
mattsoulanille Apr 19, 2021
c9dfebd
Remove wasm workers when dispose is called
mattsoulanille Apr 19, 2021
6b65163
Merge branch 'master' into wasm_fix
mattsoulanille Apr 19, 2021
c81c482
Set pool size to navigator.hardwareConcurrency
mattsoulanille Apr 19, 2021
f276012
Merge branch 'wasm_fix' of github.com:mattsoulanille/tfjs into wasm_fix
mattsoulanille Apr 19, 2021
e29ea61
Separate cc_binary for simd
mattsoulanille Apr 20, 2021
4c1369b
Compile threaded with '-pthread' copt
mattsoulanille Apr 20, 2021
58930a4
Merge branch 'master' into wasm_fix
mattsoulanille Apr 20, 2021
7b40891
Pass wasm module to emscripten
mattsoulanille Apr 22, 2021
f5754f9
Fix wasm paths for overrides map test
mattsoulanille Apr 22, 2021
3deb52e
Merge branch 'wasm_fix' of github.com:mattsoulanille/tfjs into wasm_fix
mattsoulanille Apr 22, 2021
f0a17cb
Merge branch 'master' into wasm_fix
mattsoulanille Apr 22, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 7 additions & 6 deletions tfjs-backend-wasm/scripts/build-wasm.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,18 +33,19 @@ cp -f ../../dist/bin/tfjs-backend-wasm/src/cc/tfjs-backend-wasm/tfjs-backend-was

if [[ "$1" != "--dev" ]]; then
# SIMD and threaded + SIMD builds.
yarn bazel build $BAZEL_REMOTE -c opt --copt="-msimd128" //tfjs-backend-wasm/src/cc:tfjs-backend-wasm-simd \
//tfjs-backend-wasm/src/cc:tfjs-backend-wasm-threaded-simd
yarn bazel build $BAZEL_REMOTE -c opt --copt="-msimd128" //tfjs-backend-wasm/src/cc:tfjs-backend-wasm-simd
yarn bazel build $BAZEL_REMOTE -c opt --copt="-pthread" --copt="-msimd128" //tfjs-backend-wasm/src/cc:tfjs-backend-wasm-threaded-simd

# Copy SIMD
cp -f ../../dist/bin/tfjs-backend-wasm/src/cc/tfjs-backend-wasm-simd/tfjs-backend-wasm.wasm \
cp -f ../../dist/bin/tfjs-backend-wasm/src/cc/tfjs-backend-wasm-simd/tfjs-backend-wasm-simd.wasm \
../wasm-out/tfjs-backend-wasm-simd.wasm

# Copy threaded
cp -f ../../dist/bin/tfjs-backend-wasm/src/cc/tfjs-backend-wasm-threaded-simd/tfjs-backend-wasm.js \
cp -f ../../dist/bin/tfjs-backend-wasm/src/cc/tfjs-backend-wasm-threaded-simd/tfjs-backend-wasm-threaded-simd.js \
../wasm-out/tfjs-backend-wasm-threaded-simd.js
cp -f ../../dist/bin/tfjs-backend-wasm/src/cc/tfjs-backend-wasm-threaded-simd/tfjs-backend-wasm.worker.js \
cp -f ../../dist/bin/tfjs-backend-wasm/src/cc/tfjs-backend-wasm-threaded-simd/tfjs-backend-wasm-threaded-simd.worker.js \
../wasm-out/tfjs-backend-wasm-threaded-simd.worker.js
cp -f ../../dist/bin/tfjs-backend-wasm/src/cc/tfjs-backend-wasm-threaded-simd/tfjs-backend-wasm.wasm \
cp -f ../../dist/bin/tfjs-backend-wasm/src/cc/tfjs-backend-wasm-threaded-simd/tfjs-backend-wasm-threaded-simd.wasm \
../wasm-out/tfjs-backend-wasm-threaded-simd.wasm

node ./create-worker-module.js
Expand Down
8 changes: 6 additions & 2 deletions tfjs-backend-wasm/src/backend_wasm.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import './flags_wasm';
import {backend_util, BackendTimingInfo, DataStorage, DataType, deprecationWarn, engine, env, KernelBackend, TensorInfo, util} from '@tensorflow/tfjs-core';

import {BackendWasmModule, WasmFactoryConfig} from '../wasm-out/tfjs-backend-wasm';
import {BackendWasmThreadedSimdModule} from '../wasm-out/tfjs-backend-wasm-threaded-simd';
Copy link
Copy Markdown
Collaborator

@lina128 lina128 Apr 20, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does the Simd module need a separate instance too or not?

Copy link
Copy Markdown
Member Author

@mattsoulanille mattsoulanille Apr 20, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The simd module doesn't need a separate interface because its API is the same as the non-simd API. The threaded simd module, however, does need a separate interface because it now has the PThread.terminateAllThreads function exposed (via the type declarations).

import wasmFactoryThreadedSimd from '../wasm-out/tfjs-backend-wasm-threaded-simd.js';
// @ts-ignore
import {wasmWorkerContents} from '../wasm-out/tfjs-backend-wasm-threaded-simd.worker.js';
Expand All @@ -41,7 +42,7 @@ export class BackendWasm extends KernelBackend {
private dataIdNextNumber = 1;
dataIdMap: DataStorage<TensorData>;

constructor(public wasm: BackendWasmModule) {
constructor(public wasm: BackendWasmModule | BackendWasmThreadedSimdModule) {
super();
this.wasm.tfjs.init();
this.dataIdMap = new DataStorage(this, engine());
Expand Down Expand Up @@ -159,6 +160,9 @@ export class BackendWasm extends KernelBackend {

dispose() {
this.wasm.tfjs.dispose();
if ('PThread' in this.wasm) {
this.wasm.PThread.terminateAllThreads();
}
this.wasm = null;
}

Expand Down Expand Up @@ -214,7 +218,7 @@ function createInstantiateWasmFunc(path: string) {
}
response.arrayBuffer().then(binary => {
WebAssembly.instantiate(binary, imports).then(output => {
callback(output.instance);
callback(output.instance, output.module);
});
});
});
Expand Down
53 changes: 42 additions & 11 deletions tfjs-backend-wasm/src/cc/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,18 @@ KERNELS_WITH_KEEPALIVE = glob(
exclude = ["**/*_test.cc"],
)

BASE_LINKOPTS = [
"-s ALLOW_MEMORY_GROWTH=1",
"-s DEFAULT_LIBRARY_FUNCS_TO_INCLUDE=[]",
"-s DISABLE_EXCEPTION_CATCHING=1",
"-s FILESYSTEM=0",
"-s EXIT_RUNTIME=0",
"-s EXPORTED_FUNCTIONS='[\"_malloc\", \"_free\"]'",
"-s EXTRA_EXPORTED_RUNTIME_METHODS='[\"cwrap\"]'",
"-s MODULARIZE=1",
"-s MALLOC=emmalloc",
]

# This build rule generates tfjs-backend-wasm.{js,wasm}.
#
# The ".js" at the end of the build name is significant because it determines
Expand All @@ -19,17 +31,36 @@ KERNELS_WITH_KEEPALIVE = glob(
cc_binary(
name = "tfjs-backend-wasm.js",
srcs = ["backend.cc"] + KERNELS_WITH_KEEPALIVE,
linkopts = [
"-s ALLOW_MEMORY_GROWTH=1",
"-s DEFAULT_LIBRARY_FUNCS_TO_INCLUDE=[]",
"-s DISABLE_EXCEPTION_CATCHING=1",
"-s FILESYSTEM=0",
"-s EXIT_RUNTIME=0",
"-s EXPORTED_FUNCTIONS='[\"_malloc\", \"_free\"]'",
"-s EXTRA_EXPORTED_RUNTIME_METHODS='[\"cwrap\"]'",
"-s MODULARIZE=1",
linkopts = BASE_LINKOPTS + [
"-s EXPORT_NAME=WasmBackendModule",
],
deps = [
":all_kernels",
":backend",
],
)

cc_binary(
name = "tfjs-backend-wasm-simd.js",
srcs = ["backend.cc"] + KERNELS_WITH_KEEPALIVE,
linkopts = BASE_LINKOPTS,
deps = [
":all_kernels",
":backend",
],
)

cc_binary(
name = "tfjs-backend-wasm-threaded-simd.js",
srcs = ["backend.cc"] + KERNELS_WITH_KEEPALIVE,
linkopts = BASE_LINKOPTS + [
"-s EXPORT_NAME=WasmBackendModuleThreadedSimd",
"-s MALLOC=emmalloc",
"-s USE_PTHREADS=1",
"-s PROXY_TO_PTHREAD=1",
# Many x86-64 processors have 2 threads per core, so we divide by 2.
"-s PTHREAD_POOL_SIZE=" +
"'Math.min(4, Math.max(1, (navigator.hardwareConcurrency || 1) / 2))'",
Copy link
Copy Markdown
Contributor

@qjia7 qjia7 Apr 20, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it doable that we register a flag WASM_PTHREAD_POOL_SIZE in flags_wasm.ts and initialized with 'Math.min(4, Math.max(1, (navigator.hardwareConcurrency || 1) / 2))'. Then in here and backend.cc, we both use the flag value?

],
deps = [
":all_kernels",
Expand All @@ -44,13 +75,13 @@ wasm_cc_binary(

wasm_cc_binary(
name = "tfjs-backend-wasm-simd",
cc_target = ":tfjs-backend-wasm.js",
cc_target = ":tfjs-backend-wasm-simd.js",
simd = True,
)

wasm_cc_binary(
name = "tfjs-backend-wasm-threaded-simd",
cc_target = ":tfjs-backend-wasm.js",
cc_target = ":tfjs-backend-wasm-threaded-simd.js",
simd = True,
threads = "emscripten",
)
Expand Down
4 changes: 2 additions & 2 deletions tfjs-backend-wasm/src/index_test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -165,9 +165,9 @@ describeWithFlags('wasm init', BROWSER_ENVS, () => {
{
'tfjs-backend-wasm.wasm': '/base/wasm-out/tfjs-backend-wasm.wasm',
'tfjs-backend-wasm-simd.wasm':
'/base/wasm-out/tfjs-backend-wasm.wasm',
'/base/wasm-out/tfjs-backend-wasm-simd.wasm',
'tfjs-backend-wasm-threaded-simd.wasm':
'/base/wasm-out/tfjs-backend-wasm.wasm'
'/base/wasm-out/tfjs-backend-wasm-threaded-simd.wasm'
},
usePlatformFetch);
let wasmPath: string;
Expand Down
17 changes: 7 additions & 10 deletions tfjs-backend-wasm/wasm-out/tfjs-backend-wasm-threaded-simd.d.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,16 +15,13 @@
* =============================================================================
*/

export interface BackendWasmModule extends EmscriptenModule {
// Using the tfjs namespace to avoid conflict with emscripten's API.
tfjs: {
init(): void,
registerTensor(id: number, size: number, memoryOffset: number): void,
// Disposes the data behind the data bucket.
disposeData(id: number): void,
// Disposes the backend and all of its associated data.
dispose(): void,
}
import {BackendWasmModule} from './tfjs-backend-wasm';

export interface BackendWasmThreadedSimdModule extends BackendWasmModule {
PThread: {
// Terminates all webworkers
terminateAllThreads(): void,
};
}

export interface WasmFactoryConfig {
Expand Down