We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Enhancement: enable CUB backend for ReductionKernel
Currently the CUB reduction is disabled for `ReductionKernel` (by the last `False` argument):
https://github.com/cupy/cupy/blob/c1bf1aab6eb79a1cfcd5bb374460082d8eb2f3ca/cupy/core/_reduction.pyx#L724-L727
But we should revisit this at some point.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `cupy/core/_fusion_op.py`
Content:
```
1 import string
2 
3 import numpy
4 
5 from cupy.core._fusion_variable import _TraceVariable
6 from cupy.core._fusion_variable import _TraceArray
7 from cupy.core._fusion_variable import _VariableSet
8 from cupy.core import _fusion_thread_local
9 from cupy.core import _fusion_emit_code
10 from cupy.core import _kernel
11 from cupy.core import _reduction
12 
13 
14 _dtype_to_ctype = _fusion_emit_code._dtype_to_ctype
15 
16 
17 class _UfuncRoutine:
18     """A device function for single elementwise operations.
19     """
20 
21     def __init__(
22             self, name, ufunc, routine_code, in_params, out_params,
23             compute_dtypes):
24         assert isinstance(name, str)
25         assert isinstance(ufunc, _kernel.ufunc)
26         assert isinstance(routine_code, str)
27         assert isinstance(compute_dtypes, tuple)
28         assert all(isinstance(t, numpy.dtype) for t in compute_dtypes)
29         assert isinstance(in_params, list)
30         assert all(isinstance(p, _TraceVariable) for p in in_params)
31         assert isinstance(out_params, list)
32         assert all(isinstance(p, _TraceArray) for p in out_params)
33 
34         self.name = name
35         self.in_params = in_params
36         self.out_params = out_params
37         self.preamble = ufunc._preamble
38         self.routine_code = routine_code
39         self.compute_dtypes = compute_dtypes
40 
41     def emit_code(self):
42         """Returns a CUDA device function code.
43 
44         Returns a string like:
45         ```
46         __device__ void cupy_add_0(int &in0_, float &in1_, double &out0_) {
47             typedef double in0_type;
48             typedef double in1_type;
49             typedef double out0_type;
50             double in0 = (double) in0_;
51             double in1 = (double) in1_;
52             double out0 = (double) out0_;
53             out0 = in0 + in1;
54             out0_ = out0;
55         }
56         ```
57         """
58         nin = len(self.in_params)
59         dtypes = self.compute_dtypes
60         assert len(self.in_params) == len(self.compute_dtypes[:nin])
61         in_params = [
62             (_dtype_to_ctype[p.dtype], _dtype_to_ctype[t], 'in{}'.format(i))
63             for i, (p, t) in enumerate(zip(self.in_params, dtypes[:nin]))
64         ]
65         out_params = [
66             (_dtype_to_ctype[p.dtype], _dtype_to_ctype[t], 'out{}'.format(i))
67             for i, (p, t) in enumerate(zip(self.out_params, dtypes[nin:]))
68         ]
69         params = in_params + out_params
70 
71         params_code = ', '.join(['{} &{}_'.format(t, s) for t, _, s in params])
72         typedef = ['typedef {} {}_type;'.format(t, s) for _, t, s in params]
73         read = ['{} {} = ({}) {}_;'.format(t, s, t, s) for _, t, s in params]
74         write = ['{}_ = {};'.format(s, s, s) for _, _, s in out_params]
75 
76         return _fusion_emit_code._CodeBlock(
77             '__device__ void {}({})'.format(self.name, params_code),
78             typedef + read + [self.routine_code + ';'] + write)
79 
80     def emit_call_code(self):
81         params = self.in_params + self.out_params
82         return '{op_name}({params});'.format(
83             op_name=self.name,
84             params=', '.join([var.lvar_name for var in params]))
85 
86 
87 class _ElementwiseTraceOp:
88     """Ufunc or elementwise kernel with types.
89     """
90 
91     def __init__(self, ufunc_routines, in_params, out_params, ashape):
92         # The `in_params` and `out_params` should be already broadcasted to
93         # `ashape`, but they don't guarantee to be exactly same as
94         # `param.ashape`.
95 
96         _fusion_thread_local.check_not_runtime()
97         assert isinstance(ufunc_routines, list)
98         assert all(isinstance(r, _UfuncRoutine) for r in ufunc_routines)
99         assert isinstance(ashape, tuple)
100 
101         self.ops = ufunc_routines
102         self.in_params = _VariableSet(*in_params)
103         self.out_params = _VariableSet(*out_params)
104         self.ashape = ashape
105 
106     @property
107     def params(self):
108         """Returns the set of all variable the loop uses.
109         """
110         res = _VariableSet()
111         for op in self.ops:
112             res += _VariableSet(*op.in_params)
113             res += _VariableSet(*op.out_params)
114         return res
115 
116     @staticmethod
117     def _emit_declaration(params, in_params):
118         """Returns a tuple of size 2.
119 
120         1. CUDA code: declaring local variables.
121             2. The set of arrays which require indexer.
122         """
123         _fusion_thread_local.check_not_runtime()
124 
125         indexed_arrays = _VariableSet()
126         code = []
127         for var in params:
128             if var in in_params:
129                 if isinstance(var, _TraceArray):
130                     indexed_arrays.add(var)
131                     f = '${type} ${lvar} = ${var}[${indexer}.get()];'
132                 else:
133                     f = '${type} ${lvar} = ${var};'
134             else:
135                 f = '${type} ${lvar};'
136             code.append(var.format(f))
137 
138         return code, indexed_arrays
139 
140     @staticmethod
141     def _emit_after_operation(out_params):
142         """Returns a tuple of size 2.
143         1. CUDA code: writing the results of operations back to global memory.
144         2. The set of arrays which require indexer.
145         """
146 
147         _fusion_thread_local.check_not_runtime()
148 
149         indexed_arrays = _VariableSet()
150         codes = []
151         for var in out_params:
152             if isinstance(var, _TraceArray):
153                 indexed_arrays.add(var)
154                 f = '${var}[${indexer}.get()] = ${lvar};'
155             else:
156                 f = '${var} = ${lvar};'
157             codes.append(var.format(f))
158 
159         return codes, indexed_arrays
160 
161     @staticmethod
162     def _emit_set_index(indexed_params, tid):
163         """Returns a CUDA code: setting a raw index to indexers.
164         """
165         _fusion_thread_local.check_not_runtime()
166         assert isinstance(indexed_params, _VariableSet)
167 
168         return [
169             p.format('${indexer}.set(${tid});', tid=tid)
170             for p in indexed_params
171         ]
172 
173     def emit_code(self):
174         _fusion_thread_local.check_not_runtime()
175 
176         declaration, s1 = self._emit_declaration(self.params, self.in_params)
177         operation = [op.emit_call_code() for op in self.ops]
178         after_operation, s2 = self._emit_after_operation(self.out_params)
179         index_name = 'i'
180         indexed_array = s1 + s2
181         indexer_name = next(iter(indexed_array)).indexer_name
182         indexer_setup = self._emit_set_index(indexed_array, index_name)
183 
184         return _fusion_emit_code._CodeBlock(
185             'CUPY_FOR({}, {}.size())'.format(index_name, indexer_name),
186             indexer_setup + declaration + operation + after_operation)
187 
188     def emit_preamble_codes(self):
189         return [subm.preamble for subm in self.ops if subm.preamble != '']
190 
191     def emit_submodule_codes(self):
192         return [str(subm.emit_code()) for subm in self.ops]
193 
194 
195 class _ReductionTraceOp:
196     def __init__(self, name, reduce_func, expr, in_param, out_param, axis):
197         """Reduction operation.
198         """
199         _fusion_thread_local.check_not_runtime()
200         assert isinstance(name, str)
201         assert isinstance(reduce_func, _reduction._SimpleReductionKernel)
202         assert isinstance(in_param, _TraceArray)
203         assert isinstance(out_param, _TraceArray)
204         assert isinstance(axis, tuple)
205         assert all([0 <= x < in_param.ndim for x in axis])
206 
207         self.name = name
208         self.preamble = reduce_func._preamble
209         self.in_params = _VariableSet(in_param)
210         self.out_params = _VariableSet(out_param)
211         self.block_stride_name = 'block_stride_' + name
212         self.axis = axis
213 
214         if reduce_func.identity is None:
215             self.identity = ''
216         else:
217             self.identity = str(reduce_func.identity)
218 
219         _, self.expr, self.postmap_cast_code, self.reduce_ctype = expr
220         if self.reduce_ctype is None:
221             out_param, = self.out_params
222             self.reduce_ctype = _dtype_to_ctype[out_param.dtype]
223 
224         self.premap_op = None
225         self.postmap_op = None
226 
227     @property
228     def params(self):
229         return self.in_params + self.out_params
230 
231     def emit_code(self):
232         _fusion_thread_local.check_not_runtime()
233         assert len(self.in_params) == 1
234         assert len(self.out_params) == 1
235         in_param = list(self.in_params)[0]
236         out_param = list(self.out_params)[0]
237         params = ', '.join([
238             in_param.var_name,
239             out_param.var_name,
240             in_param.indexer_name,
241             out_param.indexer_name,
242         ])
243         return '{}({}, {});'.format(
244             self.name, params, self.block_stride_name)
245 
246     def emit_preamble_codes(self):
247         preamble = self.preamble
248         return [preamble] if preamble != '' else []
249 
250     def emit_submodule_codes(self):
251         """Returns a CUDA device function code.
252 
253         The emitted code assumes that ``block_stride`` and `blockDim.x` is a
254         power of 2.
255         """
256 
257         in_param, = self.in_params
258         out_param, = self.out_params
259         op_name = '{}_op'.format(self.name)
260         postmap_name = '{}_postmap'.format(self.name)
261 
262         code = string.Template('''
263 #define ${op_name}(a, b) (${reduce_expr})
264 #define ${postmap_name}(a, out0) (${postmap_cast})
265 
266 template <typename InType, typename OutType, typename InIndexerType, typename OutIndexerType>
267 __device__ void ${name}(
268         InType in_arr, OutType out_arr,
269         InIndexerType in_ind, OutIndexerType out_ind, int block_stride) {
270     typedef ${in_type} type_in0_raw;
271     typedef ${out_type} type_out0_raw;
272     typedef ${reduce_ctype} _type_reduce;
273     extern __shared__ char _sdata_raw[];
274     _type_reduce *sdata = reinterpret_cast<_type_reduce*>(_sdata_raw);
275     unsigned int tid = threadIdx.x;
276     int _J = tid >> __popc(block_stride - 1);
277     ptrdiff_t _j = (ptrdiff_t)_J * out_ind.size();
278     int J_stride = blockDim.x >> __popc(block_stride - 1);
279     ptrdiff_t j_stride = (ptrdiff_t)J_stride * out_ind.size();
280 
281     for (ptrdiff_t _i = (ptrdiff_t)blockIdx.x * block_stride; _i < out_ind.size(); _i += (ptrdiff_t)gridDim.x * block_stride) {
282         _type_reduce s = _type_reduce(${identity});
283         ptrdiff_t i = _i + (tid & (block_stride - 1));
284         for (ptrdiff_t j = i + _j; j < in_ind.size(); j += j_stride) {
285             in_ind.set(j);
286             s = ${op_name}(s, static_cast<_type_reduce>(in_arr[in_ind.get()]));
287         }
288         sdata[tid] = s;
289         __syncthreads();
290         for (unsigned int block = blockDim.x / 2; block >= block_stride; block >>= 1) {
291             if (tid < block) {
292                 sdata[tid] = ${op_name}(sdata[tid], sdata[tid + block]);
293             }
294             __syncthreads();
295         }
296         if (tid < block_stride) {
297             s = sdata[tid];
298         }
299         if (tid < block_stride && i < out_ind.size()) {
300             out_ind.set(i);
301             ${postmap_name}(s, out_arr[out_ind.get()]);
302         }
303         __syncthreads();
304     }
305 }'''  # NOQA
306         ).substitute(
307             name=self.name,
308             op_name=op_name,
309             postmap_name=postmap_name,
310             in_type=_dtype_to_ctype[in_param.dtype],
311             out_type=_dtype_to_ctype[out_param.dtype],
312             reduce_ctype=self.reduce_ctype,
313             reduce_expr=self.expr,
314             identity=self.identity,
315             postmap_cast=self.postmap_cast_code)
316 
317         return [code]
318
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/cupy/core/_fusion_op.py b/cupy/core/_fusion_op.py
--- a/cupy/core/_fusion_op.py
+++ b/cupy/core/_fusion_op.py
@@ -205,7 +205,7 @@
         assert all([0 <= x < in_param.ndim for x in axis])
 
         self.name = name
-        self.preamble = reduce_func._preamble
+        self.preamble = reduce_func.preamble
         self.in_params = _VariableSet(in_param)
         self.out_params = _VariableSet(out_param)
         self.block_stride_name = 'block_stride_' + name

{"golden_diff": "diff --git a/cupy/core/_fusion_op.py b/cupy/core/_fusion_op.py\n--- a/cupy/core/_fusion_op.py\n+++ b/cupy/core/_fusion_op.py\n@@ -205,7 +205,7 @@\n         assert all([0 <= x < in_param.ndim for x in axis])\n \n         self.name = name\n-        self.preamble = reduce_func._preamble\n+        self.preamble = reduce_func.preamble\n         self.in_params = _VariableSet(in_param)\n         self.out_params = _VariableSet(out_param)\n         self.block_stride_name = 'block_stride_' + name\n", "issue": "Enhancement: enable CUB backend for ReductionKernel\nCurrently the CUB reduction is disabled for `ReductionKernel` (by the last `False` argument):\r\nhttps://github.com/cupy/cupy/blob/c1bf1aab6eb79a1cfcd5bb374460082d8eb2f3ca/cupy/core/_reduction.pyx#L724-L727\r\nBut we should revisit this at some point.\n", "before_files": [{"content": "import string\n\nimport numpy\n\nfrom cupy.core._fusion_variable import _TraceVariable\nfrom cupy.core._fusion_variable import _TraceArray\nfrom cupy.core._fusion_variable import _VariableSet\nfrom cupy.core import _fusion_thread_local\nfrom cupy.core import _fusion_emit_code\nfrom cupy.core import _kernel\nfrom cupy.core import _reduction\n\n\n_dtype_to_ctype = _fusion_emit_code._dtype_to_ctype\n\n\nclass _UfuncRoutine:\n    \"\"\"A device function for single elementwise operations.\n    \"\"\"\n\n    def __init__(\n            self, name, ufunc, routine_code, in_params, out_params,\n            compute_dtypes):\n        assert isinstance(name, str)\n        assert isinstance(ufunc, _kernel.ufunc)\n        assert isinstance(routine_code, str)\n        assert isinstance(compute_dtypes, tuple)\n        assert all(isinstance(t, numpy.dtype) for t in compute_dtypes)\n        assert isinstance(in_params, list)\n        assert all(isinstance(p, _TraceVariable) for p in in_params)\n        assert isinstance(out_params, list)\n        assert all(isinstance(p, _TraceArray) for p in out_params)\n\n        self.name = name\n        self.in_params = in_params\n        self.out_params = out_params\n        self.preamble = ufunc._preamble\n        self.routine_code = routine_code\n        self.compute_dtypes = compute_dtypes\n\n    def emit_code(self):\n        \"\"\"Returns a CUDA device function code.\n\n        Returns a string like:\n        ```\n        __device__ void cupy_add_0(int &in0_, float &in1_, double &out0_) {\n            typedef double in0_type;\n            typedef double in1_type;\n            typedef double out0_type;\n            double in0 = (double) in0_;\n            double in1 = (double) in1_;\n            double out0 = (double) out0_;\n            out0 = in0 + in1;\n            out0_ = out0;\n        }\n        ```\n        \"\"\"\n        nin = len(self.in_params)\n        dtypes = self.compute_dtypes\n        assert len(self.in_params) == len(self.compute_dtypes[:nin])\n        in_params = [\n            (_dtype_to_ctype[p.dtype], _dtype_to_ctype[t], 'in{}'.format(i))\n            for i, (p, t) in enumerate(zip(self.in_params, dtypes[:nin]))\n        ]\n        out_params = [\n            (_dtype_to_ctype[p.dtype], _dtype_to_ctype[t], 'out{}'.format(i))\n            for i, (p, t) in enumerate(zip(self.out_params, dtypes[nin:]))\n        ]\n        params = in_params + out_params\n\n        params_code = ', '.join(['{} &{}_'.format(t, s) for t, _, s in params])\n        typedef = ['typedef {} {}_type;'.format(t, s) for _, t, s in params]\n        read = ['{} {} = ({}) {}_;'.format(t, s, t, s) for _, t, s in params]\n        write = ['{}_ = {};'.format(s, s, s) for _, _, s in out_params]\n\n        return _fusion_emit_code._CodeBlock(\n            '__device__ void {}({})'.format(self.name, params_code),\n            typedef + read + [self.routine_code + ';'] + write)\n\n    def emit_call_code(self):\n        params = self.in_params + self.out_params\n        return '{op_name}({params});'.format(\n            op_name=self.name,\n            params=', '.join([var.lvar_name for var in params]))\n\n\nclass _ElementwiseTraceOp:\n    \"\"\"Ufunc or elementwise kernel with types.\n    \"\"\"\n\n    def __init__(self, ufunc_routines, in_params, out_params, ashape):\n        # The `in_params` and `out_params` should be already broadcasted to\n        # `ashape`, but they don't guarantee to be exactly same as\n        # `param.ashape`.\n\n        _fusion_thread_local.check_not_runtime()\n        assert isinstance(ufunc_routines, list)\n        assert all(isinstance(r, _UfuncRoutine) for r in ufunc_routines)\n        assert isinstance(ashape, tuple)\n\n        self.ops = ufunc_routines\n        self.in_params = _VariableSet(*in_params)\n        self.out_params = _VariableSet(*out_params)\n        self.ashape = ashape\n\n    @property\n    def params(self):\n        \"\"\"Returns the set of all variable the loop uses.\n        \"\"\"\n        res = _VariableSet()\n        for op in self.ops:\n            res += _VariableSet(*op.in_params)\n            res += _VariableSet(*op.out_params)\n        return res\n\n    @staticmethod\n    def _emit_declaration(params, in_params):\n        \"\"\"Returns a tuple of size 2.\n\n        1. CUDA code: declaring local variables.\n            2. The set of arrays which require indexer.\n        \"\"\"\n        _fusion_thread_local.check_not_runtime()\n\n        indexed_arrays = _VariableSet()\n        code = []\n        for var in params:\n            if var in in_params:\n                if isinstance(var, _TraceArray):\n                    indexed_arrays.add(var)\n                    f = '${type} ${lvar} = ${var}[${indexer}.get()];'\n                else:\n                    f = '${type} ${lvar} = ${var};'\n            else:\n                f = '${type} ${lvar};'\n            code.append(var.format(f))\n\n        return code, indexed_arrays\n\n    @staticmethod\n    def _emit_after_operation(out_params):\n        \"\"\"Returns a tuple of size 2.\n        1. CUDA code: writing the results of operations back to global memory.\n        2. The set of arrays which require indexer.\n        \"\"\"\n\n        _fusion_thread_local.check_not_runtime()\n\n        indexed_arrays = _VariableSet()\n        codes = []\n        for var in out_params:\n            if isinstance(var, _TraceArray):\n                indexed_arrays.add(var)\n                f = '${var}[${indexer}.get()] = ${lvar};'\n            else:\n                f = '${var} = ${lvar};'\n            codes.append(var.format(f))\n\n        return codes, indexed_arrays\n\n    @staticmethod\n    def _emit_set_index(indexed_params, tid):\n        \"\"\"Returns a CUDA code: setting a raw index to indexers.\n        \"\"\"\n        _fusion_thread_local.check_not_runtime()\n        assert isinstance(indexed_params, _VariableSet)\n\n        return [\n            p.format('${indexer}.set(${tid});', tid=tid)\n            for p in indexed_params\n        ]\n\n    def emit_code(self):\n        _fusion_thread_local.check_not_runtime()\n\n        declaration, s1 = self._emit_declaration(self.params, self.in_params)\n        operation = [op.emit_call_code() for op in self.ops]\n        after_operation, s2 = self._emit_after_operation(self.out_params)\n        index_name = 'i'\n        indexed_array = s1 + s2\n        indexer_name = next(iter(indexed_array)).indexer_name\n        indexer_setup = self._emit_set_index(indexed_array, index_name)\n\n        return _fusion_emit_code._CodeBlock(\n            'CUPY_FOR({}, {}.size())'.format(index_name, indexer_name),\n            indexer_setup + declaration + operation + after_operation)\n\n    def emit_preamble_codes(self):\n        return [subm.preamble for subm in self.ops if subm.preamble != '']\n\n    def emit_submodule_codes(self):\n        return [str(subm.emit_code()) for subm in self.ops]\n\n\nclass _ReductionTraceOp:\n    def __init__(self, name, reduce_func, expr, in_param, out_param, axis):\n        \"\"\"Reduction operation.\n        \"\"\"\n        _fusion_thread_local.check_not_runtime()\n        assert isinstance(name, str)\n        assert isinstance(reduce_func, _reduction._SimpleReductionKernel)\n        assert isinstance(in_param, _TraceArray)\n        assert isinstance(out_param, _TraceArray)\n        assert isinstance(axis, tuple)\n        assert all([0 <= x < in_param.ndim for x in axis])\n\n        self.name = name\n        self.preamble = reduce_func._preamble\n        self.in_params = _VariableSet(in_param)\n        self.out_params = _VariableSet(out_param)\n        self.block_stride_name = 'block_stride_' + name\n        self.axis = axis\n\n        if reduce_func.identity is None:\n            self.identity = ''\n        else:\n            self.identity = str(reduce_func.identity)\n\n        _, self.expr, self.postmap_cast_code, self.reduce_ctype = expr\n        if self.reduce_ctype is None:\n            out_param, = self.out_params\n            self.reduce_ctype = _dtype_to_ctype[out_param.dtype]\n\n        self.premap_op = None\n        self.postmap_op = None\n\n    @property\n    def params(self):\n        return self.in_params + self.out_params\n\n    def emit_code(self):\n        _fusion_thread_local.check_not_runtime()\n        assert len(self.in_params) == 1\n        assert len(self.out_params) == 1\n        in_param = list(self.in_params)[0]\n        out_param = list(self.out_params)[0]\n        params = ', '.join([\n            in_param.var_name,\n            out_param.var_name,\n            in_param.indexer_name,\n            out_param.indexer_name,\n        ])\n        return '{}({}, {});'.format(\n            self.name, params, self.block_stride_name)\n\n    def emit_preamble_codes(self):\n        preamble = self.preamble\n        return [preamble] if preamble != '' else []\n\n    def emit_submodule_codes(self):\n        \"\"\"Returns a CUDA device function code.\n\n        The emitted code assumes that ``block_stride`` and `blockDim.x` is a\n        power of 2.\n        \"\"\"\n\n        in_param, = self.in_params\n        out_param, = self.out_params\n        op_name = '{}_op'.format(self.name)\n        postmap_name = '{}_postmap'.format(self.name)\n\n        code = string.Template('''\n#define ${op_name}(a, b) (${reduce_expr})\n#define ${postmap_name}(a, out0) (${postmap_cast})\n\ntemplate <typename InType, typename OutType, typename InIndexerType, typename OutIndexerType>\n__device__ void ${name}(\n        InType in_arr, OutType out_arr,\n        InIndexerType in_ind, OutIndexerType out_ind, int block_stride) {\n    typedef ${in_type} type_in0_raw;\n    typedef ${out_type} type_out0_raw;\n    typedef ${reduce_ctype} _type_reduce;\n    extern __shared__ char _sdata_raw[];\n    _type_reduce *sdata = reinterpret_cast<_type_reduce*>(_sdata_raw);\n    unsigned int tid = threadIdx.x;\n    int _J = tid >> __popc(block_stride - 1);\n    ptrdiff_t _j = (ptrdiff_t)_J * out_ind.size();\n    int J_stride = blockDim.x >> __popc(block_stride - 1);\n    ptrdiff_t j_stride = (ptrdiff_t)J_stride * out_ind.size();\n\n    for (ptrdiff_t _i = (ptrdiff_t)blockIdx.x * block_stride; _i < out_ind.size(); _i += (ptrdiff_t)gridDim.x * block_stride) {\n        _type_reduce s = _type_reduce(${identity});\n        ptrdiff_t i = _i + (tid & (block_stride - 1));\n        for (ptrdiff_t j = i + _j; j < in_ind.size(); j += j_stride) {\n            in_ind.set(j);\n            s = ${op_name}(s, static_cast<_type_reduce>(in_arr[in_ind.get()]));\n        }\n        sdata[tid] = s;\n        __syncthreads();\n        for (unsigned int block = blockDim.x / 2; block >= block_stride; block >>= 1) {\n            if (tid < block) {\n                sdata[tid] = ${op_name}(sdata[tid], sdata[tid + block]);\n            }\n            __syncthreads();\n        }\n        if (tid < block_stride) {\n            s = sdata[tid];\n        }\n        if (tid < block_stride && i < out_ind.size()) {\n            out_ind.set(i);\n            ${postmap_name}(s, out_arr[out_ind.get()]);\n        }\n        __syncthreads();\n    }\n}'''  # NOQA\n        ).substitute(\n            name=self.name,\n            op_name=op_name,\n            postmap_name=postmap_name,\n            in_type=_dtype_to_ctype[in_param.dtype],\n            out_type=_dtype_to_ctype[out_param.dtype],\n            reduce_ctype=self.reduce_ctype,\n            reduce_expr=self.expr,\n            identity=self.identity,\n            postmap_cast=self.postmap_cast_code)\n\n        return [code]\n", "path": "cupy/core/_fusion_op.py"}], "after_files": [{"content": "import string\n\nimport numpy\n\nfrom cupy.core._fusion_variable import _TraceVariable\nfrom cupy.core._fusion_variable import _TraceArray\nfrom cupy.core._fusion_variable import _VariableSet\nfrom cupy.core import _fusion_thread_local\nfrom cupy.core import _fusion_emit_code\nfrom cupy.core import _kernel\nfrom cupy.core import _reduction\n\n\n_dtype_to_ctype = _fusion_emit_code._dtype_to_ctype\n\n\nclass _UfuncRoutine:\n    \"\"\"A device function for single elementwise operations.\n    \"\"\"\n\n    def __init__(\n            self, name, ufunc, routine_code, in_params, out_params,\n            compute_dtypes):\n        assert isinstance(name, str)\n        assert isinstance(ufunc, _kernel.ufunc)\n        assert isinstance(routine_code, str)\n        assert isinstance(compute_dtypes, tuple)\n        assert all(isinstance(t, numpy.dtype) for t in compute_dtypes)\n        assert isinstance(in_params, list)\n        assert all(isinstance(p, _TraceVariable) for p in in_params)\n        assert isinstance(out_params, list)\n        assert all(isinstance(p, _TraceArray) for p in out_params)\n\n        self.name = name\n        self.in_params = in_params\n        self.out_params = out_params\n        self.preamble = ufunc._preamble\n        self.routine_code = routine_code\n        self.compute_dtypes = compute_dtypes\n\n    def emit_code(self):\n        \"\"\"Returns a CUDA device function code.\n\n        Returns a string like:\n        ```\n        __device__ void cupy_add_0(int &in0_, float &in1_, double &out0_) {\n            typedef double in0_type;\n            typedef double in1_type;\n            typedef double out0_type;\n            double in0 = (double) in0_;\n            double in1 = (double) in1_;\n            double out0 = (double) out0_;\n            out0 = in0 + in1;\n            out0_ = out0;\n        }\n        ```\n        \"\"\"\n        nin = len(self.in_params)\n        dtypes = self.compute_dtypes\n        assert len(self.in_params) == len(self.compute_dtypes[:nin])\n        in_params = [\n            (_dtype_to_ctype[p.dtype], _dtype_to_ctype[t], 'in{}'.format(i))\n            for i, (p, t) in enumerate(zip(self.in_params, dtypes[:nin]))\n        ]\n        out_params = [\n            (_dtype_to_ctype[p.dtype], _dtype_to_ctype[t], 'out{}'.format(i))\n            for i, (p, t) in enumerate(zip(self.out_params, dtypes[nin:]))\n        ]\n        params = in_params + out_params\n\n        params_code = ', '.join(['{} &{}_'.format(t, s) for t, _, s in params])\n        typedef = ['typedef {} {}_type;'.format(t, s) for _, t, s in params]\n        read = ['{} {} = ({}) {}_;'.format(t, s, t, s) for _, t, s in params]\n        write = ['{}_ = {};'.format(s, s, s) for _, _, s in out_params]\n\n        return _fusion_emit_code._CodeBlock(\n            '__device__ void {}({})'.format(self.name, params_code),\n            typedef + read + [self.routine_code + ';'] + write)\n\n    def emit_call_code(self):\n        params = self.in_params + self.out_params\n        return '{op_name}({params});'.format(\n            op_name=self.name,\n            params=', '.join([var.lvar_name for var in params]))\n\n\nclass _ElementwiseTraceOp:\n    \"\"\"Ufunc or elementwise kernel with types.\n    \"\"\"\n\n    def __init__(self, ufunc_routines, in_params, out_params, ashape):\n        # The `in_params` and `out_params` should be already broadcasted to\n        # `ashape`, but they don't guarantee to be exactly same as\n        # `param.ashape`.\n\n        _fusion_thread_local.check_not_runtime()\n        assert isinstance(ufunc_routines, list)\n        assert all(isinstance(r, _UfuncRoutine) for r in ufunc_routines)\n        assert isinstance(ashape, tuple)\n\n        self.ops = ufunc_routines\n        self.in_params = _VariableSet(*in_params)\n        self.out_params = _VariableSet(*out_params)\n        self.ashape = ashape\n\n    @property\n    def params(self):\n        \"\"\"Returns the set of all variable the loop uses.\n        \"\"\"\n        res = _VariableSet()\n        for op in self.ops:\n            res += _VariableSet(*op.in_params)\n            res += _VariableSet(*op.out_params)\n        return res\n\n    @staticmethod\n    def _emit_declaration(params, in_params):\n        \"\"\"Returns a tuple of size 2.\n\n        1. CUDA code: declaring local variables.\n            2. The set of arrays which require indexer.\n        \"\"\"\n        _fusion_thread_local.check_not_runtime()\n\n        indexed_arrays = _VariableSet()\n        code = []\n        for var in params:\n            if var in in_params:\n                if isinstance(var, _TraceArray):\n                    indexed_arrays.add(var)\n                    f = '${type} ${lvar} = ${var}[${indexer}.get()];'\n                else:\n                    f = '${type} ${lvar} = ${var};'\n            else:\n                f = '${type} ${lvar};'\n            code.append(var.format(f))\n\n        return code, indexed_arrays\n\n    @staticmethod\n    def _emit_after_operation(out_params):\n        \"\"\"Returns a tuple of size 2.\n        1. CUDA code: writing the results of operations back to global memory.\n        2. The set of arrays which require indexer.\n        \"\"\"\n\n        _fusion_thread_local.check_not_runtime()\n\n        indexed_arrays = _VariableSet()\n        codes = []\n        for var in out_params:\n            if isinstance(var, _TraceArray):\n                indexed_arrays.add(var)\n                f = '${var}[${indexer}.get()] = ${lvar};'\n            else:\n                f = '${var} = ${lvar};'\n            codes.append(var.format(f))\n\n        return codes, indexed_arrays\n\n    @staticmethod\n    def _emit_set_index(indexed_params, tid):\n        \"\"\"Returns a CUDA code: setting a raw index to indexers.\n        \"\"\"\n        _fusion_thread_local.check_not_runtime()\n        assert isinstance(indexed_params, _VariableSet)\n\n        return [\n            p.format('${indexer}.set(${tid});', tid=tid)\n            for p in indexed_params\n        ]\n\n    def emit_code(self):\n        _fusion_thread_local.check_not_runtime()\n\n        declaration, s1 = self._emit_declaration(self.params, self.in_params)\n        operation = [op.emit_call_code() for op in self.ops]\n        after_operation, s2 = self._emit_after_operation(self.out_params)\n        index_name = 'i'\n        indexed_array = s1 + s2\n        indexer_name = next(iter(indexed_array)).indexer_name\n        indexer_setup = self._emit_set_index(indexed_array, index_name)\n\n        return _fusion_emit_code._CodeBlock(\n            'CUPY_FOR({}, {}.size())'.format(index_name, indexer_name),\n            indexer_setup + declaration + operation + after_operation)\n\n    def emit_preamble_codes(self):\n        return [subm.preamble for subm in self.ops if subm.preamble != '']\n\n    def emit_submodule_codes(self):\n        return [str(subm.emit_code()) for subm in self.ops]\n\n\nclass _ReductionTraceOp:\n    def __init__(self, name, reduce_func, expr, in_param, out_param, axis):\n        \"\"\"Reduction operation.\n        \"\"\"\n        _fusion_thread_local.check_not_runtime()\n        assert isinstance(name, str)\n        assert isinstance(reduce_func, _reduction._SimpleReductionKernel)\n        assert isinstance(in_param, _TraceArray)\n        assert isinstance(out_param, _TraceArray)\n        assert isinstance(axis, tuple)\n        assert all([0 <= x < in_param.ndim for x in axis])\n\n        self.name = name\n        self.preamble = reduce_func.preamble\n        self.in_params = _VariableSet(in_param)\n        self.out_params = _VariableSet(out_param)\n        self.block_stride_name = 'block_stride_' + name\n        self.axis = axis\n\n        if reduce_func.identity is None:\n            self.identity = ''\n        else:\n            self.identity = str(reduce_func.identity)\n\n        _, self.expr, self.postmap_cast_code, self.reduce_ctype = expr\n        if self.reduce_ctype is None:\n            out_param, = self.out_params\n            self.reduce_ctype = _dtype_to_ctype[out_param.dtype]\n\n        self.premap_op = None\n        self.postmap_op = None\n\n    @property\n    def params(self):\n        return self.in_params + self.out_params\n\n    def emit_code(self):\n        _fusion_thread_local.check_not_runtime()\n        assert len(self.in_params) == 1\n        assert len(self.out_params) == 1\n        in_param = list(self.in_params)[0]\n        out_param = list(self.out_params)[0]\n        params = ', '.join([\n            in_param.var_name,\n            out_param.var_name,\n            in_param.indexer_name,\n            out_param.indexer_name,\n        ])\n        return '{}({}, {});'.format(\n            self.name, params, self.block_stride_name)\n\n    def emit_preamble_codes(self):\n        preamble = self.preamble\n        return [preamble] if preamble != '' else []\n\n    def emit_submodule_codes(self):\n        \"\"\"Returns a CUDA device function code.\n\n        The emitted code assumes that ``block_stride`` and `blockDim.x` is a\n        power of 2.\n        \"\"\"\n\n        in_param, = self.in_params\n        out_param, = self.out_params\n        op_name = '{}_op'.format(self.name)\n        postmap_name = '{}_postmap'.format(self.name)\n\n        code = string.Template('''\n#define ${op_name}(a, b) (${reduce_expr})\n#define ${postmap_name}(a, out0) (${postmap_cast})\n\ntemplate <typename InType, typename OutType, typename InIndexerType, typename OutIndexerType>\n__device__ void ${name}(\n        InType in_arr, OutType out_arr,\n        InIndexerType in_ind, OutIndexerType out_ind, int block_stride) {\n    typedef ${in_type} type_in0_raw;\n    typedef ${out_type} type_out0_raw;\n    typedef ${reduce_ctype} _type_reduce;\n    extern __shared__ char _sdata_raw[];\n    _type_reduce *sdata = reinterpret_cast<_type_reduce*>(_sdata_raw);\n    unsigned int tid = threadIdx.x;\n    int _J = tid >> __popc(block_stride - 1);\n    ptrdiff_t _j = (ptrdiff_t)_J * out_ind.size();\n    int J_stride = blockDim.x >> __popc(block_stride - 1);\n    ptrdiff_t j_stride = (ptrdiff_t)J_stride * out_ind.size();\n\n    for (ptrdiff_t _i = (ptrdiff_t)blockIdx.x * block_stride; _i < out_ind.size(); _i += (ptrdiff_t)gridDim.x * block_stride) {\n        _type_reduce s = _type_reduce(${identity});\n        ptrdiff_t i = _i + (tid & (block_stride - 1));\n        for (ptrdiff_t j = i + _j; j < in_ind.size(); j += j_stride) {\n            in_ind.set(j);\n            s = ${op_name}(s, static_cast<_type_reduce>(in_arr[in_ind.get()]));\n        }\n        sdata[tid] = s;\n        __syncthreads();\n        for (unsigned int block = blockDim.x / 2; block >= block_stride; block >>= 1) {\n            if (tid < block) {\n                sdata[tid] = ${op_name}(sdata[tid], sdata[tid + block]);\n            }\n            __syncthreads();\n        }\n        if (tid < block_stride) {\n            s = sdata[tid];\n        }\n        if (tid < block_stride && i < out_ind.size()) {\n            out_ind.set(i);\n            ${postmap_name}(s, out_arr[out_ind.get()]);\n        }\n        __syncthreads();\n    }\n}'''  # NOQA\n        ).substitute(\n            name=self.name,\n            op_name=op_name,\n            postmap_name=postmap_name,\n            in_type=_dtype_to_ctype[in_param.dtype],\n            out_type=_dtype_to_ctype[out_param.dtype],\n            reduce_ctype=self.reduce_ctype,\n            reduce_expr=self.expr,\n            identity=self.identity,\n            postmap_cast=self.postmap_cast_code)\n\n        return [code]\n", "path": "cupy/core/_fusion_op.py"}]}