firefox-main/testing/web-platform/tests/webnn/conformance_tests/resample2d-gather-shape-divergence.https.any.js (file symbol)

Source code

Revision control

Copy as Markdown

Other Tools

Test Info:

// META: title=test resample2d output shape agrees with the backend across a gather
// META: global=window
// META: variant=?cpu
// META: variant=?gpu
// META: variant=?npu
// META: script=../resources/utils.js
// META: timeout=long
'use strict';
//
// Regression test for a resample2d shape divergence. The output size is
// floor(input size * scale), which WebNN validates in double precision. A
// backend that re-derives it in float32 disagrees above 2^24, where float32
// cannot represent consecutive integers: for size 2^24 + 1 and scale 1.0 it
// rounds down to 2^24, under-allocating the axis by one element. A downstream
// gather, clamped to WebNN's (larger) shape, then reads out of bounds at the
// last index.
//
// The assertion does not check for a specific value, since backends may sample
// different values near 2^24. Instead it exposes `resampled` as a second output
// and checks that gathering the last index reads the same element as a direct
// read of it. That holds only when the gather stayed in bounds, i.e. when the
// shapes agree; an under-allocated axis clamps the gather elsewhere and breaks
// the equality.
// 2^24 + 1: the smallest positive integer float32 cannot represent exactly.
// Input byte length is 4 * 16777217 ~= 64 MiB, below the tensor byte length
// limit.
const kResampledDim = 16777217;
const kLastIndex = kResampledDim - 1; // Largest index WebNN admits: 16777216.
let mlContext;
promise_setup(async () => {
assert_implements(navigator.ml, 'missing navigator.ml');
mlContext = await navigator.ml.createContext(contextOptions);
});
promise_test(async () => {
const builder = new MLGraphBuilder(mlContext);
// Rank-4 input with the large dimension on axis 2.
const input = builder.input(
'input', {dataType: 'float32', shape: [1, 1, kResampledDim, 1]});
// Identity resample (scale 1.0). WebNN's validated output shape keeps axis 2
// at kResampledDim; a float32 re-derivation would shrink it by one.
const resampled = builder.resample2d(
input, {mode: 'nearest-neighbor', scales: [1, 1], axes: [2, 3]});
assert_equals(resampled.shape[2], kResampledDim,
'resample2d output preserves the axis-2 dimension');
// Gather the last WebNN-valid element along the resampled axis.
const indices = builder.input('indices', {dataType: 'int32', shape: [1]});
const gathered = builder.gather(resampled, indices, {axis: 2});
assert_array_equals(gathered.shape, [1, 1, 1, 1], 'gather output shape');
const [inputTensor, indicesTensor, resampledTensor, gatheredTensor, mlGraph] =
await Promise.all([
mlContext.createTensor({
dataType: 'float32',
shape: [1, 1, kResampledDim, 1],
writable: true,
}),
mlContext.createTensor({dataType: 'int32', shape: [1], writable: true}),
mlContext.createTensor({
dataType: 'float32',
shape: [1, 1, kResampledDim, 1],
readable: true,
}),
mlContext.createTensor(
{dataType: 'float32', shape: [1, 1, 1, 1], readable: true}),
// Expose `resampled` as an output so its last element can be read back
// and compared against the gather result.
builder.build({'resampled': resampled, 'gathered': gathered}),
]);
// Distinct non-zero sentinels at kLastIndex and kLastIndex - 1. On a
// correctly sized backend the gather reads kLastIndex in bounds, so
// gathered[0] equals a direct read of resampled[kLastIndex] regardless of the
// values. If the axis is under-allocated, both the gather and that direct
// read touch an element the backend never produced; the resulting behavior is
// backend-defined (it may throw, clamp, or read adjacent memory). The
// distinct sentinels give a mismatch something to show - e.g. a backend that
// clamps the gather to its last valid index would read a different sentinel
// than the unwritten kLastIndex. This is best-effort: a zero-filled input
// could let both sides read 0 and hide the regression. The exact values are
// never asserted; they only need to differ.
const inputData = new Float32Array(kResampledDim);
inputData[kLastIndex - 1] = 5;
inputData[kLastIndex] = 42;
mlContext.writeTensor(inputTensor, inputData);
mlContext.writeTensor(indicesTensor, new Int32Array([kLastIndex]));
mlContext.dispatch(
mlGraph, {'input': inputTensor, 'indices': indicesTensor},
{'resampled': resampledTensor, 'gathered': gatheredTensor});
const resampledData =
new Float32Array(await mlContext.readTensor(resampledTensor));
const gatheredData =
new Float32Array(await mlContext.readTensor(gatheredTensor));
assert_equals(
gatheredData[0], resampledData[kLastIndex],
'gather at the last WebNN-valid index reads the same element as a direct ' +
'read of that index, proving the resample2d output shape matches the ' +
'backend allocation (no out-of-bounds clamp)');
}, 'resample2d output shape (from scales) matches backend allocation at the 2^24 boundary');