From c0e88d87bc077f601287c1d09a9a05006d929534 Mon Sep 17 00:00:00 2001 From: Jiewen Tan Date: Wed, 30 Aug 2023 23:29:46 +0000 Subject: [PATCH 1/7] Initial commit --- WORKSPACE | 44 +++++++++---------- setup.py | 2 +- .../csrc/runtime/pjrt_computation_client.cc | 2 +- 3 files changed, 24 insertions(+), 24 deletions(-) diff --git a/WORKSPACE b/WORKSPACE index 790f6d8f31c..7df2c2fbe65 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -30,24 +30,24 @@ python_configure( # b) get the sha256 hash of the commit by running: # curl -L https://github.com/openxla/xla/archive/.tar.gz | sha256sum # and update the sha256 with the result. -http_archive( - name = "xla", - patch_args = [ - "-l", - "-p1", - ], - patch_tool = "patch", - patches = [ - "//openxla_patches:cache_urls.diff", - "//openxla_patches:f16_abi_clang.diff", - "//openxla_patches:gpu_race_condition.diff", - "//openxla_patches:constexpr_return.diff", - ], - strip_prefix = "xla-cd2cf5c34931e4fc1cacf83bfc480a5b93f05f6d", - urls = [ - "https://github.com/openxla/xla/archive/cd2cf5c34931e4fc1cacf83bfc480a5b93f05f6d.tar.gz", - ], -) +# http_archive( +# name = "xla", +# patch_args = [ +# "-l", +# "-p1", +# ], +# patch_tool = "patch", +# patches = [ +# "//openxla_patches:cache_urls.diff", +# "//openxla_patches:f16_abi_clang.diff", +# "//openxla_patches:gpu_race_condition.diff", +# "//openxla_patches:constexpr_return.diff", +# ], +# strip_prefix = "xla-66ceee5d5967a6f0f24adde73be2e7cf0a0f67fd", +# urls = [ +# "https://github.com/openxla/xla/archive/66ceee5d5967a6f0f24adde73be2e7cf0a0f67fd.tar.gz", +# ], +# ) # For development, one often wants to make changes to the OpenXLA repository as well # as the PyTorch/XLA repository. You can override the pinned repository above with a @@ -57,10 +57,10 @@ http_archive( # bazel --override_repository=xla=/path/to/openxla # or # b) by commenting out the http_archive above and uncommenting the following: -# local_repository( -# name = "xla", -# path = "/path/to/openxla", -# ) +local_repository( + name = "xla", + path = "/workspaces/work/xla", +) # Initialize OpenXLA's external dependencies. load("@xla//:workspace4.bzl", "xla_workspace4") diff --git a/setup.py b/setup.py index a01a3ac843c..a62b1233a79 100644 --- a/setup.py +++ b/setup.py @@ -72,7 +72,7 @@ base_dir = os.path.dirname(os.path.abspath(__file__)) -_libtpu_version = '0.1.dev20230809' +_libtpu_version = '0.1.dev20230830' _libtpu_storage_path = f'https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/wheels/libtpu-nightly/libtpu_nightly-{_libtpu_version}-py3-none-any.whl' diff --git a/torch_xla/csrc/runtime/pjrt_computation_client.cc b/torch_xla/csrc/runtime/pjrt_computation_client.cc index e5a35030e66..8d8f7baad5e 100644 --- a/torch_xla/csrc/runtime/pjrt_computation_client.cc +++ b/torch_xla/csrc/runtime/pjrt_computation_client.cc @@ -29,7 +29,7 @@ #include "xla/pjrt/tfrt_cpu_pjrt_client.h" #include "xla/pjrt/tpu_client.h" #include "xla/shape.h" -#include "xla/stream_executor/tpu/tpu_initializer_framework_helper.h" +// #include "xla/stream_executor/tpu/tpu_initializer_framework_helper.h" using xla::internal::XlaBuilderFriend; From f169beb1fb1398e453750503ad3216a26870a2ac Mon Sep 17 00:00:00 2001 From: Jiewen Tan Date: Thu, 31 Aug 2023 00:46:56 +0000 Subject: [PATCH 2/7] tmp --- torch_xla/csrc/runtime/computation_client.h | 2 +- torch_xla/csrc/runtime/pjrt_computation_client.cc | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/torch_xla/csrc/runtime/computation_client.h b/torch_xla/csrc/runtime/computation_client.h index 2612e1b8ccd..a563836e2d8 100644 --- a/torch_xla/csrc/runtime/computation_client.h +++ b/torch_xla/csrc/runtime/computation_client.h @@ -43,7 +43,7 @@ class ComputationClient { const std::string& device() const { return device_; } - const xla::Shape& shape() const { return shape_; } + // const xla::Shape& shape() const { return shape_; } virtual OpaqueHandle GetOpaqueHandle() = 0; diff --git a/torch_xla/csrc/runtime/pjrt_computation_client.cc b/torch_xla/csrc/runtime/pjrt_computation_client.cc index 8d8f7baad5e..eea50940b78 100644 --- a/torch_xla/csrc/runtime/pjrt_computation_client.cc +++ b/torch_xla/csrc/runtime/pjrt_computation_client.cc @@ -29,7 +29,6 @@ #include "xla/pjrt/tfrt_cpu_pjrt_client.h" #include "xla/pjrt/tpu_client.h" #include "xla/shape.h" -// #include "xla/stream_executor/tpu/tpu_initializer_framework_helper.h" using xla::internal::XlaBuilderFriend; @@ -105,7 +104,7 @@ PjRtComputationClient::PjRtComputationClient() { TF_VLOG(1) << "Initializing TFRT TPU client..."; XLA_CHECK_OK(pjrt::LoadPjrtPlugin( "tpu", sys_util::GetEnvString(env::kEnvTpuLibraryPath, "libtpu.so"))); - tsl::Status tpu_status = tensorflow::tpu::FindAndLoadTpuLibrary(); + tsl::Status tpu_status = pjrt::InitializePjrtPlugin("tpu"); XLA_CHECK(tpu_status.ok()); client_ = std::move(xla::GetCApiClient("TPU").value()); } else if (device_type == "TPU_LEGACY") { From 3a329b8b748d61a525f2b319220818b12d0f6cef Mon Sep 17 00:00:00 2001 From: Jiewen Tan Date: Thu, 31 Aug 2023 01:57:18 +0000 Subject: [PATCH 3/7] pick a better commit --- WORKSPACE | 4 ++-- setup.py | 2 +- torch_xla/csrc/runtime/computation_client.h | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/WORKSPACE b/WORKSPACE index 7df2c2fbe65..72dc2ef08e3 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -43,9 +43,9 @@ python_configure( # "//openxla_patches:gpu_race_condition.diff", # "//openxla_patches:constexpr_return.diff", # ], -# strip_prefix = "xla-66ceee5d5967a6f0f24adde73be2e7cf0a0f67fd", +# strip_prefix = "xla-3b8a5398ea9e479311d5c0c34bb1ae255551f9d4", # urls = [ -# "https://github.com/openxla/xla/archive/66ceee5d5967a6f0f24adde73be2e7cf0a0f67fd.tar.gz", +# "https://github.com/openxla/xla/archive/3b8a5398ea9e479311d5c0c34bb1ae255551f9d4.tar.gz", # ], # ) diff --git a/setup.py b/setup.py index a62b1233a79..f3b90914144 100644 --- a/setup.py +++ b/setup.py @@ -72,7 +72,7 @@ base_dir = os.path.dirname(os.path.abspath(__file__)) -_libtpu_version = '0.1.dev20230830' +_libtpu_version = '0.1.dev20230826' _libtpu_storage_path = f'https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/wheels/libtpu-nightly/libtpu_nightly-{_libtpu_version}-py3-none-any.whl' diff --git a/torch_xla/csrc/runtime/computation_client.h b/torch_xla/csrc/runtime/computation_client.h index a563836e2d8..2612e1b8ccd 100644 --- a/torch_xla/csrc/runtime/computation_client.h +++ b/torch_xla/csrc/runtime/computation_client.h @@ -43,7 +43,7 @@ class ComputationClient { const std::string& device() const { return device_; } - // const xla::Shape& shape() const { return shape_; } + const xla::Shape& shape() const { return shape_; } virtual OpaqueHandle GetOpaqueHandle() = 0; From c5bf94e117a8156bbfedce387144231bdbbe63b8 Mon Sep 17 00:00:00 2001 From: Jiewen Tan Date: Thu, 31 Aug 2023 01:58:39 +0000 Subject: [PATCH 4/7] Revert local_repository --- WORKSPACE | 44 ++++++++++++++++++++++---------------------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/WORKSPACE b/WORKSPACE index 72dc2ef08e3..aa3de0a114f 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -30,24 +30,24 @@ python_configure( # b) get the sha256 hash of the commit by running: # curl -L https://github.com/openxla/xla/archive/.tar.gz | sha256sum # and update the sha256 with the result. -# http_archive( -# name = "xla", -# patch_args = [ -# "-l", -# "-p1", -# ], -# patch_tool = "patch", -# patches = [ -# "//openxla_patches:cache_urls.diff", -# "//openxla_patches:f16_abi_clang.diff", -# "//openxla_patches:gpu_race_condition.diff", -# "//openxla_patches:constexpr_return.diff", -# ], -# strip_prefix = "xla-3b8a5398ea9e479311d5c0c34bb1ae255551f9d4", -# urls = [ -# "https://github.com/openxla/xla/archive/3b8a5398ea9e479311d5c0c34bb1ae255551f9d4.tar.gz", -# ], -# ) +http_archive( + name = "xla", + patch_args = [ + "-l", + "-p1", + ], + patch_tool = "patch", + patches = [ + "//openxla_patches:cache_urls.diff", + "//openxla_patches:f16_abi_clang.diff", + "//openxla_patches:gpu_race_condition.diff", + "//openxla_patches:constexpr_return.diff", + ], + strip_prefix = "xla-3b8a5398ea9e479311d5c0c34bb1ae255551f9d4", + urls = [ + "https://github.com/openxla/xla/archive/3b8a5398ea9e479311d5c0c34bb1ae255551f9d4.tar.gz", + ], +) # For development, one often wants to make changes to the OpenXLA repository as well # as the PyTorch/XLA repository. You can override the pinned repository above with a @@ -57,10 +57,10 @@ python_configure( # bazel --override_repository=xla=/path/to/openxla # or # b) by commenting out the http_archive above and uncommenting the following: -local_repository( - name = "xla", - path = "/workspaces/work/xla", -) +# local_repository( +# name = "xla", +# path = "/path/to/openxla", +# ) # Initialize OpenXLA's external dependencies. load("@xla//:workspace4.bzl", "xla_workspace4") From d3ec960fd4ad6b95d4dcf3aef72e9ba7f110ca98 Mon Sep 17 00:00:00 2001 From: Jiewen Tan Date: Thu, 31 Aug 2023 02:34:26 +0000 Subject: [PATCH 5/7] Update the commit --- WORKSPACE | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/WORKSPACE b/WORKSPACE index aa3de0a114f..1b51229132e 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -43,9 +43,9 @@ http_archive( "//openxla_patches:gpu_race_condition.diff", "//openxla_patches:constexpr_return.diff", ], - strip_prefix = "xla-3b8a5398ea9e479311d5c0c34bb1ae255551f9d4", + strip_prefix = "xla-7a371ed44aba34f83d6d3d1159d2e6d0d327c603", urls = [ - "https://github.com/openxla/xla/archive/3b8a5398ea9e479311d5c0c34bb1ae255551f9d4.tar.gz", + "https://github.com/openxla/xla/archive/7a371ed44aba34f83d6d3d1159d2e6d0d327c603.tar.gz", ], ) From 7e65d3d0d0261d6f9686cda5dfb93e25f4a34837 Mon Sep 17 00:00:00 2001 From: Jiewen Tan Date: Fri, 8 Sep 2023 20:44:55 +0000 Subject: [PATCH 6/7] Increase the cpp test timeout to 1000s --- test/cpp/run_tests.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/cpp/run_tests.sh b/test/cpp/run_tests.sh index 89afbe1b9d2..92f4295af9a 100755 --- a/test/cpp/run_tests.sh +++ b/test/cpp/run_tests.sh @@ -87,7 +87,7 @@ fi if [ "$LOGFILE" != "" ]; then - bazel $BAZEL_VERB $EXTRA_FLAGS //torch_xla/csrc/runtime:all //test/cpp:all ${FILTER:+"$FILTER"} 2> $LOGFILE + bazel $BAZEL_VERB $EXTRA_FLAGS //torch_xla/csrc/runtime:all //test/cpp:all --test_timeout 600 ${FILTER:+"$FILTER"} 2> $LOGFILE else - bazel $BAZEL_VERB $EXTRA_FLAGS //torch_xla/csrc/runtime:all //test/cpp:all ${FILTER:+"$FILTER"} + bazel $BAZEL_VERB $EXTRA_FLAGS //torch_xla/csrc/runtime:all //test/cpp:all --test_timeout 600 ${FILTER:+"$FILTER"} fi From 55a08239b0cbdd5daebd15efca0f978a9369a560 Mon Sep 17 00:00:00 2001 From: Jiewen Tan Date: Fri, 8 Sep 2023 20:45:18 +0000 Subject: [PATCH 7/7] Increase the cpp test timeout to 1000s --- test/cpp/run_tests.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/cpp/run_tests.sh b/test/cpp/run_tests.sh index 92f4295af9a..a2865d41e5a 100755 --- a/test/cpp/run_tests.sh +++ b/test/cpp/run_tests.sh @@ -87,7 +87,7 @@ fi if [ "$LOGFILE" != "" ]; then - bazel $BAZEL_VERB $EXTRA_FLAGS //torch_xla/csrc/runtime:all //test/cpp:all --test_timeout 600 ${FILTER:+"$FILTER"} 2> $LOGFILE + bazel $BAZEL_VERB $EXTRA_FLAGS //torch_xla/csrc/runtime:all //test/cpp:all --test_timeout 1000 ${FILTER:+"$FILTER"} 2> $LOGFILE else - bazel $BAZEL_VERB $EXTRA_FLAGS //torch_xla/csrc/runtime:all //test/cpp:all --test_timeout 600 ${FILTER:+"$FILTER"} + bazel $BAZEL_VERB $EXTRA_FLAGS //torch_xla/csrc/runtime:all //test/cpp:all --test_timeout 1000 ${FILTER:+"$FILTER"} fi