<commit_before>from libc.stdint cimport intptr_t

from cupy_backends.cuda.api.driver cimport get_build_version
from cupy_backends.cuda.api.runtime cimport _is_hip_environment
from cupy.core.core cimport ndarray
from cupy.cuda.device cimport get_compute_capability

import contextlib
import hashlib
import importlib
import os
import shutil
import string
import subprocess
import sys
import sysconfig
import tempfile
import threading

from cupy import __version__ as _cupy_ver
from cupy._environment import (get_nvcc_path, get_cuda_path)
from cupy.cuda.compiler import (_get_bool_env_variable, CompileException)
# for some reason we can't cimport stuff from cupy.cuda.cufft...
from cupy.cuda.cufft import (CUFFT_C2C, CUFFT_C2R, CUFFT_R2C,
                             CUFFT_Z2Z, CUFFT_Z2D, CUFFT_D2Z,
                             CUFFT_CB_LD_COMPLEX, CUFFT_CB_LD_COMPLEX_DOUBLE,
                             CUFFT_CB_LD_REAL, CUFFT_CB_LD_REAL_DOUBLE,
                             CUFFT_CB_ST_COMPLEX, CUFFT_CB_ST_COMPLEX_DOUBLE,
                             CUFFT_CB_ST_REAL, CUFFT_CB_ST_REAL_DOUBLE,)
from cupy.cuda.cufft import getVersion as get_cufft_version


cdef inline _set_cupy_paths():
    # Workaround: older Cython cannot use __file__ in global scope
    global _cupy_root, _cupy_include, _source_dir
    if _cupy_root is None:
        _cupy_root = os.path.join(os.path.dirname(__file__), '..')
        _cupy_include = _cupy_root + '/core/include'
        _source_dir = _cupy_root + '/cuda/'


cdef inline _set_nvcc_path():
    # get_nvcc_path() could be a long string like "ccache nvcc ..."
    cdef str nvcc = None
    global _nvcc
    if _nvcc == []:
        nvcc = get_nvcc_path()
        if nvcc is not None:
            _nvcc = nvcc.split(' ')
        else:
            _nvcc = None


# information needed for building an external module
cdef list _cc = sysconfig.get_config_var('CXX').split(' ')
cdef str _python_include = sysconfig.get_path('include')
cdef list _nvcc = []
cdef str _cuda_path = get_cuda_path()
cdef str _cuda_include = None  # workaround for Read the Docs...
cdef str _nvprune = None
if _cuda_path is not None:
    _cuda_include = _cuda_path + '/include/'
    _nvprune = os.path.join(_cuda_path, 'bin/nvprune')
    if not os.path.isfile(_nvprune):
        _nvprune = None
cdef str _build_ver = str(get_build_version())
cdef int _cufft_ver = get_cufft_version()
cdef str _cupy_root = None
cdef str _cupy_include = None
cdef str _source_dir = None
cdef str _ext_suffix = sysconfig.get_config_var('EXT_SUFFIX')


# callback related stuff
cdef str _callback_dev_code = None
cdef str _callback_cache_dir = os.environ.get(
    'CUPY_CACHE_DIR', os.path.expanduser('~/.cupy/callback_cache')) + '/'
cdef dict _callback_mgr = {}  # keep the Python modules alive
cdef object _callback_thread_local = threading.local()


cdef class _ThreadLocal:
    cdef _CallbackManager _current_cufft_callback

    def __init__(self):
        self._current_cufft_callback = None

    @staticmethod
    cdef _ThreadLocal get():
        cdef _ThreadLocal tls
        try:
            tls = _callback_thread_local.tls
        except AttributeError:
            tls = _callback_thread_local.tls = _ThreadLocal()
        return tls


