saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-26 19:52:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[js/webgpu] Support shared memory for transpose 2d (#19267)

Browse source 
For 1024x1024, without shared memoey, 18.7ms. With shared memory 13.2ms.
This commit is contained in:
Xu Xing 2024-05-22 23:15:44 +08:00 committed by GitHub
parent 068bb3d5ee
commit f1fef19b6e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 49 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
js/web/lib/wasm/jsep/webgpu/ops/transpose.ts
View file

@ -43,8 +43,30 @@ export const createTransposeProgramInfo = (inputTensor: TensorView, permAttr: nu
const outputShape = getOutputShape(inputTensor.dims, perm);
const output = outputVariable('output', inputDataType, outputShape.length);
const input = inputVariable('a', inputDataType, inputRank);
const getShaderSource = (shaderHelper: ShaderHelper) => `
let getShaderSource;
if (perm.length === 2 && perm[0] === 1 && perm[1] === 0) {
const wgslType = output.type.value;
const workgroupSize: [number, number, number] = [16, 16, 1];
getShaderSource = (shaderHelper: ShaderHelper) => `
${shaderHelper.registerUniform('output_size', 'u32').declareVariables(input, output)}
var<workgroup> tile : array<array<${wgslType}, ${workgroupSize[0] + 1}>, ${workgroupSize[0]}>;
${shaderHelper.mainStart(workgroupSize)}
var x = workgroup_id.x * ${workgroupSize[0]}u + local_id.x;
var y = workgroup_id.y * ${workgroupSize[0]}u + local_id.y;
let width = uniforms.output_shape[0];
let height = uniforms.output_shape[1];
if (x < width && y < height) {
tile[local_id.y][local_id.x] = ${input.getByOffset('y * width + x')};
}
workgroupBarrier();
x = workgroup_id.y * ${workgroupSize[0]}u + local_id.x;
y = workgroup_id.x * ${workgroupSize[0]}u + local_id.y;
if (x < height && y < width) {
${output.setByOffset('y * height + x', 'tile[local_id.x][local_id.y]')}
}
}`;
} else {
getShaderSource = (shaderHelper: ShaderHelper) => `
${shaderHelper.registerUniform('output_size', 'u32').declareVariables(input, output)}
${permFunctionBody(perm, inputRank, input, output)}
@ -57,6 +79,7 @@ export const createTransposeProgramInfo = (inputTensor: TensorView, permAttr: nu
${output.setByOffset('global_idx', input.getByIndices('aIndices'))}
}`;
}
return {
name: 'Transpose',
shaderCache: {hint: `${permAttr}`, inputDependencies: ['rank']},

24
js/web/test/data/ops/transpose.jsonc
View file

@ -190,5 +190,29 @@
]
}
]
},
{
"name": "Transpose - perms:[1, 0]",
"operator": "Transpose",
"attributes": [{ "name": "perm", "data": [1, 0], "type": "ints" }],
"cases": [
{
"name": "T[6,4]",
"inputs": [
{
"data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24],
"dims": [6, 4],
"type": "float32"
}
],
"outputs": [
{
"data": [1, 5, 9, 13, 17, 21, 2, 6, 10, 14, 18, 22, 3, 7, 11, 15, 19, 23, 4, 8, 12, 16, 20, 24],
"dims": [4, 6],
"type": "float32"
}
]
}
]
}
]