Support using our own registered CC toolchain

[cerisier]

Currently rules_ml_toolchain defines and registers a custom CC toolchain but doesn't easily supports using a custom registered CC toolchain.

From what I experienced, this comes from internal assumptions about which clang will be available and forced cuda wrappers headers from a provided sysroot.

rules_ml_toolchain/gpu/cuda/cuda_configure.bzl

Lines 312 to 319 in 6d19ed7

	if _use_hermetic_toolchains(repository_ctx):
	if _is_linux_x86_64(repository_ctx):
	clang_major_version = _llvm_x86_64_hermetic_version
	elif _is_linux_aarch64(repository_ctx):
	clang_major_version = _llvm_aarch64_hermetic_version
	else:
	print("This OS or architecture isn't supported by Hermetic C++.")

AND

rules_ml_toolchain/gpu/cuda/cuda_configure.bzl

Lines 626 to 634 in 6d19ed7

	if _is_linux_x86_64(repository_ctx):
	hermetic_wrappers_headers = "@llvm_linux_x86_64//:cuda_wrappers_headers"
	elif _is_linux_aarch64(repository_ctx):
	hermetic_wrappers_headers = "@llvm_linux_aarch64//:cuda_wrappers_headers"
	else:
	print("This OS or architecture isn't supported by Hermetic C++.")
	# Set up BUILD file for cuda/
	repository_ctx.template(

This aligns well with the hermetically provided clang but forbids using your own hermetic CC toolchain.

---

With a few patches, I was able to use the fully hermetic @llvm BCR module (I am a maintainer) to compile both CPU and GPU (CUDA) XLA PJRT plugins via rules_ml_toolchain, in remote cross-builds (bazel from my mac, linux remote-executors).

Would you be interested in discussing how to make that a first class option in rules_ml_toolchain ?

[lalten]

Would you have these patches available somewhere @cerisier ? :)

[cerisier]

Here you go: https://github.com/cerisier/xla/commits/llvm-clean/
Extreme WIP.

What it does:

• Based on @meteorcloudy Bazel 8 patche for both XLA and rules_ml_toolchain
• Registers fully hermetic, lighweight, from source @llvm toolchain instead of rules_ml_toolchain sysroot ones.
• Wraps cuda_library so that we can use the vanilla @llvm toolchain (with --cuda-path out of the toolchain config)
• Various tweaks to .bazelrc to allow vanilla remote-exec.

On rules_ml_toolchain side:

• The patches hardcode the clang version and removes rules_ml_toolchain specific hermetic_wrappers_headers in favor of the vanilla @llvm ones, until we find a way to configure this (Thiss issue).
• Uses cccl from github HEAD instead of the one from redist (can't remember why).

Neither XLA nor rules_ml_toolchain support cross remote builds (compiling from your mac even tho remote executors are linux) so you will need to be on a linux machine to compile, ideally x86.

The patches to fix cross remote build are too involved.

bazel build //xla/pjrt/c:pjrt_c_api_gpu_plugin.so --config pjrt_cuda13 --config remote --platforms @llvm//platforms:linux_amd64 --noincompatible_disallow_empty_glob

where --config remote:

common:remote --remote_header=x-buildbuddy-api-key=$BUILDBUDDY_API_KEY
common:remote --bes_results_url=https://app.buildbuddy.io/invocation/
common:remote --bes_backend=grpcs://remote.buildbuddy.io
common:remote --remote_cache=grpcs://remote.buildbuddy.io
common:remote --remote_cache_compression
common:remote --remote_download_toplevel
common:remote --nobuild_runfile_links
common:remote --remote_timeout=3600
common:remote --noexperimental_throttle_remote_action_building
common:remote --experimental_remote_execution_keepalive
common:remote --grpc_keepalive_time=30s
common:remote --remote_executor=grpcs://remote.buildbuddy.io
common:remote --jobs=800
common:remote_llvm --extra_execution_platforms=@llvm//:rbe_platform

For CPU plugins, you don't need any of the GPU specific changes in this fork and just replacing the registered toolchains by @llvm ones is enough.

When I refer to @llvm, I refer to https://github.com/hermeticbuild/hermetic-llvm (https://registry.bazel.build/modules/llvm)

I'm extremely motivated to find a way to make GPU remote builds much affordable for everyone so I hope this issue can lead somewhere :)