From b5eacfa2727c13b5853cb04c5983cff49eee1b45 Mon Sep 17 00:00:00 2001 From: Sunita Nadampalli Date: Sun, 2 Feb 2025 14:45:40 +0000 Subject: [PATCH] add aarch64 platform build support --- codegen/fix_includes.sh | 6 +++++- setup.py | 5 ++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/codegen/fix_includes.sh b/codegen/fix_includes.sh index 9256fa2a9b8b..73a763a91cfb 100755 --- a/codegen/fix_includes.sh +++ b/codegen/fix_includes.sh @@ -6,7 +6,11 @@ sed -i 's##"torch_xla/csrc/tensor.h"#' $@ # remove the runfiles-prefix used in codegen for pytorch # `torchgen` generates relative includes path and does not support customizing the root, # so we have to fix them up. -sed -i 's#bazel-out/k8-[^/]*/bin/codegen/lazy_tensor_generator.runfiles/torch/##' $@ +if [[ $(uname -m) == "x86_64" ]]; then + sed -i 's#bazel-out/k8-[^/]*/bin/codegen/lazy_tensor_generator.runfiles/torch/##' $@ +elif [[ $(uname -m) == "aarch64" ]]; then + sed -i 's#bazel-out/aarch64-[^/]*/bin/codegen/lazy_tensor_generator.runfiles/torch/##' $@ +fi # use the generated files that are in the compilation unit sed -i 's##"bazel-out/\1"#' $@ diff --git a/setup.py b/setup.py index 66e7f1511831..5adf5e315b30 100644 --- a/setup.py +++ b/setup.py @@ -83,7 +83,10 @@ _libtpu_wheel_name += f".dev{_date}+nightly" _libtpu_storage_directory = 'libtpu-nightly-releases' -_libtpu_storage_path = f'https://storage.googleapis.com/{_libtpu_storage_directory}/wheels/libtpu/{_libtpu_wheel_name}-py3-none-linux_x86_64.whl' +import platform +platform_machine = platform.machine() + +_libtpu_storage_path = f'https://storage.googleapis.com/{_libtpu_storage_directory}/wheels/libtpu/{_libtpu_wheel_name}-py3-none-linux_{platform_machine}.whl' def _get_build_mode():