brainpy / BrainPy

Brain Dynamics Programming in Python
https://brainpy.readthedocs.io/
GNU General Public License v3.0
514 stars 93 forks source link

[math] Add new customize operators with `cupy` #653

Closed Routhleck closed 6 months ago

Routhleck commented 6 months ago

Although we can now use the flexible taichi custom operator approach, taichi on cuda does not have more fine-grained control or optimization for some scenarios. So for such scenarios, we can use cupy's RawModule(https://docs.cupy.dev/en/stable/user_guide/kernel.html#raw-kernels) or jit.rawkernel(https://docs.cupy.dev/en/stable/user_guide/kernel.html#jit-kernel-definition) to compile and run CUDA native code directly as strings in real time for finer grained control.

Routhleck commented 6 months ago

https://github.com/cupy/cupy/issues/8232

chaoming0625 commented 6 months ago

Before merging this PR, I will give some modifications.