saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[js/web] JSEP node assignment optimization (#17128)

Browse source 
### Description
Since WebGPU supports only float32 and int32, having Gather, Reshape,
Shape, Squeeze and Unsqueeze ops with other data types create additional
MemCpy ops and slow down the overall execution as all other OPs with
other tensor types will be done on CPU.

Before this patch SD Unet had these numbers:
Node(s) placed on [CPUExecutionProvider]. Number of nodes: 1141
Node(s) placed on [JsExecutionProvider]. Number of nodes: 4025
memcpy tokens: 2001

After patch:
Node(s) placed on [CPUExecutionProvider]. Number of nodes: 1735
Node(s) placed on [JsExecutionProvider]. Number of nodes: 2243
memcpu tokens: 813

It also gives more than 5X performance benefit. From 12sec for one Unet
step to 2.2sec on RTX 3090 Ti, so we are almost getting to native
performance.

UPD: with latest changes from main branch and multi-threading it went
down to 1.6sec. Will try re-exporting my model to onnx with maximum
optimizations, like using MultiHeadAttention to decrease node count.
Maybe after implementing that it can go in less than 1 sec
This commit is contained in:
Arthur Islamov 2023-08-16 05:58:05 +04:00 committed by GitHub
parent 3cdf42548f
commit ccf14e891e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 51 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
js/web/lib/wasm/jsep/webgpu/ops/gather.ts
View file

@ -100,7 +100,7 @@ export const gather = (context: ComputeContext, attributes: GatherAttributes): v
const metadata = {
name: 'Gather',
inputTypes: [GpuDataType.default, GpuDataType.default],
cacheHint: attributes.cacheKey + inputs[0].dataType.toString(10) + inputs[1].dataType.toString(10),
cacheHint: attributes.cacheKey,
};
context.compute(createGatherProgramInfo(metadata, context.inputs, attributes));

20
onnxruntime/core/providers/js/js_data_types.cc Normal file
View file

@ -0,0 +1,20 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/providers/cpu/tensor/shape_op.h"
namespace onnxruntime {
namespace js {
using SupportedTypes =
TypeList<
float,
int32_t,
uint32_t>;
const std::vector<MLDataType>& JsepSupportedDataTypes() {
static const std::vector<MLDataType> supportedDataTypes = BuildKernelDefConstraintsFromTypeList<SupportedTypes>();
return supportedDataTypes;
}
} // namespace js
} // namespace onnxruntime

10
onnxruntime/core/providers/js/js_data_types.h Normal file
View file

@ -0,0 +1,10 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/framework/data_types.h"
namespace onnxruntime {
namespace js {
std::vector<MLDataType>& JsepSupportedDataTypes();
}
} // namespace onnxruntime

17
onnxruntime/core/providers/js/operators/gather.cc
View file

@ -2,21 +2,12 @@
// Licensed under the MIT License.
#include "core/providers/js/js_kernel.h"
#include "core/providers/js/js_data_types.h"
#include "gather.h"
namespace onnxruntime {
namespace js {
using AllSupportedSize =
TypeList<
float,
double,
int64_t,
uint64_t,
int32_t,
uint32_t>;
ONNX_OPERATOR_VERSIONED_KERNEL_EX(
Gather,
kOnnxDomain,
@ -24,7 +15,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
10,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", BuildKernelDefConstraintsFromTypeList<AllSupportedSize>())
.TypeConstraint("T", JsepSupportedDataTypes())
.TypeConstraint("Tind", BuildKernelDefConstraintsFromTypeList<TypeList<int32_t, int64_t>>()),
Gather);
@ -35,7 +26,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
12,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", BuildKernelDefConstraintsFromTypeList<AllSupportedSize>())
.TypeConstraint("T", JsepSupportedDataTypes())
.TypeConstraint("Tind", BuildKernelDefConstraintsFromTypeList<TypeList<int32_t, int64_t>>()),
Gather);
@ -45,7 +36,7 @@ ONNX_OPERATOR_KERNEL_EX(
13,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", BuildKernelDefConstraintsFromTypeList<AllSupportedSize>())
.TypeConstraint("T", JsepSupportedDataTypes())
.TypeConstraint("Tind", BuildKernelDefConstraintsFromTypeList<TypeList<int32_t, int64_t>>()),
Gather);

7
onnxruntime/core/providers/js/operators/reshape.cc
View file

@ -2,6 +2,7 @@
// Licensed under the MIT License.
#include "reshape.h"
#include "core/providers/js/js_data_types.h"
namespace onnxruntime {
namespace js {
@ -12,7 +13,7 @@ ONNX_OPERATOR_KERNEL_EX(
14,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
.TypeConstraint("T", JsepSupportedDataTypes())
.TypeConstraint("shape", DataTypeImpl::GetTensorType<int64_t>())
.Alias(0, 0)
.InputMemoryType(OrtMemTypeCPU, 1),
@ -24,7 +25,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
13, 13,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
.TypeConstraint("T", JsepSupportedDataTypes())
.TypeConstraint("shape", DataTypeImpl::GetTensorType<int64_t>())
.Alias(0, 0)
.InputMemoryType(OrtMemTypeCPU, 1),
@ -36,7 +37,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
5, 12,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
.TypeConstraint("T", JsepSupportedDataTypes())
.TypeConstraint("shape", DataTypeImpl::GetTensorType<int64_t>())
.Alias(0, 0)
.InputMemoryType(OrtMemTypeCPU, 1),

7
onnxruntime/core/providers/js/operators/shape_op.cc
View file

@ -2,6 +2,7 @@
// Licensed under the MIT License.
#include "core/providers/js/js_kernel.h"
#include "core/providers/js/js_data_types.h"
#include "core/providers/cpu/tensor/shape_op.h"
namespace onnxruntime {
@ -15,7 +16,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
(*KernelDefBuilder::Create())
// properly force CPU/GPU synch inside the kernel
.OutputMemoryType(OrtMemTypeCPU, 0)
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
.TypeConstraint("T", JsepSupportedDataTypes())
.TypeConstraint("T1", DataTypeImpl::GetTensorType<int64_t>()),
Shape);
@ -27,7 +28,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
(*KernelDefBuilder::Create())
// properly force CPU/GPU synch inside the kernel
.OutputMemoryType(OrtMemTypeCPU, 0)
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
.TypeConstraint("T", JsepSupportedDataTypes())
.TypeConstraint("T1", DataTypeImpl::GetTensorType<int64_t>()),
Shape);
@ -39,7 +40,7 @@ ONNX_OPERATOR_KERNEL_EX(
(*KernelDefBuilder::Create())
// properly force CPU/GPU synch inside the kernel
.OutputMemoryType(OrtMemTypeCPU, 0)
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
.TypeConstraint("T", JsepSupportedDataTypes())
.TypeConstraint("T1", DataTypeImpl::GetTensorType<int64_t>()),
Shape);

7
onnxruntime/core/providers/js/operators/squeeze.cc
View file

@ -2,6 +2,7 @@
// Licensed under the MIT License.
#include "squeeze.h"
#include "core/providers/js/js_data_types.h"
namespace onnxruntime {
namespace js {
@ -12,7 +13,7 @@ ONNX_OPERATOR_KERNEL_EX(
13,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
.TypeConstraint("T", JsepSupportedDataTypes())
.TypeConstraint("axes", DataTypeImpl::GetTensorType<int64_t>())
.Alias(0, 0)
.InputMemoryType(OrtMemTypeCPU, 1),
@ -24,7 +25,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
11, 12,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
.TypeConstraint("T", JsepSupportedDataTypes())
.Alias(0, 0),
Squeeze);
@ -34,7 +35,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
1, 10,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
.TypeConstraint("T", JsepSupportedDataTypes())
.Alias(0, 0),
Squeeze);

7
onnxruntime/core/providers/js/operators/unsqueeze.cc
View file

@ -2,6 +2,7 @@
// Licensed under the MIT License.
#include "unsqueeze.h"
#include "core/providers/js/js_data_types.h"
namespace onnxruntime {
namespace js {
@ -12,7 +13,7 @@ ONNX_OPERATOR_KERNEL_EX(
13,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
.TypeConstraint("T", JsepSupportedDataTypes())
.TypeConstraint("axes", DataTypeImpl::GetTensorType<int64_t>())
.Alias(0, 0)
.InputMemoryType(OrtMemTypeCPU, 1),
@ -24,7 +25,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
11, 12,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
.TypeConstraint("T", JsepSupportedDataTypes())
.Alias(0, 0),
Unsqueeze);
@ -34,7 +35,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(
1, 10,
kJsExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
.TypeConstraint("T", JsepSupportedDataTypes())
.Alias(0, 0),
Unsqueeze);