Invoke tfjs operations functions with GPU resource

[BruceDai]

Hi. I have this scenario of invoking tfjs operations functions for processing GPU resource, for example, invoking tf.add(a, b) with gpuBufferA and gpuBufferB by WebGPU backend.

Current my workaround likes:

// Get TypedArray f32ArrayA
const readbackBufferA = gpuDevice.createBuffer({size: sizeInBytes, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST});
const readbackEncoder = gpuDevice.createCommandEncoder();
readbackEncoder.copyBufferToBuffer(gpuBufferA, 0, readbackBufferA, 0, sizeInBytes);
device.queue.submit([readbackEncoder.finish()]);
readbackBufferA.mapAsync(GPUMapMode.READ);
const f32ArrayA = new Float32Array(readbackBufferA.getMappedRange());

// Same to get f32ArrayB
...
const f32ArrayB = new Float32Array(readbackBufferB.getMappedRange());

// Invoke tf.add()
const result = tf.add(f32ArrayA, f32ArrayB);

I don't think it's good on memory copy / performance.

Is there any efficient solution for this scenario? Any suggestion? Thanks.

[huningxin]

This would readback the data from GPU to CPU. If tf.js uses webgpu backend, it would ideal to let the data stay on GPU.

@pyu10055 @qjia7 , is there a method to create a tf.Tensor from an existing WebGPU buffer? tf.tensor only supports CPU buffer (TypedArrayBuffer).

[qjia7]

@huningxin Currently, there is no such kind of interface to directly import an external GPU resource to TFJS. I just offline talked with @lina128. We think it's helpful to have a such kind of API, like tf.tensor (GPUData) to take GPU resource as the input to create a tensor. Thanks.