Repository: cupy/cupy

Issue:
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](https://github.com/NVIDIA/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...


Discussion:
cub complex headers are pretty much stable, as they are just a struct and a couple of functions to manipulate the members that are very unlikely to change. I am not sure if doing all this is better than just keep the headers the way they are.

Also most likely, this is going to increase the cubin generated files size. Not sure if it's affordable
> cub complex headers are pretty much stable, as they are just a struct and a couple of functions to manipulate the members that are very unlikely to change. I am not sure if doing all this is better than just keep the headers the way they are.

No, I don't mean to remove CuPy's complex headers. The jitify solution (either the native one or the proposed designs) would impose a lot of overhead for compilation (according to [a public talk](https://developer.download.nvidia.com/video/gputechconf/gtc/2019/presentation/s91043-rapids-cuda-dataframe-internals-for-c++-developers.pdf) it's ~300 ms), so for complex numbers we better not to do this and stay with the status quo (I agree it's quite stable and good). 

The use case I have in mind is custom C++ `RawKernel`s that uses Thrust, CUB (related: #3244), or other CUDA C++ libraries.

> Also most likely, this is going to increase the cubin generated files size. Not sure if it's affordable

I don't understand why cubin size would increase. Even the native jitify uses NVRTC under the hood. 
Sorry, I just replied without completely understanding the issue.
I agree that for CUB this is good.
Regarding the cubin, nvmd that😅.
Ah, actually it's a legit concern and your intuition was right, @emcastillo. A jitify-generated kernel could have a larger size due to the addition of extra things (which might not be used). They recently fixed this: https://github.com/NVIDIA/jitify/pull/40.
Yeah, I thought that it was possible for all the declarations in the templates to bloat the ptx files.
Glad to see that our Nvidia friends are smart and solved this :)
cc @devavret @harrism and @jrhemstad who are the most familiar with Jitify usage in cuDF if they can answer any questions for you 😄
libcudf does not use jitify with cub/thrust if that’s the primary goal. We have some basic kernels and jitify was used as a convenience wrapper around NVRTC.
Thanks a lot for bringing the party here, @kkraus14 😂 Now I know who to ask when I revisit this issue!
There is interest in compiling Thrust / CUB with jitify. Here's a PR that takes steps toward that: https://github.com/NVIDIA/jitify/pull/39  Not sure what the status of that PR is now.
> The use case I have in mind is custom C++ `RawKernel`s that uses Thrust, CUB (related: #3244), or other CUDA C++ libraries.

I knew there must be someone thinking of the same thing as I do! See #3728.


@harrism just to note that https://github.com/NVIDIA/jitify/pull/39 has been merged.
Hi @maddyscientist, just out of curiosity to evaluate the level of challenge, do I understand it correctly that once C/C++ header includes are mapped correctly (ex: `#include <limits.h>` -> `jitsafe_header_limits_h`) NVRTC will be able to compile it correctly? We (CuPy) already have NVRTC and the entire JIT compilation workflow in place, so I am hoping this is the last missing piece to achieve what jitify can provide.
Not sure I follow this question @leofang.  Jitify takes care of this mapping automatically, so that the code the user provides just has to `#include <limits.h>` are normal in the user kernel source code.
Thanks @maddyscientist, I meant to do this mapping under the hood for CuPy users. So suppose I have a piece of "illegal" code like below:
```python
my_bad_code = r'''
#include <limits.h>
#include <thrust/...>

// If we just use NVRTC to compile this snippet, we get errors
extern "C" __global__ void some_kernel_cannot_be_compiled_natively_by_nvrtc(...) { 
    /* ... */
}
'''
```
but I'd still like to pass this to `cupy.RawKernel`/`cupy.RawModule` and use NVRTC to compile it:
```python
mod = cupy.RawModule(code=my_bad_code, backend='nvrtc')
ker = mod.get_function('some_kernel_cannot_be_compiled_natively_by_nvrtc')
ker(grid, block, ...)  # launch it
```
My understanding for a working prototype is that behind the scene we should do something equivalent of this for users:
```python
if '#include <limits.h>' in my_bad_code: 
    # jitsafe_header_limits_h is a string read from jitify.hpp
    my_bad_code = my_bad_code.replace('#include <limits.h>', jitsafe_header_limits_h)
elif '...' in my_bad_code:
    ...
```
then pass the processed `my_bad_code` string to NVRTC and compile it. Do I miss or misunderstand anything here? (I hope the Python part is not too confusing 🙂)
@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.

Resolving patch:
diff --git a/cupy/cuda/compiler.py b/cupy/cuda/compiler.py
--- a/cupy/cuda/compiler.py
+++ b/cupy/cuda/compiler.py
@@ -9,10 +9,14 @@
 
 from cupy.cuda import device
 from cupy.cuda import function
+from cupy_backends.cuda.api import driver
 from cupy_backends.cuda.api import runtime
 from cupy_backends.cuda.libs import nvrtc
 from cupy import _util
 
+if not runtime.is_hip and driver.get_build_version() > 0:
+    from cupy.cuda.jitify import jitify
+
 
 _nvrtc_version = None
 _win32 = sys.platform.startswith('win32')
@@ -29,6 +33,10 @@ class HIPCCException(Exception):
     pass
 
 
+class JitifyException(Exception):
+    pass
+
+
 def _run_cc(cmd, cwd, backend, log_stream=None):
     # backend in ('nvcc', 'hipcc')
     try:
@@ -142,16 +150,55 @@ def _get_bool_env_variable(name, default):
         return False
 
 
+_jitify_header_source_map = None
+
+
+def _jitify_prep(source, options, cu_path):
+    # TODO(leofang): refactor this?
+    global _jitify_header_source_map
+    if _jitify_header_source_map is None:
+        from cupy.core import core
+        _jitify_header_source_map = core._get_header_source_map()
+
+    # jitify requires the 1st line to be the program name
+    source = cu_path + '\n' + source
+
+    # Upon failure, in addition to throw an error Jitify also prints the log
+    # to stdout. In principle we could intercept that by hijacking stdout's
+    # file descriptor (tested locally), but the problem is pytest also does
+    # the same thing internally, causing strange errors when running the tests.
+    # As a result, we currently maintain Jitify's default behavior for easy
+    # debugging, and wait for the upstream to address this issue
+    # (NVIDIA/jitify#79).
+
+    try:
+        name, options, headers, include_names, hdr_map = jitify(
+            source, options, _jitify_header_source_map)
+    except Exception as e:  # C++ could throw all kinds of errors
+        raise JitifyException(str(e))
+    assert name == cu_path
+    _jitify_header_source_map = hdr_map
+
+    return options, headers, include_names
+
+
 def compile_using_nvrtc(source, options=(), arch=None, filename='kern.cu',
                         name_expressions=None, log_stream=None,
-                        cache_in_memory=False):
+                        cache_in_memory=False, jitify=False):
     if not arch:
         arch = _get_arch()
 
     options += ('-arch=compute_{}'.format(arch),)
 
-    def _compile(source, options, cu_path, name_expressions, log_stream):
-        prog = _NVRTCProgram(source, cu_path,
+    def _compile(
+            source, options, cu_path, name_expressions, log_stream, jitify):
+        if jitify:
+            options, headers, include_names = _jitify_prep(
+                source, options, cu_path)
+        else:
+            headers = include_names = ()
+
+        prog = _NVRTCProgram(source, cu_path, headers, include_names,
                              name_expressions=name_expressions)
         try:
             ptx, mapping = prog.compile(options, log_stream)
@@ -171,9 +218,11 @@ def _compile(source, options, cu_path, name_expressions, log_stream):
                 cu_file.write(source)
 
             return _compile(source, options, cu_path,
-                            name_expressions, log_stream)
+                            name_expressions, log_stream, jitify)
     else:
-        return _compile(source, options, '', name_expressions, log_stream)
+        cu_path = '' if not jitify else filename
+        return _compile(source, options, cu_path, name_expressions,
+                        log_stream, jitify)
 
 
 def compile_using_nvcc(source, options=(), arch=None,
@@ -307,7 +356,7 @@ def get_cache_dir():
 def compile_with_cache(
         source, options=(), arch=None, cache_dir=None, extra_source=None,
         backend='nvrtc', *, enable_cooperative_groups=False,
-        name_expressions=None, log_stream=None):
+        name_expressions=None, log_stream=None, jitify=False):
 
     if enable_cooperative_groups:
         if backend != 'nvcc':
@@ -335,13 +384,13 @@ def compile_with_cache(
         return _compile_with_cache_cuda(
             source, options, arch, cache_dir, extra_source, backend,
             enable_cooperative_groups, name_expressions, log_stream,
-            cache_in_memory)
+            cache_in_memory, jitify)
 
 
 def _compile_with_cache_cuda(
         source, options, arch, cache_dir, extra_source=None, backend='nvrtc',
         enable_cooperative_groups=False, name_expressions=None,
-        log_stream=None, cache_in_memory=False):
+        log_stream=None, cache_in_memory=False, jitify=False):
     # NVRTC does not use extra_source. extra_source is used for cache key.
     global _empty_file_preprocess_cache
     if cache_dir is None:
@@ -360,6 +409,18 @@ def _compile_with_cache_cuda(
     if _get_bool_env_variable('CUPY_CUDA_COMPILE_WITH_DEBUG', False):
         options += ('--device-debug', '--generate-line-info')
 
+    is_jitify_requested = ('-DCUPY_USE_JITIFY' in options)
+    if jitify and not is_jitify_requested:
+        # jitify is set in RawKernel/RawModule, translate it to an option
+        # that is useless to the compiler, but can be used as part of the
+        # hash key
+        options += ('-DCUPY_USE_JITIFY',)
+    elif is_jitify_requested and not jitify:
+        # jitify is requested internally, just set the flag
+        jitify = True
+    if jitify and backend != 'nvrtc':
+        raise ValueError('jitify only works with NVRTC')
+
     env = (arch, options, _get_nvrtc_version(), backend)
     base = _empty_file_preprocess_cache.get(env, None)
     if base is None:
@@ -401,7 +462,7 @@ def _compile_with_cache_cuda(
         cu_name = '' if cache_in_memory else name + '.cu'
         ptx, mapping = compile_using_nvrtc(
             source, options, arch, cu_name, name_expressions,
-            log_stream, cache_in_memory)
+            log_stream, cache_in_memory, jitify)
         if _is_cudadevrt_needed(options):
             # for separate compilation
             ls = function.LinkState()
diff --git a/cupy_setup_build.py b/cupy_setup_build.py
--- a/cupy_setup_build.py
+++ b/cupy_setup_build.py
@@ -238,6 +238,23 @@
         'check_method': build.check_cub_version,
         'version_method': build.get_cub_version,
     })
+
+    MODULES.append({
+        'name': 'jitify',
+        'file': [
+            'cupy.cuda.jitify',
+        ],
+        'include': [
+            'cuda.h',
+            'cuda_runtime.h',
+            'nvrtc.h',
+        ],
+        'libraries': [
+            'cuda',
+            'cudart',
+            'nvrtc',
+        ],
+    })
 else:
     MODULES.append({
         'name': 'cub',
@@ -558,8 +575,6 @@ def make_extensions(options, compiler, use_cython):
         link_args = s.setdefault('extra_link_args', [])
 
         if module['name'] == 'cusolver':
-            compile_args = s.setdefault('extra_compile_args', [])
-            link_args = s.setdefault('extra_link_args', [])
             # openmp is required for cusolver
             if use_hip:
                 pass
@@ -570,6 +585,10 @@ def make_extensions(options, compiler, use_cython):
             elif compiler.compiler_type == 'msvc':
                 compile_args.append('/openmp')
 
+        # this fixes RTD (no_cuda) builds...
+        if module['name'] == 'jitify':
+            compile_args.append('--std=c++11')
+
         original_s = s
         for f in module['file']:
             s = copy.deepcopy(original_s)
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -8,17 +8,19 @@
 import cupy_setup_build
 
 
-if len(os.listdir('cupy/core/include/cupy/cub/')) == 0:
-    msg = '''
-    The folder cupy/core/include/cupy/cub/ is a git submodule but is
-    currently empty. Please use the command
-
-        git submodule update --init
-
-    to populate the folder before building from source.
-    '''
-    print(msg, file=sys.stderr)
-    sys.exit(1)
+for submodule in ('cupy/core/include/cupy/cub/',
+                  'cupy/core/include/cupy/jitify'):
+    if len(os.listdir(submodule)) == 0:
+        msg = '''
+        The folder %s is a git submodule but is
+        currently empty. Please use the command
+
+            git submodule update --init
+
+        to populate the folder before building from source.
+        ''' % submodule
+        print(msg, file=sys.stderr)
+        sys.exit(1)
 
 
 if sys.version_info[:3] == (3, 5, 0):


Test patch:
diff --git a/tests/cupy_tests/core_tests/test_raw.py b/tests/cupy_tests/core_tests/test_raw.py
--- a/tests/cupy_tests/core_tests/test_raw.py
+++ b/tests/cupy_tests/core_tests/test_raw.py
@@ -380,11 +380,20 @@ def new_target(name, default):
 
 
 @testing.parameterize(
+    # First test NVRTC
     {'backend': 'nvrtc', 'in_memory': False},
     # this run will read from in-memory cache
     {'backend': 'nvrtc', 'in_memory': True},
     # this run will force recompilation
     {'backend': 'nvrtc', 'in_memory': True, 'clean_up': True},
+    # Below is the same set of NVRTC tests, with Jitify turned on. For tests
+    # that can already pass, it shouldn't matter whether Jitify is on or not,
+    # and the only side effect is to add overhead. It doesn't make sense to
+    # test NVCC + Jitify.
+    {'backend': 'nvrtc', 'in_memory': False, 'jitify': True},
+    {'backend': 'nvrtc', 'in_memory': True, 'jitify': True},
+    {'backend': 'nvrtc', 'in_memory': True, 'clean_up': True, 'jitify': True},
+    # Finally, we test NVCC
     {'backend': 'nvcc', 'in_memory': False},
 )
 class TestRaw(unittest.TestCase):
@@ -394,6 +403,10 @@ def setUp(self):
             _util.clear_memo()
         self.dev = cupy.cuda.runtime.getDevice()
         assert self.dev != 1
+        if not hasattr(self, 'jitify'):
+            self.jitify = False
+        if cupy.cuda.runtime.is_hip and self.jitify:
+            self.skipTest('Jitify does not support ROCm/HIP')
 
         self.temporary_cache_dir_context = use_temporary_cache_dir()
         self.in_memory_context = compile_in_memory(self.in_memory)
@@ -402,14 +415,14 @@ def setUp(self):
 
         self.kern = cupy.RawKernel(
             _test_source1, 'test_sum',
-            backend=self.backend)
+            backend=self.backend, jitify=self.jitify)
         self.mod2 = cupy.RawModule(
             code=_test_source2,
-            backend=self.backend)
+            backend=self.backend, jitify=self.jitify)
         self.mod3 = cupy.RawModule(
             code=_test_source3,
             options=('-DPRECISION=2',),
-            backend=self.backend)
+            backend=self.backend, jitify=self.jitify)
 
     def tearDown(self):
         if (self.in_memory
@@ -486,12 +499,20 @@ def test_invalid_compiler_flag(self):
         if cupy.cuda.runtime.is_hip and self.backend == 'nvrtc':
             self.skipTest('hiprtc does not handle #error macro properly')
 
-        with pytest.raises(cupy.cuda.compiler.CompileException) as ex:
+        if self.jitify:
+            ex_type = cupy.cuda.compiler.JitifyException
+        else:
+            ex_type = cupy.cuda.compiler.CompileException
+
+        with pytest.raises(ex_type) as ex:
             mod = cupy.RawModule(code=_test_source3,
                                  options=('-DPRECISION=3',),
-                                 backend=self.backend)
+                                 backend=self.backend,
+                                 jitify=self.jitify)
             mod.get_function('test_sum')  # enforce compilation
