Add experimental support for building JAX CPU and GPU wheels with GCC.

The `build.py` script uses Clang compiler by default, and JAX doesn't support building with GCC officially. However, experimental GCC support is still present.

Command examples:

```
python build/build.py build --wheels=jaxlib,jax-cuda-plugin --use_clang=false
python build/build.py build --wheels=jaxlib,jax-cuda-plugin --use_clang=false --gcc_path=/use/bin/gcc
```

This change addresses the request in #25488.

PiperOrigin-RevId: 707930913

Author: Google-ML-Automation

build/build.py

@@ -277,6 +277,15 @@ def add_artifact_subcommand_arguments(parser: argparse.ArgumentParser):
         """,
   )
 
+  compile_group.add_argument(
+      "--gcc_path",
+      type=str,
+      default="",
+      help="""
+        Path to the GCC binary to use.
+        """,
+  )
+
   compile_group.add_argument(
       "--disable_mkl_dnn",
       action="store_true",
@@ -481,7 +490,13 @@ async def main():
         # versions of Clang.
         wheel_build_command.append("--config=clang")
     else:
-      logging.debug("Use Clang: False")
+      gcc_path = args.gcc_path or utils.get_gcc_path_or_exit()
+      logging.debug(
+          "Using GCC as the compiler, gcc path: %s",
+          gcc_path,
+      )
+      wheel_build_command.append(f"--repo_env=CC=\"{gcc_path}\"")
+      wheel_build_command.append(f"--repo_env=BAZEL_COMPILER=\"{gcc_path}\"")
 
     if not args.disable_mkl_dnn:
       logging.debug("Enabling MKL DNN")
@@ -515,12 +530,16 @@ async def main():
 
     if "cuda" in wheel:
       wheel_build_command.append("--config=cuda")
-      wheel_build_command.append(
+      if args.use_clang:
+        wheel_build_command.append(
             f"--action_env=CLANG_CUDA_COMPILER_PATH=\"{clang_path}\""
         )
-      if args.build_cuda_with_clang:
-        logging.debug("Building CUDA with Clang")
-        wheel_build_command.append("--config=build_cuda_with_clang")
+        if args.build_cuda_with_clang:
+          logging.debug("Building CUDA with Clang")
+          wheel_build_command.append("--config=build_cuda_with_clang")
+        else:
+          logging.debug("Building CUDA with NVCC")
+          wheel_build_command.append("--config=build_cuda_with_nvcc")
       else:
         logging.debug("Building CUDA with NVCC")
         wheel_build_command.append("--config=build_cuda_with_nvcc")

build/tools/utils.py

@@ -170,19 +170,25 @@ def get_bazel_version(bazel_path):
     return None
   return tuple(int(x) for x in match.group(1).split("."))
 
-def get_clang_path_or_exit():
-  which_clang_output = shutil.which("clang")
-  if which_clang_output:
-    # If we've found a clang on the path, need to get the fully resolved path
+def get_compiler_path_or_exit(compiler_path_flag, compiler_name):
+  which_compiler_output = shutil.which(compiler_name)
+  if which_compiler_output:
+    # If we've found a compiler on the path, need to get the fully resolved path
     # to ensure that system headers are found.
-    return str(pathlib.Path(which_clang_output).resolve())
+    return str(pathlib.Path(which_compiler_output).resolve())
   else:
     print(
-        "--clang_path is unset and clang cannot be found"
-        " on the PATH. Please pass --clang_path directly."
+        f"--{compiler_path_flag} is unset and {compiler_name} cannot be found"
+        " on the PATH. Please pass --{compiler_path_flag} directly."
     )
     sys.exit(-1)
 
+def get_gcc_path_or_exit():
+  return get_compiler_path_or_exit("gcc_path", "gcc")
+
+def get_clang_path_or_exit():
+  return get_compiler_path_or_exit("clang_path", "clang")
+
 def get_clang_major_version(clang_path):
   clang_version_proc = subprocess.run(
       [clang_path, "-E", "-P", "-"],