Repository: cupy/cupy

Issue:
`cupyx.jit` forgot to process preamble and `out0_type` identifiers for expressions cannot be inlined.
### Description

As described in the title, in the code:
https://github.com/cupy/cupy/blob/a1c775f4eb0d3560ab0a5a4a3eb50cd2dbd911bb/cupyx/jit/_compile.py#L459-L474

The translation forgot to process preamble and `out0_type` if the expression cannot be inlined.

### To Reproduce

```py
# Write the code here
import os
os.environ['CUPY_DUMP_CUDA_SOURCE_ON_ERROR'] = "1"
import cupy as cp
from cupyx import jit
@jit.rawkernel()
def clip(x, a_min, a_max, out):
    tid = jit.blockDim.x*jit.blockIdx.x+jit.threadIdx.x
    t_stride = jit.gridDim.x*jit.blockDim.x
    for j in range(x.shape[1]):
        for i in range(tid, x.shape[0], t_stride):
            out[i,j] = min(max(x[i,j],a_min),a_max)
inmat = cp.random.normal(size=(4,4))
out = cp.empty_like(inmat)
clip[(128,), (1024,)](inmat, 0.0, 0.5, out)
print(inmat, out)
```

## What I got

```
NVRTC compilation error: C:\Users\berry\AppData\Local\Temp\tmpean568ny\1b653cc5fd909314f9f9b18035b9de662ff5873c.cubin.cu(7): error: identifier "NAN" is undefined

C:\Users\berry\AppData\Local\Temp\tmpean568ny\1b653cc5fd909314f9f9b18035b9de662ff5873c.cubin.cu(7): error: identifier "out0_type" is undefined

C:\Users\berry\AppData\Local\Temp\tmpean568ny\1b653cc5fd909314f9f9b18035b9de662ff5873c.cubin.cu(12): error: identifier "NAN" is undefined

C:\Users\berry\AppData\Local\Temp\tmpean568ny\1b653cc5fd909314f9f9b18035b9de662ff5873c.cubin.cu(12): error: identifier "out0_type" is undefined

4 errors detected in the compilation of "C:\Users\berry\AppData\Local\Temp\tmpean568ny\1b653cc5fd909314f9f9b18035b9de662ff5873c.cubin.cu".

-----
Name: C:\Users\berry\AppData\Local\Temp\tmpean568ny\1b653cc5fd909314f9f9b18035b9de662ff5873c.cubin.cu
Options: -DCUPY_JIT_MODE --std=c++14 -IJ:\ram_intro\venv\lib\site-packages\cupy\_core\include -IJ:\ram_intro\venv\lib\site-packages\cupy\_core\include\cupy\_cuda\cuda-11 -IC:\NV_CUDA\v11.7\include -ftz=true -arch=sm_86
CUDA source:
01 #include <cupy/complex.cuh>
02 #include <cupy/carray.cuh>
03 #include <cupy/atomics.cuh>
04 
05 template <typename in0_type, typename in1_type>
06 __device__ double cupy_maximum_float64(in0_type in0, in1_type in1) {
07     return (isnan(in0) | isnan(in1)) ? out0_type(NAN) : out0_type(max(in0, in1));
08 }
09 
10 template <typename in0_type, typename in1_type>
11 __device__ double cupy_minimum_float64(in0_type in0, in1_type in1) {
12     return (isnan(in0) | isnan(in1)) ? out0_type(NAN) : out0_type(min(in0, in1));
13 }
14 
15 extern "C" __global__ void clip(CArray<double, 2, true, true> x, float a_min, float a_max, CArray<double, 2, true, true> out) {
16   ptrdiff_t j;
17   ptrdiff_t i;
18   unsigned int tid = ((blockDim.x * blockIdx.x) + threadIdx.x);
19   unsigned int t_stride = (gridDim.x * blockDim.x);
20   for (ptrdiff_t __it = 0, __stop = thrust::get<1>(x.get_shape()), __step = 1; __it < __stop; __it += __step) {
21     j = __it;
22     for (ptrdiff_t __it = tid, __stop = thrust::get<0>(x.get_shape()), __step = t_stride; __it < __stop; __it += __step) {
23       i = __it;
24       out._indexing(thrust::make_tuple(i, j)) = cupy_minimum_float64(cupy_maximum_float64(x._indexing(thrust::make_tuple(i, j)), (double)(a_min)), (double)(a_max));
25     }
26   }
27 }
-----
Traceback (most recent call last):
  File "J:\ram_intro\venv\lib\site-packages\cupy\cuda\compiler.py", line 686, in compile
    nvrtc.compileProgram(self.ptr, options)
  File "cupy_backends\cuda\libs\nvrtc.pyx", line 154, in cupy_backends.cuda.libs.nvrtc.compileProgram
  File "cupy_backends\cuda\libs\nvrtc.pyx", line 166, in cupy_backends.cuda.libs.nvrtc.compileProgram
  File "cupy_backends\cuda\libs\nvrtc.pyx", line 84, in cupy_backends.cuda.libs.nvrtc.check_status
cupy_backends.cuda.libs.nvrtc.NVRTCError: NVRTC_ERROR_COMPILATION (6)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "C:\Users\berry\AppData\Roaming\JetBrains\PyCharm2022.2\scratches\scratch.py", line 102, in <module>
    clip[(128,), (1024,)](inmat, 0.0, 0.5, out)
  File "J:\ram_intro\venv\lib\site-packages\cupyx\jit\_interface.py", line 138, in <lambda>
    return lambda *args, **kwargs: self(grid, block, args, **kwargs)
  File "J:\ram_intro\venv\lib\site-packages\cupyx\jit\_interface.py", line 110, in __call__
    module = core.compile_with_cache(
  File "cupy\_core\core.pyx", line 2172, in cupy._core.core.compile_with_cache
  File "cupy\_core\core.pyx", line 2232, in cupy._core.core.compile_with_cache
  File "J:\ram_intro\venv\lib\site-packages\cupy\cuda\compiler.py", line 493, in _compile_module_with_cache
    return _compile_with_cache_cuda(
  File "J:\ram_intro\venv\lib\site-packages\cupy\cuda\compiler.py", line 571, in _compile_with_cache_cuda
    ptx, mapping = compile_using_nvrtc(
  File "J:\ram_intro\venv\lib\site-packages\cupy\cuda\compiler.py", line 322, in compile_using_nvrtc
    return _compile(source, options, cu_path,
  File "J:\ram_intro\venv\lib\site-packages\cupy\cuda\compiler.py", line 306, in _compile
    compiled_obj, mapping = prog.compile(options, log_stream)
  File "J:\ram_intro\venv\lib\site-packages\cupy\cuda\compiler.py", line 705, in compile
    raise CompileException(log, self.src, self.name, options,
cupy.cuda.compiler.CompileException: C:\Users\berry\AppData\Local\Temp\tmpean568ny\1b653cc5fd909314f9f9b18035b9de662ff5873c.cubin.cu(7): error: identifier "NAN" is undefined

C:\Users\berry\AppData\Local\Temp\tmpean568ny\1b653cc5fd909314f9f9b18035b9de662ff5873c.cubin.cu(7): error: identifier "out0_type" is undefined

C:\Users\berry\AppData\Local\Temp\tmpean568ny\1b653cc5fd909314f9f9b18035b9de662ff5873c.cubin.cu(12): error: identifier "NAN" is undefined

C:\Users\berry\AppData\Local\Temp\tmpean568ny\1b653cc5fd909314f9f9b18035b9de662ff5873c.cubin.cu(12): error: identifier "out0_type" is undefined

4 errors detected in the compilation of "C:\Users\berry\AppData\Local\Temp\tmpean568ny\1b653cc5fd909314f9f9b18035b9de662ff5873c.cubin.cu".

```