-        assert 'precision not supported' in str(ex.value)
+
+        if not self.jitify:
+            assert 'precision not supported' in str(ex.value)
 
     def _generate_file(self, ext: str):
         # generate cubin/ptx by calling nvcc/hipcc
@@ -598,7 +619,7 @@ def test_get_function_failure(self):
                      'ROCm/HIP does not support dynamic parallelism')
     def test_dynamical_parallelism(self):
         ker = cupy.RawKernel(_test_source4, 'test_kernel', options=('-dc',),
-                             backend=self.backend)
+                             backend=self.backend, jitify=self.jitify)
         N = 169
         inner_chunk = 13
         x = cupy.zeros((N,), dtype=cupy.float32)
@@ -608,7 +629,7 @@ def test_dynamical_parallelism(self):
     def test_dynamical_parallelism_compile_failure(self):
         # no option for separate compilation is given should cause an error
         ker = cupy.RawKernel(_test_source4, 'test_kernel',
-                             backend=self.backend)
+                             backend=self.backend, jitify=self.jitify)
         N = 10
         inner_chunk = 2
         x = cupy.zeros((N,), dtype=cupy.float32)
@@ -630,7 +651,8 @@ def test_cuFloatComplex(self):
 
         mod = cupy.RawModule(
             code=_test_cuComplex,
-            translate_cucomplex=True)
+            translate_cucomplex=True,
+            jitify=self.jitify)
         a = cupy.random.random((N,)) + 1j*cupy.random.random((N,))
         a = a.astype(dtype)
         b = cupy.random.random((N,)) + 1j*cupy.random.random((N,))
