Support Jitify ?

This is more for my own little curiosity than a feature request, but I thought it’s worth discussing.

It is not possible to compile Thrust/CUB code in NVRTC because many system headers are missing, which was why @emcastillo struggled in #2629 and we still have to keep a copy of our own Thrust complex headers. Jitify claims to provide an easy solution to address this.

I have been looking into jitify’s implementation lately, and I think the key idea is to:

1. Recursively find out what header (and its absolute path) should be included (by repeating compiling -> inspecting error log)
2. Fix the headers: for example, provide some stubs to patch a few C++ std functions or system headers so that they can work as __device__ functions, remove static variables, fix pragma, etc

Since we have most of the toolchain in place (compile/launch/CUDA API/etc), especially #3319 is now merged, I think it is not super difficult to address this. I can think of two design choices:

1. Wrap jitify in cupy.cuda.jitify: Ideally, we still use our NVRTC wrapper to do compilation, and only take all these lookup implementations as-is so we don’t reinvent the wheel. Unfortunately it is currently not possible, as many logics are baked in jitify’s public APIs, and it requires some refactoring on the their side for us to reuse.
2. Implement the compile-error search loop by ourselves in cupy/cuda/compiler.py: This would require kinda reimplementing what’s done in jitify, but only for a small part (I hope). The pro is jitify doesn’t need to be a dependency and we just reuse what’s already available (very important for code maintenance), and the cons are that it takes more effort and that it needs to copy a significant amount of code from jitify (mainly those code patches).

I think a few RAPIDS libraries like cuDF use jitify. I will dig into how they use it…

[maddyscientist]

@harrism just to note that https://github.com/NVIDIA/jitify/pull/39 has been merged.

[maddyscientist]

@leofang if you’re just wanting some of the headers provided by Jitify, you might be better looking at using libcu++'s headers: https://github.com/NVIDIA/libcudacxx/tree/main/include/cuda/std.

These are more fully featured, and are being continually expanded to provide more coverage of the the C++ standard library, as well as some C headers.