Support for CUFFT callbacks

[vchuravy]

One interesting use-case is to pass Julia functions as callbacks to cufft
https://devblogs.nvidia.com/parallelforall/cuda-pro-tip-use-cufft-callbacks-custom-data-processing/

To register the callbacks with the cuFFT plan, the first step is to get the device function pointers from the device onto the host.

[maleadt]

This will require quite some work, because we don't have symbol resolution other than calling functions from emit_invoke (ie. no function pointers). We'd need proper JIT symbol lookup, like fptr/jl_generate_fptr, somehow dealing with multiple contexts, the fact that API calls only work at runtime vs compile-time, etc along the way.

It's also not clear what exactly the pointer should be (CUfunction_t? or more likely cuModuleGetGlobal).

[david-macmahon]

In reply to your question in #614 about how important this feature is, I think the answer is that it very important in certain circumstances. For example, we currently use CuFFT callbacks in a CUDA C program that performs long FFTs of 8-bit signed integer data (equivalent to Complex{Int8}) and then produce integrated power spectra. Callbacks are used on both inputs and outputs.

Input callbacks

CuFFT does not handle 8-bit data directly so without callbacks one has to pre-convert the input array to 32-bit floats. This increases the size of the CuFFT input buffer by a factor of at least 4 (assuming the CuFFT output buffer can be used to stage the 8-bit input data) compared to using callbacks where the input buffer holds the 8-bit samples and conversion to 32-bit floats happens "on-the-fly" so the 32-bit floats never occupy global memory. This reduced memory requirement allows for longer (or just more) FFTs to be performed.

Output callbacks

The integrated power spectra are produced by integrating (i.e. summing) the magnitude squared (aka abs2()) of the complex outputs of multiple FFTs. Without callbacks, the FFT must write the output samples to global memory and then another kernel must read the output samples that were just written to global memory before computing the magnitude squared and adding it into the integrated power spectra. With callbacks, the magnitude squared calculation and integration can happen "on-the-fly" as the output samples are produced thereby eliminating the need to read the just-written output samples from global memory. This reduces the load on the GPU global memory which can be good for overall throughput.

CuFFT callbacks allow us to do longer FFTs with (potentially) higher throughput than would be possible without callbacks. This makes CUDA.jl support of this feature important for us so we can fully replace the C program with a Julia implementation. Without this feature, we will have to keep the C program around for when the desired FFT lengths with 32-bit input buffers exceed the memory available on the GPU.

[maleadt]

I was having another look at the documentation, and:

NOTE:The callback API is available in the statically linked cuFFT library only, and only on 64 bit LINUX operating systems.

so that's a problem in itself. Or maybe that only applies to the helper routines, as cufftXtSetCallback is available in libcufft.so (we just haven't wrapped cufftXt.h -- not sure when and how the Xt functionality is available exactly).

[maleadt]

From https://webcache.googleusercontent.com/search?q=cache:RuVLoPCjAx4J:https://quabr.com/40565350/can-i-use-nvrtc-to-compile-cufft-callback-functions-that-use-compiler-generated+&cd=1&hl=en&ct=clnk&gl=be&client=firefox-b-d:

    Generate the source code for the callback function based on the parameters that I'm provided at runtime.

    Use the NVRTC library to compile the source code to PTX.

    Use the CUDA Driver API's linker function to compile the PTX to a cubin.

    Load the resulting cubin using the CUDA Driver API's module management functionality.

    Use the module API to obtain a pointer to the location of the callback function on the device.

    Pass the resulting callback function pointer to the cuFFT library via cufftXtSetCallback().

We have all of that, so it should be possible.

[david-macmahon]

I was having another look at the documentation, and:

NOTE:The callback API is available in the statically linked cuFFT library only, and only on 64 bit LINUX operating systems.

so that's a problem in itself. Or maybe that only applies to the helper routines, as cufftXtSetCallback is available in libcufft.so (we just haven't wrapped cufftXt.h -- not sure when and how the Xt functionality is available exactly).

Wow, yeah, that could be a show stopper. I checked my C program that uses CuFFT callbacks and it looks like it does statically link the CuFFT library. I'll try to build it with dynamically linking to the CuFFT library to see if that works.

