Explicitly look for program name to allow reusing program_sources

[leofang]

Hello it's me again. My PR to support Jitify in CuPy (cupy/cupy#4228) just went in. But, when working on that I notice suboptimal performance in Jitify's compile-fail-search loop (implemented in jitify::detail::load_program, which is the only Jitify function we use, as we already did the rest of Jitify functionalities ourselves). The reason, from our perspective, is that the map program_sources cannot be reused, as the current implementation assumes that the program name is looked up right after the first item is inserted to the map, and if the map is already populated the lookup would fetch a wrong name.

To work around this limitation, we currently cache all headers searched by Jitify, and pass them to load_program's 2nd argument (headers, a vector of string). But the drawback is a new program_sources map still needs to be constructed every time the search loop is launched. The compilation time is just slower than using a subprocess to call nvcc (what we did before introducing Jitify): cupy/cupy#4228 (comment).

The program-name lookup limitation is lifted in this PR by explicitly passing in a string pointer. If it's not null then we record the name, otherwise do nothing. This is an internal (and small) change that would not have any impact to Jitify's existing functionality at all, but is very important for us to reuse the existing resource to the maximal extent. I hope this is an easy PR to review. 🙂

Thank you!

[benbarsdell]

Thanks for the PR! This is an interesting use-case; I'm glad it's working for you with only this small change required.

Just a couple of formatting issues, otherwise LGTM.

[leofang]

Thanks a lot for your quick review, @benbarsdell! The requested changes are made.

btw for formatting, do you follow any style guide or use any linter? In Python we have the PEP-8 guide and tools like flake8 and autopep8, but I know literally nothing in the C++ world 😅

[benbarsdell]

I just noticed that this causes a couple of test failures (when running make test). I think the problem is that when the program is specified as a filename instead of a code string it never sets *program_name.
I think you can fix it by moving these lines to after the if statement here (as long as this doesn't break your use-case).

Regarding auto-formatting, we use clang-format (the format is defined in the .clang-format file).

[leofang]

I just noticed that this causes a couple of test failures (when running make test). I think the problem is that when the program is specified as a filename instead of a code string it never sets *program_name.
I think you can fix it by moving these lines to after the if statement here (as long as this doesn't break your use-case).

Interesting, thanks for double checking for me, @benbarsdell! I should have run make test myself...Though I've run our tests many times and never hit this bug. I guess it's because in our case when the program is a header filename, program_name is always null. I fixed it in 1391b7b.

Before applying your suggestion, I had 4 failing tests:

[  FAILED  ] JitifyTest.ConstantMemory
[  FAILED  ] JitifyTest.ConstantMemory_experimental
[  FAILED  ] JitifyTest.ThrustHeaders
[  FAILED  ] JitifyTest.ClassKernelArg

but now I still have 1 test failing (JitifyTest.ThrustHeaders). Could it be related to the issue that I am experiencing in #81? Could you check? Below is the error log (after applying the fix):

[ RUN      ] JitifyTest.ThrustHeaders
---------------------------------------
--- Source of thrust_program ---
---------------------------------------
  1 #include <thrust/iterator/counting_iterator.h>
  2 __global__ void my_kernel(thrust::counting_iterator<int> begin,
  3                           thrust::counting_iterator<int> end) {
  4 }
---------------------------------------
Found #include iterator from /usr/include/thrust/iterator/iterator_categories.h:40 [] at:
  __jitify_builtin/iterator
Found #include limits from /usr/include/thrust/detail/numeric_traits.h:20 [] at:
  __jitify_builtin/limits
Found #include cfloat from limits:4 [__jitify_builtin/limits] at:
  __jitify_builtin/cfloat
Found #include climits from limits:5 [__jitify_builtin/limits] at:
  __jitify_builtin/climits
Found #include cstdint from limits:6 [__jitify_builtin/limits] at:
  __jitify_builtin/cstdint
Found #include cstddef from /usr/include/thrust/iterator/detail/counting_iterator.inl:23 [] at:
  __jitify_builtin/cstddef
---------------------------------------------------
--- JIT compile log for thrust_program ---
---------------------------------------------------
/usr/include/thrust/iterator/iterator_facade.h(251): error: namespace "std" has no member "ptrdiff_t"

/usr/include/thrust/iterator/iterator_traits.h(56): error: namespace "std" has no member "ptrdiff_t"

/usr/include/thrust/iterator/iterator_traits.h(66): error: namespace "std" has no member "ptrdiff_t"

/usr/include/thrust/iterator/iterator_adaptor.h(147): error: "#" not expected here

/usr/include/thrust/iterator/iterator_adaptor.h(188): error: "#" not expected here

/usr/include/thrust/iterator/iterator_adaptor.h(193): error: "#" not expected here

/usr/include/thrust/iterator/iterator_adaptor.h(199): error: "#" not expected here

/usr/include/thrust/iterator/iterator_adaptor.h(207): error: "#" not expected here

/usr/include/thrust/iterator/iterator_adaptor.h(212): error: "#" not expected here

/usr/include/thrust/iterator/iterator_adaptor.h(220): error: "#" not expected here

/usr/include/thrust/iterator/iterator_adaptor.h(147): error: expected a declaration

/usr/include/thrust/iterator/iterator_adaptor.h(188): error: expected a declaration

12 errors detected in the compilation of "thrust_program".

---------------------------------------------------
unknown file: Failure
C++ exception with description "Runtime compilation failed" thrown in the test body.
[  FAILED  ] JitifyTest.ThrustHeaders (389 ms)

Regarding auto-formatting, we use clang-format (the format is defined in the .clang-format file).

Thanks!

[benbarsdell]

Aha, I think I know why you're having problems now: because your Thrust library is installed in /usr/include/thrust. There is a known issue in NVRTC < 11.0 where /usr/include is automatically added as an include path, and this circumvents Jitify's header loading functionality.
This should be fixed if you use CUDA >= 11.0 or if you move Thrust to somewhere else (or simply remove that directory and use the copy that comes in the CUDA Toolkit in /usr/local/cuda/include/thrust).

edit: fixed the directory name.

[leofang]

Oh cool. It makes sense. Glad to know. Unfortunately I don't have root access to my test machines, so I may have to spin off a docker and build from there. But it sounds like specific to my environment, so if all of the tests pass for you locally, I suppose all is good?

[leofang]

Thanks!! 🙏