cdef inline void _sanity_checks(
        str cb_load, str cb_store,
        ndarray cb_load_aux_arr, ndarray cb_store_aux_arr) except*:
    if _is_hip_environment:
        raise RuntimeError('hipFFT does not support callbacks')
    if not sys.platform.startswith('linux'):
        raise RuntimeError('cuFFT callbacks are only available on Linux')
    if not (sys.maxsize > 2**32):
        raise RuntimeError('cuFFT callbacks require 64 bit OS')
    if not cb_load and not cb_store:
        raise ValueError('need to specify either cb_load or cb_store, '
                         'or both')
    if cb_load and 'd_loadCallbackPtr' not in cb_load:
        raise ValueError('need to specify d_loadCallbackPtr in cb_load')
    if cb_store and 'd_storeCallbackPtr' not in cb_store:
        raise ValueError('need to specify d_storeCallbackPtr in cb_store')
    if _nvcc is None:
        raise RuntimeError('nvcc is required but not found')
    if _nvprune is None:
        raise RuntimeError('nvprune is required but not found')  # TODO(leofang): warn instead
    if cb_load_aux_arr is not None:
        if not cb_load:
            raise ValueError('load callback is not given')
    if cb_store_aux_arr is not None:
        if not cb_store:
            raise ValueError('store callback is not given')


cdef inline void _cythonize(str tempdir, str mod_name) except*:
    shutil.copyfile(os.path.join(_source_dir, 'cupy_cufft.h'),
                    os.path.join(tempdir, 'cupy_cufft.h'))
    shutil.copyfile(os.path.join(_source_dir, 'cufft.pxd'),
                    os.path.join(tempdir, mod_name + '.pxd'))
    shutil.copyfile(os.path.join(_source_dir, 'cufft.pyx'),
                    os.path.join(tempdir, mod_name + '.pyx'))
    p = subprocess.run(['cython', '-3', '--cplus',
                        '-E', 'use_hip=0',
                        '-E', 'CUDA_VERSION=' + _build_ver,
                        '-E', 'CUPY_CUFFT_STATIC=True',
                        os.path.join(tempdir, mod_name + '.pyx'),
                        '-o', os.path.join(tempdir, mod_name + '.cpp')],
                       env=os.environ)
    p.check_returncode()

<commit_msg>make nvprune optional<commit_after>from libc.stdint cimport intptr_t

from cupy_backends.cuda.api.driver cimport get_build_version
from cupy_backends.cuda.api.runtime cimport _is_hip_environment
from cupy.core.core cimport ndarray
from cupy.cuda.device cimport get_compute_capability

import contextlib
import hashlib
import importlib
import os
import shutil
import string
import subprocess
import sys
import sysconfig
import tempfile
import threading
import warnings

from cupy import __version__ as _cupy_ver
from cupy._environment import (get_nvcc_path, get_cuda_path)
from cupy.cuda.compiler import (_get_bool_env_variable, CompileException)
# for some reason we can't cimport stuff from cupy.cuda.cufft...
from cupy.cuda.cufft import (CUFFT_C2C, CUFFT_C2R, CUFFT_R2C,
                             CUFFT_Z2Z, CUFFT_Z2D, CUFFT_D2Z,
                             CUFFT_CB_LD_COMPLEX, CUFFT_CB_LD_COMPLEX_DOUBLE,
                             CUFFT_CB_LD_REAL, CUFFT_CB_LD_REAL_DOUBLE,
                             CUFFT_CB_ST_COMPLEX, CUFFT_CB_ST_COMPLEX_DOUBLE,
                             CUFFT_CB_ST_REAL, CUFFT_CB_ST_REAL_DOUBLE,)
from cupy.cuda.cufft import getVersion as get_cufft_version


cdef inline _set_cupy_paths():
    # Workaround: older Cython cannot use __file__ in global scope
    global _cupy_root, _cupy_include, _source_dir
    if _cupy_root is None:
        _cupy_root = os.path.join(os.path.dirname(__file__), '..')
        _cupy_include = _cupy_root + '/core/include'
        _source_dir = _cupy_root + '/cuda/'


cdef inline _set_nvcc_path():
    # get_nvcc_path() could be a long string like "ccache nvcc ..."
    cdef str nvcc = None
    global _nvcc
    if _nvcc == []:
        nvcc = get_nvcc_path()
        if nvcc is not None:
            _nvcc = nvcc.split(' ')
        else:
            _nvcc = None


