Skip to content

Commit

Permalink
[js/webgpu] following up for JSEP/WebGPU code cleanup (#15666)
Browse files Browse the repository at this point in the history
### Description
This PR resolves a part of non-critical comments from code review
comments in #14579.

- use `USE_JSEP` instead of `USE_JS` in build definition to make it less
ambiguous
- remove unused util functions from util.ts
- fix transpose.h
- other misc fixes
  • Loading branch information
fs-eire authored Apr 26, 2023
1 parent ebaafac commit b98317b
Show file tree
Hide file tree
Showing 16 changed files with 35 additions and 351 deletions.
8 changes: 4 additions & 4 deletions cmake/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ option(onnxruntime_USE_QNN "Build with QNN support" OFF)
option(onnxruntime_USE_SNPE "Build with SNPE support" OFF)
option(onnxruntime_USE_RKNPU "Build with RKNPU support" OFF)
option(onnxruntime_USE_DNNL "Build with DNNL support" OFF)
option(onnxruntime_USE_JS "Build with JavaScript implemented kernels support" OFF)
option(onnxruntime_USE_JSEP "Build with JavaScript implemented kernels support" OFF)
option(onnxruntime_BUILD_UNIT_TESTS "Build ONNXRuntime unit tests" ON)
option(onnxruntime_BUILD_CSHARP "Build C# library" OFF)
option(onnxruntime_BUILD_OBJC "Build Objective-C library" OFF)
Expand Down Expand Up @@ -662,9 +662,9 @@ if (onnxruntime_USE_NNAPI_BUILTIN)
list(APPEND ORT_PROVIDER_CMAKE_FLAGS -Donnxruntime_USE_NNAPI_BUILTIN=1)
list(APPEND ONNXRUNTIME_PROVIDER_NAMES nnapi)
endif()
if (onnxruntime_USE_JS)
list(APPEND ORT_PROVIDER_FLAGS -DUSE_JS=1)
list(APPEND ORT_PROVIDER_CMAKE_FLAGS -Donnxruntime_USE_JS=1)
if (onnxruntime_USE_JSEP)
list(APPEND ORT_PROVIDER_FLAGS -DUSE_JSEP=1)
list(APPEND ORT_PROVIDER_CMAKE_FLAGS -Donnxruntime_USE_JSEP=1)
list(APPEND ONNXRUNTIME_PROVIDER_NAMES js)
endif()
if (onnxruntime_USE_QNN)
Expand Down
6 changes: 3 additions & 3 deletions cmake/onnxruntime_providers.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,7 @@ endif()
if(onnxruntime_USE_NNAPI_BUILTIN)
set(PROVIDERS_NNAPI onnxruntime_providers_nnapi)
endif()
if(onnxruntime_USE_JS)
if(onnxruntime_USE_JSEP)
set(PROVIDERS_JS onnxruntime_providers_js)
endif()
if(onnxruntime_USE_QNN)
Expand Down Expand Up @@ -1067,8 +1067,8 @@ if (onnxruntime_USE_NNAPI_BUILTIN)
endif()
endif()

if (onnxruntime_USE_JS)
add_compile_definitions(USE_JS=1)
if (onnxruntime_USE_JSEP)
add_compile_definitions(USE_JSEP=1)

file(GLOB_RECURSE onnxruntime_providers_js_cc_srcs
"${ONNXRUNTIME_ROOT}/core/providers/js/*.h"
Expand Down
6 changes: 3 additions & 3 deletions cmake/onnxruntime_unittests.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -504,7 +504,7 @@ if(onnxruntime_USE_NNAPI_BUILTIN)
list(APPEND onnxruntime_test_providers_dependencies onnxruntime_providers_nnapi)
endif()

if(onnxruntime_USE_JS)
if(onnxruntime_USE_JSEP)
list(APPEND onnxruntime_test_providers_dependencies onnxruntime_providers_js)
endif()

Expand Down Expand Up @@ -609,7 +609,7 @@ if(onnxruntime_USE_NNAPI_BUILTIN)
list(APPEND onnxruntime_test_providers_libs onnxruntime_providers_nnapi)
endif()

if(onnxruntime_USE_JS)
if(onnxruntime_USE_JSEP)
list(APPEND onnxruntime_test_framework_src_patterns ${TEST_SRC_DIR}/providers/js/*)
list(APPEND onnxruntime_test_framework_libs onnxruntime_providers_js)
list(APPEND onnxruntime_test_providers_dependencies onnxruntime_providers_js)
Expand Down Expand Up @@ -851,7 +851,7 @@ if (onnxruntime_BUILD_WEBASSEMBLY)
if (onnxruntime_ENABLE_WEBASSEMBLY_THREADS)
set_property(TARGET onnxruntime_test_all APPEND_STRING PROPERTY LINK_FLAGS " -s USE_PTHREADS=1 -s PROXY_TO_PTHREAD=1")
endif()
if (onnxruntime_USE_JS)
if (onnxruntime_USE_JSEP)
set_property(TARGET onnxruntime_test_all APPEND_STRING PROPERTY LINK_FLAGS " --pre-js \"${ONNXRUNTIME_ROOT}/wasm/js_internal_api.js\"")
endif()
endif()
Expand Down
6 changes: 3 additions & 3 deletions cmake/onnxruntime_webassembly.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -199,7 +199,7 @@ else()
endif()

set(EXPORTED_RUNTIME_METHODS "['stackAlloc','stackRestore','stackSave','UTF8ToString','stringToUTF8','lengthBytesUTF8']")
if (onnxruntime_USE_JS)
if (onnxruntime_USE_JSEP)
set(EXPORTED_FUNCTIONS "_malloc,_free,_JsepOutput")
else()
set(EXPORTED_FUNCTIONS "_malloc,_free")
Expand All @@ -219,12 +219,12 @@ else()
--no-entry
)

if (onnxruntime_USE_JS)
if (onnxruntime_USE_JSEP)
# NOTE: "-s ASYNCIFY=1" is required for JSEP to work with WebGPU
# This flag allows async functions to be called from sync functions, in the cost of binary size and
# build time. See https://emscripten.org/docs/porting/asyncify.html for more details.

target_compile_definitions(onnxruntime_webassembly PRIVATE USE_JS=1)
target_compile_definitions(onnxruntime_webassembly PRIVATE USE_JSEP=1)
target_link_options(onnxruntime_webassembly PRIVATE
--pre-js "${ONNXRUNTIME_ROOT}/wasm/js_internal_api.js"
"SHELL:-s ASYNCIFY=1"
Expand Down
6 changes: 4 additions & 2 deletions js/web/lib/wasm/jsep/backend-webgpu.ts
Original file line number Diff line number Diff line change
Expand Up @@ -144,8 +144,10 @@ export class WebGpuBackend {
}

dispose(): void {
// TODO: uninitialization
// this.glContext.dispose();
// currently, we do not do anything in this function. In all known use cases, we don't have the requirement to
// actually dispose the WebGpuBackend instance, because it's always used as a singleton.
//
// revisit this place if we get real requirement to dispose the instance.
}

getCommandEncoder(): GPUCommandEncoder {
Expand Down
4 changes: 2 additions & 2 deletions js/web/lib/wasm/jsep/init.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ class TensorViewImpl implements TensorView {
}
}

class OpKernelContext implements ComputeContext {
class ComputeContextImpl implements ComputeContext {
readonly opKernelContext: number;
readonly inputs: readonly TensorView[];
get customData(): {[key: string]: unknown} {
Expand Down Expand Up @@ -142,7 +142,7 @@ export const init = async(module: OrtWasmModule): Promise<void> => {
// jsepRun
(kernel: number, contextDataOffset: number) => {
LOG_DEBUG('verbose', () => `[WebGPU] jsepRun: kernel=${kernel}, contextDataOffset=${contextDataOffset}`);
const context = new OpKernelContext(module, backend, contextDataOffset);
const context = new ComputeContextImpl(module, backend, contextDataOffset);
return backend.computeKernel(kernel, context);
});
}
Expand Down
Loading

0 comments on commit b98317b

Please sign in to comment.