Use clang for PyTorch/XLA only.

pytorch · Feb 6, 2025 · 5e38816 · 5e38816
1 parent d8ab058
commit 5e38816
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 5 deletions.
diff --git a/infra/ansible/config/env.yaml b/infra/ansible/config/env.yaml
@@ -2,8 +2,8 @@
 # They'll be accessible for all processes on the host, also in the development image.
 release_env:
   common:
-    CC: "/usr/lib/llvm-{{ clang_version }}/bin/clang"
-    CXX: "/usr/lib/llvm-{{ clang_version }}/bin/clang"
+    CC: gcc-10
+    CXX: g++-10
     LD_LIBRARY_PATH: "$LD_LIBRARY_PATH:/usr/local/lib"
 
   tpu:
@@ -21,8 +21,8 @@ build_env:
     # Set explicitly to 0 as setup.py defaults this flag to true if unset.
     BUILD_CPP_TESTS: "{{ build_cpp_tests }}"
     # Force GCC because clang/bazel has issues.
-    CC: "/usr/lib/llvm-{{ clang_version }}/bin/clang"
-    CXX: "/usr/lib/llvm-{{ clang_version }}/bin/clang"
+    CC: gcc-10
+    CXX: g++-10
     PYTORCH_BUILD_NUMBER: 1
     TORCH_XLA_VERSION: "{{ package_version }}"
     PYTORCH_BUILD_VERSION: "{{ package_version }}"
@@ -45,3 +45,7 @@ build_env:
     ACCELERATOR: tpu
     TPUVM_MODE: 1
     BUNDLE_LIBTPU: "{{ bundle_libtpu }}"
+
+clang_compiler:
+  CC: /usr/lib/{{ clang_version }}/bin/clang
+  CXX: /usr/lib/{{ clang_version }}/bin/clang++
diff --git a/infra/ansible/playbook.yaml b/infra/ansible/playbook.yaml
@@ -82,6 +82,7 @@
             combine(build_env[arch] | default({}, true)) |
             combine(build_env[accelerator] | default({}, true))
           }}"
+        clang_compiler: "{{ clang_compiler }}"
       when: stage == "build"
       tags: build_srcs
 
@@ -92,6 +93,7 @@
             build_env.common | default({}, true) |
             combine(build_env[arch] | default({}, true)) |
             combine(build_env[accelerator] | default({}, true))
+            combine(clang_compiler)
           }}"
       when: stage == "build_plugin"
       tags: build_plugin

diff --git a/infra/ansible/roles/build_srcs/tasks/main.yaml b/infra/ansible/roles/build_srcs/tasks/main.yaml
@@ -46,7 +46,7 @@
   ansible.builtin.command:
     cmd: python setup.py bdist_wheel
     chdir: "{{ (src_root, 'pytorch/xla') | path_join }}"
-  environment: "{{ env_vars }}"
+  environment: "{{ env_vars | combine(clang_compiler) }}"
 
 - name: Find XLA *.whl files in pytorch/xla/dist
   ansible.builtin.find: