Multiple invocation of a jax FFI binded function leads to incorrect results

[Vandermode]

Hi @dfm ,

Apologies for the follow-up, but I’ve encountered an unexpected behavior while using the FFI API, building on our previous discussion in #27763.

Your guidance helped me successfully bind my C function, and initial results matched expectations. However, further testing revealed a peculiar issue with the JAX binding that I’d like to clarify.

below is the code for performance benchmarking.

            jax_fft = jax.jit(self.vkfft_app.jax_fft, donate_argnums=0)
            jax_ifft = jax.jit(self.vkfft_app.jax_ifft, donate_argnums=0)
            # Benchmark VkFFT (normal - jax) implementation
            start_time = time.time()
            for _ in range(10):
                x_padded_vkfft = _pad(x_vkfft_jax, prop_cfg.pad_factor)
                # print('np.allclose(x_vkfft_jax, self.x)', np.allclose(self.x, x_vkfft_jax))
                # print(x_vkfft_jax.shape, x_padded_vkfft.shape)
                x_padded_vkfft = jax_fft(x_padded_vkfft)
                x_padded_vkfft = x_padded_vkfft * H_vkfft_jax
                x_padded_vkfft = jax_ifft(x_padded_vkfft)
                dest_field_jax_normal = _unpad(x_padded_vkfft, prop_cfg.pad_factor)
                jax.block_until_ready(dest_field_jax_normal)

            vkfft_jax_time = (time.time() - start_time) / 10
            print(f"VkFFT (normal - jax) implementation: {vkfft_jax_time*1000:.3f} ms per iteration")            
            
            # Benchmark VkFFT (normal) implementation
            start_time = time.time()
            for _ in range(10):
                x_padded_vkfft = cupy.pad(x_vkfft, pad_width=prop_cfg.pad_width, mode="constant", constant_values=0)
                x_padded_vkfft = self.vkfft_app.fft(x_padded_vkfft)
                x_padded_vkfft *= H_vkfft
                x_padded_vkfft = self.vkfft_app.ifft(x_padded_vkfft)
                dest_field3 = x_padded_vkfft[..., pad_y:-pad_y, pad_x:-pad_x] if pad_y > 0 else x_padded_vkfft
                cupy.cuda.stream.get_current_stream().synchronize()
            vkfft_time = (time.time() - start_time) / 10
            print(f"VkFFT (normal - cupy) implementation: {vkfft_time*1000:.3f} ms per iteration")

            np.testing.assert_allclose(dest_field_ref, dest_field3, atol=1e-3, rtol=1e-3)  # pass
            np.testing.assert_allclose(dest_field_ref, dest_field_jax_normal, atol=1e-3, rtol=1e-3) # failed. 

Observations:

1. The JAX FFI binding fails, while the CuPy ctype binding passes the same test.
2. Reducing the loop count to 1 makes the JAX binding produce correct results. Notably, failures occur starting from the second invocation.

Root Cause:

After extensive testing, I traced the issue to temporary buffers allocated automatically within my C function. This is confirmed because:

1. Small inputs (no temp buffer needed): Multiple JAX invocations succeed.
2. Large inputs (require temp buffers): Failures occur consistently after the first invocation.

Core Question:

How can we safely wrap a C function that internally allocates temporary buffers? The current behavior suggests that JAX FFI interactions with these buffers lead to incorrect results upon subsequent invocations.

[dfm]

I don't think there's enough information here to provide any feedback! I don't see why there would be any JAX-specific interactions with allocated temporary buffers - JAX just executes your C++ code!

If you can make a small reproducible example that doesn't depend on your external library, I'd be happy to help debug that.

[Vandermode]

Thanks for the reply! I think creating a small reproducible example without external lib would also be a challenging task.

I think I misclarify the problem context. My C function is actually implemented in CUDA, and it would allocate temporal buffer in GPU, so I guess there might be a conflict with the jax GPU memory management.

After more tests, I found if I add these sum statement to enforce an implicit synchronization, then the results become correct. Removing these sum statement would still result in wrong outcome.

            for _ in range(2):
                x_padded_vkfft = _pad(x_vkfft_jax, prop_cfg.pad_factor)
                print(x_padded_vkfft.sum())
                x_padded_vkfft = jax_fft(x_padded_vkfft)
                print(x_padded_vkfft.sum())
                x_padded_vkfft = x_padded_vkfft * H_vkfft_jax
                print(x_padded_vkfft.sum())
                x_padded_vkfft = jax_ifft(x_padded_vkfft)
                print(x_padded_vkfft.sum())
                dest_field_jax_normal = _unpad(x_padded_vkfft, prop_cfg.pad_factor)
                print(dest_field_jax_normal.sum())
                dest_field_jax_normal = jax.block_until_ready(dest_field_jax_normal)
                print(dest_field_jax_normal.sum())

I was wondering if you have any insight what's the problem here. Thanks!

[dfm]

You might need to make sure that you're allocating the memory using the CUDA stream from XLA to ensure the correct ordering. Take a look at the CUDA FFI example here for an example of binding the CUDA stream to your FFI handler.

[Vandermode]

Hi @dfm , I am following your advice to bind the CUDA stream, but encounter another problem:

How can I bind the following function via the FFI API?

ffi::Error VkFFTInitConfigImpl(cudaStream_t stream, void *config, void *app)

I can only pass the CUDA stream to the config, and then contruct my app from the config, and finally run the computation with app as attribute. (no way to directly pass the stream to the final compute kernel in my CUDA lib)

I was also seeking for an explicit way to control the CUDA stream in jax (because then I can explicitly pass the jax cuda stream to my config), but seems nobody has answers for this question See, #24426

// Update the handler registration
XLA_FFI_DEFINE_HANDLER_SYMBOL(VkFFTInitConfig, VkFFTInitConfigImpl,
                       ffi::Ffi::Bind()
                        .Ctx<xla::ffi::PlatformStream<CUstream>>()
                        .Attr<ffi::Pointer<void>>('config') // config
                        .Attr<ffi::Pointer<void>>('app') // output
);


XLA_FFI_DEFINE_HANDLER_SYMBOL(VkFFTForward, VkFFTForwardImpl,
                              ffi::Ffi::Bind()                  
                                .Attr<ffi::Pointer<void>>("app")
                                .Arg<ffi::AnyBuffer>() // input
                                .Arg<ffi::AnyBuffer>() // kernel
                                .Ret<ffi::AnyBuffer>() // output
);

XLA_FFI_DEFINE_HANDLER_SYMBOL(VkFFTInverse, VkFFTInverseImpl,
                              ffi::Ffi::Bind()
                                .Attr<ffi::Pointer<void>>("app")
                                .Arg<ffi::AnyBuffer>() // input
                                .Ret<ffi::AnyBuffer>() // output
);