Since Callback feature isn't documented beyond [this sample] [(https://docs.cupy.dev/en/latest/user_guide/fft.html#fft-callbacks), I might just be holding it wrong. I tried to modify that example to do a real-to-complex transform as attached below.
When run , it gives this output:
Traceback (most recent call last):
File "/workspace/test/CB-R2C-2D.py", line 28, in <module>
assert cp.allclose(b, c)
AssertionError
C2R is working properly but not R2C. Is it supported? It's not clearly documented if all possible fft combinations are supported.
Here is the modified R2C code:
#!/usr/bin/env python3
import cupy as cp
# a load callback that overwrites the input array to 1
code = r'''
__device__ cufftReal CB_ConvertInputR(
void *dataIn,
size_t offset,
void *callerInfo,
void *sharedPtr)
{
cufftReal x;
x = 1.;
return x;
}
__device__ cufftCallbackLoadR d_loadCallbackPtr = CB_ConvertInputR;
'''
a = cp.random.random((64, 128, 128)).astype(cp.float64)
# this fftn call uses callback
with cp.fft.config.set_cufft_callbacks(cb_load=code):
b = cp.fft.rfftn(a, axes=(1,2))
# this does not use
c = cp.fft.rfftn(cp.ones(shape=a.shape, dtype=cp.float64), axes=(1,2))
# result agrees
assert cp.allclose(b, c)