Can the CUDA code ported into JAX code with competitive performance

sunset1995 / DirectVoxGO

Direct voxel grid optimization for fast radiance field reconstruction.

https://sunset1995.github.io/dvgo

Other

1.05k stars 110 forks source link

Can the CUDA code ported into JAX code with competitive performance #35

Open riven314 opened 2 years ago

riven314 commented 2 years ago

This work looks promising! I am thinking if it's feasible to port the CUDA code into JAX and achieve similar speed. Do you think it's feasible to achieve this?

sunset1995 commented 2 years ago

I'm less experience with JAX. The custom cuda code is necessary to achieve the speed. But I'm not sure whether it will be difficult to adapt my custom cuda code to JAX or just call them from JAX.