@@ -694,7 +716,8 @@ def test_cuDoubleComplex(self):
 
         mod = cupy.RawModule(
             code=_test_cuComplex,
-            translate_cucomplex=True)
+            translate_cucomplex=True,
+            jitify=self.jitify)
         a = cupy.random.random((N,)) + 1j*cupy.random.random((N,))
         a = a.astype(dtype)
         b = cupy.random.random((N,)) + 1j*cupy.random.random((N,))
@@ -754,7 +777,9 @@ def test_cuDoubleComplex(self):
         assert (out == a + b).all()
 
     def test_const_memory(self):
-        mod = cupy.RawModule(code=test_const_mem, backend=self.backend)
+        mod = cupy.RawModule(code=test_const_mem,
+                             backend=self.backend,
+                             jitify=self.jitify)
         ker = mod.get_function('multiply_by_const')
         mem_ptr = mod.get_global('some_array')
         const_arr = cupy.ndarray((100,), cupy.float32, mem_ptr)
@@ -777,8 +802,10 @@ def test_template_specialization(self):
         # compile code
         name_expressions = ['my_sqrt<int>', 'my_sqrt<float>',
                             'my_sqrt<complex<double>>', 'my_func']
-        mod = cupy.RawModule(code=test_cxx_template, options=('--std=c++11',),
-                             name_expressions=name_expressions)
+        mod = cupy.RawModule(code=test_cxx_template,
+                             options=('--std=c++11',),
+                             name_expressions=name_expressions,
+                             jitify=self.jitify)
 
         dtypes = (cupy.int32, cupy.float32, cupy.complex128, cupy.float64)
         for ker_T, dtype in zip(name_expressions, dtypes):
