Repository: cupy/cupy

Issue:
Allow C++ Templating Functionality
Would it be possible to allow C++ Templating functionality to raw kernels?

Rationale:
Instead of having multiple functions based on input type, there would be one templated global function.


Discussion:
Unfortunately this is not possible despite how easy it seems. The reason is C++ name mangling. When specializing a template function, its name gets mangled (note that the template specialization can't be done with the `extern "C"` linkage to avoid name mangling), so you can't retrieve the function by name via `cuModuleGetFunction` (the actual CUDA API used in `RawKernel` and `RawModule`). The name just gets distorted by the compiler.

This is something that's bothering indeed. For example, we might want to cache the compiled template code in disk (ref: numba/numba#4389) for future reuse, but given this challenge we have no way to retrieve the kernels (but compiling is OK!) 
Example:
```python
import cupy as cp


template = '''
template <typename T>
__global__ void do_something(T* input, int N) {
    unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
    if (tid < N) {
        input[tid] += 1;
    }
}

// specialization for float
template<>
__global__ void do_something(float* input, int N);
'''

mod = cp.RawModule(code=template, options=('--std=c++11',))  # <-- compiling is OK!
func = mod.get_function('do_something')                      # <-- error
```
Error:
```python
Traceback (most recent call last):
  File "template.py", line 19, in <module>
    func = mod.get_function('do_something')
  File "cupy/core/raw.pyx", line 302, in cupy.core.raw.RawModule.get_function
  File "cupy/cuda/function.pyx", line 213, in cupy.cuda.function.Module.get_function
  File "cupy/cuda/function.pyx", line 216, in cupy.cuda.function.Module.get_function
  File "cupy/cuda/function.pyx", line 162, in cupy.cuda.function.Function.__init__
  File "cupy/cuda/driver.pyx", line 262, in cupy.cuda.driver.moduleGetFunction
  File "cupy/cuda/driver.pyx", line 124, in cupy.cuda.driver.check_status
cupy.cuda.driver.CUDADriverError: CUDA_ERROR_NOT_FOUND: named symbol not found
```

Okay, I expected as much. I sounds like if want multiple types, I'll need multiple kernels...
What is rationale for `extern C` ? Would an improvement to `cuModuleGetFunction` alleviate the issue?
`extern "C"`: https://stackoverflow.com/questions/1041866/what-is-the-effect-of-extern-c-in-c

`cuModuleGetFunction` is a CUDA driver API. There's nothing we can do to it. 

If you want to make your life *much* harder, you can always write a dispatcher that can live with the C linkage, but internally calls any C++ specializations you want. This is more or less how CuPy interfaces with some C++ template code internally, see:
- [thrust.pyx](https://github.com/cupy/cupy/blob/master/cupy/cuda/thrust.pyx) (interface with [cupy_thrust.cu](https://github.com/cupy/cupy/blob/master/cupy/cuda/cupy_thrust.cu))
- [cub.pyx](https://github.com/cupy/cupy/blob/master/cupy/cuda/cub.pyx) (interface with [cupy_cub.cu](https://github.com/cupy/cupy/blob/master/cupy/cuda/cupy_cub.cu))

But you might as well write a string template in Python, do a string replacement with the types you want, and recompile using RawKernel/RawModule. This is much easier.

@leofang When are say _string template_, are you referring to the following?
https://www.geeksforgeeks.org/template-class-in-python/

I could have something like
```
# A Python program to demonstrate working of string template 
from string import Template 
  
# List Student stores the name and marks of three students 
type = ['float', 'double', 'complex<float>', 'complex<double>'] 
  
# We are creating a basic structure to print the name and 
# marks of the students. 
t = Template(<some cuda kernel>) 
  
for i in type: 
     print (t.substitute(data_type = i))
```

Then call _t_ from inside the _extern "C"_, inside the _loaded_from_source_?
That's right. In CuPy both ufuncs and reduction kernels are done like this. For example, this is ufunc: https://github.com/cupy/cupy/blob/f00bf53ec16c67ce26b06c0afcd76da9df038349/cupy/core/_kernel.pyx#L38-L62
It's a bit more complicated because of the machinery needed for generalizations, but depending on your use cases it could be much simpler.
Thanks!
@leofang Sorry to reopen this but I wanted to ask a question related to the issue. I now have a better understanding how CuPy is using NVRTC behind the scenes. I've been reading up on it more to possibility use in other applications, mainly C++. While reading the documentation I came across a section discussing interacting with Template host code https://docs.nvidia.com/cuda/nvrtc/index.html#host-template-interface. Is possible to add this functionality to CuPy? I'm assuming this is the framework CuPy is already using.

Resolving patch:
diff --git a/cupy/cuda/compiler.py b/cupy/cuda/compiler.py
--- a/cupy/cuda/compiler.py
+++ b/cupy/cuda/compiler.py
@@ -13,6 +13,7 @@
 from cupy.cuda import runtime
 from cupy import util
 
+
 _nvrtc_version = None
 _nvrtc_max_compute_capability = None
 _win32 = sys.platform.startswith('win32')
@@ -75,6 +76,10 @@ def _is_cudadevrt_needed(options):
 
 
 def _get_cudadevrt_path():
+    global _cudadevrt
+    if _cudadevrt is not None:
+        return _cudadevrt
+
     # defer import to here to avoid circular dependency
     from cupy.cuda import get_cuda_path
     global _win32
@@ -113,7 +118,8 @@ def _get_bool_env_variable(name, default):
         return False
 
 
-def compile_using_nvrtc(source, options=(), arch=None, filename='kern.cu'):
+def compile_using_nvrtc(source, options=(), arch=None, filename='kern.cu',
+                        name_expressions=None):
     if not arch:
         arch = _get_arch()
 
@@ -125,9 +131,10 @@ def compile_using_nvrtc(source, options=(), arch=None, filename='kern.cu'):
         with open(cu_path, 'w') as cu_file:
             cu_file.write(source)
 
-        prog = _NVRTCProgram(source, cu_path)
+        prog = _NVRTCProgram(source, cu_path,
+                             name_expressions=name_expressions)
         try:
-            ptx = prog.compile(options)
+            ptx, mapping = prog.compile(options)
         except CompileException as e:
             dump = _get_bool_env_variable(
                 'CUPY_DUMP_CUDA_SOURCE_ON_ERROR', False)
@@ -135,7 +142,7 @@ def compile_using_nvrtc(source, options=(), arch=None, filename='kern.cu'):
                 e.dump(sys.stderr)
             raise
 
-        return ptx
+        return ptx, mapping
 
 
 def compile_using_nvcc(source, options=(), arch=None,
@@ -232,7 +239,7 @@ def _preprocess(source, options, arch, backend):
 
         prog = _NVRTCProgram(source, '')
         try:
-            result = prog.compile(options)
+            result, _ = prog.compile(options)
         except CompileException as e:
             dump = _get_bool_env_variable(
                 'CUPY_DUMP_CUDA_SOURCE_ON_ERROR', False)
@@ -268,7 +275,8 @@ def get_cache_dir():
 
 def compile_with_cache(
         source, options=(), arch=None, cache_dir=None, extra_source=None,
-        backend='nvrtc', *, enable_cooperative_groups=False):
+        backend='nvrtc', *, enable_cooperative_groups=False,
+        name_expressions=None):
 
     if enable_cooperative_groups:
         if backend != 'nvcc':
@@ -278,18 +286,22 @@ def compile_with_cache(
             raise ValueError(
                 'Cooperative groups is not supported in HIP.')
 
+    if name_expressions is not None:
+        if runtime.is_hip or backend != 'nvrtc':
+            raise NotImplementedError
+
     if runtime.is_hip:
         return _compile_with_cache_hipcc(
             source, options, arch, cache_dir, extra_source)
     else:
         return _compile_with_cache_cuda(
             source, options, arch, cache_dir, extra_source, backend,
-            enable_cooperative_groups)
+            enable_cooperative_groups, name_expressions)
 
 
 def _compile_with_cache_cuda(
         source, options, arch, cache_dir, extra_source=None, backend='nvrtc',
-        enable_cooperative_groups=False):
+        enable_cooperative_groups=False, name_expressions=None):
     # NVRTC does not use extra_source. extra_source is used for cache key.
     global _empty_file_preprocess_cache
     if cache_dir is None:
@@ -330,8 +342,9 @@ def _compile_with_cache_cuda(
     mod = function.Module()
     # To handle conflicts in concurrent situation, we adopt lock-free method
     # to avoid performance degradation.
+    # We force recompiling to retrieve C++ mangled names if so desired.
     path = os.path.join(cache_dir, name)
-    if os.path.exists(path):
+    if os.path.exists(path) and not name_expressions:
         with open(path, 'rb') as file:
             data = file.read()
         if len(data) >= 32:
@@ -343,16 +356,16 @@ def _compile_with_cache_cuda(
                 return mod
 
     if backend == 'nvrtc':
-        ptx = compile_using_nvrtc(source, options, arch, name + '.cu')
+        ptx, mapping = compile_using_nvrtc(
+            source, options, arch, name + '.cu', name_expressions)
         ls = function.LinkState()
         ls.add_ptr_data(ptx, 'cupy.ptx')
         # for separate compilation
         if _is_cudadevrt_needed(options):
-            global _cudadevrt
-            if _cudadevrt is None:
-                _cudadevrt = _get_cudadevrt_path()
+            _cudadevrt = _get_cudadevrt_path()
             ls.add_ptr_file(_cudadevrt)
         cubin = ls.complete()
+        mod._set_mapping(mapping)
     elif backend == 'nvcc':
         rdc = _is_cudadevrt_needed(options)
         cubin = compile_using_nvcc(source, options, arch, name + '.cu',
@@ -422,7 +435,7 @@ def dump(self, f):
 class _NVRTCProgram(object):
 
     def __init__(self, src, name='default_program', headers=(),
-                 include_names=()):
+                 include_names=(), name_expressions=None):
         self.ptr = None
 
         if isinstance(src, bytes):
@@ -433,6 +446,7 @@ def __init__(self, src, name='default_program', headers=(),
         self.src = src
         self.name = name
         self.ptr = nvrtc.createProgram(src, name, headers, include_names)
+        self.name_expressions = name_expressions
 
     def __del__(self, is_shutting_down=util.is_shutting_down):
         if is_shutting_down():
@@ -442,8 +456,16 @@ def __del__(self, is_shutting_down=util.is_shutting_down):
 
     def compile(self, options=()):
         try:
+            if self.name_expressions:
+                for ker in self.name_expressions:
+                    nvrtc.addAddNameExpression(self.ptr, ker)
             nvrtc.compileProgram(self.ptr, options)
-            return nvrtc.getPTX(self.ptr)
+            mapping = None
+            if self.name_expressions:
+                mapping = {}
+                for ker in self.name_expressions:
+                    mapping[ker] = nvrtc.getLoweredName(self.ptr, ker)
+            return nvrtc.getPTX(self.ptr), mapping
         except nvrtc.NVRTCError:
             log = nvrtc.getProgramLog(self.ptr)
             raise CompileException(log, self.src, self.name, options, 'nvrtc')


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
@@ -281,6 +281,25 @@
 }
 '''
 
+test_cxx_template = r'''
+#include <cupy/complex.cuh>
+
+template<typename T>
+__global__ void my_sqrt(T* input, int N) {
+  unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;
+  if (x < N) {
+    input[x] *= input[x];
+  }
+}
+
+__global__ void my_func(double* input, int N) {
+  unsigned int x = blockIdx.x * blockDim.x + threadIdx.x;
+  if (x < N) {
+    input[x] *= input[x];
+  }
+}
+'''
+
 if 'CUPY_CACHE_DIR' in os.environ:
     _old_cache_dir = os.environ['CUPY_CACHE_DIR']
     _is_cache_env_var_set = True
@@ -325,8 +344,6 @@ def tearDown(self):
             os.environ.pop('CUPY_CACHE_DIR')
         compiler._empty_file_preprocess_cache = {}
 
-        cupy.cuda.runtime.setDevice(self.dev)
-
     def _helper(self, kernel, dtype):
         N = 10
         x1 = cupy.arange(N**2, dtype=dtype).reshape(N, N)
@@ -606,27 +623,82 @@ def test_const_memory(self):
         ker((1,), (100,), (output_arr, cupy.int32(100)))
         assert (data == output_arr).all()
 
+    def test_template_specialization(self):
+        if self.backend == 'nvcc':
+            self.skipTest('nvcc does not support template specialization')
+
+        # 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)
+
+        dtypes = (cupy.int32, cupy.float32, cupy.complex128, cupy.float64)
+        for ker_T, dtype in zip(name_expressions, dtypes):
+            # get specialized kernels
+            ker = mod.get_function(ker_T)
+
+            # prepare inputs & expected outputs
+            in_arr = cupy.testing.shaped_random((10,), dtype=dtype)
+            out_arr = in_arr**2
+
+            # run
+            ker((1,), (10,), (in_arr, 10))
+
+            # check results
+            assert cupy.allclose(in_arr, out_arr)
+
+    def test_template_failure(self):
+        name_expressions = ['my_sqrt<int>']
+
+        # 1. nvcc is disabled for this feature
+        if self.backend == 'nvcc':
+            with pytest.raises(ValueError) as e:
+                cupy.RawModule(code=test_cxx_template, backend=self.backend,
+                               options=('--std=c++11',),
+                               name_expressions=name_expressions)
+            assert 'nvrtc' in str(e.value)
+            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',))
+        # ...try to get a specialized kernel
+        with pytest.raises(cupy.cuda.driver.CUDADriverError,
+                           match='named symbol not found'):
+            mod.get_function('my_sqrt<int>')
+
+        # 3. compile code without specifying C++ standard
+        with pytest.raises(ValueError):
+            cupy.RawModule(code=test_cxx_template,
+                           name_expressions=name_expressions)
+
+        # 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)
+        with pytest.raises(cupy.cuda.driver.CUDADriverError,
+                           match='named symbol not found'):
+            mod.get_function('my_sqrt<double>')
+
     @testing.multi_gpu(2)
     def test_context_switch_RawKernel(self):
         # run test_basic() on another device
 
         # For RawKernel, we need to launch it once to force compiling
         x1, x2, y = self._helper(self.kern, cupy.float32)
-        cupy.cuda.runtime.setDevice(1)
 
-        x1, x2, y = self._helper(self.kern, cupy.float32)
-        assert cupy.allclose(y, x1 + x2)
+        with cupy.cuda.Device(1):
+            x1, x2, y = self._helper(self.kern, cupy.float32)
+            assert cupy.allclose(y, x1 + x2)
 
     @testing.multi_gpu(2)
     def test_context_switch_RawModule1(self):
         # run test_module() on another device
         # in this test, re-compiling happens at get_function()
-        cupy.cuda.runtime.setDevice(1)
-
-        module = self.mod2
-        ker_sum = module.get_function('test_sum')
-        x1, x2, y = self._helper(ker_sum, cupy.float32)
-        assert cupy.allclose(y, x1 + x2)
+        with cupy.cuda.Device(1):
+            module = self.mod2
+            ker_sum = module.get_function('test_sum')
+            x1, x2, y = self._helper(ker_sum, cupy.float32)
+            assert cupy.allclose(y, x1 + x2)
 
     @testing.multi_gpu(2)
     def test_context_switch_RawModule2(self):
@@ -634,10 +706,10 @@ def test_context_switch_RawModule2(self):
         # in this test, re-compiling happens at kernel launch
         module = self.mod2
         ker_sum = module.get_function('test_sum')
-        cupy.cuda.runtime.setDevice(1)
 
-        x1, x2, y = self._helper(ker_sum, cupy.float32)
-        assert cupy.allclose(y, x1 + x2)
+        with cupy.cuda.Device(1):
+            x1, x2, y = self._helper(ker_sum, cupy.float32)
+            assert cupy.allclose(y, x1 + x2)
 
     @testing.multi_gpu(2)
     def test_context_switch_RawModule3(self):
@@ -645,12 +717,12 @@ def test_context_switch_RawModule3(self):
         # generate cubin in the temp dir and load it on device 0
         file_path = self._generate_file('cubin')
         mod = cupy.RawModule(path=file_path, backend=self.backend)
-        # in this test, reloading happens at get_function()
-        cupy.cuda.runtime.setDevice(1)
 
-        ker = mod.get_function('test_div')
-        x1, x2, y = self._helper(ker, cupy.float32)
-        assert cupy.allclose(y, x1 / (x2 + 1.0))
+        # in this test, reloading happens at get_function()
+        with cupy.cuda.Device(1):
+            ker = mod.get_function('test_div')
+            x1, x2, y = self._helper(ker, cupy.float32)
+            assert cupy.allclose(y, x1 / (x2 + 1.0))
 
     @testing.multi_gpu(2)
     def test_context_switch_RawModule4(self):
@@ -659,11 +731,67 @@ def test_context_switch_RawModule4(self):
         file_path = self._generate_file('cubin')
         mod = cupy.RawModule(path=file_path, backend=self.backend)
         ker = mod.get_function('test_div')
+
         # in this test, reloading happens at kernel launch
-        cupy.cuda.runtime.setDevice(1)
+        with cupy.cuda.Device(1):
+            x1, x2, y = self._helper(ker, cupy.float32)
+            assert cupy.allclose(y, x1 / (x2 + 1.0))
 
-        x1, x2, y = self._helper(ker, cupy.float32)
-        assert cupy.allclose(y, x1 / (x2 + 1.0))
+    @testing.multi_gpu(2)
+    def test_context_switch_RawModule5(self):
+        # run test_template_specialization() on another device
+        # in this test, re-compiling happens at get_function()
+        if self.backend == 'nvcc':
+            self.skipTest('nvcc does not support template specialization')
+
+        # compile code
+        name_expressions = ['my_sqrt<unsigned int>']
+        mod = cupy.RawModule(code=test_cxx_template, options=('--std=c++11',),
+                             name_expressions=name_expressions)
+
+        # switch device
+        with cupy.cuda.Device(1):
+            # get specialized kernels
+            name = name_expressions[0]
+            ker = mod.get_function(name)
+
+            # prepare inputs & expected outputs
+            in_arr = cupy.testing.shaped_random((10,), dtype=cupy.uint32)
+            out_arr = in_arr**2
+
+            # run
+            ker((1,), (10,), (in_arr, 10))
+
+            # check results
+            assert cupy.allclose(in_arr, out_arr)
+
+    @testing.multi_gpu(2)
+    def test_context_switch_RawModule6(self):
+        # run test_template_specialization() on another device
+        # in this test, re-compiling happens at kernel launch
+        if self.backend == 'nvcc':
+            self.skipTest('nvcc does not support template specialization')
+
+        # compile code
+        name_expressions = ['my_sqrt<unsigned int>']
+        mod = cupy.RawModule(code=test_cxx_template, options=('--std=c++11',),
+                             name_expressions=name_expressions)
+
+        # get specialized kernels
+        name = name_expressions[0]
+        ker = mod.get_function(name)
+
+        # switch device
+        with cupy.cuda.Device(1):
+            # prepare inputs & expected outputs
+            in_arr = cupy.testing.shaped_random((10,), dtype=cupy.uint32)
+            out_arr = in_arr**2
+
+            # run
+            ker((1,), (10,), (in_arr, 10))
+
+            # check results
+            assert cupy.allclose(in_arr, out_arr)
 
 
 _test_grid_sync = r'''