/** * @license * Copyright 2024 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ (function (global, factory) { typeof exports === 'object' && typeof module !== 'undefined' ? factory(exports, require('@tensorflow/tfjs-core')) : typeof define === 'function' && define.amd ? define(['exports', '@tensorflow/tfjs-core'], factory) : (global = typeof globalThis !== 'undefined' ? globalThis : global || self, factory(global.tf = global.tf || {}, global.tf)); })(this, (function (exports, tf) { 'use strict'; function _interopNamespaceDefault(e) { var n = Object.create(null); if (e) { Object.keys(e).forEach(function (k) { if (k !== 'default') { var d = Object.getOwnPropertyDescriptor(e, k); Object.defineProperty(n, k, d.get ? d : { enumerable: true, get: function () { return e[k]; } }); } }); } n.default = e; return n; } var tf__namespace = /*#__PURE__*/_interopNamespaceDefault(tf); /****************************************************************************** Copyright (c) Microsoft Corporation. Permission to use, copy, modify, and/or distribute this software for any purpose with or without fee is hereby granted. THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. ***************************************************************************** */ /* global Reflect, Promise */ var extendStatics = function (d, b) { extendStatics = Object.setPrototypeOf || ({ __proto__: [] } instanceof Array && function (d, b) { d.__proto__ = b; }) || function (d, b) { for (var p in b) if (Object.prototype.hasOwnProperty.call(b, p)) d[p] = b[p]; }; return extendStatics(d, b); }; function __extends(d, b) { if (typeof b !== "function" && b !== null) throw new TypeError("Class extends value " + String(b) + " is not a constructor or null"); extendStatics(d, b); function __() { this.constructor = d; } d.prototype = b === null ? Object.create(b) : (__.prototype = b.prototype, new __()); } function __awaiter(thisArg, _arguments, P, generator) { function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); } return new (P || (P = Promise))(function (resolve, reject) { function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } } function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } } function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); } step((generator = generator.apply(thisArg, _arguments || [])).next()); }); } function __generator(thisArg, body) { var _ = { label: 0, sent: function () { if (t[0] & 1) throw t[1]; return t[1]; }, trys: [], ops: [] }, f, y, t, g; return g = { next: verb(0), "throw": verb(1), "return": verb(2) }, typeof Symbol === "function" && (g[Symbol.iterator] = function () { return this; }), g; function verb(n) { return function (v) { return step([n, v]); }; } function step(op) { if (f) throw new TypeError("Generator is already executing."); while (_) try { if (f = 1, y && (t = op[0] & 2 ? y["return"] : op[0] ? y["throw"] || ((t = y["return"]) && t.call(y), 0) : y.next) && !(t = t.call(y, op[1])).done) return t; if (y = 0, t) op = [op[0] & 2, t.value]; switch (op[0]) { case 0: case 1: t = op; break; case 4: _.label++; return { value: op[1], done: false }; case 5: _.label++; y = op[1]; op = [0]; continue; case 7: op = _.ops.pop(); _.trys.pop(); continue; default: if (!(t = _.trys, t = t.length > 0 && t[t.length - 1]) && (op[0] === 6 || op[0] === 2)) { _ = 0; continue; } if (op[0] === 3 && (!t || (op[1] > t[0] && op[1] < t[3]))) { _.label = op[1]; break; } if (op[0] === 6 && _.label < t[1]) { _.label = t[1]; t = op; break; } if (t && _.label < t[2]) { _.label = t[2]; _.ops.push(op); break; } if (t[2]) _.ops.pop(); _.trys.pop(); continue; } op = body.call(thisArg, _); } catch (e) { op = [6, e]; y = 0; } finally { f = t = 0; } if (op[0] & 5) throw op[1]; return { value: op[0] ? op[1] : void 0, done: true }; } } function __values(o) { var s = typeof Symbol === "function" && Symbol.iterator, m = s && o[s], i = 0; if (m) return m.call(o); if (o && typeof o.length === "number") return { next: function () { if (o && i >= o.length) o = void 0; return { value: o && o[i++], done: !o }; } }; throw new TypeError(s ? "Object is not iterable." : "Symbol.iterator is not defined."); } function __read(o, n) { var m = typeof Symbol === "function" && o[Symbol.iterator]; if (!m) return o; var i = m.call(o), r, ar = [], e; try { while ((n === void 0 || n-- > 0) && !(r = i.next()).done) ar.push(r.value); } catch (error) { e = { error: error }; } finally { try { if (r && !r.done && (m = i["return"])) m.call(i); } finally { if (e) throw e.error; } } return ar; } function __spreadArray(to, from, pack) { if (pack || arguments.length === 2) for (var i = 0, l = from.length, ar; i < l; i++) { if (ar || !(i in from)) { if (!ar) ar = Array.prototype.slice.call(from, 0, i); ar[i] = from[i]; } } return to.concat(ar || Array.prototype.slice.call(from)); } /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var ENV = tf.env(); /** The batched dispatching calls size in the device queue. */ ENV.registerFlag('WEBGPU_DEFERRED_SUBMIT_BATCH_SIZE', function () { return 15; }); /** * Whether we forward execution to the CPU backend if tensors are small and * reside on the CPU. */ ENV.registerFlag('WEBGPU_CPU_FORWARD', function () { return true; }); /** * This flag is used to test different types of matmul programs. * * See MatMulProgramType in webgpu_util.ts for a list of available values. */ ENV.registerFlag('WEBGPU_MATMUL_PROGRAM_TYPE', function () { return -1; }); /** * Whether to use conv2dTranspose_naive which directly implement the * conv2dTranspose logic rather than using a matmul to simulate. */ ENV.registerFlag('WEBGPU_USE_NAIVE_CONV2D_TRANSPOSE', function () { return true; }); /** * Whether we use low power GPU. Otherwise, a high performance GPU will be * requested. */ ENV.registerFlag('WEBGPU_USE_LOW_POWER_GPU', function () { return false; }); /** * Threshold for input tensor size that determines whether WebGPU backend will * delegate computation to CPU. * * Default value is 1000. */ ENV.registerFlag('WEBGPU_CPU_HANDOFF_SIZE_THRESHOLD', function () { return 1000; }); /** * Whether to use a dummy canvas to make profiling tools like PIX work with * TFJS webgpu backend. */ ENV.registerFlag('WEBGPU_USE_PROFILE_TOOL', function () { return false; }); /** * Whether to use import API. */ ENV.registerFlag('WEBGPU_IMPORT_EXTERNAL_TEXTURE', function () { return true; }); /** * Whether to use conv2dNaive for debugging. */ ENV.registerFlag('WEBGPU_USE_NAIVE_CONV2D_DEBUG', function () { return false; }); /** * Threshold to increase dispatched workgroups for matmul. If too few workgroups * are dispatched, it means the hardware may be in low occupancy. * -1 means it's not set by the user. A default strategy will be applied. */ ENV.registerFlag('WEBGPU_THRESHOLD_TO_INCREASE_WORKGROUPS_FOR_MATMUL', function () { return -1; }); /** * Whether we will run im2col as a separate shader for convolution. */ ENV.registerFlag('WEBGPU_CONV_SEPARATE_IM2COL_SHADER', function () { return false; }); /** * A string used to match shader key. If any matches, print the related shader. * Seperated by comma. 'all' to print all. 'binary' to print binary(add, mul, * etc.). 'unary,conv2d' to print both unary and conv2d. */ ENV.registerFlag('WEBGPU_PRINT_SHADER', function () { return ''; }); /** Experimental flag, whether enter compile only phase. */ ENV.registerFlag('WEBGPU_ENGINE_COMPILE_ONLY', function () { return false; }); /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var AdapterInfo = /** @class */ (function () { function AdapterInfo(adapterInfo) { if (adapterInfo) { this.vendor = adapterInfo.vendor; this.architecture = adapterInfo.architecture; this.intelGPUGeneration = this.getIntelGPUGeneration(); } } AdapterInfo.prototype.getIntelGPUGeneration = function () { if (this.isIntel()) { if (this.architecture.startsWith('gen')) { return Number(this.architecture.match(/\d+/)); } else if (this.architecture.startsWith('xe')) { return 12; } } return 0; }; AdapterInfo.prototype.isIntel = function () { return this.vendor === 'intel'; }; return AdapterInfo; }()); /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var BufferManager = /** @class */ (function () { function BufferManager(device) { this.device = device; this.numUsedBuffers = 0; this.numFreeBuffers = 0; this.freeBuffers = new Map(); this.usedBuffers = new Map(); this.numBytesUsed = 0; this.numBytesAllocated = 0; } BufferManager.prototype.acquireBuffer = function (size, usage, mappedAtCreation, reuse) { if (mappedAtCreation === void 0) { mappedAtCreation = false; } if (reuse === void 0) { reuse = true; } var buffer; var key = getBufferKey(size, usage); if (reuse) { if (!this.freeBuffers.has(key)) { this.freeBuffers.set(key, []); } if (this.freeBuffers.get(key).length > 0) { buffer = this.freeBuffers.get(key).pop(); this.numFreeBuffers--; } else { buffer = this.device.createBuffer({ size: size, usage: usage, mappedAtCreation: mappedAtCreation }); this.numBytesAllocated += size; } } else { buffer = this.device.createBuffer({ size: size, usage: usage, mappedAtCreation: mappedAtCreation }); this.numBytesAllocated += size; } if (!this.usedBuffers.has(key)) { this.usedBuffers.set(key, []); } this.usedBuffers.get(key).push(buffer); this.numUsedBuffers++; this.numBytesUsed += size; return buffer; }; BufferManager.prototype.releaseBuffer = function (buffer, reuse) { if (reuse === void 0) { reuse = true; } if (this.freeBuffers.size === 0) { return; } var size = buffer.size; var usage = buffer.usage; var key = getBufferKey(size, usage); var bufferArray = this.usedBuffers.get(key); var index = bufferArray.indexOf(buffer); if (index < 0) { throw new Error('Cannot find the buffer in buffer manager'); } bufferArray[index] = bufferArray[bufferArray.length - 1]; bufferArray.pop(); this.numUsedBuffers--; this.numBytesUsed -= size; if (reuse) { this.freeBuffers.get(key).push(buffer); this.numFreeBuffers++; } else { buffer.destroy(); this.numBytesAllocated -= size; } }; BufferManager.prototype.getNumUsedBuffers = function () { return this.numUsedBuffers; }; BufferManager.prototype.getNumFreeBuffers = function () { return this.numFreeBuffers; }; BufferManager.prototype.dispose = function () { this.freeBuffers.forEach(function (buffers, key) { buffers.forEach(function (buffer) { buffer.destroy(); }); }); this.usedBuffers.forEach(function (buffers, key) { buffers.forEach(function (buffer) { buffer.destroy(); }); }); this.freeBuffers = new Map(); this.usedBuffers = new Map(); this.numUsedBuffers = 0; this.numFreeBuffers = 0; this.numBytesUsed = 0; this.numBytesAllocated = 0; }; return BufferManager; }()); function getBufferKey(size, usage) { return "".concat(size, "_").concat(usage); } /** * @license * Copyright 2022 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var TextureManager = /** @class */ (function () { function TextureManager(device) { this.device = device; this.numUsedTextures = 0; this.numFreeTextures = 0; this.freeTextures = new Map(); this.usedTextures = new Map(); this.numBytesUsed = 0; this.numBytesAllocated = 0; } TextureManager.prototype.acquireTexture = function (width, height, format, usage) { var bytesPerElement = getBytesPerElement(format); var byteSize = width * height * bytesPerElement; var key = getTextureKey(width, height, format, usage); if (!this.freeTextures.has(key)) { this.freeTextures.set(key, []); } if (!this.usedTextures.has(key)) { this.usedTextures.set(key, []); } this.numBytesUsed += byteSize; this.numUsedTextures++; if (this.freeTextures.get(key).length > 0) { this.numFreeTextures--; var newTexture_1 = this.freeTextures.get(key).shift(); this.usedTextures.get(key).push(newTexture_1); return newTexture_1; } this.numBytesAllocated += byteSize; var newTexture = this.device.createTexture({ size: [width, height], format: format, usage: usage, }); this.usedTextures.get(key).push(newTexture); return newTexture; }; TextureManager.prototype.releaseTexture = function (texture) { if (this.freeTextures.size === 0) { return; } var width = texture.width; var height = texture.height; var format = texture.format; var usage = texture.usage; var key = getTextureKey(width, height, format, usage); if (!this.freeTextures.has(key)) { this.freeTextures.set(key, []); } this.freeTextures.get(key).push(texture); this.numFreeTextures++; this.numUsedTextures--; var textureList = this.usedTextures.get(key); var textureIndex = textureList.indexOf(texture); if (textureIndex < 0) { throw new Error('Cannot release a texture that was never provided by this ' + 'texture manager'); } textureList.splice(textureIndex, 1); var bytesPerElement = getBytesPerElement(format); var byteSize = width * height * bytesPerElement; this.numBytesUsed -= byteSize; }; TextureManager.prototype.getNumUsedTextures = function () { return this.numUsedTextures; }; TextureManager.prototype.getNumFreeTextures = function () { return this.numFreeTextures; }; TextureManager.prototype.dispose = function () { this.freeTextures.forEach(function (textures, key) { textures.forEach(function (texture) { texture.destroy(); }); }); this.usedTextures.forEach(function (textures, key) { textures.forEach(function (texture) { texture.destroy(); }); }); this.freeTextures = new Map(); this.usedTextures = new Map(); this.numUsedTextures = 0; this.numFreeTextures = 0; this.numBytesUsed = 0; this.numBytesAllocated = 0; }; return TextureManager; }()); function getTextureKey(width, height, format, usage) { return "".concat(width, "_").concat(height, "_").concat(format, "_").concat(usage); } function getBytesPerElement(format) { if (format === 'rgba8unorm') { return 16; } else { throw new Error("".concat(format, " is not supported!")); } } /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ // Generates WGSL that computes strides. function symbolicallyComputeStrides(indicesArr, variableName) { if (Math.max.apply(Math, __spreadArray([], __read(indicesArr), false)) > 5) { throw new Error('Cannot symbolically compute strides for rank > 6 tensor.'); } var numCoords = indicesArr.length; var indicesStr = 'xyzwuv'; var shape = indicesArr.map(function (d) { return "".concat(variableName, ".").concat(indicesStr[d]); }); var strides = new Array(numCoords - 1); strides[numCoords - 2] = shape[numCoords - 1]; for (var i = numCoords - 3; i >= 0; --i) { strides[i] = "(".concat(strides[i + 1], " * ").concat(shape[i + 1], ")"); } return strides; } var atomicAddSnippet = function (ptr, v, type) { if (type === 'int32') { return "atomicAdd(".concat(ptr, ", bitcast(").concat(v, "));"); } else { // atomicAdd only supports uint/int type. For float, we use // atomicCompareExchangeWeak to simulate. return "\n {\n var oldValue = 0;\n loop {\n let newValueF32 = bitcast(oldValue) + (".concat(v, ");\n let newValue = bitcast(newValueF32);\n let res = atomicCompareExchangeWeak(").concat(ptr, ", oldValue, newValue);\n if res.exchanged {\n break;\n }\n oldValue = res.old_value;\n }\n }"); } }; /** * @license * Copyright 2022 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var PixelsOpType; (function (PixelsOpType) { PixelsOpType[PixelsOpType["FROM_PIXELS"] = 0] = "FROM_PIXELS"; PixelsOpType[PixelsOpType["DRAW"] = 1] = "DRAW"; })(PixelsOpType || (PixelsOpType = {})); var compileProgram = function (device, program, inputsData, output, parallelCompilation) { var outputData = { dtype: output.dtype, shape: output.shape }; var source = makeShader(inputsData, outputData, program); var module = device.createShaderModule({ code: source, label: program.constructor.name }); var printShaderString = tf.env().get('WEBGPU_PRINT_SHADER'); if (printShaderString !== '') { printShaderString = printShaderString.toLowerCase(); var printShaderArray = printShaderString.split(','); if (printShaderString === 'all' || printShaderArray.some(function (item) { return program.shaderKey.toLowerCase().includes(item); })) { console.group(program.shaderKey); console.debug(source); console.groupEnd(); } } if (parallelCompilation) { return device.createComputePipelineAsync({ compute: { module: module, entryPoint: '_start' }, label: program.constructor.name, layout: 'auto' }); } else { return device.createComputePipeline({ compute: { module: module, entryPoint: '_start' }, label: program.constructor.name, layout: 'auto' }); } }; var typeSnippet = function (component, type) { if (type === void 0) { type = 'f32'; } switch (component) { case 1: return "".concat(type); case 2: return "vec2<".concat(type, ">"); case 3: return "vec3<".concat(type, ">"); case 4: return "vec4<".concat(type, ">"); default: throw new Error("".concat(component, "-component ").concat(type, " is not supported.")); } }; function getCoordsDataType(rank) { if (rank <= 1) { return 'i32'; } else if (rank === 2) { return "vec2"; } else if (rank === 3) { return "vec3"; } else if (rank === 4) { return "vec4"; } else if (rank === 5) { return "vec5"; } else if (rank === 6) { return "vec6"; } else { throw Error("GPU for rank ".concat(rank, " is not yet supported")); } } function getCoordsXYZ(index) { if (index === 0) { return 'x'; } else if (index === 1) { return 'y'; } else if (index === 2) { return 'z'; } else if (index === 3) { return 'w'; } else if (index === 4) { return 'u'; } else if (index === 5) { return 'v'; } else { throw Error("Index ".concat(index, " is not yet supported")); } } function getMainHeaderString() { var params = []; for (var _i = 0; _i < arguments.length; _i++) { params[_i] = arguments[_i]; } var snippet; switch (params.length) { case 0: snippet = "\n fn main()\n "; break; case 1: snippet = "\n fn main(".concat(params[0], " : i32)\n "); break; default: throw Error('Unreachable'); } return snippet; } function getStartHeaderString(useGlobalIndex, program) { var snippet; snippet = "\n ".concat(getWorkgroupSizeString(program), "\n fn _start(@builtin(local_invocation_id) LocalId : vec3,\n @builtin(global_invocation_id) GlobalId : vec3,\n @builtin(local_invocation_index) LocalIndex: u32,\n @builtin(workgroup_id) WorkgroupId : vec3,\n @builtin(num_workgroups) NumWorkgroups : vec3) {\n localId = LocalId;\n localIndex = LocalIndex;\n globalId = GlobalId;\n numWorkgroups = NumWorkgroups;\n workgroupId = WorkgroupId;\n ").concat(useGlobalIndex ? "main(getGlobalIndex());" : "main();", ";\n }\n "); return snippet; } function getWorkgroupSizeString(program) { return "\n @compute @workgroup_size(".concat(program.workgroupSize[0], ", ").concat(program.workgroupSize[1], ", ").concat(program.workgroupSize[2], ")\n"); } function makeShader(inputInfo, outputData, program) { var prefixSnippets = []; var flatWorkgroupSize = program.workgroupSize[0] * program.workgroupSize[1] * program.workgroupSize[2]; program.outputComponent = program.outputComponent ? program.outputComponent : 1; prefixSnippets.push("\n\n var localId: vec3;\n var localIndex: u32;\n var globalId: vec3;\n var numWorkgroups: vec3;\n var workgroupId: vec3;\n\n // Only used when the y/z dimension of workgroup size is 1.\n fn getGlobalIndex() -> i32 {\n ".concat(isFlatDispatch(program) ? " return i32(globalId.x);" : " return i32((workgroupId.z * numWorkgroups.x * numWorkgroups.y +\n workgroupId.y * numWorkgroups.x + workgroupId.x) * ".concat(flatWorkgroupSize, "u +\n localIndex);\n "), "\n }\n ")); if (program.pixelsOpType != null) { var inoutSnippet = program.pixelsOpType === PixelsOpType.FROM_PIXELS ? "@group(0) @binding(0) var result: array<".concat(dataTypeToGPUType(outputData.dtype, program.outputComponent), ">;") : "@group(0) @binding(1) var inBuf : array<".concat(dataTypeToGPUType(inputInfo[0].dtype, program.outputComponent), ">;"); var outShapeStridesType = outputData.shape.length === 3 ? 'vec2' : 'i32'; prefixSnippets.push("\n struct Uniform {\n outShapeStrides : ".concat(outShapeStridesType, ",\n size : i32,\n numChannels : i32,\n alpha : f32,\n };\n\n ").concat(inoutSnippet, "\n @group(0) @binding(2) var uniforms: Uniform;\n ")); var useGlobalIndex_1 = isFlatDispatchLayout(program); return [ commonSnippet, prefixSnippets.join('\n'), getCoordsFromIndexSnippet(outputData.shape), program.getUserCode(), getStartHeaderString(useGlobalIndex_1, program), ].join('\n'); } var stridesLength; var stridesDataType; var uniformDeclaration = 'struct Uniforms { NAN : f32, INFINITY : f32, '; program.variableNames.forEach(function (x, i) { var perDataType = getCoordsDataType(inputInfo[i].shape.length); uniformDeclaration += "".concat(x.charAt(0).toLowerCase() + x.slice(1), "Shape : ").concat(perDataType, ", "); stridesLength = inputInfo[i].shape.length - 1; stridesDataType = getCoordsDataType(stridesLength); uniformDeclaration += "".concat(x.charAt(0).toLowerCase() + x.slice(1), "ShapeStrides: ").concat(stridesDataType, ", "); }); var outputDataType = getCoordsDataType(outputData.shape.length); uniformDeclaration += "outShape : ".concat(outputDataType, ", "); stridesLength = outputData.shape.length - 1; stridesDataType = getCoordsDataType(stridesLength); uniformDeclaration += "\n outShapeStrides: ".concat(stridesDataType, ", "); if (program.size) { uniformDeclaration += 'size : i32, '; } if (program.uniforms) { uniformDeclaration += program.uniforms; } uniformDeclaration += '};'; uniformDeclaration = insertAlignment(uniformDeclaration); prefixSnippets.push(uniformDeclaration); // Output buffer. if (program.atomic) { prefixSnippets.push("\n @group(0) @binding(0) var result: array>;\n "); } else { prefixSnippets.push("\n @group(0) @binding(0) var result: array<".concat(dataTypeToGPUType(outputData.dtype, program.outputComponent), ">;\n ")); } program.variableNames.forEach(function (x, i) { prefixSnippets.push("\n @group(0) @binding(".concat(1 + i, ") var ").concat(x, ": array<").concat(program.variableComponents ? dataTypeToGPUType(inputInfo[i].dtype, program.variableComponents[i]) : dataTypeToGPUType(inputInfo[i].dtype, program.outputComponent), ">;\n ")); }); if (uniformDeclaration !== '') { prefixSnippets.push("\n @group(0) @binding(".concat(1 + program.variableNames.length, ") var uniforms: Uniforms;\n ")); } var coordsSnippet = getOutputCoordsSnippet(outputData.shape, program.dispatchLayout); var sources = [ commonSnippet, prefixSnippets.join('\n') + isInfSnippet, getCoordsFromIndexSnippet(outputData.shape), coordsSnippet, getOutputIndexFromCoordsSnippet(outputData.shape.length) ]; if (!program.atomic) { sources.push(setOutputSnippet(outputData.shape, outputData.dtype, program.outputComponent)); } program.variableNames.forEach(function (x, i) { sources.push("".concat(getCoordsFromIndexSnippet(inputInfo[i].shape, x))); }); var inputSnippet = inputInfo .map(function (x, i) { return getInputSnippet(x, outputData.shape, program.variableComponents ? program.variableComponents[i] : program.outputComponent, program.dispatchLayout.x.length === outputData.shape.length); }) .join('\n'); sources.push(inputSnippet); sources.push(program.getUserCode()); var useGlobalIndex = isFlatDispatchLayout(program); sources.push(getStartHeaderString(useGlobalIndex, program)); var source = sources.join('\n'); return source; } function makeShaderKey(program, inputsData, output) { var key = program.shaderKey; if (program.pixelsOpType != null) { return key; } var shapes = []; var types = []; inputsData.forEach(function (element) { shapes.push(element.shape); types.push(element.dtype); }); shapes.push(output.shape); types.push(output.dtype); var broadcastDims = inputsData.map(function (d) { return tf.backend_util.getBroadcastDims(d.shape, output.shape); }); var inputShapesEqualsOutShape = inputsData.map(function (d) { return tf.util.arraysEqual(d.shape, output.shape); }).join('_'); var broadcastDimsKey = broadcastDims.map(function (d) { return d.join('_'); }).join(';'); var flatDispatchString = isFlatDispatch(program) ? 'flatDispatch' : ''; key += '_' + (program.workgroupSize ? program.workgroupSize.join(',') : '') + shapes.map(function (shape) { return shape.length; }).join(',') + types.join(',') + program.variableNames.join(',') + broadcastDimsKey + inputShapesEqualsOutShape + flatDispatchString; return key; } var commonSnippet = "\n struct vec5 {x: i32, y: i32, z: i32, w: i32, u: i32};\n struct vec6 {x: i32, y: i32, z: i32, w: i32, u: i32, v: i32};\n\n // Checks whether coordinates lie within the bounds of the shape.\n fn coordsInBounds2D(coord : vec2, shape : vec2) -> bool {\n return all(coord >= vec2(0)) && all(coord < shape);\n }\n fn coordsInBounds3D(coord : vec3, shape : vec3) -> bool {\n return all(coord >= vec3(0)) && all(coord < shape);\n }\n fn coordsInBounds4D(coord : vec4, shape : vec4) -> bool {\n return all(coord >= vec4(0)) && all(coord < shape);\n }\n\n fn getIndexFromCoords1D(coord : i32, shape : i32) -> i32 {\n return coord;\n }\n fn getIndexFromCoords2D(coords : vec2, shape : vec2) -> i32 {\n return dot(coords, vec2(shape.y, 1));\n }\n fn getIndexFromCoords3D(coords : vec3, shape : vec3) -> i32 {\n return dot(coords, vec3(shape.y * shape.z, shape.z, 1));\n }\n fn getIndexFromCoords4D(coords : vec4, shape : vec4) -> i32 {\n return dot(coords, vec4(\n shape.y * shape.z * shape.w, shape.z * shape.w, shape.w, 1));\n }\n fn getIndexFromCoords5D(coords : vec5, shape : vec5) -> i32 {\n let shapeStrides: vec5 = vec5(shape.y * shape.z * shape.w * shape.u, shape.z * shape.w * shape.u, shape.w * shape.u, shape.u, 1);\n return coords.x*shapeStrides.x + coords.y*shapeStrides.y + coords.z*shapeStrides.z + coords.w*shapeStrides.w + coords.u*shapeStrides.u;\n }\n fn getIndexFromCoords6D(coords : vec6, shape : vec6) -> i32 {\n let shapeStrides: vec6 = vec6(shape.y * shape.z * shape.w * shape.u * shape.v, shape.z * shape.w * shape.u * shape.v, shape.w * shape.u * shape.v, shape.u * shape.v, shape.v, 1);\n return coords.x*shapeStrides.x + coords.y*shapeStrides.y + coords.z*shapeStrides.z + coords.w*shapeStrides.w + coords.u*shapeStrides.u + coords.v*shapeStrides.v;\n }\n\n // NaN defination in IEEE 754-1985 is :\n // - sign = either 0 or 1.\n // - biased exponent = all 1 bits.\n // - fraction = anything except all 0 bits (since all 0 bits represents infinity).\n // https://en.wikipedia.org/wiki/IEEE_754-1985#Representation_of_non-numbers\n fn isnan(val: f32) -> bool {\n let floatToUint: u32 = bitcast(val);\n return (floatToUint & 0x7fffffffu) > 0x7f800000u;\n }\n fn isnanVec4(val : vec4) -> vec4 {\n let floatToUint: vec4 = bitcast>(val);\n return (floatToUint & vec4(0x7fffffffu)) > vec4(0x7f800000u);\n }\n"; var isInfSnippet = "\n fn isinf(val: f32) -> bool {\n return abs(val) == uniforms.INFINITY;\n }\n"; /** * Derives logical coordinates from a flat index. Performs integer division * with each stride and decrements the index until the index equals the final * dimension coordinate. */ function getCoordsFromIndexSnippet(shape, name) { if (name === void 0) { name = ''; } var rank = shape.length; var funcName = name !== '' ? "get".concat(name.charAt(0).toUpperCase() + name.slice(1), "CoordsFromIndex") : 'getCoordsFromIndex'; var stridesName = name !== '' ? "".concat(name.charAt(0).toLowerCase() + name.slice(1), "ShapeStrides") : "outShapeStrides"; if (rank <= 1) { return "fn ".concat(funcName, "(index : i32) -> i32 { return index; }"); } var strides = tf.util.computeStrides(shape); var dtype = getCoordsDataType(rank); var coords = []; for (var i = 0; i < rank; i++) { coords.push("d".concat(i)); } if (strides.length === 1) { return " fn ".concat(funcName, "(index : i32) -> vec2 {\n let d0 = index / uniforms.").concat(stridesName, "; let d1 = index - d0 * uniforms.").concat(stridesName, ";\n return vec2(d0, d1);\n }"); } var snippet; snippet = 'var index2 = index;' + strides .map(function (_, i) { var line1 = "let ".concat(coords[i], " = index2 / uniforms.").concat(stridesName, ".").concat(getCoordsXYZ(i)); var line2 = i === strides.length - 1 ? "let ".concat(coords[i + 1], " = index2 - ").concat(coords[i], " * uniforms.").concat(stridesName, ".").concat(getCoordsXYZ(i)) : "index2 = index2 - ".concat(coords[i], " * uniforms.").concat(stridesName, ".").concat(getCoordsXYZ(i)); return "".concat(line1, "; ").concat(line2, ";"); }) .join(''); return "\n fn ".concat(funcName, "(index : i32) -> ").concat(dtype, " {\n ").concat(snippet, "\n return ").concat(dtype, "(").concat(coords.join(','), ");\n }\n "); } function getInputAtCoordsSnippet(inputInfo, component) { var texName = inputInfo.name; var rank = inputInfo.shape.length; var type = getCoordsDataType(rank); var funcName = 'get' + texName.charAt(0).toUpperCase() + texName.slice(1); var dims = ['d0', 'd1', 'd2', 'd3', 'd4', 'd5'].slice(0, rank); var inputs = dims.map(function (d) { return "".concat(d, " : i32"); }).join(', '); if (rank < 1) { return "\n fn ".concat(funcName, "() -> ").concat(typeSnippet(component), " {\n return ").concat(typeSnippet(component), "(").concat(texName, "[0]);\n }\n "); } var shapeStr = "uniforms.".concat(texName.charAt(0).toLowerCase() + texName.slice(1), "Shape"); var rankStr = "".concat(rank, "D"); if (rank === 0) { rankStr = '1D'; } return "\n fn ".concat(funcName, "(").concat(inputs, ") -> ").concat(typeSnippet(component), " {\n return ").concat(typeSnippet(component), "(").concat(texName, "[getIndexFromCoords").concat(rankStr, "(").concat(type, "(").concat(dims.join(','), "),\n ").concat(shapeStr, ")").concat(component === 1 ? '' : " / ".concat(component), "]);\n }\n "); } function getInputByOutputSnippet(inputInfo, outShape, component, isFlatDispatchLayout) { var texName = inputInfo.name; var texFuncSnippet = texName.charAt(0).toUpperCase() + texName.slice(1); var funcName = 'get' + texFuncSnippet + 'ByOutput'; var inRank = inputInfo.shape.length; var outRank = outShape.length; var type = getCoordsDataType(outRank); // If the inShape equals the outShape and the dispatch layout is flat, we can // directly use |gl_GlobalInvocationID.x| as the index and don't need coords // conversion between these two shapes. if (tf.util.arraysEqual(inputInfo.shape, outShape) && isFlatDispatchLayout) { return "\n fn ".concat(funcName, "Index(globalIndex : i32) -> ").concat(typeSnippet(component), " {\n return ").concat(typeSnippet(component), "(").concat(texName, "[globalIndex]);\n }\n\n fn ").concat(funcName, "Coords(coords : ").concat(type, ") -> ").concat(typeSnippet(component), " {\n return ").concat(typeSnippet(component), "(").concat(texName, "[").concat(outRank > 1 ? 'getOutputIndexFromCoords(coords)' : 'coords').concat(component === 1 ? '' : " / ".concat(component), "]);\n }\n "); } var broadcastDims = tf.backend_util.getBroadcastDims(inputInfo.shape, outShape); var rankDiff = outRank - inRank; var coordsSnippet = ''; if (inRank === 0) { return "\n fn ".concat(funcName, "Index(globalIndex : i32) -> ").concat(typeSnippet(component), "{\n return get").concat(texFuncSnippet, "();\n }\n\n fn ").concat(funcName, "Coords(coords : ").concat(type, ") -> ").concat(typeSnippet(component), "{\n return get").concat(texFuncSnippet, "();\n }\n "); } else { if (outRank < 2 && broadcastDims.length >= 1) { coordsSnippet = 'coords = 0;'; } else { coordsSnippet = broadcastDims.map(function (d) { return "coords.".concat(getCoordsXYZ(d + rankDiff), " = 0;"); }) .join('\n'); } } var unpackedCoordsSnippet = ''; if (outRank < 2 && inRank > 0) { unpackedCoordsSnippet = 'coords'; } else { if (outRank > 1) { var coordsType = getCoordsDataType(inRank); var coordsValues = inputInfo.shape.map(function (s, i) { return "coords.".concat(getCoordsXYZ(i + rankDiff)); }) .join(', '); unpackedCoordsSnippet = "".concat(coordsType, "(").concat(coordsValues, ")"); } else { unpackedCoordsSnippet = 'coords'; } } var shapeStr = "uniforms.".concat(texName.charAt(0).toLowerCase() + texName.slice(1), "Shape"); var rankStr = "".concat(inRank, "D"); return "\n fn ".concat(funcName, "Index(globalIndex : i32) -> ").concat(typeSnippet(component), " {\n var coords = getCoordsFromIndex(globalIndex);\n ").concat(coordsSnippet, "\n return ").concat(typeSnippet(component), "(").concat(texName, "[getIndexFromCoords").concat(rankStr, "(").concat(unpackedCoordsSnippet, ", ").concat(shapeStr, ")").concat(component === 1 ? '' : " / ".concat(component), "]);\n }\n\n fn ").concat(funcName, "Coords(coordsIn : ").concat(type, ") -> ").concat(typeSnippet(component), " {\n var coords = coordsIn;\n ").concat(coordsSnippet, "\n return ").concat(typeSnippet(component), "(").concat(texName, "[getIndexFromCoords").concat(rankStr, "(").concat(unpackedCoordsSnippet, ", ").concat(shapeStr, ")").concat(component === 1 ? '' : " / ".concat(component), "]);\n }\n"); } function getInputSnippet(inputInfo, outShape, component, isFlatDispatchLayout) { var res = getInputAtCoordsSnippet(inputInfo, component); var inShape = inputInfo.shape; if (inShape.length <= outShape.length) { res += getInputByOutputSnippet(inputInfo, outShape, component, isFlatDispatchLayout); } return res; } /** * Generates getOutputCoords() function that computes output coordinates * from dispatch geometry to reduce arithmetic. */ function getOutputCoordsSnippet(outShape, dispatchLayout) { var x = dispatchLayout.x, _a = dispatchLayout.y, y = _a === void 0 ? [] : _a, _b = dispatchLayout.z, z = _b === void 0 ? [] : _b; var outRank = outShape.length; var rank = x.length + y.length + z.length; // getOutputCoords is only meaningful when the output rank is same with // dispatch layout rank. if (rank !== outRank) { return ''; } if (x.length === outRank) { var dtype_1 = getCoordsDataType(outRank); var snippet_1 = "fn getOutputCoords() -> ".concat(dtype_1, "{\n let globalIndex = getGlobalIndex();\n return getCoordsFromIndex(globalIndex);\n }\n "); return snippet_1; } var gatherDimensionsStr = ''; var dims = [x, y, z]; for (var i = 0; i < dims.length; i++) { var arr = dims[i]; if (arr.length === 0) { continue; } if (arr.length === 1) { gatherDimensionsStr += "let d".concat(arr[0], " = i32(globalId[").concat(i, "]);"); } else { var strides = symbolicallyComputeStrides(arr, 'uniforms.outShape'); gatherDimensionsStr += "var index".concat(i, " = i32(globalId[").concat(i, "]);"); for (var j = 0; j < strides.length; j++) { gatherDimensionsStr += "let d".concat(arr[j], " = index").concat(i, " / ").concat(strides[j], ";"); if (j === strides.length - 1) { gatherDimensionsStr += "let d".concat(arr[j + 1], " = ") + "index".concat(i, " - d").concat(arr[j], " * ").concat(strides[j], ";"); } else { gatherDimensionsStr += "index".concat(i, " = index").concat(i, " - d").concat(arr[j], " * ").concat(strides[j], ";"); } } } } var dimensions = []; for (var i = 0; i < rank; i++) { dimensions.push("d".concat(i)); } var dtype = getCoordsDataType(rank); var snippet = "fn getOutputCoords() -> ".concat(dtype, " {\n ").concat(gatherDimensionsStr, "\n"); if (dimensions.length === 0) { snippet += "return ".concat(dtype, "(0); }"); } else { snippet += "return ".concat(dtype, "(").concat(dimensions.join(','), "); }"); } return snippet; } function getOutputIndexFromCoordsSnippet(outRank) { var snippet = ''; switch (outRank) { case 0: case 1: snippet += "\n fn getOutputIndexFromCoords(coords : i32) -> i32 {\n return coords;\n }\n "; break; case 2: snippet += "\n fn getOutputIndexFromCoords(coords : vec2) -> i32 {\n return dot(coords, vec2(uniforms.outShapeStrides, 1));\n }\n "; break; case 3: snippet += "\n fn getOutputIndexFromCoords(coords : vec3) -> i32 {\n return dot(coords, vec3(uniforms.outShapeStrides.x, uniforms.outShapeStrides.y, 1));\n }\n "; break; case 4: snippet += "\n fn getOutputIndexFromCoords(coords : vec4) -> i32 {\n return dot(coords, vec4(\n uniforms.outShapeStrides.x, uniforms.outShapeStrides.y, uniforms.outShapeStrides.z, 1));\n }\n "; break; case 5: snippet += "\n fn getOutputIndexFromCoords(coords : vec5) -> i32 {\n return coords.x * uniforms.outShapeStrides.x +\n coords.y * uniforms.outShapeStrides.y +\n coords.z * uniforms.outShapeStrides.z +\n coords.w * uniforms.outShapeStrides.w +\n coords.u;\n }\n "; break; case 6: snippet += "\n fn getOutputIndexFromCoords(coords : vec6) -> i32 {\n return coords.x * uniforms.outShapeStrides.x +\n coords.y * uniforms.outShapeStrides.y +\n coords.z * uniforms.outShapeStrides.z +\n coords.w * uniforms.outShapeStrides.w +\n coords.u * uniforms.outShapeStrides.u +\n coords.v;\n }\n "; break; default: tf.util.assert(false, function () { return "Unsupported ".concat(outRank, "D shape"); }); break; } return snippet; } function isFlatDispatch(program) { return program.dispatch[1] === 1 && program.dispatch[2] === 1; } function dataTypeToGPUType(type, component) { if (component === void 0) { component = 1; } if (type === 'float32') { return typeSnippet(component, 'f32'); } else if (type === 'int32' || type === 'bool') { return typeSnippet(component, 'i32'); } throw new Error("type ".concat(type, " is not supported.")); } function setOutputSnippet(outShape, outBufferType, component) { var outRank = outShape.length; var gpuType = dataTypeToGPUType(outBufferType, component); var snippet = "fn setOutputAtIndex(flatIndex : i32, value : ".concat(typeSnippet(component), ") {\n result[flatIndex] = ").concat(gpuType, "(value);\n }\n\n fn setOutputAtIndexI32(flatIndex : i32, value : ").concat(typeSnippet(component, 'i32'), ") {\n result[flatIndex] = ").concat(gpuType, "(value);\n }\n "); if (outRank >= 2) { var dims = ['d0', 'd1', 'd2', 'd3', 'd4', 'd5'].slice(0, outRank); var type = getCoordsDataType(outRank); snippet += "\n fn setOutputAtCoords(".concat(dims.map(function (d) { return "".concat(d, " : i32"); }).join(', '), ", value : ").concat(typeSnippet(component), ") {\n let flatIndex = getOutputIndexFromCoords(").concat(type, "(").concat(dims.join(', '), "));\n setOutputAtIndex(flatIndex").concat(component === 1 ? '' : " / ".concat(component), ", value);\n }\n fn setOutputAtCoordsI32(").concat(dims.map(function (d) { return "".concat(d, " : i32"); }).join(', '), ", value : ").concat(typeSnippet(component, 'i32'), ") {\n let flatIndex = getOutputIndexFromCoords(").concat(type, "(").concat(dims.join(', '), "));\n setOutputAtIndexI32(flatIndex").concat(component === 1 ? '' : " / ".concat(component), ", value);\n }\n "); } return snippet; } function insertAlignment(uniformShader) { // insert alignment when current pattern is vec5 or vec6 var curInsertRe = /(\w+)\s*:\s*vec(5|6)/g; uniformShader = uniformShader.replace(curInsertRe, function (match) { return '@align(16) ' + match; }); // insert alignment when previous pattern is vec5 or vec6 var preInsertRe = /vec(5|6)\s*,\s*(\w+)/g; uniformShader = uniformShader.replace(preInsertRe, function (_, p1, p2) { return "vec".concat(p1, ", @align(16) ").concat(p2); }); return uniformShader; } function isFlatDispatchLayout(program) { if (program.dispatchLayout.hasOwnProperty('y') && program.dispatchLayout.y.length !== 0) { return false; } if (program.dispatchLayout.hasOwnProperty('z') && program.dispatchLayout.z.length !== 0) { return false; } return true; } var arrayProduct = function (arr) { var product = 1; for (var i = 0; i < arr.length; i++) { product *= arr[i]; } return product; }; function tilesFitEvenlyIntoShape(tileSize, shape) { if (tileSize.length !== shape.length) { throw new Error("Cannot compute whether rank ".concat(tileSize.length) + " tiles fit evenly into rank ".concat(shape.length, " shape") + " - ranks must match."); } return shape.every(function (dim, dimIdx) { return dim % tileSize[dimIdx] === 0; }); } // Computes dispatch geometry based on layout of output dimensions and // workgroupSize. function computeDispatch(layout, outputShape, workgroupSize, elementsPerThread) { if (workgroupSize === void 0) { workgroupSize = [1, 1, 1]; } if (elementsPerThread === void 0) { elementsPerThread = [1, 1, 1]; } var _a = __read([ Math.ceil(arrayProduct(layout.x.map(function (d) { return outputShape[d]; })) / (workgroupSize[0] * elementsPerThread[0])), layout.y ? Math.ceil(arrayProduct(layout.y.map(function (d) { return outputShape[d]; })) / (workgroupSize[1] * elementsPerThread[1])) : 1, layout.z ? Math.ceil(arrayProduct(layout.z.map(function (d) { return outputShape[d]; })) / (workgroupSize[2] * elementsPerThread[2])) : 1 ], 3), dispatchX = _a[0], dispatchY = _a[1], dispatchZ = _a[2]; return [dispatchX, dispatchY, dispatchZ]; } function computeWorkgroupInfoForMatMul(dimAOuter, dimInner, dimBOuter, transposeA) { if (transposeA === void 0) { transposeA = false; } // These are experimental values. Usually, we need to adjust the work group // size based on the input shapes to improve the EU occupancy. // TODO: WebGPU limits the maximum allowed shared memory size as 16K. To make // sure it doesn't exceed this limitations. Temporarily reduce the work group // size to [8, 8, 1] and the work per thread size is [4, 4, 1]. But we should // revisit it and find the balance between work group size and work per thread // size. var workgroupSize = [8, 8, 1]; var elementsPerThread = [4, 4, 1]; if (!transposeA) { if (dimAOuter <= 8) { elementsPerThread[1] = 1; } if (dimInner <= 16 && dimBOuter <= 16) { workgroupSize[0] = 4; } } return { workgroupSize: workgroupSize, elementsPerThread: elementsPerThread }; } function computeWorkgroupSizeForConv2d(layout, outputShape, isVec4) { if (isVec4 === void 0) { isVec4 = false; } if (isVec4) { return [8, 8, 1]; } var dim0 = arrayProduct(layout.x.map(function (d) { return outputShape[d]; })); var dim1 = arrayProduct(layout.y.map(function (d) { return outputShape[d]; })); // TODO(jiajia.qin@intel.com): More fine tune based on outputShape. // These are experimental values. Usually, we need to adjust the work group // size based on the output shape. For example, when one dimension is smaller // than 4, it will be wasteful if we assign a larger size for this dimension, // which results lots of threads doing useless work and reduces parallelism // of hardware threads. But it is always a balance between work group size // and shared memory. If one dimension is too small, such as 1, shared memory // will won't be fully utilized. if (dim0 <= 4) { return [4, 16, 1]; } if (dim1 <= 4) { return [16, 4, 1]; } return [16, 16, 1]; } function computeWorkPerThreadForConv2d(layout, outputShape, isVec4) { if (isVec4 === void 0) { isVec4 = false; } if (isVec4) { return [4, 4, 1]; } var dim0 = arrayProduct(layout.x.map(function (d) { return outputShape[d]; })); var dim1 = arrayProduct(layout.y.map(function (d) { return outputShape[d]; })); // TODO(jiajia.qin@intel.com): More fine tune based on outputShape. // The following conditions correspond to the values set in // computeWorkgroupSizeForConv2d. if (dim0 <= 4) { return [1, 2, 1]; } if (dim1 <= 4) { return [2, 1, 1]; } return [2, 2, 1]; } function flatDispatchLayout(shape) { return { x: shape.map(function (d, i) { return i; }) }; } function GPUBytesPerElement(dtype) { if (dtype === 'float32' || dtype === 'int32' || dtype === 'bool' || dtype === 'string') { return 4; } else if (dtype === 'complex64') { return 8; } else { throw new Error("Unknown dtype ".concat(dtype)); } } function isWebGPUSupported() { return !!(typeof globalThis !== 'undefined' && (globalThis.navigator) && (globalThis.navigator.gpu)); } function assertNotComplex(tensor, opName) { if (!Array.isArray(tensor)) { tensor = [tensor]; } tensor.forEach(function (t) { if (t != null) { tf.util.assert(t.dtype !== 'complex64', function () { return "".concat(opName, " does not support complex64 tensors ") + 'in the WebGPU backend.'; }); } }); } var MatMulProgramType; (function (MatMulProgramType) { MatMulProgramType[MatMulProgramType["MatMulReduceProgram"] = 0] = "MatMulReduceProgram"; MatMulProgramType[MatMulProgramType["MatMulSplitKProgram"] = 1] = "MatMulSplitKProgram"; MatMulProgramType[MatMulProgramType["MatMulSmallOutputSizeProgram"] = 2] = "MatMulSmallOutputSizeProgram"; MatMulProgramType[MatMulProgramType["MatMulPackedProgram"] = 3] = "MatMulPackedProgram"; MatMulProgramType[MatMulProgramType["MatMulMax"] = 4] = "MatMulMax"; })(MatMulProgramType || (MatMulProgramType = {})); var webgpu_util = { __proto__: null, GPUBytesPerElement: GPUBytesPerElement, get MatMulProgramType () { return MatMulProgramType; }, assertNotComplex: assertNotComplex, computeDispatch: computeDispatch, computeWorkPerThreadForConv2d: computeWorkPerThreadForConv2d, computeWorkgroupInfoForMatMul: computeWorkgroupInfoForMatMul, computeWorkgroupSizeForConv2d: computeWorkgroupSizeForConv2d, flatDispatchLayout: flatDispatchLayout, isWebGPUSupported: isWebGPUSupported, tilesFitEvenlyIntoShape: tilesFitEvenlyIntoShape }; // Empirically determined constant used to determine size threshold for handing // off execution to the CPU. var CPU_HANDOFF_SIZE_THRESHOLD = tf.env().getNumber('WEBGPU_CPU_HANDOFF_SIZE_THRESHOLD'); // Reshape dispatch, not to exceed device limits. var reshapeDispatch = function (device, program) { var MAX_COMPUTE_PER_DIMENSION_DISPATCH_SIZE = device.limits.maxComputeWorkgroupsPerDimension; var layout = program['dispatchLayout']; var dispatch = program['dispatch']; if (dispatch.every(function (d) { return d <= MAX_COMPUTE_PER_DIMENSION_DISPATCH_SIZE; })) { return dispatch; } tf.util.assert(dispatch[0] > MAX_COMPUTE_PER_DIMENSION_DISPATCH_SIZE && layout.y === undefined && layout.z === undefined, function () { return 'Dispatch size exceeds WebGPU limits in Y or Z dimension.'; }); var dispatchAverage = Math.ceil(Math.sqrt(dispatch[0])); if (dispatchAverage > MAX_COMPUTE_PER_DIMENSION_DISPATCH_SIZE) { dispatchAverage = Math.ceil(Math.cbrt(dispatch[0])); tf.util.assert(dispatchAverage <= MAX_COMPUTE_PER_DIMENSION_DISPATCH_SIZE, function () { return 'Total dispatch size exceeds WebGPU maximum.'; }); return [dispatchAverage, dispatchAverage, dispatchAverage]; } else { return [dispatchAverage, dispatchAverage, 1]; } }; var WebGPUBackend = /** @class */ (function (_super) { __extends(WebGPUBackend, _super); function WebGPUBackend(device, adapterInfo) { var _this = _super.call(this) || this; _this.commandQueueOwnedIds = new WeakSet(); _this.dispatchCountInPass = 0; _this.disposed = false; _this.downloadWaitMs = 0; _this.tensorDataPendingDisposal = []; _this.queryResolveBuffer = null; _this.querySet = null; _this.querySetCount = 2; _this.stagingPendingDisposal = []; _this.uniformPendingDisposal = []; _this.uploadWaitMs = 0; _this.hasReadSyncWarned = false; _this.hasTimestampQueryWarned = false; if (!isWebGPUSupported()) { throw new Error('WebGPU is not supported on this device'); } _this.pipelineCache = {}; _this.device = device; _this.queue = device.queue; _this.commandEncoder = null; _this.computePassEncoder = null; _this.adapterInfo = new AdapterInfo(adapterInfo); _this.supportTimestampQuery = _this.device.features.has('timestamp-query'); _this.thresholdToIncreaseWorkgroups = _this.adapterInfo.intelGPUGeneration >= 12 ? 16 : 8; _this.bufferManager = new BufferManager(_this.device); _this.textureManager = new TextureManager(_this.device); _this.tensorMap = new tf.DataStorage(_this, tf.engine()); // Profiling tools like PIX needs this dummy canvas to // trigger capturing a frame. if (tf.env().getBool('WEBGPU_USE_PROFILE_TOOL')) { _this.dummyCanvas = document.createElement('canvas'); _this.dummyCanvas.width = 1; _this.dummyCanvas.height = 1; _this.dummyContext = _this.dummyCanvas.getContext('webgpu'); _this.dummyContext.configure({ device: device, format: 'bgra8unorm', }); document.body.appendChild(_this.dummyCanvas); } return _this; } WebGPUBackend.prototype.nextDataId = function () { return WebGPUBackend.nextDataId++; }; WebGPUBackend.prototype.floatPrecision = function () { return 32; }; /** * Dispose the memory if the dataId has 0 refCount. Return true if the memory * is released or delayed in this backend, false if there are still * references. * @param dataId * @oaram force Optional, remove the data regardless of refCount */ WebGPUBackend.prototype.disposeData = function (dataId, force) { if (force === void 0) { force = false; } // No-op if already disposed. if (!this.tensorMap.has(dataId)) { return true; } var tensorData = this.tensorMap.get(dataId); if (force) { tensorData.refCount = 0; } else { tensorData.refCount--; } if (tensorData.refCount > 0) { return false; } if (tensorData.complexTensorInfos != null) { this.disposeData(tensorData.complexTensorInfos.real.dataId); this.disposeData(tensorData.complexTensorInfos.imag.dataId); } if (this.commandQueueOwnedIds.has(dataId)) { this.tensorDataPendingDisposal.push(dataId); return true; } this.releaseResource(dataId); this.tensorMap.delete(dataId); return true; }; WebGPUBackend.prototype.memory = function () { return { numBytesInGPU: this.bufferManager.numBytesUsed, numBytesAllocatedInGPU: this.bufferManager.numBytesAllocated, unreliable: false }; }; WebGPUBackend.prototype.releaseResource = function (dataId) { var tensorData = this.tensorMap.get(dataId); if (!tensorData || !tensorData.resource) { return; } // If tensor's resource is from external, do not release. if (tensorData.external) { tensorData.resource = null; return; } if (tensorData.resource instanceof GPUBuffer) { this.bufferManager.releaseBuffer(tensorData.resource); } else if (tensorData.resource instanceof GPUTexture) { this.textureManager.releaseTexture(tensorData.resource); } tensorData.resource = null; }; /** Return refCount of a `TensorData`. */ WebGPUBackend.prototype.refCount = function (dataId) { if (this.tensorMap.has(dataId)) { var tensorData = this.tensorMap.get(dataId); return tensorData.refCount; } return 0; }; /** Increase refCount of a `TensorData`. */ WebGPUBackend.prototype.incRef = function (dataId) { var tensorData = this.tensorMap.get(dataId); tensorData.refCount++; }; /** Decrease refCount of a `TensorData`. */ WebGPUBackend.prototype.decRef = function (dataId) { if (this.tensorMap.has(dataId)) { var tensorData = this.tensorMap.get(dataId); tensorData.refCount--; } }; WebGPUBackend.prototype.write = function (values, shape, dtype) { if (dtype === 'complex64' && values != null) { throw new Error("Cannot write to a complex64 dtype. " + "Please use tf.complex(real, imag)."); } var dataId = { id: this.nextDataId() }; this.tensorMap.set(dataId, { dtype: dtype, shape: shape, values: values, refCount: 1 }); return dataId; }; WebGPUBackend.prototype.move = function (dataId, values, shape, dtype, refCount) { if (dtype === 'complex64') { throw new Error("Cannot write to a complex64 dtype. " + "Please use tf.complex(real, imag)."); } this.tensorMap.set(dataId, { dtype: dtype, shape: shape, values: values, refCount: refCount }); }; WebGPUBackend.prototype.submitQueue = function () { var _this = this; this.queue.submit([this.commandEncoder.finish()]); this.commandEncoder = null; this.dispatchCountInPass = 0; this.commandQueueOwnedIds = new WeakSet(); this.tensorDataPendingDisposal.forEach(function (d) { _this.releaseResource(d); _this.tensorMap.delete(d); }); this.uniformPendingDisposal.forEach(function (b) { return _this.bufferManager.releaseBuffer(b); }); this.stagingPendingDisposal.forEach(function (b) { return _this.bufferManager.releaseBuffer(b, false); }); this.tensorDataPendingDisposal = []; this.uniformPendingDisposal = []; this.stagingPendingDisposal = []; }; WebGPUBackend.prototype.ensureCommandEncoderReady = function () { if (!this.commandEncoder) { this.commandEncoder = this.device.createCommandEncoder(); } }; WebGPUBackend.prototype.endComputePassEncoder = function () { if (this.computePassEncoder) { this.computePassEncoder.end(); this.computePassEncoder = null; } }; // Check if parallel compilation is done. WebGPUBackend.prototype.checkCompileCompletionAsync = function () { return __awaiter(this, void 0, void 0, function () { var pipelines, e_1; var _this = this; return __generator(this, function (_b) { switch (_b.label) { case 0: _b.trys.push([0, 2, , 3]); return [4 /*yield*/, Promise.all(Object.values(this.pipelineCache))]; case 1: pipelines = _b.sent(); return [3 /*break*/, 3]; case 2: e_1 = _b.sent(); // TODO: Add test case to catch this exception. throw new Error(e_1.message); case 3: Object.keys(this.pipelineCache).map(function (key, i) { _this.pipelineCache[key] = pipelines[i]; }); return [2 /*return*/]; } }); }); }; WebGPUBackend.prototype.getBufferData = function (buffer) { return __awaiter(this, void 0, void 0, function () { var size, stagingBuffer, values; return __generator(this, function (_b) { switch (_b.label) { case 0: if (tf.env().getBool('WEBGPU_ENGINE_COMPILE_ONLY')) { console.warn('The data may be invalid since WEBGPU_ENGINE_COMPILE_ONLY is true, this can only be called when WEBGPU_ENGINE_COMPILE_ONLY is false'); return [2 /*return*/, null]; } size = buffer.size; stagingBuffer = this.bufferManager.acquireBuffer(size, GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ); this.ensureCommandEncoderReady(); this.endComputePassEncoder(); this.commandEncoder.copyBufferToBuffer(buffer, 0, stagingBuffer, 0, size); this.submitQueue(); return [4 /*yield*/, stagingBuffer.mapAsync(GPUMapMode.READ)]; case 1: _b.sent(); values = stagingBuffer.getMappedRange().slice(0); stagingBuffer.unmap(); if (stagingBuffer != null) { this.bufferManager.releaseBuffer(stagingBuffer); } // Need to get texture from swapChain to enable profiling tool // to capture a frame if (tf.env().getBool('WEBGPU_USE_PROFILE_TOOL')) { tf.util.assert(this.dummyContext !== undefined, function () { return "Fail to get context for profiling tool"; }); this.dummyContext.getCurrentTexture(); } return [2 /*return*/, values]; } }); }); }; WebGPUBackend.prototype.convertAndCacheOnCPU = function (dataId, data) { var tensorData = this.tensorMap.get(dataId); tensorData.values = data; return tensorData.values; }; WebGPUBackend.prototype.readSync = function (dataId) { var _this = this; var tensorData = this.tensorMap.get(dataId); var values = tensorData.values, complexTensorInfos = tensorData.complexTensorInfos; if (values != null || tensorData.dtype === 'string') { return values; } if (tensorData.dtype === 'complex64') { var realValues = this.readSync(complexTensorInfos.real.dataId); var imagValues = this.readSync(complexTensorInfos.imag.dataId); var complexVals = tf.util.convertBackendValuesAndArrayBuffer(tf.backend_util.mergeRealAndImagArrays(realValues, imagValues).buffer, 'float32'); this.convertAndCacheOnCPU(dataId, complexVals); return complexVals; } if (!this.hasReadSyncWarned) { this.hasReadSyncWarned = true; console.warn("The performance of synchronously reading data from GPU to CPU is " + "poor on the webgpu backend, please use asynchronous APIs instead."); } var alphaModes = ['opaque', 'premultiplied']; var buffer = tensorData.resource; var bufferSize = buffer.size; tf.util.assert(bufferSize % 4 === 0, function () { return 'Because there is 4 bytes for ' + 'one pixel, buffer size must be multiple of 4.'; }); var pixelsSize = bufferSize / 4; var valsGPU = new ArrayBuffer(bufferSize); // TODO: adjust the reading window size according the `bufferSize`. var canvasWidth = 256, canvasHeight = 256; var stagingDeviceStorage = alphaModes.map(function (_) { return new OffscreenCanvas(canvasWidth, canvasHeight); }); var stagingHostStorage = new OffscreenCanvas(canvasWidth, canvasHeight); this.endComputePassEncoder(); stagingDeviceStorage .map(function (storage, index) { var context = storage.getContext('webgpu'); // TODO: use rgba8unorm format when this format is supported on Mac. // https://bugs.chromium.org/p/chromium/issues/detail?id=1298618 context.configure({ device: _this.device, format: 'bgra8unorm', usage: GPUTextureUsage.COPY_DST, alphaMode: alphaModes[index], }); return context.getCurrentTexture(); }) .map(function (texture, index) { var bytesPerRow = canvasWidth * 4; var readDataGPUToCPU = function (width, height, offset) { _this.ensureCommandEncoderReady(); _this.commandEncoder.copyBufferToTexture({ buffer: buffer, bytesPerRow: bytesPerRow, offset: offset, }, { texture: texture, }, { width: width, height: height, }); _this.submitQueue(); var context = stagingHostStorage.getContext('2d', { willReadFrequently: true, }); context.clearRect(0, 0, width, height); context.drawImage(stagingDeviceStorage[index], 0, 0); var stagingValues = context.getImageData(0, 0, width, height).data; var alphaMode = alphaModes[index]; var span = new Uint8ClampedArray(valsGPU, offset, width * height * 4); for (var k = 0; k < span.length; k += 4) { if (alphaMode === 'premultiplied') { span[k + 3] = stagingValues[k + 3]; } else { var value = stagingValues[k]; span[k] = stagingValues[k + 2]; span[k + 1] = stagingValues[k + 1]; span[k + 2] = value; } } }; var fullyReadCount = Math.floor(pixelsSize / (canvasWidth * canvasHeight)); var width = canvasWidth, height = canvasHeight, offset = 0; for (var i = 0; i < fullyReadCount; i++) { // Read the buffer data, which fully fill the whole canvas. readDataGPUToCPU(width, height, offset); offset += canvasWidth * canvasHeight * 4; } var remainSize = pixelsSize % (canvasWidth * canvasHeight); height = Math.floor(remainSize / canvasWidth); if (height > 0) { // Read the buffer data, which fully fill certain rows of canvas. readDataGPUToCPU(width, height, offset); offset += height * (canvasWidth * 4); } width = remainSize % canvasWidth; if (width > 0) { // Read the buffer data, which not fully fill one row of canvas. readDataGPUToCPU(width, 1, offset); } }); var vals = tf.util.convertBackendValuesAndArrayBuffer(valsGPU, tensorData.dtype); this.convertAndCacheOnCPU(dataId, vals); return vals; }; WebGPUBackend.prototype.read = function (dataId) { return __awaiter(this, void 0, void 0, function () { var tensorData, values, vals, ps, realValues, imagValues, data; return __generator(this, function (_b) { switch (_b.label) { case 0: if (!this.tensorMap.has(dataId)) { throw new Error("Tensor ".concat(dataId, " was not registered!")); } tensorData = this.tensorMap.get(dataId); values = tensorData.values; if (values != null) { return [2 /*return*/, values]; } if (!(tensorData.dtype === 'complex64')) return [3 /*break*/, 2]; return [4 /*yield*/, Promise.all([ this.read(tensorData.complexTensorInfos.real.dataId), this.read(tensorData.complexTensorInfos.imag.dataId) ])]; case 1: ps = _b.sent(); realValues = ps[0]; imagValues = ps[1]; vals = tf.backend_util.mergeRealAndImagArrays(realValues, imagValues); return [3 /*break*/, 4]; case 2: return [4 /*yield*/, this.getBufferData(tensorData.resource)]; case 3: data = _b.sent(); vals = tf.util.convertBackendValuesAndArrayBuffer(data, tensorData.dtype); _b.label = 4; case 4: this.convertAndCacheOnCPU(dataId, vals); return [2 /*return*/, vals]; } }); }); }; // The source GPUBuffer and destination GPUBuffer have the same size and // usage. WebGPUBackend.prototype.copyBuffer = function (srcBuffer) { var size = srcBuffer.size; var usage = srcBuffer.usage; var dstBuffer = this.bufferManager.acquireBuffer(size, usage); this.ensureCommandEncoderReady(); this.endComputePassEncoder(); this.commandEncoder.copyBufferToBuffer(srcBuffer, 0, dstBuffer, 0, size); this.submitQueue(); return dstBuffer; }; /** * Create a TF.js tensor out of an existing WebGPU buffer. */ WebGPUBackend.prototype.createTensorFromGPUData = function (webGPUData, shape, dtype) { var buffer = webGPUData.buffer; if (dtype === 'complex64') { throw new Error("Cannot write to a complex64 dtype. "); } var dataId = { id: this.nextDataId() }; this.tensorMap.set(dataId, { dtype: dtype, shape: shape, values: null, refCount: 1, external: webGPUData.zeroCopy }); var tensorData = this.tensorMap.get(dataId); var size = GPUBytesPerElement(tensorData.dtype) * tf.util.sizeFromShape(tensorData.shape); if (webGPUData.buffer.size < size) { throw new Error("GPUBuffer size(".concat(webGPUData.buffer.size, ") is smaller than tensor size(").concat(size, ")!")); } else if ((webGPUData.buffer.usage & (GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC)) !== (GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC)) { throw new Error('GPUBuffer.usage should include GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC!'); } // Do buffer copy by default. if (webGPUData.zeroCopy !== true) { buffer = this.copyBuffer(buffer); } tensorData.resource = buffer; return tf.engine().makeTensorFromDataId(dataId, shape, dtype, this); }; /** * Read tensor to a new GPUBuffer. * @param dataId The source tensor. */ WebGPUBackend.prototype.readToGPU = function (dataId) { var srcTensorData = this.tensorMap.get(dataId); var values = srcTensorData.values, dtype = srcTensorData.dtype, shape = srcTensorData.shape, resource = srcTensorData.resource; if (dtype === 'complex64') { throw new Error('Does not support reading buffer for complex64 dtype.'); } if (resource == null) { if (values != null) { throw new Error('Data is not on GPU but on CPU.'); } else { throw new Error('There is no data on GPU or CPU.'); } } var srcBuffer = resource; var size = srcBuffer.size; var usage = srcBuffer.usage; var buffer = this.bufferManager.acquireBuffer(size, usage); this.ensureCommandEncoderReady(); this.endComputePassEncoder(); this.commandEncoder.copyBufferToBuffer(resource, 0, buffer, 0, size); this.submitQueue(); var tensorInfo = this.makeTensorInfo(shape, dtype); // Make engine track this tensor, so that we can dispose it later. var tensorRef = tf.engine().makeTensorFromTensorInfo(tensorInfo); var tensorData = this.tensorMap.get(tensorInfo.dataId); tensorData.resource = buffer; return { tensorRef: tensorRef, buffer: buffer }; }; WebGPUBackend.prototype.bufferSync = function (t) { var data = this.readSync(t.dataId); if (t.dtype === 'string') { try { // Decode the bytes into string. var strings = data.map(function (d) { return tf.util.decodeString(d); }); return tf.buffer(t.shape, t.dtype, strings); } catch (_a) { throw new Error('Failed to decode encoded string bytes into utf-8'); } } return tf.buffer(t.shape, t.dtype, data); }; WebGPUBackend.prototype.time = function (f) { return __awaiter(this, void 0, void 0, function () { var oldActiveTimers, newActiveTimers, outerMostTime, flattenedActiveTimerQueries, flattenedActiveTimerNames, res, kernelMs; return __generator(this, function (_b) { switch (_b.label) { case 0: if (!this.supportTimestampQuery && !this.hasTimestampQueryWarned) { console.warn("This device doesn't support timestamp-query extension. " + "Start Chrome browser with flag " + "--enable-dawn-features=allow_unsafe_apis to try it again. " + "Otherwise, zero will be shown for the kernel time when profiling " + "mode is enabled."); this.hasTimestampQueryWarned = true; } oldActiveTimers = this.activeTimers; newActiveTimers = []; outerMostTime = false; if (this.programTimersStack == null) { this.programTimersStack = newActiveTimers; outerMostTime = true; } else { this.activeTimers.push(newActiveTimers); } this.activeTimers = newActiveTimers; f(); flattenedActiveTimerQueries = tf.util.flatten(this.activeTimers.map(function (d) { return d.query; })) .filter(function (d) { return d != null; }); flattenedActiveTimerNames = tf.util.flatten(this.activeTimers.map(function (d) { return d.name; })) .filter(function (d) { return d != null; }); this.activeTimers = oldActiveTimers; if (outerMostTime) { this.programTimersStack = null; } res = { uploadWaitMs: this.uploadWaitMs, downloadWaitMs: this.downloadWaitMs, kernelMs: null, wallMs: null }; return [4 /*yield*/, Promise.all(flattenedActiveTimerQueries)]; case 1: kernelMs = _b.sent(); res['kernelMs'] = tf.util.sum(kernelMs); res['getExtraProfileInfo'] = function () { return kernelMs.map(function (d, i) { return ({ name: flattenedActiveTimerNames[i], ms: d }); }) .map(function (d) { return "".concat(d.name, ": ").concat(d.ms); }) .join(', '); }; this.uploadWaitMs = 0; this.downloadWaitMs = 0; return [2 /*return*/, res]; } }); }); }; WebGPUBackend.prototype.makeTensorInfo = function (shape, dtype, values) { if (dtype === 'string' && values != null && values.length > 0 && tf.util.isString(values[0])) { values = values.map(function (d) { return tf.util.encodeString(d); }); } var dataId = this.write(values, shape, dtype); return { dataId: dataId, shape: shape, dtype: dtype }; }; WebGPUBackend.prototype.tensorToBinding = function (tensor) { if (!tensor) { return null; } var tensorData = this.tensorMap.get(tensor.dataId); var resource = tensorData.resource; if (resource instanceof GPUBuffer) { return { buffer: resource }; } if (resource instanceof GPUTexture) { return resource.createView(); } // GPUExternalTexture return resource; }; WebGPUBackend.prototype.uploadToGPU = function (dataId) { var tensorData = this.tensorMap.get(dataId); // Already on the GPU. if (tensorData.resource != null) { return; } var size = GPUBytesPerElement(tensorData.dtype) * tf.util.sizeFromShape(tensorData.shape); var buffer; var usage = GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST; if (tensorData.values) { buffer = this.bufferManager.acquireBuffer(size, usage, true); if (buffer.mapState === 'unmapped') { var stagingBuffer = this.bufferManager.acquireBuffer(size, GPUBufferUsage.MAP_WRITE | GPUBufferUsage.COPY_SRC, true, false); var arrayBuffer = stagingBuffer.getMappedRange(); if (tensorData.dtype === 'int32' || tensorData.dtype === 'bool') { new Int32Array(arrayBuffer).set(tensorData.values); } else { new Float32Array(arrayBuffer).set(tensorData.values); } stagingBuffer.unmap(); this.ensureCommandEncoderReady(); this.endComputePassEncoder(); this.commandEncoder.copyBufferToBuffer(stagingBuffer, 0, buffer, 0, size); this.stagingPendingDisposal.push(stagingBuffer); } else { var arrayBuffer = buffer.getMappedRange(); if (tensorData.dtype === 'int32' || tensorData.dtype === 'bool') { new Int32Array(arrayBuffer).set(tensorData.values); } else { new Float32Array(arrayBuffer).set(tensorData.values); } buffer.unmap(); } // Once uploaded, don't store the values on cpu. tensorData.values = null; } else { buffer = this.bufferManager.acquireBuffer(size, usage); } tensorData.resource = buffer; }; WebGPUBackend.prototype.makeUniforms = function (programUniform) { var currentOffset = 0; var preLength = 0; var offsets = []; var maxAlignmentOfField = 1; programUniform.forEach(function (d) { if (d.data.length === 0) { d.data = [1]; } // https://www.w3.org/TR/WGSL/#alignof var baseAlignment; switch (d.data.length) { case 1: baseAlignment = 4; break; case 2: baseAlignment = 8; break; case 3: baseAlignment = 16; break; case 4: baseAlignment = 16; break; case 5: baseAlignment = 16; break; case 6: baseAlignment = 16; break; default: tf.util.assert(false, function () { return "Unsupported ".concat(d.data.length, "D shape"); }); } if (preLength === 5 || preLength === 6) { baseAlignment = 16; } if (baseAlignment > maxAlignmentOfField) { maxAlignmentOfField = baseAlignment; } currentOffset = Math.ceil(currentOffset / baseAlignment) * baseAlignment; preLength = d.data.length; offsets.push(currentOffset); currentOffset += d.data.length * 4; }); currentOffset = Math.ceil(currentOffset / maxAlignmentOfField) * maxAlignmentOfField; var arrayBuffer = new ArrayBuffer(currentOffset); programUniform.forEach(function (d, i) { var offset = offsets[i]; if (d.type === 'int32') { new Int32Array(arrayBuffer, offset, d.data.length).set(d.data); } else if (d.type === 'uint32') { new Uint32Array(arrayBuffer, offset, d.data.length).set(d.data); } else { new Float32Array(arrayBuffer, offset, d.data.length).set(d.data); } }); var uniformBuffer = this.bufferManager.acquireBuffer(currentOffset, GPUBufferUsage.COPY_DST | GPUBufferUsage.UNIFORM); this.queue.writeBuffer(uniformBuffer, 0, arrayBuffer, 0, currentOffset); this.uniformPendingDisposal.push(uniformBuffer); return { offset: 0, size: currentOffset, buffer: uniformBuffer }; }; WebGPUBackend.prototype.runWebGPUProgram = function (program, inputs, outputDtype, programDefinedUniform, output) { var _this = this; if (!output) { output = this.makeTensorInfo(program.outputShape, outputDtype); } if (tf.util.sizeFromShape(output.shape) === 0) { // Short-circuit the computation since the result is empty (has 0 in its // shape). this.tensorMap.get(output.dataId).values = tf.util.getTypedArrayFromDType(output.dtype, 0); return output; } this.uploadToGPU(output.dataId); program.dispatch = reshapeDispatch(this.device, program); var inputsData = inputs.map(function (input, i) { if (input.dtype === 'complex64') { throw new Error("GPGPUProgram does not support complex64 input. For complex64 " + "dtypes, please separate the program into real and imaginary " + "parts."); } _this.uploadToGPU(input.dataId); return { // Returning dtype from tensorMap because it reflects dtype // of underlying buffer, rather than abstract dtype. dtype: _this.tensorMap.get(input.dataId).dtype, shape: input.shape, name: program.variableNames[i] }; }); program.shaderKey = makeShaderKey(program, inputsData, output); var parallelCompilation = tf.env().getBool('WEBGPU_ENGINE_COMPILE_ONLY'); if (!(program.shaderKey in this.pipelineCache)) { this.pipelineCache[program.shaderKey] = compileProgram(this.device, program, inputsData, output, parallelCompilation); } program.pipeline = this.pipelineCache[program.shaderKey]; if (!parallelCompilation) { this.recordAndSubmit(program, output, inputs, programDefinedUniform); } return output; }; WebGPUBackend.prototype.recordAndSubmit = function (program, output, inputs, programDefinedUniform) { var _this = this; if (program.pipeline instanceof Promise) { throw new Error('Please call checkCompileCompletionAsync to ensure parallel compilation is done!'); } // There are six kinds of uniforms: NAN, INFINITY, shapes, shape strides, // program size, program defined uniforms. var programUniform = []; var bufferShapes = []; var uniformsType = 'int32'; if (program.pixelsOpType == null) { programUniform.push({ type: 'float32', data: [NaN] }, { type: 'float32', data: [Infinity] }); bufferShapes = inputs.concat(output).map(function (d) { return d.shape; }); var uniformsType_1 = 'int32'; bufferShapes.map(function (d) { programUniform.push({ type: uniformsType_1, data: d }); var strides = tf.util.computeStrides(d); programUniform.push({ type: uniformsType_1, data: strides }); }); } else { var strides = tf.util.computeStrides(output.shape); programUniform.push({ type: uniformsType, data: strides }); } if (program.size) { var size = tf.util.sizeFromShape(program.outputShape); programUniform.push({ type: uniformsType, data: [program.outputComponent ? size / program.outputComponent : size] }); } if (programDefinedUniform) { programUniform = __spreadArray(__spreadArray([], __read(programUniform), false), __read(programDefinedUniform), false); } var bindings = __spreadArray(__spreadArray([ this.tensorToBinding(output) ], __read(inputs.map(function (t) { return _this.tensorToBinding(t); })), false), [ this.makeUniforms(programUniform) ], false); inputs.forEach(function (input) { _this.commandQueueOwnedIds.add(input.dataId); }); this.commandQueueOwnedIds.add(output.dataId); var bindGroup = this.device.createBindGroup({ layout: program.pipeline.getBindGroupLayout(0), entries: bindings.map(function (b, i) { return ({ binding: i, resource: b }); }), }); var shouldTimeProgram = this.activeTimers != null; this.ensureCommandEncoderReady(); var computePassDescriptor = {}; if (shouldTimeProgram && this.supportTimestampQuery) { this.endComputePassEncoder(); if (this.querySet == null) { this.querySet = this.device.createQuerySet({ type: 'timestamp', count: this.querySetCount, }); } computePassDescriptor.timestampWrites = { querySet: this.querySet, beginningOfPassWriteIndex: 0, endOfPassWriteIndex: 1, }; this.computePassEncoder = this.commandEncoder.beginComputePass(computePassDescriptor); } else if (!this.computePassEncoder) { this.computePassEncoder = this.commandEncoder.beginComputePass(computePassDescriptor); } this.computePassEncoder.setPipeline(program.pipeline); this.computePassEncoder.setBindGroup(0, bindGroup); this.computePassEncoder.dispatchWorkgroups(program.dispatch[0], program.dispatch[1], program.dispatch[2]); this.dispatchCountInPass++; if (shouldTimeProgram || tf.env().get('WEBGPU_DEFERRED_SUBMIT_BATCH_SIZE') <= this.dispatchCountInPass || program.pixelsOpType === PixelsOpType.DRAW) { this.endComputePassEncoder(); if (shouldTimeProgram) { this.activeTimers.push({ name: program.constructor.name, query: this.getQueryTime() }); } else { this.submitQueue(); } } }; WebGPUBackend.prototype.getQueryTime = function () { return __awaiter(this, void 0, void 0, function () { var queryStagingBuffer, arrayBuffer, time; return __generator(this, function (_b) { switch (_b.label) { case 0: if (!this.supportTimestampQuery) { return [2 /*return*/, 0]; } if (this.queryResolveBuffer == null) { this.queryResolveBuffer = this.bufferManager.acquireBuffer(this.querySetCount * 8, GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST | GPUBufferUsage.QUERY_RESOLVE); } this.commandEncoder.resolveQuerySet(this.querySet, 0, this.querySetCount, this.queryResolveBuffer, 0); queryStagingBuffer = this.bufferManager.acquireBuffer(this.querySetCount * 8, GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST); this.commandEncoder.copyBufferToBuffer(this.queryResolveBuffer, 0, queryStagingBuffer, 0, this.querySetCount * 8); this.submitQueue(); return [4 /*yield*/, queryStagingBuffer.mapAsync(GPUMapMode.READ)]; case 1: _b.sent(); arrayBuffer = new BigUint64Array(queryStagingBuffer.getMappedRange()); time = Number(arrayBuffer[1] - arrayBuffer[0]) / 1000000; queryStagingBuffer.unmap(); this.bufferManager.releaseBuffer(queryStagingBuffer); return [2 /*return*/, time]; } }); }); }; WebGPUBackend.prototype.shouldExecuteOnCPU = function (inputs, sizeThreshold) { var _this = this; if (sizeThreshold === void 0) { sizeThreshold = CPU_HANDOFF_SIZE_THRESHOLD; } return tf.env().getBool('WEBGPU_CPU_FORWARD') && inputs.every(function (input) { return _this.tensorMap.get(input.dataId).resource == null && tf.util.sizeFromShape(input.shape) < sizeThreshold; }); }; WebGPUBackend.prototype.numDataIds = function () { return this.tensorMap.numDataIds() - this.tensorDataPendingDisposal.length; }; WebGPUBackend.prototype.dispose = function () { if (this.disposed) { return; } if (this.querySet != null) { this.querySet.destroy(); } this.bufferManager.dispose(); this.textureManager.dispose(); this.disposed = true; }; return WebGPUBackend; }(tf.KernelBackend)); WebGPUBackend.nextDataId = 0; if (isWebGPUSupported()) { tf.registerBackend('webgpu', function () { return __awaiter(void 0, void 0, void 0, function () { var gpuDescriptor, adapter, deviceDescriptor, requiredFeatures, adapterLimits, device, adapterInfo, _a, _b; return __generator(this, function (_c) { switch (_c.label) { case 0: gpuDescriptor = { powerPreference: tf.env().get('WEBGPU_USE_LOW_POWER_GPU') ? 'low-power' : 'high-performance' }; return [4 /*yield*/, navigator.gpu.requestAdapter(gpuDescriptor)]; case 1: adapter = _c.sent(); deviceDescriptor = {}; requiredFeatures = []; if (adapter.features.has('timestamp-query')) { requiredFeatures.push('timestamp-query'); } if (adapter.features.has('bgra8unorm-storage')) { requiredFeatures.push(['bgra8unorm-storage']); } deviceDescriptor.requiredFeatures = requiredFeatures; adapterLimits = adapter.limits; deviceDescriptor.requiredLimits = { 'maxComputeWorkgroupStorageSize': adapterLimits.maxComputeWorkgroupStorageSize, 'maxComputeWorkgroupsPerDimension': adapterLimits.maxComputeWorkgroupsPerDimension, 'maxStorageBufferBindingSize': adapterLimits.maxStorageBufferBindingSize, 'maxBufferSize': adapterLimits.maxBufferSize, 'maxComputeWorkgroupSizeX': adapterLimits.maxComputeWorkgroupSizeX, 'maxComputeInvocationsPerWorkgroup': adapterLimits.maxComputeInvocationsPerWorkgroup, }; return [4 /*yield*/, adapter.requestDevice(deviceDescriptor)]; case 2: device = _c.sent(); if (!('info' in adapter)) return [3 /*break*/, 3]; _a = adapter.info; return [3 /*break*/, 7]; case 3: if (!('requestAdapterInfo' in adapter)) return [3 /*break*/, 5]; return [4 /*yield*/, adapter.requestAdapterInfo()]; case 4: _b = _c.sent(); return [3 /*break*/, 6]; case 5: _b = undefined; _c.label = 6; case 6: _a = _b; _c.label = 7; case 7: adapterInfo = _a; return [2 /*return*/, new WebGPUBackend(device, adapterInfo)]; } }); }); }, 3 /*priority*/); } /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var BinaryOpType; (function (BinaryOpType) { BinaryOpType[BinaryOpType["ADD"] = 0] = "ADD"; BinaryOpType[BinaryOpType["ATAN2"] = 1] = "ATAN2"; BinaryOpType[BinaryOpType["COMPLEX_MULTIPLY_IMAG"] = 2] = "COMPLEX_MULTIPLY_IMAG"; BinaryOpType[BinaryOpType["COMPLEX_MULTIPLY_REAL"] = 3] = "COMPLEX_MULTIPLY_REAL"; BinaryOpType[BinaryOpType["DIV"] = 4] = "DIV"; BinaryOpType[BinaryOpType["ELU_DER"] = 5] = "ELU_DER"; BinaryOpType[BinaryOpType["EQUAL"] = 6] = "EQUAL"; BinaryOpType[BinaryOpType["FLOOR_DIV"] = 7] = "FLOOR_DIV"; BinaryOpType[BinaryOpType["GREATER"] = 8] = "GREATER"; BinaryOpType[BinaryOpType["GREATER_EQUAL"] = 9] = "GREATER_EQUAL"; BinaryOpType[BinaryOpType["LESS"] = 10] = "LESS"; BinaryOpType[BinaryOpType["LESS_EQUAL"] = 11] = "LESS_EQUAL"; BinaryOpType[BinaryOpType["LOGICAL_AND"] = 12] = "LOGICAL_AND"; BinaryOpType[BinaryOpType["LOGICAL_OR"] = 13] = "LOGICAL_OR"; BinaryOpType[BinaryOpType["MAX"] = 14] = "MAX"; BinaryOpType[BinaryOpType["MIN"] = 15] = "MIN"; BinaryOpType[BinaryOpType["MOD"] = 16] = "MOD"; BinaryOpType[BinaryOpType["MUL"] = 17] = "MUL"; BinaryOpType[BinaryOpType["NOT_EQUAL"] = 18] = "NOT_EQUAL"; BinaryOpType[BinaryOpType["POW"] = 19] = "POW"; BinaryOpType[BinaryOpType["PRELU"] = 20] = "PRELU"; BinaryOpType[BinaryOpType["SQUARED_DIFFERENCE"] = 21] = "SQUARED_DIFFERENCE"; BinaryOpType[BinaryOpType["SUB"] = 22] = "SUB"; })(BinaryOpType || (BinaryOpType = {})); var ADD = 'let resultTemp = a + b;'; var ATAN2 = 'let resultTemp = atan2(a, b);'; // (Ar + Ai)(Br + Bi) = // ArBr + ArBi + AiBr + AiBi = ArBr - AB + ArBi + AiBr // Yr = ArBr - AB // Yi = ArBi + AiBr var COMPLEX_MULTIPLY_REAL = 'let resultTemp = areal * breal - aimag * bimag;'; var COMPLEX_MULTIPLY_IMAG = 'let resultTemp = areal * bimag + aimag * breal;'; var DIV = 'let resultTemp = a / b;'; var ELU_DER = 'let resultTemp = select(a * (b + 1.0), a, b >= b - b);'; var EQUAL = "\n let zero = sign(a) * 0 + 0;\n let one = sign(b) * 0 + 1;\n let resultTemp = select(zero, one, a == b);\n"; var FLOOR_DIV = "\n let remainder =\n select(a % b, round(a % b), (round(a) == a) & (round(b) == b));\n let quotient = (a - remainder) / b;\n let resultTemp =\n round(select(quotient, quotient - 1, sign(remainder) == -sign(b)));\n"; var GREATER = "\n let zero = sign(a) * 0 + 0;\n let one = sign(b) * 0 + 1;\n let resultTemp = select(zero, one, a > b);\n"; var GREATER_EQUAL = "\n let zero = sign(a) * 0 + 0;\n let one = sign(b) * 0 + 1;\n let resultTemp = select(zero, one, a >= b);\n"; var LESS = "\n let zero = sign(a) * 0 + 0;\n let one = sign(b) * 0 + 1;\n let resultTemp = select(zero, one, a < b);\n"; var LESS_EQUAL = "\n let zero = sign(a) * 0 + 0;\n let one = sign(b) * 0 + 1;\n let resultTemp = select(zero, one, a <= b);\n"; var LOGICAL_AND = 'return f32(a >= 1.0 && b >= 1.0);'; var LOGICAL_AND_VEC4 = "return (vec4(a >= vec4(1.0)) *\n vec4(b >= vec4(1.0)));"; var LOGICAL_OR = 'return f32(a >= 1.0 || b >= 1.0);'; var LOGICAL_OR_VEC4 = "return min(vec4(a >= vec4(1.0)) +\n vec4(b >= vec4(1.0)), vec4(1.0));"; var MAX = 'let resultTemp = max(a, b);'; var MIN = 'let resultTemp = min(a, b);'; var MOD = "\n let isNaN = b == 0.;\n var resultTemp = a % b;\n resultTemp = select((resultTemp + b) % b, resultTemp,\n (a < 0. && b < 0.) || (a >= 0. && b > 0.));\n"; var MOD_VEC4 = "\n let isNaN = !vec4(b);\n var resultTemp = vec4(a % b);\n if (!((a[0] < 0. && b[0] < 0.) || (a[0] >= 0. && b[0] > 0.))) {\n resultTemp[0] = (resultTemp[0] + b[0]) % b[0];\n }\n if (!((a[1] < 0. && b[1] < 0.) || (a[1] >= 0. && b[1] > 0.))) {\n resultTemp[1] = (resultTemp[1] + b[1]) % b[1];\n }\n if (!((a[2] < 0. && b[2] < 0.) || (a[2] >= 0. && b[2] > 0.))) {\n resultTemp[2] = (resultTemp[2] + b[2]) % b[2];\n }\n if (!((a[3] < 0. && b[3] < 0.) || (a[3] >= 0. && b[3] > 0.))) {\n resultTemp[3] = (resultTemp[3] + b[3]) % b[3];\n }\n"; var MUL = 'let resultTemp = a * b;'; var NOT_EQUAL = "\n var resultTemp = f32(a != b);\n let valueForNaN = 1.0;\n"; var NOT_EQUAL_VEC4 = "\n var resultTemp = vec4(a != b);\n let valueForNaN = 1.0;\n"; var POW = "\n let isNaN = a < 0.0 && floor(b) < b;\n if (b == 0.0) {\n return 1.0;\n }\n var resultTemp = select(sign(a) * pow(abs(a), b), pow(abs(a), b),\n round(abs(b) % 2.0) != 1.0);\n"; var POW_VEC4 = "\n let isModRound1Bool = vec4(round(abs(b) % vec4(2.0))) == vec4(1);\n let isModRound1 = vec4(isModRound1Bool);\n let multiplier = sign(a) * isModRound1 + (vec4(1.0) - isModRound1);\n var resultTemp = multiplier * pow(abs(a), b);\n\n // Ensure that a^0 = 1, including 0^0 = 1 as this correspond to TF and JS\n let isExpZero = b == vec4(0.0);\n if (isExpZero.r) {\n resultTemp.r = 1.0;\n }\n if (isExpZero.g) {\n resultTemp.g = 1.0;\n }\n if (isExpZero.b) {\n resultTemp.b = 1.0;\n }\n if (isExpZero.a) {\n resultTemp.a = 1.0;\n }\n let isNaN = (a < vec4(0.0)) & (floor(b) < b);\n"; var PRELU = "if (a < 0.0) { return b * a; } return a;"; var PRELU_VEC4 = "\n let aLessThanZero = vec4(a < vec4(0.0));\n return (aLessThanZero * (b * a)) + ((vec4(1.0) - aLessThanZero) * a);\n"; var SQUARED_DIFFERENCE = 'let resultTemp = (a - b) * (a - b);'; var SUB = 'let resultTemp = a - b;'; function getBinaryOpString(type, useVec4) { var doOpSnippet; // Ops with NaN check do { switch (type) { case BinaryOpType.ATAN2: doOpSnippet = ATAN2; break; case BinaryOpType.MAX: doOpSnippet = MAX; break; case BinaryOpType.MIN: doOpSnippet = MIN; break; case BinaryOpType.MOD: doOpSnippet = useVec4 ? MOD_VEC4 : MOD; break; case BinaryOpType.NOT_EQUAL: doOpSnippet = useVec4 ? NOT_EQUAL_VEC4 : NOT_EQUAL; break; case BinaryOpType.POW: doOpSnippet = useVec4 ? POW_VEC4 : POW; break; default: continue; } var isNaN = void 0; var dTypeN = void 0; var boolN = void 0; if (useVec4) { isNaN = 'isnanVec4'; dTypeN = 'vec4'; boolN = 'vec4'; } else { isNaN = 'isnan'; dTypeN = 'f32'; boolN = 'bool'; } return "\n let aIsNaN = ".concat(isNaN, "(a);\n let aPostLegalization = select(a, ").concat(dTypeN, "(42), aIsNaN);\n let bIsNaN = ").concat(isNaN, "(b);\n let bPostLegalization = select(b, ").concat(dTypeN, "(42), bIsNaN);\n let isNaN = false;\n let valueForNaN = uniforms.NAN;\n {\n let a = aPostLegalization;\n let b = bPostLegalization;\n ").concat(doOpSnippet, "\n return select(\n resultTemp, ").concat(dTypeN, "(valueForNaN),\n ").concat(boolN, "(isNaN) | aIsNaN | bIsNaN);\n }\n "); } while (false); // Ops without NaN check switch (type) { case BinaryOpType.ADD: doOpSnippet = ADD; break; case BinaryOpType.COMPLEX_MULTIPLY_IMAG: doOpSnippet = COMPLEX_MULTIPLY_IMAG; break; case BinaryOpType.COMPLEX_MULTIPLY_REAL: doOpSnippet = COMPLEX_MULTIPLY_REAL; break; case BinaryOpType.DIV: doOpSnippet = DIV; break; case BinaryOpType.ELU_DER: doOpSnippet = ELU_DER; break; case BinaryOpType.EQUAL: doOpSnippet = EQUAL; break; case BinaryOpType.FLOOR_DIV: doOpSnippet = FLOOR_DIV; break; case BinaryOpType.GREATER: doOpSnippet = GREATER; break; case BinaryOpType.GREATER_EQUAL: doOpSnippet = GREATER_EQUAL; break; case BinaryOpType.LESS: doOpSnippet = LESS; break; case BinaryOpType.LESS_EQUAL: doOpSnippet = LESS_EQUAL; break; case BinaryOpType.LOGICAL_AND: return useVec4 ? LOGICAL_AND_VEC4 : LOGICAL_AND; case BinaryOpType.LOGICAL_OR: return useVec4 ? LOGICAL_OR_VEC4 : LOGICAL_OR; case BinaryOpType.MUL: doOpSnippet = MUL; break; case BinaryOpType.PRELU: return useVec4 ? PRELU_VEC4 : PRELU; case BinaryOpType.SQUARED_DIFFERENCE: doOpSnippet = SQUARED_DIFFERENCE; break; case BinaryOpType.SUB: doOpSnippet = SUB; break; // throw new Error(`BinaryType ${type} is not implemented!`); } return "\n ".concat(doOpSnippet, "\n return resultTemp;\n "); } /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var UnaryOpType; (function (UnaryOpType) { UnaryOpType[UnaryOpType["ABS"] = 0] = "ABS"; UnaryOpType[UnaryOpType["ACOS"] = 1] = "ACOS"; UnaryOpType[UnaryOpType["ACOSH"] = 2] = "ACOSH"; UnaryOpType[UnaryOpType["ASIN"] = 3] = "ASIN"; UnaryOpType[UnaryOpType["ASINH"] = 4] = "ASINH"; UnaryOpType[UnaryOpType["ATAN"] = 5] = "ATAN"; UnaryOpType[UnaryOpType["ATANH"] = 6] = "ATANH"; UnaryOpType[UnaryOpType["CEIL"] = 7] = "CEIL"; UnaryOpType[UnaryOpType["COS"] = 8] = "COS"; UnaryOpType[UnaryOpType["COSH"] = 9] = "COSH"; UnaryOpType[UnaryOpType["ELU"] = 10] = "ELU"; UnaryOpType[UnaryOpType["ERF"] = 11] = "ERF"; UnaryOpType[UnaryOpType["EXP"] = 12] = "EXP"; UnaryOpType[UnaryOpType["EXPM1"] = 13] = "EXPM1"; UnaryOpType[UnaryOpType["FLOOR"] = 14] = "FLOOR"; UnaryOpType[UnaryOpType["IS_FINITE"] = 15] = "IS_FINITE"; UnaryOpType[UnaryOpType["IS_INF"] = 16] = "IS_INF"; UnaryOpType[UnaryOpType["IS_NAN"] = 17] = "IS_NAN"; UnaryOpType[UnaryOpType["LINEAR"] = 18] = "LINEAR"; UnaryOpType[UnaryOpType["LOG"] = 19] = "LOG"; UnaryOpType[UnaryOpType["LOG1P"] = 20] = "LOG1P"; UnaryOpType[UnaryOpType["LOGICAL_NOT"] = 21] = "LOGICAL_NOT"; UnaryOpType[UnaryOpType["NEG"] = 22] = "NEG"; UnaryOpType[UnaryOpType["RELU"] = 23] = "RELU"; UnaryOpType[UnaryOpType["RELU6"] = 24] = "RELU6"; UnaryOpType[UnaryOpType["LEAKYRELU"] = 25] = "LEAKYRELU"; UnaryOpType[UnaryOpType["RECIPROCAL"] = 26] = "RECIPROCAL"; UnaryOpType[UnaryOpType["ROUND"] = 27] = "ROUND"; UnaryOpType[UnaryOpType["RSQRT"] = 28] = "RSQRT"; UnaryOpType[UnaryOpType["SELU"] = 29] = "SELU"; UnaryOpType[UnaryOpType["SIGMOID"] = 30] = "SIGMOID"; UnaryOpType[UnaryOpType["SIGN"] = 31] = "SIGN"; UnaryOpType[UnaryOpType["SIN"] = 32] = "SIN"; UnaryOpType[UnaryOpType["SINH"] = 33] = "SINH"; UnaryOpType[UnaryOpType["SOFTPLUS"] = 34] = "SOFTPLUS"; UnaryOpType[UnaryOpType["SQRT"] = 35] = "SQRT"; UnaryOpType[UnaryOpType["SQUARE"] = 36] = "SQUARE"; UnaryOpType[UnaryOpType["STEP"] = 37] = "STEP"; UnaryOpType[UnaryOpType["TAN"] = 38] = "TAN"; UnaryOpType[UnaryOpType["TANH"] = 39] = "TANH"; UnaryOpType[UnaryOpType["TO_INT"] = 40] = "TO_INT"; })(UnaryOpType || (UnaryOpType = {})); var ABS = "return abs(a);"; var ACOS = "\n if (abs(a) > 1.) {\n return uniforms.NAN;\n }\n return acos(a);\n"; var ACOSH = "\n if (a < 1.) {\n return uniforms.NAN;\n }\n return acosh(a);\n"; var ASIN = "\n if (abs(a) > 1.) {\n return uniforms.NAN;\n }\n return asin(a);\n"; var ASINH = "return asinh(a);"; var ATAN = "\n if (isnan(a)) {\n return uniforms.NAN;\n }\n return atan(a);\n"; var ATANH = "\n if (abs(a) > 1.) {\n return uniforms.NAN;\n }\n if (a == 1.) {\n return uniforms.INFINITY;\n }\n if (a == -1.) {\n return -uniforms.INFINITY;\n }\n return atanh(a);\n"; var CEIL = "return ceil(a);"; var COS = "return cos(a);"; var COSH = "\n let e2x = exp(-a);\n return (e2x + 1.0 / e2x) / 2.0;\n"; var EXPM1 = "return exp(a) - 1.0;"; var ELU = "if (a >= 0.0) { return a; } return (exp(a) - 1.0);"; var ELU_VEC4 = "\n var resFloat = exp(a) - vec4(1.0);\n if (a.r >= 0.0) {\n resFloat.r = a.r;\n }\n if (a.g >= 0.0) {\n resFloat.g = a.g;\n }\n if (a.b >= 0.0) {\n resFloat.b = a.b;\n }\n if (a.a >= 0.0) {\n resFloat.a = a.a;\n }\n return resFloat;\n"; var ERF = "\n // Error function is calculated approximately with elementary function.\n // See \"Handbook of Mathematical Functions with Formulas,\n // Graphs, and Mathematical Tables\", Abramowitz and Stegun.\n let p = ".concat(tf.backend_util.ERF_P, ";\n let a1 = ").concat(tf.backend_util.ERF_A1, ";\n let a2 = ").concat(tf.backend_util.ERF_A2, ";\n let a3 = ").concat(tf.backend_util.ERF_A3, ";\n let a4 = ").concat(tf.backend_util.ERF_A4, ";\n let a5 = ").concat(tf.backend_util.ERF_A5, ";\n\n let sign = sign(a);\n let absA = abs(a);\n let t = 1.0 / (1.0 + p * absA);\n return sign * (1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * exp(-absA * absA));\n"); var EXP = "return exp(a);"; var FLOOR = "return floor(a);"; var IS_FINITE = "return f32(!isnan(a) && !isinf(a));"; var IS_INF = "return f32(isinf(a));"; var IS_NAN = "return f32(isnan(a));"; var LINEAR = "return a;"; var LOG = "if (a < 0.0) { return uniforms.NAN; }\n return log(a);"; var LOG1P = "\n if (isnan(a)) { return a; }\n return log(1.0 + a);\n"; var LOGICAL_NOT = "return f32(!(a >= 1.0));"; var NEG = "return -a;"; var LEAKYRELU = "if (a < 0.0) { return uniforms.alpha * a; } return a;"; var LEAKYRELU_VEC4 = "\n let aLessThanZero = vec4(a < vec4(0.0));\n return (aLessThanZero * (uniforms.alpha * a)) + ((vec4(1.0) - aLessThanZero) * a);\n"; var RECIPROCAL = "return 1.0 / a;"; var RELU = "return select(a, 0.0, a < 0.0);"; var RELU6 = 'return clamp(a, 0.0, 6.0);'; var RELU6_VEC4 = 'return clamp(a, vec4(0.0, 0.0, 0.0, 0.0), vec4(6.0, 6.0, 6.0, 6.0));'; var RELU_VEC4 = "\n return select(a, vec4(0.0), a < vec4(0.0));\n"; var ROUND = "return round(a);"; var RSQRT = "return inverseSqrt(a);"; // Stable and Attracting Fixed Point (0, 1) for Normalized Weights. // See: https://arxiv.org/abs/1706.02515 var SELU = "\n if (a >= 0.0) {\n return ".concat(tf.backend_util.SELU_SCALE, " * a;\n } else {\n return ").concat(tf.backend_util.SELU_SCALEALPHA, " * (exp(a) - 1.0);\n }\n"); var SIGMOID = "return 1.0 / (1.0 + exp(-1.0 * a));"; var SIGN = "return sign(a);"; var SIN = "return sin(a);"; var SINH = "\n let e2x = exp(a);\n return (e2x - 1.0 / e2x) / 2.0;\n"; var SOFTPLUS = "\n let epsilon = 1.1920928955078125e-7;\n let threshold = log(epsilon) + 2.0;\n\n let too_large = a > -threshold;\n let too_small = a < threshold;\n let exp_a = exp(a);\n\n if (too_large) {\n return a;\n } else if (too_small) {\n return exp_a;\n } else {\n return log(exp_a + 1.0);\n }\n"; var SQRT = "return sqrt(a);"; var SQUARE = "return a * a;"; var STEP = "\n if (isnan(a)) {\n return a;\n }\n\n return select(uniforms.stepAlpha, 1.0, a > 0.0);\n"; var TAN = "return tan(a);"; var TANH = "\n let e2x = exp(-2.0 * abs(a));\n return sign(a) * (1.0 - e2x) / (1.0 + e2x);\n"; var TO_INT = "return f32(i32((a)));"; function getUnaryOpString(type, useVec4) { switch (type) { case UnaryOpType.ABS: return ABS; case UnaryOpType.ACOS: return ACOS; case UnaryOpType.ACOSH: return ACOSH; case UnaryOpType.ASIN: return ASIN; case UnaryOpType.ASINH: return ASINH; case UnaryOpType.ATAN: return ATAN; case UnaryOpType.ATANH: return ATANH; case UnaryOpType.COS: return COS; case UnaryOpType.COSH: return COSH; case UnaryOpType.CEIL: return CEIL; case UnaryOpType.ELU: return useVec4 ? ELU_VEC4 : ELU; case UnaryOpType.ERF: return ERF; case UnaryOpType.EXP: return EXP; case UnaryOpType.EXPM1: return EXPM1; case UnaryOpType.FLOOR: return FLOOR; case UnaryOpType.IS_FINITE: return IS_FINITE; case UnaryOpType.IS_INF: return IS_INF; case UnaryOpType.IS_NAN: return IS_NAN; case UnaryOpType.LINEAR: return LINEAR; case UnaryOpType.LOG: return LOG; case UnaryOpType.LOG1P: return LOG1P; case UnaryOpType.LOGICAL_NOT: return LOGICAL_NOT; case UnaryOpType.NEG: return NEG; case UnaryOpType.LEAKYRELU: return useVec4 ? LEAKYRELU_VEC4 : LEAKYRELU; case UnaryOpType.RECIPROCAL: return RECIPROCAL; case UnaryOpType.RELU: return useVec4 ? RELU_VEC4 : RELU; case UnaryOpType.RELU6: return useVec4 ? RELU6_VEC4 : RELU6; case UnaryOpType.ROUND: return ROUND; case UnaryOpType.RSQRT: return RSQRT; case UnaryOpType.SELU: return SELU; case UnaryOpType.SIGMOID: return SIGMOID; case UnaryOpType.SIGN: return SIGN; case UnaryOpType.SIN: return SIN; case UnaryOpType.SINH: return SINH; case UnaryOpType.SOFTPLUS: return SOFTPLUS; case UnaryOpType.SQRT: return SQRT; case UnaryOpType.SQUARE: return SQUARE; case UnaryOpType.STEP: return STEP; case UnaryOpType.TAN: return TAN; case UnaryOpType.TANH: return TANH; case UnaryOpType.TO_INT: return TO_INT; default: throw new Error("BinaryType ".concat(type, " is not implemented!")); } } /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function activationFnSnippet(activation, hasPreluActivationWeights, packed, coordsLength) { if (hasPreluActivationWeights === void 0) { hasPreluActivationWeights = false; } if (packed === void 0) { packed = false; } if (coordsLength === void 0) { coordsLength = 3; } if (activation === null) { return ''; } var activationOpSnippet = ''; if (activation === 'linear') { activationOpSnippet = getUnaryOpString(UnaryOpType.LINEAR); } else if (activation === 'relu') { activationOpSnippet = getUnaryOpString(UnaryOpType.RELU, packed); } else if (activation === 'elu') { activationOpSnippet = getUnaryOpString(UnaryOpType.ELU, packed); } else if (activation === 'relu6') { activationOpSnippet = getUnaryOpString(UnaryOpType.RELU6, packed); } else if (activation === 'prelu') { activationOpSnippet = getBinaryOpString(BinaryOpType.PRELU, packed); } else if (activation === 'sigmoid') { activationOpSnippet = getUnaryOpString(UnaryOpType.SIGMOID, packed); } else if (activation === 'leakyrelu') { activationOpSnippet = getUnaryOpString(UnaryOpType.LEAKYRELU, packed); } else { throw new Error("Activation ".concat(activation, " has not been implemented for the WebGPU backend.")); } var elementSize = packed ? 4 : 1; var dataType = typeSnippet(elementSize); var activationFnSnippet = ''; if (hasPreluActivationWeights) { activationFnSnippet = "\n fn activation(a : ".concat(dataType, ", coords : vec").concat(coordsLength, ") -> ").concat(dataType, " {\n let b = getPreluActivationWeightsByOutputCoords(coords);\n ").concat(activationOpSnippet, "\n }"); } else { activationFnSnippet = "\n fn activation(a : ".concat(dataType, ", coords : vec").concat(coordsLength, ") -> ").concat(dataType, " {\n ").concat(activationOpSnippet, "\n }"); } return activationFnSnippet; } function biasActivationSnippet(hasBias, activation) { return "\n ".concat(hasBias ? 'value = value + getBiasByOutputCoords(coords);' : '', "\n ").concat(activation ? 'value = activation(value, coords);' : '', "\n "); } function matMulReadFnSource(transposeA, transposeB, fitAOuter, fitBOuter, fitInner, component) { if (fitAOuter === void 0) { fitAOuter = false; } if (fitInner === void 0) { fitInner = false; } if (component === void 0) { component = 1; } tf.util.assert(transposeA && component === 1 || !transposeA, function () { return "transposeA ".concat(transposeA, " is not compatible with component size ").concat(component); }); var sampleA = "\n ".concat(transposeA ? "value = getA(batch, col, row);" : "value = getA(batch, row, col);", "\n\n "); var sampleB = transposeB ? "value = getB(batch, col, row);" : "value = getB(batch, row, col);"; return "\n fn mm_readA(batch: i32, row: i32, col: i32) -> ".concat(typeSnippet(component), " {\n var value = ").concat(typeSnippet(component), "(0.0);\n ").concat(fitAOuter && fitInner ? sampleA : "\n ".concat(transposeA ? "if(row < uniforms.dimAOuter && col < uniforms.dimInner)" : "if(row < uniforms.aShape[1] && col < uniforms.aShape[2])", "\n {\n ").concat(sampleA, "\n }\n "), "\n return value;\n }\n\n fn mm_readB(batch: i32, row: i32, col: i32) -> ").concat(typeSnippet(component), " {\n var value = ").concat(typeSnippet(component), "(0.0);\n ").concat(sampleB, "\n return value;\n }\n "); } function matMulReadWriteFnSource(hasBias, activation, transposeA, transposeB, fitAOuter, fitBOuter, fitInner, component) { if (fitAOuter === void 0) { fitAOuter = false; } if (fitBOuter === void 0) { fitBOuter = false; } if (fitInner === void 0) { fitInner = false; } if (component === void 0) { component = 1; } return "\n ".concat(matMulReadFnSource(transposeA, transposeB, fitAOuter, fitBOuter, fitInner, component), "\n fn mm_write(batch: i32, row: i32, col: i32, valueIn: ").concat(typeSnippet(component), ") {\n ").concat(fitAOuter && fitBOuter ? '' : 'if (row < uniforms.dimAOuter && col < uniforms.dimBOuter)', "\n {\n var value = valueIn;\n let coords = vec3(batch, row, col);\n ").concat(biasActivationSnippet(hasBias, activation), "\n setOutputAtCoords(coords[0], coords[1], coords[2], value);\n }\n }\n "); } var writeDataToSubAVec4Snippet = function (transpose, innerElementSize) { if (transpose) { return "\n mm_Asub[inputRow][inputCol] = mm_readA(batchA,\n kStart + inputRow,\n globalRowStart + inputCol * ".concat(innerElementSize, ");\n "); } else { return "\n mm_Asub[inputRow][inputCol] = mm_readA(batchA,\n globalRow + innerRow,\n kStart + inputCol * ".concat(innerElementSize, ");\n "); } }; var calculateResultSnippet = function (transposeA, innerElementSize, rowPerThread, tileInner) { if (transposeA) { return "\n for (var k = 0; k < ".concat(tileInner, "; k++) {\n let BCached0 = mm_Bsub[k][tileCol];\n let ACached0 = mm_Asub[k][localRow];\n for (var i = 0; i < ").concat(rowPerThread, "; i++) {\n acc[i] = fma(BCached0, vec4(ACached0[i]), acc[i]);\n }\n }"); } else { var bCachedStr = ''; var accStr = ''; for (var i = 0; i < innerElementSize; i++) { bCachedStr += "let BCached".concat(i, " = mm_Bsub[k * ").concat(innerElementSize, " + ").concat(i, "][tileCol];"); accStr += "acc[i] = fma(BCached".concat(i, ", vec4(ACached[").concat(i, "]), acc[i]);"); } return "\n for (var k = 0; k < ".concat(tileInner / innerElementSize, "; k++) {\n ").concat(bCachedStr, "\n for (var i = 0; i < ").concat(rowPerThread, "; i++) {\n let ACached = mm_Asub[tileRow + i][k];\n ").concat(accStr, "\n }\n }"); } }; function makeMatMulPackedVec4Source(workPerThread, workgroupSize, transposeA, tileInner, splitK, splitedDimInner, broadcastBatch) { if (transposeA === void 0) { transposeA = false; } if (tileInner === void 0) { tileInner = 32; } if (splitK === void 0) { splitK = false; } if (splitedDimInner === void 0) { splitedDimInner = 32; } if (broadcastBatch === void 0) { broadcastBatch = false; } var tileAOuter = workgroupSize[1] * workPerThread[1]; var tileBOuter = workgroupSize[0] * workPerThread[0]; var tileAWidth = transposeA ? tileAOuter : tileInner; var tileAHight = transposeA ? tileInner : tileAOuter; var innerElementSize = tileAWidth / workgroupSize[0]; var rowPerThreadB = tileInner / workgroupSize[1]; var rowPerThread = workPerThread[1]; var colPerThread = workPerThread[0]; tf.util.assert(((transposeA && innerElementSize === 4 && workPerThread[1] === 4) || (!transposeA && (innerElementSize === 3 || innerElementSize === 4))) && tileAWidth % workgroupSize[0] === 0 && tileInner % workgroupSize[1] === 0 && workPerThread[0] === 4, function () { return "If transposeA ".concat(transposeA, " is true, innerElementSize ").concat(innerElementSize, " and workPerThread[1] ").concat(workPerThread[1], " must be 4.\n Otherwise, innerElementSize ").concat(innerElementSize, " must be 3 or 4.\n tileAWidth ").concat(tileAWidth, " must be divisible by workgroupSize[0]").concat(workgroupSize[0], ". tileInner ").concat(tileInner, " must be divisible by workgroupSize[1] ").concat(workgroupSize[1], ". colPerThread ").concat(workPerThread[0], " must be 4."); }); return "\n var mm_Asub : array, ").concat(tileAWidth / innerElementSize, ">, ").concat(tileAHight, ">;\n var mm_Bsub : array, ").concat(tileBOuter / workPerThread[0], ">, ").concat(tileInner, ">;\n\n ").concat(getMainHeaderString(), " {\n let localRow = i32(localId.y);\n let tileRow = localRow * ").concat(rowPerThread, ";\n let tileCol = i32(localId.x);\n\n let globalRow = i32(globalId.y) * ").concat(rowPerThread, ";\n let globalCol = i32(globalId.x) * ").concat(colPerThread, ";\n let batch = ").concat(splitK ? '0' : 'i32(globalId.z)', ";\n let batchA = ").concat(splitK || !broadcastBatch ? 'batch' : 'batch % uniforms.aShape[0]', ";\n let batchB = ").concat(splitK || !broadcastBatch ? 'batch' : 'batch % uniforms.bShape[0]', ";\n let globalRowStart = i32(workgroupId.y) * ").concat(tileAOuter, ";\n\n let numTiles = ").concat(splitK ? "".concat(Math.ceil(splitedDimInner / tileInner)) : "(uniforms.dimInner - 1) / ".concat(tileInner, " + 1"), ";\n var kStart = ").concat(splitK ? "i32(globalId.z) * ".concat(splitedDimInner) : '0', ";\n\n var acc: array, ").concat(rowPerThread, ">;\n\n // Loop over shared dimension.\n let tileRowB = localRow * ").concat(rowPerThreadB, ";\n for (var t = 0; t < numTiles; t++) {\n // Load one tile of A into local memory.\n for (var innerRow = 0; innerRow < ").concat(rowPerThread, "; innerRow++) {\n let inputRow = tileRow + innerRow;\n let inputCol = tileCol;\n ").concat(writeDataToSubAVec4Snippet(transposeA, innerElementSize), "\n }\n\n // Load one tile of B into local memory.\n for (var innerRow = 0; innerRow < ").concat(rowPerThreadB, "; innerRow++) {\n let inputRow = tileRowB + innerRow;\n let inputCol = tileCol;\n mm_Bsub[inputRow][inputCol] = mm_readB(batchB, kStart + inputRow, globalCol);\n }\n kStart = kStart + ").concat(tileInner, ";\n workgroupBarrier();\n\n // Compute acc values for a single thread.\n ").concat(calculateResultSnippet(transposeA, innerElementSize, rowPerThread, tileInner), "\n workgroupBarrier();\n }\n\n for (var innerRow = 0; innerRow < ").concat(rowPerThread, "; innerRow++) {\n mm_write(batch, globalRow + innerRow, globalCol, acc[innerRow]);\n }\n }"); } var writeDataToSubASnippet = function (transpose) { if (transpose) { return "\n mm_Asub[inputRow][inputCol] = mm_readA(batchA,\n kStart + inputRow,\n globalRowStart + inputCol);\n "; } else { return "\n mm_Asub[inputRow][inputCol] = mm_readA(batchA,\n globalRowStart + inputRow,\n kStart + inputCol);\n "; } }; var readDataFromSubASnippet = function (transposeA) { return transposeA ? 'let ACached = mm_Asub[k][tileRow + innerRow];' : 'let ACached = mm_Asub[tileRow + innerRow][k];'; }; // sequentialAccessByThreads means sequential data in memory is accessed by // threads, instead of a single thread (default behavior). function makeMatMulPackedSource(workPerThread, workgroupSize, transposeA, tileInner, splitK, splitedDimInner, sequentialAccessByThreads, broadcastBatch) { if (transposeA === void 0) { transposeA = false; } if (tileInner === void 0) { tileInner = 32; } if (splitK === void 0) { splitK = false; } if (splitedDimInner === void 0) { splitedDimInner = 32; } if (sequentialAccessByThreads === void 0) { sequentialAccessByThreads = false; } if (broadcastBatch === void 0) { broadcastBatch = false; } var tileAOuter = workPerThread[1] * workgroupSize[1]; var tileBOuter = workPerThread[0] * workgroupSize[0]; var tileAWidth = transposeA ? tileAOuter : tileInner; var tileAHight = transposeA ? tileInner : tileAOuter; tf.util.assert(tileAHight % workgroupSize[1] === 0 && tileAWidth % workgroupSize[0] === 0 && tileInner % workgroupSize[1] === 0, function () { return "tileAHight ".concat(tileAHight, " must be divisible by workgroupSize[1]").concat(workgroupSize[1], ", tileAWidth ").concat(tileAWidth, " must be divisible by workgroupSize[0]").concat(workgroupSize[0], ", tileInner ").concat(tileInner, " must be divisible by workgroupSize[1]").concat(workgroupSize[1]); }); var rowPerThreadA = tileAHight / workgroupSize[1]; var colPerThreadA = tileAWidth / workgroupSize[0]; var rowPerThreadB = tileInner / workgroupSize[1]; var rowPerThread = workPerThread[1]; var colPerThread = workPerThread[0]; var matmulSnippet = sequentialAccessByThreads ? "\n let localRow = i32(localId.y);\n let localCol = i32(localId.x);\n let globalRowStart = i32(workgroupId.y) * ".concat(tileAOuter, ";\n let globalColStart = i32(workgroupId.x) * ").concat(tileBOuter, ";\n\n // Loop over shared dimension.\n for (var t = 0; t < numTiles; t++) {\n // Load one tile of A into local memory.\n for (var inputRow = localRow; inputRow < ").concat(tileAHight, "; inputRow = inputRow + ").concat(workgroupSize[1], ") {\n for (var inputCol = localCol; inputCol < ").concat(tileAWidth, "; inputCol = inputCol + ").concat(workgroupSize[0], ") {\n ").concat(writeDataToSubASnippet(transposeA), "\n }\n }\n // Load one tile of B into local memory.\n for (var inputRow = localRow; inputRow < ").concat(tileInner, "; inputRow = inputRow + ").concat(workgroupSize[1], ") {\n for (var inputCol = localCol; inputCol < ").concat(tileBOuter, "; inputCol = inputCol + ").concat(workgroupSize[0], ") {\n mm_Bsub[inputRow][inputCol] = mm_readB(batchB,\n kStart + inputRow,\n globalColStart + inputCol);\n }\n }\n kStart = kStart + ").concat(tileInner, ";\n workgroupBarrier();\n\n // Compute acc values for a single thread.\n var BCached : array;\n for (var k = 0; k < ").concat(tileInner, "; k++) {\n for (var inner = 0; inner < ").concat(colPerThread, "; inner++) {\n BCached[inner] = mm_Bsub[k][localCol + inner * ").concat(workgroupSize[0], "];\n }\n for (var innerRow = 0; innerRow < ").concat(rowPerThread, "; innerRow++) {\n let ACached = ").concat(transposeA ? "mm_Asub[k][localRow + innerRow * ".concat(workgroupSize[1], "];") : "mm_Asub[localRow + innerRow * ".concat(workgroupSize[1], "][k];"), "\n for (var innerCol = 0; innerCol < ").concat(colPerThread, "; innerCol++) {\n acc[innerRow][innerCol] =\n fma(ACached, BCached[innerCol], acc[innerRow][innerCol]);\n }\n }\n }\n workgroupBarrier();\n }\n for (var innerRow = 0; innerRow < ").concat(rowPerThread, "; innerRow++) {\n let gRow = globalRowStart + localRow + innerRow * ").concat(workgroupSize[1], ";\n for (var innerCol = 0; innerCol < ").concat(colPerThread, "; innerCol++) {\n let gCol = globalColStart + localCol + innerCol * ").concat(workgroupSize[0], ";\n mm_write(batch, gRow, gCol, acc[innerRow][innerCol]);\n }\n }\n ") : "\n let tileRow = i32(localId.y) * ".concat(rowPerThread, ";\n let tileCol = i32(localId.x) * ").concat(colPerThread, ";\n\n let globalRow = i32(globalId.y) * ").concat(rowPerThread, ";\n let globalCol = i32(globalId.x) * ").concat(colPerThread, ";\n let globalRowStart = i32(workgroupId.y) * ").concat(tileAOuter, ";\n\n let tileRowA = i32(localId.y) * ").concat(rowPerThreadA, ";\n let tileColA = i32(localId.x) * ").concat(colPerThreadA, ";\n let tileRowB = i32(localId.y) * ").concat(rowPerThreadB, ";\n // Loop over shared dimension.\n for (var t = 0; t < numTiles; t++) {\n // Load one tile of A into local memory.\n for (var innerRow = 0; innerRow < ").concat(rowPerThreadA, "; innerRow++) {\n for (var innerCol = 0; innerCol < ").concat(colPerThreadA, "; innerCol++) {\n let inputRow = tileRowA + innerRow;\n let inputCol = tileColA + innerCol;\n ").concat(writeDataToSubASnippet(transposeA), "\n }\n }\n\n // Load one tile of B into local memory.\n for (var innerRow = 0; innerRow < ").concat(rowPerThreadB, "; innerRow++) {\n for (var innerCol = 0; innerCol < ").concat(colPerThread, "; innerCol++) {\n let inputRow = tileRowB + innerRow;\n let inputCol = tileCol + innerCol;\n mm_Bsub[inputRow][inputCol] = mm_readB(batchB,\n kStart + inputRow,\n globalCol + innerCol);\n }\n }\n kStart = kStart + ").concat(tileInner, ";\n workgroupBarrier();\n\n // Compute acc values for a single thread.\n var BCached : array;\n for (var k = 0; k < ").concat(tileInner, "; k++) {\n for (var inner = 0; inner < ").concat(colPerThread, "; inner++) {\n BCached[inner] = mm_Bsub[k][tileCol + inner];\n }\n\n for (var innerRow = 0; innerRow < ").concat(rowPerThread, "; innerRow++) {\n ").concat(readDataFromSubASnippet(transposeA), "\n for (var innerCol = 0; innerCol < ").concat(colPerThread, "; innerCol++) {\n acc[innerRow][innerCol] =\n fma(ACached, BCached[innerCol], acc[innerRow][innerCol]);\n }\n }\n }\n\n workgroupBarrier();\n }\n\n for (var innerRow = 0; innerRow < ").concat(rowPerThread, "; innerRow++) {\n for (var innerCol = 0; innerCol < ").concat(colPerThread, "; innerCol++) {\n mm_write(batch, globalRow + innerRow, globalCol + innerCol,\n acc[innerRow][innerCol]);\n }\n }\n "); return "\n var mm_Asub : array, ").concat(tileAHight, ">;\n var mm_Bsub : array, ").concat(tileInner, ">;\n\n ").concat(getMainHeaderString(), " {\n let batch = ").concat(splitK ? '0' : 'i32(globalId.z)', ";\n let batchA = ").concat(splitK || !broadcastBatch ? 'batch' : 'batch % uniforms.aShape[0]', ";\n let batchB = ").concat(splitK || !broadcastBatch ? 'batch' : 'batch % uniforms.bShape[0]', ";\n let numTiles = ").concat(splitK ? "".concat(Math.ceil(splitedDimInner / tileInner)) : "(uniforms.dimInner - 1) / ".concat(tileInner, " + 1"), ";\n var kStart = ").concat(splitK ? "i32(globalId.z) * ".concat(splitedDimInner) : '0', ";\n\n var acc : array, ").concat(rowPerThread, ">;\n\n // Without this initialization strange values show up in acc.\n for (var innerRow = 0; innerRow < ").concat(rowPerThread, "; innerRow++) {\n for (var innerCol = 0; innerCol < ").concat(colPerThread, "; innerCol++) {\n acc[innerRow][innerCol] = 0.0;\n }\n }\n ").concat(matmulSnippet, "\n }\n "); } var readVectorASnippet = function (transpose) { return transpose ? "\n mm_readA(batchA, colA, globalRow),\n mm_readA(batchA, colA + 1, globalRow),\n mm_readA(batchA, colA + 2, globalRow),\n mm_readA(batchA, colA + 3, globalRow)\n " : "\n mm_readA(batchA, globalRow, colA),\n mm_readA(batchA, globalRow, colA + 1),\n mm_readA(batchA, globalRow, colA + 2),\n mm_readA(batchA, globalRow, colA + 3)\n "; }; function makeVectorMatrixProductSource(workgroupSize, transposeA) { if (transposeA === void 0) { transposeA = false; } tf.util.assert(workgroupSize[1] === 1 && workgroupSize[2] === 1, function () { return "A linear work group size is required. But got ".concat(workgroupSize, "."); }); var tileSize = workgroupSize[0] * 4; return "\n var mm_Asub : array, ".concat(workgroupSize[0], ">;\n\n ").concat(getMainHeaderString(), " {\n let tileCol = i32(localId.x);\n let globalCol = i32(globalId.x);\n let globalRow = i32(globalId.y);\n\n let numTiles = (uniforms.dimInner - 1) / ").concat(tileSize, " + 1;\n let batch = i32(globalId.z);\n let batchA = batch % uniforms.aShape[0];\n let batchB = batch % uniforms.bShape[0];\n // Without this initialization strange values show up in acc.\n var acc = 0.0;\n\n // Loop over shared dimension.\n for (var t = 0; t < numTiles; t++) {\n // Load one tile of A into local memory.\n let colA = t * ").concat(tileSize, " + tileCol * 4;\n mm_Asub[tileCol] = vec4(").concat(readVectorASnippet(transposeA), ");\n workgroupBarrier();\n\n // Compute acc values for a single thread.\n for (var k = 0; k < ").concat(tileSize / 4, "; k++) {\n let rowB = t * ").concat(tileSize, " + k * 4;\n let BCached = vec4(mm_readB(batchB, rowB, globalCol),\n mm_readB(batchB, rowB + 1, globalCol),\n mm_readB(batchB, rowB + 2, globalCol),\n mm_readB(batchB, rowB + 3, globalCol));\n\n let ACached = mm_Asub[k];\n acc = acc + dot(ACached, BCached);\n }\n\n workgroupBarrier();\n }\n\n mm_write(batch, globalRow, globalCol, acc);\n }\n "); } var MatMulPackedProgram = /** @class */ (function () { function MatMulPackedProgram(aShape, outputShape, transposeA, transposeB, bias, activation, preluActivationWeights, sequentialAccessByThreads) { var _a; if (transposeA === void 0) { transposeA = false; } if (transposeB === void 0) { transposeB = false; } if (bias === void 0) { bias = null; } if (activation === void 0) { activation = null; } if (preluActivationWeights === void 0) { preluActivationWeights = null; } if (sequentialAccessByThreads === void 0) { sequentialAccessByThreads = false; } this.variableNames = ['A', 'B']; this.uniforms = "dimAOuter : i32, dimBOuter : i32, dimInner : i32,"; this.outputShape = outputShape; this.dispatchLayout = { x: [2], y: [1], z: [0] }; var dimInner = transposeA ? aShape[1] : aShape[2]; this.isVec4 = ((dimInner % 4 === 0 && !transposeA) || (outputShape[1] % 4 === 0 && transposeA)) && outputShape[2] % 4 === 0 && !transposeB; this.outputComponent = this.isVec4 ? 4 : 1; this.isVectorA = outputShape[1] === 1 && !transposeA; if (!this.isVec4 && this.isVectorA) { // For makeVectorMatrixProductSource this.elementsPerThread = [1, 1, 1]; this.workgroupSize = [32, 1, 1]; } else { var workgroupInfo = computeWorkgroupInfoForMatMul(outputShape[1], dimInner, outputShape[2], transposeA); this.workgroupSize = workgroupInfo.workgroupSize; this.elementsPerThread = workgroupInfo.elementsPerThread; } this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize, this.elementsPerThread); var addBias = bias != null; var hasPreluActivationWeights = preluActivationWeights != null; if (addBias) { this.variableNames.push('bias'); } if (hasPreluActivationWeights) { this.variableNames.push('preluActivationWeights'); } this.sequentialAccessByThreads = sequentialAccessByThreads; this.transposeA = transposeA; this.transposeB = transposeB; this.addBias = addBias; this.activation = activation; this.hasPreluActivationWeights = hasPreluActivationWeights; _a = __read(this.getShapeFit(outputShape[1], outputShape[2], dimInner), 3), this.fitAOuter = _a[0], this.fitBOuter = _a[1], this.fitInner = _a[2]; this.shaderKey = "matMulPacked_".concat(this.elementsPerThread, "_").concat(transposeA, "_").concat(transposeB, "_").concat(this.activation, "_").concat(this.fitAOuter, "_").concat(this.fitBOuter, "_").concat(this.fitInner, "_").concat(this.isVec4, "_").concat(this.isVectorA, "_").concat(this.sequentialAccessByThreads); } MatMulPackedProgram.prototype.getShapeFit = function (dimAOuter, dimBOuter, dimInner) { var tileAOuter = this.workgroupSize[1] * this.elementsPerThread[1]; var tileBOuter = this.workgroupSize[0] * this.elementsPerThread[0]; if (!this.isVec4 && this.isVectorA) { // For makeVectorMatrixProductSource this.tileInner = this.workgroupSize[0] * 4; } else { this.tileInner = tileBOuter; } var fitAOuter = dimAOuter % tileAOuter === 0; var fitBOuter = dimBOuter % tileBOuter === 0; var fitInner = dimInner % this.tileInner === 0; return [fitAOuter, fitBOuter, fitInner]; }; MatMulPackedProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(activationFnSnippet(this.activation, this.hasPreluActivationWeights, this.isVec4), "\n ").concat(matMulReadWriteFnSource(this.addBias, this.activation, false /* transposeA is implemented in makeMatMulPackedSource */, this.transposeB, this.fitAOuter, this.fitBOuter, this.fitInner, this.isVec4 ? 4 : 1), "\n ").concat(this.isVec4 ? makeMatMulPackedVec4Source(this.elementsPerThread, this.workgroupSize, this.transposeA, this.tileInner, false, null, true) : (this.isVectorA ? makeVectorMatrixProductSource(this.workgroupSize, this.transposeA) : makeMatMulPackedSource(this.elementsPerThread, this.workgroupSize, this.transposeA, this.tileInner, false, null, this.sequentialAccessByThreads, true)), "\n "); return userCode; }; return MatMulPackedProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function makeMatMulReduceSource(workgroupSizeX) { return "\n var sumValues : array;\n ").concat(getMainHeaderString(), " {\n let coords = getOutputCoords();\n let batch = coords[0];\n let batchA = batch % uniforms.aShape[0];\n let batchB = batch % uniforms.bShape[0];\n let row = coords[1];\n let col = coords[2];\n var sum = 0.0;\n let Length = uniforms.dimInner;\n for (var k = i32(localId.x); k < Length; k = k + ").concat(workgroupSizeX, ") {\n let dataA = mm_readA(batchA, row, k);\n let dataB = mm_readB(batchB, k, col);\n sum = sum + dataA * dataB;\n }\n sumValues[localId.x] = sum;\n workgroupBarrier();\n\n for(var currentSize = ").concat(workgroupSizeX / 2, "u; currentSize > 1u;\n currentSize = currentSize / 2u) {\n if (localId.x < currentSize)\n {\n sumValues[localId.x] = sumValues[localId.x] + sumValues[localId.x + currentSize];\n }\n workgroupBarrier();\n }\n\n if (localId.x == 0u) {\n sum = sumValues[0] + sumValues[1];\n mm_write(batch, row, col, sum);\n }\n }\n "); } var MatMulReduceProgram = /** @class */ (function () { function MatMulReduceProgram(outputShape, transposeA, transposeB, bias, activation, preluActivationWeights) { if (transposeA === void 0) { transposeA = false; } if (transposeB === void 0) { transposeB = false; } if (bias === void 0) { bias = null; } if (activation === void 0) { activation = null; } if (preluActivationWeights === void 0) { preluActivationWeights = null; } this.variableNames = ['A', 'B']; this.uniforms = "dimAOuter : i32, dimBOuter : i32, dimInner : i32,"; this.workgroupSize = [256, 1, 1]; this.outputShape = outputShape; this.dispatchLayout = { x: [], y: [1, 2], z: [0] }; this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); var addBias = bias != null; var hasPreluActivationWeights = preluActivationWeights != null; if (addBias) { this.variableNames.push('bias'); } if (hasPreluActivationWeights) { this.variableNames.push('preluActivationWeights'); } this.transposeA = transposeA; this.transposeB = transposeB; this.addBias = addBias; this.activation = activation; this.hasPreluActivationWeights = hasPreluActivationWeights; this.shaderKey = "matMulReduce_".concat(this.activation, "_").concat(transposeA, "_").concat(transposeB); } MatMulReduceProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(activationFnSnippet(this.activation, this.hasPreluActivationWeights), "\n ").concat(matMulReadWriteFnSource(this.addBias, this.activation, this.transposeA, this.transposeB), "\n ").concat(makeMatMulReduceSource(this.workgroupSize[0]), "\n "); return userCode; }; return MatMulReduceProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function makeMatMulSmallOutputSizeSource(workgroupSize) { var tileAOuter = workgroupSize[1]; var tileBOuter = workgroupSize[0]; var tileInner = tileAOuter > tileBOuter ? tileAOuter : tileBOuter; return "\n var mm_Asub : array, ").concat(tileAOuter, ">;\n var mm_Bsub : array, ").concat(tileInner, ">;\n\n // If the output size is small for matrix multiplication, avoid to use vec4\n // and handle some elements per thread to optimally utilize the ALU.\n // Read data from global memory to registers firstly, then store them into\n // shared memory, so it is instruction-Level parallelism for arithmetic\n // operations and others handle IO operations between barrier api, makes ALU\n // and load/store units work simultaneously, could improves the performance.\n ").concat(getMainHeaderString(), " {\n let tileRow = i32(localId.y);\n let tileCol = i32(localId.x);\n let globalRow = i32(globalId.y);\n let globalCol = i32(globalId.x);\n let batch = i32(globalId.z);\n let batchA = batch % uniforms.aShape[0];\n let batchB = batch % uniforms.bShape[0];\n\n // uniforms.dimInner should be greater than 0.\n let numTiles = (uniforms.dimInner - 1) / ").concat(tileInner, " + 1;\n var acc = 0.0;\n\n var globalColA = tileCol;\n var globalRowB = 0;\n var regA = mm_readA(batchA, globalRow, globalColA);\n var regB0 = mm_readB(batchB, globalRowB + 2 * tileRow, globalCol);\n var regB1 = mm_readB(batchB, globalRowB + 2 * tileRow + 1, globalCol);\n globalColA = globalColA + ").concat(tileInner, ";\n globalRowB = globalRowB + ").concat(tileInner, ";\n\n for (var t = 0; t < numTiles; t = t + 1) {\n mm_Asub[tileRow][tileCol] = regA;\n mm_Bsub[2 * tileRow][tileCol] = regB0;\n mm_Bsub[2 * tileRow + 1][tileCol] = regB1;\n\n workgroupBarrier();\n\n regA = mm_readA(batchA, globalRow, globalColA);\n regB0 = mm_readB(batchB, globalRowB + 2 * tileRow, globalCol);\n regB1 = mm_readB(batchB, globalRowB + 2 * tileRow + 1, globalCol);\n globalColA = globalColA + ").concat(tileInner, ";\n globalRowB = globalRowB + ").concat(tileInner, ";\n\n for (var k = 0; k < ").concat(tileInner, "; k = k + 1) {\n acc = acc + mm_Asub[tileRow][k] * mm_Bsub[k][tileCol];\n }\n workgroupBarrier();\n }\n\n mm_write(batch, globalRow, globalCol, acc);\n }\n "); } var MatMulSmallOutputSizeProgram = /** @class */ (function () { function MatMulSmallOutputSizeProgram(aShape, bShape, outputShape, transposeA, transposeB, bias, activation, preluActivationWeights) { if (transposeA === void 0) { transposeA = false; } if (transposeB === void 0) { transposeB = false; } if (bias === void 0) { bias = null; } if (activation === void 0) { activation = null; } if (preluActivationWeights === void 0) { preluActivationWeights = null; } this.variableNames = ['A', 'B']; this.uniforms = "dimAOuter : i32, dimBOuter : i32, dimInner : i32,"; this.workgroupSize = [16, 8, 1]; this.outputShape = outputShape; this.dispatchLayout = { x: [2], y: [1], z: [0] }; this.dispatch = [ Math.ceil(outputShape[2] / this.workgroupSize[0]), Math.ceil(outputShape[1] / this.workgroupSize[1]), outputShape[0] ]; var addBias = bias != null; if (addBias) { this.variableNames.push('bias'); } var hasPreluActivationWeights = preluActivationWeights != null; if (hasPreluActivationWeights) { this.variableNames.push('preluActivationWeights'); } this.transposeA = transposeA; this.transposeB = transposeB; this.addBias = addBias; this.activation = activation; this.hasPreluActivationWeights = hasPreluActivationWeights; this.shaderKey = "matMulSmallOutputSize_".concat(this.activation, "_").concat(transposeA, "_").concat(transposeB); } MatMulSmallOutputSizeProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(activationFnSnippet(this.activation, this.hasPreluActivationWeights), "\n ").concat(matMulReadWriteFnSource(this.addBias, this.activation, this.transposeA, this.transposeB), "\n ").concat(makeMatMulSmallOutputSizeSource(this.workgroupSize), "\n "); return userCode; }; return MatMulSmallOutputSizeProgram; }()); /** * @license * Copyright 2022 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var MatMulSplitKProgram = /** @class */ (function () { function MatMulSplitKProgram(outputShape, dimInner, transposeA, transposeB) { if (transposeA === void 0) { transposeA = false; } if (transposeB === void 0) { transposeB = false; } this.variableNames = ['A', 'B']; this.uniforms = "dimAOuter : i32, dimBOuter : i32, dimInner : i32,"; this.workgroupSize = [8, 8, 1]; this.atomic = true; this.splitedDimInner = 128; tf.util.assert(outputShape[0] === 1, function () { return 'MatMulSplitKProgram only supports batch = 1.'; }); this.outputShape = outputShape; this.dispatchLayout = { x: [2], y: [1], z: [0, 3] }; var isVec4 = (transposeA && this.outputShape[1] % 4 === 0 || !transposeA && dimInner % 4 === 0) && this.outputShape[2] % 4 === 0; this.elementsPerThread = [4, 4, this.splitedDimInner]; this.outputComponent = isVec4 ? 4 : 1; if (!isVec4) { if (this.outputShape[1] < 16) { this.elementsPerThread[1] = 1; } if (this.outputShape[2] < 16) { this.elementsPerThread[0] = 1; } } this.dispatch = computeDispatch(this.dispatchLayout, [ this.outputShape[0], this.outputShape[1], this.outputShape[2], dimInner ], this.workgroupSize, this.elementsPerThread); this.transposeA = transposeA; this.transposeB = transposeB; this.shaderKey = "matMulSplitK_".concat(transposeA, "_").concat(transposeB, "_").concat(this.elementsPerThread, "_").concat(this.outputComponent); } MatMulSplitKProgram.prototype.getUserCode = function () { var component = this.outputComponent; var userCode = "\n ".concat(matMulReadFnSource(false, this.transposeB, false, false, false, component), "\n fn mm_write(batch: i32, row : i32, col : i32, value : ").concat(typeSnippet(component), ") {\n if (row < uniforms.dimAOuter && col < uniforms.dimBOuter) {\n let coords = vec3(batch, row, col);\n let flatIndex = getOutputIndexFromCoords(coords);\n // The problem is that we should initialize output to zero before using.\n // Otherwise, the original value will be added to the result.\n for (var i = 0; i < ").concat(component, "; i = i + 1) {\n ").concat(atomicAddSnippet('&result[flatIndex + i]', "".concat(component > 1 ? 'value[i]' : 'value'), 'float32'), "\n }\n }\n }\n ").concat(component === 4 ? makeMatMulPackedVec4Source(this.elementsPerThread, this.workgroupSize, this.transposeA, 32, true, this.splitedDimInner) : makeMatMulPackedSource(this.elementsPerThread, this.workgroupSize, this.transposeA, 32, true, this.splitedDimInner), "\n "); return userCode; }; return MatMulSplitKProgram; }()); var BiasActivationProgram = /** @class */ (function () { function BiasActivationProgram(outputShape, bias, activation, preluActivationWeights) { if (bias === void 0) { bias = null; } if (activation === void 0) { activation = null; } if (preluActivationWeights === void 0) { preluActivationWeights = null; } this.uniforms = ''; this.variableNames = ['x']; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = outputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.addBias = bias != null; this.hasPreluActivationWeights = preluActivationWeights != null; this.activation = activation; if (this.addBias) { this.variableNames.push('bias'); } if (this.hasPreluActivationWeights) { this.variableNames.push('preluActivationWeights'); } this.shaderKey = "biasActivation_".concat(activation); } BiasActivationProgram.prototype.getUserCode = function () { return "\n ".concat(activationFnSnippet(this.activation, this.hasPreluActivationWeights), "\n ").concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n var value = getXByOutputIndex(index);\n ").concat(biasActivationSnippet(this.addBias, this.activation), "\n setOutputAtIndex(index, value);\n }\n }\n "); }; return BiasActivationProgram; }()); /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var FillProgram = /** @class */ (function () { function FillProgram(shape) { this.variableNames = []; this.outputShape = []; this.uniforms = 'value : f32,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = shape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'fill'; } FillProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n setOutputAtIndex(index, uniforms.value);\n }\n }\n "); return userCode; }; return FillProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function fill(args) { var backend = args.backend, attrs = args.attrs; var shape = attrs.shape, value = attrs.value; var dtype = attrs.dtype; dtype = dtype || tf.util.inferDtype(value); if (dtype === 'string') { // String type should be handled in CPU memory. var values = tf.util.getArrayFromDType(dtype, tf.util.sizeFromShape(shape)); values.fill(value); return backend.makeTensorInfo(shape, dtype, values); } else { var program = new FillProgram(shape); var uniformData = [{ type: 'float32', data: [value] }]; return backend.runWebGPUProgram(program, [], dtype, uniformData); } } var fillConfig = { kernelName: tf.Fill, backendName: 'webgpu', kernelFunc: fill }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function reshape(args) { var inputs = args.inputs, attrs = args.attrs; var x = inputs.x; var shape = attrs.shape; var xSize = tf.util.sizeFromShape(x.shape); var $shape = tf.util.inferFromImplicitShape(shape, xSize); var $xSize = tf.util.sizeFromShape($shape); tf.util.assert(xSize === $xSize, function () { return "The new shape (".concat($shape, ") has ").concat($xSize, " elements and the old ") + "shape (".concat(x.shape, ") has ").concat(xSize, " elements. The new shape and old ") + "shape must have the same number of elements."; }); // Backend needs to track refCount for the dataId for reshape op args.backend.incRef(x.dataId); return { dataId: x.dataId, shape: $shape, dtype: x.dtype }; } var reshapeConfig = { kernelName: tf.Reshape, backendName: 'webgpu', kernelFunc: reshape }; function batchMatMulImpl(_a) { var e_1, _b, e_2, _c; var a = _a.a, b = _a.b, transposeA = _a.transposeA, transposeB = _a.transposeB, backend = _a.backend, _d = _a.bias, bias = _d === void 0 ? null : _d, _e = _a.preluActivationWeights, preluActivationWeights = _e === void 0 ? null : _e, _f = _a.leakyreluAlpha, leakyreluAlpha = _f === void 0 ? 0 : _f, _g = _a.activation, activation = _g === void 0 ? null : _g; var aRank = a.shape.length; var bRank = b.shape.length; var innerShapeA = transposeA ? a.shape[aRank - 2] : a.shape[aRank - 1]; var innerShapeB = transposeB ? b.shape[bRank - 1] : b.shape[bRank - 2]; var outerShapeA = transposeA ? a.shape[aRank - 1] : a.shape[aRank - 2]; var outerShapeB = transposeB ? b.shape[bRank - 2] : b.shape[bRank - 1]; var outerDimsA = a.shape.slice(0, -2); var outerDimsB = b.shape.slice(0, -2); var batchDimA = tf.util.sizeFromShape(outerDimsA); var batchDimB = tf.util.sizeFromShape(outerDimsB); var outShapeOuterDims = tf.broadcast_util.assertAndGetBroadcastShape(a.shape.slice(0, -2), b.shape.slice(0, -2)); var outShape = outShapeOuterDims.concat([outerShapeA, outerShapeB]); tf.util.assert(innerShapeA === innerShapeB, function () { return "Error in matMul: inner shapes (".concat(innerShapeA, ") and (") + "".concat(innerShapeB, ") of Tensors with shapes ").concat(a.shape, " and ") + "".concat(b.shape, " and transposeA=").concat(transposeA) + " and transposeB=".concat(transposeB, " must match."); }); var a3dShape = transposeA ? [batchDimA, innerShapeA, outerShapeA] : [batchDimA, outerShapeA, innerShapeA]; var b3dShape = transposeB ? [batchDimB, outerShapeB, innerShapeB] : [batchDimB, innerShapeB, outerShapeB]; // The rest of the implementation is designed to operate on rank-3 tensors var a3d = reshape({ inputs: { x: a }, backend: backend, attrs: { shape: a3dShape } }); var b3d = reshape({ inputs: { x: b }, backend: backend, attrs: { shape: b3dShape } }); var intermediates = [a3d, b3d]; var batchDim = Math.max(batchDimA, batchDimB); var inputs = [a3d, b3d]; var dimensions = [ { type: 'int32', data: [outerShapeA] }, { type: 'int32', data: [outerShapeB] }, { type: 'int32', data: [innerShapeA] } ]; var program; var out; var outputShape = [batchDim, outerShapeA, outerShapeB]; var matmulProgramType = tf.env().get('WEBGPU_MATMUL_PROGRAM_TYPE'); if (matmulProgramType < 0) { // Usually increasing workgroups is a good way to gain more performance for // few workgroups by tiling 32x32 (default matmul algorithm). Currently, // there are three ways to increase workgroups. 1) MatMulReduceProgram, // which is used only when the output size is very small (128 for now). 2) // MatMulSplitKProgram, increasing workgroups by spliting K. 3) // MatMulSmallOutputSizeProgram, increasing workgroups by small tile size. // For different devices, the minimum optimal workgroups may be different. // So here we set a |thresholdToIncreaseWorkgroups| to indicate whether we // need to increase workgroups. And the literal number is an empirical // value. var thresholdFlagValue = tf.env().getNumber('WEBGPU_THRESHOLD_TO_INCREASE_WORKGROUPS_FOR_MATMUL'); var thresholdToIncreaseWorkgroups = thresholdFlagValue > 0 ? thresholdFlagValue : backend.thresholdToIncreaseWorkgroups; var workgroupsBy32x32 = batchDim * Math.ceil(outerShapeA / 32) * Math.ceil(outerShapeB / 32); var hasFewWorkgroups = workgroupsBy32x32 <= thresholdToIncreaseWorkgroups || (outerShapeA <= 8 && workgroupsBy32x32 <= thresholdToIncreaseWorkgroups * 2); if (hasFewWorkgroups) { if (batchDim * outerShapeA * outerShapeB <= 128) { matmulProgramType = MatMulProgramType.MatMulReduceProgram; } else if (batchDim === 1 && innerShapeB >= 2000) { matmulProgramType = MatMulProgramType.MatMulSplitKProgram; } else { matmulProgramType = MatMulProgramType.MatMulSmallOutputSizeProgram; } } else { matmulProgramType = MatMulProgramType.MatMulPackedProgram; } } switch (matmulProgramType) { case MatMulProgramType.MatMulReduceProgram: program = new MatMulReduceProgram(outputShape, transposeA, transposeB, bias, activation, preluActivationWeights); break; case MatMulProgramType.MatMulSplitKProgram: { // The output buffer must be initailzed to zero before using since we // use atomicAdd in MatMulSplitKProgram. out = fill({ backend: backend, attrs: { shape: outputShape, value: 0, dtype: a.dtype } }); program = new MatMulSplitKProgram(outputShape, innerShapeB, transposeA, transposeB); if (bias || activation) { out = backend.runWebGPUProgram(program, inputs, a.dtype, dimensions, out); var biasActivationProgram = new BiasActivationProgram(out.shape, bias, activation, preluActivationWeights); var uniformData = null; var activationInputs = [out]; if (bias) { activationInputs.push(bias); } if (preluActivationWeights) { activationInputs.push(preluActivationWeights); } if (activation === 'leakyrelu') { uniformData = [{ type: 'float32', data: [leakyreluAlpha] }]; biasActivationProgram.uniforms += ' alpha : f32,'; } var outActivated = backend.runWebGPUProgram(biasActivationProgram, activationInputs, out.dtype, uniformData); intermediates.push(out); var outReshaped_1 = reshape({ inputs: { x: outActivated }, backend: backend, attrs: { shape: outShape } }); intermediates.push(outActivated); try { for (var intermediates_1 = __values(intermediates), intermediates_1_1 = intermediates_1.next(); !intermediates_1_1.done; intermediates_1_1 = intermediates_1.next()) { var i = intermediates_1_1.value; backend.disposeData(i.dataId); } } catch (e_1_1) { e_1 = { error: e_1_1 }; } finally { try { if (intermediates_1_1 && !intermediates_1_1.done && (_b = intermediates_1.return)) _b.call(intermediates_1); } finally { if (e_1) throw e_1.error; } } return outReshaped_1; } break; } case MatMulProgramType.MatMulSmallOutputSizeProgram: program = new MatMulSmallOutputSizeProgram(a3dShape, b3dShape, outputShape, transposeA, transposeB, bias, activation, preluActivationWeights); break; case MatMulProgramType.MatMulPackedProgram: // Experiments show that sequential access is more friendly for Intel // GPUs. var sequentialAccessByThreads = backend.adapterInfo.isIntel(); program = new MatMulPackedProgram(a3dShape, outputShape, transposeA, transposeB, bias, activation, preluActivationWeights, sequentialAccessByThreads); break; default: throw new Error("Unsupported MatMulProgramType ".concat(matmulProgramType, ".")); } if (bias) { inputs.push(bias); } if (preluActivationWeights) { inputs.push(preluActivationWeights); } if (activation === 'leakyrelu') { dimensions.push({ type: 'float32', data: [leakyreluAlpha] }); program.uniforms += ' alpha : f32,'; } out = backend.runWebGPUProgram(program, inputs, a.dtype, dimensions, out); var outReshaped = reshape({ inputs: { x: out }, backend: backend, attrs: { shape: outShape } }); intermediates.push(out); try { for (var intermediates_2 = __values(intermediates), intermediates_2_1 = intermediates_2.next(); !intermediates_2_1.done; intermediates_2_1 = intermediates_2.next()) { var i = intermediates_2_1.value; backend.disposeData(i.dataId); } } catch (e_2_1) { e_2 = { error: e_2_1 }; } finally { try { if (intermediates_2_1 && !intermediates_2_1.done && (_c = intermediates_2.return)) _c.call(intermediates_2); } finally { if (e_2) throw e_2.error; } } return outReshaped; } /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function _fusedMatMul(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var a = inputs.a, b = inputs.b, bias = inputs.bias, preluActivationWeights = inputs.preluActivationWeights; var transposeA = attrs.transposeA, transposeB = attrs.transposeB, activation = attrs.activation, leakyreluAlpha = attrs.leakyreluAlpha; return batchMatMulImpl({ a: a, b: b, transposeA: transposeA, transposeB: transposeB, backend: backend, bias: bias, preluActivationWeights: preluActivationWeights, leakyreluAlpha: leakyreluAlpha, activation: activation }); } var _fusedMatMulConfig = { kernelName: tf._FusedMatMul, backendName: 'webgpu', kernelFunc: _fusedMatMul, }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var BinaryOpComplexProgram = /** @class */ (function () { function BinaryOpComplexProgram(op, aShape, bShape) { this.variableNames = ['AReal', 'AImag', 'BReal', 'BImag']; this.workgroupSize = [128, 1, 1]; this.size = true; this.outputShape = tf.backend_util.assertAndGetBroadcastShape(aShape, bShape); this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = "binaryOpComplex_".concat(op); this.op = op; } BinaryOpComplexProgram.prototype.getUserCode = function () { var opStr = getBinaryOpString(this.op, false); var userCode = "\n fn binaryOpComplex(\n areal : f32, aimag : f32, breal : f32, bimag : f32) -> f32 {\n ".concat(opStr, "\n }\n\n ").concat(getMainHeaderString('index'), " {\n if(index < uniforms.size) {\n let areal = getARealByOutputIndex(index);\n let aimag = getAImagByOutputIndex(index);\n let breal = getBRealByOutputIndex(index);\n let bimag = getBImagByOutputIndex(index);\n setOutputAtIndex(index, binaryOpComplex(areal, aimag, breal, bimag));\n }\n }\n "); return userCode; }; return BinaryOpComplexProgram; }()); /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var BinaryOpProgram = /** @class */ (function () { function BinaryOpProgram(op, aShape, bShape) { this.size = true; this.variableNames = ['A', 'B']; this.outputShape = tf.backend_util.assertAndGetBroadcastShape(aShape, bShape); this.dispatchLayout = flatDispatchLayout(this.outputShape); this.op = op; this.useSharedMemoryWithA = aShape.length <= 1 && bShape.length > 1 && aShape[0] < 128; this.useSharedMemoryWithB = bShape.length <= 1 && aShape.length > 1 && bShape[0] < 128; if (this.useSharedMemoryWithA || this.useSharedMemoryWithB) { this.outputComponent = 1; this.variableComponents = [1, 1]; // lastDimensionSize is used as sharedBuf array size, so can not be // used as uniform. this.lastDimensionSize = this.useSharedMemoryWithB ? bShape[0] : aShape[0]; this.shaderKey = "binary_".concat(op, "_").concat(this.lastDimensionSize); this.type = 'shared'; // This is an experimental value when using shared memory. // Note that the maximum of workgroup X dimension is 256. this.workgroupSize = [256, 1, 1]; } else { var aDivisibleBy4 = aShape.length > 0 && aShape[aShape.length - 1] % 4 === 0; var bDivisibleBy4 = bShape.length > 0 && bShape[bShape.length - 1] % 4 === 0; if (aDivisibleBy4 && bDivisibleBy4) { this.outputComponent = 4; this.variableComponents = [4, 4]; } else if ((aDivisibleBy4 && (tf.util.isScalarShape(bShape) || bShape[bShape.length - 1] === 1)) || (bDivisibleBy4 && (tf.util.isScalarShape(aShape) || aShape[aShape.length - 1] === 1))) { this.outputComponent = 4; this.variableComponents = aDivisibleBy4 ? [4, 1] : [1, 4]; } else { this.outputComponent = 1; this.variableComponents = [1, 1]; } this.type = 'nonshared'; this.shaderKey = "binary_".concat(op, "_").concat(this.variableComponents); // TODO(jiajia.qin@intel.com): Heuristically select a good work group // size. this.workgroupSize = [128, 1, 1]; } this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize, [this.outputComponent, 1, 1]); } BinaryOpProgram.prototype.getUserCode = function () { var userCode; var dType = this.outputComponent === 4 ? 'vec4' : 'f32'; var opFnStr = "\n fn binaryOperation(a : ".concat(dType, ", b : ").concat(dType, ") -> ").concat(dType, " {\n ").concat(getBinaryOpString(this.op, this.outputComponent === 4), "\n };\n "); if (this.type === 'shared') { var sharedIndexSnippet = this.lastDimensionSize > 1 ? "coords[".concat(this.outputShape.length - 1, "]") : '0'; var accessDataSnippet = this.useSharedMemoryWithB ? "let a = getAByOutputIndex(index);\n let b = sharedBuf[".concat(sharedIndexSnippet, "];") : "let a = sharedBuf[".concat(sharedIndexSnippet, "];\n let b = getBByOutputIndex(index);"); userCode = "\n ".concat(opFnStr, "\n var sharedBuf : array;\n ").concat(getMainHeaderString('index'), " {\n // Fill in the shared memory buffer.\n let localIndex = i32(localId.x);\n if(localIndex < ").concat(this.lastDimensionSize, ") {\n sharedBuf[localIndex] = f32(").concat(this.useSharedMemoryWithB ? 'B' : 'A', "[localIndex]);\n }\n workgroupBarrier();\n\n if(index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n ").concat(accessDataSnippet, "\n setOutputAtIndex(index, binaryOperation(a, b));\n }\n }\n "); } else { userCode = "\n ".concat(opFnStr, "\n ").concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index * ").concat(this.outputComponent, ");\n let a = ").concat(dType, "(getAByOutputCoords(coords));\n let b = ").concat(dType, "(getBByOutputCoords(coords));\n setOutputAtIndex(index, binaryOperation(a, b));\n }\n }\n "); } return userCode; }; return BinaryOpProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function identity(args) { var inputs = args.inputs; var x = inputs.x; args.backend.incRef(x.dataId); return { dataId: x.dataId, shape: x.shape, dtype: x.dtype }; } var identityConfig = { kernelName: tf.Identity, backendName: 'webgpu', kernelFunc: identity }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ /** * Complex tensors share data with their real and imaginary components. Complex * tensors' reference to the components is tracked by refCount on the individual * component. The refCounts are increased by the identity call. * * When a complex tensor is disposed, it will reduce the refCount on the * components by calling disposeData on each. */ function complex(args) { var inputs = args.inputs, backend = args.backend; var real = inputs.real, imag = inputs.imag; var complexInfo = backend.makeTensorInfo(real.shape, 'complex64'); var complex = backend.tensorMap.get(complexInfo.dataId); var realTensorInfo = identity({ inputs: { x: real }, backend: backend }); var imagTensorInfo = identity({ inputs: { x: imag }, backend: backend }); complex.complexTensorInfos = { real: realTensorInfo, imag: imagTensorInfo }; return complexInfo; } var complexConfig = { kernelName: tf.Complex, backendName: 'webgpu', kernelFunc: complex }; /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var UnaryOpProgram = /** @class */ (function () { function UnaryOpProgram(outputShape, op, uniforms) { if (uniforms === void 0) { uniforms = ''; } this.variableNames = ['A']; this.size = true; // TODO(jiajia.qin@intel.com): Heuristically select a good work group size. var workgroupSizeX = 128; this.workgroupSize = [workgroupSizeX, 1, 1]; this.outputShape = outputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.op = op; if (uniforms !== '') { this.uniforms = uniforms; } this.shaderKey = "unary_".concat(op); } UnaryOpProgram.prototype.getUserCode = function () { return "\n fn unaryOperation(a : f32) -> f32 {\n ".concat(getUnaryOpString(this.op, false), "\n }\n ").concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let a = getAByOutputIndex(index);\n setOutputAtIndex(index, unaryOperation(a));\n }\n }\n "); }; return UnaryOpProgram; }()); /** * Template that creates a `KernelFunc` for unary ops. * @param opType Op type to create `UnaryOpProgram`. * @param cpuKernelImpl Optional. Shared functionality from tfjs-backend-cpu, it * will be involved when necessary. * @param dtype Optional. If set, the result has this dtype. Otherwise, the * result has the same dtype as the first input. This is mainly used in * comparison kernels, such as Equal, Less, Greater, etc. */ function unaryKernelFunc(_a) { var opType = _a.opType, cpuKernelImpl = _a.cpuKernelImpl, dtype = _a.dtype; return function (_a) { var inputs = _a.inputs, backend = _a.backend; var x = inputs.x; var webgpuBackend = backend; var $dtype = dtype || x.dtype; if (webgpuBackend.shouldExecuteOnCPU([x]) && cpuKernelImpl != null) { var xData = webgpuBackend.tensorMap.get(x.dataId); var outValues = cpuKernelImpl(xData.values, $dtype); return webgpuBackend.makeTensorInfo(x.shape, $dtype, outValues); } var program = new UnaryOpProgram(x.shape, opType); return webgpuBackend.runWebGPUProgram(program, [x], $dtype); }; } /** * Template that creates a `KernelFunc` for binary ops. * @param opType Op type to create `BinaryOpProgram`. * @param cpuKernelImpl Optional. Shared functionality from tfjs-backend-cpu, it * will be involved when necessary. * @param dtype Optional. If set, the result has this dtype. Otherwise, the * result has the same dtype as the first input. This is mainly used in * comparison kernels, such as Equal, Less, Greater, etc. */ function binaryKernelFunc(_a) { var opType = _a.opType, cpuKernelImpl = _a.cpuKernelImpl, _b = _a.supportsComplex, supportsComplex = _b === void 0 ? false : _b, dtype = _a.dtype; return function (_a) { var _b; var inputs = _a.inputs, backend = _a.backend; var a = inputs.a, b = inputs.b; var webgpuBackend = backend; if (supportsComplex && a.dtype === 'complex64') { var aData = webgpuBackend.tensorMap.get(a.dataId); var bData = webgpuBackend.tensorMap.get(b.dataId); var real = void 0, imag = void 0; if (opType !== BinaryOpType.MUL) { _b = __read([ [aData.complexTensorInfos.real, bData.complexTensorInfos.real], [aData.complexTensorInfos.imag, bData.complexTensorInfos.imag] ].map(function (complexParts) { var _a = __read(complexParts, 2), aPart = _a[0], bPart = _a[1]; var aHandle = { dataId: aPart.dataId, dtype: aPart.dtype, shape: a.shape }; var bHandle = { dataId: bPart.dataId, dtype: bPart.dtype, shape: b.shape }; var program = new BinaryOpProgram(opType, a.shape, b.shape); return webgpuBackend.runWebGPUProgram(program, [aHandle, bHandle], tf.upcastType(aPart.dtype, bPart.dtype)); }), 2), real = _b[0], imag = _b[1]; } else { var realProgram = new BinaryOpComplexProgram(BinaryOpType.COMPLEX_MULTIPLY_REAL, a.shape, b.shape); var imagProgram = new BinaryOpComplexProgram(BinaryOpType.COMPLEX_MULTIPLY_IMAG, a.shape, b.shape); var inputs_1 = [ { dataId: aData.complexTensorInfos.real.dataId, dtype: aData.complexTensorInfos.real.dtype, shape: a.shape }, { dataId: aData.complexTensorInfos.imag.dataId, dtype: aData.complexTensorInfos.imag.dtype, shape: a.shape }, { dataId: bData.complexTensorInfos.real.dataId, dtype: bData.complexTensorInfos.real.dtype, shape: b.shape }, { dataId: bData.complexTensorInfos.imag.dataId, dtype: bData.complexTensorInfos.imag.dtype, shape: b.shape } ]; real = webgpuBackend.runWebGPUProgram(realProgram, inputs_1, 'float32'); imag = webgpuBackend.runWebGPUProgram(imagProgram, inputs_1, 'float32'); } var complexOutput = complex({ inputs: { real: real, imag: imag }, backend: webgpuBackend }); webgpuBackend.disposeData(real.dataId); webgpuBackend.disposeData(imag.dataId); // TODO: Implement CPU forwarding for complex inputs. return complexOutput; } var $dtype = dtype || tf.upcastType(a.dtype, b.dtype); if ((a.dtype === 'string' || b.dtype === 'string' || webgpuBackend.shouldExecuteOnCPU([a, b])) && cpuKernelImpl != null) { var aData = webgpuBackend.tensorMap.get(a.dataId).values; var bData = webgpuBackend.tensorMap.get(b.dataId).values; var decodedAVals = a.dtype === 'string' ? // tslint:disable-next-line: no-any tf.backend_util.fromUint8ToStringArray(aData) : aData; var decodedBVals = a.dtype === 'string' ? // tslint:disable-next-line: no-any tf.backend_util.fromUint8ToStringArray(bData) : bData; var _c = __read(cpuKernelImpl(a.shape, b.shape, decodedAVals, decodedBVals, $dtype), 2), outValues = _c[0], outShape = _c[1]; return webgpuBackend.makeTensorInfo(outShape, $dtype, outValues); } var program = new BinaryOpProgram(opType, a.shape, b.shape); return webgpuBackend.runWebGPUProgram(program, [a, b], $dtype); }; } /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the License); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an AS IS BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function simpleAbsImpl(vals) { var resultValues = new Float32Array(vals.length); for (var i = 0; i < vals.length; ++i) { resultValues[i] = Math.abs(vals[i]); } return resultValues; } /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ /** * Template that creates implementation for binary ops. Supports broadcast. */ function createSimpleBinaryKernelImpl(op) { return function (aShape, bShape, aVals, bVals, dtype) { var newShape = tf.backend_util.assertAndGetBroadcastShape(aShape, bShape); var resultRank = newShape.length; var resultStrides = tf.util.computeStrides(newShape); var resultSize = tf.util.sizeFromShape(newShape); var result = tf.util.getTypedArrayFromDType(dtype, resultSize); var aRank = aShape.length; var bRank = bShape.length; var aStrides = tf.util.computeStrides(aShape); var bStrides = tf.util.computeStrides(bShape); var aBroadcastDims = tf.backend_util.getBroadcastDims(aShape, newShape); var bBroadcastDims = tf.backend_util.getBroadcastDims(bShape, newShape); if (aBroadcastDims.length + bBroadcastDims.length === 0) { for (var i = 0; i < result.length; ++i) { result[i] = op(aVals[i % aVals.length], bVals[i % bVals.length]); } } else { var _loop_1 = function (i) { var loc = tf.util.indexToLoc(i, resultRank, resultStrides); var aLoc = loc.slice(-aRank); aBroadcastDims.forEach(function (d) { return aLoc[d] = 0; }); var aIndex = tf.util.locToIndex(aLoc, aRank, aStrides); var bLoc = loc.slice(-bRank); bBroadcastDims.forEach(function (d) { return bLoc[d] = 0; }); var bIndex = tf.util.locToIndex(bLoc, bRank, bStrides); result[i] = op(aVals[aIndex], bVals[bIndex]); }; for (var i = 0; i < result.length; ++i) { _loop_1(i); } } return [result, newShape]; }; } function castImpl(values, shape, inputType, dtype) { if (dtype === 'int32') { var resultValues = Int32Array.from(values); return [shape, 'int32', resultValues]; } if (dtype === 'bool') { // This is essentially the result of notEqual(x, 0). We avoid using // kernel notEqual to avoid circular dependency, i.e. binary_utils -> // cast -> notEqual -> binary_utils. var zero = tf.util.toTypedArray([0], inputType); var _a = __read(createSimpleBinaryKernelImpl(function (a, b) { return (a !== b) ? 1 : 0; })(shape, [], values, zero, 'bool'), 2), resultData = _a[0], resultShape = _a[1]; return [resultShape, 'bool', resultData]; } throw new Error("Error in Cast: failed to cast ".concat(inputType, " to ").concat(dtype)); } /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var addImpl = createSimpleBinaryKernelImpl((function (a, b) { return a + b; })); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ /** * Template that creates implementation for unary op. */ function createSimpleUnaryImpl(op) { return function (values, dtype, attrs) { var newValues = tf.util.getArrayFromDType(dtype, values.length); for (var i = 0; i < values.length; ++i) { newValues[i] = op(values[i], attrs); } return newValues; }; } /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the License); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an AS IS BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var ceilImpl = createSimpleUnaryImpl(function (xi) { return Math.ceil(xi); }); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function concatImpl$1(inputs, outShape, dtype, simplyConcat) { var outVals = tf.util.getArrayFromDType(dtype, tf.util.sizeFromShape(outShape)); if (simplyConcat && dtype !== 'string') { // Use built-in TypedArray.set() method for speed. var offset_1 = 0; inputs.forEach(function (input) { var size = tf.util.sizeFromShape(input.shape); outVals.set(input.vals, offset_1); offset_1 += size; }); } else { var colOffset_1 = 0; inputs.forEach(function (input) { var decodedData = dtype === 'string' ? tf.backend_util.fromUint8ToStringArray(input.vals) : input.vals; var tIdx = 0; for (var row = 0; row < input.shape[0]; ++row) { var resIdx = row * outShape[1] + colOffset_1; for (var col = 0; col < input.shape[1]; ++col) { outVals[resIdx + col] = decodedData[tIdx++]; } } colOffset_1 += input.shape[1]; }); } return outVals; } /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var equalImpl = createSimpleBinaryKernelImpl(function (a, b) { return (a === b) ? 1 : 0; }); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the License); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an AS IS BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var expImpl = createSimpleUnaryImpl(function (xi) { return Math.exp(xi); }); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the License); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an AS IS BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var expm1Impl = createSimpleUnaryImpl(function (xi) { return Math.expm1(xi); }); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the License); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an AS IS BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var floorImpl = createSimpleUnaryImpl(function (xi) { return Math.floor(xi); }); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var floorDivImpl = createSimpleBinaryKernelImpl(function (a, b) { return Math.floor(a / b); }); function gatherNdImpl(indicesData, paramsBuf, dtype, numSlices, sliceRank, sliceSize, strides, paramsShape, paramsSize) { var outBuf = tf.buffer([numSlices, sliceSize], dtype); for (var i = 0; i < numSlices; i++) { var index = []; var flattenIndex = 0; for (var j = 0; j < sliceRank; j++) { var dim = indicesData[i * sliceRank + j]; flattenIndex += dim * strides[j]; index.push(dim); } if (flattenIndex < 0 || flattenIndex >= paramsSize / sliceSize) { throw new Error("Invalid indices: ".concat(index, " does not index into ").concat(paramsShape)); } for (var k = 0; k < sliceSize; k++) { outBuf.values[i * sliceSize + k] = paramsBuf.get.apply(paramsBuf, __spreadArray([], __read(paramsBuf.indexToLoc(flattenIndex * sliceSize + k)), false)); } } return outBuf; } /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function gatherV2Impl(xBuf, indicesBuf, flattenOutputShape) { var outBuf = tf.buffer(flattenOutputShape, xBuf.dtype); for (var i = 0; i < outBuf.size; ++i) { var newLoc = outBuf.indexToLoc(i); var originalLoc = newLoc.slice(); var batchIdx = originalLoc[0]; var indicesIdx = originalLoc[2]; var indicesIndex = indicesBuf.locToIndex([batchIdx, indicesIdx]); originalLoc[2] = indicesBuf.values[indicesIndex]; var originalIndex = xBuf.locToIndex(originalLoc); if (0 <= originalIndex && originalIndex < xBuf.values.length) { outBuf.values[i] = xBuf.values[originalIndex]; } // Else, index is out of bounds, so leave the default zero val in outBuf. } return outBuf; } /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var greaterImpl = createSimpleBinaryKernelImpl(function (a, b) { return (a > b) ? 1 : 0; }); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var greaterEqualImpl = createSimpleBinaryKernelImpl(function (a, b) { return (a >= b) ? 1 : 0; }); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var lessImpl = createSimpleBinaryKernelImpl(function (a, b) { return (a < b) ? 1 : 0; }); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var lessEqualImpl = createSimpleBinaryKernelImpl(function (a, b) { return (a <= b) ? 1 : 0; }); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the License); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an AS IS BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var logImpl = createSimpleUnaryImpl(function (xi) { return Math.log(xi); }); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function maxImpl(aVals, reduceSize, outShape, dtype) { var vals = tf.util.getTypedArrayFromDType(dtype, tf.util.sizeFromShape(outShape)); for (var i = 0; i < vals.length; ++i) { var offset = i * reduceSize; var max = aVals[offset]; for (var j = 0; j < reduceSize; ++j) { var value = aVals[offset + j]; if (Number.isNaN(value) || value > max) { // comparison with NaN always return false max = value; } } vals[i] = max; } return vals; } /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var maximumImpl = createSimpleBinaryKernelImpl((function (aValue, bValue) { return Math.max(aValue, bValue); })); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var minimumImpl = createSimpleBinaryKernelImpl((function (aValue, bValue) { return Math.min(aValue, bValue); })); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var multiplyImpl = createSimpleBinaryKernelImpl((function (aValue, bValue) { return aValue * bValue; })); function negImpl(xVals, xShape, xDtype) { var minusOne = tf.util.createScalarValue(-1, xDtype); return multiplyImpl([], xShape, minusOne, xVals, xDtype); } /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var notEqualImpl = createSimpleBinaryKernelImpl((function (a, b) { return (a !== b) ? 1 : 0; })); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function transposeImpl(xVals, xShape, dtype, perm, newShape) { var xRank = xShape.length; var xSize = tf.util.sizeFromShape(xShape); var xStrides = tf.util.computeStrides(xShape); var newStrides = tf.util.computeStrides(newShape); var result = tf.util.getTypedArrayFromDType(dtype, tf.util.sizeFromShape(newShape)); for (var i = 0; i < xSize; ++i) { var loc = tf.util.indexToLoc(i, xRank, xStrides); // Permute location. var newLoc = new Array(loc.length); for (var i_1 = 0; i_1 < newLoc.length; i_1++) { newLoc[i_1] = loc[perm[i_1]]; } var newIndex = tf.util.locToIndex(newLoc, xRank, newStrides); result[newIndex] = xVals[i]; } return result; } function prodImpl(xShape, xDtype, xVals, reductionAxes) { var _a = __read(tf.backend_util.computeOutAndReduceShapes(xShape, reductionAxes), 2), outShape = _a[0], reduceShape = _a[1]; var outDtype = tf.upcastType(xDtype, 'int32'); var outVals = tf.util.makeZerosTypedArray(tf.util.sizeFromShape(outShape), outDtype); var reduceSize = tf.util.sizeFromShape(reduceShape); for (var i = 0; i < outVals.length; ++i) { var offset = i * reduceSize; var prod_1 = 1; for (var j = 0; j < reduceSize; ++j) { prod_1 *= xVals[offset + j]; } outVals[i] = prod_1; } return { outVals: outVals, outShape: outShape, outDtype: outDtype }; } tf.backend_util.RowPartitionType; /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function rangeImpl(start, stop, step, dtype) { var sameStartStop = start === stop; var increasingRangeNegativeStep = start < stop && step < 0; var decreasingRangePositiveStep = stop < start && step > 1; if (sameStartStop || increasingRangeNegativeStep || decreasingRangePositiveStep) { return tf.util.makeZerosTypedArray(0, dtype); } var numElements = Math.abs(Math.ceil((stop - start) / step)); var values = tf.util.makeZerosTypedArray(numElements, dtype); if (stop < start && step === 1) { // Auto adjust the step's sign if it hasn't been set // (or was set to 1) step = -1; } values[0] = start; for (var i = 1; i < values.length; i++) { values[i] = values[i - 1] + step; } return values; } /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the License); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an AS IS BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var rsqrtImpl = createSimpleUnaryImpl(function (xi) { return 1 / Math.sqrt(xi); }); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function scatterImpl(indices, updates, shape, outputSize, sliceSize, numUpdates, sliceRank, strides, defaultValue, sumDupeIndices) { var flattenShape = [outputSize / sliceSize, sliceSize]; var indicesData = indices.values; var updatesData = updates.values; if (outputSize === 0) { return tf.buffer(shape, updates.dtype); } var outBuf = (defaultValue instanceof tf.TensorBuffer) ? defaultValue : tf.buffer(flattenShape, updates.dtype); if (typeof defaultValue === 'string') { outBuf.values.fill(defaultValue); } else if (typeof defaultValue === 'number') { outBuf.values.fill(defaultValue); } else if (typeof defaultValue === 'boolean') { outBuf.values.fill(+defaultValue); } for (var i = 0; i < numUpdates; i++) { var index = []; var flattenIndex = 0; for (var j = 0; j < sliceRank; j++) { var dim = indicesData[i * sliceRank + j]; index.push(dim); flattenIndex += dim * strides[j]; } if (flattenIndex < 0 || flattenIndex >= outputSize / sliceSize) { throw new Error("Invalid indices: ".concat(index, " does not index into ").concat(shape)); } for (var k = 0; k < sliceSize; k++) { if (sumDupeIndices) { outBuf.values[flattenIndex * sliceSize + k] += updatesData[i * sliceSize + k]; } else { outBuf.values[flattenIndex * sliceSize + k] = updates.rank === 0 ? updatesData[0] : updatesData[i * sliceSize + k]; } } } return outBuf; } function sliceImpl(vals, begin, size, shape, dtype) { var isContinous = tf.slice_util.isSliceContinous(shape, begin, size); var length = tf.util.sizeFromShape(size); var xStrides = tf.util.computeStrides(shape); if (isContinous) { var flatOffset = tf.slice_util.computeFlatOffset(begin, xStrides); if (dtype === 'string') { return vals.slice(flatOffset, flatOffset + length); } return vals.subarray(flatOffset, flatOffset + length); } var decodedData = dtype === 'string' ? tf.backend_util.fromUint8ToStringArray(vals) : vals; var inBuf = tf.buffer(shape, dtype, decodedData); var outBuf = tf.buffer(size, dtype); for (var i = 0; i < outBuf.size; ++i) { var outLoc = outBuf.indexToLoc(i); var inLoc = outLoc.map(function (idx, j) { return idx + begin[j]; }); outBuf.set.apply(outBuf, __spreadArray([inBuf.get.apply(inBuf, __spreadArray([], __read(inLoc), false))], __read(outLoc), false)); } if (dtype === 'string') { return tf.backend_util.fromStringArrayToUint8(outBuf.values); } return outBuf.values; } function stridedSliceImpl(outShape, xBuf, strides, begin) { var outBuf = tf.buffer(outShape, xBuf.dtype); for (var i = 0; i < outBuf.size; i++) { var loc = outBuf.indexToLoc(i); var newLoc = new Array(loc.length); for (var j = 0; j < newLoc.length; j++) { newLoc[j] = loc[j] * strides[j] + begin[j]; } outBuf.set.apply(outBuf, __spreadArray([xBuf.get.apply(xBuf, __spreadArray([], __read(newLoc), false))], __read(loc), false)); } return outBuf; } /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ /** * The StringNGramsOp class creates ngrams from ragged string data. * The constructor contains all attributes related to the operation such as * padding widths and strings, and the compute function can be used to * compute the ngrams for different ragged tensor inputs. */ var StringNGramsOp = /** @class */ (function () { function StringNGramsOp(separator, nGramWidths, leftPad, rightPad, padWidth, preserveShortSequences) { this.separator = tf.util.encodeString(separator); this.nGramWidths = nGramWidths; this.leftPad = tf.util.encodeString(leftPad); this.rightPad = tf.util.encodeString(rightPad); this.padWidth = padWidth; this.preserveShort = preserveShortSequences; } StringNGramsOp.prototype.getPadWidth = function (nGramWidth) { // Ngrams can be padded with either a fixed pad width or a dynamic pad // width depending on the 'padWidth' arg, but in no case should the padding // ever be wider than 'nGramWidth' - 1. return Math.min(this.padWidth < 0 ? nGramWidth - 1 : this.padWidth, nGramWidth - 1); }; StringNGramsOp.prototype.getNumNGrams = function (length, nGramWidth) { var padWidth = this.getPadWidth(nGramWidth); return Math.max(0, ((length + 2 * padWidth) - nGramWidth) + 1); }; StringNGramsOp.prototype.createNGrams = function (data, splitIndex, output, outputStartIndex, numNGrams, nGramWidth) { var _loop_1 = function (nGramIndex) { var padWidth = this_1.getPadWidth(nGramWidth); var leftPadding = Math.max(0, padWidth - nGramIndex); var rightPadding = Math.max(0, padWidth - (numNGrams - (nGramIndex + 1))); var numTokens = nGramWidth - (leftPadding + rightPadding); var dataStartIndex = splitIndex + (leftPadding > 0 ? 0 : nGramIndex - padWidth); // Calculate the total expected size of the nGram so we can reserve the // correct amount of space in the string. var nGramSize = 0; // Size of the left padding. nGramSize += leftPadding * this_1.leftPad.length; // Size of the tokens. for (var n = 0; n < numTokens; ++n) { nGramSize += data[dataStartIndex + n].length; } // Size of the right padding. nGramSize += rightPadding * this_1.rightPad.length; // Size of the separators. var numSeparators = leftPadding + rightPadding + numTokens - 1; nGramSize += numSeparators * this_1.separator.length; // Build the nGram. output[outputStartIndex + nGramIndex] = new Uint8Array(nGramSize); var nGram = output[outputStartIndex + nGramIndex]; var nextNGramIndex = 0; var appendToNGram = function (str) { return str.forEach(function (value) { return nGram[nextNGramIndex++] = value; }); }; for (var n = 0; n < leftPadding; ++n) { appendToNGram(this_1.leftPad); appendToNGram(this_1.separator); } // Only output first numTokens - 1 pairs of data and separator for (var n = 0; n < numTokens - 1; ++n) { appendToNGram(data[dataStartIndex + n]); appendToNGram(this_1.separator); } // Handle case when there are no tokens or no right padding as these // can result in consecutive separators. if (numTokens > 0) { // If we have tokens, then output last and then pair each separator // with the right padding that follows, to ensure nGram ends either with // the token or with the right pad. appendToNGram(data[dataStartIndex + numTokens - 1]); for (var n = 0; n < rightPadding; ++n) { appendToNGram(this_1.separator); appendToNGram(this_1.rightPad); } } else { // If we don't have tokens, then the last item inserted into the nGram // has been the separator from the left padding loop above. Hence, // output right pad and separator and make sure to finish with a // padding, not a separator. for (var n = 0; n < rightPadding - 1; ++n) { appendToNGram(this_1.rightPad); appendToNGram(this_1.separator); } appendToNGram(this_1.rightPad); } }; var this_1 = this; for (var nGramIndex = 0; nGramIndex < numNGrams; ++nGramIndex) { _loop_1(nGramIndex); } }; // Data and splits together form the definition of the ragged tensor, // where data is 1 dimensional and contains the values of the tensor // and splits denotes the indices at which each row starts. StringNGramsOp.prototype.compute = function (data, splits) { var _this = this; // Validate that the splits are valid indices into data, only if there are // splits specified. var inputDataSize = data.length; var splitsSize = splits.length; if (splitsSize > 0) { var prevSplit = splits[0]; if (prevSplit !== 0) { throw new Error("First split value must be 0, got ".concat(prevSplit)); } for (var i = 1; i < splitsSize; ++i) { var validSplits = splits[i] >= prevSplit; validSplits = validSplits && (splits[i] <= inputDataSize); if (!validSplits) { throw new Error("Invalid split value ".concat(splits[i], ", must be in [").concat(prevSplit, ", ").concat(inputDataSize, "]")); } prevSplit = splits[i]; } if (prevSplit !== inputDataSize) { throw new Error("Last split value must be data size. Expected ".concat(inputDataSize, ", got ").concat(prevSplit)); } } var numBatchItems = splitsSize - 1; var nGramsSplits = tf.util.getArrayFromDType('int32', splitsSize); // If there is no data or size, return an empty ragged tensor. if (inputDataSize === 0 || splitsSize === 0) { var empty = new Array(inputDataSize); for (var i = 0; i <= numBatchItems; ++i) { nGramsSplits[i] = 0; } return [empty, nGramsSplits]; } nGramsSplits[0] = 0; var _loop_2 = function (i) { var length = splits[i] - splits[i - 1]; var numNGrams = 0; this_2.nGramWidths.forEach(function (nGramWidth) { numNGrams += _this.getNumNGrams(length, nGramWidth); }); if (this_2.preserveShort && length > 0 && numNGrams === 0) { numNGrams = 1; } nGramsSplits[i] = nGramsSplits[i - 1] + numNGrams; }; var this_2 = this; for (var i = 1; i <= numBatchItems; ++i) { _loop_2(i); } var nGrams = new Array(nGramsSplits[numBatchItems]); var _loop_3 = function (i) { var splitIndex = splits[i]; var outputStartIdx = nGramsSplits[i]; this_3.nGramWidths.forEach(function (nGramWidth) { var length = splits[i + 1] - splits[i]; var numNGrams = _this.getNumNGrams(length, nGramWidth); _this.createNGrams(data, splitIndex, nGrams, outputStartIdx, numNGrams, nGramWidth); outputStartIdx += numNGrams; }); // If we're preserving short sequences, check to see if no sequence was // generated by comparing the current output start idx to the original // one (nGramSplitsdata). If no ngrams were generated, then they will // be equal (since we increment outputStartIdx by numNGrams every // time we create a set of ngrams.) if (this_3.preserveShort && outputStartIdx === nGramsSplits[i]) { var dataLength = splits[i + 1] - splits[i]; // One legitimate reason to not have any ngrams when this.preserveShort // is true is if the sequence itself is empty. In that case, move on. if (dataLength === 0) { return "continue"; } // We don't have to worry about dynamic padding sizes here: if padding // was dynamic, every sequence would have had sufficient padding to // generate at least one nGram. var nGramWidth = dataLength + 2 * this_3.padWidth; var numNGrams = 1; this_3.createNGrams(data, splitIndex, nGrams, outputStartIdx, numNGrams, nGramWidth); } }; var this_3 = this; for (var i = 0; i < numBatchItems; ++i) { _loop_3(i); } return [nGrams, nGramsSplits]; }; return StringNGramsOp; }()); function stringNGramsImpl(data, dataSplits, separator, nGramWidths, leftPad, rightPad, padWidth, preserveShortSequences) { return new StringNGramsOp(separator, nGramWidths, leftPad, rightPad, padWidth, preserveShortSequences) .compute(data, dataSplits); } /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var subImpl = createSimpleBinaryKernelImpl((function (aValue, bValue) { return aValue - bValue; })); /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ /** * An implementation of the tile kernel shared between webgl and cpu for string * tensors only. */ function tileImpl(xBuf, reps) { var newShape = new Array(xBuf.rank); for (var i = 0; i < newShape.length; i++) { newShape[i] = xBuf.shape[i] * reps[i]; } var result = tf.buffer(newShape, xBuf.dtype); for (var i = 0; i < result.values.length; ++i) { var newLoc = result.indexToLoc(i); var originalLoc = new Array(xBuf.rank); for (var j = 0; j < originalLoc.length; j++) { originalLoc[j] = newLoc[j] % xBuf.shape[j]; } var originalIndex = xBuf.locToIndex(originalLoc); result.values[i] = xBuf.values[originalIndex]; } return result; } var comparePair = function (a, b) { var valueDiff = b.value - a.value; return valueDiff === 0 ? a.index - b.index : valueDiff; }; /** * Partitions array where all elements smaller than the (k+1) smallest element * are found to the left of it, and all larger to the right of it. * Based on the Floyd-Rivest Algorithm, ref: * https://en.wikipedia.org/wiki/Floyd%E2%80%93Rivest_algorithm * @param array: Array to partition * @param left: Left index for the interval * @param right: Right index for the interval * @param k: Desired index value, where array[k] is the (k+1)th smallest element * when left = 0 */ function select$1(array, k, left, right) { if (left === void 0) { left = 0; } if (right === void 0) { right = array.length - 1; } while (right > left) { // Use select recursively to sample a smaller set of size s // the arbitrary constants 600 and 0.5 are used in the original // version to minimize execution time. if (right - left > 600) { var n = right - left + 1; var i_1 = k - left + 1; var z = Math.log(n); var s = 0.5 * Math.exp(2 * z / 3); var sd = 0.5 * Math.sqrt(z * s * (n - s) / n) * Math.sign(i_1 - n / 2); var newLeft = Math.max(left, Math.floor(k - i_1 * s / n + sd)); var newRight = Math.min(right, Math.floor(k + (n - i_1) * s / n + sd)); select$1(array, k, newLeft, newRight); } // partition the elements between left and right around t var t = array[k]; var i = left; var j = right; tf.util.swap(array, left, k); if (comparePair(array[right], t) > 0) { tf.util.swap(array, left, right); } while (i < j) { tf.util.swap(array, i, j); i++; j--; while (comparePair(array[i], t) < 0) { i = i + 1; } while (comparePair(array[j], t) > 0) { j = j - 1; } } if (comparePair(array[left], t) === 0) { tf.util.swap(array, left, j); } else { j = j + 1; tf.util.swap(array, j, right); } // Adjust left and right towards the boundaries of the subset // containing the (k - left + 1)th smallest element. if (j <= k) { left = j + 1; } if (k <= j) { right = j - 1; } } } function topKImpl(x, xShape, xDtype, k, sorted) { // Reshape into a 2d tensor [batch, lastDim] and compute topk along lastDim. var lastDim = xShape[xShape.length - 1]; var _a = __read([x.length / lastDim, lastDim], 2), batch = _a[0], size = _a[1]; var allTopKVals = tf.util.getTypedArrayFromDType(xDtype, batch * k); var allTopKIndices = tf.util.getTypedArrayFromDType('int32', batch * k); var _loop_1 = function (b) { var offset = b * size; var vals = x.subarray(offset, offset + size); var valAndInd = new Array(vals.length); vals.forEach(function (value, index) { return valAndInd[index] = { value: value, index: index }; }); if (k < valAndInd.length) { select$1(valAndInd, k); valAndInd = valAndInd.slice(0, k); } if (sorted) { valAndInd.sort(comparePair); } var outOffset = b * k; var topKVals = allTopKVals.subarray(outOffset, outOffset + k); var topKIndices = allTopKIndices.subarray(outOffset, outOffset + k); for (var i = 0; i < k; i++) { topKVals[i] = valAndInd[i].value; topKIndices[i] = valAndInd[i].index; } }; for (var b = 0; b < batch; b++) { _loop_1(b); } // Reshape back to the original input shape, except that the last // dimension is k. var outputShape = xShape.slice(); outputShape[outputShape.length - 1] = k; return [ tf.buffer(outputShape, xDtype, allTopKVals), tf.buffer(outputShape, 'int32', allTopKIndices) ]; } /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var addImplCPU = addImpl, castImplCPU = castImpl, ceilImplCPU = ceilImpl, concatImplCPU = concatImpl$1, equalImplCPU = equalImpl, expImplCPU = expImpl, expm1ImplCPU = expm1Impl, floorImplCPU = floorImpl, floorDivImplCPU = floorDivImpl, gatherNdImplCPU = gatherNdImpl, gatherV2ImplCPU = gatherV2Impl, greaterEqualImplCPU = greaterEqualImpl, greaterImplCPU = greaterImpl, lessEqualImplCPU = lessEqualImpl, lessImplCPU = lessImpl, logImplCPU = logImpl, maxImplCPU = maxImpl, maximumImplCPU = maximumImpl, minimumImplCPU = minimumImpl, multiplyImplCPU = multiplyImpl, negImplCPU = negImpl, notEqualImplCPU = notEqualImpl, prodImplCPU = prodImpl, rangeImplCPU = rangeImpl, rsqrtImplCPU = rsqrtImpl, scatterImplCPU = scatterImpl, simpleAbsImplCPU = simpleAbsImpl, sliceImplCPU = sliceImpl, stridedSliceImplCPU = stridedSliceImpl, stringNGramsImplCPU = stringNGramsImpl, subImplCPU = subImpl, tileImplCPU = tileImpl, topKImplCPU = topKImpl, transposeImplCPU = transposeImpl; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var abs = unaryKernelFunc({ opType: UnaryOpType.ABS, cpuKernelImpl: simpleAbsImplCPU }); var absConfig = { kernelName: tf.Abs, backendName: 'webgpu', kernelFunc: abs }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var acos = unaryKernelFunc({ opType: UnaryOpType.ACOS }); var acosConfig = { kernelName: tf.Acos, backendName: 'webgpu', kernelFunc: acos }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var acosh = unaryKernelFunc({ opType: UnaryOpType.ACOSH }); var acoshConfig = { kernelName: tf.Acosh, backendName: 'webgpu', kernelFunc: acosh }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var addKernelFunc = binaryKernelFunc({ opType: BinaryOpType.ADD, cpuKernelImpl: addImplCPU, supportsComplex: true }); var addConfig = { kernelName: tf.Add, backendName: 'webgpu', kernelFunc: addKernelFunc }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var AddNPackedProgram = /** @class */ (function () { function AddNPackedProgram(shapes) { this.workPerThread = 1; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = shapes[0]; this.variableNames = shapes.map(function (_, i) { return "T".concat(i); }); this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize, [this.workPerThread, 1, 1]); this.shaderKey = 'addN'; } AddNPackedProgram.prototype.getUserCode = function () { var snippets = []; // Get target elements from every input tensor. this.variableNames.forEach(function (variable) { snippets.push("let v".concat(variable, " = get").concat(variable, "ByOutputCoords(coords);")); }); // Calculate the sum of all elements. var operation = this.variableNames .map(function (variable) { return "v".concat(variable); }) .join(' + '); var userCode = "\n ".concat(getMainHeaderString('index'), " {\n for (var i = 0; i < ").concat(this.workPerThread, "; i = i + 1) {\n let flatIndex = index * ").concat(this.workPerThread, " + i;\n if (flatIndex < uniforms.size) {\n let coords = getCoordsFromIndex(flatIndex);\n ").concat(snippets.join('\n '), "\n setOutputAtIndex(flatIndex, ").concat(operation, ");\n }\n }\n }\n "); return userCode; }; return AddNPackedProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function addN(args) { var inputs = args.inputs, backend = args.backend; var tensors = inputs; if (tensors.length === 1) { return identity({ inputs: { x: tensors[0] }, backend: backend }); } var dtype = tensors.map(function (t) { return t.dtype; }).reduce(function (d1, d2) { return tf.upcastType(d1, d2); }); var shapes = tensors.map(function (t) { return t.shape; }); var program = new AddNPackedProgram(shapes); return backend.runWebGPUProgram(program, tensors, dtype); } var addNConfig = { kernelName: tf.AddN, backendName: 'webgpu', kernelFunc: addN }; /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var TransposeSharedProgram = /** @class */ (function () { function TransposeSharedProgram(aShape, newDim) { this.variableNames = ['A']; // Note that the maximum number of workgroup invocations by webgpu is 256. this.workgroupSize = [16, 16, 1]; var outputShape = new Array(aShape.length); for (var i = 0; i < outputShape.length; i++) { outputShape[i] = aShape[newDim[i]]; } this.outputShape = outputShape; this.dispatchLayout = { x: [0], y: [1] }; this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize, [1, 1, 1]); this.shaderKey = 'transposeShared'; } TransposeSharedProgram.prototype.getUserCode = function () { var _this = this; tf.util.assert(this.workgroupSize[0] === this.workgroupSize[1], function () { return "Must be a square tile, current tile shape is ".concat(_this.workgroupSize[0], " x ").concat(_this.workgroupSize[1]); }); var tileSize = this.workgroupSize[0]; var userCode = "\n var tile : array, ").concat(this.workgroupSize[0], ">;\n ").concat(getMainHeaderString(), " {\n var x = i32(workgroupId.x) * ").concat(tileSize, " + i32(localId.x);\n var y = i32(workgroupId.y) * ").concat(tileSize, " + i32(localId.y);\n let width = uniforms.outShape[0];\n let height = uniforms.outShape[1];\n if (x < width && y < height) {\n tile[localId.y][localId.x] = f32(A[y * width + x]);\n }\n workgroupBarrier();\n\n x = i32(workgroupId.y) * ").concat(tileSize, " + i32(localId.x);\n y = i32(workgroupId.x) * ").concat(tileSize, " + i32(localId.y);\n if (x < height && y < width) {\n setOutputAtIndex((y * height + x), tile[localId.x]\n [localId.y]);\n }\n }\n "); return userCode; }; return TransposeSharedProgram; }()); /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var TransposeProgram = /** @class */ (function () { function TransposeProgram(aShape, newDim) { this.variableNames = ['A']; this.workPerThread = 1; this.workgroupSize = [64, 1, 1]; this.size = true; var outputShape = new Array(aShape.length); for (var i = 0; i < outputShape.length; i++) { outputShape[i] = aShape[newDim[i]]; } this.outputShape = outputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize, [this.workPerThread, 1, 1]); this.newDim = newDim; this.shaderKey = "transpose_".concat(newDim); } TransposeProgram.prototype.getUserCode = function () { var dtype = getCoordsDataType(this.outputShape.length); var switched = getSwitchedCoords(this.newDim); var userCode = "\n ".concat(getMainHeaderString('index'), " {\n for(var i = 0; i < ").concat(this.workPerThread, "; i = i + 1) {\n let flatIndex = index * ").concat(this.workPerThread, " + i;\n if(flatIndex < uniforms.size) {\n let coords = getCoordsFromIndex(flatIndex);\n setOutputAtIndex(flatIndex, A[getIndexFromCoords").concat(this.outputShape.length, "D(\n ").concat(dtype, "(").concat(switched, "), uniforms.aShape)]);\n }\n }\n }\n "); return userCode; }; return TransposeProgram; }()); function getSwitchedCoords(newDim) { var rank = newDim.length; if (rank > 6) { throw Error("Transpose for rank ".concat(rank, " is not yet supported")); } var switchedCoords = new Array(rank); for (var i = 0; i < newDim.length; i++) { switchedCoords[newDim[i]] = "coords.".concat(getCoordsXYZ(i)); } return switchedCoords.join(); } /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function transpose(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var perm = attrs.perm; var webgpuBackend = backend; var xRank = x.shape.length; var newShape = new Array(xRank); for (var i = 0; i < newShape.length; i++) { newShape[i] = x.shape[perm[i]]; } if (backend.shouldExecuteOnCPU([x])) { var xData = webgpuBackend.tensorMap.get(x.dataId); var values = xData.values; var outValues = transposeImplCPU(values, x.shape, x.dtype, perm, newShape); return backend.makeTensorInfo(newShape, x.dtype, outValues); } if (x.shape.length === 2 && tf.util.arraysEqual(perm, [1, 0])) { var program_1 = new TransposeSharedProgram(x.shape, perm); return webgpuBackend.runWebGPUProgram(program_1, [x], x.dtype); } var program = new TransposeProgram(x.shape, perm); return webgpuBackend.runWebGPUProgram(program, [x], x.dtype); } var transposeConfig = { kernelName: tf.Transpose, backendName: 'webgpu', kernelFunc: transpose }; var ReduceProgram = /** @class */ (function () { function ReduceProgram(reduceInfo, reduceType, maxComputeWorkgroupSizeX) { this.variableNames = ['x']; this.uniforms = 'reduceSize : i32,'; this.size = true; this.inputShape = [reduceInfo.batchSize, reduceInfo.inSize]; var _a = __read(tf.backend_util.computeOutAndReduceShapes(this.inputShape, [1]), 1), outputShape = _a[0]; this.outputShape = outputShape.length === 0 ? [1] : outputShape; // If reduceSize |reduceInfo.inSize| is very large, the I/O accessing will // become the bottleneck. Increasing workgroupSize can reduce the times of // accessing global memory. The threshold value is just to make sure the // reduceSize is large enough for a bigger workgroupSize. if (reduceInfo.inSize >= 32768 && maxComputeWorkgroupSizeX >= 512) { this.workgroupSize = [512, 1, 1]; } else if (reduceInfo.inSize >= 4096) { this.workgroupSize = [256, 1, 1]; } else { this.workgroupSize = [64, 1, 1]; } this.dispatchLayout = flatDispatchLayout(this.outputShape); // A work group only outputs a data, so we transfer [1, 1, 1] to compute // dispatch size. this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, [1, 1, 1]); this.reduceType = reduceType; this.shaderKey = "reduce_".concat(reduceType); } ReduceProgram.prototype.getUserCode = function () { var reduceOp = ""; var initValue = '0.0'; var workgroupSizeX = this.workgroupSize[0]; if (this.reduceType === 'min' || this.reduceType === 'max') { reduceOp = "\n if (isnan(candidate)) {\n bestValue = uniforms.NAN;\n } else if (!isnan(bestValue) && candidate ".concat(this.reduceType === 'min' ? '<' : '>', " bestValue)\n { bestValue = candidate; }"); initValue = 'f32(x[offset])'; } else if (this.reduceType === 'sum' || this.reduceType === 'mean') { reduceOp = ' bestValue = bestValue + candidate; '; } else if (this.reduceType === 'prod') { reduceOp = ' bestValue = bestValue * candidate; '; initValue = '1.0'; } else if (this.reduceType === 'all') { reduceOp = ' bestValue = f32(bestValue >= 1.0 && candidate >= 1.0); '; initValue = '1.0'; } else if (this.reduceType === 'any') { reduceOp = ' bestValue = f32(bestValue >= 1.0 || candidate >= 1.0); '; initValue = '0.0'; } var outputSnippet = this.reduceType === 'mean' ? // tslint:disable-next-line:max-line-length "setOutputAtIndex(outputIndex, bestValue / f32(uniforms.reduceSize));" : "setOutputAtIndex(outputIndex, bestValue);"; var sharedMemorySnippet = "\n var xBestValues : array;\n "); var userCode = "\n fn DIV_CEIL(a : u32, b : u32) -> u32 {\n return ((a - 1u) / b + 1u);\n }\n\n ".concat(sharedMemorySnippet, "\n fn getOffset(outputIndex : i32) -> i32 {\n let outputCoords = getCoordsFromIndex(outputIndex);\n let offset = ").concat(this.outputShape.length === 1 ? 'outputCoords' : 'outputCoords[0]', " * uniforms.reduceSize;\n return offset;\n }\n ").concat(getMainHeaderString('index'), " {\n let outputIndex = index / ").concat(workgroupSizeX, ";\n let offset = getOffset(outputIndex);\n var bestValue = ").concat(initValue, ";\n let Length = uniforms.reduceSize;\n let WorkPerThread = DIV_CEIL(u32(Length), ").concat(workgroupSizeX, "u);\n for (var k = i32(localId.x); k < Length && outputIndex < uniforms.size;\n k = k + ").concat(workgroupSizeX, ") {\n let candidate = f32(x[offset + k]);\n ").concat(reduceOp, "\n }\n xBestValues[localId.x] = bestValue;\n workgroupBarrier();\n\n var reduceSize = min(u32(Length), ").concat(workgroupSizeX, "u);\n for (var currentSize = reduceSize / 2u; reduceSize > 1u;\n currentSize = reduceSize / 2u) {\n let interval = DIV_CEIL(reduceSize, 2u);\n if (localId.x < currentSize) {\n let candidate = xBestValues[localId.x + interval];\n ").concat(reduceOp, "\n xBestValues[localId.x] = bestValue;\n }\n reduceSize = interval;\n workgroupBarrier();\n }\n\n if (localId.x == 0u && outputIndex < uniforms.size) {\n ").concat(outputSnippet, "\n }\n }\n "); return userCode; }; return ReduceProgram; }()); var RETURN_TYPES = { 'mean': 'float32', 'all': 'bool', 'any': 'bool', }; function reduce(x, axis, keepDims, reduceType, backend) { var xRank = x.shape.length; var toDispose = []; var origAxes = tf.util.parseAxisParam(axis, x.shape); var axes = origAxes; var permutedAxes = tf.backend_util.getAxesPermutation(axes, xRank); var input = x; if (permutedAxes != null) { input = transpose({ inputs: { x: x }, attrs: { perm: permutedAxes }, backend: backend }); axes = tf.backend_util.getInnerMostAxes(axes.length, xRank); toDispose.push(input); } tf.backend_util.assertAxesAreInnerMostDims(reduceType, axes, xRank); var _a = __read(tf.backend_util.computeOutAndReduceShapes(input.shape, axes), 2), reduceOutShape = _a[0], reduceShape = _a[1]; var resOutShape = reduceOutShape; if (keepDims) { // rather than reshape at the end, set the target shape here. resOutShape = tf.backend_util.expandShapeToKeepDim(reduceOutShape, origAxes); } var res; if ((reduceType === 'max' || reduceType === 'prod') && backend.shouldExecuteOnCPU([input])) { var xVals = backend.tensorMap.get(input.dataId).values; switch (reduceType) { case 'max': var outValues = maxImplCPU(xVals, tf.util.sizeFromShape(reduceShape), resOutShape, x.dtype); res = backend.makeTensorInfo(resOutShape, x.dtype, outValues); break; case 'prod': var _b = prodImplCPU(input.shape, input.dtype, xVals, axes), outVals = _b.outVals, outShape = _b.outShape, outDtype = _b.outDtype; res = backend.makeTensorInfo(outShape, outDtype, outVals); break; default: throw new Error("".concat(reduceType, " CPU implementation is not yet supported.")); } } else { var inSize = tf.util.sizeFromShape(reduceShape); var xSize = tf.util.sizeFromShape(input.shape); var batchSize = xSize / inSize; var reduceInfo = { windowSize: inSize, inSize: inSize, batchSize: batchSize, outSize: 1 }; var dtype = RETURN_TYPES[reduceType] || tf.sumOutType(x.dtype); var uniformData = [ { type: 'int32', data: [inSize] }, ]; var program = new ReduceProgram(reduceInfo, reduceType, backend.device.limits.maxComputeWorkgroupSizeX); var reduced = backend.runWebGPUProgram(program, [input], dtype, uniformData); toDispose.push(reduced); res = reshape({ inputs: { x: reduced }, attrs: { shape: resOutShape }, backend: backend }); } toDispose.forEach(function (t) { return backend.disposeData(t.dataId); }); return res; } /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function all(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var keepDims = attrs.keepDims, axis = attrs.axis; return reduce(x, axis, keepDims, 'all', backend); } var allConfig = { kernelName: tf.All, backendName: 'webgpu', kernelFunc: all }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function any(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var keepDims = attrs.keepDims, axis = attrs.axis; return reduce(x, axis, keepDims, 'any', backend); } var anyConfig = { kernelName: tf.Any, backendName: 'webgpu', kernelFunc: any }; var ArgMinMaxProgram = /** @class */ (function () { function ArgMinMaxProgram(inputShape, axis, reduceType) { this.workgroupSize = [64, 1, 1]; this.variableNames = ['x']; this.uniforms = 'infinityValue : f32,'; this.size = true; var axes = [axis]; this.op = reduceType === 'min' ? '<' : '>'; // |outShape| is the shape with the removed axis var _a = __read(tf.backend_util.computeOutAndReduceShapes(inputShape, axes), 2), outputShape = _a[0], reduceShape = _a[1]; this.outputShape = outputShape.length === 0 ? [1] : outputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); // The shared algorithm is mainly used for large reduce size. It fully // utilizes the threads in one workgroup to do the reduction. However, // when the reduce size is very small, it's better to use the plain // algorithm to reduce the number of workgroups to speedup. The threthold // can be further tuned. if (tf.util.sizeFromShape(reduceShape) < 32) { this.type = 'plain'; this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); } else { this.type = 'shared'; // A work group only outputs a data, so we transfer [1, 1, 1] to compute // dispatch size. this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, [1, 1, 1]); } this.inputShape = inputShape; this.shaderKey = "argMinMax_".concat(this.op, "_").concat(this.type); } ArgMinMaxProgram.prototype.getUserCode = function () { var _this = this; var workgroupSizeX = this.workgroupSize[0]; var getInputShapeLastDim = function () { if (_this.inputShape.length === 1) { return 'uniforms.xShape'; } else { return "uniforms.xShape.".concat(getCoordsXYZ(_this.inputShape.length - 1)); } }; var splitOutputCoords = function () { var snippet = ''; if (_this.outputShape.length === 1) { if (_this.inputShape.length !== 1) { snippet += 'outputCoords,'; } } else { for (var i = 0; i < _this.outputShape.length; i++) { snippet += "outputCoords.".concat(getCoordsXYZ(i), ","); } } return snippet; }; if (this.type === 'shared') { var sharedMemorySnippet = "\n var xBestIndices : array;\n var xBestValues : array;\n "); var userCode = "\n fn DIV_CEIL(a : u32, b : u32) -> u32 {\n return ((a - 1u) / b + 1u);\n }\n\n ".concat(sharedMemorySnippet, "\n\n ").concat(getMainHeaderString('index'), " {\n let outputIndex = index / ").concat(workgroupSizeX, ";\n let reduceLength = ").concat(getInputShapeLastDim(), ";\n\n var bestIndex = i32(localId.x);\n var bestValue = uniforms.infinityValue;\n let outputCoords = getCoordsFromIndex(outputIndex);\n for (var k = i32(localId.x); k < reduceLength && outputIndex < uniforms.size;\n k = k + ").concat(workgroupSizeX, ") {\n let candidate = getX(").concat(splitOutputCoords(), " k);\n if (!isnan(candidate) && candidate ").concat(this.op, " bestValue) {\n bestValue = candidate;\n bestIndex = k;\n }\n }\n xBestValues[localId.x] = bestValue;\n xBestIndices[localId.x] = bestIndex;\n workgroupBarrier();\n\n var reduceSize = min(u32(reduceLength), ").concat(workgroupSizeX, "u);\n for (var currentSize = reduceSize / 2u; reduceSize > 1u;\n currentSize = reduceSize / 2u) {\n let interval = DIV_CEIL(reduceSize, 2u);\n if (localId.x < currentSize) {\n let candidate = xBestValues[localId.x + interval];\n if (candidate ").concat(this.op, " bestValue) {\n bestValue = candidate;\n xBestValues[localId.x] = bestValue;\n xBestIndices[localId.x] = xBestIndices[localId.x + interval];\n }\n }\n reduceSize = interval;\n workgroupBarrier();\n }\n\n if (localId.x == 0u && outputIndex < uniforms.size) {\n setOutputAtIndexI32(outputIndex, xBestIndices[localId.x]);\n }\n }\n "); return userCode; } else { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let outputCoords = getCoordsFromIndex(index);\n var bestIndex = 0;\n var bestValue = getX(").concat(splitOutputCoords(), " 0);\n let reduceLength = ").concat(getInputShapeLastDim(), ";\n for (var i = 1; i < reduceLength; i++) {\n let candidate = getX(").concat(splitOutputCoords(), " i);\n if (candidate ").concat(this.op, " bestValue) {\n bestValue = candidate;\n bestIndex = i;\n }\n }\n setOutputAtIndexI32(index, bestIndex);\n }\n }\n "); return userCode; } }; return ArgMinMaxProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function argMax(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var axis = attrs.axis; var axes = tf.util.parseAxisParam(axis, x.shape); var permutedAxes = tf.backend_util.getAxesPermutation(axes, x.shape.length); var $x = x; var intermediateTensorInfos = []; if (permutedAxes != null) { $x = transpose({ inputs: { x: x }, backend: backend, attrs: { perm: permutedAxes } }); intermediateTensorInfos.push($x); axes = tf.backend_util.getInnerMostAxes(axes.length, $x.shape.length); } tf.backend_util.assertAxesAreInnerMostDims('argMax', [axes[0]], $x.shape.length); var program = new ArgMinMaxProgram($x.shape, axes[0], 'max'); var uniformData = [{ type: 'float32', data: [Number.NEGATIVE_INFINITY] }]; var out = backend.runWebGPUProgram(program, [$x], 'int32', uniformData); intermediateTensorInfos.forEach(function (t) { return backend.disposeData(t.dataId); }); return out; } var argMaxConfig = { kernelName: tf.ArgMax, backendName: 'webgpu', kernelFunc: argMax }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function argMin(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var axis = attrs.axis; var axes = tf.util.parseAxisParam(axis, x.shape); var permutedAxes = tf.backend_util.getAxesPermutation(axes, x.shape.length); var $x = x; var intermediateTensorInfos = []; if (permutedAxes != null) { $x = transpose({ inputs: { x: x }, backend: backend, attrs: { perm: permutedAxes } }); intermediateTensorInfos.push($x); axes = tf.backend_util.getInnerMostAxes(axes.length, $x.shape.length); } tf.backend_util.assertAxesAreInnerMostDims('argMin', [axes[0]], $x.shape.length); var program = new ArgMinMaxProgram($x.shape, axes[0], 'min'); var uniformData = [{ type: 'float32', data: [Number.POSITIVE_INFINITY] }]; var out = backend.runWebGPUProgram(program, [$x], 'int32', uniformData); intermediateTensorInfos.forEach(function (t) { return backend.disposeData(t.dataId); }); return out; } var argMinConfig = { kernelName: tf.ArgMin, backendName: 'webgpu', kernelFunc: argMin }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var asin = unaryKernelFunc({ opType: UnaryOpType.ASIN }); var asinConfig = { kernelName: tf.Asin, backendName: 'webgpu', kernelFunc: asin }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var asinh = unaryKernelFunc({ opType: UnaryOpType.ASINH }); var asinhConfig = { kernelName: tf.Asinh, backendName: 'webgpu', kernelFunc: asinh }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var atan = unaryKernelFunc({ opType: UnaryOpType.ATAN }); var atanConfig = { kernelName: tf.Atan, backendName: 'webgpu', kernelFunc: atan }; /** * @license * Copyright 2022 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var atan2 = binaryKernelFunc({ opType: BinaryOpType.ATAN2 }); var atan2Config = { kernelName: tf.Atan2, backendName: 'webgpu', kernelFunc: atan2 }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var atanh = unaryKernelFunc({ opType: UnaryOpType.ATANH }); var atanhConfig = { kernelName: tf.Atanh, backendName: 'webgpu', kernelFunc: atanh }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var PoolWithFilterSizeEqualsOneProgram = /** @class */ (function () { function PoolWithFilterSizeEqualsOneProgram(convInfo) { this.variableNames = ['x']; this.uniforms = "strides : vec2,"; this.workgroupSize = [256, 1, 1]; this.size = true; this.outputShape = convInfo.outShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'poolWithFilterSizeEqualsOne'; } PoolWithFilterSizeEqualsOneProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let batch = coords[0];\n let d = coords[3];\n\n let xRCCorner = coords.yz * uniforms.strides;\n let xRCorner = xRCCorner.x;\n let xCCorner = xRCCorner.y;\n\n let value = getX(batch, xRCorner, xCCorner, d);\n setOutputAtIndex(index, value);\n }\n }\n "); return userCode; }; return PoolWithFilterSizeEqualsOneProgram; }()); /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var Pool2DProgram = /** @class */ (function () { function Pool2DProgram(convInfo, poolType, computePositions, flattenPositions, includeBatchIndex) { if (computePositions === void 0) { computePositions = false; } if (flattenPositions === void 0) { flattenPositions = false; } if (includeBatchIndex === void 0) { includeBatchIndex = false; } this.variableNames = ['x']; this.uniforms = "strides : vec2, pads : vec2, dilations : vec2, convDims : vec2, filterDims : vec2,"; // TODO(jiajia.qin@intel.com): Dynamically choose different workgroupSize for // different output shapes. this.workgroupSize = [128, 1, 1]; this.size = true; if (poolType === 'avg' && computePositions) { throw new Error('Cannot compute positions for average pool.'); } this.outputShape = convInfo.outShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.poolType = poolType; this.computePositions = computePositions; this.flattenPositions = flattenPositions; this.includeBatchIndex = includeBatchIndex; this.shaderKey = "pool2D_".concat(poolType, "_").concat(computePositions, "_").concat(flattenPositions, "_").concat(includeBatchIndex); } Pool2DProgram.prototype.getUserCode = function () { var updateSnippet; if (this.poolType === 'avg') { updateSnippet = "resultValue = resultValue + value; count = count + 1.0;"; } else if (this.computePositions) { var positionStr = this.flattenPositions ? (this.includeBatchIndex ? "((batch * uniforms.xShape[1] + xR) * uniforms.xShape[2] + xC) * uniforms.xShape[3] + d" : "(xR * uniforms.xShape[2] + xC) * uniforms.xShape[3] + d") : "wR * uniforms.filterDims.y + wC"; updateSnippet = "let currMaxValue = mix(value, maxValue, maxValueFound);\n if (value >= currMaxValue) {\n maxValue = value;\n maxValueFound = 1.0;\n maxPosition = ".concat(positionStr, ";\n }"); } else { updateSnippet = "resultValue = max(value, resultValue);"; } var returnValue = "resultValue"; if (this.poolType === 'avg') { returnValue = "resultValue / max(count, 1.0)"; } var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let batch = coords[0];\n let d = coords[3];\n let xRCCorner = vec2(coords.yz) * uniforms.strides - uniforms.pads;\n let xRCorner = xRCCorner.x;\n let xCCorner = xRCCorner.y;\n\n ").concat(this.computePositions ? "var maxValue = 0.0;\n var maxValueFound = 0.0;\n var maxPosition = 0;" : "var resultValue = ".concat(this.poolType === 'avg' ? '0.0' : '-1.0 / pow(10.0, -20.0)', ";"), "\n\n var count = 0.0;\n for (var wR = 0; wR < uniforms.filterDims.x; wR = wR + uniforms.dilations.x) {\n let xR = xRCorner + wR;\n\n if (xR < 0 || xR >= uniforms.convDims.x) {\n continue;\n }\n\n for (var wC = 0; wC < uniforms.filterDims.y; wC = wC + uniforms.dilations.y) {\n let xC = xCCorner + wC;\n if (xC < 0 || xC >= uniforms.convDims.y) {\n continue;\n }\n\n let value = getX(batch, xR, xC, d);\n ").concat(updateSnippet, "\n }\n }\n\n ").concat(this.computePositions ? "setOutputAtIndexI32(index, maxPosition);" : "setOutputAtIndex(index, ".concat(returnValue, ");"), "\n }\n }\n "); return userCode; }; return Pool2DProgram; }()); var Pool3DProgram = /** @class */ (function () { function Pool3DProgram(convInfo, poolType, computePositions, flattenPositions, includeBatchIndex) { if (computePositions === void 0) { computePositions = false; } if (flattenPositions === void 0) { flattenPositions = false; } if (includeBatchIndex === void 0) { includeBatchIndex = false; } this.variableNames = ['x']; this.uniforms = "strides : vec3, pads : vec3, convDims : vec3, filterDims : vec3,"; this.workgroupSize = [128, 1, 1]; this.size = true; if (poolType === 'avg' && computePositions) { throw new Error('Cannot compute positions for average pool.'); } this.outputShape = convInfo.outShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.poolType = poolType; this.computePositions = computePositions; this.flattenPositions = flattenPositions; this.includeBatchIndex = includeBatchIndex; this.shaderKey = "pool3D_".concat(poolType, "_").concat(computePositions, "_").concat(flattenPositions, "_").concat(includeBatchIndex); } Pool3DProgram.prototype.getUserCode = function () { var updateSnippet; if (this.poolType === 'avg') { updateSnippet = "resultValue += value; count += 1.0;"; } else if (this.computePositions) { var positionStr = this.flattenPositions ? (this.includeBatchIndex ? "(((batch * uniforms.xShape.y + xD) * uniforms.xShape.z + xR) * uniforms.xShape.w + xC) * uniforms.xShape.u + ch" : "((xD * uniforms.xShape.z + xR) * uniforms.xShape.w + xC) * uniforms.xShape.u + ch") : "wD * uniforms.filterDims.y * uniforms.filterDims.y + wR * uniforms.filterDims.z + wC"; updateSnippet = "let currMaxValue = mix(value, maxValue, maxValueFound);\n if (value >= currMaxValue) {\n maxValue = value;\n maxValueFound = 1.0;\n maxPosition = ".concat(positionStr, ";\n }"); } else { updateSnippet = "resultValue = max(value, resultValue);"; } var returnValue = "resultValue"; if (this.poolType === 'avg') { returnValue = "resultValue / max(count, 1.0)"; } var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let batch = coords.x;\n let ch = coords.u;\n\n let xCorner = vec3(coords.y, coords.z, coords.w) * uniforms.strides - uniforms.pads;\n let xDCorner = xCorner.x;\n let xRCorner = xCorner.y;\n let xCCorner = xCorner.z;\n\n ").concat(this.computePositions ? "var maxValue = 0.0;\n var maxValueFound = 0.0;\n var maxPosition = 0;" : "var resultValue = ".concat(this.poolType === 'avg' ? '0.0' : '-1.0 / pow(10.0, -20.0)', ";"), "\n\n var count = 0.0;\n for (var wD = 0; wD < uniforms.filterDims.x; wD++) {\n let xD = xDCorner + wD;\n if (xD < 0 || xD >= uniforms.convDims.x) {\n continue;\n }\n\n for (var wR = 0; wR < uniforms.filterDims.y; wR++) {\n let xR = xRCorner + wR;\n if (xR < 0 || xR >= uniforms.convDims.y) {\n continue;\n }\n\n for (var wC = 0; wC < uniforms.filterDims.z; wC++) {\n let xC = xCCorner + wC;\n if (xC < 0 || xC >= uniforms.convDims.z) {\n continue;\n }\n\n let value = getX(batch, xD, xR, xC, ch);\n ").concat(updateSnippet, "\n }\n }\n }\n\n ").concat(this.computePositions ? "setOutputAtIndexI32(index, maxPosition);" : "setOutputAtIndex(index, ".concat(returnValue, ");"), "\n }\n }\n "); return userCode; }; return Pool3DProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function max(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var reductionIndices = attrs.reductionIndices, keepDims = attrs.keepDims; return reduce(x, reductionIndices, keepDims, 'max', backend); } var maxConfig = { kernelName: tf.Max, backendName: 'webgpu', kernelFunc: max }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function mean(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var keepDims = attrs.keepDims, axis = attrs.axis; return reduce(x, axis, keepDims, 'mean', backend); } var meanConfig = { kernelName: tf.Mean, backendName: 'webgpu', kernelFunc: mean }; /** * @license * Copyright 2022 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function poolImpl(x, convInfo, poolType, backend) { if (convInfo.filterWidth === 1 && convInfo.filterHeight === 1 && tf.util.arraysEqual(convInfo.inShape, convInfo.outShape)) { return identity({ inputs: { x: x }, backend: backend }); } if (convInfo.filterWidth === convInfo.inWidth && convInfo.filterHeight === convInfo.inHeight && convInfo.batchSize === 1 && convInfo.padInfo.type === 'VALID') { var length = x.shape.length; var reshapeX = reshape({ inputs: { x: x }, backend: backend, attrs: { shape: [ x.shape[length - 3] * x.shape[length - 2] /* height * width */, x.shape[length - 1] /* channel */ ] } }); var reduceX = void 0; if (poolType === 'avg') { reduceX = mean({ inputs: { x: reshapeX }, backend: backend, attrs: { axis: 0, keepDims: false } }); } else { tf.util.assert(poolType === 'max', function () { return "Invalid pool type ".concat(poolType); }); reduceX = max({ inputs: { x: reshapeX }, backend: backend, attrs: { reductionIndices: 0, keepDims: false } }); } var result = reshape({ inputs: { x: reduceX }, backend: backend, attrs: { shape: convInfo.outShape } }); backend.disposeData(reshapeX.dataId); backend.disposeData(reduceX.dataId); return result; } var program; var dimensions = [{ type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }]; if (convInfo.filterHeight === 1 && convInfo.filterWidth === 1) { program = new PoolWithFilterSizeEqualsOneProgram(convInfo); } else { if (poolType === 'avg') { program = new Pool2DProgram(convInfo, 'avg'); } else { tf.util.assert(poolType === 'max', function () { return "Invalid pool type ".concat(poolType); }); program = new Pool2DProgram(convInfo, 'max'); } dimensions.push({ type: 'int32', data: [convInfo.padInfo.top, convInfo.padInfo.left] }, { type: 'int32', data: [convInfo.dilationHeight, convInfo.dilationWidth] }, { type: 'int32', data: [convInfo.inHeight, convInfo.inWidth] }, { type: 'int32', data: [convInfo.effectiveFilterHeight, convInfo.effectiveFilterWidth] }); } return backend.runWebGPUProgram(program, [x], x.dtype, dimensions); } /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function avgPool(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var filterSize = attrs.filterSize, strides = attrs.strides, pad = attrs.pad, dimRoundingMode = attrs.dimRoundingMode; var dilations = 1; var convInfo = tf.backend_util.computePool2DInfo(x.shape, filterSize, strides, dilations, pad, dimRoundingMode); return poolImpl(x, convInfo, 'avg', backend); } var avgPoolConfig = { kernelName: tf.AvgPool, backendName: 'webgpu', kernelFunc: avgPool }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function avgPool3D(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var filterSize = attrs.filterSize, strides = attrs.strides, pad = attrs.pad, dataFormat = attrs.dataFormat, dimRoundingMode = attrs.dimRoundingMode; var dilations = [1, 1, 1]; var convInfo = tf.backend_util.computePool3DInfo(x.shape, filterSize, strides, dilations, pad, dimRoundingMode, dataFormat); var avgPoolProgram = new Pool3DProgram(convInfo, 'avg'); var dimensions = [ { type: 'int32', data: [convInfo.strideDepth, convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.padInfo.front, convInfo.padInfo.top, convInfo.padInfo.left] }, { type: 'int32', data: [convInfo.inDepth, convInfo.inHeight, convInfo.inWidth] }, { type: 'int32', data: [ convInfo.effectiveFilterDepth, convInfo.effectiveFilterHeight, convInfo.effectiveFilterWidth ] } ]; return backend.runWebGPUProgram(avgPoolProgram, [x], x.dtype, dimensions); } var avgPool3DConfig = { kernelName: tf.AvgPool3D, backendName: 'webgpu', kernelFunc: avgPool3D }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var AvgPool2DBackpropProgram = /** @class */ (function () { function AvgPool2DBackpropProgram(convInfo) { this.variableNames = ['dy']; this.uniforms = "strides : vec2, pads : vec2, dilations : vec2, filterDims : vec2,\n outHeight : i32, outWidth : i32, avgMultiplier : f32,"; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = convInfo.inShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = "avgPool2DBackprop"; } AvgPool2DBackpropProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let batch = coords[0];\n let d = coords[3];\n\n let dyRCCorner = vec2(coords.yz) - uniforms.pads;\n let dyRCorner = dyRCCorner.x;\n let dyCCorner = dyRCCorner.y;\n\n // Convolve dy(?, ?, d) with pos mask(:, :, d) to get dx(xR, xC, d).\n // ? = to be determined. : = across all values in that axis.\n var dotProd = 0.0;\n for (var wR = 0; wR < uniforms.filterDims[0]; wR = wR + uniforms.dilations[0]) {\n let dyR = f32(dyRCorner + wR) / f32(uniforms.strides[0]);\n\n if (dyR < 0.0 || dyR >= f32(uniforms.outHeight) || fract(dyR) > 0.0) {\n continue;\n }\n let idyR = i32(dyR);\n\n for (var wC = 0; wC < uniforms.filterDims[1]; wC = wC + uniforms.dilations[1]) {\n let dyC = f32(dyCCorner + wC) / f32(uniforms.strides[1]);\n\n if (dyC < 0.0 || dyC >= f32(uniforms.outWidth) || fract(dyC) > 0.0) {\n continue;\n }\n let idyC = i32(dyC);\n\n let dyValue = getDy(batch, idyR, idyC, d);\n\n dotProd = dotProd + dyValue * uniforms.avgMultiplier;\n }\n }\n setOutputAtIndex(index, dotProd);\n }\n }\n "); return userCode; }; return AvgPool2DBackpropProgram; }()); var AvgPool3DBackpropProgram = /** @class */ (function () { function AvgPool3DBackpropProgram(convInfo) { this.variableNames = ['dy']; this.uniforms = "strides : vec3, pads : vec3, filterDims : vec3,\n outDepth : i32, outHeight : i32, outWidth : i32, avgMultiplier : f32,"; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = convInfo.inShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = "avgPool3DBackprop"; } AvgPool3DBackpropProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let batch = coords.x;\n let ch = coords.u;\n\n let dyCorner = vec3(coords.y, coords.z, coords.w) - uniforms.pads;\n let dyDCorner = dyCorner.x;\n let dyRCorner = dyCorner.y;\n let dyCCorner = dyCorner.z;\n\n // Convolve dy(?, ?, ?, d) with pos mask(:, :, :, ch) to get\n // dx(xD, xR, xC, ch).\n // ? = to be determined. : = across all values in that axis.\n var dotProd = 0.0;\n for (var wD = 0; wD < uniforms.filterDims[0]; wD++) {\n let dyD = f32(dyDCorner + wD) / f32(uniforms.strides[0]);\n\n if (dyD < 0.0 || dyD >= f32(uniforms.outDepth) || fract(dyD) > 0.0) {\n continue;\n }\n let idyD = i32(dyD);\n\n for (var wR = 0; wR < uniforms.filterDims[1]; wR++) {\n let dyR = f32(dyRCorner + wR) / f32(uniforms.strides[1]);\n\n if (dyR < 0.0 || dyR >= f32(uniforms.outHeight) || fract(dyR) > 0.0) {\n continue;\n }\n let idyR = i32(dyR);\n\n for (var wC = 0; wC < uniforms.filterDims[2]; wC++) {\n let dyC = f32(dyCCorner + wC) / f32(uniforms.strides[2]);\n\n if (dyC < 0.0 || dyC >= f32(uniforms.outWidth) || fract(dyC) > 0.0) {\n continue;\n }\n let idyC = i32(dyC);\n\n let dyValue = getDy(batch, idyD, idyR, idyC, ch);\n dotProd += dyValue * uniforms.avgMultiplier;\n }\n }\n }\n setOutputAtIndex(index, dotProd);\n }\n }\n "); return userCode; }; return AvgPool3DBackpropProgram; }()); /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function avgPool3DGrad(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var dy = inputs.dy, input = inputs.input; var x = input; var filterSize = attrs.filterSize, strides = attrs.strides, pad = attrs.pad, dimRoundingMode = attrs.dimRoundingMode; var convInfo = tf.backend_util.computePool3DInfo(x.shape, filterSize, strides, 1 /* dilations */, pad, dimRoundingMode); var program = new AvgPool3DBackpropProgram(convInfo); var avgMultiplier = 1 / (convInfo.filterDepth * convInfo.filterHeight * convInfo.filterWidth); var uniformData = [ { type: 'int32', data: [convInfo.strideDepth, convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [ convInfo.effectiveFilterDepth - 1 - convInfo.padInfo.front, convInfo.effectiveFilterHeight - 1 - convInfo.padInfo.top, convInfo.effectiveFilterWidth - 1 - convInfo.padInfo.left ] }, { type: 'int32', data: [ convInfo.effectiveFilterDepth, convInfo.effectiveFilterHeight, convInfo.effectiveFilterWidth ] }, { type: 'int32', data: [convInfo.outDepth] }, { type: 'int32', data: [convInfo.outHeight] }, { type: 'int32', data: [convInfo.outWidth] }, { type: 'float32', data: [avgMultiplier] } ]; return backend.runWebGPUProgram(program, [dy], x.dtype, uniformData); } var avgPool3DGradConfig = { kernelName: tf.AvgPool3DGrad, backendName: 'webgpu', kernelFunc: avgPool3DGrad }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function avgPoolGrad(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var dy = inputs.dy, input = inputs.input; var x = input; assertNotComplex([dy, input], 'avgPoolGrad'); var filterSize = attrs.filterSize, strides = attrs.strides, pad = attrs.pad; var convInfo = tf.backend_util.computePool2DInfo(x.shape, filterSize, strides, 1 /* dilations */, pad); var program = new AvgPool2DBackpropProgram(convInfo); var avgMultiplier = 1 / (convInfo.filterHeight * convInfo.filterWidth); var uniformData = [ { type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [ convInfo.effectiveFilterHeight - 1 - convInfo.padInfo.top, convInfo.effectiveFilterWidth - 1 - convInfo.padInfo.left ] }, { type: 'int32', data: [convInfo.dilationHeight, convInfo.dilationWidth] }, { type: 'int32', data: [convInfo.effectiveFilterHeight, convInfo.effectiveFilterWidth] }, { type: 'int32', data: [convInfo.outHeight] }, { type: 'int32', data: [convInfo.outWidth] }, { type: 'float32', data: [avgMultiplier] } ]; return backend.runWebGPUProgram(program, [dy], x.dtype, uniformData); } var avgPoolGradConfig = { kernelName: tf.AvgPoolGrad, backendName: 'webgpu', kernelFunc: avgPoolGrad }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function batchMatMul(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var a = inputs.a, b = inputs.b; var transposeA = attrs.transposeA, transposeB = attrs.transposeB; return batchMatMulImpl({ a: a, b: b, transposeA: transposeA, transposeB: transposeB, backend: backend }); } var batchMatMulConfig = { kernelName: tf.BatchMatMul, backendName: 'webgpu', kernelFunc: batchMatMul, }; /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var SliceProgram = /** @class */ (function () { function SliceProgram(start, destSize) { this.variableNames = ['source']; this.workPerThread = 1; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = destSize; this.rank = destSize.length; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize, [this.workPerThread, 1, 1]); this.start = start; this.uniforms = "start : ".concat(getCoordsDataType(start.length), ", "); this.shaderKey = 'slice'; } SliceProgram.prototype.getUserCode = function () { var dtype = getCoordsDataType(this.rank); var sourceCoords = getCoords$1(this.rank); var coordSum; if (this.start.length === 1) { coordSum = this.outputShape.map(function (_, i) { return "sourceLoc = uniforms.start + coords;"; }); } else { coordSum = this.outputShape.map(function (_, i) { return "sourceLoc.".concat(coords[i], " = uniforms.start.").concat(getCoordsXYZ(i), " + coords.").concat(coords[i], ";"); }); } var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n var sourceLoc : ").concat(dtype, ";\n let coords = getCoordsFromIndex(index);\n ").concat(coordSum.join('\n'), "\n setOutputAtIndex(index, getSource(").concat(sourceCoords, "));\n }\n }\n "); return userCode; }; return SliceProgram; }()); var coords = ['x', 'y', 'z', 'w', 'u', 'v']; function getCoords$1(rank) { if (rank === 1) { return 'sourceLoc'; } else if (rank <= 6) { return coords.slice(0, rank).map(function (coord) { return "sourceLoc.".concat(coord); }).join(','); } else { throw Error("Slicing for rank ".concat(rank, " is not yet supported")); } } function slice(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var begin = attrs.begin, size = attrs.size; var _a = __read(tf.slice_util.parseSliceParams(x, begin, size), 2), $begin = _a[0], $size = _a[1]; tf.slice_util.assertParamsValid(x, $begin, $size); if (backend.shouldExecuteOnCPU([x]) || x.dtype === 'string') { var xTensorData = backend.tensorMap.get(x.dataId); var outValues = sliceImplCPU(xTensorData.values, $begin, $size, x.shape, x.dtype); return backend.makeTensorInfo($size, x.dtype, outValues); } if (tf.util.sizeFromShape($size) === 0) { return backend.makeTensorInfo($size, x.dtype, []); } // TODO(xing.xu): Add shadow slice support. var program = new SliceProgram($begin, $size); var uniformData = [{ type: 'int32', data: $begin }]; return backend.runWebGPUProgram(program, [x], x.dtype, uniformData); } var sliceConfig = { kernelName: tf.Slice, backendName: 'webgpu', kernelFunc: slice }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var batchToSpaceND = function (args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var blockShape = attrs.blockShape, crops = attrs.crops; tf.util.assert(x.shape.length <= 4, function () { return 'batchToSpaceND for rank > 4 with a WebGPU backend not ' + 'implemented yet'; }); var prod = blockShape.reduce(function (a, b) { return a * b; }); var reshaped = tf.backend_util.getReshaped(x.shape, blockShape, prod); var permuted = tf.backend_util.getPermuted(reshaped.length, blockShape.length); var reshapedPermuted = tf.backend_util.getReshapedPermuted(x.shape, blockShape, prod); var sliceBeginCoords = tf.backend_util.getSliceBeginCoords(crops, blockShape.length); var sliceSize = tf.backend_util.getSliceSize(reshapedPermuted, crops, blockShape.length); var toDispose = []; var reshapedIntermediate = reshape({ inputs: { x: x }, backend: backend, attrs: { shape: reshaped } }); var transposedIntermediate = transpose({ inputs: { x: reshapedIntermediate }, backend: backend, attrs: { perm: permuted } }); var reshapedIntermediate2 = reshape({ inputs: { x: transposedIntermediate }, backend: backend, attrs: { shape: reshapedPermuted } }); var sliced = slice({ inputs: { x: reshapedIntermediate2 }, backend: backend, attrs: { begin: sliceBeginCoords, size: sliceSize } }); toDispose.push(reshapedIntermediate); toDispose.push(transposedIntermediate); toDispose.push(reshapedIntermediate2); toDispose.forEach(function (t) { return backend.disposeData(t.dataId); }); return sliced; }; var batchToSpaceNDConfig = { kernelName: tf.BatchToSpaceND, backendName: 'webgpu', kernelFunc: batchToSpaceND }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var writeSnippet = "\n fn bincount_write(index: i32, value: f32) {\n ".concat(atomicAddSnippet('&result[index]', 'value', 'float32'), "\n }\n"); var binaryWriteSnippet = "\n fn bincount_write(index: i32, value: f32) {\n atomicStore(&result[index], bitcast(value));\n }\n"; var BincountProgram = /** @class */ (function () { function BincountProgram(shape, hasWeights, binaryOutput) { if (binaryOutput === void 0) { binaryOutput = false; } this.outputShape = []; this.variableNames = ['x']; this.uniforms = 'binCountSize : i32,'; this.workgroupSize = [64, 1, 1]; this.atomic = true; this.hasWeights = true; this.binaryOutput = false; this.outputShape = shape; this.rank = shape.length; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.binaryOutput = binaryOutput; if (binaryOutput) { this.atomic = false; } this.hasWeights = hasWeights; if (this.hasWeights) { this.variableNames.push('w'); } this.shaderKey = "bincount_".concat(this.hasWeights, "_").concat(this.binaryOutput, "_").concat(this.rank); } BincountProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(this.binaryOutput ? binaryWriteSnippet : writeSnippet, "\n ").concat(getMainHeaderString('index'), " {\n ").concat(this.rank === 1 ? "if (index < uniforms.xShape) {\n let indexVal = i32(getX(index));\n if (indexVal < uniforms.binCountSize) {\n let value = ".concat(this.binaryOutput ? 1. : (this.hasWeights ? 'getW(index)' : '1.'), ";\n bincount_write(indexVal, value);\n }\n }") : "let coord = getCoordsFromIndex(index);\n if (coordsInBounds2D(coord, uniforms.xShape)) {\n let indexVal = i32(getX(coord[0], coord[1]));\n if (indexVal < uniforms.binCountSize) {\n let value = ".concat(this.binaryOutput ? 1. : (this.hasWeights ? 'getW(coord[0], coord[1])' : '1.'), ";\n bincount_write(coord.x * uniforms.binCountSize + indexVal, value);\n }\n }"), "\n }\n "); return userCode; }; return BincountProgram; }()); /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function bincount(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, weights = inputs.weights; var size = attrs.size; var xSize = tf.util.sizeFromShape(x.shape); var weightsSize = tf.util.sizeFromShape(weights.shape); var hasWeights = weightsSize > 0; var outputSize = [size]; var dtype = weights.dtype; var output = fill({ backend: backend, attrs: { shape: outputSize, value: 0, dtype: dtype } }); var program = new BincountProgram([xSize], hasWeights); var uniformData = [{ type: 'int32', data: [size] }]; var bincountInputs = hasWeights ? [x, weights] : [x]; var res = backend.runWebGPUProgram(program, bincountInputs, dtype, uniformData, output); return res; } var bincountConfig = { kernelName: tf.Bincount, backendName: 'webgpu', kernelFunc: bincount }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var BroadcastArgsProgram = /** @class */ (function () { function BroadcastArgsProgram(shape) { this.outputShape = []; this.variableNames = ['s0', 's1']; this.uniforms = 's0Size : i32, s1Size : i32, '; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = [shape]; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'broadcastArgs'; } BroadcastArgsProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n var s0 = 1.0;\n var s1 = 1.0;\n let indexS0 = index - uniforms.size + uniforms.s0Size;\n let indexS1 = index - uniforms.size + uniforms.s1Size;\n if (indexS0 >= 0) {\n s0 = getS0(indexS0);\n }\n if (indexS1 >= 0) {\n s1 = getS1(indexS1);\n }\n\n if (s0 == 1.0) {\n setOutputAtIndex(index, s1);\n } else if (s1 == 1.0) {\n setOutputAtIndex(index, s0);\n } else if (s0 != s1) {\n setOutputAtIndex(index, uniforms.NAN);\n } else {\n setOutputAtIndex(index, s0);\n }\n }\n }\n "); return userCode; }; return BroadcastArgsProgram; }()); /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function broadcastArgs(args) { var inputs = args.inputs, backend = args.backend; var s0 = inputs.s0, s1 = inputs.s1; if (backend.shouldExecuteOnCPU([s0, s1])) { var s0TensorInfo = backend.tensorMap.get(s0.dataId); var s1TensorInfo = backend.tensorMap.get(s1.dataId); var s0Vals = s0TensorInfo.values; var s1Vals = s1TensorInfo.values; var broadcastShape = tf.backend_util.assertAndGetBroadcastShape(Array.from(s0Vals), Array.from(s1Vals)); return backend.makeTensorInfo([broadcastShape.length], 'int32', Int32Array.from(broadcastShape)); } var s0Size = tf.util.sizeFromShape(s0.shape); var s1Size = tf.util.sizeFromShape(s1.shape); var outputSize = Math.max(s0Size, s1Size); var program = new BroadcastArgsProgram(outputSize); var uniformData = [{ type: 'int32', data: [s0Size] }, { type: 'int32', data: [s1Size] }]; return backend.runWebGPUProgram(program, [s0, s1], 'int32', uniformData); } var broadcastArgsConfig = { kernelName: tf.BroadcastArgs, backendName: 'webgpu', kernelFunc: broadcastArgs }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var notEqual = binaryKernelFunc({ opType: BinaryOpType.NOT_EQUAL, dtype: 'bool', cpuKernelImpl: notEqualImplCPU }); var notEqualConfig = { kernelName: tf.NotEqual, backendName: 'webgpu', kernelFunc: notEqual }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function real(args) { var inputs = args.inputs, backend = args.backend; var input = inputs.input; var inputData = backend.tensorMap.get(input.dataId); return identity({ inputs: { x: inputData.complexTensorInfos.real }, backend: backend }); } var realConfig = { kernelName: tf.Real, backendName: 'webgpu', kernelFunc: real }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function int(input, backend) { var program = new UnaryOpProgram(input.shape, UnaryOpType.TO_INT); var output = backend.runWebGPUProgram(program, [input], 'int32'); return { dataId: output.dataId, shape: output.shape, dtype: output.dtype }; } function cast(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var dtype = attrs.dtype; // Casting to complex64. if (dtype === 'complex64') { if (x.dtype === 'complex64') { return identity({ inputs: { x: x }, backend: backend }); } // TODO: Import kernel function once zeros is modularized. var zerosTensor = tf__namespace.zeros(x.shape); var floatX = cast({ inputs: { x: x }, backend: backend, attrs: { dtype: 'float32' } }); var result = complex({ inputs: { real: floatX, imag: zerosTensor }, backend: backend }); zerosTensor.dispose(); backend.disposeData(floatX.dataId); return result; } // Casting from complex64 if (x.dtype === 'complex64') { var realPart = real({ inputs: { input: x }, backend: backend }); var result = cast({ inputs: { x: realPart }, backend: backend, attrs: { dtype: dtype } }); backend.disposeData(realPart.dataId); return result; } if (!tf.util.hasEncodingLoss(x.dtype, dtype)) { // We don't change the underlying data, since we cast to higher // precision. var result = identity({ inputs: { x: x }, backend: backend }); return { dataId: result.dataId, shape: result.shape, dtype: dtype }; } if (backend.shouldExecuteOnCPU([x])) { var values = backend.tensorMap.get(x.dataId).values; var _a = __read(castImplCPU(values, x.shape, x.dtype, dtype), 3), resultShape = _a[0], resultType = _a[1], resultData = _a[2]; return backend.makeTensorInfo(resultShape, resultType, resultData); } if (dtype === 'int32') { return int(x, backend); } if (dtype === 'bool') { var zerosTensorInfo = backend.makeTensorInfo([], 'bool', tf.util.getTypedArrayFromDType('bool', 1)); var binaryInputs = { a: x, b: zerosTensorInfo }; var result = notEqual({ inputs: binaryInputs, backend: backend }); backend.disposeData(zerosTensorInfo.dataId); return result; } throw new Error("Error in Cast: failed to cast ".concat(x.dtype, " to ").concat(dtype)); } var castConfig = { kernelName: tf.Cast, backendName: 'webgpu', kernelFunc: cast }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var ceil = unaryKernelFunc({ opType: UnaryOpType.CEIL, cpuKernelImpl: ceilImplCPU }); var ceilConfig = { kernelName: tf.Ceil, backendName: 'webgpu', kernelFunc: ceil }; /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var ClipVec4Program = /** @class */ (function () { function ClipVec4Program(outputShape) { this.variableNames = ['A']; this.uniforms = 'minVal : f32, maxVal : f32,'; this.workPerThread = 4; this.workgroupSize = [64, 1, 1]; this.outputComponent = 4; this.size = true; this.outputShape = outputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize, [this.workPerThread, 1, 1]); this.shaderKey = 'clipVec4'; } ClipVec4Program.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if(index < uniforms.size) {\n let value = getAByOutputIndex(index);\n var clampedValue = clamp(\n value, vec4(uniforms.minVal), vec4(uniforms.maxVal));\n clampedValue = select(clampedValue, value, isnanVec4(value));\n setOutputAtIndex(index, clampedValue);\n }\n }\n "); return userCode; }; return ClipVec4Program; }()); /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var ClipProgram = /** @class */ (function () { function ClipProgram(outputShape) { this.variableNames = ['A']; this.uniforms = 'minVal : f32, maxVal : f32,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = outputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'clip'; } ClipProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if(index < uniforms.size) {\n let value = getAByOutputIndex(index);\n if (isnan(value)) {\n setOutputAtIndex(index, value);\n return;\n }\n setOutputAtIndex(index, clamp(value, uniforms.minVal, uniforms.maxVal));\n }\n }\n "); return userCode; }; return ClipProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function clipByValue(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var clipValueMin = attrs.clipValueMin, clipValueMax = attrs.clipValueMax; var program; var uniformData = [ { type: 'float32', data: [clipValueMin] }, { type: 'float32', data: [clipValueMax] } ]; if (tf.util.sizeFromShape(x.shape) % 4 === 0) { program = new ClipVec4Program(x.shape); } else { program = new ClipProgram(x.shape); } return backend.runWebGPUProgram(program, [x], x.dtype, uniformData); } var clipByValueConfig = { kernelName: tf.ClipByValue, backendName: 'webgpu', kernelFunc: clipByValue }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var ComplexAbsProgram = /** @class */ (function () { function ComplexAbsProgram(shape) { this.outputShape = []; this.variableNames = ['real', 'imag']; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = shape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'complexAbs'; } ComplexAbsProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let re = abs(getRealByOutputIndex(index));\n let im = abs(getImagByOutputIndex(index));\n let mx = max(re, im);\n\n // The length function in wgsl may be not underflow-safe on some GPUs.\n // So the safe solution is to ensure underflow-safety in all cases.\n setOutputAtIndex(index, select(mx * length(vec2(1, min(re, im)/mx)), 0.0, mx == 0.0));\n }\n }\n "); return userCode; }; return ComplexAbsProgram; }()); /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ // Returns a TensorInfo with the complex shape and the dataId of the // underlying part. We need to do this because a reshaped complex tensor is // not reflected in its parts. function makeComplexComponentTensorInfo(complexTensor, complexPart) { return { dataId: complexPart.dataId, dtype: complexPart.dtype, shape: complexTensor.shape }; } function complexAbs(args) { var inputs = args.inputs, backend = args.backend; var x = inputs.x; var xData = backend.tensorMap.get(x.dataId); var program = new ComplexAbsProgram(x.shape); var programInputs = [ makeComplexComponentTensorInfo(x, xData.complexTensorInfos.real), makeComplexComponentTensorInfo(x, xData.complexTensorInfos.imag), ]; return backend.runWebGPUProgram(program, programInputs, programInputs[0].dtype); } var complexAbsConfig = { kernelName: tf.ComplexAbs, backendName: 'webgpu', kernelFunc: complexAbs }; /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var ConcatProgram = /** @class */ (function () { function ConcatProgram(shapes) { this.uniforms = ''; this.workPerThread = 1; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = tf.backend_util.computeOutShape(shapes, 1 /* axis */); this.variableNames = shapes.map(function (_, i) { return "T".concat(i); }); this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize, [this.workPerThread, 1, 1]); this.offsetLength = shapes.length - 1; for (var i = 0; i < this.offsetLength; i++) { this.uniforms += "offset".concat(i, " : i32,"); } this.shaderKey = 'concat'; } ConcatProgram.prototype.getUserCode = function () { var snippets = []; if (this.offsetLength > 0) { snippets.push("if (yC < uniforms.offset0){ setOutputAtCoords(coords.x, coords.y, getT0(yR, yC)); }"); for (var i = 1; i < this.offsetLength; i++) { snippets.push("else if (yC < uniforms.offset".concat([i], "){ ") + "setOutputAtCoords(coords.x, coords.y, getT".concat(i, "(yR, yC - uniforms.offset").concat(i - 1, ")); }")); } var lastIndex = this.offsetLength; var lastShiftIndex = this.offsetLength - 1; snippets.push("else { setOutputAtCoords(coords.x, coords.y, getT".concat(lastIndex, "(yR, yC - uniforms.offset").concat(lastShiftIndex, ")); }")); } else { snippets.push("setOutputAtCoords(coords.x, coords.y, getT0(yR, yC));"); } var userCode = "\n ".concat(getMainHeaderString('index'), " {\n for(var i = 0; i < ").concat(this.workPerThread, "; i = i + 1) {\n let flatIndex = index * ").concat(this.workPerThread, " + i;\n if(flatIndex < uniforms.size) {\n let coords = getCoordsFromIndex(flatIndex);\n let yR = coords.x;\n let yC = coords.y;\n\n ").concat(snippets.join('\n '), "\n }\n }\n }\n "); return userCode; }; return ConcatProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function imag(args) { var inputs = args.inputs, backend = args.backend; var input = inputs.input; var inputData = backend.tensorMap.get(input.dataId); return identity({ inputs: { x: inputData.complexTensorInfos.imag }, backend: backend }); } var imagConfig = { kernelName: tf.Imag, backendName: 'webgpu', kernelFunc: imag }; function concatImpl(inputs, axis, backend) { var e_1, _a; var dtype = inputs[0].dtype; if (dtype === 'complex64') { var reals = inputs.map(function (t) { return real({ inputs: { input: t }, backend: backend }); }); var imags = inputs.map(function (t) { return imag({ inputs: { input: t }, backend: backend }); }); var realConcated = concatImpl(reals, axis, backend); var imagConcated = concatImpl(imags, axis, backend); var result = complex({ inputs: { real: realConcated, imag: imagConcated }, backend: backend }); reals.forEach(function (r) { return backend.disposeData(r.dataId); }); imags.forEach(function (i) { return backend.disposeData(i.dataId); }); backend.disposeData(realConcated.dataId); backend.disposeData(imagConcated.dataId); return result; } var runOnCpu = backend.shouldExecuteOnCPU(inputs); // Run on cpu if dtype is string. For string, the backend represents it // as Uint8Array[], where each Uint8Array is a character. Given that the // computation is only on the outer array, uploading the whole data onto // gpu is wasteful. Also, currently webgpu doesn't have a design to // upload and retrieve Uint8Array[] between cpu and gpu. Therefore, we // just run the kernel on cpu if dtype is string. if (dtype === 'string') { runOnCpu = true; } if (runOnCpu) { // Any concat of n-dimensional tensors across any axis can be reduced to // a concatenation of two-dimensional tensors across the axis 1 by first // partitioning the axes of the original tensors into those less than the // axis to be concatenated and the rest. Then reshape the tensors // into a two-dimensional tensor by collapsing these two sets of axes and // concatenate the resulting matrices across the axis 1, finally reshaping // the result to have the proper shape. var tensors2D_1 = inputs.map(function (t) { var innerSize = tf.util.sizeFromShape(t.shape.slice(axis)); var shape = [-1, innerSize]; return reshape({ inputs: { x: t }, backend: backend, attrs: { shape: shape } }); }); var inputsValShapes = tensors2D_1.map(function (t) { return { vals: backend.readSync(t.dataId), shape: t.shape }; }); // Concats 2d tensors along axis=1. var outShape_1 = tf.backend_util.computeOutShape(tensors2D_1.map(function (t) { return t.shape; }), 1 /* axis */); var simplyConcat = tensors2D_1[0].shape[0] === 1; var outVals = concatImplCPU(inputsValShapes, outShape_1, dtype, simplyConcat); var finalOutShape = tf.backend_util.computeOutShape(inputs.map(function (t) { return t.shape; }), axis); var outInfo = backend.makeTensorInfo(finalOutShape, dtype, outVals); tensors2D_1.forEach(function (t) { return backend.disposeData(t.dataId); }); return outInfo; } // There is a storage buffer limitation in compute stage, one for output so // the maximum for input is limits.maxStorageBuffersPerShaderStage - 1 var maxInputNum = backend.device.limits.maxStorageBuffersPerShaderStage - 1; if (inputs.length > maxInputNum) { var reducedInputs = []; for (var i = 0; i < inputs.length; i += maxInputNum) { var subArray = inputs.slice(i, i + maxInputNum); reducedInputs.push(concatImpl(subArray, axis, backend)); } var result = concatImpl(reducedInputs, axis, backend); try { for (var reducedInputs_1 = __values(reducedInputs), reducedInputs_1_1 = reducedInputs_1.next(); !reducedInputs_1_1.done; reducedInputs_1_1 = reducedInputs_1.next()) { var i = reducedInputs_1_1.value; backend.disposeData(i.dataId); } } catch (e_1_1) { e_1 = { error: e_1_1 }; } finally { try { if (reducedInputs_1_1 && !reducedInputs_1_1.done && (_a = reducedInputs_1.return)) _a.call(reducedInputs_1); } finally { if (e_1) throw e_1.error; } } return result; } var _b = computeTensors2D(inputs, axis, backend), tensors2D = _b.tensors2D, outShape = _b.outShape; var shapes = (tensors2D).map(function (t) { return t.shape; }); var program = new ConcatProgram(shapes); var uniformData = []; var offsets = new Array(shapes.length - 1); if (offsets.length > 0) { offsets[0] = shapes[0][1]; uniformData.push({ type: 'int32', data: [offsets[0]] }); for (var i = 1; i < offsets.length; i++) { offsets[i] = offsets[i - 1] + shapes[i][1]; uniformData.push({ type: 'int32', data: [offsets[i]] }); } } var res = backend.runWebGPUProgram(program, tensors2D, tensors2D[0].dtype, uniformData); tensors2D.forEach(function (r) { return backend.disposeData(r.dataId); }); var reshapedResult = reshape({ inputs: { x: res }, backend: backend, attrs: { shape: outShape } }); backend.disposeData(res.dataId); return reshapedResult; } function computeTensors2D(inputs, axis, backend) { var outShape = tf.backend_util.computeOutShape(inputs.map(function (t) { return t.shape; }), axis); var tensors2D = inputs.map(function (t) { return reshape({ inputs: { x: t }, backend: backend, attrs: { shape: [ tf.util.sizeFromShape(t.shape.slice(0, axis)), tf.util.sizeFromShape(t.shape.slice(axis)) ] } }); }); return { tensors2D: tensors2D, outShape: outShape }; } /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function concat(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var axis = attrs.axis; var $axis = tf.util.parseAxisParam(axis, inputs[0].shape)[0]; var shapes = inputs.map(function (t) { return t.shape; }); tf.backend_util.assertParamsConsistent(shapes, $axis); var outShape = tf.backend_util.computeOutShape(inputs.map(function (t) { return t.shape; }), $axis); if (tf.util.sizeFromShape(outShape) === 0) { return backend.makeTensorInfo(outShape, inputs[0].dtype, []); } // Keep only non-empty tensors (ignore tensors with 0 in their shape). var $inputs = inputs.filter(function (t) { return tf.util.sizeFromShape(t.shape) > 0; }); if ($inputs.length === 1) { return identity({ inputs: { x: $inputs[0] }, backend: backend }); } return concatImpl($inputs, $axis, backend); } var concatConfig = { kernelName: tf.Concat, backendName: 'webgpu', kernelFunc: concat }; /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function conv2dCommonSnippet(isChannelsLast, fitAOuter, fitBOuter, fitInner, addBias, activation, hasPreluActivationWeights, innerElementSizeX, innerElementSizeW, innerElementSize) { if (addBias === void 0) { addBias = false; } if (activation === void 0) { activation = null; } if (hasPreluActivationWeights === void 0) { hasPreluActivationWeights = false; } if (innerElementSizeX === void 0) { innerElementSizeX = 4; } if (innerElementSizeW === void 0) { innerElementSizeW = 4; } if (innerElementSize === void 0) { innerElementSize = 4; } var getXSnippet = function (innerElementSize) { switch (innerElementSize) { case 1: return 'resData = f32(x[xIndex]);'; case 3: return 'resData = vec3(x[xIndex], x[xIndex + 1], x[xIndex + 2]);'; case 4: return 'resData = vec4(x[xIndex / 4]);'; default: throw new Error("innerElementSize ".concat(innerElementSize, " is not supported.")); } }; var getWSnippet = function (innerElementSize) { switch (innerElementSize) { case 1: return 'return f32(W[row * uniforms.wShape[3] + col]);'; case 4: return 'return vec4(W[(row * uniforms.wShape[3] + col) / 4]);'; default: throw new Error("innerElementSize ".concat(innerElementSize, " is not supported.")); } }; var coordASnippet = isChannelsLast ? "\n let coord = vec4(batch, xRow, xCol, xCh);\n " : "\n let coord = vec4(batch, xCh, xRow, xCol);\n "; var coordResSnippet = isChannelsLast ? "\n let coords = vec4(\n batch,\n row / outWidth,\n row % outWidth,\n col);\n " : "\n let coords = vec4(\n batch,\n row,\n col / outWidth,\n col % outWidth);\n "; var xHight = isChannelsLast ? 'uniforms.xShape[1]' : 'uniforms.xShape[2]'; var xWidth = isChannelsLast ? 'uniforms.xShape[2]' : 'uniforms.xShape[3]'; var row = isChannelsLast ? 'row' : 'col'; var col = isChannelsLast ? 'col' : 'row'; var readXSnippet = "\n let inChannels = uniforms.wShape[2];\n let outWidth = ".concat(isChannelsLast ? 'uniforms.outShape[2]' : 'uniforms.outShape[3]', ";\n let outRow = ").concat(row, " / outWidth;\n let outCol = ").concat(row, " % outWidth;\n\n let WRow = ").concat(col, " / (uniforms.filterDims[1] * inChannels);\n let WCol = ").concat(col, " / inChannels % uniforms.filterDims[1];\n let xRow = outRow * uniforms.strides[0] + uniforms.dilations[0] * WRow - uniforms.pads[0];\n let xCol = outCol * uniforms.strides[1] + uniforms.dilations[1] * WCol - uniforms.pads[1];\n let xCh = ").concat(col, " % inChannels;\n var resData = ").concat(typeSnippet(innerElementSizeX), "(0.0);\n // The bounds checking is always needed since we use it to pad zero for\n // the 'same' padding type.\n if (xRow >= 0 && xRow < ").concat(xHight, " && xCol >= 0 && xCol < ").concat(xWidth, ") {\n ").concat(coordASnippet, "\n let xIndex = getIndexFromCoords4D(coord, uniforms.xShape);\n ").concat(getXSnippet(innerElementSizeX), "\n }\n return resData;"); var sampleX = isChannelsLast ? (fitAOuter && fitInner ? "\n ".concat(readXSnippet) : "\n if (row < uniforms.dimAOuter && col < uniforms.dimInner) {\n ".concat(readXSnippet, "\n }\n return ").concat(typeSnippet(innerElementSizeX), "(0.0);")) : (fitInner && fitBOuter ? "\n ".concat(readXSnippet) : "\n if (row < uniforms.dimInner && col < uniforms.dimBOuter) {\n ".concat(readXSnippet, "\n }\n return ").concat(typeSnippet(innerElementSizeX), "(0.0);")); var sampleW = "".concat(getWSnippet(innerElementSizeW)); var resType = typeSnippet(innerElementSize); var aType = isChannelsLast ? typeSnippet(innerElementSizeX) : typeSnippet(innerElementSizeW); var bType = isChannelsLast ? typeSnippet(innerElementSizeW) : typeSnippet(innerElementSizeX); var userCode = "\n ".concat(activationFnSnippet(activation, hasPreluActivationWeights, innerElementSize === 4, 4), "\n fn mm_readA(batch: i32, row : i32, col : i32) -> ").concat(aType, " {\n ").concat(isChannelsLast ? sampleX : sampleW, "\n }\n\n fn mm_readB(batch: i32, row : i32, col : i32) -> ").concat(bType, " {\n ").concat(isChannelsLast ? sampleW : sampleX, "\n }\n\n fn mm_write(batch: i32, row : i32, col : i32, valueIn : ").concat(resType, ") {\n if (row < uniforms.dimAOuter && col < uniforms.dimBOuter)\n {\n var value = valueIn;\n let outWidth = ").concat(isChannelsLast ? 'uniforms.outShape[2]' : 'uniforms.outShape[3]', ";\n ").concat(coordResSnippet, "\n ").concat(biasActivationSnippet(addBias, activation), "\n setOutputAtCoords(coords[0], coords[1], coords[2], coords[3], value);\n }\n }"); return userCode; } var Conv2DMMProgram = /** @class */ (function () { function Conv2DMMProgram(convInfo, dimAOuter, dimBOuter, dimInner, addBias, activation, hasPreluActivationWeights, sequentialAccessByThreads) { if (addBias === void 0) { addBias = false; } if (activation === void 0) { activation = null; } if (hasPreluActivationWeights === void 0) { hasPreluActivationWeights = false; } if (sequentialAccessByThreads === void 0) { sequentialAccessByThreads = false; } this.variableNames = ['x', 'W']; this.uniforms = "filterDims : vec2, pads : vec2, strides : vec2, dilations : vec2, dimAOuter : i32, dimBOuter : i32, dimInner : i32,"; this.outputShape = convInfo.outShape; this.isChannelsLast = convInfo.dataFormat === 'channelsLast'; this.isVec4 = (((convInfo.inChannels % 4 === 0 || convInfo.inChannels % 3 === 0) && this.isChannelsLast) || (convInfo.outWidth % 4 === 0 && !this.isChannelsLast)) && convInfo.outChannels % 4 === 0; this.dispatchLayout = this.isChannelsLast ? { x: [3], y: [1, 2], z: [0] } : { x: [2, 3], y: [1], z: [0] }; this.workgroupSize = computeWorkgroupSizeForConv2d(this.dispatchLayout, this.outputShape, this.isVec4); this.elementsPerThread = computeWorkPerThreadForConv2d(this.dispatchLayout, this.outputShape, this.isVec4); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize, this.elementsPerThread); if (this.isVec4) { this.outputComponent = 4; if (this.isChannelsLast && convInfo.inChannels % 4 !== 0) { this.innerElementSize = 3; this.variableComponents = [1, 4]; } else { this.innerElementSize = 4; this.variableComponents = [4, 4]; } if (addBias) { this.variableNames.push('bias'); this.variableComponents.push(4); } if (hasPreluActivationWeights) { this.variableNames.push('preluActivationWeights'); this.variableComponents.push(4); } } else { this.innerElementSize = this.elementsPerThread[0]; if (addBias) { this.variableNames.push('bias'); } if (hasPreluActivationWeights) { this.variableNames.push('preluActivationWeights'); } } this.sequentialAccessByThreads = sequentialAccessByThreads; this.addBias = addBias; this.activation = activation; this.hasPreluActivationWeights = hasPreluActivationWeights; this.tileAOuter = this.workgroupSize[1] * this.elementsPerThread[1]; this.tileBOuter = this.workgroupSize[0] * this.elementsPerThread[0]; this.tileInner = Math.max(this.workgroupSize[0] * this.innerElementSize, this.workgroupSize[1]); this.fitAOuter = dimAOuter % this.tileAOuter === 0; this.fitBOuter = dimBOuter % this.tileBOuter === 0; this.fitInner = dimInner % this.tileInner === 0; this.shaderKey = "conv2DMM_".concat(this.elementsPerThread, "_").concat(this.activation, "}_").concat(this.fitAOuter, "_").concat(this.fitBOuter, "_").concat(this.fitInner, "_").concat(this.isVec4, "_").concat(this.innerElementSize, "_").concat(this.isChannelsLast, "_").concat(this.sequentialAccessByThreads); } Conv2DMMProgram.prototype.getUserCode = function () { var matMulSource = this.isVec4 ? makeMatMulPackedVec4Source(this.elementsPerThread, this.workgroupSize, !this.isChannelsLast, this.tileInner) : makeMatMulPackedSource(this.elementsPerThread, this.workgroupSize, !this.isChannelsLast, this.tileInner, false, null, this.sequentialAccessByThreads); var elementsSize = this.isVec4 ? [this.innerElementSize, 4, 4] : [1, 1, 1]; var userCode = "\n ".concat(conv2dCommonSnippet(this.isChannelsLast, this.fitAOuter, this.fitBOuter, this.fitInner, this.addBias, this.activation, this.hasPreluActivationWeights, elementsSize[0], elementsSize[1], elementsSize[2]), "\n ").concat(matMulSource, "\n "); return userCode; }; return Conv2DMMProgram; }()); /** * @license * Copyright 2022 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var Conv2DNaiveProgram = /** @class */ (function () { function Conv2DNaiveProgram(convInfo, addBias, activation, hasPreluActivationWeights) { if (addBias === void 0) { addBias = false; } if (activation === void 0) { activation = null; } if (hasPreluActivationWeights === void 0) { hasPreluActivationWeights = false; } this.variableNames = ['x', 'W']; this.uniforms = 'filterDims: vec2, pads: vec2, strides: vec2, dilations: vec2,'; this.workgroupSize = [4, 4, 8]; this.outputShape = convInfo.outShape; this.isChannelsLast = convInfo.dataFormat === 'channelsLast'; this.dispatchLayout = this.isChannelsLast ? { x: [2], y: [1], z: [0, 3] } : { x: [3], y: [2], z: [0, 1] }; this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.addBias = addBias; this.activation = activation; this.hasPreluActivationWeights = hasPreluActivationWeights; if (addBias) { this.variableNames.push('bias'); } if (hasPreluActivationWeights) { this.variableNames.push('preluActivationWeights'); } this.shaderKey = "conv2dnaive_".concat(this.activation, "_").concat(this.isChannelsLast); } Conv2DNaiveProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(activationFnSnippet(this.activation, this.hasPreluActivationWeights, false, 4), "\n fn readInp(batch : i32, row : i32, col : i32, chan : i32) -> f32{\n let coords = vec4(batch, row, col, chan);\n if (coordsInBounds4D(coords, uniforms.xShape)) {\n return getX(batch, row, col, chan);\n } else {\n return 0.0;\n }\n }\n fn readFilt(row : i32, col : i32, xChannel : i32, outChannel : i32) -> f32{\n let coords = vec4(row, col, xChannel, outChannel);\n if(coordsInBounds4D(coords, uniforms.wShape)) {\n return getW(row, col, xChannel, outChannel);\n } else {\n return 0.0;\n }\n }\n fn writeResult(batch : i32, row : i32, col : i32, chan : i32, valueIn : f32) {\n let coords = ").concat(this.isChannelsLast ? "vec4(batch, row, col, chan);" : "vec4(batch, chan, row, col);", "\n if (coordsInBounds4D(coords, uniforms.outShape)) {\n var value = valueIn;\n ").concat(biasActivationSnippet(this.addBias, this.activation), "\n setOutputAtCoords(coords.x, coords.y, coords.z, coords.w, value);\n }\n }\n ").concat(getMainHeaderString('index'), " {\n let coords = getOutputCoords();\n let batch = coords[0];\n let outChannel = ").concat(this.isChannelsLast ? "coords[3];" : "coords[1];", "\n let outRow = ").concat(this.isChannelsLast ? "coords[1];" : "coords[2];", "\n let outCol = ").concat(this.isChannelsLast ? "coords[2];" : "coords[3];", "\n var acc : f32 = 0.0;\n for (var row = 0; row < uniforms.filterDims[0]; row = row + 1) {\n for (var col = 0; col < uniforms.filterDims[1]; col = col + 1) {\n let xRow = outRow * uniforms.strides[0] + uniforms.dilations[0] * row - uniforms.pads[0];\n let xCol = outCol * uniforms.strides[1] + uniforms.dilations[1] * col - uniforms.pads[1];\n for (var xChannel = 0; xChannel < ").concat(this.isChannelsLast ? "uniforms.xShape[3];" : "uniforms.xShape[1];", " xChannel = xChannel + 1) {\n ").concat(this.isChannelsLast ? "let v = readInp(batch, xRow, xCol, xChannel);" : "let v = readInp(batch, xChannel, xRow, xCol);", "\n let f = readFilt(row, col, xChannel, outChannel);\n acc = acc + v * f;\n }\n }\n }\n writeResult(batch, outRow, outCol, outChannel, acc);\n }\n "); return userCode; }; return Conv2DNaiveProgram; }()); /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var Im2ColProgram = /** @class */ (function () { function Im2ColProgram(outputShape, isChannelsLast) { this.variableNames = ['x']; this.uniforms = "pads : vec2, strides : vec2, dilations : vec2, outWidth : i32, itemsPerBlockRow : i32,\n inChannels : i32,"; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = outputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.isChannelsLast = isChannelsLast; this.shaderKey = "im2col_".concat(this.isChannelsLast); } Im2ColProgram.prototype.getUserCode = function () { var rowDim = this.isChannelsLast ? 1 : 2; var colDim = this.isChannelsLast ? 2 : 3; var row = this.isChannelsLast ? 'coords[1]' : 'coords[2]'; var col = this.isChannelsLast ? 'coords[2]' : 'coords[1]'; var getXSnippet = this.isChannelsLast ? 'getX(batch, xRow, xCol, ch)' : 'getX(batch, ch, xRow, xCol)'; var userCode = "\n ".concat(getMainHeaderString('index'), " {\n let coords = getCoordsFromIndex(index);\n if(index < uniforms.size) {\n let batch = coords[0];\n let row = ").concat(row, ";\n let col = ").concat(col, ";\n let offsetY = (row / uniforms.outWidth) * uniforms.strides[0] - uniforms.pads[0];\n let xRow = offsetY + uniforms.dilations[0] * (col / uniforms.itemsPerBlockRow);\n var value = 0.0;\n if(xRow < uniforms.xShape[").concat(rowDim, "] && xRow >= 0) {\n let offsetX = (row % uniforms.outWidth) * uniforms.strides[1] -\n uniforms.pads[1];\n let xCol = offsetX + uniforms.dilations[1] * ((col %\n uniforms.itemsPerBlockRow) / uniforms.inChannels);\n let ch = col % uniforms.inChannels;\n if(xCol < uniforms.xShape[").concat(colDim, "] && xCol >= 0) {\n value = ").concat(getXSnippet, ";\n }\n }\n setOutputAtIndex(index, value);\n }\n }\n "); return userCode; }; return Im2ColProgram; }()); // conv2dByMatMul fuses height and width into one dimension to compute // batchMatMul, so bias and activation weights are also supposed to fuse the two // dimensions into one. // // This function computes the target shape for fusing height and width // dimensions. Returning null means the shape is already compatible. function getShapeForBatchMatMul(shape, isChannelsLast) { var length = shape.length; if (length >= 3) { return isChannelsLast ? __spreadArray(__spreadArray([], __read(shape.slice(0, -3) /* batch */), false), [ shape[length - 3] * shape[length - 2] /* height * width */, shape[length - 1] /* channel */ ], false) : __spreadArray(__spreadArray([], __read(shape.slice(0, -3) /* batch */), false), [ shape[length - 3] /* channel */, shape[length - 2] * shape[length - 1] /* height * width */ ], false); } else if (!isChannelsLast && length === 1 && shape[0] > 1) { return [shape[0], 1]; } else { return null; } } // For 1x1 kernels that iterate through every point in the input, convolution // can be expressed as matrix multiplication (without need for memory // remapping). function conv2dByMatMul(_a) { var e_1, _b; var x = _a.x, filter = _a.filter, convInfo = _a.convInfo, backend = _a.backend, _c = _a.bias, bias = _c === void 0 ? null : _c, _d = _a.preluActivationWeights, preluActivationWeights = _d === void 0 ? null : _d, _e = _a.leakyreluAlpha, leakyreluAlpha = _e === void 0 ? 0 : _e, _f = _a.activation, activation = _f === void 0 ? null : _f; var isChannelsLast = convInfo.dataFormat === 'channelsLast'; var transposeA = isChannelsLast ? false : true; var transposeB = false; var sameSize = isChannelsLast && convInfo.filterHeight === convInfo.inHeight && convInfo.filterWidth === convInfo.inWidth && convInfo.padInfo.type === 'VALID'; var intermediates = []; var xReshaped; var filterReshaped; if (sameSize) { var sharedDim = convInfo.inHeight * convInfo.inWidth * convInfo.inChannels; xReshaped = reshape({ inputs: { x: x }, backend: backend, attrs: { shape: [1, convInfo.batchSize, sharedDim] } }); filterReshaped = reshape({ inputs: { x: filter }, backend: backend, attrs: { shape: [1, sharedDim, convInfo.outChannels] } }); } else { xReshaped = reshape({ inputs: { x: x }, backend: backend, attrs: { shape: isChannelsLast ? [ convInfo.batchSize, convInfo.inHeight * convInfo.inWidth, convInfo.inChannels ] : [ convInfo.batchSize, convInfo.inChannels, convInfo.inHeight * convInfo.inWidth ] } }); filterReshaped = reshape({ inputs: { x: filter }, backend: backend, attrs: { shape: [1, convInfo.inChannels, convInfo.outChannels] } }); } intermediates.push(xReshaped); intermediates.push(filterReshaped); if (preluActivationWeights != null) { var targetShape = getShapeForBatchMatMul(preluActivationWeights.shape, isChannelsLast); if (targetShape != null) { preluActivationWeights = reshape({ inputs: { x: preluActivationWeights }, backend: backend, attrs: { shape: targetShape } }); intermediates.push(preluActivationWeights); } } if (bias != null) { var targetShape = getShapeForBatchMatMul(bias.shape, isChannelsLast); if (targetShape != null) { bias = reshape({ inputs: { x: bias }, backend: backend, attrs: { shape: targetShape } }); intermediates.push(bias); } } var result = batchMatMulImpl({ a: isChannelsLast ? xReshaped : filterReshaped, b: isChannelsLast ? filterReshaped : xReshaped, transposeA: transposeA, transposeB: transposeB, backend: backend, bias: bias, activation: activation, preluActivationWeights: preluActivationWeights, leakyreluAlpha: leakyreluAlpha }); var out = reshape({ inputs: { x: result }, backend: backend, attrs: { shape: convInfo.outShape } }); intermediates.push(result); try { for (var intermediates_1 = __values(intermediates), intermediates_1_1 = intermediates_1.next(); !intermediates_1_1.done; intermediates_1_1 = intermediates_1.next()) { var i = intermediates_1_1.value; backend.disposeData(i.dataId); } } catch (e_1_1) { e_1 = { error: e_1_1 }; } finally { try { if (intermediates_1_1 && !intermediates_1_1.done && (_b = intermediates_1.return)) _b.call(intermediates_1); } finally { if (e_1) throw e_1.error; } } return out; } // Implements the im2col algorithm as outlined in "High Performance // Convolutional Neural Networks for Document Processing" (Suvisoft, 2006) function conv2dWithIm2Col(_a) { var e_2, _b; var x = _a.x, filter = _a.filter, convInfo = _a.convInfo, backend = _a.backend, _c = _a.bias, bias = _c === void 0 ? null : _c, _d = _a.preluActivationWeights, preluActivationWeights = _d === void 0 ? null : _d, _e = _a.leakyreluAlpha, leakyreluAlpha = _e === void 0 ? 0 : _e, _f = _a.activation, activation = _f === void 0 ? null : _f; // Rearranges conv2d input so each block to be convolved over forms the // row of a new matrix with shape [outHeight * outWidth, // filterWidth * filterHeight * inChannels]. The filter is also rearranged so // each output channel forms a col of a new matrix with shape [ // filterWidth * filterHeight * inChannels, outChannels]. The convolution is // then computed by multiplying these matrices and reshaping the result. var filterWidth = convInfo.filterWidth, filterHeight = convInfo.filterHeight, inChannels = convInfo.inChannels, strideWidth = convInfo.strideWidth, strideHeight = convInfo.strideHeight, padInfo = convInfo.padInfo, outWidth = convInfo.outWidth, outHeight = convInfo.outHeight, dilationWidth = convInfo.dilationWidth, dilationHeight = convInfo.dilationHeight, dataFormat = convInfo.dataFormat; var isChannelsLast = dataFormat === 'channelsLast'; var sharedDim = filterWidth * filterHeight * inChannels; var numCols = outHeight * outWidth; var x2ColShape = isChannelsLast ? [convInfo.batchSize, numCols, sharedDim] : [convInfo.batchSize, sharedDim, numCols]; var im2ColProgram = new Im2ColProgram(x2ColShape, isChannelsLast); var dimensions = [ { type: 'int32', data: [padInfo.top, padInfo.left] }, { type: 'int32', data: [strideHeight, strideWidth] }, { type: 'int32', data: [dilationHeight, dilationWidth] }, { type: 'int32', data: [outWidth] }, { type: 'int32', data: [inChannels * filterWidth] }, { type: 'int32', data: [inChannels] } ]; var x2Col = backend.runWebGPUProgram(im2ColProgram, [x], x.dtype, dimensions); var intermediates = []; intermediates.push(x2Col); var filterReshaped = reshape({ inputs: { x: filter }, backend: backend, attrs: { shape: [1, sharedDim, -1] } }); intermediates.push(filterReshaped); if (preluActivationWeights != null) { var targetShape = getShapeForBatchMatMul(preluActivationWeights.shape, isChannelsLast); if (targetShape != null) { preluActivationWeights = reshape({ inputs: { x: preluActivationWeights }, backend: backend, attrs: { shape: targetShape } }); intermediates.push(preluActivationWeights); } } if (bias != null) { var targetShape = getShapeForBatchMatMul(bias.shape, isChannelsLast); if (targetShape != null) { bias = reshape({ inputs: { x: bias }, backend: backend, attrs: { shape: targetShape } }); intermediates.push(bias); } } var transposeA = isChannelsLast ? false : true; var transposeB = false; var result = batchMatMulImpl({ a: isChannelsLast ? x2Col : filterReshaped, b: isChannelsLast ? filterReshaped : x2Col, transposeA: transposeA, transposeB: transposeB, backend: backend, bias: bias, activation: activation, preluActivationWeights: preluActivationWeights, leakyreluAlpha: leakyreluAlpha }); var out = reshape({ inputs: { x: result }, backend: backend, attrs: { shape: convInfo.outShape } }); intermediates.push(result); try { for (var intermediates_2 = __values(intermediates), intermediates_2_1 = intermediates_2.next(); !intermediates_2_1.done; intermediates_2_1 = intermediates_2.next()) { var i = intermediates_2_1.value; backend.disposeData(i.dataId); } } catch (e_2_1) { e_2 = { error: e_2_1 }; } finally { try { if (intermediates_2_1 && !intermediates_2_1.done && (_b = intermediates_2.return)) _b.call(intermediates_2); } finally { if (e_2) throw e_2.error; } } return out; } function conv2DImpl(_a) { var e_3, _b; var x = _a.x, filter = _a.filter, convInfo = _a.convInfo, backend = _a.backend, _c = _a.bias, bias = _c === void 0 ? null : _c, _d = _a.preluActivationWeights, preluActivationWeights = _d === void 0 ? null : _d, _e = _a.leakyreluAlpha, leakyreluAlpha = _e === void 0 ? 0 : _e, _f = _a.activation, activation = _f === void 0 ? null : _f; var hasBias = bias != null; var hasPreluActivationWeights = preluActivationWeights != null; var isChannelsLast = convInfo.dataFormat === 'channelsLast'; var sameSize = isChannelsLast && convInfo.filterHeight === convInfo.inHeight && convInfo.filterWidth === convInfo.inWidth && convInfo.padInfo.type === 'VALID'; var useNaiveConv2d = tf.env().getBool('WEBGPU_USE_NAIVE_CONV2D_DEBUG'); if (!useNaiveConv2d && (sameSize || (convInfo.filterHeight === 1 && convInfo.filterWidth === 1 && convInfo.dilationHeight === 1 && convInfo.dilationWidth === 1 && convInfo.strideHeight === 1 && convInfo.strideWidth === 1 && (convInfo.padInfo.type === 'SAME' || convInfo.padInfo.type === 'VALID')))) { return conv2dByMatMul({ x: x, filter: filter, convInfo: convInfo, backend: backend, bias: bias, activation: activation, preluActivationWeights: preluActivationWeights, leakyreluAlpha: leakyreluAlpha }); } var thresholdFlagValue = tf.env().getNumber('WEBGPU_THRESHOLD_TO_INCREASE_WORKGROUPS_FOR_MATMUL'); var thresholdToIncreaseWorkgroups = thresholdFlagValue > -1 ? thresholdFlagValue : backend.thresholdToIncreaseWorkgroups; var workgroupsBy32x32 = convInfo.batchSize * Math.ceil((convInfo.outHeight * convInfo.outWidth) / 32) * Math.ceil(convInfo.outChannels / 32); if (tf.env().getBool('WEBGPU_CONV_SEPARATE_IM2COL_SHADER') || workgroupsBy32x32 <= thresholdToIncreaseWorkgroups) { return conv2dWithIm2Col({ x: x, filter: filter, convInfo: convInfo, backend: backend, bias: bias, preluActivationWeights: preluActivationWeights, leakyreluAlpha: leakyreluAlpha, activation: activation }); } var program; var padInfo = [convInfo.padInfo.top, convInfo.padInfo.left]; var dimensions = [ { type: 'int32', data: [convInfo.filterHeight, convInfo.filterWidth] }, { type: 'int32', data: __spreadArray([], __read(padInfo), false) }, { type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.dilationHeight, convInfo.dilationWidth] } ]; if (useNaiveConv2d) { program = new Conv2DNaiveProgram(convInfo, hasBias, activation, hasPreluActivationWeights); } else { var dimAOuter = isChannelsLast ? convInfo.outHeight * convInfo.outWidth : convInfo.outChannels; var dimBOuter = isChannelsLast ? convInfo.outChannels : convInfo.outHeight * convInfo.outWidth; var dimInner = convInfo.filterHeight * convInfo.filterWidth * convInfo.inChannels; dimensions.push({ type: 'int32', data: [dimAOuter] }, { type: 'int32', data: [dimBOuter] }, { type: 'int32', data: [dimInner] }); // Experiments show that sequential access is more friendly for Intel GPUs. var sequentialAccessByThreads = backend.adapterInfo.isIntel(); program = new Conv2DMMProgram(convInfo, dimAOuter, dimBOuter, dimInner, hasBias, activation, hasPreluActivationWeights, sequentialAccessByThreads); } var intermediates = []; var inputVar = [x, filter]; if (hasBias) { if (!isChannelsLast && bias.shape.length === 1) { bias = reshape({ inputs: { x: bias }, backend: backend, attrs: { shape: [bias.shape[0], 1, 1] } }); intermediates.push(bias); } inputVar.push(bias); } if (hasPreluActivationWeights) { if (!isChannelsLast && preluActivationWeights.shape.length === 1) { preluActivationWeights = reshape({ inputs: { x: preluActivationWeights }, backend: backend, attrs: { shape: [preluActivationWeights.shape[0], 1, 1] } }); intermediates.push(preluActivationWeights); } inputVar.push(preluActivationWeights); } if (activation === 'leakyrelu') { dimensions.push({ type: 'float32', data: [leakyreluAlpha] }); program.uniforms += ' alpha : f32,'; } var out = backend.runWebGPUProgram(program, inputVar, x.dtype, dimensions); try { for (var intermediates_3 = __values(intermediates), intermediates_3_1 = intermediates_3.next(); !intermediates_3_1.done; intermediates_3_1 = intermediates_3.next()) { var i = intermediates_3_1.value; backend.disposeData(i.dataId); } } catch (e_3_1) { e_3 = { error: e_3_1 }; } finally { try { if (intermediates_3_1 && !intermediates_3_1.done && (_b = intermediates_3.return)) _b.call(intermediates_3); } finally { if (e_3) throw e_3.error; } } return out; } /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function conv2d(args) { var inputs = args.inputs, attrs = args.attrs, backend = args.backend; var x = inputs.x, filter = inputs.filter; var strides = attrs.strides, pad = attrs.pad, dataFormat = attrs.dataFormat, dilations = attrs.dilations, dimRoundingMode = attrs.dimRoundingMode; var $dataFormat = tf.backend_util.convertConv2DDataFormat(dataFormat); var convInfo = tf.backend_util.computeConv2DInfo(x.shape, filter.shape, strides, dilations, pad, dimRoundingMode, false /* depthwise */, $dataFormat); return conv2DImpl({ x: x, filter: filter, convInfo: convInfo, backend: backend }); } var conv2DConfig = { kernelName: tf.Conv2D, backendName: 'webgpu', kernelFunc: conv2d }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var Conv2DDerInputProgram = /** @class */ (function () { function Conv2DDerInputProgram(convInfo) { this.variableNames = ['dy', 'W']; this.uniforms = 'filterDims : vec2, pads : vec2, strides : vec2, outBackprop : vec4,'; this.workgroupSize = [64, 1, 1]; this.size = false; this.isVec4 = false; this.workPerThread = 1; this.outputShape = convInfo.inShape; this.isChannelsLast = convInfo.dataFormat === 'channelsLast'; this.isVec4 = this.isChannelsLast && convInfo.outChannels % 4 === 0 && convInfo.inChannels % 4 === 0; if (this.isVec4) { // TODO: Expand to any value. this.workPerThread = 2; this.outputComponent = 4; this.workgroupSize = [4, 4, 4]; this.dispatchLayout = { x: [3], y: [2], z: [0, 1] }; this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize, [4, this.workPerThread, 1]); } else { this.size = true; this.workPerThread = 1; this.workgroupSize = [64, 1, 1]; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); } this.shaderKey = "conv2DDerInput_".concat(this.isChannelsLast, "_").concat(this.isVec4, "_").concat(this.workPerThread); } Conv2DDerInputProgram.prototype.getUserCode = function () { var rowDim = this.isChannelsLast ? 1 : 2; var colDim = this.isChannelsLast ? 2 : 3; var channelDim = this.isChannelsLast ? 3 : 1; var vec4Snippet = "\n ".concat(getMainHeaderString(), " {\n let batch = i32(globalId.z) / uniforms.outShape[1];\n let r = i32(globalId.z) % uniforms.outShape[1];\n let c = i32(globalId.y) * ").concat(this.workPerThread, ";\n let d1 = i32(globalId.x) * 4;\n\n let dyCorner = vec2(r, c) - uniforms.pads;\n\n // Convolve dy(?, ?, d2) with w(:, :, d1, d2) to compute dx(xR, xC, d1).\n // ? = to be determined. : = across all values in that axis.\n var dotProd: array, ").concat(this.workPerThread, ">;\n for (var i = 0; i < ").concat(this.workPerThread, "; i++) {\n dotProd[i] = vec4(0.0);\n }\n for (var wR = 0; wR < uniforms.filterDims.x; wR = wR + 1) {\n let dyR = f32(dyCorner.x + wR) / f32(uniforms.strides.x);\n let wRPerm = uniforms.filterDims.x - 1 - wR;\n if (dyR < 0.0 || dyR >= f32(uniforms.outBackprop[1]) ||\n fract(dyR) > 0.0) {\n continue;\n }\n let idyR = i32(dyR);\n\n for (var wC = 0; wC < uniforms.filterDims.y; wC = wC + 1) {\n let dyC = f32(dyCorner.y + wC) / f32(uniforms.strides.y);\n let dyC2 = f32(dyCorner.y + 1 + wC) / f32(uniforms.strides.y);\n let wCPerm = uniforms.filterDims.y - 1 - wC;\n var bDyCVal = true;\n var bDyCVal2 = true;\n if (dyC < 0.0 || dyC >= f32(uniforms.outBackprop[2]) ||\n fract(dyC) > 0.0) {\n bDyCVal = false;\n }\n if (dyC2 < 0.0 || dyC2 >= f32(uniforms.outBackprop[2]) ||\n fract(dyC2) > 0.0) {\n bDyCVal2 = false;\n }\n\n let idyC = i32(dyC);\n let idyC2 = i32(dyC2);\n if (bDyCVal && bDyCVal2) {\n let d2Length = uniforms.outBackprop[3];\n for (var d2 = 0; d2 < d2Length; d2 = d2 + 4) {\n let wValue0 = getW(wRPerm, wCPerm, d1, d2);\n let wValue1 = getW(wRPerm, wCPerm, d1 + 1, d2);\n let wValue2 = getW(wRPerm, wCPerm, d1 + 2, d2);\n let wValue3 = getW(wRPerm, wCPerm, d1 + 3, d2);\n var xValue = getDy(batch, idyR, idyC, d2);\n let tmpval = vec4(dot(xValue, wValue0),\n dot(xValue, wValue1),\n dot(xValue, wValue2),\n dot(xValue, wValue3));\n dotProd[0] = dotProd[0] + tmpval;\n xValue = getDy(batch, idyR, idyC2, d2);\n dotProd[1] = dotProd[1] + vec4(dot(xValue, wValue0),\n dot(xValue, wValue1),\n dot(xValue, wValue2),\n dot(xValue, wValue3));\n }\n } else if (bDyCVal) {\n let d2Length = uniforms.outBackprop[3];\n for (var d2 = 0; d2 < d2Length; d2 = d2 + 4) {\n let wValue0 = getW(wRPerm, wCPerm, d1, d2);\n let wValue1 = getW(wRPerm, wCPerm, d1 + 1, d2);\n let wValue2 = getW(wRPerm, wCPerm, d1 + 2, d2);\n let wValue3 = getW(wRPerm, wCPerm, d1 + 3, d2);\n var xValue = getDy(batch, idyR, idyC, d2);\n let tmpval = vec4(dot(xValue, wValue0),\n dot(xValue, wValue1),\n dot(xValue, wValue2),\n dot(xValue, wValue3));\n dotProd[0] = dotProd[0] + tmpval;\n }\n } else if (bDyCVal2) {\n let d2Length = uniforms.outBackprop[3];\n for (var d2 = 0; d2 < d2Length; d2 = d2 + 4) {\n let wValue0 = getW(wRPerm, wCPerm, d1, d2);\n let wValue1 = getW(wRPerm, wCPerm, d1 + 1, d2);\n let wValue2 = getW(wRPerm, wCPerm, d1 + 2, d2);\n let wValue3 = getW(wRPerm, wCPerm, d1 + 3, d2);\n var xValue = getDy(batch, idyR, idyC2, d2);\n let tmpval = vec4(dot(xValue, wValue0),\n dot(xValue, wValue1),\n dot(xValue, wValue2),\n dot(xValue, wValue3));\n dotProd[1] = dotProd[1] + tmpval;\n }\n }\n }\n }\n\n for (var i = 0; i < ").concat(this.workPerThread, "; i = i + 1) {\n let coords = vec4(batch, r, c + i, d1);\n if (coordsInBounds4D(coords, uniforms.outShape)) {\n setOutputAtCoords(coords[0], coords[1], coords[2], coords[3], dotProd[i]);\n }\n }\n }\n "); return this.isVec4 ? "\n ".concat(vec4Snippet, "\n ") : "\n ".concat(getMainHeaderString('index'), " {\n if(index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let batch = coords[0];\n let d1 = coords[").concat(channelDim, "];\n\n let dyCorner = vec2(coords[").concat(rowDim, "], coords[").concat(colDim, "]) - uniforms.pads;\n let dyRCorner = dyCorner.x;\n let dyCCorner = dyCorner.y;\n\n // Convolve dy(?, ?, d2) with w(:, :, d1, d2) to compute dx(xR, xC, d1).\n // ? = to be determined. : = across all values in that axis.\n var dotProd = 0.0;\n for (var wR = 0; wR < uniforms.filterDims.x; wR = wR + 1) {\n let dyR = (f32(dyRCorner) + f32(wR)) / f32(uniforms.strides.x);\n let wRPerm = uniforms.filterDims.x - 1 - wR;\n if (dyR < 0.0 || dyR >= f32(uniforms.outBackprop[1]) || fract(dyR) > 0.0 ||\n wRPerm < 0) {\n continue;\n }\n let idyR = i32(dyR);\n\n for (var wC = 0; wC < uniforms.filterDims.y; wC = wC + 1) {\n let dyC = (f32(dyCCorner) + f32(wC)) / f32(uniforms.strides.y);\n let wCPerm = uniforms.filterDims.y - 1 - wC;\n if (dyC < 0.0 || dyC >= f32(uniforms.outBackprop[2]) ||\n fract(dyC) > 0.0 || wCPerm < 0) {\n continue;\n }\n let idyC = i32(dyC);\n\n for (var d2 = 0; d2 < uniforms.outBackprop[3]; d2 = d2 + 1) {\n let xValue = ").concat(this.isChannelsLast ? 'getDy(batch, idyR, idyC, d2)' : 'getDy(batch, d2, idyR, idyC)', ";\n let wValue = getW(wRPerm, wCPerm, d1, d2);\n dotProd = dotProd + xValue * wValue;\n }\n }\n }\n setOutputAtIndex(index, dotProd);\n }\n }\n "); }; return Conv2DDerInputProgram; }()); var Conv2DDerFilterProgram = /** @class */ (function () { function Conv2DDerFilterProgram(convInfo) { this.variableNames = ['x', 'dy']; this.uniforms = 'pads : vec2, strides : vec2, batchSize : i32, outHeight : i32, outWidth : i32, inHeight : i32, inWidth : i32,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = convInfo.filterShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.isChannelsLast = convInfo.dataFormat === 'channelsLast'; this.shaderKey = "conv2DDerFilter_".concat(this.isChannelsLast); } Conv2DDerFilterProgram.prototype.getUserCode = function () { return "\n ".concat(getMainHeaderString('index'), " {\n if(index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let wR = coords[0];\n let wC = coords[1];\n let d1 = coords[2];\n let d2 = coords[3];\n\n // Convolve x(?, ?, d1) with dy(:, :, d2) to get dw(wR, wC, d1, d2).\n // ? = to be determined. : = across all values in that axis.\n var dotProd = 0.0;\n for (var b = 0; b < uniforms.batchSize; b = b + 1) {\n for (var yR = 0; yR < uniforms.outHeight; yR = yR + 1) {\n let xR = wR + yR * uniforms.strides[0] - uniforms.pads[0];\n if (xR < 0 || xR >= uniforms.inHeight) {\n continue;\n }\n\n for (var yC = 0; yC < uniforms.outWidth; yC = yC + 1) {\n let xC = wC + yC * uniforms.strides[1] - uniforms.pads[1];\n\n if (xC < 0 || xC >= uniforms.inWidth) {\n continue;\n }\n\n if (").concat(this.isChannelsLast, ") {\n let dyValue = getDy(b, yR, yC, d2);\n let xValue = getX(b, xR, xC, d1);\n dotProd = dotProd + xValue * dyValue;\n } else {\n let dyValue = getDy(b, d2, yR, yC);\n let xValue = getX(b, d1, xR, xC);\n dotProd = dotProd + xValue * dyValue;\n }\n }\n }\n }\n setOutputAtIndex(index, dotProd);\n }\n }\n "); }; return Conv2DDerFilterProgram; }()); var Conv3DDerFilterProgram = /** @class */ (function () { function Conv3DDerFilterProgram(convInfo) { this.variableNames = ['x', 'dy']; this.uniforms = "pads : vec3, strides : vec3, batchSize : i32, outDepth : i32,\n outHeight : i32, outWidth : i32, inDepth : i32, inHeight : i32, inWidth : i32,"; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = convInfo.filterShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = "conv3DDerFilter"; } Conv3DDerFilterProgram.prototype.getUserCode = function () { return "\n ".concat(getMainHeaderString('index'), " {\n if(index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let wF = coords.x;\n let wR = coords.y;\n let wC = coords.z;\n let d1 = coords.w;\n let d2 = coords.u;\n\n var dotProd = 0.0;\n for (var b = 0; b < uniforms.batchSize; b++) {\n for (var yF = 0; yF < uniforms.outDepth; yF++) {\n let xF = wF + yF * uniforms.strides[0] - uniforms.pads[0];\n if (xF < 0 || xF >= uniforms.inDepth) {\n continue;\n }\n\n for (var yR = 0; yR < uniforms.outHeight; yR++) {\n let xR = wR + yR * uniforms.strides[1] - uniforms.pads[1];\n if (xR < 0 || xR >= uniforms.inHeight) {\n continue;\n }\n\n for (var yC = 0; yC < uniforms.outWidth; yC++) {\n let xC = wC + yC * uniforms.strides[2] - uniforms.pads[2];\n if (xC < 0 || xC >= uniforms.inWidth) {\n continue;\n }\n\n let dyValue = getDy(b, yF, yR, yC, d2);\n let xValue = getX(b, xF, xR, xC, d1);\n dotProd += xValue * dyValue;\n }\n }\n }\n }\n setOutputAtIndex(index, dotProd);\n }\n }\n "); }; return Conv3DDerFilterProgram; }()); var Conv3DDerInputProgram = /** @class */ (function () { function Conv3DDerInputProgram(convInfo) { this.variableNames = ['dy', 'W']; this.uniforms = "filterDims : vec3, pads : vec3, strides : vec3,\n outDepth : i32, outHeight : i32, outWidth : i32, outChannels : i32,"; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = convInfo.inShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = "conv3DDerInput"; } Conv3DDerInputProgram.prototype.getUserCode = function () { return "\n ".concat(getMainHeaderString('index'), " {\n if(index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let batch = coords.x;\n let d1 = coords.u;\n\n let dyCorner = vec3(coords.y, coords.z, coords.w) - uniforms.pads;\n let dyFCorner = dyCorner.x;\n let dyRCorner = dyCorner.y;\n let dyCCorner = dyCorner.z;\n\n var dotProd = 0.0;\n for (var wF = 0; wF < uniforms.filterDims[0]; wF++) {\n let dyF = f32(dyFCorner + wF) / f32(uniforms.strides[0]);\n if (dyF < 0.0 || dyF >= f32(uniforms.outDepth) || fract(dyF) > 0.0) {\n continue;\n }\n let idyF = i32(dyF);\n\n let wFPerm = uniforms.filterDims[0] - 1 - wF;\n\n for (var wR = 0; wR < uniforms.filterDims[1]; wR++) {\n let dyR = f32(dyRCorner + wR) / f32(uniforms.strides[1]);\n\n if (dyR < 0.0 || dyR >= f32(uniforms.outHeight) || fract(dyR) > 0.0) {\n continue;\n }\n let idyR = i32(dyR);\n\n let wRPerm = uniforms.filterDims[1] - 1 - wR;\n\n for (var wC = 0; wC < uniforms.filterDims[2]; wC++) {\n let dyC = f32(dyCCorner + wC) / f32(uniforms.strides[2]);\n\n if (dyC < 0.0 || dyC >= f32(uniforms.outWidth) || fract(dyC) > 0.0) {\n continue;\n }\n let idyC = i32(dyC);\n\n let wCPerm = uniforms.filterDims[2] - 1 - wC;\n\n for (var d2 = 0; d2 < uniforms.outChannels; d2++) {\n let xValue = getDy(batch, idyF, idyR, idyC, d2);\n let wValue = getW(wFPerm, wRPerm, wCPerm, d1, d2);\n dotProd += xValue * wValue;\n }\n }\n }\n }\n setOutputAtIndex(index, dotProd);\n }\n }\n "); }; return Conv3DDerInputProgram; }()); /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function conv2DBackpropFilter(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, dy = inputs.dy; var strides = attrs.strides, pad = attrs.pad, dataFormat = attrs.dataFormat, dimRoundingMode = attrs.dimRoundingMode, filterShape = attrs.filterShape; var $dataFormat = tf.backend_util.convertConv2DDataFormat(dataFormat); var convInfo = tf.backend_util.computeConv2DInfo(x.shape, filterShape, strides, 1 /* dilations */, pad, dimRoundingMode, false /* depthwise */, $dataFormat); var program = new Conv2DDerFilterProgram(convInfo); var uniformData = [ { type: 'int32', data: [convInfo.padInfo.top, convInfo.padInfo.left] }, { type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.batchSize] }, { type: 'int32', data: [convInfo.outHeight] }, { type: 'int32', data: [convInfo.outWidth] }, { type: 'int32', data: [convInfo.inHeight] }, { type: 'int32', data: [convInfo.inWidth] } ]; return backend.runWebGPUProgram(program, [x, dy], x.dtype, uniformData); } var conv2DBackpropFilterConfig = { kernelName: tf.Conv2DBackpropFilter, backendName: 'webgpu', kernelFunc: conv2DBackpropFilter }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function conv2dTransposeCommonSnippet(innerElementSize) { if (innerElementSize === void 0) { innerElementSize = 4; } var getWSnippet = function (innerElementSize) { switch (innerElementSize) { case 1: return 'return W[getIndexFromCoords4D(coord, uniforms.wShape)];'; case 4: return "\n let coord1 = vec4(coordX, coordY, col + 1, rowInner);\n let coord2 = vec4(coordX, coordY, col + 2, rowInner);\n let coord3 = vec4(coordX, coordY, col + 3, rowInner);\n let v0 = W[getIndexFromCoords4D(coord, uniforms.wShape)];\n let v1 = W[getIndexFromCoords4D(coord1, uniforms.wShape)];\n let v2 = W[getIndexFromCoords4D(coord2, uniforms.wShape)];\n let v3 = W[getIndexFromCoords4D(coord3, uniforms.wShape)];\n return vec4(v0, v1, v2, v3);\n "; default: throw new Error("innerElementSize ".concat(innerElementSize, " is not supported.")); } }; var readASnippet = "\n let outRow = row / uniforms.outShape[2];\n let outCol = row % uniforms.outShape[2];\n\n let WRow = col / (uniforms.filterDims[1] * uniforms.outBackprop[3]);\n let WCol = col / uniforms.outBackprop[3] % uniforms.filterDims[1];\n let xR = f32(outRow - uniforms.pads[0] + WRow) / f32(uniforms.strides[0]);\n let xC = f32(outCol - uniforms.pads[1] + WCol) / f32(uniforms.strides[1]);\n if (xR < 0.0 || xR >= f32(uniforms.outBackprop[1]) || fract(xR) > 0.0) {\n return ".concat(typeSnippet(innerElementSize), "(0.0);\n }\n if (xC < 0.0 || xC >= f32(uniforms.outBackprop[2]) || fract(xC) > 0.0) {\n return ").concat(typeSnippet(innerElementSize), "(0.0);\n }\n let coord = vec4(\n batch,\n i32(xR),\n i32(xC),\n col % uniforms.outBackprop[3]);\n return x[getIndexFromCoords4D(coord, uniforms.xShape)/").concat(innerElementSize, "];"); var sampleA = "if (row < uniforms.dimAOuter && col < uniforms.dimInner) {\n ".concat(readASnippet, "\n }\n return ").concat(typeSnippet(innerElementSize), "(0.0);"); var userCode = "\n fn mm_readA(batch: i32, row : i32, col : i32) -> ".concat(typeSnippet(innerElementSize), " {\n ").concat(sampleA, "\n }\n\n fn mm_readB(batch: i32, row : i32, col : i32) -> ").concat(typeSnippet(innerElementSize), " {\n let coordX = uniforms.filterDims.x - 1 -\n row / (uniforms.filterDims[1] * uniforms.outBackprop[3]);\n let coordY = uniforms.filterDims.y - 1 -\n (row / uniforms.outBackprop[3]) % uniforms.filterDims[1];\n if (row < uniforms.dimInner && col < uniforms.dimBOuter &&\n coordX >= 0 && coordY >= 0) {\n let rowInner = row % uniforms.outBackprop[3];\n let coord = vec4(coordX, coordY, col, rowInner);\n ").concat(getWSnippet(innerElementSize), "\n }\n return ").concat(typeSnippet(innerElementSize), "(0.0);\n }\n\n fn mm_write(batch: i32, row : i32, col : i32, valueInput : ").concat(typeSnippet(innerElementSize), ") {\n if (row < uniforms.dimAOuter && col < uniforms.dimBOuter) {\n var value = valueInput;\n let outCoord = vec4(\n batch,\n row / uniforms.outShape[2],\n row % uniforms.outShape[2],\n col);\n result[getIndexFromCoords4D(outCoord, uniforms.outShape)/").concat(innerElementSize, "] = value;\n }\n }"); return userCode; } var Conv2DDerInputMMProgram = /** @class */ (function () { function Conv2DDerInputMMProgram(convInfo) { this.variableNames = ['x', 'W']; this.uniforms = 'filterDims : vec2, pads : vec2, strides : vec2, outBackprop : vec4, dimAOuter : i32, dimBOuter : i32, dimInner : i32,'; this.outputShape = convInfo.inShape; tf.util.assert(convInfo.dataFormat === 'channelsLast', function () { return 'TODO: NCHW is unimplemented'; }); this.isVec4 = convInfo.inChannels % 4 === 0 && convInfo.outChannels % 4 === 0; this.dispatchLayout = { x: [3], y: [1, 2], z: [0] }; this.workgroupSize = computeWorkgroupSizeForConv2d(this.dispatchLayout, this.outputShape, this.isVec4); this.elementsPerThread = computeWorkPerThreadForConv2d(this.dispatchLayout, this.outputShape, this.isVec4); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize, this.elementsPerThread); if (this.isVec4) { this.outputComponent = 4; this.variableComponents = [4, 1]; } this.shaderKey = "conv2DDerInputMM_".concat(this.isVec4, "_").concat(this.elementsPerThread); } Conv2DDerInputMMProgram.prototype.getUserCode = function () { var matMulSource = this.isVec4 ? makeMatMulPackedVec4Source(this.elementsPerThread, this.workgroupSize) : makeMatMulPackedSource(this.elementsPerThread, this.workgroupSize); var userCode = "\n ".concat(conv2dTransposeCommonSnippet(this.isVec4 ? 4 : 1), "\n ").concat(matMulSource, "\n "); return userCode; }; return Conv2DDerInputMMProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function conv2DBackpropInput(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var dy = inputs.dy, filter = inputs.filter; var inputShape = attrs.inputShape, strides = attrs.strides, pad = attrs.pad, dataFormat = attrs.dataFormat, dimRoundingMode = attrs.dimRoundingMode; var $dataFormat = tf.backend_util.convertConv2DDataFormat(dataFormat); var convInfo = tf.backend_util.computeConv2DInfo(inputShape, filter.shape, strides, 1 /* dilations */, pad, dimRoundingMode, false, $dataFormat); var dimensions = [ { type: 'int32', data: [convInfo.filterHeight, convInfo.filterWidth] }, { type: 'int32', data: [ convInfo.filterHeight - 1 - convInfo.padInfo.top, convInfo.filterWidth - 1 - convInfo.padInfo.left ] }, { type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [ convInfo.batchSize, convInfo.outHeight, convInfo.outWidth, convInfo.outChannels ] }, ]; var program; // TODO: Experiment when to use Conv2DDerInputMMProgram algorithm. if (tf.env().getBool('WEBGPU_USE_NAIVE_CONV2D_TRANSPOSE') || convInfo.dataFormat !== 'channelsLast') { program = new Conv2DDerInputProgram(convInfo); } else { program = new Conv2DDerInputMMProgram(convInfo); var dimAOuter = convInfo.inHeight * convInfo.inWidth; var dimBOuter = convInfo.inChannels; var dimInner = convInfo.filterHeight * convInfo.filterWidth * convInfo.outChannels; dimensions.push({ type: 'uint32', data: [dimAOuter] }, { type: 'uint32', data: [dimBOuter] }, { type: 'uint32', data: [dimInner] }); } return backend.runWebGPUProgram(program, [dy, filter], 'float32', dimensions); } var conv2DBackpropInputConfig = { kernelName: tf.Conv2DBackpropInput, backendName: 'webgpu', kernelFunc: conv2DBackpropInput, }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var Conv3DNaiveProgram = /** @class */ (function () { function Conv3DNaiveProgram(convInfo) { this.variableNames = ['x', 'W']; this.uniforms = 'filterDims: vec3, pads: vec3, strides: vec3, dilations: vec3,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = convInfo.outShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = "conv3dnaive"; } Conv3DNaiveProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getOutputCoords();\n let batch = coords.x;\n let d2 = coords.u;\n\n let xFRCCorner = vec3(coords.y, coords.z, coords.w) * uniforms.strides - uniforms.pads;\n let xFCorner = xFRCCorner.x;\n let xRCorner = xFRCCorner.y;\n let xCCorner = xFRCCorner.z;\n\n let inputDepthNearestVec4 = (uniforms.xShape.u / 4) * 4;\n let inputDepthVec4Remainder = uniforms.xShape.u % 4;\n\n var dotProd = 0.0;\n for (var wF = 0; wF < uniforms.filterDims[0]; wF++) {\n let xF = xFCorner + wF * uniforms.dilations[0];\n if (xF < 0 || xF >= uniforms.xShape.y) {\n continue;\n }\n\n for (var wR = 0; wR < uniforms.filterDims[1]; wR++) {\n let xR = xRCorner + wR * uniforms.dilations[1];\n if (xR < 0 || xR >= uniforms.xShape.z) {\n continue;\n }\n\n for (var wC = 0; wC < uniforms.filterDims[2]; wC++) {\n let xC = xCCorner + wC * uniforms.dilations[2];\n if (xC < 0 || xC >= uniforms.xShape.w) {\n continue;\n }\n\n for (var d1 = 0; d1 < inputDepthNearestVec4; d1 += 4) {\n let xValues = vec4(\n getX(batch, xF, xR, xC, d1),\n getX(batch, xF, xR, xC, d1 + 1),\n getX(batch, xF, xR, xC, d1 + 2),\n getX(batch, xF, xR, xC, d1 + 3)\n );\n let wValues = vec4(\n getW(wF, wR, wC, d1, d2),\n getW(wF, wR, wC, d1 + 1, d2),\n getW(wF, wR, wC, d1 + 2, d2),\n getW(wF, wR, wC, d1 + 3, d2)\n );\n\n dotProd += dot(xValues, wValues);\n }\n\n if (inputDepthVec4Remainder == 1) {\n dotProd += getX(batch, xF, xR, xC, inputDepthNearestVec4) *\n getW(wF, wR, wC, inputDepthNearestVec4, d2);\n } else if (inputDepthVec4Remainder == 2) {\n let xValues = vec2(\n getX(batch, xF, xR, xC, inputDepthNearestVec4),\n getX(batch, xF, xR, xC, inputDepthNearestVec4 + 1)\n );\n let wValues = vec2(\n getW(wF, wR, wC, inputDepthNearestVec4, d2),\n getW(wF, wR, wC, inputDepthNearestVec4 + 1, d2)\n );\n dotProd += dot(xValues, wValues);\n } else if (inputDepthVec4Remainder == 3) {\n let xValues = vec3(\n getX(batch, xF, xR, xC, inputDepthNearestVec4),\n getX(batch, xF, xR, xC, inputDepthNearestVec4 + 1),\n getX(batch, xF, xR, xC, inputDepthNearestVec4 + 2)\n );\n let wValues = vec3(\n getW(wF, wR, wC, inputDepthNearestVec4, d2),\n getW(wF, wR, wC, inputDepthNearestVec4 + 1, d2),\n getW(wF, wR, wC, inputDepthNearestVec4 + 2, d2)\n );\n dotProd += dot(xValues, wValues);\n }\n }\n }\n }\n setOutputAtIndex(index, dotProd);\n }\n }"); return userCode; }; return Conv3DNaiveProgram; }()); function conv3D(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, filter = inputs.filter; var strides = attrs.strides, pad = attrs.pad, dilations = attrs.dilations; var convInfo = tf.backend_util.computeConv3DInfo(x.shape, filter.shape, strides, dilations, pad); var padInfo = [convInfo.padInfo.front, convInfo.padInfo.top, convInfo.padInfo.left]; var dimensions = [ { type: 'int32', data: [convInfo.filterDepth, convInfo.filterHeight, convInfo.filterWidth] }, { type: 'int32', data: __spreadArray([], __read(padInfo), false) }, { type: 'int32', data: [convInfo.strideDepth, convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [ convInfo.dilationDepth, convInfo.dilationHeight, convInfo.dilationWidth ] } ]; var program = new Conv3DNaiveProgram(convInfo); var dtype = tf.upcastType(x.dtype, filter.dtype); return backend.runWebGPUProgram(program, [x, filter], dtype, dimensions); } var conv3DConfig = { kernelName: tf.Conv3D, backendName: 'webgpu', kernelFunc: conv3D, }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function conv3DBackpropFilterV2(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, dy = inputs.dy; var strides = attrs.strides, pad = attrs.pad, filterShape = attrs.filterShape; var convInfo = tf.backend_util.computeConv3DInfo(x.shape, filterShape, strides, 1 /* dilations */, pad); var program = new Conv3DDerFilterProgram(convInfo); var uniformData = [ { type: 'int32', data: [convInfo.padInfo.front, convInfo.padInfo.top, convInfo.padInfo.left] }, { type: 'int32', data: [convInfo.strideDepth, convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.batchSize] }, { type: 'int32', data: [convInfo.outDepth] }, { type: 'int32', data: [convInfo.outHeight] }, { type: 'int32', data: [convInfo.outWidth] }, { type: 'int32', data: [convInfo.inDepth] }, { type: 'int32', data: [convInfo.inHeight] }, { type: 'int32', data: [convInfo.inWidth] } ]; return backend.runWebGPUProgram(program, [x, dy], dy.dtype, uniformData); } var conv3DBackpropFilterV2Config = { kernelName: tf.Conv3DBackpropFilterV2, backendName: 'webgpu', kernelFunc: conv3DBackpropFilterV2 }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function conv3DBackpropInputV2(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var dy = inputs.dy, filter = inputs.filter; var strides = attrs.strides, pad = attrs.pad, inputShape = attrs.inputShape; var convInfo = tf.backend_util.computeConv3DInfo(inputShape, filter.shape, strides, 1 /* dilations */, pad); var program = new Conv3DDerInputProgram(convInfo); var uniformData = [ { type: 'int32', data: [convInfo.filterDepth, convInfo.filterHeight, convInfo.filterWidth] }, { type: 'int32', data: [ convInfo.filterDepth - 1 - convInfo.padInfo.front, convInfo.filterHeight - 1 - convInfo.padInfo.top, convInfo.filterWidth - 1 - convInfo.padInfo.left ] }, { type: 'int32', data: [convInfo.strideDepth, convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.outDepth] }, { type: 'int32', data: [convInfo.outHeight] }, { type: 'int32', data: [convInfo.outWidth] }, { type: 'int32', data: [convInfo.outChannels] } ]; return backend.runWebGPUProgram(program, [dy, filter], dy.dtype, uniformData); } var conv3DBackpropInputV2Config = { kernelName: tf.Conv3DBackpropInputV2, backendName: 'webgpu', kernelFunc: conv3DBackpropInputV2, }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var cos = unaryKernelFunc({ opType: UnaryOpType.COS }); var cosConfig = { kernelName: tf.Cos, backendName: 'webgpu', kernelFunc: cos }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var cosh = unaryKernelFunc({ opType: UnaryOpType.COSH }); var coshConfig = { kernelName: tf.Cosh, backendName: 'webgpu', kernelFunc: cosh }; var CropAndResizeProgram = /** @class */ (function () { function CropAndResizeProgram(channnel, boxShape, cropSize, method) { this.variableNames = ['Image', 'Boxes', 'BoxInd']; this.uniforms = 'extrapolationValue : f32,'; this.workgroupSize = [64, 1, 1]; this.size = true; var _a = __read(boxShape, 1), numBoxes = _a[0]; this.outputShape = [numBoxes, cropSize[0], cropSize[1], channnel]; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.methodId = method === 'bilinear' ? 1 : 0; this.cropHeightBiggerThan1 = this.outputShape[1] > 1; this.cropWidthBiggerThan1 = this.outputShape[2] > 1; this.shaderKey = "cropAndResize_".concat(this.methodId, "_").concat(this.cropHeightBiggerThan1, "_").concat(this.cropWidthBiggerThan1); } CropAndResizeProgram.prototype.getUserCode = function () { var _a = __read(["f32(uniforms.imageShape[1] - 1)", "f32(uniforms.imageShape[2] - 1)"], 2), inputHeightFloat = _a[0], inputWidthFloat = _a[1]; var _b = __read(this.cropHeightBiggerThan1 ? [ "(".concat(inputHeightFloat, " / f32(uniforms.outShape[1] - 1))"), '(y2-y1) * height_ratio', "y1*".concat(inputHeightFloat, " + f32(y)*(height_scale)"), ] : [ '0.0', '0.0', "0.5 * (y1+y2) * ".concat(inputHeightFloat), ], 3), heightRatio = _b[0], heightScale = _b[1], inY = _b[2]; var _c = __read(this.cropWidthBiggerThan1 ? [ "(".concat(inputWidthFloat, " / f32(uniforms.outShape[2] - 1))"), '(x2-x1) * width_ratio', "x1*".concat(inputWidthFloat, " + f32(x)*(width_scale)"), ] : [ '0.0', '0.0', "0.5 * (x1+x2) * ".concat(inputWidthFloat), ], 3), widthRatio = _c[0], widthScale = _c[1], inX = _c[2]; // Reference implementation // tslint:disable-next-line:max-line-length // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/crop_and_resize_op_gpu.cu.cc var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let height_ratio = f32(").concat(heightRatio, ");\n let width_ratio = f32(").concat(widthRatio, ");\n let b = coords[0];\n let y = coords[1];\n let x = coords[2];\n let d = coords[3];\n // get box vals\n let y1 = getBoxes(b, 0);\n let x1 = getBoxes(b, 1);\n let y2 = getBoxes(b, 2);\n let x2 = getBoxes(b, 3);\n // get image in batch index\n let bInd = i32(round(getBoxInd(b)));\n if(bInd < 0 || bInd >= uniforms.outShape[0]) {\n return;\n }\n let height_scale = ").concat(heightScale, ";\n let width_scale = ").concat(widthScale, ";\n let in_y = ").concat(inY, ";\n if( in_y < 0.0 || in_y > ").concat(inputHeightFloat, " ) {\n setOutputAtIndex(index, uniforms.extrapolationValue);\n return;\n }\n let in_x = ").concat(inX, ";\n if( in_x < 0.0 || in_x > ").concat(inputWidthFloat, " ) {\n setOutputAtIndex(index, uniforms.extrapolationValue);\n return;\n }\n let sourceFracIndexCR = vec2(in_x,in_y);\n if(").concat(this.methodId, " == 1) {\n // Compute the four integer indices.\n let sourceFloorCR = vec2(sourceFracIndexCR);\n let sourceCeilCR = vec2(ceil(sourceFracIndexCR));\n let topLeft = getImage(bInd, sourceFloorCR.y, sourceFloorCR.x, d);\n let bottomLeft = getImage(bInd, sourceCeilCR.y, sourceFloorCR.x, d);\n let topRight = getImage(bInd, sourceFloorCR.y, sourceCeilCR.x, d);\n let bottomRight = getImage(bInd, sourceCeilCR.y, sourceCeilCR.x, d);\n let fracCR = sourceFracIndexCR - vec2(sourceFloorCR);\n let top = topLeft + (topRight - topLeft) * fracCR.x;\n let bottom = bottomLeft + (bottomRight - bottomLeft) * fracCR.x;\n let newValue = top + (bottom - top) * fracCR.y;\n setOutputAtIndex(index, newValue);\n } else {\n // Compute the coordinators of nearest neighbor point.\n let sourceNearestCR = vec2(floor(\n sourceFracIndexCR + vec2(0.5,0.5)));\n let newValue = getImage(\n bInd, sourceNearestCR.y, sourceNearestCR.x, d);\n setOutputAtIndex(index, newValue);\n }\n }\n }\n "); return userCode; }; return CropAndResizeProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var cropAndResize = function (args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var image = inputs.image, boxes = inputs.boxes, boxInd = inputs.boxInd; var cropSize = attrs.cropSize, method = attrs.method, extrapolationValue = attrs.extrapolationValue; var program = new CropAndResizeProgram(image.shape[3], boxes.shape, cropSize, method); var uniformData = [{ type: 'float32', data: [extrapolationValue] }]; return backend.runWebGPUProgram(program, [image, boxes, boxInd], 'float32', uniformData); }; var cropAndResizeConfig = { kernelName: tf.CropAndResize, backendName: 'webgpu', kernelFunc: cropAndResize }; /** * @license * Copyright 2022 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var CumOpType; (function (CumOpType) { CumOpType["Prod"] = "*"; CumOpType["Sum"] = "+"; })(CumOpType || (CumOpType = {})); var CumProgram = /** @class */ (function () { function CumProgram(op, shape, exclusive, reverse) { this.variableNames = ['x']; // pow(i32, i32) is not supported, use pow(f32, f32) instead. this.uniforms = 'index : f32,'; this.size = true; this.workgroupSize = [128, 1, 1]; this.outputShape = shape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.exclusive = exclusive; this.reverse = reverse; this.op = op; this.shaderKey = "cum_".concat(this.op, "_").concat(this.exclusive, "_").concat(this.reverse); } CumProgram.prototype.getUserCode = function () { var rank = this.outputShape.length; var initVal = this.op === CumOpType.Prod ? '1.0' : '0.0'; var val = this.exclusive ? initVal : "getX(".concat(getCoords(rank, 'coords', this.op), ")"); var length = this.outputShape[this.outputShape.length - 1]; var condition = ''; var idxString = ''; // When exclusive is set, the cum op becomes roll op that copies the // value from the previous index based on the direction specified by the // reverse flag. if (this.exclusive) { condition = this.reverse ? "end != ".concat(length - 1) : 'end != 0'; idxString = this.reverse ? 'end + 1' : 'end - 1'; } else { condition = this.reverse ? "end + pow2 < ".concat(length) : 'end >= pow2'; idxString = (this.reverse ? 'end + pow2' : 'end - pow2'); } return "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n var coords = getCoordsFromIndex(index);\n\n let end = ").concat(getFinalCoord(rank, 'coords', this.op), ";\n var val = ").concat(val, ";\n let pow2 = i32(pow(2.0, uniforms.index));\n if (").concat(condition, ") {\n let idx = ").concat(idxString, ";\n ").concat(getFinalCoord(rank, 'coords', this.op), " = idx;\n val ").concat(this.op, "= getX(").concat(getCoords(rank, 'coords', this.op), ");\n }\n setOutputAtIndex(index, val);\n }\n }\n "); }; return CumProgram; }()); function getCoords(rank, name, op) { if (rank === 1) { return "".concat(name); } else if (rank === 2) { return "".concat(name, ".x, ").concat(name, ".y"); } else if (rank === 3) { return "".concat(name, ".x, ").concat(name, ".y, ").concat(name, ".z"); } else if (rank === 4) { return "".concat(name, ".x, ").concat(name, ".y, ").concat(name, ".z, ").concat(name, ".w"); } else { throw Error("Cumulative ".concat(op, " for rank ").concat(rank, " is not yet supported")); } } function getFinalCoord(rank, name, op) { if (rank === 1) { return "".concat(name); } else if (rank === 2) { return "".concat(name, ".y"); } else if (rank === 3) { return "".concat(name, ".z"); } else if (rank === 4) { return "".concat(name, ".w"); } else { throw Error("Cumulative ".concat(op, " for rank ").concat(rank, " is not yet supported")); } } /** * @license * Copyright 2022 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function cumImpl(op, x, backend, axis, exclusive, reverse) { var xRank = x.shape.length; var permutation = tf.backend_util.getAxesPermutation([axis], xRank); var permutedX = x; if (permutation != null) { permutedX = transpose({ inputs: { x: x }, backend: backend, attrs: { perm: permutation } }); } var permutedAxis = tf.backend_util.getInnerMostAxes(1, xRank)[0]; if (permutedAxis !== xRank - 1) { throw new Error("WebGPU cumprod shader expects an inner-most axis=".concat(x.shape.length - 1, " ") + "but got axis=".concat(axis)); } var size = permutedX.shape[permutedAxis]; var result = identity({ inputs: { x: permutedX }, backend: backend }); // Use cum parallel algorithm, inspired by: // https://developer.nvidia.com/gpugems/gpugems3/part-vi-gpu-computing/chapter-39-parallel-prefix-sum-scan-cuda // Note: although the algorithm is called sum, it works for any associtative // operator with an identity. for (var i = 0; i <= Math.ceil(Math.log2(size)) - 1; i++) { var program = new CumProgram(op, permutedX.shape, false, reverse); var prevResult = result; var uniformData = [{ type: 'float32', data: [i] }]; result = backend.runWebGPUProgram(program, [result], result.dtype, uniformData); backend.disposeData(prevResult.dataId); } // For exclusive cum, shift the end result in the direction of product or sum // and add 1 for product or 0 for sum to the front index. if (exclusive) { var program = new CumProgram(op, permutedX.shape, exclusive, reverse); var prevResult = result; var uniformData = [{ type: 'float32', data: [0] }]; result = backend.runWebGPUProgram(program, [result], result.dtype, uniformData); backend.disposeData(prevResult.dataId); } if (permutation != null) { var reversePermutation = tf.backend_util.getUndoAxesPermutation(permutation); var reverseTransposedResult = transpose({ inputs: { x: result }, backend: backend, attrs: { perm: reversePermutation } }); backend.disposeData(result.dataId); backend.disposeData(permutedX.dataId); return reverseTransposedResult; } return result; } /** * @license * Copyright 2022 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function cumprod(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var axis = attrs.axis, exclusive = attrs.exclusive, reverse = attrs.reverse; return cumImpl(CumOpType.Prod, x, backend, axis, exclusive, reverse); } var cumprodConfig = { kernelName: tf.Cumprod, backendName: 'webgpu', kernelFunc: cumprod }; /** * @license * Copyright 2022 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function cumsum(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var axis = attrs.axis, exclusive = attrs.exclusive, reverse = attrs.reverse; return cumImpl(CumOpType.Sum, x, backend, axis, exclusive, reverse); } var cumsumConfig = { kernelName: tf.Cumsum, backendName: 'webgpu', kernelFunc: cumsum }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function denseBincount(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, weights = inputs.weights; var size = attrs.size, binaryOutput = attrs.binaryOutput; var xRankOne = x.shape.length === 1; var weightsSize = tf.util.sizeFromShape(weights.shape); var hasWeights = weightsSize > 0; var dtype = weights.dtype; var xSize = xRankOne ? [x.shape[0]] : [x.shape[0], x.shape[1]]; var outputSize = xRankOne ? [size] : [x.shape[0], size]; var output = fill({ backend: backend, attrs: { shape: outputSize, value: 0, dtype: dtype } }); var program = new BincountProgram(xSize, hasWeights, binaryOutput); var uniformData = [{ type: 'int32', data: [size] }]; var bincountInputs = hasWeights ? [x, weights] : [x]; var res = backend.runWebGPUProgram(program, bincountInputs, dtype, uniformData, output); return res; } var denseBincountConfig = { kernelName: tf.DenseBincount, backendName: 'webgpu', kernelFunc: denseBincount }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var DepthToSpaceProgram = /** @class */ (function () { function DepthToSpaceProgram(outputShape, dataFormat) { this.variableNames = ['x']; this.workgroupSize = [64, 1, 1]; this.size = true; this.uniforms = 'blockSize : i32,'; this.outputShape = outputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = "depthToSpace_".concat(dataFormat); this.dataFormat = dataFormat; } DepthToSpaceProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let b = coords[0];\n let h = ").concat(this.getHeightCoordString(), ";\n let w = ").concat(this.getWidthCoordString(), ";\n let d = ").concat(this.getDepthCoordString(), ";\n\n let in_h = h / uniforms.blockSize;\n let offset_h = h % uniforms.blockSize;\n let in_w = w / uniforms.blockSize;\n let offset_w = w % uniforms.blockSize;\n let offset_d = (offset_h * uniforms.blockSize + offset_w) *\n ").concat(this.getOutputDepthSize(), ";\n let in_d = d + offset_d;\n\n let rlt = ").concat(this.getInputSamplingString(), ";\n setOutputAtIndex(index, rlt);\n }\n }"); return userCode; }; DepthToSpaceProgram.prototype.getHeightCoordString = function () { if (this.dataFormat === 'NHWC') { return "coords[1]"; } else { return "coords[2]"; } }; DepthToSpaceProgram.prototype.getWidthCoordString = function () { if (this.dataFormat === 'NHWC') { return "coords[2]"; } else { return "coords[3]"; } }; DepthToSpaceProgram.prototype.getDepthCoordString = function () { if (this.dataFormat === 'NHWC') { return "coords[3]"; } else { return "coords[1]"; } }; DepthToSpaceProgram.prototype.getOutputDepthSize = function () { if (this.dataFormat === 'NHWC') { return "uniforms.outShape[3]"; } else { return "uniforms.outShape[1]"; } }; DepthToSpaceProgram.prototype.getInputSamplingString = function () { if (this.dataFormat === 'NHWC') { return "getX(b, in_h, in_w, in_d)"; } else { return "getX(b, in_d, in_h, in_w)"; } }; return DepthToSpaceProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function depthToSpace(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var blockSize = attrs.blockSize, dataFormat = attrs.dataFormat; var batchSize = x.shape[0]; var inputHeight = (dataFormat === 'NHWC') ? x.shape[1] : x.shape[2]; var inputWidth = (dataFormat === 'NHWC') ? x.shape[2] : x.shape[3]; var inputDepth = (dataFormat === 'NHWC') ? x.shape[3] : x.shape[1]; var outputHeight = inputHeight * blockSize; var outputWidth = inputWidth * blockSize; var outputDepth = inputDepth / (blockSize * blockSize); var outputShape = (dataFormat === 'NHWC') ? [batchSize, outputHeight, outputWidth, outputDepth] : [batchSize, outputDepth, outputHeight, outputWidth]; var uniformData = [ { type: 'int32', data: [blockSize] }, ]; var program = new DepthToSpaceProgram(outputShape, dataFormat); return backend.runWebGPUProgram(program, [x], x.dtype, uniformData); } var depthToSpaceConfig = { kernelName: tf.DepthToSpace, backendName: 'webgpu', kernelFunc: depthToSpace }; /** * @license * Copyright 2022 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var DepthwiseConv2DNCHWSharedProgram = /** @class */ (function () { function DepthwiseConv2DNCHWSharedProgram(outputShape, filterHeight, filterWidth, addBias, activation, hasPreluActivation) { if (addBias === void 0) { addBias = false; } if (activation === void 0) { activation = null; } if (hasPreluActivation === void 0) { hasPreluActivation = false; } this.variableNames = ['x', 'W']; this.uniforms = "pads : vec2, inDims : vec2,"; this.workgroupSize = [16, 16, 1]; this.outputShape = outputShape; this.dispatchLayout = { x: [3], y: [2], z: [0, 1] }; this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); if (addBias) { this.variableNames.push('bias'); } if (hasPreluActivation) { this.variableNames.push('preluActivationWeights'); } this.addBias = addBias; this.activation = activation; this.hasPreluActivation = hasPreluActivation; this.filterHeight = filterHeight; this.filterWidth = filterWidth; this.shaderKey = "depthwiseNCHW_".concat(this.activation, "_").concat(this.filterHeight, "_").concat(this.filterWidth); } DepthwiseConv2DNCHWSharedProgram.prototype.getUserCode = function () { var filterSize = this.filterWidth * this.filterHeight; var flatWorkgroupSize = this.workgroupSize[0] * this.workgroupSize[1] * this.workgroupSize[2]; var tileAHeight = this.workgroupSize[1] + this.filterHeight - 1; var tileAWidth = this.workgroupSize[0] + this.filterWidth - 1; var userCode = "\n ".concat(activationFnSnippet(this.activation, this.hasPreluActivation, false, 4), "\n\n var mm_Asub : array, ").concat(tileAHeight, ">;\n var mm_Bsub : array, ").concat(this.filterHeight, ">;\n fn readX(batch : i32, channel : i32, row : i32, col : i32) -> f32 {\n var value = 0.0;\n if (row >=0 && row < uniforms.inDims[0] && col >=0 && col < uniforms.inDims[1])\n {\n value = getX(batch, channel, row, col);\n }\n return value;\n }\n\n ").concat(getMainHeaderString(), " {\n let coords = getOutputCoords();\n let batch = coords[0];\n let xRCCorner = vec2(coords.zw) - uniforms.pads;\n let channelMul = uniforms.wShape[3];\n let d1 = coords[1] / channelMul;\n let q = coords[1] % channelMul;\n\n let inputRowStart = xRCCorner.x;\n let inputColStart = xRCCorner.y;\n\n let localRow = i32(localId.y);\n let localCol = i32(localId.x);\n\n // Load one tile of X into local memory.\n for (var inputRow = localRow; inputRow < ").concat(tileAHeight, "; inputRow = inputRow + ").concat(this.workgroupSize[1], ") {\n for (var inputCol = localCol; inputCol < ").concat(tileAWidth, "; inputCol = inputCol + ").concat(this.workgroupSize[0], ") {\n let rowOffset = inputRow - localRow;\n let colOffset = inputCol - localCol;\n mm_Asub[inputRow][inputCol] = readX(batch, d1, inputRowStart + rowOffset, inputColStart + colOffset);\n }\n }\n\n // Load one tile of W into local memory.\n var wIndex = i32(localIndex);\n ").concat(filterSize < flatWorkgroupSize ? "if (wIndex < ".concat(filterSize, ")") : "for(; wIndex < ".concat(filterSize, "; wIndex = wIndex + ").concat(flatWorkgroupSize, ")"), "\n\n {\n let wRow = wIndex / ").concat(this.filterWidth, ";\n let wCol = wIndex % ").concat(this.filterWidth, ";\n mm_Bsub[wRow][wCol] = getW(wRow, wCol, d1, q);\n }\n\n workgroupBarrier();\n\n var value = 0.0;\n for (var wR = 0; wR < ").concat(this.filterHeight, "; wR = wR + 1) {\n for (var wC = 0; wC < ").concat(this.filterWidth, "; wC = wC + 1) {\n let xVal = mm_Asub[localRow + wR][localCol + wC];\n let wVal = mm_Bsub[wR][wC];\n value = fma(xVal, wVal, value);\n }\n }\n ").concat(biasActivationSnippet(this.addBias, this.activation), "\n if (coordsInBounds4D(coords, uniforms.outShape)) {\n setOutputAtCoords(coords[0], coords[1], coords[2], coords[3], value);\n }\n }\n "); return userCode; }; return DepthwiseConv2DNCHWSharedProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var DepthwiseConv2DVec4Program = /** @class */ (function () { function DepthwiseConv2DVec4Program(convInfo, addBias, activation, hasPreluActivation) { if (addBias === void 0) { addBias = false; } if (activation === void 0) { activation = null; } if (hasPreluActivation === void 0) { hasPreluActivation = false; } this.variableNames = ['x', 'W']; this.uniforms = 'pads : vec2, inDims : vec2, virtualWidth : i32,'; this.workgroupSize = [64, 1, 1]; this.workPerThread = 4; this.outputComponent = 4; this.outputShape = convInfo.outShape; this.virtualWidth = Math.ceil(this.outputShape[2] / this.workPerThread) * this.workPerThread; var virtualOutputShape = [ this.outputShape[0], this.outputShape[1], this.virtualWidth, this.outputShape[3] ]; this.dispatchLayout = flatDispatchLayout(virtualOutputShape); this.dispatch = computeDispatch(this.dispatchLayout, virtualOutputShape, this.workgroupSize, [this.outputComponent * this.workPerThread, 1, 1]); tf.util.assert(convInfo.dataFormat === 'channelsLast', function () { return 'TODO: NCHW is unimplemented'; }); if (addBias) { this.variableNames.push('bias'); } if (hasPreluActivation) { this.variableNames.push('preluActivationWeights'); } this.convInfo = convInfo; this.addBias = addBias; this.activation = activation; this.hasPreluActivation = hasPreluActivation; this.shaderKey = "depthwiseVec4_".concat(activation, "_").concat(this.convInfo.filterHeight, "_").concat(this.convInfo.filterWidth, "_").concat(this.convInfo.strideHeight, "_").concat(this.convInfo.strideWidth, "_").concat(this.workPerThread); } DepthwiseConv2DVec4Program.prototype.getUserCode = function () { var xNumber = (this.workPerThread - 1) * this.convInfo.strideWidth + this.convInfo.filterWidth; var strideHeight = this.convInfo.strideHeight; var strideWidth = this.convInfo.strideWidth; var userCode = "\n ".concat(activationFnSnippet(this.activation, this.hasPreluActivation, true, 4), "\n fn readX(batch : i32, row : i32, col : i32, channel : i32) -> vec4 {\n var value = vec4(0.0);\n if (col >=0 && col < uniforms.inDims[1]) {\n value = getX(batch, row, col, channel);\n }\n return value;\n }\n\n ").concat(getMainHeaderString('index'), " {\n let width0 = uniforms.outShape[3] / ").concat(this.outputComponent, ";\n let d1 = (index % width0) * ").concat(this.outputComponent, ";\n var index1 = index / width0;\n let width1 = uniforms.virtualWidth / ").concat(this.workPerThread, ";\n let c = (index1 % width1) * ").concat(this.workPerThread, ";\n index1 = index1 / width1;\n let r = index1 % uniforms.outShape[1];\n let batch = index1 / uniforms.outShape[1];\n\n let xRCCorner = vec2(r, c) * vec2(").concat(strideHeight, ", ").concat(strideWidth, ") - uniforms.pads;\n\n let xRCorner = xRCCorner.x;\n let xCCorner = xRCCorner.y;\n var xVals : array, ").concat(xNumber, ">;\n var dotProd : array, ").concat(this.workPerThread, ">;\n for (var i = 0; i < ").concat(this.workPerThread, "; i++) {\n dotProd[i] = vec4(0.0);\n }\n\n // Use constant instead of uniform can give better performance.\n for (var wR = 0; wR < ").concat(this.convInfo.filterHeight, "; wR = wR + 1) {\n let xR = xRCorner + wR;\n if (xR >=0 && xR < uniforms.inDims[0]) {\n for (var i = 0; i < ").concat(xNumber, "; i++) {\n xVals[i] = readX(batch, xR, xCCorner + i, d1);\n }\n for (var wC = 0; wC < ").concat(this.convInfo.filterWidth, "; wC = wC + 1) {\n let wValue = getW(wR, wC, d1, 0);\n for (var i = 0; i < ").concat(this.workPerThread, "; i++) {\n dotProd[i] = fma(xVals[i * ").concat(strideWidth, " + wC], wValue, dotProd[i]);\n }\n }\n }\n }\n\n for (var i = 0; i < ").concat(this.workPerThread, "; i = i + 1) {\n let coords = vec4(batch, r, c + i, d1);\n if (coordsInBounds4D(coords, uniforms.outShape)) {\n var value = dotProd[i];\n ").concat(biasActivationSnippet(this.addBias, this.activation), "\n setOutputAtCoords(coords[0], coords[1], coords[2], coords[3], value);\n }\n }\n }\n "); return userCode; }; return DepthwiseConv2DVec4Program; }()); /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var DepthwiseConv2DProgram = /** @class */ (function () { function DepthwiseConv2DProgram(convInfo, addBias, activation, hasPreluActivation) { if (addBias === void 0) { addBias = false; } if (activation === void 0) { activation = null; } if (hasPreluActivation === void 0) { hasPreluActivation = false; } this.variableNames = ['x', 'W']; this.uniforms = "pads : vec2, inDims : vec2, filterHeight : i32,\n filterWidth : i32, strides : vec2, dilations : vec2,"; // This is an experimental value. this.workgroupSize = [256, 1, 1]; this.size = true; this.outputShape = convInfo.outShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.isChannelsLast = convInfo.dataFormat === 'channelsLast'; if (addBias) { this.variableNames.push('bias'); } if (hasPreluActivation) { this.variableNames.push('preluActivationWeights'); } this.convInfo = convInfo; this.addBias = addBias; this.activation = activation; this.hasPreluActivation = hasPreluActivation; this.shaderKey = "depthwise_".concat(this.activation, "_").concat(this.isChannelsLast); } DepthwiseConv2DProgram.prototype.getUserCode = function () { var getXSnippet = this.isChannelsLast ? 'getX(batch, xR, xC, d1);' : 'getX(batch, d1, xR, xC);'; var userCode = "\n ".concat(activationFnSnippet(this.activation, this.hasPreluActivation, false, 4), "\n\n ").concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getOutputCoords();\n let batch = coords[0];\n let xRCCorner = vec2(coords.").concat(this.isChannelsLast ? 'yz' : 'zw', ") * uniforms.strides - uniforms.pads;\n let d2 = coords[").concat(this.isChannelsLast ? 3 : 1, "];\n let channelMul = uniforms.wShape[3];\n let d1 = d2 / channelMul;\n let q = d2 % channelMul;\n\n let inputRowStart = xRCCorner.x;\n let inputColStart = xRCCorner.y;\n let inputRowEnd = inputRowStart + uniforms.filterHeight *\n uniforms.dilations[0];\n let inputColEnd = inputColStart + uniforms.filterWidth *\n uniforms.dilations[1];\n\n // Convolve x(?, ?, d1)|x(d1, ?, ?) with w(:, :, d1, q) to get\n // y(yR, yC, d2)|y(d2, yR, yC). ? = to be determined. : = across all\n // values in that axis. x(?, ?, d1) and y(yR, yC, d2) is for NHWC.\n // x(d1, ?, ?) and y(d2, yR, yC) is for NCHW.\n var value = 0.0;\n\n // Extract if checking out of for loop for performance.\n if (inputRowStart >= 0 && inputColStart >= 0 &&\n inputRowEnd < uniforms.inDims[0] &&\n inputColEnd < uniforms.inDims[1]) {\n for (var wR = 0; wR < uniforms.filterHeight; wR = wR + 1) {\n let xR = inputRowStart + wR * uniforms.dilations[0];\n\n for (var wC = 0; wC < uniforms.filterWidth; wC = wC + 1) {\n let xC = inputColStart + wC * uniforms.dilations[1];\n\n let xVal = ").concat(getXSnippet, ";\n let wVal = getW(wR, wC, d1, q);\n value = value + xVal * wVal;\n }\n }\n } else {\n for (var wR = 0; wR < uniforms.filterHeight; wR = wR + 1) {\n let xR = inputRowStart + wR * uniforms.dilations[0];\n\n if (xR < 0 || xR >= uniforms.inDims[0]) {\n continue;\n }\n\n for (var wC = 0; wC < uniforms.filterWidth; wC = wC + 1) {\n let xC = inputColStart + wC * uniforms.dilations[1];\n\n if (xC < 0 || xC >= uniforms.inDims[1]) {\n continue;\n }\n\n let xVal = ").concat(getXSnippet, ";\n let wVal = getW(wR, wC, d1, q);\n value = value + xVal * wVal;\n }\n }\n }\n ").concat(biasActivationSnippet(this.addBias, this.activation), "\n setOutputAtCoords(coords[0], coords[1], coords[2], coords[3], value);\n }\n }\n "); return userCode; }; return DepthwiseConv2DProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function depthwiseConv2dNative(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, filter = inputs.filter; var strides = attrs.strides, pad = attrs.pad, dataFormat = attrs.dataFormat, dilations = attrs.dilations, dimRoundingMode = attrs.dimRoundingMode; var $dataFormat = tf.backend_util.convertConv2DDataFormat(dataFormat); var $dilations = dilations; if ($dilations == null) { $dilations = [1, 1]; } var convInfo = tf.backend_util.computeConv2DInfo(x.shape, filter.shape, strides, $dilations, pad, dimRoundingMode, true /* depthwise */, $dataFormat); var dimensions = [ { type: 'int32', data: [convInfo.padInfo.top, convInfo.padInfo.left] }, { type: 'int32', data: [convInfo.inHeight, convInfo.inWidth] }, ]; var isChannelsLast = convInfo.dataFormat === 'channelsLast'; var program; if (!isChannelsLast && convInfo.inHeight > 16 && convInfo.inWidth > 16 && convInfo.strideHeight === 1 && convInfo.strideWidth === 1 && convInfo.dilationWidth === 1 && convInfo.dilationHeight === 1 && convInfo.inChannels === convInfo.outChannels) { program = new DepthwiseConv2DNCHWSharedProgram(convInfo.outShape, convInfo.filterHeight, convInfo.filterWidth); } else if (isChannelsLast && convInfo.outHeight > 4 && convInfo.outWidth > 4 && convInfo.strideWidth <= 2 && convInfo.inChannels === convInfo.outChannels && convInfo.dilationHeight === 1 && convInfo.dilationWidth === 1 && convInfo.inChannels % 4 === 0) { program = new DepthwiseConv2DVec4Program(convInfo); dimensions.push({ type: 'int32', data: [program.virtualWidth] }); } else { program = new DepthwiseConv2DProgram(convInfo); dimensions.push({ type: 'int32', data: [convInfo.filterHeight] }, { type: 'int32', data: [convInfo.filterWidth] }, { type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.dilationHeight, convInfo.dilationWidth] }); } return backend.runWebGPUProgram(program, [x, filter], x.dtype, dimensions); } var depthwiseConv2dNativeConfig = { kernelName: tf.DepthwiseConv2dNative, backendName: 'webgpu', kernelFunc: depthwiseConv2dNative, }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var DepthwiseConv2DDerFilterProgram = /** @class */ (function () { function DepthwiseConv2DDerFilterProgram(convInfo) { this.variableNames = ['x', 'dy']; this.uniforms = "strides : vec2, pads : vec2, filterDims : vec2, outHeight : i32,\n outWidth : i32, inHeight : i32, inWidth : i32, batchSize : i32, channelMul : i32,"; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = convInfo.filterShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = "depthwise_conv2d_backprop_filter"; } DepthwiseConv2DDerFilterProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let wR = coords[0];\n let wC = coords[1];\n let d1 = coords[2];\n let dm = coords[3];\n let d2 = d1 * uniforms.channelMul + dm;\n\n var dotProd = 0.0;\n for (var b = 0; b < uniforms.batchSize; b++) {\n for (var yR = 0; yR < uniforms.outHeight; yR++) {\n let xR = wR + yR * uniforms.strides[0] - uniforms.pads[0];\n\n if (xR < 0 || xR >= uniforms.inHeight) {\n continue;\n }\n\n for (var yC = 0; yC < uniforms.outWidth; yC++) {\n let xC = wC + yC * uniforms.strides[1] - uniforms.pads[1];\n\n if (xC < 0 || xC >= uniforms.inWidth) {\n continue;\n }\n\n let dyValue = getDy(b, yR, yC, d2);\n let xValue = getX(b, xR, xC, d1);\n dotProd += xValue * dyValue;\n }\n }\n }\n setOutputAtIndex(index, dotProd);\n }\n }\n "); return userCode; }; return DepthwiseConv2DDerFilterProgram; }()); var DepthwiseConv2DDerInputProgram = /** @class */ (function () { function DepthwiseConv2DDerInputProgram(convInfo) { this.variableNames = ['dy', 'W']; this.uniforms = "strides : vec2, pads : vec2, filterDims : vec2,\n outHeight : i32, outWidth : i32, channelMul : i32,"; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = convInfo.inShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = "depthwise_conv2d_backprop_input"; } DepthwiseConv2DDerInputProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let batch = coords[0];\n let d1 = coords[3];\n let dyCorner = coords.yz - uniforms.pads;\n let dyRCorner = dyCorner.x;\n let dyCCorner = dyCorner.y;\n\n var dotProd = 0.0;\n for (var wR = 0; wR < uniforms.filterDims[0]; wR++) {\n let dyR = f32(dyRCorner + wR) / f32(uniforms.strides[0]);\n\n if (dyR < 0.0 || dyR >= f32(uniforms.outHeight) || fract(dyR) > 0.0) {\n continue;\n }\n\n let idyR = i32(dyR);\n let wRPerm = uniforms.filterDims[0] - 1 - wR;\n\n for (var wC = 0; wC < uniforms.filterDims[1]; wC++) {\n let dyC = f32(dyCCorner + wC) / f32(uniforms.strides[1]);\n\n if (dyC < 0.0 || dyC >= f32(uniforms.outWidth) || fract(dyC) > 0.0) {\n continue;\n }\n\n let idyC = i32(dyC);\n let wCPerm = uniforms.filterDims[1] - 1 - wC;\n\n for (var dm = 0; dm < uniforms.channelMul; dm++) {\n let d2 = d1 * uniforms.channelMul + dm;\n let xValue = getDy(batch, idyR, idyC, d2);\n let wValue = getW(wRPerm, wCPerm, d1, dm);\n dotProd += xValue * wValue;\n }\n }\n }\n setOutputAtIndex(index, dotProd);\n }\n }\n "); return userCode; }; return DepthwiseConv2DDerInputProgram; }()); /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function depthwiseConv2dNativeBackpropFilter(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, dy = inputs.dy; var strides = attrs.strides, dilations = attrs.dilations, pad = attrs.pad, dimRoundingMode = attrs.dimRoundingMode, filterShape = attrs.filterShape; var convInfo = tf.backend_util.computeConv2DInfo(x.shape, filterShape, strides, dilations, pad, dimRoundingMode, true /* depthwise */); var program = new DepthwiseConv2DDerFilterProgram(convInfo); var uniformData = [ { type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.padInfo.top, convInfo.padInfo.left] }, { type: 'int32', data: [convInfo.filterHeight, convInfo.filterWidth] }, { type: 'int32', data: [convInfo.outHeight] }, { type: 'int32', data: [convInfo.outWidth] }, { type: 'int32', data: [convInfo.inHeight] }, { type: 'int32', data: [convInfo.inWidth] }, { type: 'int32', data: [convInfo.batchSize] }, { type: 'int32', data: [convInfo.outChannels / convInfo.inChannels] } ]; return backend.runWebGPUProgram(program, [x, dy], 'float32', uniformData); } var depthwiseConv2dNativeBackpropFilterConfig = { kernelName: tf.DepthwiseConv2dNativeBackpropFilter, backendName: 'webgpu', kernelFunc: depthwiseConv2dNativeBackpropFilter }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function depthwiseConv2dNativeBackpropInput(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var dy = inputs.dy, filter = inputs.filter; var strides = attrs.strides, dilations = attrs.dilations, pad = attrs.pad, dimRoundingMode = attrs.dimRoundingMode, inputShape = attrs.inputShape; var convInfo = tf.backend_util.computeConv2DInfo(inputShape, filter.shape, strides, dilations, pad, dimRoundingMode, true /* depthwise */); var program = new DepthwiseConv2DDerInputProgram(convInfo); var uniformData = [ { type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [ convInfo.filterHeight - 1 - convInfo.padInfo.top, convInfo.filterWidth - 1 - convInfo.padInfo.left ] }, { type: 'int32', data: [convInfo.filterHeight, convInfo.filterWidth] }, { type: 'int32', data: [convInfo.outHeight] }, { type: 'int32', data: [convInfo.outWidth] }, { type: 'int32', data: [convInfo.outChannels / convInfo.inChannels] } ]; return backend.runWebGPUProgram(program, [dy, filter], dy.dtype, uniformData); } var depthwiseConv2dNativeBackpropInputConfig = { kernelName: tf.DepthwiseConv2dNativeBackpropInput, backendName: 'webgpu', kernelFunc: depthwiseConv2dNativeBackpropInput }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var DiagProgram = /** @class */ (function () { function DiagProgram(size) { this.variableNames = ['x']; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = [size, size]; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'diag'; } DiagProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getOutputCoords();\n let value = select(0.0, getX(coords[0]), coords[0] == coords[1]);\n setOutputAtIndex(index, value);\n }\n }\n "); return userCode; }; return DiagProgram; }()); function diag(args) { var inputs = args.inputs, backend = args.backend; var x = inputs.x; var outShape = __spreadArray(__spreadArray([], __read(x.shape), false), __read(x.shape), false); var xSize = tf.util.sizeFromShape(x.shape); var flat = reshape({ inputs: { x: x }, backend: backend, attrs: { shape: [xSize] } }); var program = new DiagProgram(xSize); var res = backend.runWebGPUProgram(program, [flat], flat.dtype); var out = reshape({ inputs: { x: res }, backend: backend, attrs: { shape: outShape } }); backend.disposeData(flat.dataId); backend.disposeData(res.dataId); return out; } var diagConfig = { kernelName: tf.Diag, backendName: 'webgpu', kernelFunc: diag }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var Dilation2DProgram = /** @class */ (function () { function Dilation2DProgram(convInfo) { this.variableNames = ['x', 'w']; this.uniforms = 'filterDims: vec2, pads: vec2, strides: vec2, dilations: vec2'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = convInfo.outShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'dilation2d'; } Dilation2DProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let neg_infinity = -3.4e38;\n let coords = getOutputCoords();\n let batch = coords.x;\n let d1 = coords.w;\n let outTopLeftCorner = coords.yz * uniforms.strides - uniforms.pads;\n let hBeg = outTopLeftCorner.x;\n let wBeg = outTopLeftCorner.y;\n\n var curVal = neg_infinity;\n for (var h = 0; h < uniforms.filterDims[0]; h = h + 1) {\n let hIn = hBeg + h * uniforms.dilations[0];\n\n if (hIn >= 0 && hIn < uniforms.xShape[1]) {\n for (var w = 0; w < uniforms.filterDims[1]; w = w + 1) {\n let wIn = wBeg + w * uniforms.dilations[1];\n\n if (wIn >= 0 && wIn < uniforms.xShape[2]) {\n let val = getX(batch, hIn, wIn, d1) + getW(h, w, d1);\n if (val > curVal) {\n curVal = val;\n }\n }\n }\n }\n }\n\n setOutputAtIndex(index, curVal);\n }\n }\n "); return userCode; }; return Dilation2DProgram; }()); function dilation2D(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, filter = inputs.filter; var strides = attrs.strides, pad = attrs.pad, dilations = attrs.dilations; var convInfo = tf.backend_util.computeDilation2DInfo(x.shape, filter.shape, strides, pad, 'NHWC' /* dataFormat */, dilations); var padInfo = [convInfo.padInfo.top, convInfo.padInfo.left]; var uniformData = [ { type: 'int32', data: [convInfo.filterHeight, convInfo.filterWidth] }, { type: 'int32', data: __spreadArray([], __read(padInfo), false) }, { type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.dilationHeight, convInfo.dilationWidth] } ]; var program = new Dilation2DProgram(convInfo); var out = backend.runWebGPUProgram(program, [x, filter], x.dtype, uniformData); return out; } var dilation2DConfig = { kernelName: tf.Dilation2D, backendName: 'webgpu', kernelFunc: dilation2D }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var Dilation2DBackpropInputProgram = /** @class */ (function () { function Dilation2DBackpropInputProgram(convInfo, outputDtype) { this.variableNames = ['x', 'w', 'dy']; this.uniforms = 'filterDims: vec2, pads: vec2, strides: vec2, dilations: vec2, dySize: i32,'; this.workgroupSize = [64, 1, 1]; this.atomic = true; this.outputShape = convInfo.inShape; this.dispatchLayout = flatDispatchLayout(convInfo.outShape); this.dispatch = computeDispatch(this.dispatchLayout, convInfo.outShape, this.workgroupSize); if (outputDtype !== 'float32' && outputDtype !== 'int32') { throw new Error("Dilation2DBackpropInput only supports float32 and int32\n types, does not support ".concat(outputDtype, " type.")); } this.type = outputDtype; this.shaderKey = 'dilation2DBackpropInput'; } Dilation2DBackpropInputProgram.prototype.getUserCode = function () { // This implementation follows the TF c++ cuda implementation: // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/dilation_ops_gpu.cu.cc var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.dySize) {\n let coords = getDyCoordsFromIndex(index);\n let b = coords[0];\n let r = coords[1];\n let c = coords[2];\n let d = coords[3];\n\n let dyCorner = vec2(r, c) * uniforms.strides - uniforms.pads;\n var curVal = -3.4e38; // neg_infinity\n var xRMax = 0;\n var xCMax = 0;\n\n // In the case of multiple argmax branches, we only back-propagate\n // along the last branch, i.e., the one with largest value of\n // 'wR * uniforms.filterDims[1] + wC', similarly to the max-pooling\n // backward routines.\n for (var wR = 0; wR < uniforms.filterDims[0]; wR++) {\n let xR = dyCorner.x + wR * uniforms.dilations[0];\n\n if (xR >= 0 && xR < uniforms.xShape[1]) {\n for (var wC = 0; wC < uniforms.filterDims[1]; wC++) {\n let xC = dyCorner.y + wC * uniforms.dilations[1];\n\n if (xC >= 0 && xC < uniforms.xShape[2]) {\n let val = getX(b, xR, xC, d) + getW(wR, wC, d);\n if (val > curVal) {\n curVal = val;\n xRMax = xR;\n xCMax = xC;\n }\n }\n }\n }\n }\n\n let flatIndexIn = d + uniforms.xShape[3] *\n (xCMax + uniforms.xShape[2] * (xRMax + uniforms.xShape[1] * b));\n let value = getDy(b, r, c, d);\n ").concat(atomicAddSnippet('&result[flatIndexIn]', 'value', this.type), "\n }\n }\n "); return userCode; }; return Dilation2DBackpropInputProgram; }()); var Dilation2DBackpropFilterProgram = /** @class */ (function () { function Dilation2DBackpropFilterProgram(convInfo, shape, outputDtype) { this.variableNames = ['x', 'w', 'dy']; this.uniforms = 'filterDims: vec2, pads: vec2, strides: vec2, dilations: vec2, dySize: i32,'; this.workgroupSize = [64, 1, 1]; this.atomic = true; this.outputShape = convInfo.filterShape; this.dispatchLayout = flatDispatchLayout(convInfo.outShape); this.dispatch = computeDispatch(this.dispatchLayout, convInfo.outShape, this.workgroupSize); if (outputDtype !== 'float32' && outputDtype !== 'int32') { throw new Error("Dilation2DBackpropFilter only supports float32 and int32\n types, does not support ".concat(outputDtype, " type.")); } this.type = outputDtype; this.shaderKey = 'dilation2DBackpropFilter'; } Dilation2DBackpropFilterProgram.prototype.getUserCode = function () { // This implementation follows the TF c++ cuda implementation: // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/dilation_ops_gpu.cu.cc var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.dySize) {\n let coords = getDyCoordsFromIndex(index);\n let b = coords[0];\n let r = coords[1];\n let c = coords[2];\n let d = coords[3];\n\n let dyCorner = vec2(r, c) * uniforms.strides - uniforms.pads;\n var curVal = -3.4e38; // neg_infinity\n var wRMax = 0;\n var wCMax = 0;\n\n // In the case of multiple argmax branches, we only back-propagate\n // along the last branch, i.e., the one with largest value of\n // 'wR * uniforms.filterDims[1] + wC', similarly to the max-pooling\n // backward routines.\n for (var wR = 0; wR < uniforms.filterDims[0]; wR++) {\n let xR = dyCorner.x + wR * uniforms.dilations[0];\n\n if (xR >= 0 && xR < uniforms.xShape[1]) {\n for (var wC = 0; wC < uniforms.filterDims[1]; wC++) {\n let xC = dyCorner.y + wC * uniforms.dilations[1];\n\n if (xC >= 0 && xC < uniforms.xShape[2]) {\n let val = getX(b, xR, xC, d) + getW(wR, wC, d);\n if (val > curVal) {\n curVal = val;\n wRMax = wR;\n wCMax = wC;\n }\n }\n }\n }\n }\n\n let flatIndexIn = d + uniforms.wShape[2] * (wCMax + wRMax * uniforms.wShape[1]);\n let value = getDy(b, r, c, d);\n ").concat(atomicAddSnippet('&result[flatIndexIn]', 'value', this.type), "\n }\n }\n "); return userCode; }; return Dilation2DBackpropFilterProgram; }()); /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function dilation2DBackpropFilter(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, filter = inputs.filter, dy = inputs.dy; var strides = attrs.strides, pad = attrs.pad, dilations = attrs.dilations; var convInfo = tf.backend_util.computeDilation2DInfo(x.shape, filter.shape, strides, pad, 'NHWC' /* dataFormat */, dilations); var dtype = filter.dtype; var program = new Dilation2DBackpropFilterProgram(convInfo, filter.shape, dtype); var uniformData = [ { type: 'int32', data: [convInfo.filterHeight, convInfo.filterWidth] }, { type: 'int32', data: [convInfo.padInfo.top, convInfo.padInfo.left] }, { type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.dilationHeight, convInfo.dilationWidth] }, { type: 'int32', data: [tf.util.sizeFromShape(convInfo.outShape)] } ]; var output = fill({ backend: backend, attrs: { shape: filter.shape, value: 0, dtype: dtype } }); return backend.runWebGPUProgram(program, [x, filter, dy], dtype, uniformData, output); } var dilation2DBackpropFilterConfig = { kernelName: tf.Dilation2DBackpropFilter, backendName: 'webgpu', kernelFunc: dilation2DBackpropFilter }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function dilation2DBackpropInput(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, filter = inputs.filter, dy = inputs.dy; var strides = attrs.strides, pad = attrs.pad, dilations = attrs.dilations; var convInfo = tf.backend_util.computeDilation2DInfo(x.shape, filter.shape, strides, pad, 'NHWC' /* dataFormat */, dilations); var dtype = x.dtype; var program = new Dilation2DBackpropInputProgram(convInfo, dtype); var uniformData = [ { type: 'int32', data: [convInfo.filterHeight, convInfo.filterWidth] }, { type: 'int32', data: [convInfo.padInfo.top, convInfo.padInfo.left] }, { type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.dilationHeight, convInfo.dilationWidth] }, { type: 'int32', data: [tf.util.sizeFromShape(convInfo.outShape)] } ]; var output = fill({ backend: backend, attrs: { shape: convInfo.inShape, value: 0, dtype: dtype } }); return backend.runWebGPUProgram(program, [x, filter, dy], dtype, uniformData, output); } var dilation2DBackpropInputConfig = { kernelName: tf.Dilation2DBackpropInput, backendName: 'webgpu', kernelFunc: dilation2DBackpropInput }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var DrawProgram = /** @class */ (function () { function DrawProgram(outShape, type, textureFormat) { this.variableNames = ['Image']; this.uniforms = 'alpha: f32,'; this.workgroupSize = [64, 1, 1]; this.pixelsOpType = PixelsOpType.DRAW; this.size = true; this.outputShape = outShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.type = type; this.textureFormat = textureFormat; this.shaderKey = "draw_".concat(type, "_").concat(textureFormat); } DrawProgram.prototype.getUserCode = function () { var calculateResult; var value = this.type === 'float32' ? 'value' : 'value / 255.0'; calculateResult = "\n if (uniforms.numChannels == 1) {\n rgba[0] = ".concat(value, ";\n rgba[1] = ").concat(value, ";\n rgba[2] = ").concat(value, ";\n } else {\n rgba[d] = ").concat(value, ";\n }"); var userCode = "\n @group(0) @binding(0) var outImage : texture_storage_2d<".concat(this.textureFormat, ", write>;\n ").concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n var rgba = vec4(0.0, 0.0, 0.0, uniforms.alpha);\n for (var d = 0; d < uniforms.numChannels; d = d + 1) {\n let value = f32(inBuf[index * uniforms.numChannels + d]);\n ").concat(calculateResult, "\n }\n rgba.x = rgba.x * rgba.w;\n rgba.y = rgba.y * rgba.w;\n rgba.z = rgba.z * rgba.w;\n let coords = getCoordsFromIndex(index);\n textureStore(outImage, vec2(coords.yx), rgba);\n }\n }\n "); return userCode; }; return DrawProgram; }()); function draw(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var image = inputs.image; var canvas = attrs.canvas, options = attrs.options; var _a = __read(image.shape.slice(0, 2), 2), height = _a[0], width = _a[1]; var imageOptions = (options || {}).imageOptions; var alpha = (imageOptions === null || imageOptions === void 0 ? void 0 : imageOptions.alpha) || 1; // 'rgba8unorm' should work on macOS according to // https://bugs.chromium.org/p/chromium/issues/detail?id=1298618. But // failed on macOS/M2. So use 'bgra8unorm' first when available. var format = backend.device.features.has('bgra8unorm-storage') ? 'bgra8unorm' : 'rgba8unorm'; var outShape = [height, width]; var program = new DrawProgram(outShape, image.dtype, format); canvas.width = width; canvas.height = height; var backendName = 'webgpu'; var gpuContext = canvas.getContext(backendName); var canvasWebGPU; if (!gpuContext) { canvasWebGPU = new OffscreenCanvas(width, height); gpuContext = canvasWebGPU.getContext(backendName); } var numChannels = image.shape.length === 3 ? image.shape[2] : 1; gpuContext.configure({ device: backend.device, format: format, usage: GPUTextureUsage.STORAGE_BINDING, alphaMode: 'premultiplied' }); var outputDtype = 'int32'; var output = backend.makeTensorInfo(outShape, outputDtype); var info = backend.tensorMap.get(output.dataId); info.resource = gpuContext.getCurrentTexture(); info.external = true; var uniformData = [{ type: 'uint32', data: [numChannels] }, { type: 'float32', data: [alpha] }]; backend.runWebGPUProgram(program, [image], outputDtype, uniformData, output); if (canvasWebGPU) { var canvas2dContext = canvas.getContext('2d'); if (!canvas2dContext) { throw new Error("Please make sure this canvas has only been used for 2d or webgpu context!"); } canvas2dContext.drawImage(canvasWebGPU, 0, 0); } backend.disposeData(output.dataId); return image; } var drawConfig = { kernelName: tf.Draw, backendName: 'webgpu', kernelFunc: draw }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var multiplyKernelFunc = binaryKernelFunc({ opType: BinaryOpType.MUL, cpuKernelImpl: multiplyImplCPU, supportsComplex: true }); var multiplyConfig = { kernelName: tf.Multiply, backendName: 'webgpu', kernelFunc: multiplyKernelFunc }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function sum(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var axis = attrs.axis, keepDims = attrs.keepDims; return reduce(x, axis, keepDims, 'sum', backend); } var sumConfig = { kernelName: tf.Sum, backendName: 'webgpu', kernelFunc: sum }; function einsum(args) { var e_1, _a, e_2, _b; var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var equation = attrs.equation; var tensors = inputs; var _c = tf.backend_util.decodeEinsumEquation(equation, tensors.length), allDims = _c.allDims, summedDims = _c.summedDims, idDims = _c.idDims; tf.backend_util.checkEinsumDimSizes(allDims.length, idDims, tensors); var _d = tf.backend_util.getEinsumComputePath(summedDims, idDims), path = _d.path, steps = _d.steps; var nSteps = steps.length; var out = null; var numDimsRemaining = allDims.length; var tensorsToDispose = []; for (var i = 0; i < nSteps; ++i) { try { for (var _e = (e_1 = void 0, __values(steps[i])), _f = _e.next(); !_f.done; _f = _e.next()) { var idTerm = _f.value; var _g = tf.backend_util.getEinsumPermutation(numDimsRemaining, idDims[idTerm]), perm = _g.permutationIndices, dimsToExpand = _g.expandDims; var x = void 0; if (tf.backend_util.isIdentityPermutation(perm)) { x = tensors[idTerm]; } else { x = transpose({ inputs: { x: tensors[idTerm] }, backend: backend, attrs: { perm: perm } }); tensorsToDispose.push(x); } var targetShape = x.shape.slice(); for (var k = 0; k < dimsToExpand.length; ++k) { targetShape.splice(dimsToExpand[k], 0, 1); } if (!tf.util.arraysEqual(x.shape, targetShape)) { x = reshape({ inputs: { x: x }, backend: backend, attrs: { shape: targetShape } }); tensorsToDispose.push(x); } if (out === null) { out = x; } else { // tslint:disable-next-line: no-unnecessary-type-assertion out = multiplyKernelFunc({ inputs: { a: x, b: out }, backend: backend }); tensorsToDispose.push(out); } } } catch (e_1_1) { e_1 = { error: e_1_1 }; } finally { try { if (_f && !_f.done && (_a = _e.return)) _a.call(_e); } finally { if (e_1) throw e_1.error; } } if (i < nSteps - 1) { if (path[i] >= 0) { out = sum({ inputs: { x: out }, backend: backend, attrs: { axis: path[i] - (allDims.length - numDimsRemaining), keepDims: false } }); tensorsToDispose.push(out); } numDimsRemaining--; } } try { // Clean up intermediate tensors. for (var tensorsToDispose_1 = __values(tensorsToDispose), tensorsToDispose_1_1 = tensorsToDispose_1.next(); !tensorsToDispose_1_1.done; tensorsToDispose_1_1 = tensorsToDispose_1.next()) { var tensorInfo = tensorsToDispose_1_1.value; if (tensorInfo === out) { continue; } backend.disposeData(tensorInfo.dataId); } } catch (e_2_1) { e_2 = { error: e_2_1 }; } finally { try { if (tensorsToDispose_1_1 && !tensorsToDispose_1_1.done && (_b = tensorsToDispose_1.return)) _b.call(tensorsToDispose_1); } finally { if (e_2) throw e_2.error; } } return out; } var einsumConfig = { kernelName: tf.Einsum, backendName: 'webgpu', kernelFunc: einsum }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var elu = unaryKernelFunc({ opType: UnaryOpType.ELU }); var eluConfig = { kernelName: tf.Elu, backendName: 'webgpu', kernelFunc: elu }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var eluGrad = function (args) { var inputs = args.inputs, backend = args.backend; var dy = inputs.dy, y = inputs.y; var program = new BinaryOpProgram(BinaryOpType.ELU_DER, dy.shape, y.shape); return backend.runWebGPUProgram(program, [dy, y], dy.dtype); }; var eluGradConfig = { kernelName: tf.EluGrad, backendName: 'webgpu', kernelFunc: eluGrad }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var equal = binaryKernelFunc({ opType: BinaryOpType.EQUAL, dtype: 'bool', cpuKernelImpl: equalImplCPU }); var equalConfig = { kernelName: tf.Equal, backendName: 'webgpu', kernelFunc: equal }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var erf = unaryKernelFunc({ opType: UnaryOpType.ERF }); var erfConfig = { kernelName: tf.Erf, backendName: 'webgpu', kernelFunc: erf }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var exp = unaryKernelFunc({ opType: UnaryOpType.EXP, cpuKernelImpl: expImplCPU, dtype: 'float32', }); var expConfig = { kernelName: tf.Exp, backendName: 'webgpu', kernelFunc: exp }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the License); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an AS IS BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function expandDims(args) { var inputs = args.inputs, attrs = args.attrs, backend = args.backend; var dim = attrs.dim; var input = inputs.input; var inputRank = input.shape.length; var newShape = input.shape.slice(); var $dim = dim; if (dim < 0) { // Negative value is counted from the tail of rank. tf.util.assert(-(inputRank + 1) <= dim, function () { return "Axis must be in the interval [".concat(-(inputRank + 1), ", ").concat(inputRank, "]"); }); $dim = inputRank + dim + 1; } newShape.splice($dim, 0, 1); return reshape({ inputs: { x: input }, backend: backend, attrs: { shape: newShape } }); } var expandDimsConfig = { kernelName: tf.ExpandDims, backendName: 'webgpu', kernelFunc: expandDims, }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var expm1 = unaryKernelFunc({ opType: UnaryOpType.EXPM1, cpuKernelImpl: expm1ImplCPU }); var expm1Config = { kernelName: tf.Expm1, backendName: 'webgpu', kernelFunc: expm1 }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var FFTProgram = /** @class */ (function () { function FFTProgram(component, shape) { this.variableNames = ['real', 'imag']; this.outputShape = []; this.uniforms = 'exponentMultiplier : f32, denominator: f32,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = shape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.component = component; this.shaderKey = "fft_".concat(component); } FFTProgram.prototype.getUserCode = function () { var opString = this.component === 'real' ? 'return real * expR - imag * expI;' : 'return real * expI + imag * expR;'; var userCode = "\n fn unaryOpComplex(real: f32, expR: f32, imag: f32, expI: f32) -> f32 {\n ".concat(opString, "\n }\n\n fn mulMatDFT(batch: i32, index: i32) -> f32 {\n let indexRatio = f32(index) / f32(uniforms.realShape[1]);\n let exponentMultiplierTimesIndexRatio =\n uniforms.exponentMultiplier * indexRatio;\n\n var result = 0.0;\n\n for (var i = 0; i < uniforms.realShape[1]; i = i + 1) {\n // x = (-2|2 * PI / N) * index * i;\n let x = exponentMultiplierTimesIndexRatio * f32(i);\n let expR = cos(x);\n let expI = sin(x);\n let real = getReal(batch, i);\n let imag = getImag(batch, i);\n\n result = result +\n unaryOpComplex(real, expR, imag, expI) / uniforms.denominator;\n }\n\n return result;\n }\n\n ").concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getOutputCoords();\n setOutputAtIndex(index, mulMatDFT(coords[0], coords[1]));\n }\n }\n "); return userCode; }; return FFTProgram; }()); /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function fftImpl(x, inverse, backend) { var xData = backend.tensorMap.get(x.dataId); var inputSize = tf.util.sizeFromShape(x.shape); // Collapse all outer dimensions to a single batch dimension. var innerDimensionSize = x.shape[x.shape.length - 1]; var batch = inputSize / innerDimensionSize; var toDispose = []; var input2D = reshape({ inputs: { x: x }, backend: backend, attrs: { shape: [batch, innerDimensionSize] } }); toDispose.push(input2D); var xShape = input2D.shape; var realProgram = new FFTProgram('real', xShape); var imagProgram = new FFTProgram('imag', xShape); var inputs = [ { dataId: xData.complexTensorInfos.real.dataId, dtype: xData.complexTensorInfos.real.dtype, shape: xShape }, { dataId: xData.complexTensorInfos.imag.dataId, dtype: xData.complexTensorInfos.imag.dtype, shape: xShape } ]; var exponentMultiplier = inverse ? 2.0 * Math.PI : -2.0 * Math.PI; var denominator = inverse ? xShape[1] : 1.0; var uniformData = [ { type: 'float32', data: [exponentMultiplier] }, { type: 'float32', data: [denominator] } ]; var realPart = backend.runWebGPUProgram(realProgram, inputs, 'float32', uniformData); toDispose.push(realPart); var imagPart = backend.runWebGPUProgram(imagProgram, inputs, 'float32', uniformData); toDispose.push(imagPart); var complexOutput = complex({ inputs: { real: realPart, imag: imagPart }, backend: backend }); toDispose.push(complexOutput); var complexOutputReshaped = reshape({ inputs: { x: complexOutput }, backend: backend, attrs: { shape: x.shape } }); toDispose.forEach(function (t) { return backend.disposeData(t.dataId); }); return complexOutputReshaped; } /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function fft(args) { var inputs = args.inputs, backend = args.backend; var input = inputs.input; return fftImpl(input, false /* inverse */, backend); } var fftConfig = { kernelName: tf.FFT, backendName: 'webgpu', kernelFunc: fft }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var FlipLeftRightProgram = /** @class */ (function () { function FlipLeftRightProgram(imageShape) { this.outputShape = []; this.variableNames = ['x']; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = imageShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'flipLeftRight'; } FlipLeftRightProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let coordX = uniforms.xShape[2] - coords[2] - 1;\n let outputValue = getX(coords[0], coords[1], coordX, coords[3]);\n setOutputAtIndex(index, outputValue);\n }\n }\n "); return userCode; }; return FlipLeftRightProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var flipLeftRightConfig = { kernelName: tf.FlipLeftRight, backendName: 'webgpu', kernelFunc: function (_a) { var inputs = _a.inputs, backend = _a.backend; var image = inputs.image; var webgpuBackend = backend; var program = new FlipLeftRightProgram(image.shape); var output = webgpuBackend.runWebGPUProgram(program, [image], image.dtype); return output; } }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var floor = unaryKernelFunc({ opType: UnaryOpType.FLOOR, cpuKernelImpl: floorImplCPU }); var floorConfig = { kernelName: tf.Floor, backendName: 'webgpu', kernelFunc: floor }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var floorDiv = binaryKernelFunc({ opType: BinaryOpType.FLOOR_DIV, cpuKernelImpl: floorDivImplCPU, dtype: 'int32' }); var floorDivConfig = { kernelName: tf.FloorDiv, backendName: 'webgpu', kernelFunc: floorDiv }; /** * @license * Copyright 2022 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var FromPixelsProgram = /** @class */ (function () { function FromPixelsProgram(outputShape, numChannels, importVideo) { if (importVideo === void 0) { importVideo = false; } this.pixelsOpType = PixelsOpType.FROM_PIXELS; this.outputShape = [0]; this.variableNames = []; this.workgroupSize = [256, 1, 1]; // The empirical value. this.outputShape = outputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize, [numChannels, 1, 1]); this.importVideo = importVideo; this.shaderKey = "fromPixels_".concat(this.importVideo); } FromPixelsProgram.prototype.getUserCode = function () { var textureLoad = this.importVideo ? 'textureLoad(src, vec2(coords.yx));' : 'textureLoad(src, vec2(coords.yx), 0)'; var textureType = this.importVideo ? 'texture_external' : 'texture_2d'; return "\n @binding(1) @group(0) var src: ".concat(textureType, ";\n ").concat(getMainHeaderString('index'), " {\n let flatIndex = index * uniforms.numChannels;\n if (flatIndex < uniforms.size) {\n let coords = getCoordsFromIndex(flatIndex);\n let values = ").concat(textureLoad, ";\n for (var i = 0; i < uniforms.numChannels; i = i + 1) {\n result[flatIndex + i] = i32(floor(255.0 * values[i]));\n }\n }\n }\n "); }; return FromPixelsProgram; }()); var fromPixelsConfig = { kernelName: tf.FromPixels, backendName: 'webgpu', kernelFunc: fromPixels, }; var fromPixels2DContext; var willReadFrequently = tf.env().getBool('CANVAS2D_WILL_READ_FREQUENTLY_FOR_GPU'); function fromPixels(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var pixels = inputs.pixels; var numChannels = attrs.numChannels; if (pixels == null) { throw new Error('pixels passed to tf.browser.fromPixels() can not be null'); } var isVideo = typeof (HTMLVideoElement) !== 'undefined' && pixels instanceof HTMLVideoElement; var isImage = typeof (HTMLImageElement) !== 'undefined' && pixels instanceof HTMLImageElement; var isCanvas = (typeof (HTMLCanvasElement) !== 'undefined' && pixels instanceof HTMLCanvasElement) || (typeof (OffscreenCanvas) !== 'undefined' && pixels instanceof OffscreenCanvas); var isImageBitmap = typeof (ImageBitmap) !== 'undefined' && pixels instanceof ImageBitmap; var _a = __read(isVideo ? [ pixels.videoWidth, pixels.videoHeight ] : [pixels.width, pixels.height], 2), width = _a[0], height = _a[1]; var outputShape = [height, width, numChannels]; var importVideo = tf.env().getBool('WEBGPU_IMPORT_EXTERNAL_TEXTURE') && isVideo; var isVideoOrImage = isVideo || isImage; if (isImageBitmap || isCanvas || isVideoOrImage) { var resource = void 0; if (importVideo) { resource = backend.device.importExternalTexture({ source: pixels }); } else { if (isVideoOrImage) { var newWillReadFrequently = tf.env().getBool('CANVAS2D_WILL_READ_FREQUENTLY_FOR_GPU'); if (fromPixels2DContext == null || newWillReadFrequently !== willReadFrequently) { willReadFrequently = newWillReadFrequently; fromPixels2DContext = document.createElement('canvas').getContext('2d', { willReadFrequently: willReadFrequently }); } fromPixels2DContext.canvas.width = width; fromPixels2DContext.canvas.height = height; fromPixels2DContext.drawImage(pixels, 0, 0, width, height); pixels = fromPixels2DContext.canvas; } var usage = GPUTextureUsage.COPY_DST | GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.TEXTURE_BINDING; var format = 'rgba8unorm'; var texture = backend.textureManager.acquireTexture(outputShape[1], outputShape[0], format, usage); backend.queue.copyExternalImageToTexture({ source: pixels }, { texture: texture }, [outputShape[1], outputShape[0]]); resource = texture; } var size = tf.util.sizeFromShape(outputShape); var strides = tf.util.computeStrides(outputShape); var program = new FromPixelsProgram(outputShape, numChannels, importVideo); var uniformData = [ { type: 'uint32', data: [size] }, { type: 'uint32', data: [numChannels] }, { type: 'uint32', data: __spreadArray([], __read(strides), false) } ]; var input = backend.makeTensorInfo([height, width], 'int32'); var info = backend.tensorMap.get(input.dataId); info.resource = resource; var result = backend.runWebGPUProgram(program, [input], 'int32', uniformData); backend.disposeData(input.dataId); return result; } // TODO: Encoding should happen on GPU once we no longer have to download // image data to the CPU. var imageData = pixels.data; var pixelArray = imageData; if (numChannels != null && numChannels !== 4) { pixelArray = new Uint8Array(pixels.width * pixels.height * numChannels); var dataLength = imageData.length; var j = 0; for (var i = 0; i < dataLength; i++) { if (i % 4 < numChannels) { pixelArray[j++] = imageData[i]; } } } var output = backend.makeTensorInfo(outputShape, 'int32', new Int32Array(pixelArray)); backend.uploadToGPU(output.dataId); return output; } /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var BatchNormProgram = /** @class */ (function () { function BatchNormProgram(xShape, meanShape, varianceShape, offsetShape, scaleShape) { this.uniforms = 'varianceEpsilon : f32,'; // This is an experimental value. this.workgroupSize = [128, 1, 1]; this.size = true; this.variableNames = ['x', 'mean', 'variance']; tf.backend_util.assertAndGetBroadcastShape(xShape, meanShape); tf.backend_util.assertAndGetBroadcastShape(xShape, varianceShape); this.outputShape = xShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); if (offsetShape != null) { tf.backend_util.assertAndGetBroadcastShape(xShape, offsetShape); this.variableNames.push('offset'); } if (scaleShape != null) { tf.backend_util.assertAndGetBroadcastShape(xShape, scaleShape); this.variableNames.push('scale'); } this.offsetShape = offsetShape; this.scaleShape = scaleShape; this.shaderKey = 'batchNorm'; } BatchNormProgram.prototype.getUserCode = function () { var offsetSnippet = '0.0'; if (this.offsetShape != null) { offsetSnippet = 'getOffsetByOutputIndex(index)'; } var scaleSnippet = '1.0'; if (this.scaleShape != null) { scaleSnippet = 'getScaleByOutputIndex(index)'; } var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size)\n {\n let xValue = getXByOutputIndex(index);\n let meanValue = getMeanByOutputIndex(index);\n let varianValue = getVarianceByOutputIndex(index);\n let offsetValue = ").concat(offsetSnippet, ";\n let scaleValue = ").concat(scaleSnippet, ";\n let inv = scaleValue * inverseSqrt(varianValue + f32(uniforms.varianceEpsilon));\n setOutputAtIndex(index,dot(vec3(xValue, -meanValue, offsetValue), vec3(inv, inv, 1.0)));\n }\n }\n "); return userCode; }; return BatchNormProgram; }()); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var fusedBatchNormConfig = { kernelName: tf.FusedBatchNorm, backendName: 'webgpu', kernelFunc: function (_a) { var inputs = _a.inputs, attrs = _a.attrs, backend = _a.backend; var x = inputs.x, scale = inputs.scale, offset = inputs.offset, mean = inputs.mean, variance = inputs.variance; var varianceEpsilon = attrs.varianceEpsilon; var webGPUBackend = backend; var batchNormInputs = [x, mean, variance]; var offsetShape = null; if (offset != null) { offsetShape = offset.shape; batchNormInputs.push(offset); } var scaleShape = null; if (scale != null) { scaleShape = scale.shape; batchNormInputs.push(scale); } var program = new BatchNormProgram(x.shape, mean.shape, variance.shape, offsetShape, scaleShape); var uniformData = [{ type: 'float32', data: [varianceEpsilon] }]; return webGPUBackend.runWebGPUProgram(program, batchNormInputs, x.dtype, uniformData); } }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function fusedConv2d(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, filter = inputs.filter, bias = inputs.bias, preluActivationWeights = inputs.preluActivationWeights; var strides = attrs.strides, pad = attrs.pad, dataFormat = attrs.dataFormat, dilations = attrs.dilations, dimRoundingMode = attrs.dimRoundingMode, activation = attrs.activation, leakyreluAlpha = attrs.leakyreluAlpha; var $dataFormat = tf.backend_util.convertConv2DDataFormat(dataFormat); var convInfo = tf.backend_util.computeConv2DInfo(x.shape, filter.shape, strides, dilations, pad, dimRoundingMode, false /* depthwise */, $dataFormat); return conv2DImpl({ x: x, filter: filter, convInfo: convInfo, backend: backend, bias: bias, preluActivationWeights: preluActivationWeights, leakyreluAlpha: leakyreluAlpha, activation: activation }); } var fusedConv2DConfig = { kernelName: tf.FusedConv2D, backendName: 'webgpu', kernelFunc: fusedConv2d, }; /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function fusedDepthwiseConv2D(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, filter = inputs.filter, bias = inputs.bias, preluActivationWeights = inputs.preluActivationWeights; var strides = attrs.strides, pad = attrs.pad, dilations = attrs.dilations, dimRoundingMode = attrs.dimRoundingMode, activation = attrs.activation, leakyreluAlpha = attrs.leakyreluAlpha; var $dilations = dilations; if ($dilations == null) { $dilations = [1, 1]; } tf.util.assert(tf.backend_util.eitherStridesOrDilationsAreOne(strides, $dilations), function () { return 'Error in depthwiseConv2d: Either strides or dilations must be ' + "1. Got strides ".concat(strides, " and dilations '").concat($dilations, "'"); }); var convInfo = tf.backend_util.computeConv2DInfo(x.shape, filter.shape, strides, $dilations, pad, dimRoundingMode, true /* depthwise */); var programInputs = [x, filter]; var hasBias = bias != null; var hasPreluActivationWeights = preluActivationWeights != null; if (hasBias) { programInputs.push(bias); } if (hasPreluActivationWeights) { programInputs.push(preluActivationWeights); } var dimensions = [ { type: 'int32', data: [convInfo.padInfo.top, convInfo.padInfo.left] }, { type: 'int32', data: [convInfo.inHeight, convInfo.inWidth] }, ]; var program; if (convInfo.outHeight > 4 && convInfo.outWidth > 4 && convInfo.strideWidth <= 2 && convInfo.inChannels === convInfo.outChannels && convInfo.dilationHeight === 1 && convInfo.dilationWidth === 1 && convInfo.inChannels % 4 === 0) { program = new DepthwiseConv2DVec4Program(convInfo, hasBias, activation, hasPreluActivationWeights); dimensions.push({ type: 'int32', data: [program.virtualWidth] }); } else { program = new DepthwiseConv2DProgram(convInfo, hasBias, activation, hasPreluActivationWeights); dimensions.push({ type: 'int32', data: [convInfo.filterHeight] }, { type: 'int32', data: [convInfo.filterWidth] }, { type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.dilationHeight, convInfo.dilationWidth] }); } if (activation === 'leakyrelu') { dimensions.push({ type: 'float32', data: [leakyreluAlpha] }); program.uniforms += ' alpha : f32,'; } var result = backend.runWebGPUProgram(program, programInputs, 'float32', dimensions); return result; } var fusedDepthwiseConv2DConfig = { kernelName: tf.FusedDepthwiseConv2D, backendName: 'webgpu', kernelFunc: fusedDepthwiseConv2D, }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var GatherNDProgram = /** @class */ (function () { function GatherNDProgram(sliceDim, shape) { this.variableNames = ['A', 'indices']; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = shape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = "gathernd_".concat(sliceDim); this.sliceDim = sliceDim; this.uniforms = "sliceDim : i32, strides : ".concat(getCoordsDataType(sliceDim), ","); } GatherNDProgram.prototype.getUserCode = function () { var strideString; if (this.sliceDim > 1) { strideString = 'uniforms.strides[j]'; } else { strideString = 'uniforms.strides'; } var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n var flattenIndex = 0;\n for (var j = 0; j < uniforms.sliceDim; j = j + 1) {\n let indexTemp = i32(round(getIndices(coords[0], j)));\n let strideNum = ").concat(strideString, ";\n flattenIndex = flattenIndex + indexTemp * strideNum;\n }\n\n setOutputAtIndex(index, getA(flattenIndex, coords[1]));\n }\n }\n "); return userCode; }; return GatherNDProgram; }()); function gatherNd(args) { var inputs = args.inputs, backend = args.backend; var params = inputs.params, indices = inputs.indices; var indicesShape = indices.shape; var sliceRank = indicesShape[indicesShape.length - 1]; var paramsSize = tf.util.sizeFromShape(params.shape); var _a = __read(tf.backend_util.prepareAndValidate(params, indices), 4), resultShape = _a[0], numSlices = _a[1], sliceSize = _a[2], strides = _a[3]; var flattenIndices = reshape({ inputs: { x: indices }, backend: backend, attrs: { shape: [numSlices, sliceRank] } }); var flattenX = reshape({ inputs: { x: params }, backend: backend, attrs: { shape: [(tf.util.sizeFromShape(params.shape) / sliceSize), sliceSize] } }); if (backend.shouldExecuteOnCPU([params, indices]) || params.dtype === 'string') { var indicesData = backend.readSync(indices.dataId); var paramsBuf = backend.bufferSync(params); var outValue = gatherNdImplCPU(indicesData, paramsBuf, params.dtype, numSlices, sliceRank, sliceSize, strides, params.shape, paramsSize); return backend.makeTensorInfo(resultShape, params.dtype, outValue.values); } var program = new GatherNDProgram(sliceRank, [numSlices, sliceSize]); var uniformData = [{ type: 'int32', data: [sliceRank] }, { type: 'int32', data: strides }]; var res = backend.runWebGPUProgram(program, [flattenX, flattenIndices], flattenX.dtype, uniformData); var reshaped = reshape({ inputs: { x: res }, backend: backend, attrs: { shape: resultShape } }); backend.disposeData(flattenIndices.dataId); backend.disposeData(flattenX.dataId); backend.disposeData(res.dataId); return reshaped; } var gatherNdConfig = { kernelName: tf.GatherNd, backendName: 'webgpu', kernelFunc: gatherNd }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var GatherProgram = /** @class */ (function () { function GatherProgram(aShape, outputShape) { this.variableNames = ['A', 'indices']; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = aShape.slice(); this.aShape = aShape; this.outputShape = outputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = "gather"; } GatherProgram.prototype.getUserCode = function () { var sourceCoords = getSourceCoords$1(this.aShape); var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let resRC = getCoordsFromIndex(index);\n let indexZ = i32(getIndices(resRC.x, resRC.z));\n let inBounds = select(0.0, 1.0, indexZ >= 0 && indexZ < uniforms.aShape[2]);\n setOutputAtIndex(index, inBounds * getA(").concat(sourceCoords, "));\n }\n }\n "); return userCode; }; return GatherProgram; }()); // The input and output are always flattened into rank 4 tensors. function getSourceCoords$1(aShape) { var currentCoords = ['resRC.x', 'resRC.y', 'resRC.z', 'resRC.w']; var sourceCoords = []; for (var i = 0; i < aShape.length; i++) { if (i === 2) { sourceCoords.push('indexZ'); } else { sourceCoords.push("".concat(currentCoords[i])); } } return sourceCoords.join(); } /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function gatherV2(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, indices = inputs.indices; var axis = attrs.axis, batchDims = attrs.batchDims; // Unlike WebGL, WebGPU won't check if index is out of bound by calling // backend.readSync() function in debug mode. var parsedAxis = tf.util.parseAxisParam(axis, x.shape)[0]; var shapeInfo = tf.backend_util.segment_util.collectGatherOpShapeInfo(x, indices, parsedAxis, batchDims); var indicesSize = tf.util.sizeFromShape(indices.shape); var toDispose = []; var flattenX = reshape({ inputs: { x: x }, backend: backend, attrs: { shape: [ shapeInfo.batchSize, shapeInfo.outerSize, shapeInfo.dimSize, shapeInfo.sliceSize ] } }); var flattenIndex = reshape({ inputs: { x: indices }, backend: backend, attrs: { shape: [shapeInfo.batchSize, indicesSize / shapeInfo.batchSize] } }); toDispose.push(flattenX); toDispose.push(flattenIndex); var flattenOutputShape = [ shapeInfo.batchSize, shapeInfo.outerSize, indicesSize / shapeInfo.batchSize, shapeInfo.sliceSize ]; if (backend.shouldExecuteOnCPU([x, indices])) { var indicesTensorData = backend.tensorMap.get(flattenIndex.dataId); var indicesValues = indicesTensorData.values; var indicesBuffer = tf.buffer(flattenIndex.shape, flattenIndex.dtype, indicesValues); var flattenXTensorData = backend.tensorMap.get(flattenX.dataId); var xValues = flattenXTensorData.values; var xBuffer = tf.buffer(flattenX.shape, flattenX.dtype, xValues); var outBuf = gatherV2ImplCPU(xBuffer, indicesBuffer, flattenOutputShape); toDispose.forEach(function (t) { return backend.disposeData(t.dataId); }); return backend.makeTensorInfo(shapeInfo.outputShape, outBuf.dtype, outBuf.values); } var program = new GatherProgram(flattenX.shape, flattenOutputShape); var res = backend.runWebGPUProgram(program, [flattenX, flattenIndex], flattenX.dtype); toDispose.push(res); var reshaped = reshape({ inputs: { x: res }, backend: backend, attrs: { shape: shapeInfo.outputShape } }); toDispose.forEach(function (t) { return backend.disposeData(t.dataId); }); return reshaped; } var gatherV2Config = { kernelName: tf.GatherV2, backendName: 'webgpu', kernelFunc: gatherV2 }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var greater = binaryKernelFunc({ opType: BinaryOpType.GREATER, cpuKernelImpl: greaterImplCPU, dtype: 'bool', }); var greaterConfig = { kernelName: tf.Greater, backendName: 'webgpu', kernelFunc: greater }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var greaterEqual = binaryKernelFunc({ opType: BinaryOpType.GREATER_EQUAL, dtype: 'bool', cpuKernelImpl: greaterEqualImplCPU }); var greaterEqualConfig = { kernelName: tf.GreaterEqual, backendName: 'webgpu', kernelFunc: greaterEqual }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function ifft(args) { var inputs = args.inputs, backend = args.backend; var input = inputs.input; return fftImpl(input, true /* inverse */, backend); } var ifftConfig = { kernelName: tf.IFFT, backendName: 'webgpu', kernelFunc: ifft }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var isFinite = unaryKernelFunc({ opType: UnaryOpType.IS_FINITE, dtype: 'bool' }); var isFiniteConfig = { kernelName: tf.IsFinite, backendName: 'webgpu', kernelFunc: isFinite }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var isInf = unaryKernelFunc({ opType: UnaryOpType.IS_INF, dtype: 'bool' }); var isInfConfig = { kernelName: tf.IsInf, backendName: 'webgpu', kernelFunc: isInf }; /** * @license * Copyright 2022 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var isNaN = unaryKernelFunc({ opType: UnaryOpType.IS_NAN, dtype: 'bool' }); var isNaNConfig = { kernelName: tf.IsNan, backendName: 'webgpu', kernelFunc: isNaN }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function leakyRelu(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var alpha = attrs.alpha; var uniformData = [{ type: 'float32', data: [alpha] }]; var program = new UnaryOpProgram(x.shape, UnaryOpType.LEAKYRELU, 'alpha : f32,'); return backend.runWebGPUProgram(program, [x], 'float32', uniformData); } var leakyReluConfig = { kernelName: tf.LeakyRelu, backendName: 'webgpu', kernelFunc: leakyRelu }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var less = binaryKernelFunc({ opType: BinaryOpType.LESS, dtype: 'bool', cpuKernelImpl: lessImplCPU }); var lessConfig = { kernelName: tf.Less, backendName: 'webgpu', kernelFunc: less }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var lessEqual = binaryKernelFunc({ opType: BinaryOpType.LESS_EQUAL, dtype: 'bool', cpuKernelImpl: lessEqualImplCPU }); var lessEqualConfig = { kernelName: tf.LessEqual, backendName: 'webgpu', kernelFunc: lessEqual }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var LinSpaceProgram = /** @class */ (function () { function LinSpaceProgram(shape) { this.variableNames = []; this.outputShape = []; this.uniforms = 'start : f32, step : f32,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = [shape]; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'linSpace'; } LinSpaceProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n setOutputAtIndex(index, uniforms.start + f32(index) * uniforms.step);\n }\n }\n "); return userCode; }; return LinSpaceProgram; }()); /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function linSpace(args) { var backend = args.backend, attrs = args.attrs; var start = attrs.start, stop = attrs.stop, num = attrs.num; var step = (stop - start) / (num - 1); var program = new LinSpaceProgram(num); var uniformData = [{ type: 'float32', data: [start] }, { type: 'float32', data: [step] }]; return backend.runWebGPUProgram(program, [], 'float32', uniformData); } var linSpaceConfig = { kernelName: tf.LinSpace, backendName: 'webgpu', kernelFunc: linSpace }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var log = unaryKernelFunc({ opType: UnaryOpType.LOG, cpuKernelImpl: logImplCPU }); var logConfig = { kernelName: tf.Log, backendName: 'webgpu', kernelFunc: log }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var log1p = unaryKernelFunc({ opType: UnaryOpType.LOG1P }); var log1pConfig = { kernelName: tf.Log1p, backendName: 'webgpu', kernelFunc: log1p }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var logicalAnd = binaryKernelFunc({ opType: BinaryOpType.LOGICAL_AND, dtype: 'bool' }); var logicalAndConfig = { kernelName: tf.LogicalAnd, backendName: 'webgpu', kernelFunc: logicalAnd }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var logicalNot = unaryKernelFunc({ opType: UnaryOpType.LOGICAL_NOT }); var logicalNotConfig = { kernelName: tf.LogicalNot, backendName: 'webgpu', kernelFunc: logicalNot }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var logicalOr = binaryKernelFunc({ opType: BinaryOpType.LOGICAL_OR }); var logicalOrConfig = { kernelName: tf.LogicalOr, backendName: 'webgpu', kernelFunc: logicalOr }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var powOperatorSnippet = "\n var powValue = 0.0;\n let basis = uniforms.bias + uniforms.alpha * sum;\n if (uniforms.beta == 0.5) {\n powValue = inverseSqrt(basis);\n } else if (uniforms.beta == 1.0) {\n powValue = 1.0 / basis;\n } else {\n powValue = exp(log(basis) * (-uniforms.beta));\n }\n"; var LRNProgram = /** @class */ (function () { function LRNProgram(xShape) { this.outputShape = []; this.variableNames = ['x']; this.uniforms = 'radius : i32, bias : f32, alpha : f32, beta : f32,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = xShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'lrn'; } LRNProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getOutputCoords();\n let b = coords[0];\n let r = coords[1];\n let c = coords[2];\n let d = coords[3];\n\n let x = getX(b, r, c, d);\n var sum = 0.0;\n for (var i = -uniforms.radius; i <= uniforms.radius; i = i + 1) {\n let idx = d + i;\n if (idx >= 0 && idx < uniforms.xShape[3]) {\n let z = getX(b, r, c, idx);\n sum = sum + z * z;\n }\n }\n ").concat(powOperatorSnippet, "\n\n setOutputAtIndex(index, x * powValue);\n }\n }\n "); return userCode; }; return LRNProgram; }()); var LRNSharedProgram = /** @class */ (function () { function LRNSharedProgram(xShape, radius) { var _this = this; this.outputShape = []; this.variableNames = ['x']; this.uniforms = 'radius : i32, bias : f32, alpha : f32, beta : f32,'; this.workgroupSize = [256, 1, 1]; this.maxAllowRadius = 16; tf.util.assert(radius <= this.maxAllowRadius, function () { return "Radius must be less than or equal to ".concat(_this.maxAllowRadius, ", current radius is ").concat(radius); }); this.outputShape = xShape; // The reason why not using this.workgroupSize[0] + 2 * maxAllowRadius here // is to make sure that there is only one time global memory load access for // each thread. this.elementsPerWorkgroup = this.workgroupSize[0] - 2 * this.maxAllowRadius; this.dispatchLayout = { x: [3], y: [2], z: [0, 1] }; this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, [ this.elementsPerWorkgroup, this.workgroupSize[1], this.workgroupSize[2] ]); this.shaderKey = 'lrn_shared'; } LRNSharedProgram.prototype.getUserCode = function () { var userCode = "\n var lrnSub: array;\n const elementsPerWorkgroup = ").concat(this.elementsPerWorkgroup, ";\n const maxAllowRadius = ").concat(this.maxAllowRadius, ";\n\n ").concat(getMainHeaderString(), " {\n let localDepth = i32(localId.x);\n let workgroupDepth = i32(workgroupId.x) * elementsPerWorkgroup;\n let xDepth = workgroupDepth + localDepth - maxAllowRadius;\n let b = i32(globalId.z) / uniforms.xShape[1];\n let r = i32(globalId.z) - b * uniforms.xShape[1];\n let c = i32(globalId.y);\n let d = workgroupDepth + localDepth;\n\n var x = 0.0;\n if (xDepth >= 0 && xDepth < uniforms.xShape[3]) {\n x = getX(b, r, c, xDepth);\n }\n lrnSub[localDepth] = x;\n workgroupBarrier();\n\n if (localDepth < elementsPerWorkgroup && d < uniforms.outShape[3]) {\n var sum = 0.0;\n let index = localDepth + maxAllowRadius;\n for (var i = -uniforms.radius; i <= uniforms.radius; i = i + 1) {\n let z = lrnSub[index + i];\n sum = sum + z * z;\n }\n ").concat(powOperatorSnippet, "\n\n setOutputAtCoords(b, r, c, d, lrnSub[index] * powValue);\n }\n } "); return userCode; }; return LRNSharedProgram; }()); /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function lrn(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var depthRadius = attrs.depthRadius, bias = attrs.bias, alpha = attrs.alpha, beta = attrs.beta; // When the adjacent channels is less than or equal to 16, which could cover // most cases, we use shared memory version to get better performance. // The theoretical adjacent channels may be very large, but the shared memory // size of hardware is limited, so we use the naive version when the adjacent // channels is large. var program; if (depthRadius > 16) { program = new LRNProgram(x.shape); } else { program = new LRNSharedProgram(x.shape, depthRadius); } var uniformData = [ { type: 'int32', data: [depthRadius] }, { type: 'float32', data: [bias] }, { type: 'float32', data: [alpha] }, { type: 'float32', data: [beta] } ]; var res = backend.runWebGPUProgram(program, [x], x.dtype, uniformData); return res; } var lrnConfig = { kernelName: tf.LRN, backendName: 'webgpu', kernelFunc: lrn }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var LRNGradProgram = /** @class */ (function () { function LRNGradProgram(inputShape) { this.outputShape = []; this.variableNames = ['inputImage', 'outputImage', 'dy']; this.uniforms = 'depthRadius : i32, bias : f32, alpha : f32, beta : f32,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = inputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'lrn_grad'; } LRNGradProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getOutputCoords();\n let b = coords[0];\n let r = coords[1];\n let c = coords[2];\n\n let MIN_DEPTH_BEGIN = 0;\n let MAX_DEPTH_END = uniforms.outShape[3];\n var result = 0.0;\n for (var d = MIN_DEPTH_BEGIN; d < MAX_DEPTH_END; d++) {\n let depthBegin = max(MIN_DEPTH_BEGIN, d - uniforms.depthRadius);\n let depthEnd = min(MAX_DEPTH_END, d + uniforms.depthRadius + 1);\n\n var norm = 0.0;\n for (var k = MIN_DEPTH_BEGIN; k < MAX_DEPTH_END; k++) {\n if (k < depthBegin) {\n continue;\n } else if (k >= depthBegin && k < depthEnd) {\n norm += getInputImage(b, r, c, k) * getInputImage(b, r, c, k);\n } else {\n break;\n }\n }\n\n norm = uniforms.alpha * norm + uniforms.bias;\n\n for (var k = MIN_DEPTH_BEGIN; k < MAX_DEPTH_END; k++) {\n if (k < depthBegin) {\n continue;\n } else if (k >= depthBegin && k < depthEnd) {\n var dyi = -2.0 * uniforms.alpha * uniforms.beta\n * getInputImage(b, r, c, k) * getOutputImage(b, r, c, d) / norm;\n if (k == d) {\n dyi += pow(norm, -1.0 * uniforms.beta);\n }\n if (k == coords[3]) {\n dyi *= getDy(b, r, c, d);\n result += dyi;\n }\n } else {\n break;\n }\n }\n }\n\n setOutputAtIndex(index, result);\n }\n }\n "); return userCode; }; return LRNGradProgram; }()); /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function lrnGrad(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, y = inputs.y, dy = inputs.dy; var depthRadius = attrs.depthRadius, bias = attrs.bias, alpha = attrs.alpha, beta = attrs.beta; var program = new LRNGradProgram(x.shape); var uniformData = [ { type: 'int32', data: [depthRadius] }, { type: 'float32', data: [bias] }, { type: 'float32', data: [alpha] }, { type: 'float32', data: [beta] } ]; var res = backend.runWebGPUProgram(program, [x, y, dy], x.dtype, uniformData); return res; } var lrnGradConfig = { kernelName: tf.LRNGrad, backendName: 'webgpu', kernelFunc: lrnGrad }; /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var maximum = binaryKernelFunc({ opType: BinaryOpType.MAX, cpuKernelImpl: maximumImplCPU, }); var maximumConfig = { kernelName: tf.Maximum, backendName: 'webgpu', kernelFunc: maximum }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function maxPool(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var filterSize = attrs.filterSize, strides = attrs.strides, pad = attrs.pad, dimRoundingMode = attrs.dimRoundingMode; var dilations = 1; var convInfo = tf.backend_util.computePool2DInfo(x.shape, filterSize, strides, dilations, pad, dimRoundingMode); return poolImpl(x, convInfo, 'max', backend); } var maxPoolConfig = { kernelName: tf.MaxPool, backendName: 'webgpu', kernelFunc: maxPool }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function maxPool3d(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var filterSize = attrs.filterSize, strides = attrs.strides, pad = attrs.pad, dataFormat = attrs.dataFormat, dimRoundingMode = attrs.dimRoundingMode; var dilations = [1, 1, 1]; var convInfo = tf.backend_util.computePool3DInfo(x.shape, filterSize, strides, dilations, pad, dimRoundingMode, dataFormat); var maxPoolProgram = new Pool3DProgram(convInfo, 'max'); var dimensions = [ { type: 'int32', data: [convInfo.strideDepth, convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.padInfo.front, convInfo.padInfo.top, convInfo.padInfo.left] }, { type: 'int32', data: [convInfo.inDepth, convInfo.inHeight, convInfo.inWidth] }, { type: 'int32', data: [ convInfo.effectiveFilterDepth, convInfo.effectiveFilterHeight, convInfo.effectiveFilterWidth ] } ]; return backend.runWebGPUProgram(maxPoolProgram, [x], x.dtype, dimensions); } var maxPool3DConfig = { kernelName: tf.MaxPool3D, backendName: 'webgpu', kernelFunc: maxPool3d }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var MaxPool2DBackpropProgram = /** @class */ (function () { function MaxPool2DBackpropProgram(convInfo) { this.variableNames = ['dy', 'maxPos']; this.uniforms = "strides : vec2, pads : vec2, dilations : vec2, filterDims : vec2,\n outHeight : i32, outWidth : i32"; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = convInfo.inShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'maxPool2DBackprop'; } MaxPool2DBackpropProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let batch = coords[0];\n let d = coords[3];\n\n let dyRCCorner = vec2(coords.yz) - uniforms.pads;\n let dyRCorner = dyRCCorner.x;\n let dyCCorner = dyRCCorner.y;\n\n // Convolve dy(?, ?, d) with pos mask(:, :, d) to get dx(xR, xC, d).\n // ? = to be determined. : = across all values in that axis.\n var dotProd = 0.0;\n let lastIndex = uniforms.filterDims[0] * uniforms.filterDims[1] - 1;\n for (var wR = 0; wR < uniforms.filterDims[0]; wR += uniforms.dilations[0]) {\n let dyR = f32(dyRCorner + wR) / f32(uniforms.strides[0]);\n\n if (dyR < 0.0 || dyR >= f32(uniforms.outHeight) || fract(dyR) > 0.0) {\n continue;\n }\n let idyR = i32(dyR);\n\n for (var wC = 0; wC < uniforms.filterDims[1]; wC += uniforms.dilations[1]) {\n let dyC = f32(dyCCorner + wC) / f32(uniforms.strides[1]);\n\n if (dyC < 0.0 || dyC >= f32(uniforms.outWidth) || fract(dyC) > 0.0) {\n continue;\n }\n let idyC = i32(dyC);\n\n let dyValue = getDy(batch, idyR, idyC, d);\n let maxPosValue = lastIndex - i32(getMaxPos(batch, idyR, idyC, d));\n\n // Get the current value, check it against the value from the\n // position matrix.\n let curPosValue = wR * uniforms.filterDims[1] + wC;\n let mask = select(0.0, 1.0, maxPosValue == curPosValue);\n dotProd += dyValue * mask;\n }\n }\n setOutputAtIndex(index, dotProd);\n }\n }\n "); return userCode; }; return MaxPool2DBackpropProgram; }()); var MaxPool3DBackpropProgram = /** @class */ (function () { function MaxPool3DBackpropProgram(convInfo) { this.variableNames = ['dy', 'maxPos']; this.uniforms = "strides : vec3, pads : vec3, filterDims : vec3,\n outDepth : i32, outHeight : i32, outWidth : i32"; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = convInfo.inShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'maxPool3DBackprop'; } MaxPool3DBackpropProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let batch = coords.x;\n let ch = coords.u;\n\n let dyCorner = vec3(coords.y, coords.z, coords.w) - uniforms.pads;\n let dyDCorner = dyCorner.x;\n let dyRCorner = dyCorner.y;\n let dyCCorner = dyCorner.z;\n\n // Convolve dy(?, ?, ?, ch) with pos mask(:, :, :, d) to get\n // dx(xD, xR, xC, ch).\n // ? = to be determined. : = across all values in that axis.\n var dotProd = 0.0;\n let lastIndex = uniforms.filterDims[0] * uniforms.filterDims[1] * uniforms.filterDims[2] - 1;\n\n for (var wD = 0; wD < uniforms.filterDims[0]; wD++) {\n let dyD = f32(dyDCorner + wD) / f32(uniforms.strides[0]);\n\n if (dyD < 0.0 || dyD >= f32(uniforms.outDepth) || fract(dyD) > 0.0) {\n continue;\n }\n let idyD = i32(dyD);\n\n for (var wR = 0; wR < uniforms.filterDims[1]; wR++) {\n let dyR = f32(dyRCorner + wR) / f32(uniforms.strides[1]);\n\n if (dyR < 0.0 || dyR >= f32(uniforms.outHeight) || fract(dyR) > 0.0) {\n continue;\n }\n let idyR = i32(dyR);\n\n for (var wC = 0; wC < uniforms.filterDims[2]; wC++) {\n let dyC = f32(dyCCorner + wC) / f32(uniforms.strides[2]);\n\n if (dyC < 0.0 || dyC >= f32(uniforms.outWidth) || fract(dyC) > 0.0) {\n continue;\n }\n let idyC = i32(dyC);\n\n let dyValue = getDy(batch, idyD, idyR, idyC, ch);\n let maxPosValue = lastIndex - i32(getMaxPos(batch, idyD, idyR, idyC, ch));\n\n // Get the current value, check it against the value from the\n // position matrix.\n let curPosValue = wD * uniforms.filterDims[1] * uniforms.filterDims[2] + wR * uniforms.filterDims[2] + wC;\n let mask = select(0.0, 1.0, maxPosValue == curPosValue);\n dotProd += dyValue * mask;\n }\n }\n }\n\n setOutputAtIndex(index, dotProd);\n }\n }\n "); return userCode; }; return MaxPool3DBackpropProgram; }()); /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function maxPool3DGrad(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var dy = inputs.dy, input = inputs.input; var x = input; var filterSize = attrs.filterSize, strides = attrs.strides, pad = attrs.pad, dimRoundingMode = attrs.dimRoundingMode; var dilations = [1, 1, 1]; var convInfo = tf.backend_util.computePool3DInfo(x.shape, filterSize, strides, dilations, pad, dimRoundingMode); var maxPool3dPositionsProgram = new Pool3DProgram(convInfo, 'max', true /* get positions */); var uniformData = [ { type: 'int32', data: [convInfo.strideDepth, convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.padInfo.front, convInfo.padInfo.top, convInfo.padInfo.left] }, { type: 'int32', data: [convInfo.inDepth, convInfo.inHeight, convInfo.inWidth] }, { type: 'int32', data: [ convInfo.effectiveFilterDepth, convInfo.effectiveFilterHeight, convInfo.effectiveFilterWidth ] } ]; var maxPool3dPositions = backend.runWebGPUProgram(maxPool3dPositionsProgram, [x], 'int32', uniformData); var maxPool3dBackpropProgram = new MaxPool3DBackpropProgram(convInfo); uniformData = [ { type: 'int32', data: [convInfo.strideDepth, convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [ convInfo.effectiveFilterDepth - 1 - convInfo.padInfo.front, convInfo.effectiveFilterHeight - 1 - convInfo.padInfo.top, convInfo.effectiveFilterWidth - 1 - convInfo.padInfo.left ] }, { type: 'int32', data: [ convInfo.effectiveFilterDepth, convInfo.effectiveFilterHeight, convInfo.effectiveFilterWidth ] }, { type: 'int32', data: [convInfo.outDepth] }, { type: 'int32', data: [convInfo.outHeight] }, { type: 'int32', data: [convInfo.outWidth] } ]; var result = backend.runWebGPUProgram(maxPool3dBackpropProgram, [dy, maxPool3dPositions], x.dtype, uniformData); backend.disposeData(maxPool3dPositions.dataId); return result; } var maxPool3DGradConfig = { kernelName: tf.MaxPool3DGrad, backendName: 'webgpu', kernelFunc: maxPool3DGrad }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function maxPoolGrad(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var dy = inputs.dy, input = inputs.input, output = inputs.output; var x = input; assertNotComplex([input, output], 'maxPoolGrad'); var filterSize = attrs.filterSize, strides = attrs.strides, pad = attrs.pad, dimRoundingMode = attrs.dimRoundingMode; var convInfo = tf.backend_util.computePool2DInfo(x.shape, filterSize, strides, 1 /* dilations */, pad, dimRoundingMode); var maxPoolPositionsProgram = new Pool2DProgram(convInfo, 'max', true); var uniformData = [ { type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.padInfo.top, convInfo.padInfo.left] }, { type: 'int32', data: [convInfo.dilationHeight, convInfo.dilationWidth] }, { type: 'int32', data: [convInfo.inHeight, convInfo.inWidth] }, { type: 'int32', data: [convInfo.effectiveFilterHeight, convInfo.effectiveFilterWidth] } ]; var maxPoolPositions = backend.runWebGPUProgram(maxPoolPositionsProgram, [x], 'int32', uniformData); var maxPoolBackpropProgram = new MaxPool2DBackpropProgram(convInfo); uniformData = [ { type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [ convInfo.effectiveFilterHeight - 1 - convInfo.padInfo.top, convInfo.effectiveFilterWidth - 1 - convInfo.padInfo.left ] }, { type: 'int32', data: [convInfo.dilationHeight, convInfo.dilationWidth] }, { type: 'int32', data: [convInfo.effectiveFilterHeight, convInfo.effectiveFilterWidth] }, { type: 'int32', data: [convInfo.outHeight] }, { type: 'int32', data: [convInfo.outWidth] } ]; var result = backend.runWebGPUProgram(maxPoolBackpropProgram, [dy, maxPoolPositions], x.dtype, uniformData); backend.disposeData(maxPoolPositions.dataId); return result; } var maxPoolGradConfig = { kernelName: tf.MaxPoolGrad, backendName: 'webgpu', kernelFunc: maxPoolGrad }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function maxPoolWithArgmax(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var filterSize = attrs.filterSize, strides = attrs.strides, pad = attrs.pad, includeBatchInIndex = attrs.includeBatchInIndex; var x = inputs.x; tf.util.assert(x.shape.length === 4, function () { return "Error in maxPool: input must be rank 4 but got rank ".concat(x.shape.length, "."); }); var dilations = [1, 1]; tf.util.assert(tf.backend_util.eitherStridesOrDilationsAreOne(strides, dilations), function () { return 'Error in maxPool: Either strides or dilations must be 1. ' + "Got strides ".concat(strides, " and dilations '").concat(dilations, "'"); }); var convInfo = tf.backend_util.computePool2DInfo(x.shape, filterSize, strides, dilations, pad); var uniformData = [ { type: 'int32', data: [convInfo.strideHeight, convInfo.strideWidth] }, { type: 'int32', data: [convInfo.padInfo.top, convInfo.padInfo.left] }, { type: 'int32', data: [convInfo.dilationHeight, convInfo.dilationWidth] }, { type: 'int32', data: [convInfo.inHeight, convInfo.inWidth] }, { type: 'int32', data: [convInfo.effectiveFilterHeight, convInfo.effectiveFilterWidth] } ]; var program = new Pool2DProgram(convInfo, 'max', false); var poolOutput = backend.runWebGPUProgram(program, [x], x.dtype, uniformData); program = new Pool2DProgram(convInfo, 'max', true, true, includeBatchInIndex); var indexOutput = backend.runWebGPUProgram(program, [x], 'int32', uniformData); return [poolOutput, indexOutput]; } var maxPoolWithArgmaxConfig = { kernelName: tf.MaxPoolWithArgmax, backendName: 'webgpu', kernelFunc: maxPoolWithArgmax }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function min(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var axis = attrs.axis, keepDims = attrs.keepDims; return reduce(x, axis, keepDims, 'min', backend); } var minConfig = { kernelName: tf.Min, backendName: 'webgpu', kernelFunc: min }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var minimum = binaryKernelFunc({ opType: BinaryOpType.MIN, cpuKernelImpl: minimumImplCPU, }); var minimumConfig = { kernelName: tf.Minimum, backendName: 'webgpu', kernelFunc: minimum }; /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var MirrorPadProgram = /** @class */ (function () { function MirrorPadProgram(xShape, paddings, mode) { var _this = this; this.uniforms = ''; this.variableNames = ['x']; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = paddings.map(function (p, i) { return p[0] /* beforePad */ + xShape[i] + p[1]; } /* afterPad */); this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.xShape = xShape; paddings.map(function (_, i) { _this.uniforms += " pad".concat(i, " : vec2,"); }); this.offset = mode === 'reflect' ? 0 : 1; this.shaderKey = "mirrorPad_".concat(mode); } MirrorPadProgram.prototype.getUserCode = function () { var rank = this.xShape.length; // The length of paddings are same with the rank of the input tensor. var start = this.xShape.map(function (_, i) { return "uniforms.pad".concat(i, "[0]"); }).join(','); var end = this.xShape .map(function (_, i) { return "uniforms.pad".concat(i, "[0] + uniforms.xShape").concat(rank > 1 ? "[".concat(i, "]") : ''); }) .join(','); var shaderStart = rank === 1 ? 'start' : 'start[i]'; var shaderEnd = rank === 1 ? 'end' : 'end[i]'; var shaderOutC = rank === 1 ? 'outC' : 'outC[i]'; var dtype = getCoordsDataType(rank); var unpackedCoords = rank > 1 ? ['coords[0]', 'coords[1]', 'coords[2]', 'coords[3]'].slice(0, rank) : 'coords'; return "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let start = ").concat(dtype, "(").concat(start, ");\n let end = ").concat(dtype, "(").concat(end, ");\n var outC = getCoordsFromIndex(index);\n for (var i = 0; i < ").concat(rank, "; i = i + 1) {\n if (").concat(shaderOutC, " < ").concat(shaderStart, ") {\n ").concat(shaderOutC, " = ").concat(shaderStart, " * 2 - ").concat(shaderOutC, " - ").concat(this.offset, ";\n } else if(").concat(shaderOutC, " >= ").concat(shaderEnd, ") {\n ").concat(shaderOutC, " = (").concat(shaderEnd, " - 1) * 2 - ").concat(shaderOutC, " + ").concat(this.offset, ";\n }\n }\n let coords = outC - start;\n setOutputAtIndex(index, getX(").concat(unpackedCoords, "));\n }\n }\n "); }; return MirrorPadProgram; }()); /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var mirrorPadConfig = { kernelName: tf.MirrorPad, backendName: 'webgpu', kernelFunc: function (_a) { var inputs = _a.inputs, attrs = _a.attrs, backend = _a.backend; var x = inputs.x; var paddings = attrs.paddings, mode = attrs.mode; var webGPUBackend = backend; var uniformData = paddings.map(function (p) { return { type: 'int32', data: [p[0], p[1]] }; }); var program = new MirrorPadProgram(x.shape, paddings, mode); var output = webGPUBackend.runWebGPUProgram(program, [x], x.dtype, uniformData); return output; } }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var mod = binaryKernelFunc({ opType: BinaryOpType.MOD }); var modConfig = { kernelName: tf.Mod, backendName: 'webgpu', kernelFunc: mod }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var MultinomialProgram = /** @class */ (function () { function MultinomialProgram(batchSize, numSamples) { this.variableNames = ['probs']; this.outputShape = []; this.uniforms = 'seed : f32, numOutcomes: i32,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = [batchSize, numSamples]; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'multinomial'; } MultinomialProgram.prototype.getUserCode = function () { var userCode = "\n //Based on the work of Dave Hoskins\n //https://www.shadertoy.com/view/4djSRW\n fn random (seed : f32, resultUV : vec2) -> f32 {\n let HASHSCALE1 = 443.8975;\n let p = resultUV * seed;\n var p3 = fract(vec3(p.xyx) * HASHSCALE1);\n p3 = p3 + dot(p3, p3.yzx + 19.19);\n return fract((p3.x + p3.y) * p3.z);\n }\n\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getOutputCoords();\n let batch = coords[0];\n\n let resUV = vec2(f32(coords[1]) / f32(uniforms.outShape[1]),\n f32(coords[0]) / f32(uniforms.outShape[0]));\n let r = random(uniforms.seed, resUV);\n var cdf = 0.0;\n for (var i = 0; i < uniforms.numOutcomes - 1; i = i + 1) {\n cdf = cdf + getProbs(batch, i);\n\n if (r < cdf) {\n setOutputAtIndexI32(index, i);\n return;\n }\n }\n\n // If no other event happened, last event happened.\n setOutputAtIndexI32(index, uniforms.numOutcomes - 1);\n }\n }\n "); return userCode; }; return MultinomialProgram; }()); /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var SoftmaxProgram = /** @class */ (function () { function SoftmaxProgram(outputShape) { this.variableNames = ['logits']; this.outputShape = outputShape; // [rows, cols] this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = [this.outputShape[0], 1, 1]; if (this.outputShape[1] >= 4096) { this.workgroupSize = [256, 1, 1]; } else { this.workgroupSize = [64, 1, 1]; } this.shaderKey = 'softmax'; } SoftmaxProgram.prototype.getUserCode = function () { var userCode = "\n var buf : array;\n var rowMaxShared : f32;\n var rowSumShared : f32;\n const blockSize = ").concat(this.workgroupSize[0], ";\n ").concat(getMainHeaderString('index'), " {\n let row = index / blockSize;\n let tid = i32(localId.x);\n let cols = uniforms.outShape[1];\n\n var threadMax = -3.402823e+38f;\n for (var col = tid; col < cols; col += blockSize) {\n let value = getLogits(row, col);\n threadMax = max(threadMax, value);\n }\n if (tid < cols) {\n buf[tid] = threadMax;\n }\n workgroupBarrier();\n\n var reduceSize = min(cols, blockSize);\n for (var currSize = reduceSize >> 1; currSize > 0; currSize = reduceSize >> 1) {\n reduceSize = currSize + (reduceSize & 1);\n if (tid < currSize) {\n buf[tid] = max(buf[tid], buf[tid + reduceSize]);\n }\n workgroupBarrier();\n }\n\n if (tid == 0) {\n rowMaxShared = buf[0];\n }\n workgroupBarrier();\n\n var threadSum = 0.0;\n for (var col = tid; col < cols; col += blockSize) {\n let subExp = exp(getLogits(row, col) - rowMaxShared);\n threadSum += subExp;\n }\n buf[tid] = threadSum;\n workgroupBarrier();\n\n for (var currSize = blockSize >> 1; currSize > 0; currSize = currSize >> 1) {\n if (tid < currSize) {\n buf[tid] = buf[tid] + buf[tid + currSize];\n }\n workgroupBarrier();\n }\n\n if (tid == 0) {\n rowSumShared = buf[0];\n }\n workgroupBarrier();\n\n for (var col = tid; col < cols; col += blockSize) {\n let value = exp(getLogits(row, col) - rowMaxShared) / rowSumShared;\n setOutputAtCoords(row, col, value);\n }\n }\n "); return userCode; }; return SoftmaxProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function softmax(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var logits = inputs.logits; var dim = attrs.dim; var logitsReshaped = reshape({ inputs: { x: logits }, backend: backend, attrs: { shape: [ tf.util.sizeFromShape(logits.shape) / logits.shape[dim], logits.shape[dim] ] } }); var program = new SoftmaxProgram(logitsReshaped.shape); var res = backend.runWebGPUProgram(program, [logitsReshaped], logits.dtype); var resReshaped = reshape({ inputs: { x: res }, backend: backend, attrs: { shape: logits.shape } }); backend.disposeData(logitsReshaped.dataId); backend.disposeData(res.dataId); return resReshaped; } var softmaxConfig = { kernelName: tf.Softmax, backendName: 'webgpu', kernelFunc: softmax }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function multinomial(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var logits = inputs.logits; var numSamples = attrs.numSamples, seed = attrs.seed, normalized = attrs.normalized; var probs = normalized ? logits : softmax({ inputs: { logits: logits }, backend: backend, attrs: { dim: logits.shape.length - 1 } }); var batchSize = probs.shape[0]; var numOutcomes = probs.shape[1]; var program = new MultinomialProgram(batchSize, numSamples); var uniformData = [{ type: 'float32', data: [seed] }, { type: 'int32', data: [numOutcomes] }]; var res = backend.runWebGPUProgram(program, [probs], 'int32', uniformData); if (!normalized) { backend.disposeData(probs.dataId); } return res; } var multinomialConfig = { kernelName: tf.Multinomial, backendName: 'webgpu', kernelFunc: multinomial }; // This doesn't use unaryKernelFunc because negImplCPU is not of type // SimpleUnaryKernelImplCPU. function neg(args) { var inputs = args.inputs, backend = args.backend; var x = inputs.x; if (backend.shouldExecuteOnCPU([x])) { var xData = backend.tensorMap.get(x.dataId); var _a = __read(negImplCPU(xData.values, x.shape, x.dtype), 2), outValues = _a[0], newShape = _a[1]; return backend.makeTensorInfo(newShape, x.dtype, outValues); } var program = new UnaryOpProgram(x.shape, UnaryOpType.NEG); return backend.runWebGPUProgram(program, [x], x.dtype); } var negConfig = { kernelName: tf.Neg, backendName: 'webgpu', kernelFunc: neg }; /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function nonMaxSuppressionV3(args) { console.warn('tf.nonMaxSuppression() in webgpu locks the UI thread. ' + 'Call tf.nonMaxSuppressionAsync() instead'); var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var boxes = inputs.boxes, scores = inputs.scores; var maxOutputSize = attrs.maxOutputSize, iouThreshold = attrs.iouThreshold, scoreThreshold = attrs.scoreThreshold; var boxesVals = backend.readSync(boxes.dataId); var scoresVals = backend.readSync(scores.dataId); var selectedIndices = tf.kernel_impls.nonMaxSuppressionV3Impl(boxesVals, scoresVals, maxOutputSize, iouThreshold, scoreThreshold).selectedIndices; return backend.makeTensorInfo([selectedIndices.length], 'int32', new Int32Array(selectedIndices)); } var nonMaxSuppressionV3Config = { kernelName: tf.NonMaxSuppressionV3, backendName: 'webgpu', kernelFunc: nonMaxSuppressionV3 }; /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function nonMaxSuppressionV5(args) { console.warn('tf.nonMaxSuppression() in webgpu locks the UI thread. ' + 'Call tf.nonMaxSuppressionAsync() instead'); var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var boxes = inputs.boxes, scores = inputs.scores; var maxOutputSize = attrs.maxOutputSize, iouThreshold = attrs.iouThreshold, scoreThreshold = attrs.scoreThreshold, softNmsSigma = attrs.softNmsSigma; var boxesVals = backend.readSync(boxes.dataId); var scoresVals = backend.readSync(scores.dataId); var maxOutputSizeVal = maxOutputSize; var iouThresholdVal = iouThreshold; var scoreThresholdVal = scoreThreshold; var softNmsSigmaVal = softNmsSigma; var _a = tf.kernel_impls.nonMaxSuppressionV5Impl(boxesVals, scoresVals, maxOutputSizeVal, iouThresholdVal, scoreThresholdVal, softNmsSigmaVal), selectedIndices = _a.selectedIndices, selectedScores = _a.selectedScores; return [ backend.makeTensorInfo([selectedIndices.length], 'int32', new Int32Array(selectedIndices)), backend.makeTensorInfo([selectedScores.length], 'float32', new Float32Array(selectedScores)) ]; } var nonMaxSuppressionV5Config = { kernelName: tf.NonMaxSuppressionV5, backendName: 'webgpu', kernelFunc: nonMaxSuppressionV5 }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var OneHotProgram = /** @class */ (function () { function OneHotProgram(numIndices, depth) { this.variableNames = ['x']; this.uniforms = 'onValue : f32, offValue : f32,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = [numIndices, depth]; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'onehot'; } OneHotProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if(index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n setOutputAtIndex(index, mix(uniforms.offValue, uniforms.onValue,\n f32(i32(round(getX(coords.x))) == coords.y)));\n }\n }\n "); return userCode; }; return OneHotProgram; }()); function oneHot(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var indices = inputs.indices; var dtype = attrs.dtype, depth = attrs.depth, onValue = attrs.onValue, offValue = attrs.offValue; var indicesSize = tf.util.sizeFromShape(indices.shape); var program = new OneHotProgram(indicesSize, depth); var reshaped = reshape({ inputs: { x: indices }, backend: backend, attrs: { shape: [indicesSize] } }); var uniformData = [{ type: 'float32', data: [onValue] }, { type: 'float32', data: [offValue] }]; var result = backend.runWebGPUProgram(program, [reshaped], dtype, uniformData); backend.disposeData(reshaped.dataId); var outShape = __spreadArray(__spreadArray([], __read(indices.shape), false), [depth], false); var out = reshape({ inputs: { x: result }, backend: backend, attrs: { shape: outShape } }); backend.disposeData(result.dataId); return out; } var oneHotConfig = { kernelName: tf.OneHot, backendName: 'webgpu', kernelFunc: oneHot }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function zerosLike(args) { var inputs = args.inputs, backend = args.backend; var x = inputs.x; if (x.dtype === 'complex64') { var realPart = real({ inputs: { input: x }, backend: backend }); var r = zerosLike({ inputs: { x: realPart }, backend: backend }); var imagPart = imag({ inputs: { input: x }, backend: backend }); var i = zerosLike({ inputs: { x: imagPart }, backend: backend }); var result = complex({ inputs: { real: r, imag: i }, backend: backend }); backend.disposeData(realPart.dataId); backend.disposeData(r.dataId); backend.disposeData(imagPart.dataId); backend.disposeData(i.dataId); return result; } else { return fill({ attrs: { shape: x.shape, dtype: x.dtype, value: x.dtype === 'string' ? '' : 0 }, backend: backend }); } } var zerosLikeConfig = { kernelName: tf.ZerosLike, backendName: 'webgpu', kernelFunc: zerosLike }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function onesLike(args) { var inputs = args.inputs, backend = args.backend; var x = inputs.x; if (x.dtype === 'string') { throw new Error('onesLike is not supported under string dtype'); } else if (x.dtype === 'complex64') { var realPart = real({ inputs: { input: x }, backend: backend }); var r = onesLike({ inputs: { x: realPart }, backend: backend }); var imagPart = imag({ inputs: { input: x }, backend: backend }); var i = zerosLike({ inputs: { x: imagPart }, backend: backend }); var result = complex({ inputs: { real: r, imag: i }, backend: backend }); backend.disposeData(realPart.dataId); backend.disposeData(r.dataId); backend.disposeData(imagPart.dataId); backend.disposeData(i.dataId); return result; } else { return fill({ attrs: { shape: x.shape, dtype: x.dtype, value: 1 }, backend: backend }); } } var onesLikeConfig = { kernelName: tf.OnesLike, backendName: 'webgpu', kernelFunc: onesLike }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function pack(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var axis = attrs.axis; if (inputs.length === 1) { return expandDims({ inputs: { input: inputs[0] }, backend: backend, attrs: { dim: axis } }); } var shape = inputs[0].shape; var dtype = inputs[0].dtype; inputs.forEach(function (t) { tf.util.assertShapesMatch(shape, t.shape, 'All tensors passed to stack must have matching shapes'); tf.util.assert(dtype === t.dtype, function () { return 'All tensors passed to stack must have matching dtypes'; }); }); var intermediateTensorInfos = []; var expandedTensors = inputs.map(function (t) { var expandedT = expandDims({ inputs: { input: t }, backend: backend, attrs: { dim: axis } }); intermediateTensorInfos.push(expandedT); return expandedT; }); var result = concat({ inputs: expandedTensors, backend: backend, attrs: { axis: axis } }); intermediateTensorInfos.forEach(function (t) { return backend.disposeData(t.dataId); }); return result; } var packConfig = { kernelName: tf.Pack, backendName: 'webgpu', kernelFunc: pack }; /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function padCommon(shape, fillZero) { if (fillZero === void 0) { fillZero = false; } var rank = shape.length; var type = getCoordsDataType(rank); var start = shape.map(function (_, i) { return "uniforms.pad".concat(i, "[0]"); }).join(','); var end = shape .map(function (_, i) { return "uniforms.pad".concat(i, "[0] + uniforms.xShape").concat(rank > 1 ? "[".concat(i, "]") : ''); }) .join(','); var startValue = rank > 1 ? "".concat(type, "(").concat(start, ")") : "".concat(start); var endValue = rank > 1 ? "".concat(type, "(").concat(end, ")") : "".concat(end); var leftPadCondition = rank > 1 ? "any(paddedCoords < start)" : "paddedCoords < start"; var rightPadCondition = rank > 1 ? "any(paddedCoords >= end)" : "paddedCoords >= end"; var unpackedCoords = rank > 1 ? ['coords[0]', 'coords[1]', 'coords[2]', 'coords[3]'].slice(0, rank) : 'coords'; return "\n let start = ".concat(startValue, ";\n let end = ").concat(endValue, ";\n if (").concat(leftPadCondition, " || ").concat(rightPadCondition, ") {\n setOutputAtIndex(index, ").concat(fillZero ? 0.0 : 'uniforms.constantValue', ");\n } else {\n let coords = paddedCoords - start;\n setOutputAtIndex(index, getX(").concat(unpackedCoords, "));\n }\n "); } var PadProgram = /** @class */ (function () { function PadProgram(xShape, paddings) { var _this = this; this.variableNames = ['x']; this.uniforms = 'constantValue : f32,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = paddings.map(function (p, i) { return p[0] /* beforePad */ + xShape[i] + p[1]; } /* afterPad */); this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); paddings.map(function (_, i) { _this.uniforms += " pad".concat(i, " : vec2,"); }); this.xShape = xShape; this.shaderKey = 'pad'; } PadProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let paddedCoords = getCoordsFromIndex(index);\n ").concat(padCommon(this.xShape), "\n }\n }\n "); return userCode; }; return PadProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var padV2 = function (args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var paddings = attrs.paddings, constantValue = attrs.constantValue; if (paddings.every(function (p) { return tf.util.arraysEqual(p, [0, 0]); })) { return identity({ inputs: { x: x }, backend: backend }); } if (tf.util.sizeFromShape(x.shape) === 0) { // Short-circuit the computation, since x doesn't have value, only // the shape is used to compute output shape to pad. var outputShape = paddings.map(function (p, i) { return p[0] /* beforePad */ + x.shape[i] + p[1]; } /* afterPad */); return fill({ backend: backend, attrs: { shape: outputShape, value: constantValue, dtype: x.dtype } }); } var uniformData = [{ type: 'float32', data: [constantValue] }]; paddings.map(function (p) { return uniformData.push({ type: 'int32', data: [p[0], p[1]] }); }); var program = new PadProgram(x.shape, paddings); return backend.runWebGPUProgram(program, [x], x.dtype, uniformData); }; var padV2Config = { kernelName: tf.PadV2, backendName: 'webgpu', kernelFunc: padV2 }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var pow = binaryKernelFunc({ opType: BinaryOpType.POW, }); var powConfig = { kernelName: tf.Pow, backendName: 'webgpu', kernelFunc: pow }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function prelu(args) { var inputs = args.inputs, backend = args.backend; var x = inputs.x, alpha = inputs.alpha; var program = new BinaryOpProgram(BinaryOpType.PRELU, x.shape, alpha.shape); return backend.runWebGPUProgram(program, [x, alpha], 'float32'); } var preluConfig = { kernelName: tf.Prelu, backendName: 'webgpu', kernelFunc: prelu }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function prod(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var axis = attrs.axis, keepDims = attrs.keepDims; return reduce(x, axis, keepDims, 'prod', backend); } var prodConfig = { kernelName: tf.Prod, backendName: 'webgpu', kernelFunc: prod }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var range = function (args) { var backend = args.backend, attrs = args.attrs; var start = attrs.start, stop = attrs.stop, step = attrs.step, dtype = attrs.dtype; var values = rangeImplCPU(start, stop, step, dtype); return backend.makeTensorInfo([values.length], dtype, values); }; var rangeConfig = { kernelName: tf.Range, backendName: 'webgpu', kernelFunc: range }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var realDiv = binaryKernelFunc({ opType: BinaryOpType.DIV }); var realDivConfig = { kernelName: tf.RealDiv, backendName: 'webgpu', kernelFunc: realDiv }; /** * @license * Copyright 2022 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var reciprocal = unaryKernelFunc({ opType: UnaryOpType.RECIPROCAL }); var reciprocalConfig = { kernelName: tf.Reciprocal, backendName: 'webgpu', kernelFunc: reciprocal }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var relu = unaryKernelFunc({ opType: UnaryOpType.RELU }); var reluConfig = { kernelName: tf.Relu, backendName: 'webgpu', kernelFunc: relu }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var relu6 = unaryKernelFunc({ opType: UnaryOpType.RELU6 }); var relu6Config = { kernelName: tf.Relu6, backendName: 'webgpu', kernelFunc: relu6 }; /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var ResizeBilinearProgram = /** @class */ (function () { function ResizeBilinearProgram(inputShape, newHeight, newWidth) { this.variableNames = ['x']; this.uniforms = 'adjustHeightWidth : vec2, halfPixelCenters : f32,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = [inputShape[0], newHeight, newWidth, inputShape[3]]; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = "resizeBilinear"; } ResizeBilinearProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let b = coords[0];\n let d = coords[3];\n let rc = coords.yz;\n\n let effectiveInSize = vec2(\n f32(uniforms.xShape.y) - uniforms.adjustHeightWidth[0],\n f32(uniforms.xShape.z) - uniforms.adjustHeightWidth[1]);\n\n let effectiveOutSize = vec2(\n f32(uniforms.outShape.y) - uniforms.adjustHeightWidth[0],\n f32(uniforms.outShape.z) - uniforms.adjustHeightWidth[1]);\n\n let effectiveInputOverOutputRatioRC =\n effectiveInSize / effectiveOutSize;\n\n // Fractional source index\n let sourceFracIndexRC =\n (vec2(rc) + vec2(uniforms.halfPixelCenters)) *\n effectiveInputOverOutputRatioRC - vec2(uniforms.halfPixelCenters);\n\n // Compute the four integer indices.\n let sourceFloorRC = vec2(sourceFracIndexRC);\n let sourceCeilRC = vec2(\n min(vec2(uniforms.xShape.yz) - vec2(1.0), ceil(sourceFracIndexRC)));\n\n let topLeft = getX(b, sourceFloorRC.x, sourceFloorRC.y, d);\n let bottomLeft = getX(b, sourceCeilRC.x, sourceFloorRC.y, d);\n let topRight = getX(b, sourceFloorRC.x, sourceCeilRC.y, d);\n let bottomRight = getX(b, sourceCeilRC.x, sourceCeilRC.y, d);\n\n let fracRC = sourceFracIndexRC - vec2(sourceFloorRC);\n\n let top = topLeft + (topRight - topLeft) * fracRC.y;\n let bottom = bottomLeft + (bottomRight - bottomLeft) * fracRC.y;\n let newValue = top + (bottom - top) * fracRC.x;\n\n setOutputAtIndex(index, newValue);\n }\n }\n "); return userCode; }; return ResizeBilinearProgram; }()); function resizeBilinear(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var images = inputs.images; var alignCorners = attrs.alignCorners, size = attrs.size, halfPixelCenters = attrs.halfPixelCenters; var _a = __read(size, 2), newHeight = _a[0], newWidth = _a[1]; var adjustHeight = alignCorners && newHeight > 1 ? 1.0 : 0.0; var adjustWidth = alignCorners && newWidth > 1 ? 1.0 : 0.0; var halfPixelCentersValue = halfPixelCenters ? 0.5 : 0.0; var uniformData = [ { type: 'float32', data: [adjustHeight, adjustWidth] }, { type: 'float32', data: [halfPixelCentersValue] } ]; var program = new ResizeBilinearProgram(images.shape, newHeight, newWidth); return backend.runWebGPUProgram(program, [images], 'float32', uniformData); } var resizeBilinearConfig = { kernelName: tf.ResizeBilinear, backendName: 'webgpu', kernelFunc: resizeBilinear }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var ResizeBilinearBackpropProgram = /** @class */ (function () { function ResizeBilinearBackpropProgram(inputShape, alignCorners) { this.variableNames = ['dy']; this.uniforms = "effectiveXSize : vec2, effectiveYSize : vec2, heightScale : f32, widthScale : f32,\n invHeightScale : f32, invWidthScale : f32, winHeight : i32, winWidth : i32,"; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = inputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.alignCorners = alignCorners; this.shaderKey = "resizeBilinearBackprop_".concat(alignCorners); } ResizeBilinearBackpropProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getOutputCoords();\n let b = coords[0];\n let d = coords[3];\n let r = coords[1];\n let c = coords[2];\n\n var accumulator = 0.0;\n\n // Compute bounds for where in dy we will look\n let startRLerp = floor(f32(r) * uniforms.invHeightScale);\n let startDyR = i32(startRLerp - f32(uniforms.winHeight / 2));\n\n let startCLerp = floor(f32(c) * uniforms.invWidthScale);\n let startDyC = i32(startCLerp - f32(uniforms.winWidth / 2));\n\n // Loop over dy\n for (var dyROffset = 0; dyROffset < uniforms.winHeight; dyROffset++) {\n let dyR = startDyR + dyROffset;\n\n // Guard against the window exceeding the bounds of dy\n if (dyR < 0 || dyR >= uniforms.dyShape[1]) {\n continue;\n }\n\n for (var dyCOffset = 0; dyCOffset < uniforms.winWidth; dyCOffset++) {\n let dyC = startDyC + dyCOffset;\n\n // Guard against the window exceeding the bounds of dy\n if (dyC < 0 || dyC >= uniforms.dyShape[2]) {\n continue;\n }\n\n let dxR = f32(dyR) * uniforms.heightScale;\n let topDxRIndex = i32(floor(dxR));\n let bottomDxRIndex = i32(min(ceil(dxR), f32(uniforms.outShape[1] - 1)));\n let dxRLerp = dxR - f32(topDxRIndex);\n let inverseDxRLerp = 1.0 - dxRLerp;\n\n let dxC = f32(dyC) * uniforms.widthScale;\n let leftDxCIndex = i32(floor(dxC));\n let rightDxCIndex = i32(min(ceil(dxC), f32(uniforms.outShape[2] - 1)));\n let dxCLerp = dxC - f32(leftDxCIndex);\n let inverseDxCLerp = 1.0 - dxCLerp;\n\n if (r == topDxRIndex && c == leftDxCIndex) {\n // topLeft\n accumulator +=\n getDy(b, dyR, dyC, d) * inverseDxRLerp * inverseDxCLerp;\n }\n\n if (r == topDxRIndex && c == rightDxCIndex) {\n // topRight\n accumulator += getDy(b, dyR, dyC, d) * inverseDxRLerp * dxCLerp;\n }\n\n if (r == bottomDxRIndex && c == leftDxCIndex) {\n // bottomLeft\n accumulator += getDy(b, dyR, dyC, d) * dxRLerp * inverseDxCLerp;\n }\n\n if (r == bottomDxRIndex && c == rightDxCIndex) {\n // bottomRight\n accumulator += getDy(b, dyR, dyC, d) * dxRLerp * dxCLerp;\n }\n }\n }\n // End loop over dy\n\n setOutputAtIndex(index, accumulator);\n }\n }\n "); return userCode; }; return ResizeBilinearBackpropProgram; }()); function resizeBilinearGrad(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var images = inputs.images, dy = inputs.dy; var alignCorners = attrs.alignCorners; var _a = __read(images.shape, 3), xHeight = _a[1], xWidth = _a[2]; var _b = __read(dy.shape, 3), yHeight = _b[1], yWidth = _b[2]; var effectiveXSize = [ (alignCorners && yHeight > 1) ? xHeight - 1 : xHeight, (alignCorners && yWidth > 1) ? xWidth - 1 : xWidth ]; var effectiveYSize = [ (alignCorners && yHeight > 1) ? yHeight - 1 : yHeight, (alignCorners && yWidth > 1) ? yWidth - 1 : yWidth ]; var heightScale = effectiveXSize[0] / effectiveYSize[0]; var widthScale = effectiveXSize[1] / effectiveYSize[1]; var invHeightScale = 1 / heightScale; var invWidthScale = 1 / widthScale; // This defines the size of the window of values around a particular // index in dy that we want to search for contributions to dx. var winHeight = (Math.ceil(invHeightScale) * 2) + 2; var winWidth = (Math.ceil(invWidthScale) * 2) + 2; var program = new ResizeBilinearBackpropProgram(images.shape, alignCorners); var uniformData = [ { type: 'int32', data: effectiveXSize }, { type: 'int32', data: effectiveYSize }, { type: 'float32', data: [heightScale] }, { type: 'float32', data: [widthScale] }, { type: 'float32', data: [invHeightScale] }, { type: 'float32', data: [invWidthScale] }, { type: 'int32', data: [winHeight] }, { type: 'int32', data: [winWidth] } ]; return backend.runWebGPUProgram(program, [dy], dy.dtype, uniformData); } var resizeBilinearGradConfig = { kernelName: tf.ResizeBilinearGrad, backendName: 'webgpu', kernelFunc: resizeBilinearGrad }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var ResizeNearestNeighborProgram = /** @class */ (function () { function ResizeNearestNeighborProgram(inputShape, newHeight, newWidth, halfPixelCenters) { this.variableNames = ['x']; this.uniforms = 'adjustHeightWidth : vec2, roundBase : f32,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = [inputShape[0], newHeight, newWidth, inputShape[3]]; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.halfPixelCenters = halfPixelCenters; this.shaderKey = "resizeNearest_".concat(halfPixelCenters); } ResizeNearestNeighborProgram.prototype.getUserCode = function () { var sourceFracIndexRC; if (this.halfPixelCenters) { sourceFracIndexRC = "max((vec2(rc) + vec2(0.5)) * effectiveInputOverOutputRatioRC" + ", vec2(0.0))"; } else { sourceFracIndexRC = "vec2(rc) * effectiveInputOverOutputRatioRC"; } var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let b = coords[0];\n let d = coords[3];\n let rc = coords.yz;\n\n let effectiveInSize = vec2(\n f32(uniforms.xShape.y) - uniforms.adjustHeightWidth[0],\n f32(uniforms.xShape.z) - uniforms.adjustHeightWidth[1]);\n\n let effectiveOutSize = vec2(\n f32(uniforms.outShape.y) - uniforms.adjustHeightWidth[0],\n f32(uniforms.outShape.z) - uniforms.adjustHeightWidth[1]);\n\n let effectiveInputOverOutputRatioRC =\n effectiveInSize / effectiveOutSize;\n\n // Fractional source index\n let sourceFracIndexRC = ").concat(sourceFracIndexRC, ";\n\n // Compute the coordinators of nearest neighbor point.\n let inputShapeRC = vec2(f32(uniforms.xShape.y), f32(uniforms.xShape.z));\n let sourceNearestRC = vec2(\n min(inputShapeRC - 1.0, floor(sourceFracIndexRC + uniforms.roundBase)));\n let newValue = getX(b, sourceNearestRC.x, sourceNearestRC.y, d);\n\n setOutputAtIndex(index, newValue);\n }\n }\n "); return userCode; }; return ResizeNearestNeighborProgram; }()); function resizeNearestNeighbor(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var images = inputs.images; var alignCorners = attrs.alignCorners, halfPixelCenters = attrs.halfPixelCenters, size = attrs.size; var _a = __read(size, 2), newHeight = _a[0], newWidth = _a[1]; var adjustHeight = alignCorners && newHeight > 1 ? 1.0 : 0.0; var adjustWidth = alignCorners && newWidth > 1 ? 1.0 : 0.0; // When align corners is false, we rounds the value with floor. var roundBase = alignCorners ? 0.5 : 0.0; var uniformData = [ { type: 'float32', data: [adjustHeight, adjustWidth] }, { type: 'float32', data: [roundBase] } ]; var program = new ResizeNearestNeighborProgram(images.shape, newHeight, newWidth, halfPixelCenters); return backend.runWebGPUProgram(program, [images], images.dtype, uniformData); } var resizeNearestNeighborConfig = { kernelName: tf.ResizeNearestNeighbor, backendName: 'webgpu', kernelFunc: resizeNearestNeighbor }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var ResizeNearestNeigborBackpropProgram = /** @class */ (function () { function ResizeNearestNeigborBackpropProgram(inputShape, alignCorners) { this.variableNames = ['dy']; this.uniforms = "effectiveXSize : vec2, effectiveYSize : vec2, invHeightScale : f32, invWidthScale : f32,\n winHeight : i32, winWidth : i32,"; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = inputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.alignCorners = alignCorners; this.shaderKey = "resizeNearestNeigborBackprop_".concat(alignCorners); } ResizeNearestNeigborBackpropProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getOutputCoords();\n let b = coords[0];\n let d = coords[3];\n let r = coords[1];\n let c = coords[2];\n\n var accumulator = 0.0;\n\n // Compute bounds for where in dy we will look\n let startRLerp = floor(f32(r) * uniforms.invHeightScale);\n let startDyR = i32(floor(startRLerp - f32(uniforms.winHeight / 2)));\n\n let startCLerp = floor(f32(c) * uniforms.invWidthScale);\n let startDyC = i32(floor(startCLerp - f32(uniforms.winWidth / 2)));\n\n // Loop over dy\n for (var dyROffset = 0; dyROffset < uniforms.winHeight; dyROffset++) {\n let dyR = startDyR + dyROffset;\n\n // Guard against the window exceeding the bounds of dy\n if (dyR < 0 || dyR >= uniforms.dyShape[1]) {\n continue;\n }\n\n for (var dyCOffset = 0; dyCOffset < uniforms.winWidth; dyCOffset++) {\n let dyC = startDyC + dyCOffset;\n\n // Guard against the window exceeding the bounds of dy\n if (dyC < 0 || dyC >= uniforms.dyShape[2]) {\n continue;\n }\n\n let sourceFracRow = f32(uniforms.effectiveXSize[0]) *\n (f32(dyR) / f32(uniforms.effectiveYSize[0]));\n\n let sourceFracCol = f32(uniforms.effectiveXSize[1]) *\n (f32(dyC) / f32(uniforms.effectiveYSize[1]));\n\n let sourceNearestRow =\n i32(min(f32(uniforms.outShape[1] - 1),\n ").concat(this.alignCorners ? 'floor(sourceFracRow + 0.5)' : 'floor(sourceFracRow)', "));\n\n let sourceNearestCol =\n i32(min(f32(uniforms.outShape[2] - 1),\n ").concat(this.alignCorners ? 'floor(sourceFracCol + 0.5)' : 'floor(sourceFracCol)', "));\n\n if (r == sourceNearestRow && c == sourceNearestCol) {\n accumulator += getDy(b, dyR, dyC, d);\n }\n }\n }\n // End loop over dy\n\n setOutputAtIndex(index, accumulator);\n }\n }\n "); return userCode; }; return ResizeNearestNeigborBackpropProgram; }()); function resizeNearestNeighborGrad(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var images = inputs.images, dy = inputs.dy; var alignCorners = attrs.alignCorners; var _a = __read(images.shape, 3), xHeight = _a[1], xWidth = _a[2]; var _b = __read(dy.shape, 3), yHeight = _b[1], yWidth = _b[2]; var effectiveXSize = [ (alignCorners && yHeight > 1) ? xHeight - 1 : xHeight, (alignCorners && yWidth > 1) ? xWidth - 1 : xWidth ]; var effectiveYSize = [ (alignCorners && yHeight > 1) ? yHeight - 1 : yHeight, (alignCorners && yWidth > 1) ? yWidth - 1 : yWidth ]; var heightScale = effectiveXSize[0] / effectiveYSize[0]; var widthScale = effectiveXSize[1] / effectiveYSize[1]; var invHeightScale = 1 / heightScale; var invWidthScale = 1 / widthScale; // This defines the size of the window of values around a particular // index in dy that we want to search for contributions to dx. var winHeight = (Math.ceil(invHeightScale) * 2) + 2; var winWidth = (Math.ceil(invWidthScale) * 2) + 2; var program = new ResizeNearestNeigborBackpropProgram(images.shape, alignCorners); var uniformData = [ { type: 'int32', data: effectiveXSize }, { type: 'int32', data: effectiveYSize }, { type: 'float32', data: [invHeightScale] }, { type: 'float32', data: [invWidthScale] }, { type: 'int32', data: [winHeight] }, { type: 'int32', data: [winWidth] } ]; return backend.runWebGPUProgram(program, [dy], dy.dtype, uniformData); } var resizeNearestNeighborGradConfig = { kernelName: tf.ResizeNearestNeighborGrad, backendName: 'webgpu', kernelFunc: resizeNearestNeighborGrad }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var ReverseProgram = /** @class */ (function () { function ReverseProgram(xShape) { this.variableNames = ['x']; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = xShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.uniforms = " axis : vec4,"; this.shaderKey = 'reverse'; } ReverseProgram.prototype.getUserCode = function () { var reverseCoordsSnippet = "\n // Using uniform variables as judging conditions, so the function has\n // coherent execution within all threads.\n fn getReverseCoords(coords : vec4) -> vec4 {\n var reverseCoords = coords;\n if (uniforms.axis[0] == 1) {\n reverseCoords[0] = uniforms.xShape[0] - coords[0] - 1;\n }\n if (uniforms.axis[1] == 1) {\n reverseCoords[1] = uniforms.xShape[1] - coords[1] - 1;\n }\n if (uniforms.axis[2] == 1) {\n reverseCoords[2] = uniforms.xShape[2] - coords[2] - 1;\n }\n if (uniforms.axis[3] == 1) {\n reverseCoords[3] = uniforms.xShape[3] - coords[3] - 1;\n }\n\n return reverseCoords;\n }\n "; var userCode = "\n ".concat(reverseCoordsSnippet, "\n ").concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let reverseCoords = getReverseCoords(coords);\n setOutputAtIndex(index, getX(reverseCoords[0],\n reverseCoords[1], reverseCoords[2], reverseCoords[3]));\n }\n }\n "); return userCode; }; return ReverseProgram; }()); /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function reverse(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var dims = attrs.dims; var xRank = x.shape.length; if (xRank === 0) { return identity({ inputs: { x: x }, backend: backend }); } var xShape = x.shape; var xShape4D = [1, 1, 1, 1]; xShape.forEach(function (d, i) { var index = i + 4 - xRank; xShape4D[index] = d; }); var axes = tf.util.parseAxisParam(dims, x.shape); var dims4D = [0, 0, 0, 0]; axes.forEach(function (ax) { var index = ax + 4 - xRank; dims4D[index] = 1; }); var uniformData = [{ type: 'int32', data: dims4D }]; var xReshaped = reshape({ inputs: { x: x }, backend: backend, attrs: { shape: xShape4D } }); var program = new ReverseProgram(xShape4D); var values = backend.runWebGPUProgram(program, [xReshaped], xReshaped.dtype, uniformData); backend.disposeData(xReshaped.dataId); var result = reshape({ inputs: { x: values }, backend: backend, attrs: { shape: xShape } }); backend.disposeData(values.dataId); return result; } var reverseConfig = { kernelName: tf.Reverse, backendName: 'webgpu', kernelFunc: reverse }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var RotateProgram = /** @class */ (function () { function RotateProgram(imageShape, fillValue) { this.outputShape = []; this.variableNames = ['x']; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = imageShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.uniforms = "centerX : f32, centerY : f32, sinRadians : f32,\n cosRadians : f32,"; this.shaderKey = 'rotate'; this.outputShape = imageShape; if (typeof fillValue === 'number') { this.uniforms += " fillValue : f32,"; this.fillSnippet = "var outputValue = uniforms.fillValue;"; this.shaderKey += '_float'; } else { this.uniforms += " fillValue : vec3,"; this.fillSnippet = "var outputValue = uniforms.fillValue[coords[3]];"; this.shaderKey += '_vec3'; } } RotateProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let coordXFloat = (f32(coords[2]) - uniforms.centerX) *\n uniforms.cosRadians - (f32(coords[1]) - uniforms.centerY) *\n uniforms.sinRadians;\n let coordYFloat = (f32(coords[2]) - uniforms.centerX) *\n uniforms.sinRadians + (f32(coords[1]) - uniforms.centerY) *\n uniforms.cosRadians;\n let coordX = i32(round(coordXFloat + uniforms.centerX));\n let coordY = i32(round(coordYFloat + uniforms.centerY));\n ").concat(this.fillSnippet, "\n if(coordX >= 0 && coordX < uniforms.xShape[2] && coordY >= 0 &&\n coordY < uniforms.xShape[1]) {\n outputValue = getX(coords[0], coordY, coordX, coords[3]);\n }\n setOutputAtIndex(index, outputValue);\n }\n }\n "); return userCode; }; return RotateProgram; }()); var rotateWithOffsetConfig = { kernelName: tf.RotateWithOffset, backendName: 'webgpu', kernelFunc: function (_a) { var inputs = _a.inputs, attrs = _a.attrs, backend = _a.backend; var image = inputs.image; var radians = attrs.radians, fillValue = attrs.fillValue, center = attrs.center; var webgpuBackend = backend; var program = new RotateProgram(image.shape, fillValue); var _b = __read(tf.backend_util.getImageCenter(center, image.shape[1], image.shape[2]), 2), centerX = _b[0], centerY = _b[1]; var uniformData = [ { type: 'float32', data: [centerX] }, { type: 'float32', data: [centerY] }, { type: 'float32', data: [Math.sin(radians)] }, { type: 'float32', data: [Math.cos(radians)] } ]; if (typeof fillValue === 'number') { uniformData.push({ type: 'float32', data: [Number.parseFloat(fillValue.toFixed(2))] }); } else { uniformData.push({ type: 'float32', data: fillValue }); } var output = webgpuBackend.runWebGPUProgram(program, [image], image.dtype, uniformData); return output; } }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var round = unaryKernelFunc({ opType: UnaryOpType.ROUND }); var roundConfig = { kernelName: tf.Round, backendName: 'webgpu', kernelFunc: round }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var rsqrt = unaryKernelFunc({ opType: UnaryOpType.RSQRT, cpuKernelImpl: rsqrtImplCPU }); var rsqrtConfig = { kernelName: tf.Rsqrt, backendName: 'webgpu', kernelFunc: rsqrt }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var ScatterProgram = /** @class */ (function () { function ScatterProgram(flattenXShape, sliceDim, indicesRank, updatesRank, strides, shape, outputDtype, sumDupeIndices) { if (sumDupeIndices === void 0) { sumDupeIndices = true; } this.variableNames = ['updates', 'indices']; this.workgroupSize = [64, 1, 1]; this.atomic = true; this.outputShape = shape; this.type = outputDtype; this.sumDupeIndices = sumDupeIndices; this.dispatchLayout = flatDispatchLayout(flattenXShape); // Dispatching based on |updates| shape instead of output shape. this.dispatch = computeDispatch(this.dispatchLayout, flattenXShape, this.workgroupSize); this.sliceDimGreaterThanOne = sliceDim > 1; this.shaderKey = "scatter_".concat(indicesRank, "_").concat(updatesRank, "_").concat(this.sliceDimGreaterThanOne, "_").concat(outputDtype, "_").concat(sumDupeIndices, "_").concat(strides.length); var stridesType = getCoordsDataType(strides.length); this.uniforms = "sliceDim : i32, strides: ".concat(stridesType, ", updatesSize: i32,"); this.updatesRank = updatesRank; this.indicesRank = indicesRank; } ScatterProgram.prototype.getUserCode = function () { var indicesString = ''; if (this.indicesRank === 1) { indicesString = 'coords[0]'; } else if (this.indicesRank === 2) { indicesString = 'coords[0], j'; } var indicesSnippet = "getIndices(".concat(indicesString, ")"); var strideString = this.sliceDimGreaterThanOne ? 'uniforms.strides[j]' : 'uniforms.strides'; var outCoordsString = ''; var getUpdatesCoordsFromFlatIndex = ''; if (this.dispatchLayout.x.length === 1) { outCoordsString = 'flattenedIndex'; getUpdatesCoordsFromFlatIndex = "\n fn getUpdatesCoordsFromFlatIndex(index : i32) -> i32 {\n return index;\n }\n "; } else if (this.dispatchLayout.x.length === 2) { outCoordsString = 'vec2(flattenedIndex, coords[1])'; getUpdatesCoordsFromFlatIndex = "\n fn getUpdatesCoordsFromFlatIndex(index : i32) -> vec2 {\n // N.B. |updates| could be a scalar tensor, conceptually representing a\n // 2D tensor with all values equal to that. By design, its size must be\n // the same as |outShape[1]| in one dimension, and |indicesShape[0]|\n // gives the other.\n let sliceSize = uniforms.outShape[1];\n let d0 = index / sliceSize;\n let d1 = index - d0 * sliceSize;\n return vec2(d0, d1);\n }\n "; } var updatesString = Array.from({ length: this.updatesRank }, function (_, idx) { return "coords[".concat(idx, "]"); }); var updatesSnippet = "getUpdates(".concat(updatesString.join(', '), ")"); var userCode = "\n ".concat(getUpdatesCoordsFromFlatIndex, "\n ").concat(getMainHeaderString('index'), " {\n if (index < uniforms.updatesSize) {\n let coords = getUpdatesCoordsFromFlatIndex(index);\n var flattenedIndex = 0;\n for (var j = 0; j < uniforms.sliceDim; j = j + 1) {\n let indexInside = i32(round(").concat(indicesSnippet, "));\n flattenedIndex = flattenedIndex + indexInside * ").concat(strideString, ";\n }\n let updateValue =\n ").concat(dataTypeToGPUType(this.type), "(").concat(updatesSnippet, ");\n let flatIndex = getOutputIndexFromCoords(").concat(outCoordsString, ");\n\n ").concat(this.sumDupeIndices ? atomicAddSnippet('&result[flatIndex]', 'updateValue', this.type) : "atomicStore(&result[flatIndex], bitcast(updateValue));", "\n }\n }"); return userCode; }; return ScatterProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function scatterNd(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var indices = inputs.indices, updates = inputs.updates; var shape = attrs.shape; var _a = tf.backend_util.calculateShapes(updates, indices, shape), sliceRank = _a.sliceRank, numUpdates = _a.numUpdates, sliceSize = _a.sliceSize, strides = _a.strides, outputSize = _a.outputSize; var flattenShape = [outputSize / sliceSize, sliceSize]; if (outputSize === 0) { return backend.makeTensorInfo(shape, indices.dtype); } var flattenIndices = reshape({ inputs: { x: indices }, backend: backend, attrs: { shape: [numUpdates, sliceRank] } }); var flattenX = reshape({ inputs: { x: updates }, backend: backend, attrs: { shape: [numUpdates, sliceSize] } }); var type = flattenX.dtype; var output = fill({ backend: backend, attrs: { shape: flattenShape, value: 0, dtype: type } }); var size = tf.util.sizeFromShape(flattenX.shape); var uniformData = [ { type: 'int32', data: [sliceRank] }, { type: 'int32', data: strides }, { type: 'int32', data: [size] } ]; var program = new ScatterProgram(flattenX.shape, sliceRank, flattenIndices.shape.length, flattenX.shape.length, strides, flattenShape, type); var res = backend.runWebGPUProgram(program, [flattenX, flattenIndices], type, uniformData, output); var reshaped = reshape({ inputs: { x: res }, backend: backend, attrs: { shape: shape } }); backend.disposeData(flattenIndices.dataId); backend.disposeData(flattenX.dataId); backend.disposeData(res.dataId); return reshaped; } var scatterNdConfig = { kernelName: tf.ScatterNd, backendName: 'webgpu', kernelFunc: scatterNd }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var SearchSortedProgram = /** @class */ (function () { function SearchSortedProgram(outputShape, side) { this.outputShape = []; this.variableNames = ['sortedSequence', 'values']; this.uniforms = 'numInputs : i32,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = outputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.side = side; this.shaderKey = "search_sorted_".concat(side); } SearchSortedProgram.prototype.getUserCode = function () { var boundComparator = this.side === 'left' ? '<' : '<='; var userCode = "\n fn findBound(batch: i32, value: f32) -> i32 {\n var left = i32(0);\n var right = uniforms.numInputs;\n while (left < right) {\n var mid = (left + right) / 2;\n if (getSortedSequence(batch, mid) ".concat(boundComparator, " value) {\n left = mid + 1;\n } else {\n right = mid;\n }\n }\n return right;\n }\n\n ").concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let value = getValuesByOutputIndex(index);\n setOutputAtIndexI32(index, findBound(coords[0], value));\n }\n }\n "); return userCode; }; return SearchSortedProgram; }()); /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function searchSorted(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var sortedSequence = inputs.sortedSequence, values = inputs.values; var side = attrs.side; var program = new SearchSortedProgram([values.shape[0], values.shape[1]], side); var uniformData = [{ type: 'int32', data: [sortedSequence.shape[1]] }]; return backend.runWebGPUProgram(program, [sortedSequence, values], 'int32', uniformData); } var searchSortedConfig = { kernelName: tf.SearchSorted, backendName: 'webgpu', kernelFunc: searchSorted, }; /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var SelectProgram = /** @class */ (function () { function SelectProgram(cRank, shape, rank) { this.variableNames = ['c', 'a', 'b']; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = shape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.cRank = cRank; this.rank = rank; this.shaderKey = 'select'; } SelectProgram.prototype.getUserCode = function () { // TODO(WGSL): below code can be merged with getUserCode. var cCoords; var abCoords; if (this.rank > 4) { throw Error("Where for rank ".concat(this.rank, " is not yet supported")); } if (this.rank === 1) { abCoords = "resRC"; cCoords = "resRC"; } else { var currentCoords = ['resRC.x', 'resRC.y', 'resRC.z', 'resRC.w']; var cCoordVars = []; var abCoordVars = []; for (var i = 0; i < this.outputShape.length; i++) { abCoordVars.push("".concat(currentCoords[i])); if (i < this.cRank) { cCoordVars.push("".concat(currentCoords[i])); } } cCoords = cCoordVars.join(); abCoords = abCoordVars.join(); } var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let resRC = getCoordsFromIndex(index);\n let cVal = getC(").concat(cCoords, ");\n if (cVal >= 1.0) {\n setOutputAtIndex(index, getA(").concat(abCoords, "));\n } else {\n setOutputAtIndex(index, getB(").concat(abCoords, "));\n }\n }\n }\n "); return userCode; }; return SelectProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function select(args) { var inputs = args.inputs, backend = args.backend; var condition = inputs.condition, t = inputs.t, e = inputs.e; var program = new SelectProgram(condition.shape.length, t.shape, t.shape.length); return backend.runWebGPUProgram(program, [condition, t, e], tf.upcastType(t.dtype, e.dtype)); } var selectConfig = { kernelName: tf.Select, backendName: 'webgpu', kernelFunc: select }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var selu = unaryKernelFunc({ opType: UnaryOpType.SELU }); var seluConfig = { kernelName: tf.Selu, backendName: 'webgpu', kernelFunc: selu }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var sigmoid = unaryKernelFunc({ opType: UnaryOpType.SIGMOID }); var sigmoidConfig = { kernelName: tf.Sigmoid, backendName: 'webgpu', kernelFunc: sigmoid, }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var sign = unaryKernelFunc({ opType: UnaryOpType.SIGN }); var signConfig = { kernelName: tf.Sign, backendName: 'webgpu', kernelFunc: sign }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var sin = unaryKernelFunc({ opType: UnaryOpType.SIN }); var sinConfig = { kernelName: tf.Sin, backendName: 'webgpu', kernelFunc: sin }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var sinh = unaryKernelFunc({ opType: UnaryOpType.SINH }); var sinhConfig = { kernelName: tf.Sinh, backendName: 'webgpu', kernelFunc: sinh }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var softplus = unaryKernelFunc({ opType: UnaryOpType.SOFTPLUS }); var softplusConfig = { kernelName: tf.Softplus, backendName: 'webgpu', kernelFunc: softplus }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var SpaceToBatchNDProgram = /** @class */ (function () { function SpaceToBatchNDProgram(xShape, paddedXShape, paddings, reshapedPaddedXShape, newDim, paddedXShapeStridesShapeLength) { var _this = this; this.variableNames = ['x']; this.outputShape = []; this.uniforms = ''; this.workgroupSize = [64, 1, 1]; this.size = true; var outputShape = new Array(reshapedPaddedXShape.length); for (var i = 0; i < outputShape.length; i++) { outputShape[i] = reshapedPaddedXShape[newDim[i]]; } this.outputShape = outputShape; this.newDim = newDim; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.xShape = xShape; this.paddedXShape = paddedXShape; this.uniforms += "reshapedPaddedXShape : ".concat(getCoordsDataType(reshapedPaddedXShape.length), ", paddedXShapeStrides : ").concat(getCoordsDataType(paddedXShapeStridesShapeLength), ", "); paddings.map(function (_, i) { _this.uniforms += " pad".concat(i, " : vec2,"); }); this.shaderKey = "spaceToBatchND_".concat(newDim); } SpaceToBatchNDProgram.prototype.getUserCode = function () { var dtype = getCoordsDataType(this.outputShape.length); var switched = getSwitchedCoords(this.newDim); var userCode = "\n ".concat(getCoordsFromIndexSnippet(this.paddedXShape, 'PaddedX'), "\n ").concat(getMainHeaderString('index'), " {\n if(index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n let switchedIndex = getIndexFromCoords").concat(this.outputShape.length, "D(").concat(dtype, "(").concat(switched, "), uniforms.reshapedPaddedXShape);\n let paddedCoords = getPaddedXCoordsFromIndex(switchedIndex);\n ").concat(padCommon(this.xShape, true), "\n }\n }\n "); return userCode; }; return SpaceToBatchNDProgram; }()); var spaceToBatchND = function (args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var blockShape = attrs.blockShape, paddings = attrs.paddings; tf.util.assert(x.shape.length <= 4, function () { return 'spaceToBatchND for rank > 4 with a WebGPU backend not ' + 'implemented yet'; }); var prod = blockShape.reduce(function (a, b) { return a * b; }); var completePaddings = [[0, 0]]; completePaddings.push.apply(completePaddings, __spreadArray([], __read(paddings), false)); for (var i = 1 + blockShape.length; i < x.shape.length; ++i) { completePaddings.push([0, 0]); } var paddedXShape = completePaddings.map(function (p, i) { return p[0] /* beforePad */ + x.shape[i] + p[1]; } /* afterPad */); var reshapedPaddedShape = tf.backend_util.getReshaped(paddedXShape, blockShape, prod, false); var permutedReshapedPaddedPermutation = tf.backend_util.getPermuted(reshapedPaddedShape.length, blockShape.length, false); var flattenShape = tf.backend_util.getReshapedPermuted(paddedXShape, blockShape, prod, false); var paddedXShapeStrides = tf.util.computeStrides(paddedXShape); var program = new SpaceToBatchNDProgram(x.shape, paddedXShape, completePaddings, reshapedPaddedShape, permutedReshapedPaddedPermutation, paddedXShapeStrides.length); var uniformData = [ { type: 'int32', data: reshapedPaddedShape }, { type: 'int32', data: paddedXShapeStrides } ]; completePaddings.map(function (p) { return uniformData.push({ type: 'int32', data: [p[0], p[1]] }); }); var paddedXT = backend.runWebGPUProgram(program, [x], x.dtype, uniformData); var result = reshape({ inputs: { x: paddedXT }, backend: backend, attrs: { shape: flattenShape } }); backend.disposeData(paddedXT.dataId); return result; }; var spaceToBatchNDConfig = { kernelName: tf.SpaceToBatchND, backendName: 'webgpu', kernelFunc: spaceToBatchND }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var SparseSegmentSumProgram = /** @class */ (function () { function SparseSegmentSumProgram(outShape, sparseSize, outputDtype) { this.variableNames = ['input', 'indices', 'segmentIds']; this.outputShape = []; this.uniforms = 'segmentSize : i32, sparseSize : i32,'; this.workgroupSize = [64, 1, 1]; this.atomic = true; this.outputShape = outShape; this.type = outputDtype; this.dispatchLayout = flatDispatchLayout([sparseSize]); this.dispatch = computeDispatch(this.dispatchLayout, [sparseSize], this.workgroupSize); this.shaderKey = 'sparseSegmentSum'; } SparseSegmentSumProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.sparseSize) {\n let indexInSegmentIds = index / uniforms.segmentSize;\n let indexInSegment = index % uniforms.segmentSize;\n let indexInInput = indices[indexInSegmentIds];\n let segmentId = segmentIds[indexInSegmentIds];\n\n let value = input[indexInInput * uniforms.segmentSize + indexInSegment];\n let outIndex = segmentId * uniforms.segmentSize + indexInSegment;\n ").concat(atomicAddSnippet('&result[outIndex]', 'value', this.type), "\n }\n }\n "); return userCode; }; return SparseSegmentSumProgram; }()); var SparseSegmentIdCountProgram = /** @class */ (function () { function SparseSegmentIdCountProgram(outShape, segmentIdsShape) { this.variableNames = ['segmentIds']; this.outputShape = []; this.workgroupSize = [64, 1, 1]; this.atomic = true; this.outputShape = [outShape]; this.dispatchLayout = flatDispatchLayout(segmentIdsShape); this.dispatch = computeDispatch(this.dispatchLayout, segmentIdsShape, this.workgroupSize); this.shaderKey = 'sparseSegmentIdCountProgram'; } SparseSegmentIdCountProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.segmentIdsShape) {\n let segmentId = segmentIds[index];\n ").concat(atomicAddSnippet('&result[segmentId]', '1', 'int32'), "\n }\n }\n "); return userCode; }; return SparseSegmentIdCountProgram; }()); var SparseSegmentMeanProgram = /** @class */ (function () { function SparseSegmentMeanProgram(outShape, outputDtype) { this.variableNames = ['segmentSum', 'sameSegmentIdCount']; this.outputShape = []; this.uniforms = 'segmentSize : i32'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = outShape; this.type = outputDtype; this.dispatchLayout = flatDispatchLayout(outShape); this.dispatch = computeDispatch(this.dispatchLayout, outShape, this.workgroupSize); this.shaderKey = 'sparseSegmentMean'; } SparseSegmentMeanProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let segmentId = index / uniforms.segmentSize;\n let count = sameSegmentIdCount[segmentId];\n if (count != 0) {\n ").concat(this.type === 'float32' ? 'setOutputAtIndex(index, segmentSum[index] / f32(count));' : 'setOutputAtIndexI32(index, segmentSum[index] / count);', "\n }\n }\n }\n "); return userCode; }; return SparseSegmentMeanProgram; }()); /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function sparseSegmentReduce(input, indices, segmentIds, isSum, backend) { if (isSum === void 0) { isSum = false; } var inputSize = tf.util.sizeFromShape(input.shape); var segmentSize = inputSize / input.shape[0]; var dtype = input.dtype; // Note that the current implementation assumes that segmentIds values are // sorted. var numIndices = tf.util.sizeFromShape(indices.shape); var $segmentIds = backend.readSync(segmentIds.dataId); var lastSegmentIdPlusOne = numIndices > 0 ? $segmentIds[numIndices - 1] + 1 : 0; var outputRows = lastSegmentIdPlusOne; var program; var outputShape = input.shape.slice(); outputShape[0] = outputRows; var sparseSize = numIndices * segmentSize; var sparseSegmentSum = fill({ backend: backend, attrs: { shape: outputShape, value: 0, dtype: dtype } }); program = new SparseSegmentSumProgram(outputShape, sparseSize, dtype); var uniformData = [ { type: 'int32', data: [segmentSize] }, { type: 'int32', data: [sparseSize] } ]; var $sparseSegmentSum = backend.runWebGPUProgram(program, [input, indices, segmentIds], dtype, uniformData, sparseSegmentSum); if (isSum) { return $sparseSegmentSum; } var sparseSegmentIdCount = fill({ backend: backend, attrs: { shape: [outputRows], value: 0, dtype: 'int32' } }); program = new SparseSegmentIdCountProgram(outputRows, segmentIds.shape); var $sparseSegmentIdCount = backend.runWebGPUProgram(program, [segmentIds], 'int32', null, sparseSegmentIdCount); var sparseSegmentMean = fill({ backend: backend, attrs: { shape: outputShape, value: 0, dtype: dtype } }); program = new SparseSegmentMeanProgram(outputShape, dtype); uniformData = [{ type: 'int32', data: [segmentSize] }]; var $sparseSegmentMean = backend.runWebGPUProgram(program, [$sparseSegmentSum, $sparseSegmentIdCount], dtype, uniformData, sparseSegmentMean); backend.disposeData($sparseSegmentSum.dataId); backend.disposeData($sparseSegmentIdCount.dataId); return $sparseSegmentMean; } /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function sparseSegmentMean(args) { var inputs = args.inputs, backend = args.backend; var data = inputs.data, indices = inputs.indices, segmentIds = inputs.segmentIds; return sparseSegmentReduce(data, indices, segmentIds, false, backend); } var sparseSegmentMeanConfig = { kernelName: tf.SparseSegmentMean, backendName: 'webgpu', kernelFunc: sparseSegmentMean, }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function sparseSegmentSum(args) { var inputs = args.inputs, backend = args.backend; var data = inputs.data, indices = inputs.indices, segmentIds = inputs.segmentIds; return sparseSegmentReduce(data, indices, segmentIds, true, backend); } var sparseSegmentSumConfig = { kernelName: tf.SparseSegmentSum, backendName: 'webgpu', kernelFunc: sparseSegmentSum, }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var TileProgram = /** @class */ (function () { function TileProgram(aShape, reps) { this.variableNames = ['A']; this.workgroupSize = [64, 1, 1]; this.size = true; var outputShape = new Array(aShape.length); for (var i = 0; i < outputShape.length; i++) { outputShape[i] = aShape[i] * reps[i]; } this.outputShape = outputShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.rank = this.outputShape.length; this.shaderKey = 'tile'; } TileProgram.prototype.getUserCode = function () { var sourceCoords = getSourceCoords(this.rank, 'uniforms.'); var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let resRC = getCoordsFromIndex(index);\n setOutputAtIndex(index, getA(").concat(sourceCoords, "));\n }\n }\n "); return userCode; }; return TileProgram; }()); function getSourceCoords(rank, uniformPrefix) { if (uniformPrefix === void 0) { uniformPrefix = ''; } if (rank >= 5) { throw Error("Tile for rank ".concat(rank, " is not yet supported")); } if (rank === 1) { return "(resRC % ".concat(uniformPrefix, "aShape)"); } var currentCoords = ['resRC.x', 'resRC.y', 'resRC.z', 'resRC.w']; var sourceCoords = []; for (var i = 0; i < rank; i++) { sourceCoords.push("(".concat(currentCoords[i], " % ").concat(uniformPrefix, "aShape[").concat(i, "])")); } return sourceCoords.join(); } /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function tile(params) { var inputs = params.inputs, backend = params.backend, attrs = params.attrs; var x = inputs.x; var reps = attrs.reps; // tile gpu program cannot handle rank >= 5 case. if (backend.shouldExecuteOnCPU([x]) || x.dtype === 'string' || x.shape.length >= 5) { // Even thought string tensor is always on CPU, just to be consistent on how // to access tensor data. var data = backend.readSync(x.dataId); var value = x.dtype === 'string' ? data.map(function (d) { return tf.util.decodeString(d); }) : data; var buf = tf.buffer(x.shape, x.dtype, value); var outBuf = tileImplCPU(buf, reps); return backend.makeTensorInfo(outBuf.shape, outBuf.dtype, outBuf.values); } var program = new TileProgram(x.shape, reps); var output = backend.runWebGPUProgram(program, [x], x.dtype); return output; } var tileConfig = { kernelName: tf.Tile, backendName: 'webgpu', kernelFunc: tile, }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function sparseToDense(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var sparseIndices = inputs.sparseIndices, sparseValues = inputs.sparseValues, defaultValue = inputs.defaultValue; var outputShape = attrs.outputShape; var _a = tf.backend_util.calculateShapes(sparseValues, sparseIndices, outputShape), sliceRank = _a.sliceRank, numUpdates = _a.numUpdates, sliceSize = _a.sliceSize, strides = _a.strides, outputSize = _a.outputSize; var sumDupeIndices = false; if (sparseValues.dtype === 'string') { var indicesBuf = backend.bufferSync(sparseIndices); var updatesBuf = backend.bufferSync(sparseValues); var $defaultValue_1 = tf.util.decodeString(backend.readSync(defaultValue.dataId)[0]); var outBuf = scatterImplCPU(indicesBuf, updatesBuf, outputShape, outputSize, sliceSize, numUpdates, sliceRank, strides, $defaultValue_1, sumDupeIndices); return backend.makeTensorInfo(outputShape, outBuf.dtype, outBuf.values); } var flattenShape = [outputSize / sliceSize, sliceSize]; var $sparseIndices = reshape({ inputs: { x: sparseIndices }, backend: backend, attrs: { shape: [numUpdates, sliceRank] } }); var $sparseValues = sparseValues.shape.length ? reshape({ inputs: { x: sparseValues }, backend: backend, attrs: { shape: [numUpdates, sliceSize] } }) : identity({ inputs: { x: sparseValues }, backend: backend }); var type = $sparseValues.dtype; var zero = backend.makeTensorInfo([], type, tf.util.makeZerosTypedArray(1, type)); // Fill output tensor with the default value. var $defaultValue = reshape({ inputs: { x: defaultValue }, backend: backend, attrs: { shape: Array(flattenShape.length).fill(1) } }); var $denseValues = tile({ inputs: { x: $defaultValue }, backend: backend, attrs: { reps: flattenShape } }); var size = tf.util.sizeFromShape([numUpdates, sliceSize]); var uniformData = [ { type: 'int32', data: [sliceRank] }, { type: 'int32', data: strides }, { type: 'int32', data: [size] }, ]; switch (numUpdates) { case 0: break; case 1: { var program = new ScatterProgram([numUpdates, sliceSize], sliceRank, $sparseIndices.shape.length, $sparseValues.shape.length, strides, flattenShape, type, sumDupeIndices); backend.runWebGPUProgram(program, [$sparseValues, $sparseIndices], type, uniformData, $denseValues); } break; default: { // First replace the default value with 0 at indices. var program = new ScatterProgram([numUpdates, sliceSize], sliceRank, $sparseIndices.shape.length, zero.shape.length, strides, flattenShape, type, sumDupeIndices); backend.runWebGPUProgram(program, [zero, $sparseIndices], type, uniformData, $denseValues); } { // Then replace 0 with the (sum of) sparse value(s) at indices. var program = new ScatterProgram([numUpdates, sliceSize], sliceRank, $sparseIndices.shape.length, $sparseValues.shape.length, strides, flattenShape, type); backend.runWebGPUProgram(program, [$sparseValues, $sparseIndices], type, uniformData, $denseValues); } } var denseValues = reshape({ inputs: { x: $denseValues }, backend: backend, attrs: { shape: outputShape } }); backend.disposeData($sparseIndices.dataId); backend.disposeData($sparseValues.dataId); backend.disposeData($defaultValue.dataId); backend.disposeData(zero.dataId); backend.disposeData($denseValues.dataId); return denseValues; } var sparseToDenseConfig = { kernelName: tf.SparseToDense, backendName: 'webgpu', kernelFunc: sparseToDense }; function splitV(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var numOrSizeSplits = attrs.numOrSizeSplits, axis = attrs.axis; var $axis = tf.util.parseAxisParam(axis, x.shape)[0]; var splitSizes = tf.backend_util.prepareSplitSize(x, numOrSizeSplits, $axis); var xRank = x.shape.length; var begin = new Array(xRank).fill(0); var size = x.shape.slice(); return splitSizes.map(function (s) { var sliceSize = __spreadArray([], __read(size), false); sliceSize[$axis] = s; var sliceT = slice({ inputs: { x: x }, backend: backend, attrs: { begin: begin, size: sliceSize } }); begin[$axis] += s; return sliceT; }); } var splitVConfig = { kernelName: tf.SplitV, backendName: 'webgpu', kernelFunc: splitV }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var sqrt = unaryKernelFunc({ opType: UnaryOpType.SQRT }); var sqrtConfig = { kernelName: tf.Sqrt, backendName: 'webgpu', kernelFunc: sqrt }; /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var squareConfig = { kernelName: tf.Square, backendName: 'webgpu', kernelFunc: function (_a) { var inputs = _a.inputs, backend = _a.backend; var x = inputs.x; var webGPUBackend = backend; var program = new UnaryOpProgram(x.shape, UnaryOpType.SQUARE); return webGPUBackend.runWebGPUProgram(program, [x], x.dtype); } }; /** * @license * Copyright 2020 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var squaredDifference = binaryKernelFunc({ opType: BinaryOpType.SQUARED_DIFFERENCE, }); var squaredDifferenceConfig = { kernelName: tf.SquaredDifference, backendName: 'webgpu', kernelFunc: squaredDifference }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function step(_a) { var inputs = _a.inputs, attrs = _a.attrs, backend = _a.backend; var x = inputs.x; var program = new UnaryOpProgram(x.shape, UnaryOpType.STEP, 'stepAlpha : f32,'); var uniformData = [{ type: 'float32', data: [attrs.alpha] }]; return backend.runWebGPUProgram(program, [x], x.dtype, uniformData); } var stepConfig = { kernelName: tf.Step, backendName: 'webgpu', kernelFunc: step }; /** * @license * Copyright 2019 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var StridedSliceProgram = /** @class */ (function () { function StridedSliceProgram(destSize) { this.variableNames = ['x']; // TODO(xing.xu): Increase the workPerThread. this.workPerThread = 1; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = destSize; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize, [this.workPerThread, 1, 1]); var dtype = getCoordsDataType(this.outputShape.length); this.uniforms = "begin : ".concat(dtype, ", strides : ").concat(dtype, ", "); this.shaderKey = 'stridedSlice'; } StridedSliceProgram.prototype.getUserCode = function () { var _this = this; var rank = this.outputShape.length; var newCoords = ''; if (rank === 1) { newCoords = 'coords * uniforms.strides + uniforms.begin'; } else { var outputAxis_1 = 0; newCoords = this.outputShape .map(function (_, i) { outputAxis_1++; return _this.outputShape.length === 1 ? "coords * uniforms.strides[".concat(i, "] + uniforms.begin[").concat(i, "]") : "coords[".concat(outputAxis_1 - 1, "] * uniforms.strides[").concat(i, "] + uniforms.begin[").concat(i, "]"); }) .join(','); } var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n setOutputAtIndex(index, getX(").concat(newCoords, "));\n }\n }\n "); return userCode; }; return StridedSliceProgram; }()); /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function stridedSlice(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var begin = attrs.begin, end = attrs.end, strides = attrs.strides, beginMask = attrs.beginMask, endMask = attrs.endMask, ellipsisMask = attrs.ellipsisMask, newAxisMask = attrs.newAxisMask, shrinkAxisMask = attrs.shrinkAxisMask; var _a = tf.slice_util.sliceInfo(x.shape, begin, end, strides, beginMask, endMask, ellipsisMask, newAxisMask, shrinkAxisMask), finalShapeSparse = _a.finalShapeSparse, finalShape = _a.finalShape, isIdentity = _a.isIdentity, sliceDim0 = _a.sliceDim0, isSimpleSlice = _a.isSimpleSlice, $begin = _a.begin, $end = _a.end, $strides = _a.strides; var result; if (isIdentity) { // Optimization #1, slice is a no-op plus reshape result = reshape({ inputs: { x: x }, backend: backend, attrs: { shape: finalShape } }); } else if (sliceDim0 || isSimpleSlice) { // Optimization #2, slice is memory contiguous (only occurs in dim 0) tf.util.assert(x.shape.length >= 1, function () { return "Input must have rank at least 1, got: ".concat(x.shape.length); }); var size = tf.slice_util.computeOutShape($begin, $end, $strides); // To tolerate begin[0] > end[0] (a 0-output slice), we min(begin, end). var sliced = slice({ inputs: { x: x }, backend: backend, attrs: { begin: $begin, size: size } }); result = reshape({ inputs: { x: sliced }, backend: backend, attrs: { shape: finalShape } }); backend.disposeData(sliced.dataId); } else { var shouldExecuteOnCPU = backend.shouldExecuteOnCPU([x]); if (shouldExecuteOnCPU) { var values = backend.readSync(x.dataId); var xBuf = tf.buffer(x.shape, x.dtype, values); var resultValues = stridedSliceImplCPU(finalShapeSparse, xBuf, $strides, $begin); result = backend.makeTensorInfo(finalShape, x.dtype, resultValues.values); } else { var program = new StridedSliceProgram(finalShapeSparse); var uniformData = [{ type: 'int32', data: $begin }, { type: 'int32', data: $strides }]; var resultValues = backend.runWebGPUProgram(program, [x], x.dtype, uniformData); result = reshape({ inputs: { x: resultValues }, backend: backend, attrs: { shape: finalShape } }); backend.disposeData(resultValues.dataId); } } return result; } var stridedSliceConfig = { kernelName: tf.StridedSlice, backendName: 'webgpu', kernelFunc: stridedSlice }; function stringNGrams(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var separator = attrs.separator, nGramWidths = attrs.nGramWidths, leftPad = attrs.leftPad, rightPad = attrs.rightPad, padWidth = attrs.padWidth, preserveShortSequences = attrs.preserveShortSequences; var data = inputs.data, dataSplits = inputs.dataSplits; var $data = backend.readSync(data.dataId); var $dataSplits = backend.readSync(dataSplits.dataId); var _a = __read(stringNGramsImplCPU($data, $dataSplits, separator, nGramWidths, leftPad, rightPad, padWidth, preserveShortSequences), 2), nGrams = _a[0], nGramsSplits = _a[1]; return [ backend.makeTensorInfo([nGrams.length], 'string', nGrams), backend.makeTensorInfo(dataSplits.shape, 'int32', nGramsSplits), ]; } var stringNGramsConfig = { kernelName: tf.StringNGrams, backendName: 'webgpu', kernelFunc: stringNGrams, }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var sub = binaryKernelFunc({ opType: BinaryOpType.SUB, cpuKernelImpl: subImplCPU, supportsComplex: true }); var subConfig = { kernelName: tf.Sub, backendName: 'webgpu', kernelFunc: sub }; /** * @license * Copyright 2022 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var tan = unaryKernelFunc({ opType: UnaryOpType.TAN }); var tanConfig = { kernelName: tf.Tan, backendName: 'webgpu', kernelFunc: tan }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var tanh = unaryKernelFunc({ opType: UnaryOpType.TANH }); var tanhConfig = { kernelName: tf.Tanh, backendName: 'webgpu', kernelFunc: tanh }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function tensorScatterUpdate(args) { var inputs = args.inputs, backend = args.backend; args.attrs; var tensor = inputs.tensor, indices = inputs.indices, updates = inputs.updates; var _b = tf.backend_util.calculateShapes(updates, indices, tensor.shape), sliceRank = _b.sliceRank, numUpdates = _b.numUpdates, sliceSize = _b.sliceSize, strides = _b.strides, outputSize = _b.outputSize; var flattenShape = [outputSize / sliceSize, sliceSize]; if (outputSize === 0) { return backend.makeTensorInfo(tensor.shape, indices.dtype); } var toDispose = []; var flattenIndices = reshape({ inputs: { x: indices }, backend: backend, attrs: { shape: [numUpdates, sliceRank] } }); toDispose.push(flattenIndices); var flattenX = reshape({ inputs: { x: updates }, backend: backend, attrs: { shape: [numUpdates, sliceSize] } }); toDispose.push(flattenX); var flattenTensor = reshape({ inputs: { x: tensor }, backend: backend, attrs: { shape: flattenShape } }); toDispose.push(flattenTensor); var output = tile({ inputs: { x: flattenTensor }, backend: backend, attrs: { reps: Array(flattenShape.length).fill(1) } }); var program = new ScatterProgram([numUpdates, sliceSize], sliceRank, flattenIndices.shape.length, flattenX.shape.length, strides, flattenShape, tensor.dtype, false); var size = tf.util.sizeFromShape([numUpdates, sliceSize]); var uniformData = [ { type: 'int32', data: [sliceRank] }, { type: 'int32', data: strides }, { type: 'int32', data: [size] }, ]; var res = backend.runWebGPUProgram(program, [flattenX, flattenIndices], flattenTensor.dtype, uniformData, output); toDispose.push(res); var reshaped = reshape({ inputs: { x: res }, backend: backend, attrs: { shape: tensor.shape } }); toDispose.forEach(function (t) { return backend.disposeData(t.dataId); }); return reshaped; } var tensorScatterUpdateConfig = { kernelName: tf.TensorScatterUpdate, backendName: 'webgpu', kernelFunc: tensorScatterUpdate }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ // Based on Algorithm 2 of Bitonic Top K, ref: // https://anilshanbhag.in/static/papers/gputopk_sigmod18.pdf // The original algorithm is based on computing the top K only, however // since for TFJS we require the indices of the top K values as well then the // algorithm found here is a bit modified. Rather than producing the values // at each step, the indices containing the top K are generated instead. // The output values are not generated to reduce the number of outputs in the // GPU, the values can easily be retrieved from the indices using a gather // op. var SwapProgram = /** @class */ (function () { function SwapProgram(shape) { this.variableNames = ['x', 'indices']; this.workgroupSize = [256, 1, 1]; this.size = true; this.outputShape = shape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.uniforms = "inputSize : i32, firstPass : i32, negativeInf : f32,\n dir : i32, inc : i32,"; this.shaderKey = 'swap'; } SwapProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let outC = getCoordsFromIndex(index);\n let batch = outC[0];\n let elemIdx = outC[1];\n // We compare elements pair-wise within a group of size 2 * inc.\n // The comparing rule for each group alternates between ascending\n // and descending. Within each group, we compare each pair at\n // positions i and i+inc. To decide whether an element at position i\n // is x0 or x1, we mod it by 2 * inc, if the result is smaller than\n // inc, it is in the first half of the group, we denote it as x0,\n // otherwise we denote it as x1.\n // For example, as shown in the Bitonic top K paper referenced\n // above, Figure5(a) shows that element[1] is in the second half of\n // the group when group size is 2, but it is in the first half of\n // the group when group size is 4.\n let isFirstInPair = elemIdx % (2 * uniforms.inc) < uniforms.inc;\n var i = 0;\n if (isFirstInPair) {\n i = elemIdx;\n } else {\n i = elemIdx - uniforms.inc;\n }\n\n var i0 = 0;\n if (uniforms.firstPass == 1) {\n i0 = i;\n } else {\n i0 = i32(getIndices(batch, i));\n }\n\n var i1 = 0;\n if (uniforms.firstPass == 1) {\n i1 = i + uniforms.inc;\n } else {\n i1 = i32(getIndices(batch, i + uniforms.inc));\n }\n\n var x0 = f32(0.0);\n var x1 = f32(0.0);\n if (i0 < uniforms.inputSize) {\n x0 = getX(batch, i0);\n } else {\n x0 = uniforms.negativeInf;\n }\n if (i1 < uniforms.inputSize) {\n x1 = getX(batch, i1);\n } else {\n x1 = uniforms.negativeInf;\n }\n\n let reverse = elemIdx % (2 * uniforms.dir) >= uniforms.dir;\n let isGreater = x0 > x1 || (x0 == x1 && i1 > i0);\n if (reverse == isGreater) {\n // Elements in opposite order of direction\n let iTemp = i0;\n i0 = i1;\n i1 = iTemp;\n }\n if (isFirstInPair) {\n setOutputAtIndex(index, f32(i0));\n } else {\n setOutputAtIndex(index, f32(i1));\n }\n }\n }\n "); return userCode; }; return SwapProgram; }()); var MergeProgram = /** @class */ (function () { function MergeProgram(shape) { this.variableNames = ['x', 'indices']; this.workgroupSize = [256, 1, 1]; this.size = true; this.outputShape = shape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); // |n| Size of the original input of TopK // |firstPass| indicates if this is the first time swap is being used which // means no indices input containing the top K is present yet. // |k| Top k elements desired this.uniforms = "inputSize : i32, firstPass : i32, k : i32,"; this.shaderKey = 'merge'; } MergeProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let outC = getCoordsFromIndex(index);\n let batch = outC[0];\n let elemIdx = outC[1];\n // The output size is half of the previous size.\n // If the previous sequence is | | | | _ _ _ _ | | | | _ _ _ _\n // (k=4), we only need to output the indices at positions |, the\n // indices at positions _ can be thrown away, see Figure5(b) After\n // Phase 2 (Merge phase) in the Bitonic Top K paper referenced\n // above.\n // For example, the paper shows we only need to output the orange\n // bars. The output sequence should look like this | | | | | | | |.\n // Because the sequence is halved, to map the output index back to\n // the previous sequence to find the corresponding value, we need\n // to double the index. When we double the index, we basically\n // interpolate a position, so 2i looks like\n // | _ | _ | _ | _ | _ | _ | _. We move the | to the first k\n // position of each 2k positions by - elemIdx % k. E.g. for output\n // at index 4,5,6,7, we want to get the corresponding element at\n // original index 8,9,10,11, for output at index 8,9,10,11,\n // we want to get the corresponding element at original index\n // 16,17,18,19, so on and so forth.\n\n var i = 0;\n if (elemIdx < uniforms.k) {\n i = elemIdx;\n } else {\n i = elemIdx * 2 - elemIdx % uniforms.k;\n }\n var i0 = 0;\n if (uniforms.firstPass == 1) {\n i0 = i;\n } else {\n i0 = i32(getIndices(batch, i));\n }\n var i1 = 0;\n if (uniforms.firstPass == 1) {\n i1 = i + uniforms.k;\n } else {\n i1 = i32(getIndices(batch, i + uniforms.k));\n }\n\n let x0 = getX(batch, i0);\n var x1 = f32(0.0);\n if (i1 < uniforms.inputSize) {\n x1 = getX(batch, i1);\n } else {\n x1 = x0;\n }\n\n if (x0 >= x1) {\n setOutputAtIndex(index, f32(i0));\n } else {\n setOutputAtIndex(index, f32(i1));\n }\n }\n }\n "); return userCode; }; return MergeProgram; }()); function disposeIntermediateTensorInfoOrNull(backend, tensorInfo) { if (tensorInfo !== null) { backend.disposeData(tensorInfo.dataId); } } function roundUpToPow2(num) { var pow2 = 1; while (pow2 < num) { pow2 *= 2; } return pow2; } // Based on Algorithm 2 of Bitonic Top K, ref: // https://anilshanbhag.in/static/papers/gputopk_sigmod18.pdf function topK(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x; var k = attrs.k, sorted = attrs.sorted; var xShape = x.shape; var lastDim = xShape[xShape.length - 1]; if (backend.shouldExecuteOnCPU([x])) { var xVals = backend.readSync(x.dataId); var _a = __read(topKImplCPU(xVals, xShape, x.dtype, k, sorted), 2), allTopKVals = _a[0], allTopKIndices = _a[1]; return [ backend.makeTensorInfo(allTopKVals.shape, allTopKVals.dtype, allTopKVals.values), backend.makeTensorInfo(allTopKIndices.shape, allTopKIndices.dtype, allTopKIndices.values) ]; } if (k === 0) { xShape[xShape.length - 1] = 0; return [ backend.makeTensorInfo(xShape, x.dtype, []), backend.makeTensorInfo(xShape, 'int32', []) ]; } if (lastDim === 1 /* firstPass */) { return [ x, fill({ attrs: { shape: xShape, dtype: 'int32', value: 0 }, backend: backend }) ]; } // Reshape into a 2d tensor [batch, lastDim] and compute topk along lastDim. var xSize = tf.util.sizeFromShape(xShape); var batch = xSize / lastDim; var x2D = reshape({ inputs: { x: x }, attrs: { shape: [batch, lastDim] }, backend: backend }); var kPow2 = roundUpToPow2(k); var lastDimPow2 = roundUpToPow2(lastDim); // Only the indices containing the top K are kept at every step to reduce // number of outputs in the GPU algorithms, so once the final set of indices // is computed then gather is used to grab the corresponding values // from the original input. var indices = null; // GPU algorithm always takes in an indices input but this input is not used // on the first run of a GPU algorithm, therefore if indices is null we simply // pass in x2D instead of it but the value will not actually be used var getInputs = function () { return indices === null ? [x2D, x2D] : [x2D, indices]; }; var runSwap = function (dir, inc, shape) { var inputs = getInputs(); var program = new SwapProgram(shape); var firstPass = indices === null ? 1 : 0; var uniformDataSwap = [ { type: 'int32', data: [lastDim] }, { type: 'int32', data: [firstPass] }, { type: 'float32', data: [Number.NEGATIVE_INFINITY] }, { type: 'int32', data: [dir] }, { type: 'int32', data: [inc] } ]; var prevIndices = indices; indices = backend.runWebGPUProgram(program, inputs, 'int32', uniformDataSwap); disposeIntermediateTensorInfoOrNull(backend, prevIndices); }; // Step 1: local sort for (var len = 1; len < kPow2; len *= 2) { var dir = len * 2; for (var inc = len; inc >= 1; inc /= 2) { runSwap(dir, inc, [batch, lastDimPow2]); } } // Step 2: merge for (var indicesSize = lastDimPow2; indicesSize > kPow2; indicesSize /= 2) { var inputs_1 = getInputs(); var mergeProgram = new MergeProgram([batch, indicesSize / 2]); var firstPass = indices === null ? 1 : 0; var uniformDataMerge = [ { type: 'int32', data: [lastDim] }, { type: 'int32', data: [firstPass] }, { type: 'int32', data: [kPow2] } ]; var prevIndices_1 = indices; indices = backend.runWebGPUProgram(mergeProgram, inputs_1, 'int32', uniformDataMerge); disposeIntermediateTensorInfoOrNull(backend, prevIndices_1); // Step 3: rebuild var len = kPow2 / 2; var dir = len * 2; for (var inc = len; inc >= 1; inc /= 2) { runSwap(dir, inc, indices.shape); } } // Keep only the requested top K results instead of kPow2 var prevIndices = indices; indices = slice({ inputs: { x: indices }, backend: backend, attrs: { begin: 0, size: [batch, k] } }); disposeIntermediateTensorInfoOrNull(backend, prevIndices); // Gather values on last dimension var values = gatherV2({ inputs: { x: x2D, indices: indices }, backend: backend, attrs: { axis: 1, batchDims: 1 } }); disposeIntermediateTensorInfoOrNull(backend, x2D); // Reshape back to the original input shape, except that the last // dimension is k. var newShape = xShape.slice(0, -1); newShape.push(k); prevIndices = indices; indices = reshape({ inputs: { x: indices }, attrs: { shape: newShape }, backend: backend }); disposeIntermediateTensorInfoOrNull(backend, prevIndices); var prevValues = values; values = reshape({ inputs: { x: values }, attrs: { shape: newShape }, backend: backend }); disposeIntermediateTensorInfoOrNull(backend, prevValues); return [values, indices]; } var topKConfig = { kernelName: tf.TopK, backendName: 'webgpu', kernelFunc: topK }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var TransformProgram = /** @class */ (function () { function TransformProgram(outShape) { this.variableNames = ['Image', 'Transforms']; this.uniforms = 'interpolationModeId : i32, fillModeId : i32, fillValue : f32,'; this.workgroupSize = [64, 1, 1]; this.size = true; this.outputShape = outShape; this.dispatchLayout = flatDispatchLayout(this.outputShape); this.dispatch = computeDispatch(this.dispatchLayout, this.outputShape, this.workgroupSize); this.shaderKey = 'transform'; } TransformProgram.prototype.getUserCode = function () { var userCode = "\n fn mapCoord(outCoord : f32, len : f32) -> f32{\n var inCoord = outCoord;\n if(uniforms.fillModeId == 2) {\n if (inCoord < 0.0) {\n if (len <= 1.0) {\n inCoord = 0.0;\n } else {\n let sz2 = 2.0 * len;\n if (inCoord < sz2) {\n inCoord = sz2 * f32(i32(f32(-inCoord / sz2))) +\n inCoord;\n }\n if (inCoord < -len) {\n inCoord = inCoord + sz2;\n } else {\n inCoord = -inCoord - 1.0;\n }\n }\n } else if (inCoord > len - 1.0) {\n if (len <= 1.0) {\n inCoord = 0.0;\n } else {\n let sz2 = 2.0 * len;\n inCoord = inCoord - sz2 * f32(i32(f32(inCoord / sz2)));\n if (inCoord >= len) {\n inCoord = sz2 - inCoord - 1.0;\n }\n }\n }\n return clamp(inCoord, 0.0, len - 1.0);\n } else if (uniforms.fillModeId == 3) {\n if (inCoord < 0.0) {\n if (len <= 1.0) {\n inCoord = 0.0;\n } else {\n let sz = len - 1.0;\n inCoord = inCoord + len * (f32(i32(f32(-inCoord / sz))) + 1.0);\n }\n } else if (inCoord > len - 1.0) {\n if (len <= 1.0) {\n inCoord = 0.0;\n } else {\n let sz = len - 1.0;\n inCoord = inCoord - len * f32(i32(f32(inCoord / sz)));\n }\n }\n return clamp(inCoord, 0.0, len - 1.0);\n } else if (uniforms.fillModeId == 4) {\n return clamp(outCoord, 0.0, len - 1.0);\n }\n return outCoord;\n }\n fn readWithFillValue(batch : i32, coordY : i32, coordX : i32,\n channel : i32) -> f32 {\n var outputValue : f32;\n if (0 <= coordY && coordY < uniforms.imageShape[1] && 0 <= coordX && coordX < uniforms.imageShape[2]) {\n outputValue = getImage(batch, coordY, coordX, channel);\n } else {\n outputValue = uniforms.fillValue;\n }\n return outputValue;\n }\n\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.size) {\n let coords = getCoordsFromIndex(index);\n var outputValue : f32;\n let batch = coords[0];\n let x = coords[2];\n let y = coords[1];\n let channel = coords[3];\n let xf = f32(x);\n let yf = f32(y);\n let a1 = getTransforms(batch, 0);\n let a2 = getTransforms(batch, 1);\n let a3 = getTransforms(batch, 2);\n let b1 = getTransforms(batch, 3);\n let b2 = getTransforms(batch, 4);\n let b3 = getTransforms(batch, 5);\n let c1 = getTransforms(batch, 6);\n let c2 = getTransforms(batch, 7);\n let projection = c1 * xf + c2 * yf + 1.0;\n if (projection == 0.0) {\n outputValue = uniforms.fillValue;\n } else {\n let inX = (a1 * xf + a2 * yf + a3) / projection;\n let inY = (b1 * xf + b2 * yf + b3) / projection;\n let mapX = mapCoord(inX, f32(uniforms.imageShape[2]));\n let mapY = mapCoord(inY, f32(uniforms.imageShape[1]));\n\n if (uniforms.interpolationModeId == 1) {\n let coordY = i32(round(mapY));\n let coordX = i32(round(mapX));\n outputValue = readWithFillValue(batch, coordY, coordX,\n channel);\n } else {\n let yFloor = floor(mapY);\n let xFloor = floor(mapX);\n let yCeil = yFloor + 1.0;\n let xCeil = xFloor + 1.0;\n let valueYFloor = (xCeil - mapX) *\n readWithFillValue(batch, i32(yFloor), i32(xFloor), channel) +\n (mapX - xFloor) *\n readWithFillValue(batch, i32(yFloor), i32(xCeil), channel);\n let valueYCeil = (xCeil - mapX) *\n readWithFillValue(batch, i32(yCeil), i32(xFloor), channel) +\n (mapX - xFloor) *\n readWithFillValue(batch, i32(yCeil), i32(xCeil), channel);\n outputValue = (yCeil - mapY) * valueYFloor +\n (mapY - yFloor) * valueYCeil;\n }\n }\n setOutputAtIndex(index, outputValue);\n }\n }\n "); return userCode; }; return TransformProgram; }()); function transform(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var image = inputs.image, transforms = inputs.transforms; var interpolation = attrs.interpolation, fillMode = attrs.fillMode, fillValue = attrs.fillValue, outputShape = attrs.outputShape; var _a = __read(image.shape, 4), batch = _a[0], imageHeight = _a[1], imageWidth = _a[2], numChannels = _a[3]; var _b = __read(outputShape != null ? outputShape : [imageHeight, imageWidth], 2), outHeight = _b[0], outWidth = _b[1]; var outShape = [batch, outHeight, outWidth, numChannels]; var program = new TransformProgram(outShape); var interpolationModeId = interpolation === 'nearest' ? 1 : 2; var fillModeId; switch (fillMode) { case 'constant': fillModeId = 1; break; case 'reflect': fillModeId = 2; break; case 'wrap': fillModeId = 3; break; case 'nearest': fillModeId = 4; break; default: fillModeId = 1; break; } var uniformData = [ { type: 'int32', data: [interpolationModeId] }, { type: 'int32', data: [fillModeId] }, { type: 'float32', data: [fillValue] } ]; return backend.runWebGPUProgram(program, [image, transforms], 'float32', uniformData); } var transformConfig = { kernelName: tf.Transform, backendName: 'webgpu', kernelFunc: transform }; /** * @license * Copyright 2021 Google LLC. All Rights Reserved. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function unpack(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var value = inputs.value; var axis = attrs.axis; if (axis < 0) { axis += value.shape.length; } var x = value; var xRank = x.shape.length; var num = value.shape[axis]; var outShape = new Array(xRank - 1); var outIndex = 0; for (var i = 0; i < xRank; i++) { if (i !== axis) { outShape[outIndex++] = x.shape[i]; } } var toDispose = []; var begin = new Array(xRank).fill(0); var size = x.shape.slice(); size[axis] = 1; var res = new Array(num); for (var i = 0; i < res.length; i++) { begin[axis] = i; var sliced = slice({ inputs: { x: x }, backend: backend, attrs: { begin: begin, size: size } }); var reshaped = reshape({ inputs: { x: sliced }, backend: backend, attrs: { shape: outShape } }); res[i] = reshaped; toDispose.push(sliced); } toDispose.forEach(function (t) { return backend.disposeData(t.dataId); }); return res; } var unpackConfig = { kernelName: tf.Unpack, backendName: 'webgpu', kernelFunc: unpack }; /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ var UnsortedSegmentSumProgram = /** @class */ (function () { function UnsortedSegmentSumProgram(inShape, outShape, outputDtype) { this.outputShape = []; this.variableNames = ['x', 'segmentIds']; this.uniforms = 'numSegments : i32, xSize: i32,'; this.workgroupSize = [64, 1, 1]; this.atomic = true; this.outputShape = outShape; this.dispatchLayout = flatDispatchLayout(inShape); this.dispatch = computeDispatch(this.dispatchLayout, inShape, this.workgroupSize); if (outputDtype !== 'float32' && outputDtype !== 'int32') { throw new Error("UnsortedSegmentSum only supports float32 and int32\n types, does not support ".concat(outputDtype, " type.")); } this.type = outputDtype; this.shaderKey = 'unsortedSegmentSum'; } UnsortedSegmentSumProgram.prototype.getUserCode = function () { var userCode = "\n ".concat(getMainHeaderString('index'), " {\n if (index < uniforms.xSize) {\n let coords = getXCoordsFromIndex(index);\n let b = coords[0];\n let inCol = coords[1];\n\n let segmentId = i32(getSegmentIds(inCol));\n if (segmentId >= 0) {\n let flatIndex = b * uniforms.numSegments + segmentId % uniforms.numSegments;\n let value = getX(b, inCol);\n\n ").concat(atomicAddSnippet('&result[flatIndex]', 'value', this.type), "\n }\n }\n }\n "); return userCode; }; return UnsortedSegmentSumProgram; }()); /** * @license * Copyright 2023 Google LLC. * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. * ============================================================================= */ function unsortedSegmentSum(args) { var inputs = args.inputs, backend = args.backend, attrs = args.attrs; var x = inputs.x, segmentIds = inputs.segmentIds; var numSegments = attrs.numSegments; var xRank = x.shape.length; var toDispose = []; var axis = 0; var permutation = tf.backend_util.getAxesPermutation([axis], xRank); var permutedX = x; if (permutation != null) { permutedX = transpose({ inputs: { x: x }, backend: backend, attrs: { perm: permutation } }); toDispose.push(permutedX); axis = tf.backend_util.getInnerMostAxes(1, xRank)[0]; } var outShape = tf.backend_util.segment_util.computeOutShape(permutedX.shape, axis, numSegments); var inSize = tf.util.sizeFromShape([permutedX.shape[axis]]); var a2D = reshape({ inputs: { x: permutedX }, backend: backend, attrs: { shape: [-1, inSize] } }); toDispose.push(a2D); var dtype = x.dtype; var shape = [a2D.shape[0], numSegments]; var output = fill({ backend: backend, attrs: { shape: shape, value: 0, dtype: dtype } }); var program = new UnsortedSegmentSumProgram(a2D.shape, shape, dtype); var uniformData = [ { type: 'int32', data: [numSegments] }, { type: 'int32', data: [tf.util.sizeFromShape(a2D.shape)] } ]; var segResult = backend.runWebGPUProgram(program, [a2D, segmentIds], dtype, uniformData, output); var reshaped = reshape({ inputs: { x: segResult }, backend: backend, attrs: { shape: outShape } }); toDispose.push(segResult); var result = reshaped; if (permutation != null) { toDispose.push(reshaped); var perm = tf.backend_util.getUndoAxesPermutation(permutation); result = transpose({ inputs: { x: result }, backend: backend, attrs: { perm: perm } }); } toDispose.forEach(function (t) { return backend.disposeData(t.dataId); }); return result; } var unsortedSegmentSumConfig = { kernelName: tf.UnsortedSegmentSum, backendName: 'webgpu', kernelFunc: unsortedSegmentSum }; var e_1, _a; // List all kernel configs here var kernelConfigs = [ _fusedMatMulConfig, absConfig, acosConfig, acoshConfig, addConfig, addNConfig, allConfig, anyConfig, argMaxConfig, argMinConfig, asinConfig, asinhConfig, atanConfig, atan2Config, atanhConfig, avgPoolConfig, avgPool3DConfig, avgPool3DGradConfig, avgPoolGradConfig, batchMatMulConfig, batchToSpaceNDConfig, bincountConfig, broadcastArgsConfig, castConfig, ceilConfig, clipByValueConfig, complexConfig, complexAbsConfig, concatConfig, conv2DConfig, conv2DBackpropFilterConfig, conv2DBackpropInputConfig, conv3DConfig, conv3DBackpropFilterV2Config, conv3DBackpropInputV2Config, cosConfig, coshConfig, cropAndResizeConfig, cumprodConfig, cumsumConfig, denseBincountConfig, depthToSpaceConfig, depthwiseConv2dNativeBackpropFilterConfig, depthwiseConv2dNativeBackpropInputConfig, depthwiseConv2dNativeConfig, diagConfig, dilation2DConfig, dilation2DBackpropFilterConfig, dilation2DBackpropInputConfig, drawConfig, einsumConfig, eluConfig, eluGradConfig, equalConfig, erfConfig, expConfig, expandDimsConfig, expm1Config, fftConfig, fillConfig, flipLeftRightConfig, fromPixelsConfig, floorConfig, floorDivConfig, fusedBatchNormConfig, fusedConv2DConfig, fusedDepthwiseConv2DConfig, gatherNdConfig, gatherV2Config, greaterConfig, greaterEqualConfig, identityConfig, ifftConfig, imagConfig, isFiniteConfig, isInfConfig, isNaNConfig, leakyReluConfig, lessConfig, lessEqualConfig, linSpaceConfig, log1pConfig, logConfig, logicalAndConfig, logicalNotConfig, logicalOrConfig, lrnConfig, lrnGradConfig, maxConfig, maximumConfig, maxPoolConfig, maxPoolGradConfig, maxPool3DConfig, maxPool3DGradConfig, maxPoolWithArgmaxConfig, meanConfig, minConfig, minimumConfig, mirrorPadConfig, modConfig, multinomialConfig, multiplyConfig, negConfig, nonMaxSuppressionV3Config, nonMaxSuppressionV5Config, notEqualConfig, oneHotConfig, onesLikeConfig, packConfig, padV2Config, powConfig, preluConfig, prodConfig, rangeConfig, realConfig, realDivConfig, reciprocalConfig, reluConfig, relu6Config, reshapeConfig, resizeBilinearConfig, resizeBilinearGradConfig, resizeNearestNeighborConfig, resizeNearestNeighborGradConfig, reverseConfig, rotateWithOffsetConfig, roundConfig, rsqrtConfig, scatterNdConfig, searchSortedConfig, selectConfig, seluConfig, sigmoidConfig, signConfig, sinConfig, sinhConfig, sliceConfig, stepConfig, stridedSliceConfig, stringNGramsConfig, softmaxConfig, softplusConfig, spaceToBatchNDConfig, sparseSegmentMeanConfig, sparseSegmentSumConfig, sparseToDenseConfig, splitVConfig, sqrtConfig, squareConfig, squaredDifferenceConfig, subConfig, sumConfig, tanConfig, tanhConfig, tensorScatterUpdateConfig, tileConfig, topKConfig, transformConfig, transposeConfig, unpackConfig, unsortedSegmentSumConfig, zerosLikeConfig ]; try { for (var kernelConfigs_1 = __values(kernelConfigs), kernelConfigs_1_1 = kernelConfigs_1.next(); !kernelConfigs_1_1.done; kernelConfigs_1_1 = kernelConfigs_1.next()) { var kernelConfig = kernelConfigs_1_1.value; tf.registerKernel(kernelConfig); } } catch (e_1_1) { e_1 = { error: e_1_1 }; } finally { try { if (kernelConfigs_1_1 && !kernelConfigs_1_1.done && (_a = kernelConfigs_1.return)) _a.call(kernelConfigs_1); } finally { if (e_1) throw e_1.error; } } exports.WebGPUBackend = WebGPUBackend; exports.webgpu_util = webgpu_util; })); //# sourceMappingURL=tf-backend-webgpu.js.map