Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 55 additions & 4 deletions js/web/lib/onnxjs/util.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1260,6 +1260,8 @@ export class PoolConvUtil {
* @param pads Padding for the beginning and ending along each axis.
* @param autoPad DEPRECATED attribute supported for legacy models. Specifies how to implicitly calculate pads in each
* dimension. Can take values NOTSET, SAME_UPPER, SAME_LOWER, or VALID.
* @param ceilMode When set to 1, use ceil() instead of floor() to compute the output spatial size (and apply the
* "shrink the last window if it starts entirely in padding" rule). Defaults to 0 (floor).
*/
static computePoolOutputShape(
isGlobalOperator: boolean,
Expand All @@ -1269,6 +1271,7 @@ export class PoolConvUtil {
kernelShape: number[],
pads: number[],
autoPad?: string,
ceilMode = 0,
): number[] {
if (inputDims.length <= 0) {
throw new Error('input shape must be of size greater than 0');
Expand All @@ -1286,6 +1289,7 @@ export class PoolConvUtil {
kernelShape,
pads,
autoPad,
ceilMode,
);
return outputDims;
}
Expand Down Expand Up @@ -1332,6 +1336,7 @@ export class PoolConvUtil {
kernelShape: readonly number[],
pads: number[],
autoPad?: string,
ceilMode = 0,
) {
if (isGlobalOperator) {
for (let dim = 0; dim < inputDims.length - 2; dim++) {
Expand All @@ -1349,12 +1354,44 @@ export class PoolConvUtil {
dim,
dim + inputDims.length - 2,
autoPad,
ceilMode,
),
);
}
}
}

// Computes the output spatial size for a single dimension.
// Produces results identical to the C++ PoolAttributes::ComputeOutputSize
// (onnxruntime/core/providers/cpu/nn/pool_attributes.h), including the ceil_mode
// "shrink the last window if it starts entirely in the trailing padding" rule. The JS
// signature takes a pre-computed `numerator` (equal to `inSize + padHead + padTail - dkernel`,
// matching the C++ `in_size + pad_head + pad_tail - dilation * (kernel - 1) - 1`) instead of
// the raw pooling attributes, but the computed output size is the same.
// Keep in sync with the onnxjs/jsep copy.
// NOTE: In this onnxjs copy the ceilMode path exists for shape-test parity with the jsep copy;
// the onnxjs WebGL pooling caller (backends/webgl/ops/pool.ts) intentionally does NOT pass
// ceilMode (legacy path still throws on ceil_mode != 0), so it always uses the floor default.
private static computeOutputSize(
numerator: number,
stride: number,
inSize: number,
padHead: number,
ceilMode: number,
): number {
let outSize = Math.floor(numerator / stride) + 1;
// Match C++ `ceil_mode == 1` exactly so out-of-spec ceil_mode values do not diverge.
if (ceilMode === 1) {
outSize = Math.ceil(numerator / stride) + 1;
// Ensure the last pooling window starts inside the image (ref: https://github.com/onnx/onnx/pull/5741).
// inSize and padHead are needed here to reconstruct the last window's start position.
if ((outSize - 1) * stride >= inSize + padHead) {
outSize -= 1;
}
}
return outSize;
}

// helper for computeShapeHelper() and adjustPadsBasedOnAutoPad()
// adjusts pad value for given 'autoPad' string and computes output shape along a particular dimension
private static adjustPadAndReturnShape(
Expand All @@ -1366,30 +1403,44 @@ export class PoolConvUtil {
padHeadIndex: number,
padTailIndex: number,
autoPad?: string,
ceilMode = 0,
): number {
const dkernel = dilation * (kernel - 1) + 1;
if (autoPad && autoPad !== 'NOTSET') {
switch (autoPad) {
case 'VALID':
pads[padHeadIndex] = 0;
pads[padTailIndex] = 0;
return Math.floor((inSize - dkernel) / stride + 1);
return PoolConvUtil.computeOutputSize(inSize - dkernel, stride, inSize, 0, ceilMode);
case 'SAME_LOWER':
case 'SAME_UPPER':
if (dilation !== 1) {
throw new Error('Dilation not supported for SAME_UPPER or SAME_LOWER');
} else {
const legacyTargetSize = (inSize + stride - 1) / stride;
// Integer division to match C++ pool_attributes.h ComputeSizePadDilations; float division mis-rounds SAME_* pads.
const legacyTargetSize = Math.floor((inSize + stride - 1) / stride);
const padNeeded = (legacyTargetSize - 1) * stride + kernel - inSize;
pads[padHeadIndex] = autoPad === 'SAME_LOWER' ? Math.floor((padNeeded + 1) / 2) : Math.floor(padNeeded / 2);
pads[padTailIndex] = padNeeded - pads[padHeadIndex];
return Math.floor((inSize + padNeeded - kernel) / stride + 1);
return PoolConvUtil.computeOutputSize(
inSize + pads[padHeadIndex] + pads[padTailIndex] - dkernel,
stride,
inSize,
pads[padHeadIndex],
ceilMode,
);
}
default:
throw new Error('Unsupported AutoPad type');
}
} else {
return Math.floor((inSize + pads[padHeadIndex] + pads[padTailIndex] - dkernel) / stride + 1);
return PoolConvUtil.computeOutputSize(
inSize + pads[padHeadIndex] + pads[padTailIndex] - dkernel,
stride,
inSize,
pads[padHeadIndex],
ceilMode,
);
}
}
}
Expand Down
56 changes: 52 additions & 4 deletions js/web/lib/wasm/jsep/util.ts
Original file line number Diff line number Diff line change
Expand Up @@ -366,6 +366,8 @@ export class PoolConvUtil {
* @param pads Padding for the beginning and ending along each axis.
* @param autoPad DEPRECATED attribute supported for legacy models. Specifies how to implicitly calculate pads in each
* dimension. Can take values NOTSET, SAME_UPPER, SAME_LOWER, or VALID.
* @param ceilMode When set to 1, use ceil() instead of floor() to compute the output spatial size (and apply the
* "shrink the last window if it starts entirely in padding" rule). Defaults to 0 (floor).
*/
static computePoolOutputShape(
isGlobalOperator: boolean,
Expand All @@ -375,6 +377,7 @@ export class PoolConvUtil {
kernelShape: number[],
pads: number[],
autoPad?: string,
ceilMode = 0,
): number[] {
if (inputDims.length <= 0) {
throw new Error('input shape must be of size greater than 0');
Expand All @@ -392,6 +395,7 @@ export class PoolConvUtil {
kernelShape,
pads,
autoPad,
ceilMode,
);
return outputDims;
}
Expand Down Expand Up @@ -438,6 +442,7 @@ export class PoolConvUtil {
kernelShape: readonly number[],
pads: number[],
autoPad?: string,
ceilMode = 0,
) {
if (isGlobalOperator) {
for (let dim = 0; dim < inputDims.length - 2; dim++) {
Expand All @@ -455,12 +460,41 @@ export class PoolConvUtil {
dim,
dim + inputDims.length - 2,
autoPad,
ceilMode,
),
);
}
}
}

// Computes the output spatial size for a single dimension.
// Produces results identical to the C++ PoolAttributes::ComputeOutputSize
// (onnxruntime/core/providers/cpu/nn/pool_attributes.h), including the ceil_mode
// "shrink the last window if it starts entirely in the trailing padding" rule. The JS
// signature takes a pre-computed `numerator` (equal to `inSize + padHead + padTail - dkernel`,
// matching the C++ `in_size + pad_head + pad_tail - dilation * (kernel - 1) - 1`) instead of
// the raw pooling attributes, but the computed output size is the same.
// Keep in sync with the onnxjs/jsep copy.
private static computeOutputSize(
numerator: number,
stride: number,
inSize: number,
padHead: number,
ceilMode: number,
): number {
let outSize = Math.floor(numerator / stride) + 1;
// Match C++ `ceil_mode == 1` exactly so out-of-spec ceil_mode values do not diverge.
if (ceilMode === 1) {
outSize = Math.ceil(numerator / stride) + 1;
// Ensure the last pooling window starts inside the image (ref: https://github.com/onnx/onnx/pull/5741).
// inSize and padHead are needed here to reconstruct the last window's start position.
if ((outSize - 1) * stride >= inSize + padHead) {
outSize -= 1;
}
}
return outSize;
}

// helper for computeShapeHelper() and adjustPadsBasedOnAutoPad()
// adjusts pad value for given 'autoPad' string and computes output shape along a particular dimension
private static adjustPadAndReturnShape(
Expand All @@ -472,30 +506,44 @@ export class PoolConvUtil {
padHeadIndex: number,
padTailIndex: number,
autoPad?: string,
ceilMode = 0,
): number {
const dkernel = dilation * (kernel - 1) + 1;
if (autoPad && autoPad !== 'NOTSET') {
switch (autoPad) {
case 'VALID':
pads[padHeadIndex] = 0;
pads[padTailIndex] = 0;
return Math.floor((inSize - dkernel) / stride + 1);
return PoolConvUtil.computeOutputSize(inSize - dkernel, stride, inSize, 0, ceilMode);
case 'SAME_LOWER':
case 'SAME_UPPER':
if (dilation !== 1) {
throw new Error('Dilation not supported for SAME_UPPER or SAME_LOWER');
} else {
const legacyTargetSize = (inSize + stride - 1) / stride;
// Integer division to match C++ pool_attributes.h ComputeSizePadDilations; float division mis-rounds SAME_* pads.
const legacyTargetSize = Math.floor((inSize + stride - 1) / stride);
const padNeeded = (legacyTargetSize - 1) * stride + kernel - inSize;
pads[padHeadIndex] = autoPad === 'SAME_LOWER' ? Math.floor((padNeeded + 1) / 2) : Math.floor(padNeeded / 2);
pads[padTailIndex] = padNeeded - pads[padHeadIndex];
return Math.floor((inSize + padNeeded - kernel) / stride + 1);
return PoolConvUtil.computeOutputSize(
inSize + pads[padHeadIndex] + pads[padTailIndex] - dkernel,
stride,
inSize,
pads[padHeadIndex],
ceilMode,
);
}
default:
throw new Error('Unsupported AutoPad type');
}
} else {
return Math.floor((inSize + pads[padHeadIndex] + pads[padTailIndex] - dkernel) / stride + 1);
return PoolConvUtil.computeOutputSize(
inSize + pads[padHeadIndex] + pads[padTailIndex] - dkernel,
stride,
inSize,
pads[padHeadIndex],
ceilMode,
);
}
}
}
Expand Down
25 changes: 20 additions & 5 deletions js/web/lib/wasm/jsep/webgpu/ops/pool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,8 @@ import {
} from './common';

// TODO: support:
// - ceil_mode "test_maxpool_2d_ceil"
// - ceil_mode kernel execution "test_maxpool_2d_ceil" (output SHAPE already honors ceil_mode
// via PoolConvUtil.computePoolOutputShape; the WebGPU kernel padding/divisor handling is pending)
// - storage_order "test_maxpool_with_argmax_2d_precomputed_strides"
// - [MaxPool] dilations "test_maxpool_2d_dilations"
// - [MaxPool] output[1] "test_maxpool_with_argmax_2d_precomputed_pads"
Expand Down Expand Up @@ -56,6 +57,7 @@ const getAdjustedPoolAttributesAndOutputShape = <AttributeType extends AveragePo
kernelShape,
pads,
attributes.autoPad,
attributes.ceilMode,
);

const newAttributes = Object.assign({}, attributes);
Expand Down Expand Up @@ -391,9 +393,15 @@ export const parseAveragePoolAttributes = (attributes: Record<string, unknown>):
const countIncludePad = (attributes.count_include_pad as number) === 0 ? false : true;

const attr = parsePoolCommonAttributes(attributes);
// TODO: support attribute 'ceil_mode'
// ceil_mode is honored by the output-shape math (PoolConvUtil.computePoolOutputShape now
// threads ceilMode and matches the C++ reference), but end-to-end execution stays gated
// here: the WebGPU pooling kernel must also apply the ceil_mode trailing-padding handling
// (and, for AveragePool, the count_include_pad divisor) before this throw can be removed.
// Tracked follow-up: remove this guard + add kernel ceil_mode padding support.
if (attr.ceilMode !== 0) {
throw new Error('using ceil() in shape computation is not yet supported for AveragePool');
throw new Error(
'ceil_mode output-shape is computed, but ceil_mode kernel execution (padding/divisor) is not yet implemented in the WebGPU AveragePool kernel',
);
}
const averagePoolAttributes = { countIncludePad, ...attr, cacheKey: '' };
return { ...averagePoolAttributes, cacheKey: createAveragePoolShaderKeyFromAttributes(averagePoolAttributes) };
Expand Down Expand Up @@ -491,12 +499,19 @@ export const parseMaxPoolAttributes = (attributes: Record<string, unknown>): Max
const dilations = attributes.dilations as [number, number];

const attr = parsePoolCommonAttributes(attributes);
// TODO: support attribute 'ceil_mode' and 'storage_order'
// TODO: support attribute 'storage_order'
if (storageOrder !== 0) {
throw new Error('column major storage order is not yet supported for MaxPool');
}
// ceil_mode is honored by the output-shape math (PoolConvUtil.computePoolOutputShape now
// threads ceilMode and matches the C++ reference), but end-to-end execution stays gated
// here: the WebGPU pooling kernel must also apply the ceil_mode trailing-padding handling
// before this throw can be removed. Tracked follow-up: remove this guard + add kernel
// ceil_mode padding support.
if (attr.ceilMode !== 0) {
throw new Error('using ceil() in shape computation is not yet supported for MaxPool');
throw new Error(
'ceil_mode output-shape is computed, but ceil_mode kernel execution (padding) is not yet implemented in the WebGPU MaxPool kernel',
);
}
const maxPoolAttributes = { storageOrder, dilations, ...attr, cacheKey: '' };
return { ...maxPoolAttributes, cacheKey: createMaxPoolShaderKeyFromAttributes(maxPoolAttributes) };
Expand Down
2 changes: 2 additions & 0 deletions js/web/test/unittests/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,4 +13,6 @@ if (typeof window !== 'undefined') {

require('./backends/wasm/test-model-metadata');

require('./pool-output-shape');

require('./opset');
Loading
Loading