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);