@@ -814,7 +841,9 @@ def test_template_failure(self):
             return  # the rest of tests do not apply to nvcc
 
         # 2. compile code without specializations
-        mod = cupy.RawModule(code=test_cxx_template, options=('--std=c++11',))
+        mod = cupy.RawModule(code=test_cxx_template,
+                             options=('--std=c++11',),
+                             jitify=self.jitify)
         # ...try to get a specialized kernel
         match = ('named symbol not found' if not cupy.cuda.runtime.is_hip else
                  'hipErrorNotFound')
@@ -824,11 +853,14 @@ def test_template_failure(self):
         # 3. compile code without specifying C++ standard
         with pytest.raises(ValueError):
             cupy.RawModule(code=test_cxx_template,
-                           name_expressions=name_expressions)
+                           name_expressions=name_expressions,
+                           jitify=self.jitify)
 
         # 4. try to fetch something we didn't specialize for
-        mod = cupy.RawModule(code=test_cxx_template, options=('--std=c++11',),
-                             name_expressions=name_expressions)
+        mod = cupy.RawModule(code=test_cxx_template,
+                             options=('--std=c++11',),
+                             name_expressions=name_expressions,
+                             jitify=self.jitify)
         if cupy.cuda.runtime.is_hip:
             msg = 'hipErrorNotFound'
         else:
