[ROCm] Implement RNN support

[Ruturaj4]

Created from: ROCm#171

[Ruturaj4]

@dfm and @superbobry could you please take a look?

[superbobry]

@dfm want to have a look as well?

[dfm]

This looks good overall - thanks! My main high level comment is that it would be useful to move as much of the #ifdef JAX_GPU_HIP logic into vendor.h rather than in rnn_kernels.cc directly. It's ok to have some, but the more we can move, the better. Can you look into redefining some of the macros in vendor.h to consolidate the logic there?

[dfm]

Since gpu_rnn is in jaxlib, these changes will cause problems with version skew. JAX always needs to work with the most recent stable release of jaxlib. Perhaps you could protect this using hasattr(gpu_rnn, "miopen_rnn_fwd_lowering")?

[dfm]

Similarly, this needs to be protected against old version of jaxlib.

[Ruturaj4]

@dfm still not sure why this error wouldn't go away. I have protections in place. Probably it is how you test this in your internal CI?

Seems like you are getting the jaxlib from upstream and that is why the related tests fail?

[Ruturaj4]

This looks good overall - thanks! My main high level comment is that it would be useful to move as much of the #ifdef JAX_GPU_HIP logic into vendor.h rather than in rnn_kernels.cc directly. It's ok to have some, but the more we can move, the better. Can you look into redefining some of the macros in vendor.h to consolidate the logic there?

@dfm thanks. I see what you mean. However, miopen apis are quiet different from cudnn. For e.g.

#ifdef JAX_GPU_HIP
  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetDropoutDescriptor(
      dropout_desc, handle.get(), d.dropout, dropout_states_dev, state_size, 123, false, false,
      MIOPEN_RNG_PSEUDO_XORWOW)));
#else // JAX_GPU_CUDA
  JAX_RETURN_IF_ERROR(JAX_AS_STATUS(gpudnnSetDropoutDescriptor(
      dropout_desc, handle.get(), d.dropout, nullptr, state_size, 123)));
#endif // JAX_GPU_HIP

I checked to see how many of JAX_GPU_HIP I can move, however, seems like it is very difficult to do anything considering the significant differences between the apis. What do you think?

[dfm]

Seems like you are getting the jaxlib from upstream and that is why the related tests fail?

Yes! We require that jax (the Python package) always be compatible with the currently released jaxlib. You'll probably need some sort of version guard, or you can protect the jax/jaxlib boundary using something like hasattr(gpu_rnn, "...").

Also: It looks like this has introduced some build issues for the CUDA CI. Can you take a look at those too?

[Ruturaj4]

@dfm I just fixed the patch. Could you please approve? thanks!