How to bind a void pointer in the new FFI API

[Vandermode]

Hi @dfm , I am binding a C library to jax via the new FFI API.

However, It's unclear for me how to bind a void pointer (the app below, actually a pointer to a C struct that stores configs of the computation to be run)

I tried .Ctx and .Attr tags, but none of them works. Note the pointer pointed to a highly complex struct defined in https://github.com/DTolm/VkFFT/blob/066a17c17068c0f11c9298d848c2976c71fad1c1/vkFFT/vkFFT/vkFFT_Structs/vkFFT_Structs.h#L324

namespace ffi = xla::ffi;


// Forward FFT implementation function
ffi::Error VkFFTForwardImpl(void *app, ffi::AnyBuffer input,
                            ffi::AnyBuffer kernel, // This could be an empty buffer
                            ffi::Result<ffi::AnyBuffer> output)
{
    // Cast the void* to VkFFTApplication*
    VkFFTApplication *vkapp = static_cast<VkFFTApplication*>(app);
    
    // Only use kernel if it has elements
    void *kernel_ptr = (kernel.element_count() > 0) ? kernel.untyped_data() : NULL;

    // Call the VkFFT fft function with possibly NULL kernel
    int result = fft(vkapp,
                     input.untyped_data(),
                     output->untyped_data(),
                     kernel_ptr);

    if (result != 0)
    {
        return ffi::Error::Internal("VkFFT forward execution failed with code: " + std::to_string(result));
    }

    return ffi::Error::Success();
}

// Inverse FFT implementation function
ffi::Error VkFFTInverseImpl(void *app, ffi::AnyBuffer input,
                            ffi::Result<ffi::AnyBuffer> output)
{
    // Cast the void* to VkFFTApplication*
    VkFFTApplication *vkapp = static_cast<VkFFTApplication*>(app);

    // Call the VkFFT ifft function
    int result = ifft(vkapp,
                      input.untyped_data(),
                      output->untyped_data());

    if (result != 0)
    {
        return ffi::Error::Internal("VkFFT inverse execution failed with code: " + std::to_string(result));
    }

    return ffi::Error::Success();
}

// Update the handler registration
XLA_FFI_DEFINE_HANDLER_SYMBOL(VkFFTForward, VkFFTForwardImpl,
                              ffi::Ffi::Bind()
                                  .Ctx<void *>()
                                  .Arg<ffi::AnyBuffer>() // input
                                  .OptionalArg<ffi::AnyBuffer>() // kernel
                                  .Ret<ffi::AnyBuffer>() // output
);

XLA_FFI_DEFINE_HANDLER_SYMBOL(VkFFTInverse, VkFFTInverseImpl,
                              ffi::Ffi::Bind()
                                  .Ctx<void *>()
                                  .Arg<ffi::AnyBuffer>() // input
                                  .Ret<ffi::AnyBuffer>() // output
);

[dfm]

Good question! I believe that the following should work:

ffi::Error Impl(void* ptr) {...}

XLA_FFI_DEFINE_HANDLER_SYMBOL(..., Impl,
                              ffi::Ffi::Bind()
                                  .Attr<ffi::Pointer<void>>()
);

Then, from Python, you would pass the pointer as a np.int64. The decoding code in XLA is here: https://github.com/openxla/xla/blob/145f836bd5175dc5dd262f716a0c59af2b0297a0/xla/ffi/api/ffi.h#L1016-L1037

[Vandermode]

Thanks for the hint!
I can now run the FFI binding in python after compilation.
However, the results don't align with my expectation, below is the test code in python

