Function parallelCompilationCommon

tfjs-backend-webgpu/src/backend_webgpu_test.ts:353–383

(webGPUBackend: WebGPUBackend)

});

async function parallelCompilationCommon(webGPUBackend: WebGPUBackend) {
  const startNumBytes = (tf.memory() as WebGPUMemoryInfo).numBytesInGPU;
  const startTensor = tf.memory().numTensors;
  const startDataBuckets = webGPUBackend.numDataIds();

  const a1 = tf.tensor1d([1, 1, 1]);
  const b1 = tf.tensor1d([1, 1, 1]);

  // Parallel compile.
  tf.env().set('WEBGPU_ENGINE_COMPILE_ONLY', true);
  const c1 = tf.add(a1, b1);
  await webGPUBackend.checkCompileCompletionAsync();

  // Actual inference.
  tf.env().set('WEBGPU_ENGINE_COMPILE_ONLY', false);
  const c2 = tf.add(a1, b1);
  expectArraysEqual(await c2.data(), [2, 2, 2]);

  tf.dispose([a1, b1, c1, c2]);
  const endNumBytes = (tf.memory() as WebGPUMemoryInfo).numBytesInGPU;
  const endTensor = tf.memory().numTensors;
  const endDataBuckets = webGPUBackend.numDataIds();

  // We only check numBytesInGPU. For parallel compilation,
  // numBytesInGPUAllocated will be more because of the two pass
  // uploadToGPU, but they will all be freed, resulting in endNumbytes equal
  // to startNumBytes.
  expect(startNumBytes).toEqual(endNumBytes);
  expect(startTensor).toEqual(endTensor);
  expect(endDataBuckets).toEqual(startDataBuckets);
}

describeWebGPU('parallel compilation', () => {
  let prevBackend: string;