tensorflow / tfjs

A WebGL accelerated JavaScript library for training and deploying ML models.
https://js.tensorflow.org
Apache License 2.0
18.25k stars 1.92k forks source link

WebGPU Compute shader execution? #8314

Open DennisSmolek opened 6 days ago

DennisSmolek commented 6 days ago

While looking into the WebGPU backend and execution example I am left with a few questions.

I am currently working on porting the Open Image Denoise models to work on Tensorflow.js. This has been done by someone already but they aren't up to sharing yet. This has also been done with Cuda Compute and some Rust/HLSL compute shaders.

The pipeline currently would be:

  1. Go from CPU to GPU Execute main rendering pipeline with compute, vertex, fragment shaders to get noisy image.
  2. Push image onto shared tensor storage buffer as input (GPU).
  3. Return to CPU
  4. Execute Tensorflow.js (with buffer that is still on the GPU)
  5. Tensorflow runs UFilter denoise on buffer and stores the new image data in the storage buffer (still on GPU)
  6. Return to CPU
  7. Issue fullscreen quad renderpass using the image on the storage buffer (GPU)
  8. Return to CPU

With native libraries you can execute OIDN directly in compute shaders (but is a total pain to setup) and other examples (CUDA/HLSL) also execute the DNN on the compute shaders without the CPU return.

I am curious to see if there is any existing methods to reduce the round trips between the CPU/GPU. Even simply executing the tensorflow process from a compute shader would be massive as the only thing returned to the CPU would be the final buffer...

I don't have the skills nor understanding yet to make this work, but is something I figured I would ask.