def test_jax():
    size = 2**12
    input = np.random.rand(size).astype(np.float32)
    print(input.shape)        
    input = input.astype(np.complex64)
    gd_np = np.fft.fftn(input)
    
    # move data to GPU
    input = jnp.asarray(input)
    print(input.sharding)

    d_gpu = input
    app = VkFFTApp(input.shape, dtype=np.complex64, ndim=1, inplace=True)    
    # d_gpu = app.jax_fft(d_gpu) # test failed!
    d_gpu = app.fft(d_gpu)  # This passed the test
    
    gd0 = d_gpu.get() if isinstance(d_gpu, cupy.ndarray) else d_gpu
    
    print(np.abs((gd0 - gd_np) / gd_np).max())
    print('np.allclose(gd0, input): ',  np.allclose(gd0, input, rtol=1e-6, atol=gd_np.max()*1e-6))
    print('np.allclose(gd0, gd_np): ',  np.allclose(gd0, gd_np, rtol=1e-6, atol=gd_np.max()*1e-6))
    print('np.allclose(gd_np, input): ',  np.allclose(gd_np, input, rtol=1e-6, atol=gd_np.max()*1e-6))


jax.ffi.register_ffi_target("vkfft_fft", jax.ffi.pycapsule(library.VkFFTForward), platform="CUDA")
jax.ffi.register_ffi_target("vkfft_ifft", jax.ffi.pycapsule(library.VkFFTInverse), platform="CUDA")

class VkFFTApp(VkFFTApp_cuda):
...
    def jax_fft(self, x: jnp.ndarray, kernel: Optional[jnp.ndarray] = None) -> jnp.ndarray:
        """Apply forward FFT to input array.
        
        Args:
            x: Input array
            kernel: Optional convolution kernel
            
        Returns:
            Output array after FFT
        """
        # Define output type
        output_type = jax.ShapeDtypeStruct(x.shape, x.dtype)        
        app_void_ptr_address = ctypes.cast(self.app, ctypes.c_void_p).value
        # Prepare kernel if provided
        if kernel is not None:
            return jax.ffi.ffi_call("vkfft_fft", output_type)(x, kernel, app=app_void_ptr_address)
        else:
            # Pass an empty array as kernel
            empty_kernel = jnp.zeros((0,), dtype=self.dtype)
            return jax.ffi.ffi_call("vkfft_fft", output_type)(x, empty_kernel, app=app_void_ptr_address)

you can see app.fft works while app.jax_fft fails.
In fact, the vkfft_fft just wraps the fft call in C code where fft's ctype binding app.fft passes the test while its FFI wrapped version not.

Btw, do I really need to use FFI interface to expose this C function to jax?
Can jax utilizes a ctype exposed function directly (I have a cupy version of vkfft binding via ctype)

I want to use the auto-diff system (as well as the custom sharding) provided by jax, that's why I manage to make the c binding works with the FFI interface rather than using the cupy version directly.

An alternative way is to implement all the thing in python, by using the cupy function wrapped by jax interface, with dlpack to avoid data copy, see an example with pytorch here https://docs.cupy.dev/en/stable/user_guide/interoperability.html#using-custom-kernels-in-pytorch
though, it's unclear for me if this approach can work for JIT (since JIT is necessary for custom sharding)

I know the questions are too many, Thank you very much!

[Vandermode]

Hi! I know where is the problem!

Since the jax array is immutable, it fails in inplace mode, but works in out-of-place mode.
However, I do like to use the in-place computing feature of my C function, because it can save 2x memory.
I know donate_argnums in jit can be used for inplace computation, but does this works for FFI ? or is there any alternative to use FFI inplace compute?

def test_jax():
    size = 2**12
    input = np.random.rand(size).astype(np.float32)
    print(input.shape)        
    input = input.astype(np.complex64)
    gd_np = np.fft.fftn(input)
    
    # move data to GPU
    input = jnp.asarray(input)
    print(input.sharding)

    d_gpu = input
    app = VkFFTApp(input.shape, dtype=np.complex64, ndim=1, inplace=False)
    d_gpu = app.jax_fft(d_gpu) # test passed!

    # app = VkFFTApp(input.shape, dtype=np.complex64, ndim=1, inplace=True)
    # app.jax_fft(d_gpu) # test failed!
    
    # app = VkFFTApp(input.shape, dtype=np.complex64, ndim=1, inplace=True)
    # app.fft(d_gpu)  # This passed the test (I pass jax array's __cuda_array_interface__['data'][0] pointer to the underlying c function, dlpack would also work)
    
    gd0 = d_gpu.get() if isinstance(d_gpu, cupy.ndarray) else d_gpu
    
    print(np.abs((gd0 - gd_np) / gd_np).max())
    print('np.allclose(gd0, input): ',  np.allclose(gd0, input, rtol=1e-6, atol=gd_np.max()*1e-6))
    print('np.allclose(gd0, gd_np): ',  np.allclose(gd0, gd_np, rtol=1e-6, atol=gd_np.max()*1e-6))
    print('np.allclose(gd_np, input): ',  np.allclose(gd_np, input, rtol=1e-6, atol=gd_np.max()*1e-6))

