Skip to content

Commit

Permalink
make linters happy
Browse files Browse the repository at this point in the history
  • Loading branch information
@leofang
leofang committed Mar 15, 2024
1 parent 7b964eb commit 6bc3fba
Show file tree
Hide file tree
Showing 8 changed files with 61 additions and 56 deletions.
4 changes: 2 additions & 2 deletions cupy/cuda/cufft.pxd
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,8 +51,8 @@ cpdef get_current_plan()
cpdef int getVersion() except? -1
cpdef void check_result(int result) except*
cpdef void setJITCallback(
intptr_t plan, bytes callback, int callback_type,
intptr_t caller_info) except*
intptr_t plan, bytes callback, int callback_type,
intptr_t caller_info) except*


cdef class Plan1d:
Expand Down
2 changes: 1 addition & 1 deletion cupy/cuda/cufft.pyx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ ctypedef Result (*F_cufftXtSetJITCallback)(
cdef F_cufftXtSetJITCallback cufftXtSetJITCallback


################ SoftLink utilities ################
# ****************** SoftLink utilities ******************

cdef SoftLink _L = None

Expand Down
4 changes: 2 additions & 2 deletions cupyx/jit/_interface.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,6 @@ def get_lto_ir(self, args):
)
self._cached_codes[in_types] = result

fname = result.func_name
options = (
'-DCUPY_JIT_MODE', '--std=c++14', '-dlto',
f'-arch=compute_{_get_arch()}')
Expand All @@ -189,7 +188,8 @@ def get_lto_ir(self, args):
# this statement is still valid now that we fix the CCCL issues
if major_version >= 12:
# Starting with CUDA 12.0, even without using jitify, some
# tests cause an error if the following option is not included.
# tests cause an error if the following option is not
# included.
options += ('--device-as-default-execution-space',)
program = _NVRTCProgram(
source, cu_path, headers, include_names, method='lto')
Expand Down
47 changes: 25 additions & 22 deletions examples/cufft/callback/jit_r2c_c2r_string.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
import os

import cupy as cp
import numpy as np

