[js/webgpu] Fix Expand/Gather when input type is bool (#18999)

### Description
Also update the op test suite.

### Motivation and Context
Previously the *total* size in case `Expand - last dim is not divisible
by 4` was a multiple of 4, even though the *last dimension* was not, so
the bug has never been caught.
This commit is contained in:
Jiajie Hu 2024-01-06 00:16:15 +08:00 committed by GitHub
parent 7f0aac0d8a
commit 447a3a7c70
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 48 additions and 7 deletions

View file

@ -47,7 +47,7 @@ const createExpandProgramInfo = (inputs: readonly TensorView[]): ProgramInfo =>
const outputShape: number[] = calculateOutputShape(inputShape, shape);
const dataType = inputs[0].dataType;
const components = dataType === DataType.bool ? 4 : 1;
const outputSize = ShapeUtil.size(outputShape) / components;
const outputSize = Math.ceil(ShapeUtil.size(outputShape) / components);
const enableInputShapeUniform = enableShapesUniforms(inputShape.length);
const enableOutputShapeUniform = enableShapesUniforms(outputShape.length);

View file

@ -31,7 +31,7 @@ const createGatherProgramInfo = (inputs: readonly TensorView[], attributes: Gath
const axisDimLimit = inputShape[axis];
const components = inputs[0].dataType === DataType.bool ? 4 : 1;
const outputSize = ShapeUtil.size(outputShape) / components;
const outputSize = Math.ceil(ShapeUtil.size(outputShape) / components);
const enableInputShapesUniforms = enableShapesUniforms(inputs[0].dims.length);
const inputShapeOrRank = enableInputShapesUniforms ? inputs[0].dims.length : inputs[0].dims;

View file

@ -168,20 +168,39 @@
"name": "Expand - last dim is not divisible by 4",
"inputs": [
{
"data": [true, false, false, true, true, true, false, false, false, true, true, true],
"dims": [2, 6],
"data": [true, false, false, true, true, true],
"dims": [1, 6],
"type": "bool"
},
{
"data": [2, 1],
"data": [3, 1],
"dims": [2],
"type": "int64"
}
],
"outputs": [
{
"data": [true, false, false, true, true, true, false, false, false, true, true, true],
"dims": [2, 6],
"data": [
true,
false,
false,
true,
true,
true,
true,
false,
false,
true,
true,
true,
true,
false,
false,
true,
true,
true
],
"dims": [3, 6],
"type": "bool"
}
]

View file

@ -99,6 +99,28 @@
"operator": "Gather",
"attributes": [],
"cases": [
{
"name": "data[4] indices[]",
"inputs": [
{
"data": [false, true, false, false],
"dims": [4],
"type": "bool"
},
{
"data": [1],
"dims": [],
"type": "int32"
}
],
"outputs": [
{
"data": [true],
"dims": [],
"type": "bool"
}
]
},
{
"name": "data[2,4] indices[1]",
"inputs": [