[Vandermode]

here is the C fft implementation, it can work in in-place mode when in and out are the same pointer

int fft(VkFFTApplication *app, void *in, void *out, void *kernel)
{
    VkFFTLaunchParams par = {};
    par.buffer = (void *const *)&out;
    if (in)
        par.inputBuffer = (void *const *)∈
    par.outputBuffer = (void *const *)&out;

    if (app->configuration.performConvolution)
    {
        *((void **)app->configuration.kernel) = kernel;
        par.kernel = app->configuration.kernel;
    }

    return VkFFTAppend(app, -1, &par);
}

[dfm]

Many questions here - let me try to answer them all!

Btw, do I really need to use FFI interface to expose this C function to jax?

If you can already call your function from Python, another option would be to use jax.pure_callback. You can typically get better performance and more features (e.g. in-place operations as discussed below) via the FFI, at the cost of an extra interface code.

or is there any alternative to use FFI inplace compute?

Yep! The parameter you're looking for is input_output_aliases for ffi_call.

[Vandermode]

Thanks @dfm ! it indeed works!

    def jax_fft(self, x: jnp.ndarray, kernel: Optional[jnp.ndarray] = None) -> jnp.ndarray:
        # Define output type
        output_type = jax.ShapeDtypeStruct(x.shape, x.dtype)        
        app_void_ptr_address = ctypes.cast(self.app, ctypes.c_void_p).value
        # Prepare kernel if provided
        input_output_aliases = {}
        input_output_aliases[0] = 0
        
        if kernel is not None:
            return jax.ffi.ffi_call("vkfft_fft", output_type, input_output_aliases=input_output_aliases)(x, kernel, app=app_void_ptr_address)
        else:
            # Pass an empty array as kernel
            empty_kernel = jnp.zeros((0,), dtype=self.dtype)
            return jax.ffi.ffi_call("vkfft_fft", output_type, input_output_aliases=input_output_aliases)(x, empty_kernel, app=app_void_ptr_address)

May I ask one additional question?
after setting input_output_aliases, it seems the inplace compute should still be executed with d_gpu = app.jax_fft(d_gpu) rather than app.jax_fft(d_gpu) (I expect d_gpu would be modified inplace directly without the need of the output argument)
Simply performing app.jax_fft(d_gpu) seems doesn't change anything, the d_gpu variable is still same as the input.

Moreover, If I use d_gpu_ = app.jax_fft(d_gpu), then d_gpu would not be inplace modified, it keeps the same as the input.

Could you please guide me what would happend for the above situation? Do I have to keep the input, output the same argument to use the inplace compute?

    app = VkFFTApp(input.shape, dtype=np.complex64, ndim=1, inplace=True)
    # d_gpu_ = app.jax_fft(d_gpu) # test failed!
    d_gpu = app.jax_fft(d_gpu) # test passed!
    app.jax_fft(d_gpu) # test failed!
    
    # app = VkFFTApp(input.shape, dtype=np.complex64, ndim=1, inplace=True)
    # app.fft(d_gpu)  # This passed the test

[dfm]

That's right! The values won't be updated in place when run in eager mode like that, but if you put d_gpu = app.jax_fft(d_gpu) under a JIT (potentially combined with donation), you won't see any copies.

[Vandermode]

d_gpu = jax.jit(app.jax_fft, donate_argnums=0)(d_gpu) # This passed the test

this works perfectly!

Thanks a bunch! you save my days!