@@ -837,7 +869,9 @@ def test_template_failure(self):
             mod.get_function('my_sqrt<double>')
 
     def test_raw_pointer(self):
-        mod = cupy.RawModule(code=test_cast, backend=self.backend)
+        mod = cupy.RawModule(code=test_cast,
+                             backend=self.backend,
+                             jitify=self.jitify)
         ker = mod.get_function('my_func')
 
         a = cupy.ones((100,), dtype=cupy.float64)
@@ -939,7 +973,8 @@ def test_context_switch_RawModule5(self):
         with cupy.cuda.Device(0):
             mod = cupy.RawModule(code=test_cxx_template,
                                  options=('--std=c++11',),
-                                 name_expressions=name_expressions)
+                                 name_expressions=name_expressions,
+                                 jitify=self.jitify)
 
             # get specialized kernels
             mod.get_function(name)
@@ -972,7 +1007,8 @@ def test_context_switch_RawModule6(self):
         with cupy.cuda.Device(0):
             mod = cupy.RawModule(code=test_cxx_template,
                                  options=('--std=c++11',),
-                                 name_expressions=name_expressions)
+                                 name_expressions=name_expressions,
+                                 jitify=self.jitify)
 
             # get specialized kernels
             ker = mod.get_function(name)
@@ -995,7 +1031,8 @@ def test_compile_kernel(self):
         kern = cupy.RawKernel(
             _test_compile_src, 'test_op',
             options=('-DOP=+',),
-            backend=self.backend)
+            backend=self.backend,
+            jitify=self.jitify)
         log = io.StringIO()
         with use_temporary_cache_dir():
             kern.compile(log_stream=log)