# information needed for building an external module
cdef list _cc = sysconfig.get_config_var('CXX').split(' ')
cdef str _python_include = sysconfig.get_path('include')
cdef list _nvcc = []
cdef str _cuda_path = get_cuda_path()
cdef str _cuda_include = None  # workaround for Read the Docs...
cdef str _nvprune = None
if _cuda_path is not None:
    _cuda_include = _cuda_path + '/include/'
    _nvprune = os.path.join(_cuda_path, 'bin/nvprune')
    if not os.path.isfile(_nvprune):
        _nvprune = None
cdef str _build_ver = str(get_build_version())
cdef int _cufft_ver = get_cufft_version()
cdef str _cupy_root = None
cdef str _cupy_include = None
cdef str _source_dir = None
cdef str _ext_suffix = sysconfig.get_config_var('EXT_SUFFIX')


# callback related stuff
cdef str _callback_dev_code = None
cdef str _callback_cache_dir = os.environ.get(
    'CUPY_CACHE_DIR', os.path.expanduser('~/.cupy/callback_cache')) + '/'
cdef dict _callback_mgr = {}  # keep the Python modules alive
cdef object _callback_thread_local = threading.local()


cdef class _ThreadLocal:
    cdef _CallbackManager _current_cufft_callback

    def __init__(self):
        self._current_cufft_callback = None

    @staticmethod
    cdef _ThreadLocal get():
        cdef _ThreadLocal tls
        try:
            tls = _callback_thread_local.tls
        except AttributeError:
            tls = _callback_thread_local.tls = _ThreadLocal()
        return tls


cdef inline void _sanity_checks(
        str cb_load, str cb_store,
        ndarray cb_load_aux_arr, ndarray cb_store_aux_arr) except*:
    if _is_hip_environment:
        raise RuntimeError('hipFFT does not support callbacks')
    if not sys.platform.startswith('linux'):
        raise RuntimeError('cuFFT callbacks are only available on Linux')
    if not (sys.maxsize > 2**32):
        raise RuntimeError('cuFFT callbacks require 64 bit OS')
    if not cb_load and not cb_store:
        raise ValueError('need to specify either cb_load or cb_store, '
                         'or both')
    if cb_load and 'd_loadCallbackPtr' not in cb_load:
        raise ValueError('need to specify d_loadCallbackPtr in cb_load')
    if cb_store and 'd_storeCallbackPtr' not in cb_store:
        raise ValueError('need to specify d_storeCallbackPtr in cb_store')
    if _nvcc is None:
        raise RuntimeError('nvcc is required but not found')
    if _nvprune is None:
        warnings.warn('nvprune is not found', RuntimeWarning)
    if cb_load_aux_arr is not None:
        if not cb_load:
            raise ValueError('load callback is not given')
    if cb_store_aux_arr is not None:
        if not cb_store:
            raise ValueError('store callback is not given')


cdef inline void _cythonize(str tempdir, str mod_name) except*:
    shutil.copyfile(os.path.join(_source_dir, 'cupy_cufft.h'),
                    os.path.join(tempdir, 'cupy_cufft.h'))
    shutil.copyfile(os.path.join(_source_dir, 'cufft.pxd'),
                    os.path.join(tempdir, mod_name + '.pxd'))
    shutil.copyfile(os.path.join(_source_dir, 'cufft.pyx'),
                    os.path.join(tempdir, mod_name + '.pyx'))
    p = subprocess.run(['cython', '-3', '--cplus',
                        '-E', 'use_hip=0',
                        '-E', 'CUDA_VERSION=' + _build_ver,
                        '-E', 'CUPY_CUFFT_STATIC=True',
                        os.path.join(tempdir, mod_name + '.pyx'),
                        '-o', os.path.join(tempdir, mod_name + '.cpp')],
                       env=os.environ)
    p.check_returncode()


{'fraction_non_alphanumeric': 0.08409913673071567, 'fraction_numerical': 0.002413441009932238, 'mean_word_length': 2.8464833987861478, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2448634', 'n_tokens_mistral': 3649, 'n_tokens_neox': 3584, 'n_words': 919}