How to add barrier into a jitted function?

[knightXun]

for example:

def exodus_fwd(input, mem_init, alpha, mem_sub, theta, low_mem, use_lowmem, max_spikeNum):
  input.block_until_ready()   # During the testing phase, the input data is asynchronous, and synchronous operations are required here.
  (spike, membrane) = _exodus_fwd_p.bind(input, mem_init, alpha, mem_sub, theta, low_mem, use_lowmem, max_spikeNum)
  return (spike, membrane), (spike, membrane, mem_init, alpha, mem_sub, theta, low_mem, use_lowmem)

If I add jit onto exodus_fwd, Jax can not transfer input.block_until_ready() to a primitive, I want to find a similiar primitive to replace the input.block_until_ready()

[dfm]

Can you say more about what your motivation is for including this here? Why do you want to block there?

[knightXun]

I don't know, but blocking there make the computation right, if remove this the computation is wrong

[dfm]

This sounds like it's probably a bug in your _exodus_fwd_p implementation, but since I don't know what's happening in there, it's hard to give too much concrete advice. If you can narrow it down to a minimal reproducer that doesn't depend on external code, please open a bug report!

[knightXun]

I now suspect that there might be a bug in the implementation of the operator _exodus_fwd_p: calling _exodus_fwd_p multiple times with the same data yields different results.
Now I want to know how to ensure that the CUDA operators are reentrant when I extend CUDA operators for JAX?

[knightXun]

@dfm it seems xla's bug. I pretty sure about that

[dfm]

Thanks for the update. Please put together a minimal, reproducible example that demonstrates the bug, and open an issue on either the JAX issue tracker or the XLA one as appropriate, with a clear description. I'd say that it's very important to include some sample code that we can execute without external dependencies. Thanks!