[js/webgpu] Optimize InstanceNormalization (#17491)
### Description
<!-- Describe your changes. -->
In previous implementation, there are two loops to iterate H * W
elements to calculate the `mean` and `squaredNorm` value in one thread,
meanwhile it outputs H * W elements in one thread. That results it's
very very slow when H * W is a large value. And usually, H * W does be a
large value in a model. For example, in the `candy-8` model, the shapes
of [H, W] are [224,224], [112,112], [56,56] for `InstanceNormalization`
op. And in my ADL, `[1,224,224,32]` consumes 17 ms. See below:
```
[profiling] kernel "23848328|[InstanceNormalization] 23848328" input[0]: [1,224,224,32] | float32, input[1]: [32] | float32, input[2]: [32] | float32, output[0]: [1,224,224,32] | float32, execution time: 17007914 ns
```
In this PR, it uses workgroup memory to optimize the original algorithm.
The advantage is that it can parallelly utilize the 64 (workgroupSize)
threads in one workgroup to calculate `mean` and `squaredNorm` value.
Meanwhile, it only outputs `H * W / workgroupSize` outputs for one
thread, which greatly reduces the overhead for one thread. With this
optimization, `[1,224,224,32]` becomes 3 ms and the main overhead is the
extra two `transpose`. The `createInstanceNormProgramInfo` only needs
`0.64` ms. See below:
```
[profiling] kernel "23003600|[InstanceNormalization] 23003600" input[0]: [1,224,224,32] | float32, output[0]: [1,32,224,224] | float32, execution time: 1543792 ns
program-manager.ts:115
[profiling] kernel "23003600|[InstanceNormalization] 23003600" input[0]: [1,32,224,224] | float32, input[1]: [32] | float32, input[2]: [32] | float32, output[0]: [1,32,224,224] | float32, execution time: 642652 ns
program-manager.ts:115
[profiling] kernel "23003600|[InstanceNormalization] 23003600" input[0]: [1,32,224,224] | float32, output[0]: [1,224,224,32] | float32, execution time: 991608 ns
```
This PR currently only applies the new algorithm to NCHW format. For
NHWC format, one way is to transpose the input so that it can use the
new algorithm. But the disadvantage is that 2 extra transpose are added.
@dakenf also gives another way to optimize NHWC. Details see
[here](https://github.com/microsoft/onnxruntime/blob/d45a96616da9843b037210f2d48d6b4e5bdae5c6/js/web/lib/wasm/jsep/webgpu/ops/instance-norm.ts).
I checked @dakenf's method. The perf is similar with transpose +
optimized NCHW. But on different GPUs, one is a little better than
another or vice versa. So I prefer this PR only does the NCHW part.
@dakenf can submit his optimization on NHWC.
2023-09-15 00:03:18 +00:00
|
|
|
[
|
|
|
|
|
{
|
|
|
|
|
"name": "Simple test with NHWC",
|
|
|
|
|
"operator": "InstanceNormalization",
|
|
|
|
|
"inputShapeDefinitions": "rankOnly",
|
|
|
|
|
"opset": { "domain": "", "version": 17 },
|
|
|
|
|
"cases": [
|
|
|
|
|
{
|
|
|
|
|
"name": "Simple test",
|
|
|
|
|
"inputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 9, 8, 7, 6, 5, 4],
|
|
|
|
|
"dims": [1, 4, 2, 2],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3, 4],
|
|
|
|
|
"dims": [4],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [4, 5, 6, 7],
|
|
|
|
|
"dims": [4],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
],
|
|
|
|
|
"outputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [
|
|
|
|
|
2.6583645343780518, 3.552788257598877, 4.447211742401123, 5.341635704040527, 2.3167295455932617,
|
|
|
|
|
4.105576515197754, 5.8944244384765625, 7.683271408081055, 6, 10.242595672607422, 6, 1.7574005126953125,
|
|
|
|
|
12.36654281616211, 8.788846969604492, 5.211153030395508, 1.633458137512207
|
|
|
|
|
],
|
|
|
|
|
"dims": [1, 4, 2, 2],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
]
|
|
|
|
|
}
|
|
|
|
|
]
|
|
|
|
|
},
|
2024-01-09 22:56:00 +00:00
|
|
|
{
|
|
|
|
|
"name": "Simple test with NHWC, components 1",
|
|
|
|
|
"operator": "InstanceNormalization",
|
|
|
|
|
"inputShapeDefinitions": "rankOnly",
|
|
|
|
|
"opset": { "domain": "", "version": 17 },
|
|
|
|
|
"cases": [
|
|
|
|
|
{
|
|
|
|
|
"name": "Simple test",
|
|
|
|
|
"inputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 9, 8, 7, 6, 5],
|
|
|
|
|
"dims": [1, 5, 3, 1],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3, 4, 5],
|
|
|
|
|
"dims": [5],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [4, 5, 6, 7, 8],
|
|
|
|
|
"dims": [5],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
],
|
|
|
|
|
"outputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [
|
|
|
|
|
2.775264263153076, 4, 5.224735260009766, 2.5505285263061523, 5, 7.449470520019531, 2.325794219970703, 6,
|
|
|
|
|
9.674205780029297, 11.898944854736328, 7, 2.1010589599609375, 14.123676300048828, 8, 1.876321792602539
|
|
|
|
|
],
|
|
|
|
|
"dims": [1, 5, 3, 1],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
]
|
|
|
|
|
}
|
|
|
|
|
]
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"name": "Simple test with NHWC, components 2",
|
|
|
|
|
"operator": "InstanceNormalization",
|
|
|
|
|
"inputShapeDefinitions": "rankOnly",
|
|
|
|
|
"opset": { "domain": "", "version": 17 },
|
|
|
|
|
"cases": [
|
|
|
|
|
{
|
|
|
|
|
"name": "Simple test",
|
|
|
|
|
"inputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 9, 8],
|
|
|
|
|
"dims": [2, 6, 1, 1],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3, 4, 5, 6],
|
|
|
|
|
"dims": [6],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [4, 5, 6, 7, 8, 9],
|
|
|
|
|
"dims": [6],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
],
|
|
|
|
|
"outputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [4, 5, 6, 7, 8, 9, 4, 5, 6, 7, 8, 9],
|
|
|
|
|
"dims": [2, 6, 1, 1],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
]
|
|
|
|
|
}
|
|
|
|
|
]
|
|
|
|
|
},
|
[js/webgpu] Optimize InstanceNormalization (#17491)
### Description
<!-- Describe your changes. -->
In previous implementation, there are two loops to iterate H * W
elements to calculate the `mean` and `squaredNorm` value in one thread,
meanwhile it outputs H * W elements in one thread. That results it's
very very slow when H * W is a large value. And usually, H * W does be a
large value in a model. For example, in the `candy-8` model, the shapes
of [H, W] are [224,224], [112,112], [56,56] for `InstanceNormalization`
op. And in my ADL, `[1,224,224,32]` consumes 17 ms. See below:
```
[profiling] kernel "23848328|[InstanceNormalization] 23848328" input[0]: [1,224,224,32] | float32, input[1]: [32] | float32, input[2]: [32] | float32, output[0]: [1,224,224,32] | float32, execution time: 17007914 ns
```
In this PR, it uses workgroup memory to optimize the original algorithm.
The advantage is that it can parallelly utilize the 64 (workgroupSize)
threads in one workgroup to calculate `mean` and `squaredNorm` value.
Meanwhile, it only outputs `H * W / workgroupSize` outputs for one
thread, which greatly reduces the overhead for one thread. With this
optimization, `[1,224,224,32]` becomes 3 ms and the main overhead is the
extra two `transpose`. The `createInstanceNormProgramInfo` only needs
`0.64` ms. See below:
```
[profiling] kernel "23003600|[InstanceNormalization] 23003600" input[0]: [1,224,224,32] | float32, output[0]: [1,32,224,224] | float32, execution time: 1543792 ns
program-manager.ts:115
[profiling] kernel "23003600|[InstanceNormalization] 23003600" input[0]: [1,32,224,224] | float32, input[1]: [32] | float32, input[2]: [32] | float32, output[0]: [1,32,224,224] | float32, execution time: 642652 ns
program-manager.ts:115
[profiling] kernel "23003600|[InstanceNormalization] 23003600" input[0]: [1,32,224,224] | float32, output[0]: [1,224,224,32] | float32, execution time: 991608 ns
```
This PR currently only applies the new algorithm to NCHW format. For
NHWC format, one way is to transpose the input so that it can use the
new algorithm. But the disadvantage is that 2 extra transpose are added.
@dakenf also gives another way to optimize NHWC. Details see
[here](https://github.com/microsoft/onnxruntime/blob/d45a96616da9843b037210f2d48d6b4e5bdae5c6/js/web/lib/wasm/jsep/webgpu/ops/instance-norm.ts).
I checked @dakenf's method. The perf is similar with transpose +
optimized NCHW. But on different GPUs, one is a little better than
another or vice versa. So I prefer this PR only does the NCHW part.
@dakenf can submit his optimization on NHWC.
2023-09-15 00:03:18 +00:00
|
|
|
{
|
|
|
|
|
"name": "Simple test with NCHW",
|
|
|
|
|
"operator": "InstanceNormalization",
|
|
|
|
|
"opset": { "domain": "", "version": 17 },
|
|
|
|
|
"cases": [
|
|
|
|
|
{
|
|
|
|
|
"name": "Simple test",
|
|
|
|
|
"inputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 9, 8, 7, 6, 5, 4],
|
|
|
|
|
"dims": [1, 4, 2, 2],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3, 4],
|
|
|
|
|
"dims": [4],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [4, 5, 6, 7],
|
|
|
|
|
"dims": [4],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
],
|
|
|
|
|
"outputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [
|
|
|
|
|
2.6583645343780518, 3.552788257598877, 4.447211742401123, 5.341635704040527, 2.3167295455932617,
|
|
|
|
|
4.105576515197754, 5.8944244384765625, 7.683271408081055, 6, 10.242595672607422, 6, 1.7574005126953125,
|
|
|
|
|
12.36654281616211, 8.788846969604492, 5.211153030395508, 1.633458137512207
|
|
|
|
|
],
|
|
|
|
|
"dims": [1, 4, 2, 2],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
]
|
|
|
|
|
}
|
|
|
|
|
]
|
2024-01-09 22:56:00 +00:00
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"name": "Simple test with NCHW, components 1",
|
|
|
|
|
"operator": "InstanceNormalization",
|
|
|
|
|
"opset": { "domain": "", "version": 17 },
|
|
|
|
|
"cases": [
|
|
|
|
|
{
|
|
|
|
|
"name": "Simple test",
|
|
|
|
|
"inputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 9, 8, 7, 6, 5],
|
|
|
|
|
"dims": [1, 5, 3, 1],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3, 4, 5],
|
|
|
|
|
"dims": [5],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [4, 5, 6, 7, 8],
|
|
|
|
|
"dims": [5],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
],
|
|
|
|
|
"outputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [
|
|
|
|
|
2.775264263153076, 4, 5.224735260009766, 2.5505285263061523, 5, 7.449470520019531, 2.325794219970703, 6,
|
|
|
|
|
9.674205780029297, 11.898944854736328, 7, 2.1010589599609375, 14.123676300048828, 8, 1.876321792602539
|
|
|
|
|
],
|
|
|
|
|
"dims": [1, 5, 3, 1],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
]
|
|
|
|
|
}
|
|
|
|
|
]
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"name": "Simple test with NCHW, components 2",
|
|
|
|
|
"operator": "InstanceNormalization",
|
|
|
|
|
"opset": { "domain": "", "version": 17 },
|
|
|
|
|
"cases": [
|
|
|
|
|
{
|
|
|
|
|
"name": "Simple test",
|
|
|
|
|
"inputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 9, 8, 7, 6, 5, 4, 3, 2],
|
|
|
|
|
"dims": [1, 3, 6, 1],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3],
|
|
|
|
|
"dims": [3],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [4, 5, 6],
|
|
|
|
|
"dims": [3],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
],
|
|
|
|
|
"outputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [
|
|
|
|
|
2.5361523628234863, 3.1216912269592285, 3.70723032951355, 4.292769432067871, 4.878308296203613,
|
|
|
|
|
5.4638471603393555, 1.8666191101074219, 3.9555397033691406, 6.044460296630859, 8.133380889892578,
|
|
|
|
|
6.044460296630859, 3.9555397033691406, 10.3915433883667, 8.634925842285156, 6.878308296203613,
|
|
|
|
|
5.121691703796387, 3.365074634552002, 1.6084575653076172
|
|
|
|
|
],
|
|
|
|
|
"dims": [1, 3, 6, 1],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
]
|
|
|
|
|
}
|
|
|
|
|
]
|
2024-03-29 20:13:39 +00:00
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"name": "Simple test with NHWC, components 1, buffer reuse",
|
|
|
|
|
"operator": "InstanceNormalization",
|
|
|
|
|
"inputShapeDefinitions": "rankOnly",
|
|
|
|
|
"opset": {
|
|
|
|
|
"domain": "",
|
|
|
|
|
"version": 17
|
|
|
|
|
},
|
|
|
|
|
"cases": [
|
|
|
|
|
{
|
|
|
|
|
"name": "Simple test",
|
|
|
|
|
"inputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3, 4, 5, 6],
|
|
|
|
|
"dims": [2, 3, 1, 1],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3],
|
|
|
|
|
"dims": [3],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [4, 5, 6],
|
|
|
|
|
"dims": [3],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
],
|
|
|
|
|
"outputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [4, 5, 6, 4, 5, 6],
|
|
|
|
|
"dims": [2, 3, 1, 1],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
]
|
|
|
|
|
}
|
|
|
|
|
]
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"name": "Simple test with NHWC, components 2, buffer reuse",
|
|
|
|
|
"operator": "InstanceNormalization",
|
|
|
|
|
"inputShapeDefinitions": "rankOnly",
|
|
|
|
|
"opset": {
|
|
|
|
|
"domain": "",
|
|
|
|
|
"version": 17
|
|
|
|
|
},
|
|
|
|
|
"cases": [
|
|
|
|
|
{
|
|
|
|
|
"name": "Simple test",
|
|
|
|
|
"inputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 9, 8, 7, 6, 5, 4, 3, 2],
|
|
|
|
|
"dims": [1, 6, 1, 3],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [1, 2, 3, 4, 5, 6],
|
|
|
|
|
"dims": [6],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
},
|
|
|
|
|
{
|
|
|
|
|
"data": [4, 5, 6, 7, 8, 9],
|
|
|
|
|
"dims": [6],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
],
|
|
|
|
|
"outputs": [
|
|
|
|
|
{
|
|
|
|
|
"data": [
|
|
|
|
|
2.775264263153076, 4, 5.224735260009766, 2.5505285263061523, 5, 7.449470520019531, 2.325794219970703, 6,
|
|
|
|
|
9.674205780029297, 11.898944854736328, 7, 2.1010589599609375, 14.123676300048828, 8, 1.876321792602539,
|
|
|
|
|
16.348413467407227, 9, 1.6515865325927734
|
|
|
|
|
],
|
|
|
|
|
"dims": [1, 6, 1, 3],
|
|
|
|
|
"type": "float32"
|
|
|
|
|
}
|
|
|
|
|
]
|
|
|
|
|
}
|
|
|
|
|
]
|
[js/webgpu] Optimize InstanceNormalization (#17491)
### Description
<!-- Describe your changes. -->
In previous implementation, there are two loops to iterate H * W
elements to calculate the `mean` and `squaredNorm` value in one thread,
meanwhile it outputs H * W elements in one thread. That results it's
very very slow when H * W is a large value. And usually, H * W does be a
large value in a model. For example, in the `candy-8` model, the shapes
of [H, W] are [224,224], [112,112], [56,56] for `InstanceNormalization`
op. And in my ADL, `[1,224,224,32]` consumes 17 ms. See below:
```
[profiling] kernel "23848328|[InstanceNormalization] 23848328" input[0]: [1,224,224,32] | float32, input[1]: [32] | float32, input[2]: [32] | float32, output[0]: [1,224,224,32] | float32, execution time: 17007914 ns
```
In this PR, it uses workgroup memory to optimize the original algorithm.
The advantage is that it can parallelly utilize the 64 (workgroupSize)
threads in one workgroup to calculate `mean` and `squaredNorm` value.
Meanwhile, it only outputs `H * W / workgroupSize` outputs for one
thread, which greatly reduces the overhead for one thread. With this
optimization, `[1,224,224,32]` becomes 3 ms and the main overhead is the
extra two `transpose`. The `createInstanceNormProgramInfo` only needs
`0.64` ms. See below:
```
[profiling] kernel "23003600|[InstanceNormalization] 23003600" input[0]: [1,224,224,32] | float32, output[0]: [1,32,224,224] | float32, execution time: 1543792 ns
program-manager.ts:115
[profiling] kernel "23003600|[InstanceNormalization] 23003600" input[0]: [1,32,224,224] | float32, input[1]: [32] | float32, input[2]: [32] | float32, output[0]: [1,32,224,224] | float32, execution time: 642652 ns
program-manager.ts:115
[profiling] kernel "23003600|[InstanceNormalization] 23003600" input[0]: [1,32,224,224] | float32, output[0]: [1,224,224,32] | float32, execution time: 991608 ns
```
This PR currently only applies the new algorithm to NCHW format. For
NHWC format, one way is to transpose the input so that it can use the
new algorithm. But the disadvantage is that 2 extra transpose are added.
@dakenf also gives another way to optimize NHWC. Details see
[here](https://github.com/microsoft/onnxruntime/blob/d45a96616da9843b037210f2d48d6b4e5bdae5c6/js/web/lib/wasm/jsep/webgpu/ops/instance-norm.ts).
I checked @dakenf's method. The perf is similar with transpose +
optimized NCHW. But on different GPUs, one is a little better than
another or vice versa. So I prefer this PR only does the NCHW part.
@dakenf can submit his optimization on NHWC.
2023-09-15 00:03:18 +00:00
|
|
|
}
|
|
|
|
|
]
|