File size: 3,308 Bytes
17f039f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | export const EXPECTED_ORT_VERSION = "1.27.0";
const replacements = [
["var acc: array<vec4<${n}>, rowPerThread>;", "var acc: array<vec4<f32>, rowPerThread>;", 1],
[
"var acc : array<array<${n}, colPerThread>, rowPerThread>;",
"var acc : array<array<f32, colPerThread>, rowPerThread>;",
1,
],
["BCached0 * ACached0[i] + acc[i]", "vec4<f32>(BCached0) * f32(ACached0[i]) + acc[i]", 1],
["BCached1 * ACached1[i] + acc[i]", "vec4<f32>(BCached1) * f32(ACached1[i]) + acc[i]", 1],
["BCached2 * ACached2[i] + acc[i]", "vec4<f32>(BCached2) * f32(ACached2[i]) + acc[i]", 1],
["BCached3 * ACached3[i] + acc[i]", "vec4<f32>(BCached3) * f32(ACached3[i]) + acc[i]", 1],
["BCached0 * ACached.x + acc[i]", "vec4<f32>(BCached0) * f32(ACached.x) + acc[i]", 1],
["BCached1 * ACached.y + acc[i]", "vec4<f32>(BCached1) * f32(ACached.y) + acc[i]", 1],
["BCached2 * ACached.z + acc[i]", "vec4<f32>(BCached2) * f32(ACached.z) + acc[i]", 1],
["BCached3 * ACached.w + acc[i]", "vec4<f32>(BCached3) * f32(ACached.w) + acc[i]", 1],
["ACached * BCached[innerCol]", "f32(ACached) * f32(BCached[innerCol])", 2],
[
"mm_write(batch, globalRow + innerRow, globalCol, acc[innerRow]);",
"mm_write(batch, globalRow + innerRow, globalCol, vec4<${n}>(acc[innerRow]));",
1,
],
[
"mm_write(batch, gRow, gCol, acc[innerRow][innerCol]);",
"mm_write(batch, gRow, gCol, ${n}(acc[innerRow][innerCol]));",
1,
],
[
"mm_write(batch, globalRow + innerRow, globalCol + innerCol,\n acc[innerRow][innerCol]);",
"mm_write(batch, globalRow + innerRow, globalCol + innerCol,\n ${n}(acc[innerRow][innerCol]));",
1,
],
];
const countOccurrences = (text, needle) => text.split(needle).length - 1;
export function patchOrtWebGpuF32Accumulate(source) {
const originalCounts = replacements.map(([from]) => countOccurrences(source, from));
if (originalCounts.every((count) => count === 0)) {
const patchedCounts = replacements.map(([, to]) => countOccurrences(source, to));
const complete = patchedCounts.every((count, index) => count === replacements[index][2]);
if (!complete) {
throw new Error(
"onnxruntime-web WebGPU MatMul bundle is neither the expected upstream form nor the fully patched form."
);
}
return source;
}
let patched = source;
for (let index = 0; index < replacements.length; index += 1) {
const [from, to, expectedCount] = replacements[index];
const actualCount = originalCounts[index];
if (actualCount !== expectedCount) {
throw new Error(
`Unexpected onnxruntime-web bundle shape for replacement ${index + 1}: `
+ `expected ${expectedCount}, found ${actualCount}.`
);
}
patched = patched.split(from).join(to);
}
for (let index = 0; index < replacements.length; index += 1) {
const [from, to, expectedCount] = replacements[index];
const originalCount = countOccurrences(patched, from);
const patchedCount = countOccurrences(patched, to);
if (originalCount !== 0 || patchedCount !== expectedCount) {
throw new Error(
`Failed to verify onnxruntime-web replacement ${index + 1}: `
+ `original=${originalCount}, patched=${patchedCount}, expected=${expectedCount}.`
);
}
}
return patched;
}
|