Re-enable GPU builds for windows after openxla migration and triton integration

[cloudhan]

Description

Build on windows is not possible since 0.4 release along with the introduction of openxla and triton. This is a meta thread to track the work needs to enable the building again.

• Workaround msvc build by replacing forward decl with decl openxla/xla#2264
• Fix nvcc options passing openxla/xla#2842
• Switch to use compiler param file for msvc openxla/xla#2843
• Guard clang pragma openxla/xla#2844
• Add dlfcn-win32 as dep openxla/xla#2853
• Fix build with msvc  openxla/triton#2 This fix compiling flags issue for MSVC.
• Use of preprocessor directive in function-like macro argument list is undefined behavior openxla/xla#2845
• Linking errors when building on windows #14466 This one needs careful consideration. I don't know how to solve the weak symbol of xla::runtime::ffi::GetXlaFfiStream
• triton seem to be not buildable outside of google openxla/xla#2855

This fixes #13646 #14165 #14369 #15876 #14950 #14939

What jax/jaxlib version are you using?

jaxlib 0.4.9 rc

Which accelerator(s) are you using?

No response

Additional system info

Native Windows 11

NVIDIA GPU info

not apply

[cloudhan]

@hawkinsp do you mind coordinating this?

[cloudhan]

cl.exe version 19.35.32217.1 ICE, 19.34 and older are not affected.
https://developercommunity.visualstudio.com/t/ICE-for-invoking-template-with-same-id-n/10362085

Introduced by legit change: openxla/xla@0636d61

[adam-hartshorne]

Looks like dlfcn-win32 dependency is fixed,

triton-lang/triton@323843c

[cloudhan]

@hawkinsp @chsigg only one in openxla/trition is needed. All other issues is fixed at the moment.

bazel 5.1.1 is needed, otherwise, nvcc fatal : Stray '"' due to -DBAZEL_CURRENT_REPOSITORY =\"xla\"

[jlperla]

Hi all. Excited to see the movement here. Is there any chance for a cpu-only windows JAX installable through standard pip/conda for the fall? I would love to teach with it for a course at that time (yes, I know that I can tell people to use the colab, WSL, and an unofficial conda build, but those aren't the same thing as windows support with a replicable pip requirements file/etc. I can give to low-tech users).

[hawkinsp]

@jlperla Yes, I think we'll be able to do that. I'm hoping we can ship Windows CPU (only) wheels in our next release.

[hawkinsp]

With XLA head,

• [XLA:GPU] Fix compilation failures under MSVC. openxla/xla#5501
• [XLA:runtime] Fix compilation problem under MSVC. openxla/xla#5499
  and the patch in Build Fail with invalid numeric argument '/Wno-error' #14369 (comment)
  I can build a jaxlib CUDA wheel on Windows. I didn't test it, but it builds.

I make no promises to keep testing it, but I figured I'd get it building again to see if that helps folks from the community who want to work with it.

[hawkinsp]

And with openxla/xla@972cd21 all the necessary changes have landed! jaxlib should build on Windows with CUDA enabled using XLA head.

[metab0t]

It seems that xla has removed support for native GPU on Windows:
openxla/xla#6127
openxla/xla#6104