### Installation

Wheel (`pip install cupy-***`)

### Environment

```
# Paste the output here
OS                           : Windows-10-10.0.22621-SP0
Python Version               : 3.10.8
CuPy Version                 : 11.2.0
CuPy Platform                : NVIDIA CUDA
NumPy Version                : 1.23.4
SciPy Version                : 1.9.3
Cython Build Version         : 0.29.32
Cython Runtime Version       : None
CUDA Root                    : C:\NV_CUDA\v11.7
nvcc PATH                    : C:\NV_CUDA\v11.7\bin\nvcc.EXE
CUDA Build Version           : 11070
CUDA Driver Version          : 11080
CUDA Runtime Version         : 11070
cuBLAS Version               : (available)
NCCL Build Version           : None
NCCL Runtime Version         : None
cuTENSOR Version             : None
cuSPARSELt Build Version     : None
Device 0 Name                : NVIDIA GeForce RTX 3090
Device 0 Compute Capability  : 86
Device 0 PCI Bus ID          : 0000:01:00.0
```



### Additional Information

Although I haven't tested the snippet on the latest commit, I believe the problem hasn't been fixed yet upon looking into the source code.


Discussion:

Resolving patch:
diff --git a/cupyx/jit/_compile.py b/cupyx/jit/_compile.py
--- a/cupyx/jit/_compile.py
+++ b/cupyx/jit/_compile.py
@@ -456,12 +456,13 @@ def _call_ufunc(
             if f'in{i}_type' in op.routine:
                 can_use_inline_expansion = False
 
+        env.generated.add_code(ufunc._preamble)
+
         if can_use_inline_expansion:
             # Code pass for readable generated code
             for i, x in enumerate(in_params):
                 expr = expr.replace(f'in{i}', x.code)
             expr = '(' + expr.replace('out0_type', str(out_type)) + ')'
-            env.generated.add_code(ufunc._preamble)
         else:
             template_typenames = ', '.join([
                 f'typename in{i}_type' for i in range(ufunc.nin)])
@@ -469,6 +470,7 @@ def _call_ufunc(
             params = ', '.join([f'in{i}_type in{i}' for i in range(ufunc.nin)])
             ufunc_code = f"""template <{template_typenames}>
 __device__ {out_type} {ufunc_name}({params}) {{
+    typedef {out_type} out0_type;
     return {expr};
 }}
 """


Test patch:
diff --git a/tests/cupy_tests/functional_tests/test_vectorize.py b/tests/cupy_tests/functional_tests/test_vectorize.py
--- a/tests/cupy_tests/functional_tests/test_vectorize.py
+++ b/tests/cupy_tests/functional_tests/test_vectorize.py
@@ -63,6 +63,19 @@ def my_pow(x, y):
         x1[x1 == 0] = 1
         return f(x1, x2)
 
+    @testing.for_all_dtypes_combination(
+        names=('dtype1', 'dtype2'), no_bool=True, no_complex=True)
+    @testing.numpy_cupy_allclose(rtol=1e-5)
+    def test_vectorize_minmax(self, xp, dtype1, dtype2):
+        def my_minmax(x, y):
+            return max(x, y) - min(x, y)
+
+        f = xp.vectorize(my_minmax)
+        x1 = testing.shaped_random((20, 30), xp, dtype1, seed=0)
+        x2 = testing.shaped_random((20, 30), xp, dtype2, seed=1)
+        x1[x1 == 0] = 1
+        return f(x1, x2)
+
     def run_div(self, func, xp, dtypes):
         dtype1, dtype2 = dtypes
         f = xp.vectorize(func)