From ccf14e891e91900c9336b3444c4618f655d4fffd Mon Sep 17 00:00:00 2001 From: Arthur Islamov Date: Wed, 16 Aug 2023 05:58:05 +0400 Subject: [PATCH] [js/web] JSEP node assignment optimization (#17128) ### 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 --- js/web/lib/wasm/jsep/webgpu/ops/gather.ts | 2 +- .../core/providers/js/js_data_types.cc | 20 +++++++++++++++++++ onnxruntime/core/providers/js/js_data_types.h | 10 ++++++++++ .../core/providers/js/operators/gather.cc | 17 ++++------------ .../core/providers/js/operators/reshape.cc | 7 ++++--- .../core/providers/js/operators/shape_op.cc | 7 ++++--- .../core/providers/js/operators/squeeze.cc | 7 ++++--- .../core/providers/js/operators/unsqueeze.cc | 7 ++++--- 8 files changed, 51 insertions(+), 26 deletions(-) create mode 100644 onnxruntime/core/providers/js/js_data_types.cc create mode 100644 onnxruntime/core/providers/js/js_data_types.h diff --git a/js/web/lib/wasm/jsep/webgpu/ops/gather.ts b/js/web/lib/wasm/jsep/webgpu/ops/gather.ts index 113bf7c7cc..a915a4bbd9 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/gather.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/gather.ts @@ -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)); diff --git a/onnxruntime/core/providers/js/js_data_types.cc b/onnxruntime/core/providers/js/js_data_types.cc new file mode 100644 index 0000000000..69d5bd4f9d --- /dev/null +++ b/onnxruntime/core/providers/js/js_data_types.cc @@ -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& JsepSupportedDataTypes() { + static const std::vector supportedDataTypes = BuildKernelDefConstraintsFromTypeList(); + return supportedDataTypes; +} +} // namespace js +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/js/js_data_types.h b/onnxruntime/core/providers/js/js_data_types.h new file mode 100644 index 0000000000..d6b6ac0040 --- /dev/null +++ b/onnxruntime/core/providers/js/js_data_types.h @@ -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& JsepSupportedDataTypes(); +} +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/js/operators/gather.cc b/onnxruntime/core/providers/js/operators/gather.cc index ec1ae71243..e9c6f5c792 100644 --- a/onnxruntime/core/providers/js/operators/gather.cc +++ b/onnxruntime/core/providers/js/operators/gather.cc @@ -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()) + .TypeConstraint("T", JsepSupportedDataTypes()) .TypeConstraint("Tind", BuildKernelDefConstraintsFromTypeList>()), Gather); @@ -35,7 +26,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( 12, kJsExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", BuildKernelDefConstraintsFromTypeList()) + .TypeConstraint("T", JsepSupportedDataTypes()) .TypeConstraint("Tind", BuildKernelDefConstraintsFromTypeList>()), Gather); @@ -45,7 +36,7 @@ ONNX_OPERATOR_KERNEL_EX( 13, kJsExecutionProvider, (*KernelDefBuilder::Create()) - .TypeConstraint("T", BuildKernelDefConstraintsFromTypeList()) + .TypeConstraint("T", JsepSupportedDataTypes()) .TypeConstraint("Tind", BuildKernelDefConstraintsFromTypeList>()), Gather); diff --git a/onnxruntime/core/providers/js/operators/reshape.cc b/onnxruntime/core/providers/js/operators/reshape.cc index d8959c89f3..bd6772649a 100644 --- a/onnxruntime/core/providers/js/operators/reshape.cc +++ b/onnxruntime/core/providers/js/operators/reshape.cc @@ -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()) .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()) .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()) .Alias(0, 0) .InputMemoryType(OrtMemTypeCPU, 1), diff --git a/onnxruntime/core/providers/js/operators/shape_op.cc b/onnxruntime/core/providers/js/operators/shape_op.cc index ec0de3c04a..733580a469 100644 --- a/onnxruntime/core/providers/js/operators/shape_op.cc +++ b/onnxruntime/core/providers/js/operators/shape_op.cc @@ -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()), 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()), 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()), Shape); diff --git a/onnxruntime/core/providers/js/operators/squeeze.cc b/onnxruntime/core/providers/js/operators/squeeze.cc index a51fecfd8b..e858ade348 100644 --- a/onnxruntime/core/providers/js/operators/squeeze.cc +++ b/onnxruntime/core/providers/js/operators/squeeze.cc @@ -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()) .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); diff --git a/onnxruntime/core/providers/js/operators/unsqueeze.cc b/onnxruntime/core/providers/js/operators/unsqueeze.cc index c7f90a6727..1485e800e5 100644 --- a/onnxruntime/core/providers/js/operators/unsqueeze.cc +++ b/onnxruntime/core/providers/js/operators/unsqueeze.cc @@ -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()) .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);