We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Use `warpSize` to get the warp size in ROCm environment
CuPy, currently, has a hard-coded warp size of 64 for ROCm environment. However, AMD gpu's warp size is not necessarily 64; for example, Navi21 that ROCm 5.0 supports has the warp size of 32. The following can be a good solution.

https://github.com/pytorch/pytorch/blob/2cefbb71cf65696e28ee1bdfce06d6846285e5fb/c10/macros/Macros.h#L322

Thanks @amathews-amd for pointing this out.

--- 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/__init__.py`
Content:
```
1 # mypy: ignore-errors
2 
3 from cupy._core import core  # NOQA
4 from cupy._core import fusion  # NOQA
5 from cupy._core import internal  # NOQA
6 
7 
8 # internal APIs for testing and developement
9 from cupy._core._accelerator import set_elementwise_accelerators  # NOQA
10 from cupy._core._accelerator import set_reduction_accelerators  # NOQA
11 from cupy._core._accelerator import set_routine_accelerators  # NOQA
12 from cupy._core._accelerator import get_elementwise_accelerators  # NOQA
13 from cupy._core._accelerator import get_reduction_accelerators  # NOQA
14 from cupy._core._accelerator import get_routine_accelerators  # NOQA
15 
16 
17 # import class and function
18 from cupy._core._kernel import create_ufunc  # NOQA
19 from cupy._core._kernel import ElementwiseKernel  # NOQA
20 from cupy._core._kernel import ufunc  # NOQA
21 from cupy._core._reduction import create_reduction_func  # NOQA
22 from cupy._core._reduction import ReductionKernel  # NOQA
23 from cupy._core._routines_binary import bitwise_and  # NOQA
24 from cupy._core._routines_binary import bitwise_or  # NOQA
25 from cupy._core._routines_binary import bitwise_xor  # NOQA
26 from cupy._core._routines_binary import invert  # NOQA
27 from cupy._core._routines_binary import left_shift  # NOQA
28 from cupy._core._routines_binary import right_shift  # NOQA
29 from cupy._core._routines_linalg import _mat_ptrs  # NOQA
30 from cupy._core._routines_linalg import dot  # NOQA
31 from cupy._core._routines_linalg import get_compute_type  # NOQA
32 from cupy._core._routines_linalg import matmul  # NOQA
33 from cupy._core._routines_linalg import set_compute_type  # NOQA
34 from cupy._core._routines_linalg import tensordot_core  # NOQA
35 from cupy._core._routines_logic import create_comparison  # NOQA
36 from cupy._core._routines_logic import equal  # NOQA
37 from cupy._core._routines_logic import greater  # NOQA
38 from cupy._core._routines_logic import greater_equal  # NOQA
39 from cupy._core._routines_logic import less  # NOQA
40 from cupy._core._routines_logic import less_equal  # NOQA
41 from cupy._core._routines_logic import not_equal  # NOQA
42 from cupy._core._routines_manipulation import array_split  # NOQA
43 from cupy._core._routines_manipulation import broadcast  # NOQA
44 from cupy._core._routines_manipulation import broadcast_to  # NOQA
45 from cupy._core._routines_manipulation import concatenate_method  # NOQA
46 from cupy._core._routines_manipulation import moveaxis  # NOQA
47 from cupy._core._routines_manipulation import rollaxis  # NOQA
48 from cupy._core._routines_manipulation import size  # NOQA'
49 from cupy._core._routines_math import absolute  # NOQA
50 from cupy._core._routines_math import add  # NOQA
51 from cupy._core._routines_math import angle, angle_deg  # NOQA
52 from cupy._core._routines_math import conjugate  # NOQA
53 from cupy._core._routines_math import divide  # NOQA
54 from cupy._core._routines_math import floor_divide  # NOQA
55 from cupy._core._routines_math import multiply  # NOQA
56 from cupy._core._routines_math import negative  # NOQA
57 from cupy._core._routines_math import positive  # NOQA
58 from cupy._core._routines_math import power  # NOQA
59 from cupy._core._routines_math import remainder  # NOQA
60 from cupy._core._routines_math import sqrt  # NOQA
61 from cupy._core._routines_math import subtract  # NOQA
62 from cupy._core._routines_math import true_divide  # NOQA
63 from cupy._core._routines_statistics import nanmax  # NOQA
64 from cupy._core._routines_statistics import nanmin  # NOQA
65 from cupy._core.core import _internal_ascontiguousarray  # NOQA
66 from cupy._core.core import _internal_asfortranarray  # NOQA
67 from cupy._core.core import array  # NOQA
68 from cupy._core.core import ascontiguousarray  # NOQA
69 from cupy._core.core import asfortranarray  # NOQA
70 from cupy._core.core import divmod  # NOQA
71 from cupy._core.core import elementwise_copy  # NOQA
72 from cupy._core.core import ndarray  # NOQA
73 from cupy._core.dlpack import fromDlpack  # NOQA
74 from cupy._core.dlpack import from_dlpack  # NOQA
75 from cupy._core.internal import complete_slice  # NOQA
76 from cupy._core.internal import get_size  # NOQA
77 from cupy._core.raw import RawKernel  # NOQA
78 from cupy._core.raw import RawModule  # NOQA
79
```

Path: `cupyx/jit/_interface.py`
Content:
```
1 import functools
2 import warnings
3 
4 import numpy
5 
6 from cupy_backends.cuda.api import runtime
7 import cupy
8 from cupy._core import core
9 from cupyx.jit import _compile
10 from cupyx.jit import _cuda_typerules
11 from cupyx.jit import _cuda_types
12 from cupyx.jit import _internal_types
13 from cupyx.jit._cuda_types import Scalar
14 
15 
16 class _CudaFunction:
17     """JIT cupy function object
18     """
19 
20     def __init__(self, func, mode, device=False, inline=False):
21         self.attributes = []
22 
23         if device:
24             self.attributes.append('__device__')
25         else:
26             self.attributes.append('__global__')
27 
28         if inline:
29             self.attributes.append('inline')
30 
31         self.name = getattr(func, 'name', func.__name__)
32         self.func = func
33         self.mode = mode
34 
35     def __call__(self, *args, **kwargs):
36         raise NotImplementedError
37 
38     def _emit_code_from_types(self, in_types, ret_type=None):
39         return _compile.transpile(
40             self.func, self.attributes, self.mode, in_types, ret_type)
41 
42 
43 class _JitRawKernel:
44     """JIT CUDA kernel object.
45 
46     The decorator :func:``cupyx.jit.rawkernel`` converts the target function
47     to an object of this class. This class is not inteded to be instantiated
48     by users.
49     """
50 
51     def __init__(self, func, mode, device):
52         self._func = func
53         self._mode = mode
54         self._device = device
55         self._cache = {}
56         self._cached_codes = {}
57 
58     def __call__(
59             self, grid, block, args, shared_mem=0, stream=None):
60         """Calls the CUDA kernel.
61 
62         The compilation will be deferred until the first function call.
63         CuPy's JIT compiler infers the types of arguments at the call
64         time, and will cache the compiled kernels for speeding up any
65         subsequent calls.
66 
67         Args:
68             grid (tuple of int): Size of grid in blocks.
69             block (tuple of int): Dimensions of each thread block.
70             args (tuple):
71                 Arguments of the kernel. The type of all elements must be
72                 ``bool``, ``int``, ``float``, ``complex``, NumPy scalar or
73                 ``cupy.ndarray``.
74             shared_mem (int):
75                 Dynamic shared-memory size per thread block in bytes.
76             stream (cupy.cuda.Stream): CUDA stream.
77 
78         .. seealso:: :ref:`jit_kernel_definition`
79         """
80         in_types = []
81         for x in args:
82             if isinstance(x, cupy.ndarray):
83                 t = _cuda_types.CArray.from_ndarray(x)
84             elif numpy.isscalar(x):
85                 t = _cuda_typerules.get_ctype_from_scalar(self._mode, x)
86             else:
87                 raise TypeError(f'{type(x)} is not supported for RawKernel')
88             in_types.append(t)
89         in_types = tuple(in_types)
90         device_id = cupy.cuda.get_device_id()
91 
92         kern, enable_cg = self._cache.get((in_types, device_id), (None, None))
93         if kern is None:
94             result = self._cached_codes.get(in_types)
95             if result is None:
96                 result = _compile.transpile(
97                     self._func,
98                     ['extern "C"', '__global__'],
99                     self._mode,
100                     in_types,
101                     _cuda_types.void,
102                 )
103                 self._cached_codes[in_types] = result
104 
105             fname = result.func_name
106             enable_cg = result.enable_cooperative_groups
107             options = ('-DCUPY_JIT_MODE', '--std=c++14')
108             backend = result.backend
109             if backend == 'nvcc':
110                 options += ('-DCUPY_JIT_NVCC',)
111             module = core.compile_with_cache(
112                 source=result.code,
113                 options=options,
114                 backend=backend)
115             kern = module.get_function(fname)
116             self._cache[(in_types, device_id)] = (kern, enable_cg)
117 
118         new_args = []
119         for a, t in zip(args, in_types):
120             if isinstance(t, Scalar):
121                 if t.dtype.char == 'e':
122                     a = numpy.float32(a)
123                 else:
124                     a = t.dtype.type(a)
125             new_args.append(a)
126 
127         kern(grid, block, tuple(new_args), shared_mem, stream, enable_cg)
128 
129     def __getitem__(self, grid_and_block):
130         """Numba-style kernel call.
131 
132         .. seealso:: :ref:`jit_kernel_definition`
133         """
134         grid, block = grid_and_block
135         if not isinstance(grid, tuple):
136             grid = (grid, 1, 1)
137         if not isinstance(block, tuple):
138             block = (block, 1, 1)
139         return lambda *args, **kwargs: self(grid, block, args, **kwargs)
140 
141     @property
142     def cached_codes(self):
143         """Returns a dict that has input types as keys and codes values.
144 
145         This proprety method is for debugging purpose.
146         The return value is not guaranteed to keep backward compatibility.
147         """
148         if len(self._cached_codes) == 0:
149             warnings.warn(
150                 'No codes are cached because compilation is deferred until '
151                 'the first function call.')
152         return dict([(k, v.code) for k, v in self._cached_codes.items()])
153 
154     @property
155     def cached_code(self):
156         """Returns `next(iter(self.cached_codes.values()))`.
157 
158         This proprety method is for debugging purpose.
159         The return value is not guaranteed to keep backward compatibility.
160         """
161         codes = self.cached_codes
162         if len(codes) > 1:
163             warnings.warn(
164                 'The input types of the kernel could not be inferred. '
165                 'Please use `.cached_codes` instead.')
166         return next(iter(codes.values()))
167 
168 
169 def rawkernel(*, mode='cuda', device=False):
170     """A decorator compiles a Python function into CUDA kernel.
171     """
172     cupy._util.experimental('cupyx.jit.rawkernel')
173 
174     def wrapper(func):
175         return functools.update_wrapper(
176             _JitRawKernel(func, mode, device), func)
177     return wrapper
178 
179 
180 threadIdx = _internal_types.Data('threadIdx', _cuda_types.dim3)
181 blockDim = _internal_types.Data('blockDim', _cuda_types.dim3)
182 blockIdx = _internal_types.Data('blockIdx', _cuda_types.dim3)
183 gridDim = _internal_types.Data('gridDim', _cuda_types.dim3)
184 
185 warpsize = _internal_types.Data(
186     '64' if runtime.is_hip else '32', _cuda_types.uint32)
187 warpsize.__doc__ = r"""Returns the number of threads in a warp.
188 
189 In CUDA this is always 32, and in ROCm/HIP always 64.
190 
191 .. seealso:: :obj:`numba.cuda.warpsize`
192 """
193
```
--- 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/__init__.py b/cupy/_core/__init__.py
--- a/cupy/_core/__init__.py
+++ b/cupy/_core/__init__.py
@@ -18,6 +18,7 @@
 from cupy._core._kernel import create_ufunc  # NOQA
 from cupy._core._kernel import ElementwiseKernel  # NOQA
 from cupy._core._kernel import ufunc  # NOQA
+from cupy._core._kernel import _get_warpsize  # NOQA
 from cupy._core._reduction import create_reduction_func  # NOQA
 from cupy._core._reduction import ReductionKernel  # NOQA
 from cupy._core._routines_binary import bitwise_and  # NOQA
diff --git a/cupyx/jit/_interface.py b/cupyx/jit/_interface.py
--- a/cupyx/jit/_interface.py
+++ b/cupyx/jit/_interface.py
@@ -3,7 +3,6 @@
 
 import numpy
 
-from cupy_backends.cuda.api import runtime
 import cupy
 from cupy._core import core
 from cupyx.jit import _compile
@@ -182,11 +181,8 @@
 blockIdx = _internal_types.Data('blockIdx', _cuda_types.dim3)
 gridDim = _internal_types.Data('gridDim', _cuda_types.dim3)
 
-warpsize = _internal_types.Data(
-    '64' if runtime.is_hip else '32', _cuda_types.uint32)
+warpsize = _internal_types.Data('warpSize', _cuda_types.int32)
 warpsize.__doc__ = r"""Returns the number of threads in a warp.
 
-In CUDA this is always 32, and in ROCm/HIP always 64.
-
 .. seealso:: :obj:`numba.cuda.warpsize`
 """

{"golden_diff": "diff --git a/cupy/_core/__init__.py b/cupy/_core/__init__.py\n--- a/cupy/_core/__init__.py\n+++ b/cupy/_core/__init__.py\n@@ -18,6 +18,7 @@\n from cupy._core._kernel import create_ufunc  # NOQA\n from cupy._core._kernel import ElementwiseKernel  # NOQA\n from cupy._core._kernel import ufunc  # NOQA\n+from cupy._core._kernel import _get_warpsize  # NOQA\n from cupy._core._reduction import create_reduction_func  # NOQA\n from cupy._core._reduction import ReductionKernel  # NOQA\n from cupy._core._routines_binary import bitwise_and  # NOQA\ndiff --git a/cupyx/jit/_interface.py b/cupyx/jit/_interface.py\n--- a/cupyx/jit/_interface.py\n+++ b/cupyx/jit/_interface.py\n@@ -3,7 +3,6 @@\n \n import numpy\n \n-from cupy_backends.cuda.api import runtime\n import cupy\n from cupy._core import core\n from cupyx.jit import _compile\n@@ -182,11 +181,8 @@\n blockIdx = _internal_types.Data('blockIdx', _cuda_types.dim3)\n gridDim = _internal_types.Data('gridDim', _cuda_types.dim3)\n \n-warpsize = _internal_types.Data(\n-    '64' if runtime.is_hip else '32', _cuda_types.uint32)\n+warpsize = _internal_types.Data('warpSize', _cuda_types.int32)\n warpsize.__doc__ = r\"\"\"Returns the number of threads in a warp.\n \n-In CUDA this is always 32, and in ROCm/HIP always 64.\n-\n .. seealso:: :obj:`numba.cuda.warpsize`\n \"\"\"\n", "issue": "Use `warpSize` to get the warp size in ROCm environment\nCuPy, currently, has a hard-coded warp size of 64 for ROCm environment. However, AMD gpu's warp size is not necessarily 64; for example, Navi21 that ROCm 5.0 supports has the warp size of 32. The following can be a good solution.\r\n\r\nhttps://github.com/pytorch/pytorch/blob/2cefbb71cf65696e28ee1bdfce06d6846285e5fb/c10/macros/Macros.h#L322\r\n\r\nThanks @amathews-amd for pointing this out.\n", "before_files": [{"content": "# mypy: ignore-errors\n\nfrom cupy._core import core  # NOQA\nfrom cupy._core import fusion  # NOQA\nfrom cupy._core import internal  # NOQA\n\n\n# internal APIs for testing and developement\nfrom cupy._core._accelerator import set_elementwise_accelerators  # NOQA\nfrom cupy._core._accelerator import set_reduction_accelerators  # NOQA\nfrom cupy._core._accelerator import set_routine_accelerators  # NOQA\nfrom cupy._core._accelerator import get_elementwise_accelerators  # NOQA\nfrom cupy._core._accelerator import get_reduction_accelerators  # NOQA\nfrom cupy._core._accelerator import get_routine_accelerators  # NOQA\n\n\n# import class and function\nfrom cupy._core._kernel import create_ufunc  # NOQA\nfrom cupy._core._kernel import ElementwiseKernel  # NOQA\nfrom cupy._core._kernel import ufunc  # NOQA\nfrom cupy._core._reduction import create_reduction_func  # NOQA\nfrom cupy._core._reduction import ReductionKernel  # NOQA\nfrom cupy._core._routines_binary import bitwise_and  # NOQA\nfrom cupy._core._routines_binary import bitwise_or  # NOQA\nfrom cupy._core._routines_binary import bitwise_xor  # NOQA\nfrom cupy._core._routines_binary import invert  # NOQA\nfrom cupy._core._routines_binary import left_shift  # NOQA\nfrom cupy._core._routines_binary import right_shift  # NOQA\nfrom cupy._core._routines_linalg import _mat_ptrs  # NOQA\nfrom cupy._core._routines_linalg import dot  # NOQA\nfrom cupy._core._routines_linalg import get_compute_type  # NOQA\nfrom cupy._core._routines_linalg import matmul  # NOQA\nfrom cupy._core._routines_linalg import set_compute_type  # NOQA\nfrom cupy._core._routines_linalg import tensordot_core  # NOQA\nfrom cupy._core._routines_logic import create_comparison  # NOQA\nfrom cupy._core._routines_logic import equal  # NOQA\nfrom cupy._core._routines_logic import greater  # NOQA\nfrom cupy._core._routines_logic import greater_equal  # NOQA\nfrom cupy._core._routines_logic import less  # NOQA\nfrom cupy._core._routines_logic import less_equal  # NOQA\nfrom cupy._core._routines_logic import not_equal  # NOQA\nfrom cupy._core._routines_manipulation import array_split  # NOQA\nfrom cupy._core._routines_manipulation import broadcast  # NOQA\nfrom cupy._core._routines_manipulation import broadcast_to  # NOQA\nfrom cupy._core._routines_manipulation import concatenate_method  # NOQA\nfrom cupy._core._routines_manipulation import moveaxis  # NOQA\nfrom cupy._core._routines_manipulation import rollaxis  # NOQA\nfrom cupy._core._routines_manipulation import size  # NOQA'\nfrom cupy._core._routines_math import absolute  # NOQA\nfrom cupy._core._routines_math import add  # NOQA\nfrom cupy._core._routines_math import angle, angle_deg  # NOQA\nfrom cupy._core._routines_math import conjugate  # NOQA\nfrom cupy._core._routines_math import divide  # NOQA\nfrom cupy._core._routines_math import floor_divide  # NOQA\nfrom cupy._core._routines_math import multiply  # NOQA\nfrom cupy._core._routines_math import negative  # NOQA\nfrom cupy._core._routines_math import positive  # NOQA\nfrom cupy._core._routines_math import power  # NOQA\nfrom cupy._core._routines_math import remainder  # NOQA\nfrom cupy._core._routines_math import sqrt  # NOQA\nfrom cupy._core._routines_math import subtract  # NOQA\nfrom cupy._core._routines_math import true_divide  # NOQA\nfrom cupy._core._routines_statistics import nanmax  # NOQA\nfrom cupy._core._routines_statistics import nanmin  # NOQA\nfrom cupy._core.core import _internal_ascontiguousarray  # NOQA\nfrom cupy._core.core import _internal_asfortranarray  # NOQA\nfrom cupy._core.core import array  # NOQA\nfrom cupy._core.core import ascontiguousarray  # NOQA\nfrom cupy._core.core import asfortranarray  # NOQA\nfrom cupy._core.core import divmod  # NOQA\nfrom cupy._core.core import elementwise_copy  # NOQA\nfrom cupy._core.core import ndarray  # NOQA\nfrom cupy._core.dlpack import fromDlpack  # NOQA\nfrom cupy._core.dlpack import from_dlpack  # NOQA\nfrom cupy._core.internal import complete_slice  # NOQA\nfrom cupy._core.internal import get_size  # NOQA\nfrom cupy._core.raw import RawKernel  # NOQA\nfrom cupy._core.raw import RawModule  # NOQA\n", "path": "cupy/_core/__init__.py"}, {"content": "import functools\nimport warnings\n\nimport numpy\n\nfrom cupy_backends.cuda.api import runtime\nimport cupy\nfrom cupy._core import core\nfrom cupyx.jit import _compile\nfrom cupyx.jit import _cuda_typerules\nfrom cupyx.jit import _cuda_types\nfrom cupyx.jit import _internal_types\nfrom cupyx.jit._cuda_types import Scalar\n\n\nclass _CudaFunction:\n    \"\"\"JIT cupy function object\n    \"\"\"\n\n    def __init__(self, func, mode, device=False, inline=False):\n        self.attributes = []\n\n        if device:\n            self.attributes.append('__device__')\n        else:\n            self.attributes.append('__global__')\n\n        if inline:\n            self.attributes.append('inline')\n\n        self.name = getattr(func, 'name', func.__name__)\n        self.func = func\n        self.mode = mode\n\n    def __call__(self, *args, **kwargs):\n        raise NotImplementedError\n\n    def _emit_code_from_types(self, in_types, ret_type=None):\n        return _compile.transpile(\n            self.func, self.attributes, self.mode, in_types, ret_type)\n\n\nclass _JitRawKernel:\n    \"\"\"JIT CUDA kernel object.\n\n    The decorator :func:``cupyx.jit.rawkernel`` converts the target function\n    to an object of this class. This class is not inteded to be instantiated\n    by users.\n    \"\"\"\n\n    def __init__(self, func, mode, device):\n        self._func = func\n        self._mode = mode\n        self._device = device\n        self._cache = {}\n        self._cached_codes = {}\n\n    def __call__(\n            self, grid, block, args, shared_mem=0, stream=None):\n        \"\"\"Calls the CUDA kernel.\n\n        The compilation will be deferred until the first function call.\n        CuPy's JIT compiler infers the types of arguments at the call\n        time, and will cache the compiled kernels for speeding up any\n        subsequent calls.\n\n        Args:\n            grid (tuple of int): Size of grid in blocks.\n            block (tuple of int): Dimensions of each thread block.\n            args (tuple):\n                Arguments of the kernel. The type of all elements must be\n                ``bool``, ``int``, ``float``, ``complex``, NumPy scalar or\n                ``cupy.ndarray``.\n            shared_mem (int):\n                Dynamic shared-memory size per thread block in bytes.\n            stream (cupy.cuda.Stream): CUDA stream.\n\n        .. seealso:: :ref:`jit_kernel_definition`\n        \"\"\"\n        in_types = []\n        for x in args:\n            if isinstance(x, cupy.ndarray):\n                t = _cuda_types.CArray.from_ndarray(x)\n            elif numpy.isscalar(x):\n                t = _cuda_typerules.get_ctype_from_scalar(self._mode, x)\n            else:\n                raise TypeError(f'{type(x)} is not supported for RawKernel')\n            in_types.append(t)\n        in_types = tuple(in_types)\n        device_id = cupy.cuda.get_device_id()\n\n        kern, enable_cg = self._cache.get((in_types, device_id), (None, None))\n        if kern is None:\n            result = self._cached_codes.get(in_types)\n            if result is None:\n                result = _compile.transpile(\n                    self._func,\n                    ['extern \"C\"', '__global__'],\n                    self._mode,\n                    in_types,\n                    _cuda_types.void,\n                )\n                self._cached_codes[in_types] = result\n\n            fname = result.func_name\n            enable_cg = result.enable_cooperative_groups\n            options = ('-DCUPY_JIT_MODE', '--std=c++14')\n            backend = result.backend\n            if backend == 'nvcc':\n                options += ('-DCUPY_JIT_NVCC',)\n            module = core.compile_with_cache(\n                source=result.code,\n                options=options,\n                backend=backend)\n            kern = module.get_function(fname)\n            self._cache[(in_types, device_id)] = (kern, enable_cg)\n\n        new_args = []\n        for a, t in zip(args, in_types):\n            if isinstance(t, Scalar):\n                if t.dtype.char == 'e':\n                    a = numpy.float32(a)\n                else:\n                    a = t.dtype.type(a)\n            new_args.append(a)\n\n        kern(grid, block, tuple(new_args), shared_mem, stream, enable_cg)\n\n    def __getitem__(self, grid_and_block):\n        \"\"\"Numba-style kernel call.\n\n        .. seealso:: :ref:`jit_kernel_definition`\n        \"\"\"\n        grid, block = grid_and_block\n        if not isinstance(grid, tuple):\n            grid = (grid, 1, 1)\n        if not isinstance(block, tuple):\n            block = (block, 1, 1)\n        return lambda *args, **kwargs: self(grid, block, args, **kwargs)\n\n    @property\n    def cached_codes(self):\n        \"\"\"Returns a dict that has input types as keys and codes values.\n\n        This proprety method is for debugging purpose.\n        The return value is not guaranteed to keep backward compatibility.\n        \"\"\"\n        if len(self._cached_codes) == 0:\n            warnings.warn(\n                'No codes are cached because compilation is deferred until '\n                'the first function call.')\n        return dict([(k, v.code) for k, v in self._cached_codes.items()])\n\n    @property\n    def cached_code(self):\n        \"\"\"Returns `next(iter(self.cached_codes.values()))`.\n\n        This proprety method is for debugging purpose.\n        The return value is not guaranteed to keep backward compatibility.\n        \"\"\"\n        codes = self.cached_codes\n        if len(codes) > 1:\n            warnings.warn(\n                'The input types of the kernel could not be inferred. '\n                'Please use `.cached_codes` instead.')\n        return next(iter(codes.values()))\n\n\ndef rawkernel(*, mode='cuda', device=False):\n    \"\"\"A decorator compiles a Python function into CUDA kernel.\n    \"\"\"\n    cupy._util.experimental('cupyx.jit.rawkernel')\n\n    def wrapper(func):\n        return functools.update_wrapper(\n            _JitRawKernel(func, mode, device), func)\n    return wrapper\n\n\nthreadIdx = _internal_types.Data('threadIdx', _cuda_types.dim3)\nblockDim = _internal_types.Data('blockDim', _cuda_types.dim3)\nblockIdx = _internal_types.Data('blockIdx', _cuda_types.dim3)\ngridDim = _internal_types.Data('gridDim', _cuda_types.dim3)\n\nwarpsize = _internal_types.Data(\n    '64' if runtime.is_hip else '32', _cuda_types.uint32)\nwarpsize.__doc__ = r\"\"\"Returns the number of threads in a warp.\n\nIn CUDA this is always 32, and in ROCm/HIP always 64.\n\n.. seealso:: :obj:`numba.cuda.warpsize`\n\"\"\"\n", "path": "cupyx/jit/_interface.py"}], "after_files": [{"content": "# mypy: ignore-errors\n\nfrom cupy._core import core  # NOQA\nfrom cupy._core import fusion  # NOQA\nfrom cupy._core import internal  # NOQA\n\n\n# internal APIs for testing and developement\nfrom cupy._core._accelerator import set_elementwise_accelerators  # NOQA\nfrom cupy._core._accelerator import set_reduction_accelerators  # NOQA\nfrom cupy._core._accelerator import set_routine_accelerators  # NOQA\nfrom cupy._core._accelerator import get_elementwise_accelerators  # NOQA\nfrom cupy._core._accelerator import get_reduction_accelerators  # NOQA\nfrom cupy._core._accelerator import get_routine_accelerators  # NOQA\n\n\n# import class and function\nfrom cupy._core._kernel import create_ufunc  # NOQA\nfrom cupy._core._kernel import ElementwiseKernel  # NOQA\nfrom cupy._core._kernel import ufunc  # NOQA\nfrom cupy._core._kernel import _get_warpsize  # NOQA\nfrom cupy._core._reduction import create_reduction_func  # NOQA\nfrom cupy._core._reduction import ReductionKernel  # NOQA\nfrom cupy._core._routines_binary import bitwise_and  # NOQA\nfrom cupy._core._routines_binary import bitwise_or  # NOQA\nfrom cupy._core._routines_binary import bitwise_xor  # NOQA\nfrom cupy._core._routines_binary import invert  # NOQA\nfrom cupy._core._routines_binary import left_shift  # NOQA\nfrom cupy._core._routines_binary import right_shift  # NOQA\nfrom cupy._core._routines_linalg import _mat_ptrs  # NOQA\nfrom cupy._core._routines_linalg import dot  # NOQA\nfrom cupy._core._routines_linalg import get_compute_type  # NOQA\nfrom cupy._core._routines_linalg import matmul  # NOQA\nfrom cupy._core._routines_linalg import set_compute_type  # NOQA\nfrom cupy._core._routines_linalg import tensordot_core  # NOQA\nfrom cupy._core._routines_logic import create_comparison  # NOQA\nfrom cupy._core._routines_logic import equal  # NOQA\nfrom cupy._core._routines_logic import greater  # NOQA\nfrom cupy._core._routines_logic import greater_equal  # NOQA\nfrom cupy._core._routines_logic import less  # NOQA\nfrom cupy._core._routines_logic import less_equal  # NOQA\nfrom cupy._core._routines_logic import not_equal  # NOQA\nfrom cupy._core._routines_manipulation import array_split  # NOQA\nfrom cupy._core._routines_manipulation import broadcast  # NOQA\nfrom cupy._core._routines_manipulation import broadcast_to  # NOQA\nfrom cupy._core._routines_manipulation import concatenate_method  # NOQA\nfrom cupy._core._routines_manipulation import moveaxis  # NOQA\nfrom cupy._core._routines_manipulation import rollaxis  # NOQA\nfrom cupy._core._routines_manipulation import size  # NOQA'\nfrom cupy._core._routines_math import absolute  # NOQA\nfrom cupy._core._routines_math import add  # NOQA\nfrom cupy._core._routines_math import angle, angle_deg  # NOQA\nfrom cupy._core._routines_math import conjugate  # NOQA\nfrom cupy._core._routines_math import divide  # NOQA\nfrom cupy._core._routines_math import floor_divide  # NOQA\nfrom cupy._core._routines_math import multiply  # NOQA\nfrom cupy._core._routines_math import negative  # NOQA\nfrom cupy._core._routines_math import positive  # NOQA\nfrom cupy._core._routines_math import power  # NOQA\nfrom cupy._core._routines_math import remainder  # NOQA\nfrom cupy._core._routines_math import sqrt  # NOQA\nfrom cupy._core._routines_math import subtract  # NOQA\nfrom cupy._core._routines_math import true_divide  # NOQA\nfrom cupy._core._routines_statistics import nanmax  # NOQA\nfrom cupy._core._routines_statistics import nanmin  # NOQA\nfrom cupy._core.core import _internal_ascontiguousarray  # NOQA\nfrom cupy._core.core import _internal_asfortranarray  # NOQA\nfrom cupy._core.core import array  # NOQA\nfrom cupy._core.core import ascontiguousarray  # NOQA\nfrom cupy._core.core import asfortranarray  # NOQA\nfrom cupy._core.core import divmod  # NOQA\nfrom cupy._core.core import elementwise_copy  # NOQA\nfrom cupy._core.core import ndarray  # NOQA\nfrom cupy._core.dlpack import fromDlpack  # NOQA\nfrom cupy._core.dlpack import from_dlpack  # NOQA\nfrom cupy._core.internal import complete_slice  # NOQA\nfrom cupy._core.internal import get_size  # NOQA\nfrom cupy._core.raw import RawKernel  # NOQA\nfrom cupy._core.raw import RawModule  # NOQA\n", "path": "cupy/_core/__init__.py"}, {"content": "import functools\nimport warnings\n\nimport numpy\n\nimport cupy\nfrom cupy._core import core\nfrom cupyx.jit import _compile\nfrom cupyx.jit import _cuda_typerules\nfrom cupyx.jit import _cuda_types\nfrom cupyx.jit import _internal_types\nfrom cupyx.jit._cuda_types import Scalar\n\n\nclass _CudaFunction:\n    \"\"\"JIT cupy function object\n    \"\"\"\n\n    def __init__(self, func, mode, device=False, inline=False):\n        self.attributes = []\n\n        if device:\n            self.attributes.append('__device__')\n        else:\n            self.attributes.append('__global__')\n\n        if inline:\n            self.attributes.append('inline')\n\n        self.name = getattr(func, 'name', func.__name__)\n        self.func = func\n        self.mode = mode\n\n    def __call__(self, *args, **kwargs):\n        raise NotImplementedError\n\n    def _emit_code_from_types(self, in_types, ret_type=None):\n        return _compile.transpile(\n            self.func, self.attributes, self.mode, in_types, ret_type)\n\n\nclass _JitRawKernel:\n    \"\"\"JIT CUDA kernel object.\n\n    The decorator :func:``cupyx.jit.rawkernel`` converts the target function\n    to an object of this class. This class is not inteded to be instantiated\n    by users.\n    \"\"\"\n\n    def __init__(self, func, mode, device):\n        self._func = func\n        self._mode = mode\n        self._device = device\n        self._cache = {}\n        self._cached_codes = {}\n\n    def __call__(\n            self, grid, block, args, shared_mem=0, stream=None):\n        \"\"\"Calls the CUDA kernel.\n\n        The compilation will be deferred until the first function call.\n        CuPy's JIT compiler infers the types of arguments at the call\n        time, and will cache the compiled kernels for speeding up any\n        subsequent calls.\n\n        Args:\n            grid (tuple of int): Size of grid in blocks.\n            block (tuple of int): Dimensions of each thread block.\n            args (tuple):\n                Arguments of the kernel. The type of all elements must be\n                ``bool``, ``int``, ``float``, ``complex``, NumPy scalar or\n                ``cupy.ndarray``.\n            shared_mem (int):\n                Dynamic shared-memory size per thread block in bytes.\n            stream (cupy.cuda.Stream): CUDA stream.\n\n        .. seealso:: :ref:`jit_kernel_definition`\n        \"\"\"\n        in_types = []\n        for x in args:\n            if isinstance(x, cupy.ndarray):\n                t = _cuda_types.CArray.from_ndarray(x)\n            elif numpy.isscalar(x):\n                t = _cuda_typerules.get_ctype_from_scalar(self._mode, x)\n            else:\n                raise TypeError(f'{type(x)} is not supported for RawKernel')\n            in_types.append(t)\n        in_types = tuple(in_types)\n        device_id = cupy.cuda.get_device_id()\n\n        kern, enable_cg = self._cache.get((in_types, device_id), (None, None))\n        if kern is None:\n            result = self._cached_codes.get(in_types)\n            if result is None:\n                result = _compile.transpile(\n                    self._func,\n                    ['extern \"C\"', '__global__'],\n                    self._mode,\n                    in_types,\n                    _cuda_types.void,\n                )\n                self._cached_codes[in_types] = result\n\n            fname = result.func_name\n            enable_cg = result.enable_cooperative_groups\n            options = ('-DCUPY_JIT_MODE', '--std=c++14')\n            backend = result.backend\n            if backend == 'nvcc':\n                options += ('-DCUPY_JIT_NVCC',)\n            module = core.compile_with_cache(\n                source=result.code,\n                options=options,\n                backend=backend)\n            kern = module.get_function(fname)\n            self._cache[(in_types, device_id)] = (kern, enable_cg)\n\n        new_args = []\n        for a, t in zip(args, in_types):\n            if isinstance(t, Scalar):\n                if t.dtype.char == 'e':\n                    a = numpy.float32(a)\n                else:\n                    a = t.dtype.type(a)\n            new_args.append(a)\n\n        kern(grid, block, tuple(new_args), shared_mem, stream, enable_cg)\n\n    def __getitem__(self, grid_and_block):\n        \"\"\"Numba-style kernel call.\n\n        .. seealso:: :ref:`jit_kernel_definition`\n        \"\"\"\n        grid, block = grid_and_block\n        if not isinstance(grid, tuple):\n            grid = (grid, 1, 1)\n        if not isinstance(block, tuple):\n            block = (block, 1, 1)\n        return lambda *args, **kwargs: self(grid, block, args, **kwargs)\n\n    @property\n    def cached_codes(self):\n        \"\"\"Returns a dict that has input types as keys and codes values.\n\n        This proprety method is for debugging purpose.\n        The return value is not guaranteed to keep backward compatibility.\n        \"\"\"\n        if len(self._cached_codes) == 0:\n            warnings.warn(\n                'No codes are cached because compilation is deferred until '\n                'the first function call.')\n        return dict([(k, v.code) for k, v in self._cached_codes.items()])\n\n    @property\n    def cached_code(self):\n        \"\"\"Returns `next(iter(self.cached_codes.values()))`.\n\n        This proprety method is for debugging purpose.\n        The return value is not guaranteed to keep backward compatibility.\n        \"\"\"\n        codes = self.cached_codes\n        if len(codes) > 1:\n            warnings.warn(\n                'The input types of the kernel could not be inferred. '\n                'Please use `.cached_codes` instead.')\n        return next(iter(codes.values()))\n\n\ndef rawkernel(*, mode='cuda', device=False):\n    \"\"\"A decorator compiles a Python function into CUDA kernel.\n    \"\"\"\n    cupy._util.experimental('cupyx.jit.rawkernel')\n\n    def wrapper(func):\n        return functools.update_wrapper(\n            _JitRawKernel(func, mode, device), func)\n    return wrapper\n\n\nthreadIdx = _internal_types.Data('threadIdx', _cuda_types.dim3)\nblockDim = _internal_types.Data('blockDim', _cuda_types.dim3)\nblockIdx = _internal_types.Data('blockIdx', _cuda_types.dim3)\ngridDim = _internal_types.Data('gridDim', _cuda_types.dim3)\n\nwarpsize = _internal_types.Data('warpSize', _cuda_types.int32)\nwarpsize.__doc__ = r\"\"\"Returns the number of threads in a warp.\n\n.. seealso:: :obj:`numba.cuda.warpsize`\n\"\"\"\n", "path": "cupyx/jit/_interface.py"}]}