Expand All @@ -8,48 +6,53 @@
#include <cufftXt.h>
struct cb_params {
unsigned window_N;
unsigned signal_size;
unsigned window_N;
unsigned signal_size;
};
__device__ cufftComplex cufftJITCallbackLoadComplex(void *input,
size_t index,
void *info,
void *sharedmem) {
const cb_params* params = static_cast<const cb_params*>(info);
cufftComplex* cb_output = static_cast<cufftComplex*>(input);
const unsigned sample = index % params->signal_size;
const cb_params* params = static_cast<const cb_params*>(info);
cufftComplex* cb_output = static_cast<cufftComplex*>(input);
const unsigned sample = index % params->signal_size;
return (sample < params->window_N) ? cb_output[index] : cufftComplex{0.f, 0.f};
return (sample < params->window_N) ? cb_output[index] \
: cufftComplex{0.f, 0.f};
}
"""


# Problem input parameters
batches = 830
signal_size = 328
window_size = 32
batches = 830
signal_size = 328
window_size = 32
complex_signal_size = signal_size // 2 + 1

# Wave parameters
waves = 12
signal_max_A = 20.
signal_max_f = 500.
sampling_dt = 1e-3
waves = 12
signal_max_A = 20.
signal_max_f = 500.
sampling_dt = 1e-3

# Precision threshold
threshold = 1e-6;
threshold = 1e-6

# Initialize the input signal as a composite of sine waves
# with random amplitudes and frequencies
wave_amplitudes = signal_max_A * cp.random.random((batches, waves), dtype=cp.float32)
wave_frequencies = signal_max_f * cp.random.random((batches, waves), dtype=cp.float32)
wave_amplitudes = signal_max_A * \
cp.random.random((batches, waves), dtype=cp.float32)
wave_frequencies = signal_max_f * \
cp.random.random((batches, waves), dtype=cp.float32)

# Compose the signal
input_signals = cp.empty((batches, signal_size), dtype=cp.float32)
time = 0.
for s in range(signal_size):
input_signals[..., s] = cp.sum(wave_amplitudes[...] * cp.sin(2 * cp.pi * wave_frequencies[...] * time), axis=-1)
input_signals[..., s] = cp.sum(
wave_amplitudes[...] *
cp.sin(2 * cp.pi * wave_frequencies[...] * time), axis=-1)
time += sampling_dt

# Define a structure used to pass in the window size
Expand All @@ -59,7 +62,7 @@
{'names': ('window_N', 'signal_size'),
'formats': (np.uint32, np.uint32),
'itemsize': 8,
}, align=True
}, align=True
)
cb_params_h = np.empty(1, dtype=cb_params_dtype)
cb_params_h['window_N'] = window_size
Expand All @@ -72,11 +75,11 @@
memptr_d.copy_from_host(cb_params_h.ctypes.data, cb_params_h.nbytes)

# Transform signal forward
print("Transforming signal with rfft (cufftExecR2C)");
print("Transforming signal with rfft (cufftExecR2C)")
out = cp.fft.rfft(input_signals)

# Apply window via load callback and inverse-transform the signal
print("Transforming signal with irfft (cufftExecC2R)");
print("Transforming signal with irfft (cufftExecC2R)")
with cp.fft.config.set_cufft_callbacks(cb_load=code,
cb_load_data=memptr_d,
cb_ver='jit'):
Expand Down
9 changes: 4 additions & 5 deletions examples/cufft/callback/jit_string.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
import os

import cupy as cp

# a load callback that overwrites the input array to 1
Expand All @@ -23,13 +21,14 @@

# this fftn call uses callback
with cp.fft.config.set_cufft_callbacks(cb_load=code, cb_ver='jit'):
b = cp.fft.fftn(a, axes=(1,2))
b = cp.fft.fftn(a, axes=(1, 2))

# this does not use
c = cp.fft.fftn(cp.ones(shape=a.shape, dtype=cp.complex64), axes=(1,2))
c = cp.fft.fftn(cp.ones(shape=a.shape, dtype=cp.complex64), axes=(1, 2))

# result agrees
assert cp.allclose(b, c)

# "static" plans are also cached, but are distinct from their no-callback counterparts
# "static" plans are also cached, but are distinct from their no-callback
# counterparts
cp.fft.config.get_plan_cache().show_info()
5 changes: 2 additions & 3 deletions examples/cufft/callback/jit_string2.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
import os

import cupy as cp

# a load callback that overwrites the input array to 1
Expand Down Expand Up @@ -31,5 +29,6 @@
# result agrees
assert cp.allclose(b, c)

# "static" plans are also cached, but are distinct from their no-callback counterparts
# "static" plans are also cached, but are distinct from their no-callback
# counterparts
cp.fft.config.get_plan_cache().show_info()
7 changes: 4 additions & 3 deletions examples/cufft/callback/legacy.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,14 @@

# this fftn call uses callback
with cp.fft.config.set_cufft_callbacks(cb_load=code):
b = cp.fft.fftn(a, axes=(1,2))
b = cp.fft.fftn(a, axes=(1, 2))

# this does not use
c = cp.fft.fftn(cp.ones(shape=a.shape, dtype=cp.complex64), axes=(1,2))
c = cp.fft.fftn(cp.ones(shape=a.shape, dtype=cp.complex64), axes=(1, 2))

# result agrees
assert cp.allclose(b, c)

# "static" plans are also cached, but are distinct from their no-callback counterparts
# "static" plans are also cached, but are distinct from their no-callback
# counterparts
cp.fft.config.get_plan_cache().show_info()
39 changes: 21 additions & 18 deletions tests/cupy_tests/fft_tests/test_callback.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
import Cython
except ImportError:
Cython = None
except:
else:
if Cython.__version__ < '0.29.0':
Cython = None
import numpy as np
Expand Down Expand Up @@ -294,7 +294,7 @@ def _test_load_store_helper(self, xp, dtype, fft_func):
load_types = ('x', 'cufftReal', 'cufftCallbackLoadR',
'cufftJITCallbackLoadReal')
store_types = ('x.y', 'cufftComplex', 'cufftCallbackStoreC',
'cufftJITCallbackStoreComplex')
'cufftJITCallbackStoreComplex')
else: # float64
load_types = ('x', 'cufftDoubleReal', 'cufftCallbackLoadD',
'cufftJITCallbackLoadDoubleReal')
Expand Down Expand Up @@ -506,21 +506,24 @@ def test_irfft_load_store_aux(self, xp, dtype):
@testing.parameterize(*(
testing.product_dict(
[
{'shape': (3, 4), 's': None, 'axes': None, 'norm': None},
{'shape': (3, 4), 's': (1, None), 'axes': None, 'norm': None},
{'shape': (3, 4), 's': (1, 5), 'axes': None, 'norm': None},
{'shape': (3, 4), 's': None, 'axes': (-2, -1), 'norm': None},
{'shape': (3, 4), 's': None, 'axes': None, 'norm': 'ortho'},
{'shape': (2, 3, 4), 's': None, 'axes': None, 'norm': None},
{'shape': (2, 3, 4), 's': (1, 4, None), 'axes': None, 'norm': None},
{'shape': (2, 3, 4), 's': (1, 4, 10), 'axes': None, 'norm': None},
{'shape': (2, 3, 4), 's': None, 'axes': (-3, -2, -1), 'norm': None},
{'shape': (2, 3, 4), 's': None, 'axes': None, 'norm': 'ortho'},
{'shape': (2, 3, 4), 's': (2, 3), 'axes': (0, 1, 2), 'norm': 'ortho'},
{'shape': (3, 4), 's': None, 'axes': None, 'norm': None},
{'shape': (3, 4), 's': (1, None), 'axes': None, 'norm': None},
{'shape': (3, 4), 's': (1, 5), 'axes': None, 'norm': None},
{'shape': (3, 4), 's': None, 'axes': (-2, -1), 'norm': None},
{'shape': (3, 4), 's': None, 'axes': None, 'norm': 'ortho'},
{'shape': (2, 3, 4), 's': None, 'axes': None, 'norm': None},
{'shape': (2, 3, 4), 's': (1, 4, None),
'axes': None, 'norm': None},
{'shape': (2, 3, 4), 's': (1, 4, 10), 'axes': None, 'norm': None},
{'shape': (2, 3, 4), 's': None,
'axes': (-3, -2, -1), 'norm': None},
{'shape': (2, 3, 4), 's': None, 'axes': None, 'norm': 'ortho'},
{'shape': (2, 3, 4), 's': (2, 3), 'axes': (
0, 1, 2), 'norm': 'ortho'},
],
testing.product(
{'cb_ver': cb_ver_for_test,},
{'cb_ver': cb_ver_for_test, },
),
)
))
Expand Down Expand Up @@ -692,7 +695,7 @@ def _test_load_store_helper(self, xp, dtype, fft_func):
load_types = ('x', 'cufftReal', 'cufftCallbackLoadR',
'cufftJITCallbackLoadReal')
store_types = ('x.y', 'cufftComplex', 'cufftCallbackStoreC',
'cufftJITCallbackStoreComplex')
'cufftJITCallbackStoreComplex')
else: # float64
load_types = ('x', 'cufftDoubleReal', 'cufftCallbackLoadD',
'cufftJITCallbackLoadDoubleReal')
Expand All @@ -710,7 +713,7 @@ def _test_load_store_helper(self, xp, dtype, fft_func):
'x.x', 'cufftDoubleComplex', 'cufftCallbackLoadZ',
'cufftJITCallbackLoadDoubleComplex')
store_types = ('x', 'cufftDoubleReal', 'cufftCallbackStoreD',
'cufftJITCallbackStoreDoubleReal')
'cufftJITCallbackStoreDoubleReal')
cb_load = _set_load_cb(load_code, *load_types, cb_ver=self.cb_ver)
cb_store = _set_store_cb(store_code, *store_types, cb_ver=self.cb_ver)

Expand Down Expand Up @@ -789,7 +792,7 @@ def _test_load_store_aux_helper(self, xp, dtype, fft_func):
load_types = ('x', 'cufftReal', 'cufftCallbackLoadR',
'cufftJITCallbackLoadReal', 'float')
store_types = ('x.y', 'cufftComplex', 'cufftCallbackStoreC',
'cufftJITCallbackStoreComplex', 'float')
'cufftJITCallbackStoreComplex', 'float')
else: # float64
load_types = ('x', 'cufftDoubleReal', 'cufftCallbackLoadD',
'cufftJITCallbackLoadDoubleReal', 'double')
Expand All @@ -807,7 +810,7 @@ def _test_load_store_aux_helper(self, xp, dtype, fft_func):
'x.x', 'cufftDoubleComplex', 'cufftCallbackLoadZ',
'cufftJITCallbackLoadDoubleComplex', 'double')
store_types = ('x', 'cufftDoubleReal', 'cufftCallbackStoreD',
'cufftJITCallbackStoreDoubleReal', 'double')
'cufftJITCallbackStoreDoubleReal', 'double')
cb_load = _set_load_cb(load_code, *load_types, cb_ver=self.cb_ver)
cb_store = _set_store_cb(store_code, *store_types, cb_ver=self.cb_ver)

Expand Down

0 comments on commit 6bc3fba

Please sign in to comment.