@@ -1009,7 +1046,8 @@ def test_compile_module(self):
         module = cupy.RawModule(
             code=_test_compile_src,
             backend=self.backend,
-            options=('-DOP=+',))
+            options=('-DOP=+',),
+            jitify=self.jitify)
         log = io.StringIO()
         with use_temporary_cache_dir():
             module.compile(log_stream=log)
@@ -1170,3 +1208,134 @@ def test_raw_picklable(self):
         s = subprocess.run([sys.executable,
                             self.temp_dir + '/TestRawPicklable.py'])
         s.check_returncode()  # raise if unsuccess
+
+
+# a slightly more realistic kernel involving std utilities
+std_code = r'''
+#include <type_traits>
+
+template<typename T,
+         typename = typename std::enable_if<std::is_integral<T>::value>::type>
+__global__ void shift (T* a, int N) {
+    unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid < N) {
+        a[tid] += 100;
+    }
+}
+'''
+
+
+@testing.parameterize(*testing.product({
+    'jitify': (False, True),
+}))
+@unittest.skipIf(cupy.cuda.runtime.is_hip,
+                 'Jitify does not support ROCm/HIP')
+class TestRawJitify(unittest.TestCase):
+
+    def setUp(self):
+        self.temporary_dir_context = use_temporary_cache_dir()
+        self.temp_dir = self.temporary_dir_context.__enter__()
+
+    def tearDown(self):
+        self.temporary_dir_context.__exit__(*sys.exc_info())
+
+    def _helper(self, header, options=()):
+        code = header
+        code += _test_source1
+        mod1 = cupy.RawModule(code=code,
+                              backend='nvrtc',
+                              options=options,
+                              jitify=self.jitify)
+
+        N = 10
+        x1 = cupy.arange(N**2, dtype=cupy.float32).reshape(N, N)
+        x2 = cupy.ones((N, N), dtype=cupy.float32)
+        y = cupy.zeros((N, N), dtype=cupy.float32)
+        ker = mod1.get_function('test_sum')
+        ker((N,), (N,), (x1, x2, y, N**2))
+        assert cupy.allclose(x1 + x2, y)
+
+    def _helper2(self, type_str):
+        mod2 = cupy.RawModule(code=std_code,
+                              jitify=self.jitify,
+                              name_expressions=['shift<%s>' % type_str, ],
+                              options=('--std=c++11',))
+        ker = mod2.get_function('shift<%s>' % type_str)
+        N = 256
+        a = cupy.random.random_integers(0, 7, N).astype(cupy.int32)
+        b = a.copy()
+        ker((1,), (N,), (a, N))
+        assert cupy.allclose(a, b+100)
+
+    def test_jitify1(self):
+        # simply prepend an unused header
+        hdr = '#include <cupy/cub/cub/block/block_reduce.cuh>\n'
+
+        if self.jitify:
+            # Jitify will make it work
+            self._helper(hdr)
+        else:
+            # NVRTC cannot find C++ std headers without Jitify
+            with pytest.raises(cupy.cuda.compiler.CompileException) as ex:
+                self._helper(hdr)
+            assert 'cannot open source file' in str(ex.value)
+
+    def test_jitify2(self):
+        # NVRTC cannot compile any code involving std
+        if self.jitify:
+            # Jitify will make it work
+            self._helper2('int')
+        else:
+            with pytest.raises(cupy.cuda.compiler.CompileException) as ex:
+                self._helper2('int')
+            assert 'cannot open source file' in str(ex.value)
+
+    def test_jitify3(self):
+        # We supply a type impossible to specialize. Jitify is still able to
+        # locate the headers, but when it comes to the actual compilation,
+        # NVRTC fails (raising the same exception) with different error
+        # messages.
+        ex_type = cupy.cuda.compiler.CompileException
+        with pytest.raises(ex_type) as ex:
+            self._helper2('float')
+        if self.jitify:
+            assert 'Error in parsing name expression' in str(ex.value)
+        else:
+            assert 'cannot open source file' in str(ex.value)
+
+    def test_jitify4(self):
+        # ensure JitifyException is raised with a broken code
+        code = r'''
+        __global__ void i_am_broken() {
+        '''
+
+        if self.jitify:
+            ex_type = cupy.cuda.compiler.JitifyException
+        else:
+            ex_type = cupy.cuda.compiler.CompileException
+
+        with pytest.raises(ex_type):
+            mod = cupy.RawModule(code=code, jitify=self.jitify)
+            ker = mod.get_function('i_am_broken')  # noqa
+        # if Jitify could redirect its output, we would be able to check
+        # the error log here as well (NVIDIA/jitify#79)
+
+    def test_jitify5(self):
+        # If including a header that does not exist, Jitify would attempt to
+        # comment it out and proceed. If this header is actually unused, then
+        # everything would run just fine.
+
+        hdr = 'I_INCLUDE_SOMETHING.h'
+        with open(self.temp_dir + '/' + hdr, 'w') as f:
+            dummy = '#include <cupy/I_DO_NOT_EXIST_WAH_HA_HA.h>\n'
+            f.write(dummy)
+        hdr = '#include "' + hdr + '"\n'
+
+        if self.jitify:
+            # Jitify would print a warning "[jitify] File not found" to stdout,
+            # but as mentioned above and elsewhere, we can't capture it.
+            self._helper(hdr, options=('-I'+self.temp_dir,))
+        else:
+            with pytest.raises(cupy.cuda.compiler.CompileException) as ex:
+                self._helper(hdr, options=('-I'+self.temp_dir,))
+            assert 'cannot open source file' in str(ex.value)