Thanks; I have another, currently private, repo that targets both pure CPU inference and wgpu (to vulkan.) It's somewhat similar but... different. Shares some common pieces but needs to be refactored to share more.
But it's been hard to make it competitive with CUDA. At least on this Spark and my only non-NVIDIA machine (which only has 16GB unified relatively slow RAM.)
I know a little bit about this problem space from previous work (we were working on performance-portable deep learning back around 2016). The infrastructure has improved but as far as I can tell not many teams have really "squeezed the toothpaste tube" and worked through performance issues systematically. These days a small team and robots can probably do it though.