triton seem to be not buildable outside of google

[cloudhan]

The Issues section is not opened in openxla/triton repo. So I report it here. This is related with openxla/triton and cannot be upstreamed.

I need to apply the following path before producing a jaxlib whl on windows

diff --git a/lib/Target/LLVMIR/LLVMIRTranslation.cpp b/lib/Target/LLVMIR/LLVMIRTranslation.cpp
index b2913e5..ba78043 100644
--- a/lib/Target/LLVMIR/LLVMIRTranslation.cpp
+++ b/lib/Target/LLVMIR/LLVMIRTranslation.cpp
@@ -25,7 +25,7 @@
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Linker/Linker.h"
 #include "llvm/Support/SourceMgr.h"
-#include "third_party/py/triton/google/find_cuda.h"
+// #include "third_party/py/triton/google/find_cuda.h"
 #include <dlfcn.h>
 #include <filesystem>
 #include <iterator>
@@ -164,8 +164,9 @@ static std::map<std::string, std::string> getExternLibs(mlir::ModuleOp module) {
       }
       return std::filesystem::path(fileinfo.dli_fname);
     }();
-    static const auto runtime_path = (
-        fs::path(PathToLibdevice()) / "libdevice.10.bc");
+    static const auto runtime_path =
+        this_library_path.parent_path().parent_path() / "third_party" / "cuda" /
+        "lib" / "libdevice.10.bc";
     if (fs::exists(runtime_path)) {
       externLibs.try_emplace(libdevice, runtime_path.string());
     } else {

And it seems that the find_cuda.h is only available within google, according to
https://github.com/openxla/triton/blob/5b63e5b265a2ff9784b084d901b9feff5a4fc0fc/BUILD#L486-L488

[cheshire]

We do have OSS CI for both JAX and TF, so it does build in OSS checkout. I do not think that we maintain a Windows GPU CI though.

[cloudhan]

Then how can I access the third_party/py/triton/google/find_cuda.h, I didn't find it across the whole openxla org....

I do not think that we maintain a Windows GPU CI though.

I know, in case you want to use an old version of jaxlib on windows, you can access it from https://github.com/cloudhan/jax-windows-builder ;) There is 180GB outbound bandwidth per month at the moment. I think you should take the idea, of officially adding one Win CI for jax, seriously. Some day.

[hawkinsp]

I think there's something worth figuring out here.

openxla/xla is pinning commit 1627e0 of the openxla/triton repository. That commit does not have the find_cuda.h include. I think the find_cuda.h include shouldn't have been checked into the openxla/triton repository, but it's not breaking the build because we're pinning an older commit.

There's a second question, which is: how are we locating libdevice? I think we should make XLA's logic for finding libdevice propagate to triton, otherwise you run the risk of one finding it but the other missing it.

A hack that would work, looking at the code, would be to just set TRITON_LIBDEVICE_PATH from JAX Python, although that's not my favorite idea.

[cloudhan]

A hack that would work, looking at the code, would be to just set TRITON_LIBDEVICE_PATH from JAX Python, although that's not my favorite idea.

It is a fork, then why not just add a new API called SetLibdevicePath or InitBlahBlah and just let xla call it to do the setup?

[hawkinsp]

Well, it's a fork, but it's trying to be a minimal fork. The primary goal of the fork is to synchronize Triton against LLVM head.

@chsigg any input?

[chsigg]

You are right Peter, I didn't intend to push the find_cuda.h change here. XLA is still pinned to 1627e0 so it builds fine. Whether this older commit finds the libdevice.10.bc or whether it ever reaches that code, I don't know.

It is a fork, then why not just add a new API called SetLibdevicePath or InitBlahBlah and just let xla call it to do the setup?

That's pretty much we do with this change internally, but apparently it was not needed in OSS.

The simplest approach likely would be to just revert this google-only change to get back to what we had in 1627e0, before we bump the version that XLA is pinned to.

[cloudhan]

@chsigg would it be acceptable to change it to, like:

index b2913e5..40d1aee 100644
--- a/lib/Target/LLVMIR/LLVMIRTranslation.cpp
+++ b/lib/Target/LLVMIR/LLVMIRTranslation.cpp
@@ -25,7 +25,9 @@
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Linker/Linker.h"
 #include "llvm/Support/SourceMgr.h"
+#if USE_FIND_CUDA
 #include "third_party/py/triton/google/find_cuda.h"
+#endif
 #include <dlfcn.h>
 #include <filesystem>
 #include <iterator>
@@ -164,8 +166,14 @@ static std::map<std::string, std::string> getExternLibs(mlir::ModuleOp module) {
       }
       return std::filesystem::path(fileinfo.dli_fname);
     }();
+#if USE_FIND_CUDA
     static const auto runtime_path = (
         fs::path(PathToLibdevice()) / "libdevice.10.bc");
+#else
+    static const auto runtime_path =
+        this_library_path.parent_path().parent_path() / "third_party" / "cuda" /
+        "lib" / "libdevice.10.bc";
+#endif
     if (fs::exists(runtime_path)) {
       externLibs.try_emplace(libdevice, runtime_path.string());
     } else {

and let target //third_party/py/triton/google:find_cuda propagate a define USE_FIND_CUDA.

[cloudhan]

This problem is resolved.