[maleadt]

I'm also not sure how to look up the function pointer of a kernel, cuModuleGetGlobal doesn't seem to work:

julia> @device_code_ptx kernel = @cuda launch=false identity(nothing)
// PTX CompilerJob of kernel identity(Nothing) for sm_75

//
// Generated by LLVM NVPTX Back-End
//

.version 6.3
.target sm_75
.address_size 64

        // .globl       _Z19julia_identity_2585v // -- Begin function _Z19julia_identity_2585v
.weak .global .align 8 .u64 exception_flag;
                                        // @_Z19julia_identity_2585v
.visible .entry _Z19julia_identity_2585v()
{


// %bb.0:                               // %top
        ret;
                                        // -- End function
}

julia> kernel_global = CuGlobal{Ptr{Cvoid}}(kernel.mod, "_Z19julia_identity_2585v")
ERROR: CUDA error: named symbol not found (code 500, ERROR_NOT_FOUND)

We could always emit a global datastructure that points to the functions in the module, a la https://forums.developer.nvidia.com/t/how-can-i-use-device-function-pointer-in-cuda/14405/18, but that seems cumbersome.

EDIT: looks like in C too you need to assign the device function to a global variable, https://github.com/zchee/cuda-sample/blob/05555eef0d49ebdde999f5430f185a225ef00dcd/7_CUDALibraries/simpleCUFFT_callback/simpleCUFFT_callback.cu#L48-L57.

[david-macmahon]

I don't know that much about the Driver API's module handling, but does CuGlobal() end up calling cuModuleGetGlobal()? Not sure how that differs from cuModuleGetFunction().

FWIW, the C program I use is available here. It uses the runtime API and the callback functions are defined as __device__ functions and then the device function pointers come from cudaMemcpyFromSymbol() (probably just following examples in the CuFFT docs).

[maleadt]

cuModuleGetFunction only works for kernels, not for device functions.

Looks like we'll need to emit something like:

static __device__ void callback_f()
{
    return;
}

typedef void (*callback_t)();

 __device__ callback_t callback_alias = callback_f;

.func _ZN61_INTERNAL_39_tmpxft_0000b070_00000000_7_test_cpp1_ii_16bd2dde10callback_fEv
()
;
.global .align 8 .u64 callback_alias = _ZN61_INTERNAL_39_tmpxft_0000b070_00000000_7_test_cpp1_ii_16bd2dde10callback_fEv;

.func _ZN61_INTERNAL_39_tmpxft_0000b070_00000000_7_test_cpp1_ii_16bd2dde10callback_fEv()
{
        ret;
}

[david-macmahon]

I noticed the sample you linked to also uses cudaMemcpyFromSymbol(). I also noticed its Makefile has:

LIBRARIES += -lcufft_static -lculibos

so it too is using the statically linked CuFFT library.

[maleadt]

julia> kernel = nothing

julia> @device_code_ptx kernel = cufunction(identity, Tuple{Nothing}; kernel=false)
// PTX CompilerJob of function identity(Nothing) for sm_75

//
// Generated by LLVM NVPTX Back-End
//

.version 6.3
.target sm_75
.address_size 64

        // .globl       julia_identity_4839     // -- Begin function julia_identity_4839
.visible .func julia_identity_4839
()
;
.visible .global .align 8 .u64 entrypoint = julia_identity_4839;
.weak .global .align 8 .u64 exception_flag;
                                        // @julia_identity_4839
.visible .func julia_identity_4839()
{


// %bb.0:                               // %top
        ret;
                                        // -- End function
}

julia> kernel
CUDA.DeviceFunction{identity, Tuple{Nothing}}(CuContext(0x00000000050fb8a0, instance 12bb30420e1458f8), CuModule(Ptr{Nothing} @0x000000000952d580, CuContext(0x00000000050fb8a0, instance 12bb30420e1458f8)), CuPtr{Nothing}(0x0000000000000001))

julia> CuGlobal{Ptr{Cvoid}}(kernel.mod, "entrypoint")[]
Ptr{Nothing} @0x0000000000000001

That 0x1 doesn't look right...