diff --git a/js/web/docs/webgpu-operators.md b/js/web/docs/webgpu-operators.md index ba266e4e85..57cdfb0021 100644 --- a/js/web/docs/webgpu-operators.md +++ b/js/web/docs/webgpu-operators.md @@ -21,6 +21,7 @@ Do not modify directly.* | AveragePool | ai.onnx(7-9,10,11+); com.ms.internal.nhwc(11+) | need perf optimization; need implementing activation | | Ceil | ai.onnx(6-12,13+) | | | Clip | ai.onnx(6-10,11,12,13+) | | +| Concat | ai.onnx(1-3,4-10,11-12,13+) | | | Conv | ai.onnx(1-10,11+); com.ms.internal.nhwc(11+) | need perf optimization; conv3d not supported; need implementing activation | | Cos | ai.onnx(7+) | | | Cosh | ai.onnx(9+) | | diff --git a/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts b/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts index bb4d4d558d..b9fc9d26fb 100644 --- a/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts +++ b/js/web/lib/wasm/jsep/webgpu/op-resolve-rules.ts @@ -2,6 +2,7 @@ // Licensed under the MIT License. import * as binaryOps from './ops/binary-op'; +import {concat, parseConcatAttributes} from './ops/concat'; import {conv, parseConvAttributes} from './ops/conv'; import {gemm, parseGemmAttributes} from './ops/gemm'; import {matMul} from './ops/matmul'; @@ -29,6 +30,7 @@ export const WEBGPU_OP_RESOLVE_RULES: Map = new ['Ceil', [unaryOps.ceil]], ['ClipV10', [unaryOps.clipV10]], ['Clip', [unaryOps.clip]], + ['Concat', [concat, parseConcatAttributes]], ['Conv', [conv, parseConvAttributes]], ['Cos', [unaryOps.cos]], ['Cosh', [unaryOps.cosh]], diff --git a/js/web/lib/wasm/jsep/webgpu/ops/concat.ts b/js/web/lib/wasm/jsep/webgpu/ops/concat.ts index d00f176871..b86815131c 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/concat.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/concat.ts @@ -3,7 +3,7 @@ import {TensorView} from '../../tensor'; import {ShapeUtil} from '../../util'; -import {AttributeWithCacheKey} from '../attribute-with-cache-key'; +import {AttributeWithCacheKey, createAttributeWithCacheKey} from '../attribute-with-cache-key'; import {ComputeContext, GpuDataType, ProgramInfo, ProgramInfoLoader, ProgramMetadata} from '../types'; import {createIndicesHelper, IndicesHelper, ShaderHelper} from './common'; @@ -155,3 +155,6 @@ export const concat = (context: ComputeContext, attributes: ConcatAttributes): v validateInputs(context.inputs); context.compute(createConcatProgramInfoLoader(context.inputs, attributes)); }; + +export const parseConcatAttributes = (attributes: Record): ConcatAttributes => + createAttributeWithCacheKey({axis: attributes.axis as number}); diff --git a/onnxruntime/core/providers/js/js_execution_provider.cc b/onnxruntime/core/providers/js/js_execution_provider.cc index a206cab143..9d38176c13 100644 --- a/onnxruntime/core/providers/js/js_execution_provider.cc +++ b/onnxruntime/core/providers/js/js_execution_provider.cc @@ -230,6 +230,11 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnn class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 12, float, MaxPool); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, float, GlobalMaxPool); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 1, 3, Concat); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 4, 10, Concat); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 11, 12, Concat); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kJsExecutionProvider, kOnnxDomain, 13, Concat); + std::unique_ptr RegisterKernels() { auto kernel_registry = std::make_unique(); @@ -386,6 +391,11 @@ std::unique_ptr RegisterKernels() { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, }; for (auto& function_table_entry : function_table) { diff --git a/onnxruntime/core/providers/js/operators/concat.cc b/onnxruntime/core/providers/js/operators/concat.cc new file mode 100644 index 0000000000..7d50d78c82 --- /dev/null +++ b/onnxruntime/core/providers/js/operators/concat.cc @@ -0,0 +1,46 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "concat.h" + +namespace onnxruntime { +namespace js { + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Concat, + kOnnxDomain, + 1, 3, + kJsExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + Concat); + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Concat, + kOnnxDomain, + 4, 10, + kJsExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + Concat); + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Concat, + kOnnxDomain, + 11, 12, + kJsExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + Concat); + +ONNX_OPERATOR_KERNEL_EX( + Concat, + kOnnxDomain, + 13, + kJsExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + Concat); + +} // namespace js +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/js/operators/concat.h b/onnxruntime/core/providers/js/operators/concat.h new file mode 100644 index 0000000000..7a73530799 --- /dev/null +++ b/onnxruntime/core/providers/js/operators/concat.h @@ -0,0 +1,22 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/framework/op_kernel.h" +#include "core/framework/tensor.h" +#include "core/providers/cpu/tensor/concatbase.h" +#include "core/providers/js/js_kernel.h" + +namespace onnxruntime { +namespace js { + +class Concat final : public JsKernel, public ConcatBase { + public: + Concat(const OpKernelInfo& info) : JsKernel(info), ConcatBase(info) { + JSEP_INIT_KERNEL_ATTRIBUTE(Concat, ({"axis" : $1}), static_cast(axis_)); + } +}; + +} // namespace js +} // namespace onnxruntime