diff --git a/js/web/lib/wasm/jsep/webgpu/ops/split.ts b/js/web/lib/wasm/jsep/webgpu/ops/split.ts index 76a69e4bb8061..bcc63237f0429 100644 --- a/js/web/lib/wasm/jsep/webgpu/ops/split.ts +++ b/js/web/lib/wasm/jsep/webgpu/ops/split.ts @@ -20,6 +20,7 @@ export interface SplitAttributes extends AttributeWithCacheKey { readonly axis: number; readonly numOutputs: number; readonly splitSizes: number[]; + readonly isUnevenSplitAllowed?: boolean; } const validateInputs = (inputs: readonly TensorView[]): void => { @@ -38,7 +39,12 @@ const createSplitAttributesFromInputs = ( inputs[1].getBigInt64Array().forEach((v) => splitSizes.push(Number(v))); numOutputs = splitSizes.length; } - return createAttributeWithCacheKey({ numOutputs, axis: attributes.axis, splitSizes }); + return createAttributeWithCacheKey({ + numOutputs, + axis: attributes.axis, + splitSizes, + isUnevenSplitAllowed: attributes.isUnevenSplitAllowed, + }); }; const calculateOutputIndexImpl = (numberOfTensors: number): string => ` @@ -130,17 +136,70 @@ export const createSplitProgramInfo = (inputs: readonly TensorView[], attributes export const split = (context: ComputeContext, attributes: SplitAttributes): void => { validateInputs(context.inputs); - const updatedAttributes = + let updatedAttributes = context.inputs.length === 1 ? attributes : createSplitAttributesFromInputs(context.inputs, attributes); - context.compute(createSplitProgramInfo(context.inputs, updatedAttributes), { inputs: [0] }); + if (updatedAttributes.splitSizes.length === 0) { + const inputShape = context.inputs[0].dims; + const axis = ShapeUtil.normalizeAxis(updatedAttributes.axis, inputShape.length); + const splitDim = inputShape[axis]; + const numOutputs = updatedAttributes.numOutputs; + if (numOutputs < 1) { + throw new Error('numOutputs must be positive'); + } + if (!updatedAttributes.isUnevenSplitAllowed && splitDim % numOutputs !== 0) { + throw new Error('The split dimension must be evenly divisible by the number of outputs'); + } + const splitSize = Math.ceil(splitDim / numOutputs); + const lastSplitSize = splitDim - splitSize * (numOutputs - 1); + if (updatedAttributes.isUnevenSplitAllowed && lastSplitSize < 0) { + throw new Error('numOutputs cannot produce a negative split size'); + } + const splitSizes = new Array(numOutputs).fill(splitSize); + splitSizes[numOutputs - 1] = lastSplitSize; + updatedAttributes = createAttributeWithCacheKey({ + numOutputs, + axis: updatedAttributes.axis, + splitSizes, + isUnevenSplitAllowed: updatedAttributes.isUnevenSplitAllowed, + }); + } + const inputShape = context.inputs[0].dims; + const axis = ShapeUtil.normalizeAxis(updatedAttributes.axis, inputShape.length); + const inputSize = ShapeUtil.size(inputShape); + const outputIndices: number[] = []; + const nonEmptySplitSizes: number[] = []; + for (let i = 0; i < updatedAttributes.numOutputs; i++) { + const splitSize = updatedAttributes.splitSizes[i]; + if (inputSize === 0 || splitSize === 0) { + const outputShape = inputShape.slice(); + outputShape[axis] = splitSize; + context.output(i, outputShape); + } else { + outputIndices.push(i); + nonEmptySplitSizes.push(splitSize); + } + } + if (outputIndices.length === 0) { + return; + } + if (outputIndices.length !== updatedAttributes.numOutputs) { + updatedAttributes = createAttributeWithCacheKey({ + numOutputs: outputIndices.length, + axis: updatedAttributes.axis, + splitSizes: nonEmptySplitSizes, + isUnevenSplitAllowed: updatedAttributes.isUnevenSplitAllowed, + }); + } + context.compute(createSplitProgramInfo(context.inputs, updatedAttributes), { inputs: [0], outputs: outputIndices }); }; export const parseSplitAttributes = (attributes: Record): SplitAttributes => { const axis = attributes.axis as number; const splitSizes: number[] = attributes.splitSizes as number[]; const numOutputs = (attributes.numOutputs as number) < 0 ? splitSizes.length : (attributes.numOutputs as number); - if (numOutputs !== splitSizes.length) { + const isUnevenSplitAllowed = !!attributes.isUnevenSplitAllowed; + if (splitSizes.length > 0 && numOutputs !== splitSizes.length) { throw new Error('numOutputs and splitSizes length must be equal'); } - return createAttributeWithCacheKey({ axis, numOutputs, splitSizes }); + return createAttributeWithCacheKey({ axis, numOutputs, splitSizes, isUnevenSplitAllowed }); }; diff --git a/js/web/test/data/ops/split-num-outputs.jsonc b/js/web/test/data/ops/split-num-outputs.jsonc new file mode 100644 index 0000000000000..4e3ecd482518c --- /dev/null +++ b/js/web/test/data/ops/split-num-outputs.jsonc @@ -0,0 +1,88 @@ +[ + { + "name": "Split 18 with num_outputs and empty optional split input", + "operator": "Split", + "opset": { "domain": "", "version": 18 }, + "attributes": [{ "name": "num_outputs", "data": 2, "type": "int" }], + "cases": [ + { + "name": "T[4]", + "inputs": [ + { "data": [1, 2, 3, 4], "dims": [4], "type": "float32" }, + { "data": null, "type": "int64" } + ], + "outputs": [ + { "data": [1, 2], "dims": [2], "type": "float32" }, + { "data": [3, 4], "dims": [2], "type": "float32" } + ] + }, + { + "name": "T[5]", + "inputs": [ + { "data": [1, 2, 3, 4, 5], "dims": [5], "type": "float32" }, + { "data": null, "type": "int64" } + ], + "outputs": [ + { "data": [1, 2, 3], "dims": [3], "type": "float32" }, + { "data": [4, 5], "dims": [2], "type": "float32" } + ] + }, + { + "name": "T[1] with a zero-sized final output", + "inputs": [ + { "data": [1], "dims": [1], "type": "float32" }, + { "data": null, "type": "int64" } + ], + "outputs": [ + { "data": [1], "dims": [1], "type": "float32" }, + { "data": [], "dims": [0], "type": "float32" } + ] + } + ] + }, + { + "name": "Split 18 with num_outputs and a zero-sized final output", + "operator": "Split", + "opset": { "domain": "", "version": 18 }, + "attributes": [{ "name": "num_outputs", "data": 6, "type": "int" }], + "cases": [ + { + "name": "T[10]", + "inputs": [ + { "data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], "dims": [10], "type": "float32" }, + { "data": null, "type": "int64" } + ], + "outputs": [ + { "data": [1, 2], "dims": [2], "type": "float32" }, + { "data": [3, 4], "dims": [2], "type": "float32" }, + { "data": [5, 6], "dims": [2], "type": "float32" }, + { "data": [7, 8], "dims": [2], "type": "float32" }, + { "data": [9, 10], "dims": [2], "type": "float32" }, + { "data": [], "dims": [0], "type": "float32" } + ] + } + ] + }, + { + "name": "Split 18 with num_outputs, negative axis and empty optional split input", + "operator": "Split", + "opset": { "domain": "", "version": 18 }, + "attributes": [ + { "name": "axis", "data": -1, "type": "int" }, + { "name": "num_outputs", "data": 2, "type": "int" } + ], + "cases": [ + { + "name": "T[2,5]", + "inputs": [ + { "data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], "dims": [2, 5], "type": "float32" }, + { "data": null, "type": "int64" } + ], + "outputs": [ + { "data": [1, 2, 3, 6, 7, 8], "dims": [2, 3], "type": "float32" }, + { "data": [4, 5, 9, 10], "dims": [2, 2], "type": "float32" } + ] + } + ] + } +] diff --git a/js/web/test/suite-test-list.jsonc b/js/web/test/suite-test-list.jsonc index ee63fe68ee227..db2b40571083b 100644 --- a/js/web/test/suite-test-list.jsonc +++ b/js/web/test/suite-test-list.jsonc @@ -1409,6 +1409,7 @@ //"softmax.jsonc", "sin.jsonc", //"split.jsonc", + "split-num-outputs.jsonc", "sqrt.jsonc", "sub.jsonc", "sub_int32.jsonc", diff --git a/js/web/test/unittests/opset.ts b/js/web/test/unittests/opset.ts index 57433c68a6381..f00d398cd1692 100644 --- a/js/web/test/unittests/opset.ts +++ b/js/web/test/unittests/opset.ts @@ -8,6 +8,15 @@ import { WEBGL_OP_RESOLVE_RULES } from '../../lib/onnxjs/backends/webgl/op-resol import { Graph } from '../../lib/onnxjs/graph'; import { OpSet, resolveOperator } from '../../lib/onnxjs/opset'; import { Tensor } from '../../lib/onnxjs/tensor'; +import { DataType } from '../../lib/wasm/wasm-common'; +import { TensorView } from '../../lib/wasm/jsep/tensor-view'; +import { + ComputeContext, + ComputeContextInputsOutputsMapping, + ProgramInfo, + TensorInfo, +} from '../../lib/wasm/jsep/webgpu/types'; +import { parseSplitAttributes, split, SplitAttributes } from '../../lib/wasm/jsep/webgpu/ops/split'; function createTestGraphNode(name: string, opType: string): Graph.Node { return { name, opType, inputs: [], outputs: [], attributes: new Attribute(null) }; @@ -106,3 +115,106 @@ describe('#UnitTest# - resolve rules', () => { checkConsistency(webglCheckOnlyRules); }); }); + +describe('#UnitTest# - JSEP Split runtime shapes', () => { + const runSplit = ( + dims: number[], + attributes: SplitAttributes, + splitSizes?: bigint[] | null, + ): readonly TensorInfo[] => { + const inputs = [{ dims, dataType: DataType.float }] as unknown as TensorView[]; + if (splitSizes !== undefined) { + inputs.push({ + dims: splitSizes === null ? [] : [splitSizes.length], + dataType: splitSizes === null ? 0 : DataType.int64, + getBigInt64Array: () => BigInt64Array.from(splitSizes ?? []), + } as unknown as TensorView); + } + const outputs: TensorInfo[] = []; + const context = { + inputs, + output: (index: number, shape: readonly number[]) => { + outputs[index] = { dims: shape, dataType: DataType.float }; + return 0; + }, + compute: (program: ProgramInfo, mapping: ComputeContextInputsOutputsMapping) => { + const programOutputs = program.getRunData(inputs).outputs; + programOutputs.forEach((output, index) => { + expect(output.dims.every((dim) => dim > 0)).to.equal(true); + outputs[mapping.outputs![index]] = output; + }); + return []; + }, + } as unknown as ComputeContext; + split(context, attributes); + return outputs; + }; + + const attributes = parseSplitAttributes({ axis: 0, numOutputs: 2, splitSizes: [], isUnevenSplitAllowed: true }); + + it('infers even split sizes without graph shapes', () => { + expect(runSplit([4], attributes).map((output) => output.dims)).to.deep.equal([[2], [2]]); + }); + + it('ignores the empty optional split input placeholder', () => { + expect(runSplit([4], attributes, null).map((output) => output.dims)).to.deep.equal([[2], [2]]); + }); + + it('recomputes uneven split sizes for changing input shapes', () => { + expect(runSplit([5], attributes).map((output) => output.dims)).to.deep.equal([[3], [2]]); + expect(runSplit([7], attributes).map((output) => output.dims)).to.deep.equal([[4], [3]]); + expect(attributes.splitSizes).to.deep.equal([]); + }); + + it('normalizes a negative split axis', () => { + const negativeAxis = parseSplitAttributes({ axis: -1, numOutputs: 2, splitSizes: [], isUnevenSplitAllowed: true }); + expect(runSplit([2, 5], negativeAxis).map((output) => output.dims)).to.deep.equal([ + [2, 3], + [2, 2], + ]); + }); + + it('preserves explicit split sizes from an input tensor', () => { + expect(runSplit([5], attributes, [1n, 4n]).map((output) => output.dims)).to.deep.equal([[1], [4]]); + }); + + it('preserves explicit split sizes from an attribute', () => { + const explicitSizes = parseSplitAttributes({ axis: 0, numOutputs: 2, splitSizes: [1, 4] }); + expect(runSplit([5], explicitSizes).map((output) => output.dims)).to.deep.equal([[1], [4]]); + }); + + it('requires even splitting without the num_outputs attribute', () => { + const evenOnly = parseSplitAttributes({ axis: 0, numOutputs: 2, splitSizes: [] }); + expect(runSplit([4], evenOnly).map((output) => output.dims)).to.deep.equal([[2], [2]]); + expect(() => runSplit([5], evenOnly)).to.throw('evenly divisible'); + }); + + it('allows a zero-sized final output', () => { + const sixOutputs = parseSplitAttributes({ axis: 0, numOutputs: 6, splitSizes: [], isUnevenSplitAllowed: true }); + expect(runSplit([10], sixOutputs).map((output) => output.dims)).to.deep.equal([[2], [2], [2], [2], [2], [0]]); + expect(runSplit([10], sixOutputs, null).map((output) => output.dims)).to.deep.equal([[2], [2], [2], [2], [2], [0]]); + expect(runSplit([1], attributes).map((output) => output.dims)).to.deep.equal([[1], [0]]); + }); + + it('allows zero-sized outputs for an empty input', () => { + expect(runSplit([0], attributes).map((output) => output.dims)).to.deep.equal([[0], [0]]); + expect(runSplit([0, 4], attributes).map((output) => output.dims)).to.deep.equal([ + [0, 4], + [0, 4], + ]); + }); + + it('preserves output positions when explicit split sizes include empty outputs', () => { + const explicitSizes = parseSplitAttributes({ axis: 0, numOutputs: 3, splitSizes: [0, 2, 0] }); + expect(runSplit([2], explicitSizes).map((output) => output.dims)).to.deep.equal([[0], [2], [0]]); + }); + + it('rejects invalid output counts', () => { + const invalidCount = parseSplitAttributes({ axis: 0, numOutputs: 0, splitSizes: [], isUnevenSplitAllowed: true }); + expect(() => runSplit([4], invalidCount)).to.throw('numOutputs must be positive'); + const threeOutputs = parseSplitAttributes({ axis: 0, numOutputs: 3, splitSizes: [], isUnevenSplitAllowed: true }); + expect(() => runSplit([1], threeOutputs)).to.throw('negative split size'); + const tooManyChunks = parseSplitAttributes({ axis: 0, numOutputs: 7, splitSizes: [], isUnevenSplitAllowed: true }); + expect(() => runSplit([10], tooManyChunks)).to.throw('negative split size'); + }); +}); diff --git a/onnxruntime/core/providers/coreml/builders/impl/split_op_builder.cc b/onnxruntime/core/providers/coreml/builders/impl/split_op_builder.cc index 145682f0f3524..4bf93f9a6b4a0 100644 --- a/onnxruntime/core/providers/coreml/builders/impl/split_op_builder.cc +++ b/onnxruntime/core/providers/coreml/builders/impl/split_op_builder.cc @@ -63,7 +63,7 @@ Status SplitOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, std::unique_ptr split_op = model_builder.CreateOperation(node, "split"); AddOperationInput(*split_op, "axis", model_builder.AddScalarConstant(split_op->type(), "axis", axis)); - if (input_defs.size() > 1) { + if (input_defs.size() > 1 && input_defs[1]->Exists()) { // if "split" is explicitly provided as an input const auto& const_init = *model_builder.GetConstantInitializer(input_defs[1]->Name()); const Initializer unpacked_tensor(model_builder.GetGraphViewer().GetGraph(), const_init, @@ -107,7 +107,7 @@ Status SplitOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, auto* coreml_splitnd = layer->mutable_splitnd(); coreml_splitnd->set_axis(axis); - if (input_defs.size() > 1) { + if (input_defs.size() > 1 && input_defs[1]->Exists()) { // if "split" is explicitly provided as an input // const auto& split_tensor = *model_builder.GetInitializerTensors().at(input_defs[1]->Name()); const auto& const_init = *model_builder.GetConstantInitializer(input_defs[1]->Name()); diff --git a/onnxruntime/core/providers/cpu/tensor/split.cc b/onnxruntime/core/providers/cpu/tensor/split.cc index 92911103280e7..fba93b34030c7 100644 --- a/onnxruntime/core/providers/cpu/tensor/split.cc +++ b/onnxruntime/core/providers/cpu/tensor/split.cc @@ -54,7 +54,6 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL( BuildKernelDefConstraintsFromTypeList()), Split_1_13); -// TODO: support unequal split and num_outputs ONNX_CPU_OPERATOR_KERNEL( Split, 18, diff --git a/onnxruntime/core/providers/cpu/tensor/split.h b/onnxruntime/core/providers/cpu/tensor/split.h index 02e425475228f..fe6c5545d343f 100644 --- a/onnxruntime/core/providers/cpu/tensor/split.h +++ b/onnxruntime/core/providers/cpu/tensor/split.h @@ -20,7 +20,7 @@ namespace onnxruntime { class SplitBase { public: /* - * \param num_outputs must >=0 + * \param num_outputs must >0 */ #ifdef SHARED_PROVIDER Status PrepareForCompute(const TensorShape& input_shape, int num_outputs, int64_t& axis, int& before_dims, @@ -30,6 +30,9 @@ class SplitBase { inline Status PrepareForCompute(const TensorShape& input_shape, int num_outputs, int64_t& axis, int& before_dims, int& after_dims_including_split_axis, int& after_dims_excluding_split, std::vector& split_sizes) const { + if (num_outputs <= 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Number of outputs must be positive."); + } auto input_dims = input_shape.GetDims(); const auto num_dimensions = gsl::narrow_cast(input_shape.NumDimensions()); axis = HandleNegativeAxis(axis_, num_dimensions); @@ -42,16 +45,7 @@ class SplitBase { : onnxruntime::narrow(input_shape.SizeFromDimension(SafeInt(axis) + 1)); if (num_outputs_ != -1) { - if (num_outputs_ > split_dim_size) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Invalid num_outputs value of ", num_outputs_, - ". Size of dimension being split is ", split_dim_size); - } - int32_t size = onnxruntime::narrow(std::ceil(static_cast(split_dim_size) / num_outputs)); - int32_t remainder = split_dim_size % size; - split_sizes = std::vector(num_outputs, size); - if (remainder) { - split_sizes.back() = remainder; - } + ORT_RETURN_IF_ERROR(PrepareSplitSizesForNumOutputs(split_dim_size, num_outputs, split_sizes)); } if (split_sizes.empty()) { @@ -90,11 +84,38 @@ class SplitBase { #endif // SHARED_PROVIDER protected: + Status PrepareSplitSizesForNumOutputs(int64_t split_dim_size, int num_outputs, + std::vector& split_sizes) const { + if (num_outputs <= 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Number of outputs must be positive."); + } + if (num_outputs_ != num_outputs) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Number of outputs (", num_outputs, + ") does not match num_outputs (", num_outputs_, ")."); + } + const int64_t split_size = split_dim_size / num_outputs + (split_dim_size % num_outputs != 0); + const int64_t last_split_size = + SafeInt(split_dim_size) - SafeInt(split_size) * (num_outputs - 1); + if (last_split_size < 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Invalid num_outputs value of ", num_outputs_, + ". Split sizes would produce a negative final split."); + } + split_sizes.assign(static_cast(num_outputs), split_size); + split_sizes.back() = last_split_size; + return Status::OK(); + } + template SplitBase(const KernelInfoType& info, uint32_t opset) : opset_{opset} { axis_ = info.template GetAttrOrDefault("axis", 0); size_t num_inputs = info.GetInputCount(); + if constexpr (requires { info.node().InputExists(size_t{1}); }) { + has_split_input_ = info.node().InputExists(1); + } else { + has_split_input_ = num_inputs > 1 && info.node().InputDefs()[1]->Exists(); + } + if (num_inputs == 1) { // optional if (info.GetAttrs("split", split_sizes_).IsOK()) { @@ -109,7 +130,7 @@ class SplitBase { // the ONNX type/shape inferencing handles the check that num_outputs is > 0 // ORT_ENFORCE(num_outputs_ != 0, "Invalid value in 'num_outputs' attribute of 0."); - if (num_outputs_ != -1 && info.GetInputCount() == 2) { + if (num_outputs_ != -1 && has_split_input_) { ORT_THROW("If 'num_outputs' is specified, the 'split' input should not be provided."); } } @@ -118,6 +139,7 @@ class SplitBase { const uint32_t opset_; int64_t axis_; std::vector split_sizes_; + bool has_split_input_ = false; int64_t split_size_sum_ = -1; int64_t num_outputs_ = -1; }; diff --git a/onnxruntime/core/providers/cuda/tensor/split.cc b/onnxruntime/core/providers/cuda/tensor/split.cc index 5ef4e9e8250b5..d3143df0016db 100644 --- a/onnxruntime/core/providers/cuda/tensor/split.cc +++ b/onnxruntime/core/providers/cuda/tensor/split.cc @@ -4,6 +4,7 @@ #include #include +#include "core/common/safeint.h" #include "core/providers/cuda/tensor/split.h" #include "core/providers/cuda/tensor/split_impl.h" @@ -51,7 +52,7 @@ ONNX_OPERATOR_KERNEL_EX(Split, // Self-contained split preparation that replaces SplitBase::PrepareForCompute. // The base class method is not available in plugin builds because SplitBase -// depends on CPU provider internals. Keep logic in sync with SplitBase. +// depends on CPU provider internals. Share num_outputs partitioning with SplitBase. Status SplitKernel::PrepareForComputeLocal(const TensorShape& input_shape, int num_outputs, int64_t& axis, @@ -59,6 +60,9 @@ Status SplitKernel::PrepareForComputeLocal(const TensorShape& input_shape, int& after_dims_including_split_axis, int& after_dims_excluding_split, std::vector& split_sizes) const { + if (num_outputs <= 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Number of outputs must be positive."); + } auto input_dims = input_shape.GetDims(); const auto num_dimensions = gsl::narrow_cast(input_shape.NumDimensions()); axis = HandleNegativeAxis(axis_, num_dimensions); @@ -68,21 +72,10 @@ Status SplitKernel::PrepareForComputeLocal(const TensorShape& input_shape, after_dims_including_split_axis = onnxruntime::narrow(input_shape.SizeFromDimension(onnxruntime::narrow(axis))); after_dims_excluding_split = (axis + 1 == num_dimensions) ? 1 - : onnxruntime::narrow(input_shape.SizeFromDimension(onnxruntime::narrow(axis + 1))); + : onnxruntime::narrow(input_shape.SizeFromDimension(SafeInt(axis) + 1)); if (num_outputs_ != -1) { - if (num_outputs_ > split_dim_size) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Invalid num_outputs value of ", num_outputs_, - ". Size of dimension being split is ", split_dim_size); - } - - int64_t size = (split_dim_size + num_outputs_ - 1) / num_outputs_; - int64_t remainder = split_dim_size % size; - - split_sizes = std::vector(num_outputs, size); - if (remainder) { - split_sizes.back() = remainder; - } + ORT_RETURN_IF_ERROR(PrepareSplitSizesForNumOutputs(split_dim_size, num_outputs, split_sizes)); } if (split_sizes.empty()) { diff --git a/onnxruntime/core/providers/js/operators/split.h b/onnxruntime/core/providers/js/operators/split.h index 3f6cfcb8921f3..71063a00eccd3 100644 --- a/onnxruntime/core/providers/js/operators/split.h +++ b/onnxruntime/core/providers/js/operators/split.h @@ -13,6 +13,7 @@ namespace js { class Split : public JsKernel, public SplitBase { public: Split(const OpKernelInfo& info, uint32_t opset) : JsKernel(info), SplitBase(info, opset) { + const bool is_uneven_split_allowed = num_outputs_ >= 0; std::vector split_sizes; if (split_sizes_.size() > 0) { ORT_ENFORCE(split_sizes_.size() == info.node().OutputDefs().size(), @@ -25,11 +26,7 @@ class Split : public JsKernel, public SplitBase { if (num_outputs_ < 0) { num_outputs_ = split_sizes.size(); } - } else if (split_sizes_.size() == 0 && info.GetInputCount() < 2) { - // Compute split_sizes from input shape and num_outputs. - // TODO: Shape might not be known at this point, better to handle this in javascript - auto total_split_size = info.node().InputDefs()[0]->Shape()->dim(gsl::narrow_cast(axis_)).dim_value(); - int64_t split_size_sum = 0; + } else { if (num_outputs_ < 0) { num_outputs_ = info.node().OutputDefs().size(); } else { @@ -37,23 +34,17 @@ class Split : public JsKernel, public SplitBase { "Number of outputs (", info.node().OutputDefs().size(), ") does not match num_outputs (", num_outputs_, ")"); } - for (auto output : info.node().OutputDefs()) { - auto split_size = output->Shape()->dim(gsl::narrow_cast(axis_)).dim_value(); - split_sizes.push_back(gsl::narrow_cast(split_size)); - split_size_sum += split_size; - } - ORT_ENFORCE(split_size_sum == total_split_size, - "Sum of split sizes (", split_size_sum, ") does not match input size (", total_split_size, ")"); } - // else: let javascript handle all other cases, ie. split_sizes come as input[1] JSEP_INIT_KERNEL_ATTRIBUTE(Split, ({"axis" : $1, "numOutputs" : $2, - "splitSizes" : $3 ? Array.from(HEAP32.subarray(Number($3), Number($4))) : []}), + "splitSizes" : $3 ? Array.from(HEAP32.subarray(Number($3), Number($4))) : [], + "isUnevenSplitAllowed" : !!$5}), static_cast(axis_), static_cast(num_outputs_), JSEP_HEAP32_INDEX_START(split_sizes), - JSEP_HEAP32_INDEX_END(split_sizes)); + JSEP_HEAP32_INDEX_END(split_sizes), + static_cast(is_uneven_split_allowed)); } }; diff --git a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc index 0684b6d573968..63decd7a168de 100644 --- a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc @@ -813,6 +813,7 @@ TEST(SplitOperatorTest, Split18_NumOutputs_EvenSplit) { RunTest(axis, {}, input, outputs, {kTensorrtExecutionProvider}, false, true, num_outputs, true); RunTest(axis, {}, input, outputs, {kTensorrtExecutionProvider}, false, true, num_outputs, false); + RunTest(axis, {}, input, outputs, {kTensorrtExecutionProvider}, false, true, num_outputs, true, {}, false); } TEST(SplitOperatorTest, Split18_NumOutputs_UnevenSplit) { @@ -841,6 +842,56 @@ TEST(SplitOperatorTest, Split18_NumOutputs_UnevenSplit) { RunTest(axis, {}, input, outputs, {kTensorrtExecutionProvider, kQnnExecutionProvider}, false, true, num_outputs, true); RunTest(axis, {}, input, outputs, {kTensorrtExecutionProvider, kQnnExecutionProvider}, false, true, num_outputs, false); + RunTest(axis, {}, input, outputs, {kTensorrtExecutionProvider, kQnnExecutionProvider}, false, true, + num_outputs, true, {}, false); +} + +TEST(SplitOperatorTest, Split18_NumOutputs_ZeroFinalOutput) { + const ShapeAndFloatData input = {{10}, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f, 10.f}}; + const std::vector outputs = { + {{2}, {1.f, 2.f}}, {{2}, {3.f, 4.f}}, {{2}, {5.f, 6.f}}, {{2}, {7.f, 8.f}}, {{2}, {9.f, 10.f}}, {{0}, {}}}; + const std::unordered_set excluded_providers = { + kTensorrtExecutionProvider, kQnnExecutionProvider, kCoreMLExecutionProvider, kDmlExecutionProvider}; + + RunTest(0, {}, input, outputs, excluded_providers, false, true, 6); + RunTest(0, {}, input, outputs, excluded_providers, false, true, 6, true, {}, false); + + const ShapeAndFloatData outer_axis_input = + {{10, 2}, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f, 10.f, 11.f, 12.f, 13.f, 14.f, 15.f, 16.f, 17.f, 18.f, 19.f, 20.f}}; + const std::vector outer_axis_outputs = { + {{2, 2}, {1.f, 2.f, 3.f, 4.f}}, {{2, 2}, {5.f, 6.f, 7.f, 8.f}}, {{2, 2}, {9.f, 10.f, 11.f, 12.f}}, {{2, 2}, {13.f, 14.f, 15.f, 16.f}}, {{2, 2}, {17.f, 18.f, 19.f, 20.f}}, {{0, 2}, {}}}; + RunTest(0, {}, outer_axis_input, outer_axis_outputs, excluded_providers, false, true, 6); + RunTest(0, {}, outer_axis_input, outer_axis_outputs, excluded_providers, false, true, 6, true, {}, false); +} + +TEST(SplitOperatorTest, Split18_NumOutputs_OneElement) { + const ShapeAndFloatData input = {{1}, {1.f}}; + const std::vector outputs = {{{1}, {1.f}}, {{0}, {}}}; + const std::unordered_set excluded_providers = { + kTensorrtExecutionProvider, kQnnExecutionProvider, kCoreMLExecutionProvider, kDmlExecutionProvider}; + + RunTest(0, {}, input, outputs, excluded_providers, false, true, 2); + RunTest(0, {}, input, outputs, excluded_providers, false, true, 2, true, {}, false); +} + +TEST(SplitOperatorTest, Split18_NumOutputs_ZeroSplitDimension) { + const ShapeAndFloatData input = {{2, 0}, {}}; + const std::vector outputs = {{{2, 0}, {}}, {{2, 0}, {}}}; + const std::unordered_set excluded_providers = { + kTensorrtExecutionProvider, kQnnExecutionProvider, kCoreMLExecutionProvider, kDmlExecutionProvider}; + + RunTest(1, {}, input, outputs, excluded_providers, false, true, 2); + RunTest(1, {}, input, outputs, excluded_providers, false, true, 2, true, {}, false); +} + +TEST(SplitOperatorTest, Split18_NumOutputs_LargeSplitDimension) { + const ShapeAndFloatData input = {{0, 16777217}, {}}; + const std::vector outputs = {{{0, 8388609}, {}}, {{0, 8388608}, {}}}; + const std::unordered_set excluded_providers = { + kTensorrtExecutionProvider, kQnnExecutionProvider, kCoreMLExecutionProvider, kDmlExecutionProvider}; + + RunTest(1, {}, input, outputs, excluded_providers, false, true, 2); + RunTest(1, {}, input, outputs, excluded_providers, false, true, 2, true, {}, false); } TEST(SplitOperatorTest, Split18_InvalidNumOutputs) { @@ -877,9 +928,9 @@ TEST(SplitOperatorTest, Split18_InvalidNumOutputs) { num_outputs = 3; RunTest(axis, {}, input, outputs, excluded_providers, true, true, num_outputs, false, - "Invalid num_outputs value of 3. Size of dimension being split is 2"); + "Number of outputs (2) does not match num_outputs (3)"); RunTest(axis, {}, input, outputs, excluded_providers, true, true, num_outputs, true, - "Invalid num_outputs value of 3. Size of dimension being split is 2"); + "Number of outputs (2) does not match num_outputs (3)"); } TEST(SplitOperatorTest, Split18_NumOutputsEvenSplitAxis1) {