diff --git a/.bazelrc b/.bazelrc index f938902fcacca7..c316e9524acfad 100644 --- a/.bazelrc +++ b/.bazelrc @@ -175,6 +175,13 @@ common --enable_workspace common --incompatible_enable_cc_toolchain_resolution common --repo_env USE_HERMETIC_CC_TOOLCHAIN=1 +# external_include_paths feature forces passing external repositories headers +# via -isystem alters how Clang resolves their file paths. This feature with +# old versions of rules_cc triggering errors. +# Following lines can be removed after rules_cc will be updated to >= 0.2.22 +common --features=-external_include_paths +common --host_features=-external_include_paths + # TODO: Migrate for https://github.com/bazelbuild/bazel/issues/7260 common:clang_local --noincompatible_enable_cc_toolchain_resolution common:clang_local --noincompatible_enable_android_toolchain_resolution diff --git a/MODULE.bazel b/MODULE.bazel index 61ea33e7a9a983..3f16e0882cae4f 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -36,7 +36,7 @@ bazel_dep(name = "google_cloud_cpp", version = "3.0.0-rc1", repo_name = "com_git bazel_dep(name = "crc32c", version = "1.1.0", repo_name = "com_github_google_crc32c") bazel_dep(name = "brotli", version = "1.1.0", repo_name = "org_brotli") bazel_dep(name = "rules_cc", version = "0.2.25") -bazel_dep(name = "curl", version = "8.11.0.bcr.5") +bazel_dep(name = "curl", version = "8.22.0") bazel_dep(name = "rules_webtesting", version = "0.4.1", repo_name = "io_bazel_rules_webtesting") bazel_dep(name = "rules_closure", version = "0.15.0", repo_name = "io_bazel_rules_closure") bazel_dep(name = "grpc", version = "1.78.0") @@ -86,9 +86,9 @@ bazel_dep(name = "libprotobuf-mutator", version = "1.3", repo_name = "libprotobu bazel_dep(name = "rules_ml_toolchain") archive_override( module_name = "rules_ml_toolchain", - integrity = "sha256-izPVqzV2I7ox7EYzRzcRY6ogeLcHUr4LRP2xMrkWh6A=", - strip_prefix = "rules_ml_toolchain-236a498d8624bef88e35ce2518a833204b0db19f", - urls = ["https://github.com/google-ml-infra/rules_ml_toolchain/archive/236a498d8624bef88e35ce2518a833204b0db19f.tar.gz"], + integrity = "sha256-Z38HhfFL91M5oxh8un+pXhofiIoYVo7cLl+sh9UT/Hw=", + strip_prefix = "rules_ml_toolchain-f56a19573fec3543050a5f4bcd9a07a1de08e8ee", + urls = ["https://github.com/google-ml-infra/rules_ml_toolchain/archive/f56a19573fec3543050a5f4bcd9a07a1de08e8ee.tar.gz"], ) bazel_dep(name = "xla", repo_name = "xla") diff --git a/RELEASE.md b/RELEASE.md index ff95217c3bb665..b64e7745610f17 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -3562,7 +3562,7 @@ This release introduces several vulnerability fixes: * Add `keepdims` argument to all `GlobalPooling` layers. * Add `include_preprocessing` argument to `MobileNetV3` architectures to control the inclusion of `Rescaling` layer in the model. - * Add optional argument (`force`) to `make_(train|test|predict)_funtion` + * Add optional argument (`force`) to `make_(train|test|predict)_function` methods to skip the cached function and generate a new one. This is useful to regenerate in a single call the compiled training function when any `.trainable` attribute of any model's layer has changed. diff --git a/ci/official/envs/ci_default b/ci/official/envs/ci_default index 21940c931b65c1..d9f5e3bf69c9fc 100644 --- a/ci/official/envs/ci_default +++ b/ci/official/envs/ci_default @@ -37,8 +37,10 @@ TFCI_BAZEL_COMMON_ARGS= TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX= TFCI_BUILD_PIP_PACKAGE_BASE_ARGS= TFCI_BUILD_PIP_PACKAGE_WHEEL_NAME_ARG= -# Mostly meant to be used for Windows. -# For Windows, the same wheel is uploaded to both tensorflow_cpu, and tensorflow +# Extra project names to publish the same wheel under, space separated. +# Used by the platforms that have no CUDA build to differentiate from, and so +# upload one wheel as both tensorflow and tensorflow_cpu: Windows, Linux arm64 +# and macOS arm64. TFCI_BUILD_PIP_PACKAGE_ADDITIONAL_WHEEL_NAMES= TFCI_DOCKER_ARGS= TFCI_DOCKER_ENABLE= diff --git a/ci/official/envs/linux_arm64 b/ci/official/envs/linux_arm64 index 724286027e0876..e19c06b1a0f856 100644 --- a/ci/official/envs/linux_arm64 +++ b/ci/official/envs/linux_arm64 @@ -18,6 +18,9 @@ TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=linux_arm64 # to tensorflow_cpu. These ARM builds are supposed to have the name "tensorflow" # despite lacking Nvidia CUDA support. TFCI_BUILD_PIP_PACKAGE_WHEEL_NAME_ARG="--repo_env=WHEEL_NAME=tensorflow" +# The same wheel is also published as tensorflow_cpu, so that the package name +# works on every platform TensorFlow ships. See the note in ci_default. +TFCI_BUILD_PIP_PACKAGE_ADDITIONAL_WHEEL_NAMES="tensorflow_cpu" TFCI_DOCKER_ENABLE=1 TFCI_DOCKER_IMAGE=us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-arm64:latest TFCI_DOCKER_PULL_ENABLE=1 diff --git a/ci/official/envs/macos_arm64 b/ci/official/envs/macos_arm64 index 96d8c14655cea8..c5224267eca3be 100644 --- a/ci/official/envs/macos_arm64 +++ b/ci/official/envs/macos_arm64 @@ -15,6 +15,9 @@ TFCI_BAZEL_COMMON_ARGS="--repo_env=HERMETIC_PYTHON_VERSION=$TFCI_PYTHON_VERSION --repo_env=USE_PYWRAP_RULES=True --config release_macos_arm64" TFCI_BAZEL_TARGET_SELECTING_CONFIG_PREFIX=macos_arm64 TFCI_BUILD_PIP_PACKAGE_WHEEL_NAME_ARG="--repo_env=WHEEL_NAME=tensorflow" +# The same wheel is also published as tensorflow_cpu, so that the package name +# works on every platform TensorFlow ships. See the note in ci_default. +TFCI_BUILD_PIP_PACKAGE_ADDITIONAL_WHEEL_NAMES="tensorflow_cpu" TFCI_INDEX_HTML_ENABLE=1 TFCI_LIB_SUFFIX="-cpu-darwin-arm64" TFCI_MACOS_BAZEL_TEST_DIR_ENABLE=1 diff --git a/ci/official/utilities/collect_additional_wheel.sh b/ci/official/utilities/collect_additional_wheel.sh new file mode 100755 index 00000000000000..b77be3be07ddb0 --- /dev/null +++ b/ci/official/utilities/collect_additional_wheel.sh @@ -0,0 +1,58 @@ +#!/usr/bin/env bash +# +# Copyright 2026 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# +# Usage: collect_additional_wheel.sh +# This script is aware of TFCI_ variables, so it only needs the project name of +# the wheel to collect. +# +# Copies the wheel that the additional-WHEEL_NAME build in wheel.sh just +# produced out of bazel-bin and into TFCI_OUTPUT_DIR, and gives it the same +# platform tag that rename_and_verify_wheels.sh gives the primary wheel. +# +# rename_and_verify_wheels.sh cannot do this itself: it expects exactly one +# wheel to be present, and it has already run against the primary wheel by the +# time this script is called. +set -exo pipefail + +wheel_name="$1" + +# `find -printf` is a GNU extension, and macOS ships BSD find, so let `ls -t` +# do the sorting instead and take the newest match. +built_wheel=$("$TFCI_FIND_BIN" ./bazel-bin/tensorflow/tools/pip_package \ + -iname "${wheel_name}*.whl" -exec ls -t {} + | head -n 1) +# find exits 0 when it matches nothing, so say what went wrong rather than +# leaving `cp` to report an empty path. +if [[ -z "$built_wheel" ]]; then + echo "Error: the ${wheel_name} build produced no wheel under bazel-bin" >&2 + exit 1 +fi +cp "$built_wheel" "$TFCI_OUTPUT_DIR" + +# Repair the wheel with auditwheel, the same way the primary wheel was +# repaired. Bazel tags Linux wheels `linux_`, and PyPI only accepts the +# `manylinux_*` tag that auditwheel puts on in its place. +if [[ "$TFCI_WHL_AUDIT_ENABLE" == "1" ]]; then + cd "$TFCI_OUTPUT_DIR" + python3 -m auditwheel repair --plat "$TFCI_WHL_AUDIT_PLAT" --wheel-dir . \ + "$(basename "$built_wheel")" + # Drop the unrepaired wheel, keeping the newest one, the same way + # rename_and_verify_wheels.sh does for the primary wheel. If the wheel was + # already named correctly then auditwheel did not rename it, nothing is left + # over and `rm` gets no arguments, which is why it needs -f. Matching on the + # project name leaves the primary wheel alone. + ls -t "${wheel_name}"-*.whl | tail -n +2 | xargs rm -f +fi diff --git a/ci/official/wheel.sh b/ci/official/wheel.sh index b1327912e037f5..abcc5584e8eef3 100755 --- a/ci/official/wheel.sh +++ b/ci/official/wheel.sh @@ -62,8 +62,9 @@ if [[ -n "$TFCI_BUILD_PIP_PACKAGE_ADDITIONAL_WHEEL_NAMES" ]]; then echo "Building for additional WHEEL_NAME: ${wheel_name}" CURRENT_WHEEL_NAME_ARG="--repo_env=WHEEL_NAME=${wheel_name}" tfrun bazel $TFCI_BAZEL_BAZELRC_ARGS build $TFCI_BAZEL_COMMON_ARGS --config=cuda_wheel //tensorflow/tools/pip_package:wheel $TFCI_BUILD_PIP_PACKAGE_BASE_ARGS $CURRENT_WHEEL_NAME_ARG $HERMETIC_CUDA_UMD_BUILD_FLAGS - # Copy the wheel that was just created - tfrun bash -c "$TFCI_FIND_BIN ./bazel-bin/tensorflow/tools/pip_package -iname "${wheel_name}*.whl" -printf '%T+ %p\n' | sort | tail -n 1 | awk '{print \$2}' | xargs -I {} cp {} $TFCI_OUTPUT_DIR" + # Copy the wheel that was just created, and platform-tag it to match the + # primary wheel. + tfrun bash ./ci/official/utilities/collect_additional_wheel.sh "${wheel_name}" done fi diff --git a/tensorflow/compiler/aot/BUILD b/tensorflow/compiler/aot/BUILD index add8bc7291af36..d0877d9af1ac55 100644 --- a/tensorflow/compiler/aot/BUILD +++ b/tensorflow/compiler/aot/BUILD @@ -148,7 +148,6 @@ cc_library( "@xla//xla/client:client_library", "@xla//xla/client:compile_only_client", "@xla//xla/hlo/builder:xla_computation", - "@xla//xla/pjrt:pjrt_compiler", "@xla//xla/service:compiler", "@xla//xla/service/cpu:cpu_aot_compilation_result", "@xla//xla/service/cpu:cpu_compiler", diff --git a/tensorflow/compiler/aot/compile.cc b/tensorflow/compiler/aot/compile.cc index 792ad9bcc5cad5..efc6869f6235de 100644 --- a/tensorflow/compiler/aot/compile.cc +++ b/tensorflow/compiler/aot/compile.cc @@ -42,7 +42,6 @@ limitations under the License. #include "xla/client/client_library.h" #include "xla/client/compile_only_client.h" #include "xla/hlo/builder/xla_computation.h" -#include "xla/pjrt/pjrt_compiler.h" #include "xla/service/compiler.h" #include "xla/service/cpu/cpu_aot_compilation_result.h" #include "xla/shape.h" @@ -175,10 +174,8 @@ absl::Status CompileGraph(GraphDef graph_def, const tf2xla::Config& config, graph_def, config, &computation, flags.debug_info, flags.debug_info_path_begin_marker)); } else { - TF_ASSIGN_OR_RETURN(auto* pjrt_compiler, - xla::GetDefaultPjRtCompiler(xla::CpuName())); TF_RETURN_IF_ERROR(ConvertGraphDefToXla(std::move(graph_def), config, - pjrt_compiler, &computation)); + client, &computation)); } if (flags.experimental_quantize && *quantize_xla) { diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index b9410f089772bc..a93dc10be415fd 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -156,7 +156,6 @@ cc_library( "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:lib", - "//tensorflow/core/tfrt/common:pjrt_cpu_client_registration", "@com_google_absl//absl/memory", "@xla//xla/stream_executor:platform_manager", ] + if_libtpu( @@ -292,6 +291,7 @@ cc_library( # Public visibility is needed for external TF/XLA backends. visibility = ["//visibility:public"], deps = [ + ":common", ":device_compilation_cluster_signature", ":device_compilation_profiler", ":device_compiler", @@ -304,11 +304,13 @@ cc_library( ":variable_info_util", ":xla_compile_util", ":xla_device_compiler_client", + ":xla_device_context", ":xla_launch_util", ":xla_tensor", "//tensorflow/compiler/jit/ops:xla_ops", "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:layout_util", + "//tensorflow/compiler/tf2xla:tf2xla_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops", @@ -344,6 +346,7 @@ cc_library( "//tensorflow/core/kernels/data:options_dataset_op", "//tensorflow/core/kernels/data:prefetch_dataset_op", "//tensorflow/core/platform:stream_executor_no_cuda", + "//tensorflow/core/profiler/lib:traceme", "//tensorflow/core/tfrt/common:async_value_tensor", "//tensorflow/core/tfrt/common:create_pjrt_client_util", "//tensorflow/core/tfrt/common:global_state", @@ -351,21 +354,28 @@ cc_library( "//tensorflow/core/tpu:tpu_defs", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", + "@com_google_absl//absl/memory", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", - "@xla//xla:shape_util", + "@xla//xla:executable_run_options", "@xla//xla/client:client_library", "@xla//xla/client:local_client", + "@xla//xla/hlo/ir:hlo", "@xla//xla/pjrt:pjrt_client", + "@xla//xla/pjrt:tf_pjrt_client", "@xla//xla/service:compiler", - "@xla//xla/service:shaped_buffer", - "@xla//xla/stream_executor:platform", + "@xla//xla/service:executable", + "@xla//xla/service:stream_pool", + "@xla//xla/service/gpu:gpu_executable_run_options", "@xla//xla/stream_executor:platform_manager", + "@xla//xla/stream_executor:stream", + "@xla//xla/stream_executor/host:host_platform_id", "@xla//xla/stream_executor/integrations:tf_allocator_adapter", "@xla//xla/tsl/framework:device_id_utils", ], @@ -393,9 +403,11 @@ cc_library( deps = [ ":device_compilation_cluster_signature", ":device_compilation_profiler", - ":device_compiler", ":jit_compilation_passes", + ":variable_info", + ":variable_info_util", ":xla_device_no_jit_rewrite_registration", + ":xla_launch_util", ":xla_tensor", "//tensorflow/compiler/jit/ops:xla_ops", "//tensorflow/compiler/tf2xla:layout_util", @@ -434,7 +446,6 @@ cc_library( "//tensorflow/core/kernels/data:prefetch_dataset_op", "//tensorflow/core/platform:stream_executor_no_cuda", "//tensorflow/core/tfrt/common:async_value_tensor", - "@com_google_absl//absl/status", "@com_google_absl//absl/types:optional", "@xla//xla/client:local_client", "@xla//xla/pjrt:pjrt_client", @@ -465,12 +476,13 @@ cc_library( "//tensorflow/compiler/mlir/tensorflow:dump_graph", "//tensorflow/core:framework_internal", "//tensorflow/core:framework_types_hdr", + "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/tpu/kernels:sparse_core_xla_flags_defaults", "@com_google_absl//absl/base", "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/log", "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:optional", "@xla//xla:parse_flags_from_env", ], ) @@ -484,8 +496,10 @@ cc_library( deps = [ "//tensorflow/core:framework_internal", "//tensorflow/core:framework_types_hdr", + "//tensorflow/core:lib", "//tensorflow/core/protobuf:for_core_protos_cc", "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/types:optional", ], ) @@ -615,6 +629,7 @@ cc_library( ":pjrt_tensor_buffer_util", ":variable_info", ":variable_info_util", + ":xla_tensor", "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_resource", @@ -629,7 +644,6 @@ cc_library( "//tensorflow/core/common_runtime/gpu:gpu_serving_device_selector", "//tensorflow/core/tfrt/common:async_value_tensor", "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/base", "@com_google_absl//absl/cleanup", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", @@ -646,13 +660,18 @@ cc_library( "@xla//xla/client:local_client", "@xla//xla/hlo/ir:hlo", "@xla//xla/pjrt:pjrt_client", - "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/pjrt:pjrt_common", "@xla//xla/pjrt:pjrt_executable", - "@xla//xla/pjrt/plugin/xla_cpu:cpu_execute_options", - "@xla//xla/runtime:device_id", "@xla//xla/service:executable", + "@xla//xla/service:maybe_owning_device_memory", "@xla//xla/service:shaped_buffer", - "@xla//xla/stream_executor:device_address_allocator", + "@xla//xla/service:transfer_manager", + "@xla//xla/stream_executor:device_memory", + "@xla//xla/stream_executor:device_memory_allocator", + "@xla//xla/stream_executor:event", + "@xla//xla/stream_executor:platform", + "@xla//xla/stream_executor:platform_manager", + "@xla//xla/stream_executor/host:host_platform_id", "@xla//xla/tsl/framework:device_id_utils", "@xla//xla/tsl/framework:serving_device_selector_policies", ], @@ -856,41 +875,30 @@ cc_library( visibility = [":internal"], deps = [ ":compilability_check_util", - ":device_compilation_profiler", ":device_compiler", - ":pjrt_device_compiler_client", ":variable_info", ":variable_info_util", ":xla_device_no_jit_rewrite_registration", ":xla_launch_util", "//tensorflow/compiler/tf2xla:xla_compiler", - "//tensorflow/compiler/tf2xla:xla_resource", "//tensorflow/core:framework", + "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/common_runtime:core_cpu_internal", "//tensorflow/core/common_runtime/eager:tensor_handle", "//tensorflow/core/platform:refcount", - "//tensorflow/core/platform:statusor", - "//tensorflow/core/tpu:tpu_defs", "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", - "@xla//xla:shape_util", "@xla//xla:status_macros", - "@xla//xla/hlo/ir:hlo", + "@xla//xla/client:executable_build_options", + "@xla//xla/client:local_client", "@xla//xla/hlo/translate:portable_api", - "@xla//xla/pjrt:compiled_memory_stats", - "@xla//xla/pjrt:pjrt_client", "@xla//xla/service:hlo_graph_dumper", - "@xla//xla/service:hlo_module_config", - "@xla//xla/service:hlo_proto_util", "@xla//xla/stream_executor:platform", "@xla//xla/stream_executor/host:host_platform_id", - "@xla//xla/tsl/platform:errors", - "@xla//xla/tsl/platform:statusor", ], alwayslink = 1, ) @@ -1976,12 +1984,9 @@ tf_cuda_cc_test( ], deps = [ ":device_compilation_cluster_signature", - ":device_compilation_profiler", ":device_compiler", ":device_compiler_client", - ":device_executable_persistor", ":tf_graph_to_hlo_compiler_proto_cc", - ":xla_activity_listener", ":xla_compile_util", ":xla_device_compiler_client", ":xla_gpu_device", @@ -1999,17 +2004,14 @@ tf_cuda_cc_test( "//tensorflow/core/kernels:ops_testutil", "//tensorflow/core/platform:env", "//tensorflow/core/platform:errors", - "//tensorflow/core/platform:path", "//tensorflow/core/platform:refcount", "//tensorflow/core/platform:status", "//tensorflow/core/platform:status_matchers", "//tensorflow/core/platform:statusor", "//tensorflow/core/platform:test", "@com_google_absl//absl/algorithm:container", - "@com_google_absl//absl/base", "@com_google_absl//absl/functional:any_invocable", "@com_google_absl//absl/log:check", - "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", "@com_google_absl//absl/time", @@ -2020,7 +2022,6 @@ tf_cuda_cc_test( "@xla//xla/hlo/builder:xla_computation", "@xla//xla/stream_executor:platform_manager", "@xla//xla/tsl/lib/strings:proto_serialization", - "@xla//xla/tsl/platform:errors", "@xla//xla/tsl/platform:statusor", ], ) @@ -2074,10 +2075,8 @@ tf_cuda_cc_test( "@xla//xla/client:client_library", "@xla//xla/client:local_client", "@xla//xla/pjrt:pjrt_client", - "@xla//xla/service/cpu:cpu_transfer_manager", - "@xla//xla/stream_executor:device_address_allocator", + "@xla//xla/stream_executor:device_memory_allocator", "@xla//xla/stream_executor:platform", - "@xla//xla/stream_executor:platform_manager", "@xla//xla/stream_executor/host:host_platform_id", "@xla//xla/tsl/platform:statusor", ], diff --git a/tensorflow/compiler/jit/device_compiler_test.cc b/tensorflow/compiler/jit/device_compiler_test.cc index 68b6652857da6d..d42f797526be83 100644 --- a/tensorflow/compiler/jit/device_compiler_test.cc +++ b/tensorflow/compiler/jit/device_compiler_test.cc @@ -15,7 +15,7 @@ limitations under the License. #include "tensorflow/compiler/jit/device_compiler.h" -#include +#include #include #include #include @@ -26,10 +26,7 @@ limitations under the License. #include #include #include "absl/algorithm/container.h" -#include "absl/base/casts.h" #include "absl/log/check.h" -#include "absl/status/status.h" -#include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/synchronization/notification.h" @@ -39,37 +36,32 @@ limitations under the License. #include "tensorflow/cc/ops/function_ops.h" #include "tensorflow/cc/ops/math_ops.h" #include "tensorflow/compiler/jit/device_compilation_cluster_signature.h" -#include "tensorflow/compiler/jit/device_compilation_profiler.h" #include "tensorflow/compiler/jit/device_compiler_client.h" -#include "tensorflow/compiler/jit/device_executable_persistor.h" #include "tensorflow/compiler/jit/tests/device_compiler_test_helper.h" #include "tensorflow/compiler/jit/tf_graph_to_hlo_compiler.pb.h" -#include "tensorflow/compiler/jit/xla_activity_listener.h" #include "tensorflow/compiler/jit/xla_compile_util.h" #include "tensorflow/compiler/jit/xla_device_compiler_client.h" #include "tensorflow/compiler/tf2xla/xla_argument.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/client/client_library.h" #include "xla/client/local_client.h" +#include "xla/hlo/builder/xla_computation.h" #include "xla/stream_executor/platform_manager.h" #include "xla/tsl/lib/core/status_test_util.h" #include "xla/tsl/lib/strings/proto_serialization.h" -#include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/statusor.h" #include "tensorflow/core/framework/fake_input.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/graph_to_functiondef.h" #include "tensorflow/core/framework/node_def_builder.h" -#include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/resource_base.h" -#include "tensorflow/core/framework/tensor_shape.h" -#include "tensorflow/core/framework/types.h" #include "tensorflow/core/kernels/ops_testutil.h" +#include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/file_system.h" -#include "tensorflow/core/platform/path.h" -#include "tensorflow/core/platform/refcount.h" +#include "tensorflow/core/platform/errors.h" +#include "tensorflow/core/platform/status.h" +#include "tensorflow/core/platform/status_matchers.h" +#include "tensorflow/core/platform/statusor.h" #include "tensorflow/core/platform/test.h" #include "tsl/platform/path.h" @@ -194,6 +186,7 @@ class DeviceCompilerTest : public ::testing::Test { XlaCompiler::Options GetDefaultXlaOptions() { XlaCompiler::Options options; options.device_type = DeviceType(DEVICE_GPU_XLA_JIT); + options.client = xla_device_compiler_->client(); options.flib_def = flib_def_.get(); return options; } @@ -628,6 +621,7 @@ TEST_F(OpsTestBase, CompileSingleOpSuccess) { XlaCompiler::Options options; options.device_type = DeviceType(DEVICE_GPU_XLA_JIT); + options.client = GetLocalClient(); options.flib_def = flib_def.get(); std::vector args(1); diff --git a/tensorflow/compiler/jit/device_context_test.cc b/tensorflow/compiler/jit/device_context_test.cc index 0b3701db0132bb..33bba30f3db3e1 100644 --- a/tensorflow/compiler/jit/device_context_test.cc +++ b/tensorflow/compiler/jit/device_context_test.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include +#include #include #include @@ -38,6 +39,10 @@ static bool Initialized = [] { class DeviceContextTest : public ::testing::Test { public: void SetDevice(const std::string& device_type) { + auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api; + rollout_config.AllowForDeviceInXlaLaunch(DeviceType(device_type)); + rollout_config.AllowForDeviceInXlaCompileOnDemand(DeviceType(device_type)); + auto device_factory = DeviceFactory::GetFactory(device_type); SessionOptions options; std::vector> devices; diff --git a/tensorflow/compiler/jit/device_executable_persistor_test.cc b/tensorflow/compiler/jit/device_executable_persistor_test.cc index 576dc2462894fe..c33dcb9538ffd6 100644 --- a/tensorflow/compiler/jit/device_executable_persistor_test.cc +++ b/tensorflow/compiler/jit/device_executable_persistor_test.cc @@ -122,6 +122,7 @@ class DeviceExecutionPersistorTest : public ::testing::Test { XlaCompiler::Options DefaultXlaOptions() { XlaCompiler::Options options; options.device_type = DeviceType(DEVICE_CPU_XLA_JIT); + options.client = xla_compiler_client_->client(); options.flib_def = flib_def_.get(); return options; } @@ -129,9 +130,7 @@ class DeviceExecutionPersistorTest : public ::testing::Test { XlaCompiler::Options DefaultPjRtOptions() { XlaCompiler::Options options; options.device_type = DeviceType(DEVICE_CPU_XLA_JIT); - if (pjrt_compiler_client_ != nullptr) { - options.client = pjrt_compiler_client_->client(); - } + options.client = nullptr; options.flib_def = flib_def_.get(); return options; } diff --git a/tensorflow/compiler/jit/flags.cc b/tensorflow/compiler/jit/flags.cc index a65f1ae7c7edf8..a0a0d45736f1e8 100644 --- a/tensorflow/compiler/jit/flags.cc +++ b/tensorflow/compiler/jit/flags.cc @@ -15,22 +15,18 @@ limitations under the License. #include "tensorflow/compiler/jit/flags.h" -#include -#include #include +#include // NOLINT #include -#include #include #include "absl/base/call_once.h" -#include "absl/log/log.h" #include "absl/strings/numbers.h" #include "absl/strings/str_split.h" -#include "absl/strings/string_view.h" #include "absl/strings/strip.h" #include "tensorflow/compiler/mlir/tensorflow/utils/dump_graph.h" #include "xla/parse_flags_from_env.h" -#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/platform/macros.h" #include "tensorflow/core/tpu/kernels/sparse_core_xla_flags_defaults.h" #include "tensorflow/core/util/command_line_flags.h" @@ -274,12 +270,6 @@ void AllocateAndParseFlags() { ops_flags->tf_xla_use_device_api.enabled_for_compile_and_run_ = true; ops_flags->tf_xla_use_device_api.enabled_for_all_ = false; ops_flags->tf_xla_use_device_api.enabled_for_gpu_ = true; - ops_flags->tf_xla_use_device_api.AllowForDeviceInXlaLaunch( - DeviceType(DEVICE_GPU)); - ops_flags->tf_xla_use_device_api.AllowForDeviceInXlaCompileOnDemand( - DeviceType(DEVICE_GPU)); - ops_flags->tf_xla_use_device_api.AllowForDeviceInXlaCompileAndRun( - DeviceType(DEVICE_GPU)); call_module_flags = new XlaCallModuleFlags; // The `enable_mlir_bridge` flag allows the user to explicitly request that @@ -424,11 +414,6 @@ void AllocateAndParseFlags() { AppendMarkForCompilationPassFlagsInternal(flag_list); xla::ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", *flag_list); - if (!ops_flags->tf_xla_use_device_api.enabled_for_gpu_) { - LOG(WARNING) - << "tf_xla_enable_device_api_for_gpu is ignored; assuming true"; - } - mlir_flags = new MlirCommonFlags; if (!enable_mlir_bridge_is_explicit) { mlir_flags->tf_mlir_enable_mlir_bridge = diff --git a/tensorflow/compiler/jit/flags.h b/tensorflow/compiler/jit/flags.h index a2a1b810444dc4..96154b892ae5b0 100644 --- a/tensorflow/compiler/jit/flags.h +++ b/tensorflow/compiler/jit/flags.h @@ -22,7 +22,9 @@ limitations under the License. #include #include "absl/container/flat_hash_set.h" +#include "absl/types/optional.h" #include "tensorflow/core/framework/types.h" +#include "tensorflow/core/platform/types.h" #include "tensorflow/core/protobuf/config.pb.h" #include "tensorflow/core/util/command_line_flags.h" @@ -88,7 +90,7 @@ struct MarkForCompilationPassFlags { bool tf_xla_disable_resource_variable_safety_checks_for_debugging; // If true names of clustered operations will be computed deterministically - // so that they remain stable from run to run of auto clustering. + // so that they remain stable from run to run of auto clusteing. bool tf_xla_deterministic_cluster_names; // If non-empty, JIT-compiled executables are saved to and loaded from the @@ -107,7 +109,7 @@ struct MarkForCompilationPassFlags { // for debugging. Defaults to false. bool tf_xla_disable_strict_signature_checks; - // Specifies the persistence cache prefix. Default is "xla_compile_cache" + // Specifies the persistance cache prefix. Default is "xla_compile_cache" std::string tf_xla_persistent_cache_prefix; }; @@ -161,9 +163,7 @@ struct XlaOpsCommonFlags { } bool IsEnabledInXlaLaunchForDevice(const DeviceType& device_type) const { - if (device_type.type_string() == "GPU") { - return true; - } + if (!enabled_for_gpu_ && device_type.type_string() == "GPU") return false; return enabled_for_all_ || (enabled_for_xla_launch_ && xla_launch_allowed_devices_.contains(device_type.type_string())); @@ -180,9 +180,7 @@ struct XlaOpsCommonFlags { bool IsEnabledInXlaCompileOnDemandForDevice( const DeviceType& device_type) const { - if (device_type.type_string() == "GPU") { - return true; - } + if (!enabled_for_gpu_ && device_type.type_string() == "GPU") return false; return enabled_for_all_ || (enabled_for_compile_on_demand_ && xla_compile_on_demand_allowed_devices_.contains( @@ -200,14 +198,14 @@ struct XlaOpsCommonFlags { bool IsEnabledInXlaCompileAndRunForDevice( const DeviceType& device_type) const { - if (device_type.type_string() == "GPU") { - return true; - } + if (!enabled_for_gpu_ && device_type.type_string() == "GPU") return false; return enabled_for_all_ || (enabled_for_compile_and_run_ && xla_compile_and_run_allowed_devices_.contains( device_type.type_string())); } + bool IsEnabledForGpu() const { return enabled_for_gpu_; } + // If true, uses Device API (PjRt) for single device compilation and // execution of functions marked for JIT compilation i.e. jit_compile=True. // Defaults to false. @@ -279,7 +277,7 @@ struct BuildXlaOpsPassFlags { // TensorCore sequencing will be used. bool tf_xla_disable_full_embedding_pipelining; - // Whether to enable automatic embedding pipelining when summary ops are + // Whether to enable automatical embedding pipelining when summary ops are // detected in the graph. bool tf_xla_disable_full_embedding_pipelining_with_summaries; diff --git a/tensorflow/compiler/jit/get_compiler_ir.cc b/tensorflow/compiler/jit/get_compiler_ir.cc index 88635fb9aa122f..e50869f2870e5c 100644 --- a/tensorflow/compiler/jit/get_compiler_ir.cc +++ b/tensorflow/compiler/jit/get_compiler_ir.cc @@ -19,40 +19,31 @@ limitations under the License. #include #include #include +#include #include #include "absl/algorithm/container.h" -#include "absl/container/flat_hash_map.h" #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "absl/strings/ascii.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "tensorflow/compiler/jit/compilability_check_util.h" -#include "tensorflow/compiler/jit/device_compilation_profiler.h" #include "tensorflow/compiler/jit/device_compiler.h" -#include "tensorflow/compiler/jit/pjrt_device_compiler_client.h" #include "tensorflow/compiler/jit/variable_info.h" #include "tensorflow/compiler/jit/variable_info_util.h" #include "tensorflow/compiler/jit/xla_compiler_options_util.h" #include "tensorflow/compiler/jit/xla_launch_util.h" #include "tensorflow/compiler/jit/xla_platform_info.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "tensorflow/compiler/tf2xla/xla_resource.h" -#include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/ir/hlo_print_options.h" +#include "xla/client/executable_build_options.h" +#include "xla/client/local_client.h" #include "xla/hlo/translate/portable_api.h" -#include "xla/pjrt/compiled_memory_stats.h" -#include "xla/pjrt/pjrt_client.h" #include "xla/service/hlo_graph_dumper.h" -#include "xla/service/hlo_module_config.h" -#include "xla/service/hlo_proto_util.h" -#include "xla/shape.h" #include "xla/status_macros.h" #include "xla/stream_executor/host/host_platform_id.h" #include "xla/stream_executor/platform.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" #include "tensorflow/core/common_runtime/eager/tensor_handle.h" #include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/function.h" @@ -63,27 +54,49 @@ limitations under the License. #include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/platform/refcount.h" -#include "tensorflow/core/tpu/tpu_defs.h" +#include "tensorflow/core/platform/statusor.h" +#include "tsl/platform/errors.h" +#include "tsl/platform/statusor.h" namespace tensorflow { -static absl::StatusOr> -BuildExecutable(xla::PjRtClient* pjrt_client, - const XlaCompiler::CompilationResult& result, - const XlaCompiler::Options& options, - const bool xla_embed_ir_in_executable = false) { - xla::CompileOptions compile_options = GetPjRtCompileOptions(options, result); - compile_options.executable_build_options.mutable_debug_options() - ->set_xla_embed_ir_in_executable(xla_embed_ir_in_executable); +static absl::StatusOr> BuildExecutable( + xla::LocalClient* local_client, + const XlaCompiler::CompilationResult& result, + const XlaCompiler::Options& options, + const bool xla_embed_ir_in_executable = false) { + std::vector argument_layouts( + result.xla_input_shapes.size()); + for (int i = 0, end = result.xla_input_shapes.size(); i < end; ++i) { + argument_layouts[i] = &result.xla_input_shapes[i]; + } + xla::ExecutableBuildOptions build_options; + if (result.collective_info) { + build_options.set_num_replicas(result.collective_info->group_size); + } + build_options.set_device_ordinal( + options.device_ordinal != -1 ? options.device_ordinal + : local_client->default_device_ordinal()); + build_options.set_result_layout(result.xla_output_shape); + build_options.set_device_allocator(options.device_allocator.get()); + build_options.set_alias_passthrough_params(options.alias_passthrough_params); + build_options.mutable_debug_options()->set_xla_detailed_logging( + options.detailed_logging); + // If the embed_ir_in_executable is set, hlo_proto will be dumped in + // executable. The hlo_proto contains HLO modules and buffer assignment. + build_options.mutable_debug_options()->set_xla_embed_ir_in_executable( + xla_embed_ir_in_executable); TF_ASSIGN_OR_RETURN( - std::unique_ptr executable, - pjrt_client->CompileAndLoad(*result.computation, compile_options)); - return executable; + std::vector> executables, + local_client->Compile(*result.computation, argument_layouts, + build_options)); + TF_RET_CHECK(executables.size() == 1); + return std::move(executables[0]); } static absl::StatusOr BuildHLOString( IrExportStage stage, const XlaCompiler::CompilationResult& result, - xla::PjRtClient* pjrt_client, const XlaCompiler::Options& options) { + xla::LocalClient* local_client, const XlaCompiler::Options& options) { switch (stage) { case IrExportStage::STABLEHLO: case IrExportStage::STABLEHLO_SERIALIZED: @@ -122,41 +135,33 @@ static absl::StatusOr BuildHLOString( return new_module->ToString(opts); } case IrExportStage::OPTIMIZED_HLO: - case IrExportStage::OPTIMIZED_HLO_SERIALIZED: - case IrExportStage::OPTIMIZED_HLO_PROTO_SERIALIZED: - case IrExportStage::OPTIMIZED_HLO_DOT: { - TF_RET_CHECK(pjrt_client != nullptr); - bool xla_embed_ir_in_executable = - (stage == IrExportStage::OPTIMIZED_HLO_PROTO_SERIALIZED); - TF_ASSIGN_OR_RETURN(std::unique_ptr executable, - BuildExecutable(pjrt_client, result, options, - xla_embed_ir_in_executable)); - TF_ASSIGN_OR_RETURN(std::vector> modules, - executable->GetHloModules()); - TF_RET_CHECK(!modules.empty()); + case IrExportStage::OPTIMIZED_HLO_SERIALIZED: { + TF_ASSIGN_OR_RETURN(std::unique_ptr executable, + BuildExecutable(local_client, result, options)); + xla::Executable* new_executable = executable->executable(); if (stage == IrExportStage::OPTIMIZED_HLO_SERIALIZED) { - return modules[0]->ToProto().SerializeAsString(); - } else if (stage == IrExportStage::OPTIMIZED_HLO_PROTO_SERIALIZED) { - xla::HloProto hlo_proto = xla::MakeHloProto(*modules[0]); - absl::StatusOr memory_stats = - executable->GetCompiledMemoryStats(); - if (memory_stats.ok() && - !memory_stats->serialized_buffer_assignment.empty()) { - hlo_proto.mutable_buffer_assignment()->ParseFromString( - memory_stats->serialized_buffer_assignment); - } - return hlo_proto.SerializeAsString(); - } else if (stage == IrExportStage::OPTIMIZED_HLO_DOT) { - absl::StatusOr graph = xla::RenderGraph( - *modules[0]->entry_computation(), "Visualization", - /*debug_options=*/{}, xla::RenderedGraphFormat::kDot, - /*hlo_render_options=*/{}); - TF_RETURN_IF_ERROR(graph.status()); - return *graph; + return new_executable->module().ToProto().SerializeAsString(); } else { - return modules[0]->ToString(); + return new_executable->module().ToString(); } } + case IrExportStage::OPTIMIZED_HLO_PROTO_SERIALIZED: { + TF_ASSIGN_OR_RETURN(std::unique_ptr executable, + BuildExecutable(local_client, result, options, + /*xla_embed_ir_in_executable=*/true)); + return executable->executable()->hlo_proto()->SerializeAsString(); + } + case IrExportStage::OPTIMIZED_HLO_DOT: { + TF_ASSIGN_OR_RETURN(std::unique_ptr executable, + BuildExecutable(local_client, result, options)); + absl::StatusOr graph = xla::RenderGraph( + *executable->executable()->module().entry_computation(), + "Visualization", + /*debug_options=*/{}, xla::RenderedGraphFormat::kDot, + /*hlo_render_options=*/{}); + TF_RETURN_IF_ERROR(graph.status()); + return *graph; + } } } @@ -313,7 +318,7 @@ absl::StatusOr> PrepareXlaCompilerArgs( absl::StatusOr CompileAndBuildHLOString( IrExportStage stage, const XlaCompiler::Options& options, - xla::PjRtClient* pjrt_client, const NameAttrList& function, + xla::LocalClient* local_client, const NameAttrList& function, const std::vector& args) { XlaCompiler::CompileOptions compile_options; compile_options.always_return_tuple = false; @@ -324,7 +329,7 @@ absl::StatusOr CompileAndBuildHLOString( TF_RETURN_IF_ERROR( compiler.CompileFunction(compile_options, function, args, &result)); - return BuildHLOString(stage, result, pjrt_client, options); + return BuildHLOString(stage, result, local_client, options); } /** @@ -347,8 +352,8 @@ absl::StatusOr GetCompilerIr( absl::Span input_arg_shape_and_dtype, absl::Span input_handles, CompilerArgSource compiler_arg_source) { - using PjRtDeviceCompiler = - DeviceCompiler; + using XlaDeviceCompiler = + DeviceCompiler; se::Stream* stream = nullptr; if (const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info = @@ -370,32 +375,33 @@ absl::StatusOr GetCompilerIr( compiler_arg_source)); XlaPlatformInfo platform_info = XlaPlatformInfoFromDevice(dev); - DeviceType compilation_device_type = platform_info.device_type(); - if (platform_info.device_type() == DEVICE_TPU) { - compilation_device_type = DeviceType(DEVICE_TPU_XLA_JIT); - } else { + auto compilation_device_type = platform_info.device_type(); + if (platform_info.device_type() != DEVICE_TPU) { TF_ASSIGN_OR_RETURN(compilation_device_type, GetCompilationDeviceType(platform_info.device_type())); } - PjRtDeviceCompiler* pjrt_device_compiler = nullptr; - DeviceCompilationProfiler* profiler = nullptr; - TF_RETURN_IF_ERROR(GetOrCreatePjRtDeviceCompilerAndProfiler( - platform_info, dev->resource_manager(), flr, &pjrt_device_compiler, - &profiler)); - core::ScopedUnref pjrt_device_compiler_ref(pjrt_device_compiler); - core::ScopedUnref profiler_ref(profiler); + XlaDeviceCompiler* xla_device_compiler; + TF_RETURN_IF_ERROR(dev->resource_manager()->LookupOrCreate( + dev->resource_manager()->default_container(), "xla_device_compiler", + &xla_device_compiler, [&](XlaDeviceCompiler** xla_device_compiler) { + return BuildXlaDeviceCompiler(dev, flr, platform_info, + compilation_device_type, + xla_device_compiler); + })); + core::ScopedUnref xla_device_compiler_ref(xla_device_compiler); XlaCompiler::Options options; if (platform_info.device_type() == DEVICE_TPU && stream == nullptr) { - options = GenerateCompilerOptionsForTfrtTpu(compilation_device_type, *flr); + options = GenerateCompilerOptionsForTfrtTpu(*xla_device_compiler, *flr); } else { - options = GenerateCompilerOptionsForPjRt(*flr, dev, platform_info, - pjrt_device_compiler); + options = GenerateCompilerOptions(*xla_device_compiler, *flr, dev, stream, + platform_info, + /*has_ref_vars=*/false); } - return CompileAndBuildHLOString( - stage, options, pjrt_device_compiler->client(), function, args); + return CompileAndBuildHLOString(stage, options, xla_device_compiler->client(), + function, args); } absl::StatusOr GetCompilerIr( @@ -405,6 +411,9 @@ absl::StatusOr GetCompilerIr( absl::Span input_arg_shape_and_dtype, absl::Span input_handles, CompilerArgSource compiler_arg_source) { + using XlaDeviceCompiler = + DeviceCompiler; + TF_RETURN_IF_ERROR( ValidateGetCompilerIrTfrtTpu(platform_name, /*stream=*/nullptr, stage)); @@ -436,16 +445,27 @@ absl::StatusOr GetCompilerIr( /*pjrt_device_metadata=*/nullptr, /*device_allocator=*/nullptr); DeviceType compilation_device_type = platform_info.device_type(); - if (platform_info.device_type() == DEVICE_TPU) { - compilation_device_type = DeviceType(DEVICE_TPU_XLA_JIT); - } else { + if (platform_info.device_type() != DEVICE_TPU) { TF_ASSIGN_OR_RETURN(compilation_device_type, GetCompilationDeviceType(platform_info.device_type())); } + XlaDeviceCompiler* xla_device_compiler; + const std::string xla_device_compiler_name = absl::StrCat( + absl::AsciiStrToLower(platform_name), "_xla_device_compiler"); + TF_RETURN_IF_ERROR( + context->HostCPU()->resource_manager()->LookupOrCreate( + context->HostCPU()->resource_manager()->default_container(), + xla_device_compiler_name, &xla_device_compiler, + [&](XlaDeviceCompiler** xla_device_compiler) { + return BuildXlaDeviceCompiler(/*dev=*/nullptr, flr, platform_info, + compilation_device_type, + xla_device_compiler); + })); + core::ScopedUnref xla_device_compiler_ref(xla_device_compiler); XlaCompiler::Options options; if (platform_info.device_type() == DEVICE_TPU) { - options = GenerateCompilerOptionsForTfrtTpu(compilation_device_type, *flr); + options = GenerateCompilerOptionsForTfrtTpu(*xla_device_compiler, *flr); } else { options.device_type = compilation_device_type; options.flib_def = flr->GetFunctionLibraryDefinition(); @@ -455,7 +475,7 @@ absl::StatusOr GetCompilerIr( options.alias_passthrough_params = !platform_info.is_on_xla_device(); } - return CompileAndBuildHLOString(stage, options, /*pjrt_client=*/nullptr, + return CompileAndBuildHLOString(stage, options, /*local_client=*/nullptr, function, args); } diff --git a/tensorflow/compiler/jit/kernels/BUILD b/tensorflow/compiler/jit/kernels/BUILD index 4bf1b48d4657c6..16df5fefeb8c51 100644 --- a/tensorflow/compiler/jit/kernels/BUILD +++ b/tensorflow/compiler/jit/kernels/BUILD @@ -86,6 +86,7 @@ cc_library( "@xla//xla/hlo/ir:hlo", "@xla//xla/pjrt:pjrt_client", "@xla//xla/service:executable", + "@xla//xla/service/gpu:gpu_executable_run_options", "@xla//xla/stream_executor/host:host_platform_id", "@xla//xla/tsl/concurrency:async_value", "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc", diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc index 4b8e481d718b2b..322daf4594c42b 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.cc +++ b/tensorflow/compiler/jit/kernels/xla_ops.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -37,6 +38,7 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/synchronization/mutex.h" #include "absl/types/span.h" +#include "tensorflow/compiler/jit/device_compilation_profiler.h" #include "tensorflow/compiler/jit/device_compiler.h" #include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h" #include "tensorflow/compiler/jit/flags.h" @@ -46,32 +48,49 @@ limitations under the License. #include "tensorflow/compiler/jit/xla_activity.pb.h" #include "tensorflow/compiler/jit/xla_activity_listener.h" #include "tensorflow/compiler/jit/xla_compile_util.h" +#include "tensorflow/compiler/jit/xla_compiler_options_util.h" +#include "tensorflow/compiler/jit/xla_host_recv_device_context.h" +#include "tensorflow/compiler/jit/xla_host_send_device_context.h" #include "tensorflow/compiler/jit/xla_launch_util.h" #include "tensorflow/compiler/jit/xla_platform_info.h" +#include "tensorflow/compiler/tf2xla/tf2xla_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "xla/client/local_client.h" +#include "xla/executable_run_options.h" +#include "xla/hlo/ir/hlo_input_output_alias_config.h" #include "xla/pjrt/pjrt_client.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/macros.h" -#include "xla/tsl/platform/statusor.h" +#include "xla/service/executable.h" +#include "xla/service/gpu/gpu_executable_run_options.h" +#include "xla/stream_executor/host/host_platform_id.h" +#include "xla/tsl/concurrency/async_value_ref.h" #include "xla/tsl/protobuf/error_codes.pb.h" #include "tensorflow/core/framework/allocator.h" +#include "tensorflow/core/framework/control_flow.h" #include "tensorflow/core/framework/device.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/op_requires.h" +#include "tensorflow/core/framework/rendezvous.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/lib/monitoring/counter.h" #include "tensorflow/core/platform/env.h" +#include "tensorflow/core/platform/errors.h" #include "tensorflow/core/platform/mutex.h" -#include "tensorflow/core/platform/status.h" +#include "tensorflow/core/platform/refcount.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" #include "tensorflow/core/platform/threadpool.h" #include "tensorflow/core/platform/tstring.h" +#include "tensorflow/core/platform/types.h" +#include "tensorflow/core/util/device_name_utils.h" +#include "tensorflow/core/util/stream_executor_util.h" #include "tsl/platform/thread_annotations.h" +#include "tsl/profiler/lib/traceme.h" // OP_REQUIRES_OK_RETURN is the same as OP_REQUIRES_OK except that // in error case, it returns RET instead of void. @@ -87,6 +106,8 @@ limitations under the License. namespace tensorflow { namespace { +using XlaDeviceCompiler = + DeviceCompiler; using PjRtDeviceCompiler = DeviceCompiler; @@ -180,11 +201,172 @@ class ExecutableClosureStore { void operator=(const ExecutableClosureStore&) = delete; }; +using XlaExecutableClosure = + ExecutableClosure; +using XlaExecutableClosureStore = + ExecutableClosureStore; using PjRtExecutableClosure = ExecutableClosure; using PjRtExecutableClosureStore = ExecutableClosureStore; +se::Stream* GetStream(OpKernelContext* ctx) { + return ctx->op_device_context() ? ctx->op_device_context()->stream() + : nullptr; +} + +XlaComputationLaunchContext GetLaunchContext( + const XlaPlatformInfo& platform_info, OpKernelContext* ctx, + xla::LocalClient* client, + stream_executor::DeviceAddressAllocator* allocator) { + se::Stream* stream = GetStream(ctx); + int device_ordinal = stream ? stream->parent()->device_ordinal() + : client->default_device_ordinal(); + XlaComputationLaunchContext launch_context( + client, allocator, device_ordinal, + /*allocate_xla_tensors=*/platform_info.is_on_xla_device(), + /*use_multiple_streams=*/platform_info.UseMultipleStreams()); + return launch_context; +} + +absl::Status GetTaskName(const absl::string_view device_name, + std::string* task_name) { + std::string ignored; + if (!DeviceNameUtils::SplitDeviceName(device_name, task_name, &ignored)) { + return absl::InvalidArgumentError( + absl::StrCat("Unable to parse device name: ", device_name)); + } + + return absl::OkStatus(); +} + +// Provide SendDeviceMemoryFunction for XLA host callbacks. This callback +// handles transferring from device to host. +xla::SendDeviceMemoryFunction GetSendDeviceMemoryFunction( + OpKernelContext* ctx, const std::string& program_key) { + return + [ctx, program_key]( + int64_t channel_id, se::Stream* stream, const xla::Shape& shape, + const stream_executor::DeviceAddressBase& device_memory_base, + const absl::flat_hash_map& frontend_attrs) + -> absl::StatusOr>> { + auto iter = frontend_attrs.find("_xla_host_transfer_rendezvous"); + + // Generate the Rendezvous key. + const std::string& rendezvous_key_base = + absl::StrCat(program_key, iter->second); + + const std::string& src_device = ctx->device()->name(); + std::string task_prefix; + TF_RETURN_IF_ERROR(GetTaskName(src_device, &task_prefix)); + const std::string dst_device = + absl::StrCat(task_prefix, "/device:CPU:0"); + const std::string& rendezvous_key = + Rendezvous::CreateKey(src_device, /*src_incarnation=*/1, dst_device, + rendezvous_key_base, FrameAndIter(0, 0)); + VLOG(2) << "Rendezvous Key for receiving at host: " << rendezvous_key; + + RendezvousInterface::ParsedKey parsed_key; + TF_RETURN_IF_ERROR(Rendezvous::ParseKey(rendezvous_key, &parsed_key)); + + TF_ASSIGN_OR_RETURN(auto event, stream->parent()->CreateEvent()); + tsl::AsyncValueRef> done_event = + tsl::MakeConstructedAsyncValueRef>( + std::move(event)); + + Rendezvous::Args args; + // Rendezvous::Args owns the device context pointer. + args.device_context = new XlaHostRecvDeviceContext( + stream, device_memory_base, shape, done_event); + + Tensor host_tensor; + TF_RETURN_IF_ERROR( + ctx->rendezvous()->Send(parsed_key, args, host_tensor, false)); + + return std::move(done_event); + }; +} + +// Provide RecvDeviceMemoryFunction for XLA host callbacks. This callback +// handles transferring from host to device. +xla::RecvDeviceMemoryFunction GetRecvDeviceMemoryFunction( + OpKernelContext* ctx, const std::string& program_key) { + return + [ctx, program_key]( + int64_t channel_id, se::Stream* stream, const xla::Shape& shape, + stream_executor::DeviceAddressBase* device_memory_base, + const absl::flat_hash_map& frontend_attrs) + -> absl::StatusOr>> { + auto iter = frontend_attrs.find("_xla_host_transfer_rendezvous"); + + // Generate the Rendezvous key. + const std::string& rendezvous_key_base = + absl::StrCat(program_key, iter->second); + + const std::string& dst_device = ctx->device()->name(); + std::string task_prefix; + TF_RETURN_IF_ERROR(GetTaskName(dst_device, &task_prefix)); + const std::string src_device = + absl::StrCat(task_prefix, "/device:CPU:0"); + const std::string& rendezvous_key = + Rendezvous::CreateKey(src_device, /*src_incarnation=*/1, dst_device, + rendezvous_key_base, FrameAndIter(0, 0)); + VLOG(2) << "Rendezvous Key for sending from host: " << rendezvous_key; + + RendezvousInterface::ParsedKey parsed_key; + TF_RETURN_IF_ERROR(Rendezvous::ParseKey(rendezvous_key, &parsed_key)); + + TF_ASSIGN_OR_RETURN(auto event, stream->parent()->CreateEvent()); + tsl::AsyncValueRef> done_event = + tsl::MakeConstructedAsyncValueRef>( + std::move(event)); + + Rendezvous::Args args; + // Rendezvous::Args owns the device context pointer. + args.device_context = new XlaHostSendDeviceContext( + stream, device_memory_base, shape, done_event); + + Tensor device_tensor; + bool is_dead; + TF_RETURN_IF_ERROR(ctx->rendezvous()->Recv( + parsed_key, args, &device_tensor, /*is_dead=*/&is_dead)); + + return std::move(done_event); + }; +} + +absl::StatusOr RunExecutable( + const XlaPlatformInfo& platform_info, + const XlaComputationLaunchContext& launch_context, + std::vector execution_inputs, + xla::ExecutableRunOptions run_options, xla::LocalExecutable* executable, + OpKernelContext* ctx, stream_executor::DeviceAddressAllocator* allocator) { + VLOG(2) << "Executing Xla Computation."; + Env* env = Env::Default(); + auto start_time = env->NowMicros(); + + se::Stream* stream = GetStream(ctx); + run_options.set_stream(GetStream(ctx)); + run_options.set_allocator(allocator); + run_options.set_intra_op_thread_pool(&ctx->eigen_cpu_device()); + run_options.set_rng_seed(GetXLARandomSeed()); + + absl::StatusOr execution_output; + bool run_synchronous = + !stream || platform_info.platform_id() == se::host::kHostPlatformId; + if (run_synchronous) { + execution_output = + executable->Run(std::move(execution_inputs), run_options); + } else { + execution_output = + executable->RunAsync(std::move(execution_inputs), run_options); + } + + auto elapsed = env->NowMicros() - start_time; + VLOG(2) << "Elapsed time for Xla Executable Run: " << elapsed << "us"; + return execution_output; +} + absl::StatusOr< std::pair, ResourceVarsSnapshot>> GetXlaCompilerArgsAndSnapshotVariables( @@ -209,6 +391,59 @@ GetXlaCompilerArgsAndSnapshotVariables( return result; } +absl::Status CompileToLocalExecutable( + OpKernelContext* ctx, const NameAttrList& function, bool has_ref_vars, + const XlaPlatformInfo& platform_info, + const std::vector& args, + DeviceCompileMode compile_mode, bool may_alias_resource_update, + xla::LocalClient** client, + const XlaCompiler::CompilationResult** compilation_result, + xla::LocalExecutable** executable) { + // We store information about the JIT-compiled XLA computation + // in the ResourceMgr. + ResourceMgr* rm = ctx->resource_manager(); + if (!rm) { + return absl::InternalError("No resource manager."); + } + + TF_ASSIGN_OR_RETURN(DeviceType compilation_device_type, + GetCompilationDeviceType(platform_info.device_type())); + + XlaDeviceCompiler* xla_device_compiler; + TF_RETURN_IF_ERROR(rm->LookupOrCreate( + rm->default_container(), "xla_device_compiler", &xla_device_compiler, + [&](XlaDeviceCompiler** xla_device_compiler) { + return BuildXlaDeviceCompiler(ctx->device(), ctx->function_library(), + platform_info, compilation_device_type, + xla_device_compiler); + })); + DeviceCompilationProfiler* profiler; + TF_RETURN_IF_ERROR(rm->LookupOrCreate( + rm->default_container(), "device_compilation_profiler", &profiler, + [](DeviceCompilationProfiler** profiler) { + *profiler = new DeviceCompilationProfiler(); + return absl::OkStatus(); + })); + // Hold the reference to the XLA device compiler and profiler during + // evaluation. (We could probably free them sooner because the ResourceMgr + // will retain references, but this is more obviously correct.) + core::ScopedUnref xla_device_compiler_ref(xla_device_compiler); + core::ScopedUnref profiler_ref(profiler); + + *client = static_cast(xla_device_compiler->client()); + + XlaCompiler::Options options = GenerateCompilerOptions( + *xla_device_compiler, *ctx->function_library(), ctx->device(), + GetStream(ctx), platform_info, has_ref_vars); + + XlaCompiler::CompileOptions compile_options = + GenerateCompileOptions(has_ref_vars, may_alias_resource_update); + + return xla_device_compiler->CompileIfNeeded( + options, function, args, compile_options, compile_mode, profiler, + compilation_result, executable); +} + absl::Status GetUpdatedVariables( const OpKernelContext* ctx, absl::Span inputs, absl::Span variable_indices, @@ -282,6 +517,9 @@ void XlaLocalLaunchBase::ComputeAsync(OpKernelContext* ctx, DoneCallback done) { std::vector xla_compiler_args; const XlaCompiler::CompilationResult* compilation_result; + xla::LocalClient* client; // Not owned. + xla::LocalExecutable* executable; // Not owned. + xla::PjRtClient* pjrt_client; // Not owned. xla::PjRtLoadedExecutable* pjrt_executable; // Not owned. @@ -310,26 +548,66 @@ void XlaLocalLaunchBase::ComputeAsync(OpKernelContext* ctx, DoneCallback done) { xla_compiler_args = std::move(status_or_xla_compiler_args.value()); } - VLOG(2) << "Compiling using PJRT"; - absl::Status status = CompileToPjRtLoadedExecutable( - *ctx, platform_info_, function_, xla_compiler_args, - DeviceCompileMode::kStrict, has_ref_vars_, - /*may_alias_resource_update=*/true, &compilation_result, &pjrt_client, - &pjrt_executable); - OP_REQUIRES_OK_ASYNC(ctx, status, done); + bool use_pjrt = GetXlaOpsCommonFlags() + ->tf_xla_use_device_api.IsEnabledInXlaLaunchForDevice( + platform_info_.device_type()); + if (use_pjrt) { + VLOG(2) << "Compiling using PJRT"; + absl::Status status = CompileToPjRtLoadedExecutable( + *ctx, platform_info_, function_, xla_compiler_args, + DeviceCompileMode::kStrict, has_ref_vars_, + /*may_alias_resource_update=*/true, &compilation_result, &pjrt_client, + &pjrt_executable); + OP_REQUIRES_OK_ASYNC(ctx, status, done); + + VLOG(2) << "Compiled using PJRT: " << status; + VLOG(2) << "pjrt_executable != nullptr: " << (pjrt_executable != nullptr); + VLOG(2) << "compilation_result != nullptr: " + << (compilation_result != nullptr); + VLOG(2) << "Executing using PJRT."; + + auto run_pjrt_cluster = [ctx, pjrt_client, pjrt_executable, + compilation_result, done, inputs, + resources = resources_]() { + // Separate scope so that VariableInfo locks are released before done() is + // called. + { + std::vector variable_infos; + OP_REQUIRES_OK_ASYNC( + ctx, + GetUpdatedVariables(ctx, inputs, resources, *compilation_result, + &variable_infos), + done); + OP_REQUIRES_OK_ASYNC(ctx, LockVariables(absl::MakeSpan(variable_infos)), + done); + OP_REQUIRES_OK_ASYNC( + ctx, + RunPjRtExecutable(inputs, variable_infos, *compilation_result, + pjrt_client, pjrt_executable, ctx), + done); + } + VLOG(2) << "Done executing with PJRT."; + done(); + }; + + RunInThreadPoolIfCollectivesPresent(*compilation_result, run_pjrt_cluster); + return; + } - VLOG(2) << "Compiled using PJRT: " << status; - VLOG(2) << "pjrt_executable != nullptr: " << (pjrt_executable != nullptr); - VLOG(2) << "compilation_result != nullptr: " - << (compilation_result != nullptr); - VLOG(2) << "Executing using PJRT."; + absl::Status status = CompileToLocalExecutable( + ctx, function_, /*has_ref_vars=*/has_ref_vars_, platform_info_, + xla_compiler_args, DeviceCompileMode::kStrict, + /*may_alias_resource_update=*/true, &client, &compilation_result, + &executable); + OP_REQUIRES_OK_ASYNC(ctx, status, done); - auto run_pjrt_cluster = [ctx, pjrt_client, pjrt_executable, - compilation_result, done, inputs, - resources = resources_]() { - // Separate scope so that VariableInfo locks are released before done() is + // Continuation of the execution, may be run in a different thread. + auto run_xla_cluster = [ctx, client, executable, compilation_result, done, + inputs, resources = resources_]() { + // Separate scope so that VariableInfo locks are released before done is // called. { + auto platform_info = XlaPlatformInfoFromDevice(ctx->device()); std::vector variable_infos; OP_REQUIRES_OK_ASYNC( ctx, @@ -338,17 +616,61 @@ void XlaLocalLaunchBase::ComputeAsync(OpKernelContext* ctx, DoneCallback done) { done); OP_REQUIRES_OK_ASYNC(ctx, LockVariables(absl::MakeSpan(variable_infos)), done); + absl::flat_hash_map resource_var_ptrs; + for (int i = 0; i < resources.size(); i++) { + resource_var_ptrs[resources[i]] = variable_infos[i].var()->tensor(); + } + + std::shared_ptr allocator = + GetAllocator(ctx->device(), GetStream(ctx), platform_info); + XlaComputationLaunchContext launch_context = + GetLaunchContext(platform_info, ctx, client, allocator.get()); + + const xla::HloInputOutputAliasConfig& input_output_alias = + executable->executable()->module().input_output_alias_config(); + absl::StatusOr> execution_inputs = + launch_context.PopulateInputs( + ctx, compilation_result, resource_var_ptrs, + /*missing_ctx_input_prefix=*/0, input_output_alias); + OP_REQUIRES_OK_ASYNC(ctx, execution_inputs.status(), done); + + xla::gpu::GpuExecutableRunOptions gpu_options; + xla::DeviceAssignment device_assignment; + xla::ExecutableRunOptions run_options; + if (compilation_result->collective_info.has_value()) { + OP_REQUIRES_OK_ASYNC(ctx, + ResolveDeviceAssignment( + ctx, *compilation_result->collective_info, + run_options, device_assignment, gpu_options), + done); + } + + // Hardcode run id to always be zero: TF distributed strategy + // differentiates between subsequent runs using dependency edges. This + // is safe, as only TF dist-strat can produce distributed ops, and we + // can rely on TF dist-strat invariants. + xla::RunId run_id(0); + run_options.set_run_id(run_id); + + absl::StatusOr execution_output = RunExecutable( + platform_info, launch_context, std::move(*execution_inputs), + run_options, executable, ctx, allocator.get()); + OP_REQUIRES_ASYNC(ctx, execution_output.ok(), execution_output.status(), + done); + OP_REQUIRES_OK_ASYNC( ctx, - RunPjRtExecutable(inputs, variable_infos, *compilation_result, - pjrt_client, pjrt_executable, ctx), + launch_context.PopulateOutputs( + ctx, compilation_result, execution_output->ConsumeResult(), + /*missing_ctx_input_prefix=*/0, absl::MakeSpan(variable_infos), + input_output_alias, resource_var_ptrs), done); + VLOG(1) << "Done"; } - VLOG(2) << "Done executing with PJRT."; done(); }; - RunInThreadPoolIfCollectivesPresent(*compilation_result, run_pjrt_cluster); + RunInThreadPoolIfCollectivesPresent(*compilation_result, run_xla_cluster); } namespace { @@ -440,6 +762,8 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { VLOG(3) << "XlaCompileOp " << def().name() << (must_compile_ ? "(must-compile)" : ""); const XlaCompiler::CompilationResult* kernel = nullptr; + xla::LocalClient* client = nullptr; + xla::LocalExecutable* executable = nullptr; xla::PjRtClient* pjrt_client = nullptr; xla::PjRtLoadedExecutable* pjrt_executable = nullptr; ResourceVarsSnapshot variables_snapshot; @@ -459,9 +783,14 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { : DeviceCompileMode::kLazy; }(); + bool use_pjrt = + GetXlaOpsCommonFlags() + ->tf_xla_use_device_api.IsEnabledInXlaCompileAndRunForDevice( + platform_info_.device_type()); + if (GetXlaOpsCommonFlags()->tf_xla_always_defer_compilation || cannot_compile_cluster) { - pjrt_executable = nullptr; + executable = nullptr; } else { auto args_and_variables_snapshot = GetXlaCompilerArgsAndSnapshotVariables( resources_, constants_, inputs, ctx); @@ -472,12 +801,19 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { // Do not alias resource updates as locking variables in XlaCompile and // unlocking them in XlaRun may lead to deadlocks. - VLOG(2) << "Using PJRT for compilation. Function name: " - << function_.name(); - absl::Status status = CompileToPjRtLoadedExecutable( - *ctx, platform_info_, function_, args, compile_mode, has_ref_vars_, - /*may_alias_resource_update=*/false, &kernel, &pjrt_client, - &pjrt_executable); + absl::Status status; + if (use_pjrt) { + VLOG(2) << "Using PJRT for compilation. Function name: " + << function_.name(); + status = CompileToPjRtLoadedExecutable( + *ctx, platform_info_, function_, args, compile_mode, has_ref_vars_, + /*may_alias_resource_update=*/false, &kernel, &pjrt_client, + &pjrt_executable); + } else { + status = CompileToLocalExecutable( + ctx, function_, has_ref_vars_, platform_info_, args, compile_mode, + /*may_alias_resource_update=*/false, &client, &kernel, &executable); + } if (compile_mode != DeviceCompileMode::kLazy || status.code() != error::UNIMPLEMENTED) { OP_REQUIRES_OK(ctx, status); @@ -490,6 +826,7 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { BroadcastOptimizationRemark( XlaOptimizationRemark::UNIMPLEMENTED_OPERATION, status.ToString()) .IgnoreError(); + executable = nullptr; pjrt_executable = nullptr; mutex_lock guard(cannot_compile_cluster_mu_); cannot_compile_cluster_ = true; @@ -502,7 +839,7 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { Allocator* cpu_allocator = ctx->device()->GetAllocator(host_alloc_attrs); // Async compilation returns nullptr executable without an error. - if (!pjrt_executable) { + if (!executable && !pjrt_executable) { DCHECK(!must_compile_); Tensor compilation_key(cpu_allocator, DT_STRING, TensorShape({})); Tensor compilation_successful(cpu_allocator, DT_BOOL, TensorShape({})); @@ -517,12 +854,21 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { // hit. This is because we at least need new snapshots of the resource // variables. Tensor compilation_key(cpu_allocator, DT_STRING, TensorShape({})); - PjRtExecutableClosureStore::KeyT key = - PjRtExecutableClosureStore::Global()->Produce(PjRtExecutableClosure( - pjrt_client, pjrt_executable, kernel, std::move(variables_snapshot), - constants_.size())); - compilation_key.flat()(0) = key; - VLOG(2) << "Compiled with PJRT. compilation_key: " << key; + if (use_pjrt) { + PjRtExecutableClosureStore::KeyT key = + PjRtExecutableClosureStore::Global()->Produce(PjRtExecutableClosure( + pjrt_client, pjrt_executable, kernel, std::move(variables_snapshot), + constants_.size())); + compilation_key.flat()(0) = key; + VLOG(2) << "Compiled with PJRT. compilation_key: " << key; + } else { + XlaExecutableClosureStore::KeyT key = + XlaExecutableClosureStore::Global()->Produce(XlaExecutableClosure( + client, executable, kernel, std::move(variables_snapshot), + constants_.size())); + compilation_key.flat()(0) = key; + VLOG(2) << "Compiled with XLA. compilation_key: " << key; + } Tensor compilation_successful(cpu_allocator, DT_BOOL, TensorShape({})); compilation_successful.flat()(0) = true; @@ -531,40 +877,123 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { ctx->set_output(1, compilation_successful); } -XlaRunOp::XlaRunOp(OpKernelConstruction* ctx) : OpKernel(ctx) {} +XlaRunOp::XlaRunOp(OpKernelConstruction* ctx) + : OpKernel(ctx), platform_info_(XlaPlatformInfoFromDevice(ctx->device())) {} void XlaRunOp::Compute(OpKernelContext* ctx) { VLOG(3) << "XlaRunOp " << def().name(); Tensor key_tensor = ctx->input(ctx->num_inputs() - 1); - const PjRtExecutableClosureStore::KeyT& key = key_tensor.flat()(0); - PjRtExecutableClosure closure = - PjRtExecutableClosureStore::Global()->Consume(key); + bool use_pjrt = + GetXlaOpsCommonFlags() + ->tf_xla_use_device_api.IsEnabledInXlaCompileAndRunForDevice( + platform_info_.device_type()); + + if (use_pjrt) { + const PjRtExecutableClosureStore::KeyT& key = key_tensor.flat()(0); + PjRtExecutableClosure closure = + PjRtExecutableClosureStore::Global()->Consume(key); + + // Fetch inputs from the OpKernelContext. Inputs are the same as the ones + // for XlaCompile, except that the must-be-constant inputs that appear in + // the beginning are stripped off and the closure key is appended as the + // last input. So the inputs look like: input tensors, resource variables, + // closure key tensor. + std::vector inputs = InputsFromContext(ctx); + absl::flat_hash_map variable_snapshots; + for (const auto& [variable_index, variable_tensor] : + closure.resource_var_snapshots()) { + variable_snapshots.emplace(variable_index, variable_tensor.has_value() + ? &variable_tensor.value() + : nullptr); + } - // Fetch inputs from the OpKernelContext. Inputs are the same as the ones - // for XlaCompile, except that the must-be-constant inputs that appear in - // the beginning are stripped off and the closure key is appended as the - // last input. So the inputs look like: input tensors, resource variables, - // closure key tensor. - std::vector inputs = InputsFromContext(ctx); - absl::flat_hash_map variable_snapshots; - for (const auto& [variable_index, variable_tensor] : - closure.resource_var_snapshots()) { - variable_snapshots.emplace(variable_index, variable_tensor.has_value() - ? &variable_tensor.value() - : nullptr); + { + absl::StatusOr> updated_variables = + GatherVariableInfo(ctx, *closure.compilation_result(), + closure.num_constant_args()); + OP_REQUIRES_OK(ctx, updated_variables.status()); + OP_REQUIRES_OK(ctx, LockVariables(absl::MakeSpan(*updated_variables))); + OP_REQUIRES_OK( + ctx, RunPjRtExecutable(closure.num_constant_args(), inputs, + variable_snapshots, *updated_variables, + *closure.compilation_result(), + closure.client(), closure.executable(), ctx)); + } + + OP_REQUIRES_OK(ctx, absl::OkStatus()); + return; + } + + const XlaExecutableClosureStore::KeyT& key = key_tensor.flat()(0); + + XlaExecutableClosure closure = + XlaExecutableClosureStore::Global()->Consume(key); + std::shared_ptr allocator = + GetAllocator(ctx->device(), GetStream(ctx), platform_info_); + XlaComputationLaunchContext launch_context = + GetLaunchContext(platform_info_, ctx, closure.client(), allocator.get()); + + // We're missing the must-be-constant inputs, tell `PopulateInputs` + // about this. We don't actually need these inputs because they've + // already been baked into the compiled kernel. + const xla::HloInputOutputAliasConfig& input_output_alias = + closure.executable()->executable()->module().input_output_alias_config(); + absl::StatusOr> execution_inputs; + absl::flat_hash_map snapshot_ptrs; + { + tsl::profiler::TraceMe hlo_module_activity( + [&] { + return absl::StrCat( + "Populate Inputs (", + closure.compilation_result()->xla_input_shapes.size(), ")"); + }, + tsl::profiler::TraceMeLevel::kInfo); + + for (const auto& [variable_index, variable_tensor] : + closure.resource_var_snapshots()) { + snapshot_ptrs.emplace(variable_index, variable_tensor.has_value() + ? &variable_tensor.value() + : nullptr); + } + execution_inputs = launch_context.PopulateInputs( + ctx, closure.compilation_result(), snapshot_ptrs, + /*missing_ctx_input_prefix=*/closure.num_constant_args(), + input_output_alias); + OP_REQUIRES_OK(ctx, execution_inputs.status()); } - absl::StatusOr> updated_variables = - GatherVariableInfo(ctx, *closure.compilation_result(), - closure.num_constant_args()); - OP_REQUIRES_OK(ctx, updated_variables.status()); - OP_REQUIRES_OK(ctx, LockVariables(absl::MakeSpan(*updated_variables))); + xla::ExecutableRunOptions run_options; + + // Host callbacks used for HLO send/recv. + xla::SendDeviceMemoryFunction send_function = + GetSendDeviceMemoryFunction(ctx, key); + run_options.set_send_device_memory_function(&send_function); + xla::RecvDeviceMemoryFunction recv_function = + GetRecvDeviceMemoryFunction(ctx, key); + run_options.set_recv_device_memory_function(&recv_function); + + absl::StatusOr execution_output = RunExecutable( + platform_info_, launch_context, std::move(*execution_inputs), run_options, + closure.executable(), ctx, allocator.get()); + OP_REQUIRES(ctx, execution_output.ok(), execution_output.status()); + + tsl::profiler::TraceMe hlo_module_activity( + [&] { + return absl::StrCat("Populate Outputs (", ctx->num_outputs(), ")"); + }, + tsl::profiler::TraceMeLevel::kInfo); + + absl::StatusOr> variable_infos = GatherVariableInfo( + ctx, *closure.compilation_result(), closure.num_constant_args()); + OP_REQUIRES_OK(ctx, variable_infos.status()); + OP_REQUIRES_OK(ctx, LockVariables(absl::MakeSpan(*variable_infos))); OP_REQUIRES_OK( ctx, - RunPjRtExecutable(closure.num_constant_args(), inputs, variable_snapshots, - *updated_variables, *closure.compilation_result(), - closure.client(), closure.executable(), ctx)); + launch_context.PopulateOutputs( + ctx, closure.compilation_result(), execution_output->ConsumeResult(), + /*missing_ctx_input_prefix=*/closure.num_constant_args(), + absl::MakeSpan(*variable_infos), input_output_alias, snapshot_ptrs)); } XlaMergeOp::XlaMergeOp(OpKernelConstruction* ctx) : OpKernel(ctx) {} diff --git a/tensorflow/compiler/jit/kernels/xla_ops.h b/tensorflow/compiler/jit/kernels/xla_ops.h index 96ca28cdac5f55..1e42fe5ff38741 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.h +++ b/tensorflow/compiler/jit/kernels/xla_ops.h @@ -16,17 +16,25 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_H_ #define TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_H_ +#include #include +#include "tensorflow/compiler/jit/device_compiler.h" +#include "tensorflow/compiler/jit/xla_device.h" +#include "tensorflow/compiler/jit/xla_launch_util.h" #include "tensorflow/compiler/jit/xla_platform_info.h" +#include "xla/stream_executor/integrations/tf_allocator_adapter.h" +#include "tensorflow/core/framework/allocator.h" #include "tensorflow/core/framework/attr_value.pb.h" +#include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/platform/mutex.h" -#include "tsl/platform/thread_annotations.h" +#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/util/stream_executor_util.h" namespace tensorflow { + // XlaLocalLaunchBase is almost the same as XlaLocalLaunchOp. // The only difference is that it does not require arguments to follow // the "constants, then regular args, then resources" order. @@ -68,8 +76,8 @@ class XlaLocalLaunchBase : public AsyncOpKernel { // Once all inputs are present, and their shapes are known, the op can // use a 'DeviceCompiler' to compile and execute code which is specific // to the shapes of input Tensors. -// XlaLocalLaunchOp uses xla::PjRtClient and -// xla::PjRtLoadedExecutable, and passes arguments into/out of XLA in device +// XlaLocalLaunchOp uses xla::LocalClient::Compile() and +// xla::LocalExecutable::Run(), and passes arguments into/out of XLA in device // memory. class XlaLocalLaunchOp : public XlaLocalLaunchBase { public: @@ -117,6 +125,9 @@ class XlaRunOp : public OpKernel { explicit XlaRunOp(OpKernelConstruction* ctx); void Compute(OpKernelContext* ctx) override; + + private: + const XlaPlatformInfo platform_info_; }; class XlaMergeOp : public OpKernel { diff --git a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc index c69b9ce20350f8..f1a648f15592cf 100644 --- a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc +++ b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc @@ -17,40 +17,68 @@ limitations under the License. #include "tensorflow/compiler/jit/xla_compile_on_demand_op.h" +#include +#include +#include +#include #include +#include "absl/container/flat_hash_map.h" #include "absl/log/check.h" #include "absl/log/log.h" -#include "absl/status/status.h" +#include "absl/memory/memory.h" #include "absl/types/span.h" #include "tensorflow/compiler/jit/device_compilation_cluster_signature.h" #include "tensorflow/compiler/jit/device_compilation_profiler.h" #include "tensorflow/compiler/jit/device_compiler.h" +#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/jit/variable_info.h" #include "tensorflow/compiler/jit/variable_info_util.h" +#include "tensorflow/compiler/jit/xla_compile_util.h" #include "tensorflow/compiler/jit/xla_compiler_options_util.h" #include "tensorflow/compiler/jit/xla_launch_util.h" #include "tensorflow/compiler/jit/xla_platform_info.h" +#include "tensorflow/compiler/tf2xla/const_analysis.h" +#include "tensorflow/compiler/tf2xla/tf2xla_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" +#include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "xla/client/local_client.h" +#include "xla/executable_run_options.h" +#include "xla/hlo/ir/hlo_input_output_alias_config.h" #include "xla/pjrt/pjrt_client.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" +#include "xla/pjrt/tf_pjrt_client.h" +#include "xla/service/executable.h" +#include "xla/service/gpu/gpu_executable_run_options.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/op_requires.h" #include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/refcount.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/statusor.h" +#include "tensorflow/core/tfrt/common/pjrt_util.h" +#include "tsl/platform/errors.h" namespace tensorflow { namespace { +using XlaDeviceCompiler = + DeviceCompiler; using PjRtDeviceCompiler = DeviceCompiler; -XlaCompiler::CompileOptions GetCompileOptions() { +XlaCompiler::CompileOptions GetCompileOptions(bool for_pjrt = false) { XlaCompiler::CompileOptions compile_options; compile_options.is_entry_computation = true; - compile_options.always_return_tuple = true; - compile_options.use_tuple_arg = false; + // Optimization: where possible, have the computation return a naked array + // rather than a one-element tuple. + compile_options.always_return_tuple = false; + if (for_pjrt) { + compile_options.use_tuple_arg = false; + compile_options.always_return_tuple = true; + } + return compile_options; } @@ -83,6 +111,70 @@ XlaCompileOnDemandOp::XlaCompileOnDemandOp(OpKernelConstruction* ctx) function_(GetDeviceCompilerFunction(ctx->def())), canonical_function_(Canonicalize(function_)) {} +absl::Status XlaCompileOnDemandOp::Run( + const ResourceVarsSnapshot& variable_args, + const XlaCompiler::CompilationResult* result, + const XlaDeviceCompiler* xla_device_compiler, + xla::LocalExecutable* executable, OpKernelContext* ctx) { + xla::LocalClient* client = + static_cast(xla_device_compiler->client()); + + se::Stream* stream = + ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr; + std::shared_ptr allocator_ptr = + GetAllocator(ctx->device(), stream, platform_info_); + stream_executor::DeviceAddressAllocator* allocator = allocator_ptr.get(); + XlaComputationLaunchContext launch_context( + client, allocator, client->default_device_ordinal(), + /*allocate_xla_tensors=*/platform_info_.xla_device_metadata() != nullptr, + platform_info_.xla_device_metadata() + ? platform_info_.xla_device_metadata()->UseMultipleStreams() + : false); + + absl::flat_hash_map snapshot_ptrs; + for (auto& p : variable_args) { + snapshot_ptrs.emplace(p.first, + p.second.has_value() ? &p.second.value() : nullptr); + } + + const xla::HloInputOutputAliasConfig& input_output_alias = + executable->executable()->module().input_output_alias_config(); + absl::StatusOr> execution_inputs = + launch_context.PopulateInputs(ctx, result, snapshot_ptrs, + /*missing_ctx_input_prefix=*/0, + input_output_alias); + TF_RETURN_IF_ERROR(execution_inputs.status()); + + VLOG(2) << "Executing computation: " << name(); + xla::ExecutableRunOptions run_options; + xla::gpu::GpuExecutableRunOptions gpu_options; + xla::DeviceAssignment device_assignment; + if (result->collective_info) { + TF_RETURN_IF_ERROR(ResolveDeviceAssignment(ctx, *result->collective_info, + run_options, device_assignment, + gpu_options)); + } + + run_options.set_stream(stream); + run_options.set_allocator(allocator); + run_options.set_intra_op_thread_pool(&ctx->eigen_cpu_device()); + run_options.set_rng_seed(GetXLARandomSeed()); + + absl::StatusOr run_result = + executable->Run(std::move(execution_inputs).value(), run_options); + TF_RETURN_IF_ERROR(run_result.status()); + xla::ExecutionOutput execution_output = std::move(run_result).value(); + absl::StatusOr> variable_infos = + GatherVariableInfo(ctx, *result, 0); + TF_RETURN_IF_ERROR(variable_infos.status()); + TF_RETURN_IF_ERROR(LockVariables(absl::MakeSpan(*variable_infos))); + TF_RETURN_IF_ERROR(launch_context.PopulateOutputs( + ctx, result, execution_output.ConsumeResult(), + /*missing_ctx_input_prefix=*/0, absl::MakeSpan(*variable_infos), + input_output_alias, snapshot_ptrs)); + return absl::OkStatus(); +} + absl::Status XlaCompileOnDemandOp::Compile( const std::vector& args, OpKernelContext* ctx, PjRtDeviceCompiler** pjrt_device_compiler, @@ -98,7 +190,7 @@ absl::Status XlaCompileOnDemandOp::Compile( platform_info_, *pjrt_device_compiler); // No detailed logging for on demand op. options.detailed_logging = false; - XlaCompiler::CompileOptions compile_options = GetCompileOptions(); + XlaCompiler::CompileOptions compile_options = GetCompileOptions(true); return (*pjrt_device_compiler) ->CompileSingleOpIfNeeded(options, function_, canonical_function_, args, @@ -106,9 +198,52 @@ absl::Status XlaCompileOnDemandOp::Compile( executable); } +absl::Status XlaCompileOnDemandOp::Compile( + const std::vector& args, OpKernelContext* ctx, + XlaDeviceCompiler** xla_device_compiler, + DeviceCompilationProfiler** profiler, + const XlaCompiler::CompilationResult** result, + xla::LocalExecutable** executable) { + // We store information about the JIT-compiled XLA computation + // in the ResourceMgr. + ResourceMgr* rm = ctx->resource_manager(); + CHECK(rm); + + TF_ASSIGN_OR_RETURN(DeviceType compilation_device_type, + GetCompilationDeviceType(platform_info_.device_type())); + + TF_RETURN_IF_ERROR(rm->LookupOrCreate( + rm->default_container(), "xla_device_compiler", xla_device_compiler, + [&](XlaDeviceCompiler** xla_device_compiler) { + return BuildXlaDeviceCompiler(ctx->device(), ctx->function_library(), + platform_info_, compilation_device_type, + xla_device_compiler); + })); + + TF_RETURN_IF_ERROR(rm->LookupOrCreate( + rm->default_container(), "device_compilation_profiler", profiler, + [](DeviceCompilationProfiler** profiler) { + *profiler = new DeviceCompilationProfiler(); + return absl::OkStatus(); + })); + + XlaCompiler::Options options = GenerateCompilerOptions( + **xla_device_compiler, *ctx->function_library(), ctx->device(), + ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr, + platform_info_, /*has_ref_vars=*/true); + // No detailed logging from on demand op. + options.detailed_logging = false; + XlaCompiler::CompileOptions compile_options = GetCompileOptions(); + + return (*xla_device_compiler) + ->CompileSingleOpIfNeeded(options, function_, canonical_function_, args, + compile_options, ctx, *profiler, result, + executable); +} + void XlaCompileOnDemandOp::Compute(OpKernelContext* ctx) { - const XlaCompiler::CompilationResult* result = nullptr; - DeviceCompilationProfiler* profiler = nullptr; + const XlaCompiler::CompilationResult* result; + DeviceCompilationProfiler* profiler; OP_REQUIRES(ctx, ctx->function_library(), absl::InternalError("Function library missing")); @@ -120,35 +255,70 @@ void XlaCompileOnDemandOp::Compute(OpKernelContext* ctx) { std::vector variable_indices = GetResourceVariableIndicesFromContext(ctx); - std::vector variables; - std::vector args; - // Lock variables for the whole duration of compile + execute. - OP_REQUIRES_OK(ctx, GetAndLockVariablesAndBuildXlaCompilerArguments( - *ctx, inputs, *constant_indices_or, variable_indices, - &variables, &args)); - - PjRtDeviceCompiler* pjrt_device_compiler = nullptr; - xla::PjRtLoadedExecutable* pjrt_executable = nullptr; - absl::Status status = Compile(args, ctx, &pjrt_device_compiler, &profiler, - &result, &pjrt_executable); - // Hold the reference to the XLA device compiler and profiler during - // evaluation. (We could probably free them sooner because the ResourceMgr - // will retain references, but this is more obviously correct.) - // We must also ensure these pointers are freed even if compilation fails. - core::ScopedUnref pjrt_device_compiler_ref(pjrt_device_compiler); - core::ScopedUnref profiler_ref(profiler); - OP_REQUIRES_OK(ctx, status); - - VLOG(2) << "Compiled op with PJRT: " << ctx->status(); - VLOG(2) << "result != nullptr: " << (result != nullptr); - VLOG(2) << "pjrt_executable != nullptr: " << (pjrt_executable != nullptr); - VLOG(2) << "Executing with PJRT ..."; - - OP_REQUIRES_OK(ctx, RunPjRtExecutable(inputs, variables, *result, - pjrt_device_compiler->client(), - pjrt_executable, ctx)); - - VLOG(2) << "Completed executing with PJRT!"; + bool use_pjrt = + GetXlaOpsCommonFlags() + ->tf_xla_use_device_api.IsEnabledInXlaCompileOnDemandForDevice( + platform_info_.device_type()); + if (use_pjrt) { + std::vector variables; + std::vector args; + // Lock variables for the whole duration of compile + execute. + OP_REQUIRES_OK(ctx, GetAndLockVariablesAndBuildXlaCompilerArguments( + *ctx, inputs, *constant_indices_or, + variable_indices, &variables, &args)); + + PjRtDeviceCompiler* pjrt_device_compiler = nullptr; + xla::PjRtLoadedExecutable* pjrt_executable = nullptr; + absl::Status status = Compile(args, ctx, &pjrt_device_compiler, &profiler, + &result, &pjrt_executable); + // Hold the reference to the XLA device compiler and profiler during + // evaluation. (We could probably free them sooner because the ResourceMgr + // will retain references, but this is more obviously correct.) + // We must also ensure these pointers are freed even if compilation fails. + core::ScopedUnref pjrt_device_compiler_ref(pjrt_device_compiler); + core::ScopedUnref profiler_ref(profiler); + OP_REQUIRES_OK(ctx, status); + + VLOG(2) << "Compiled op with PJRT: " << ctx->status(); + VLOG(2) << "result != nullptr: " << (result != nullptr); + VLOG(2) << "pjrt_executable != nullptr: " << (pjrt_executable != nullptr); + VLOG(2) << "Executing with PJRT ..."; + + OP_REQUIRES_OK(ctx, RunPjRtExecutable(inputs, variables, *result, + pjrt_device_compiler->client(), + pjrt_executable, ctx)); + + VLOG(2) << "Completed executing with PJRT!"; + } else { + ResourceVarsSnapshot variable_args; + std::vector args; + // Lock variables only for generating XlaCompiler::Arguments and then + // release them. + { + std::vector variables; + OP_REQUIRES_OK(ctx, GetAndLockVariablesAndBuildXlaCompilerArguments( + *ctx, inputs, *constant_indices_or, + variable_indices, &variables, &args)); + OP_REQUIRES_OK(ctx, SnapshotResourceVariables(ctx, variable_indices, + variables, &variable_args)); + } + + XlaDeviceCompiler* xla_device_compiler = nullptr; + xla::LocalExecutable* executable = nullptr; + absl::Status status = Compile(args, ctx, &xla_device_compiler, &profiler, + &result, &executable); + // Hold the reference to the XLA device compiler and profiler during + // evaluation. (We could probably free them sooner because the ResourceMgr + // will retain references, but this is more obviously correct.) + // We must also ensure these pointers are freed even if compilation fails. + core::ScopedUnref xla_device_compiler_ref(xla_device_compiler); + core::ScopedUnref profiler_ref(profiler); + OP_REQUIRES_OK(ctx, status); + + // Locks are acquired again when populating the `ctx` outputs. + OP_REQUIRES_OK( + ctx, Run(variable_args, result, xla_device_compiler, executable, ctx)); + } } } // namespace tensorflow diff --git a/tensorflow/compiler/jit/xla_compile_on_demand_op.h b/tensorflow/compiler/jit/xla_compile_on_demand_op.h index e368de9f5e2236..113972706271e7 100644 --- a/tensorflow/compiler/jit/xla_compile_on_demand_op.h +++ b/tensorflow/compiler/jit/xla_compile_on_demand_op.h @@ -21,16 +21,20 @@ limitations under the License. #include -#include "absl/status/status.h" #include "tensorflow/compiler/jit/device_compilation_cluster_signature.h" #include "tensorflow/compiler/jit/device_compilation_profiler.h" -#include "tensorflow/compiler/jit/device_compiler.h" +#include "tensorflow/compiler/jit/variable_info.h" +#include "tensorflow/compiler/jit/variable_info_util.h" +#include "tensorflow/compiler/jit/xla_launch_util.h" #include "tensorflow/compiler/jit/xla_platform_info.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" +#include "xla/client/local_client.h" #include "xla/pjrt/pjrt_client.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/lib/core/status.h" namespace tensorflow { @@ -43,6 +47,14 @@ class XlaCompileOnDemandOp : public OpKernel { void Compute(OpKernelContext* ctx) override; private: + absl::Status Compile(const std::vector& args, + OpKernelContext* ctx, + DeviceCompiler** + xla_device_compiler, + DeviceCompilationProfiler** profiler, + const XlaCompiler::CompilationResult** result, + xla::LocalExecutable** executable); + absl::Status Compile(const std::vector& args, OpKernelContext* ctx, DeviceCompiler* + xla_device_compiler, + xla::LocalExecutable* executable, OpKernelContext* ctx); + const XlaPlatformInfo platform_info_; // Canonicalized function to compile derived from the Op attributes. diff --git a/tensorflow/compiler/jit/xla_compile_util.cc b/tensorflow/compiler/jit/xla_compile_util.cc index 021d7f55911ad6..c9c80d244b35b9 100644 --- a/tensorflow/compiler/jit/xla_compile_util.cc +++ b/tensorflow/compiler/jit/xla_compile_util.cc @@ -78,6 +78,12 @@ absl::StatusOr> CreateSingleOpGraph( return graph; } +bool UsePjRtForSingleDeviceCompilation(const DeviceType& device_type) { + const auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api; + return rollout_config.IsEnabledInXlaLaunchForDevice(device_type) || + rollout_config.IsEnabledInXlaCompileOnDemandForDevice(device_type) || + rollout_config.IsEnabledInXlaCompileAndRunForDevice(device_type); +} std::string GetPjRtDeviceCompilerResourceName(const DeviceType& device_type) { return absl::StrCat(kPjRtDeviceCompilerResourceName, "_", diff --git a/tensorflow/compiler/jit/xla_compile_util.h b/tensorflow/compiler/jit/xla_compile_util.h index 81dc51169a91cf..d722ba8e784e61 100644 --- a/tensorflow/compiler/jit/xla_compile_util.h +++ b/tensorflow/compiler/jit/xla_compile_util.h @@ -45,6 +45,9 @@ absl::StatusOr> CreateSingleOpGraph( const NodeDef& node_def, absl::Span args, absl::Span result_types); +// Checks if single device compilation and execution with PJRT is enabled for +// `device_type` in either the XlaLaunch op or the XlaCompileOnDemand op. +bool UsePjRtForSingleDeviceCompilation(const DeviceType& device_type); // Gets the resource name of the PjRt DeviceCompiler for `device_type`. std::string GetPjRtDeviceCompilerResourceName(const DeviceType& device_type); diff --git a/tensorflow/compiler/jit/xla_compile_util_test.cc b/tensorflow/compiler/jit/xla_compile_util_test.cc index fa529c603e8fa9..7e55498ec42d33 100644 --- a/tensorflow/compiler/jit/xla_compile_util_test.cc +++ b/tensorflow/compiler/jit/xla_compile_util_test.cc @@ -75,6 +75,49 @@ TEST_F(OpsTestBase, CreateSingleOpGraph) { EXPECT_EQ(retval_input_node->name(), "identity_op"); } +TEST(XlaCompileUtilTest, PjRtXlaLaunchFlagTest) { + EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); + + // Flag is turned on, but no device is allowlisted. + auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api; + rollout_config.enabled_for_xla_launch_ = true; + + EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); + + // Flag is turned on, some device is allowlisted, but the requested one isn't. + rollout_config.AllowForDeviceInXlaLaunch(DeviceType(DEVICE_GPU)); + EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); + + // Flag is turned on and the requested device is allowlisted. + rollout_config.AllowForDeviceInXlaLaunch(DeviceType(DEVICE_CPU)); + EXPECT_TRUE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); + + // The requested device is allowlisted, but the flag is turned off. + rollout_config.enabled_for_xla_launch_ = false; + EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); +} + +TEST(XlaCompileUtilTest, PjRtXlaCompileOnDemandFlagTest) { + EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); + + // Flag is turned on, but no device is allowlisted. + auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api; + rollout_config.enabled_for_compile_on_demand_ = true; + + EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); + + // Flag is turned on, some device is allowlisted, but the requested one isn't. + rollout_config.AllowForDeviceInXlaCompileOnDemand(DeviceType(DEVICE_GPU)); + EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); + + // Flag is turned on and the requested device is allowlisted. + rollout_config.AllowForDeviceInXlaCompileOnDemand(DeviceType(DEVICE_CPU)); + EXPECT_TRUE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); + + // The requested device is allowlisted, but the flag is turned off. + rollout_config.enabled_for_compile_on_demand_ = false; + EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU))); +} TEST(XlaCompileUtilTest, PjRtDeviceCompilerResourceName) { EXPECT_EQ(GetPjRtDeviceCompilerResourceName(DeviceType(DEVICE_TPU)), diff --git a/tensorflow/compiler/jit/xla_compiler_options_util.cc b/tensorflow/compiler/jit/xla_compiler_options_util.cc index 31c5f80053deb3..e17815cf4e91f4 100644 --- a/tensorflow/compiler/jit/xla_compiler_options_util.cc +++ b/tensorflow/compiler/jit/xla_compiler_options_util.cc @@ -20,15 +20,19 @@ limitations under the License. #include "tensorflow/compiler/jit/device_compiler.h" #include "tensorflow/compiler/jit/xla_platform_info.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" +#include "xla/client/local_client.h" #include "xla/pjrt/pjrt_client.h" +#include "xla/stream_executor/host/host_platform_id.h" +#include "xla/stream_executor/stream.h" #include "xla/tsl/framework/device_id_utils.h" -#include "xla/tsl/framework/device_type.h" #include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/types.h" namespace tensorflow { namespace { +using XlaDeviceCompiler = + DeviceCompiler; using PjRtDeviceCompiler = DeviceCompiler; @@ -49,12 +53,41 @@ inline void LogOptions(const XlaCompiler::Options& options) { } } // namespace +XlaCompiler::Options GenerateCompilerOptions( + const XlaDeviceCompiler& xla_device_compiler, + const FunctionLibraryRuntime& function_library, DeviceBase* device, + se::Stream* stream, const XlaPlatformInfo& platform_info, + bool has_ref_vars) { + XlaCompiler::Options options; + options.client = static_cast(xla_device_compiler.client()); + if (stream != nullptr) { + options.device_ordinal = stream->parent()->device_ordinal(); + } + options.device_type = xla_device_compiler.device_type(); + options.flib_def = function_library.GetFunctionLibraryDefinition(); + options.graph_def_version = function_library.graph_def_version(); + options.allow_cpu_custom_calls = + (platform_info.platform_id() == se::host::kHostPlatformId); + options.device_allocator = GetAllocator(device, stream, platform_info); + if (platform_info.xla_device_metadata()) { + options.shape_determination_fns = + platform_info.xla_device_metadata()->default_shape_determination_fns(); + } + // If reference variables are not present in the graph, we can safely alias + // passthrough parameters without performing a copy. + options.alias_passthrough_params = + !has_ref_vars && !platform_info.is_on_xla_device(); + + LogOptions(options); + return options; +} + XlaCompiler::Options GenerateCompilerOptionsForTfrtTpu( - const tsl::DeviceType& device_type, + const XlaDeviceCompiler& xla_device_compiler, const FunctionLibraryRuntime& function_library) { XlaCompiler::Options options; // TODO(b/238830423): consider device_ordinal and shape_determination_fns. - options.device_type = device_type; + options.device_type = xla_device_compiler.device_type(); options.flib_def = function_library.GetFunctionLibraryDefinition(); options.graph_def_version = function_library.graph_def_version(); options.allow_cpu_custom_calls = false; @@ -104,9 +137,6 @@ XlaCompiler::Options GenerateCompilerOptionsForPjRt( } else if (pjrt_device_compiler != nullptr) { options.device_type = pjrt_device_compiler->device_type(); } - if (pjrt_device_compiler != nullptr) { - options.client = pjrt_device_compiler->client(); - } // TODO(b/255826209): Confirm below options are correctly set after testing. options.allow_cpu_custom_calls = false; options.alias_passthrough_params = false; diff --git a/tensorflow/compiler/jit/xla_compiler_options_util.h b/tensorflow/compiler/jit/xla_compiler_options_util.h index 96ca45311d9756..8c6fb8e0eaf2ab 100644 --- a/tensorflow/compiler/jit/xla_compiler_options_util.h +++ b/tensorflow/compiler/jit/xla_compiler_options_util.h @@ -19,16 +19,26 @@ limitations under the License. #include "tensorflow/compiler/jit/device_compiler.h" #include "tensorflow/compiler/jit/xla_platform_info.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" +#include "xla/client/local_client.h" #include "xla/pjrt/pjrt_client.h" -#include "xla/tsl/framework/device_type.h" +#include "xla/stream_executor/stream.h" #include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/function.h" namespace tensorflow { +// Returns created options for the XLA compiler. +XlaCompiler::Options GenerateCompilerOptions( + const DeviceCompiler& + xla_device_compiler, + const FunctionLibraryRuntime& function_library, DeviceBase* device, + se::Stream* stream, const XlaPlatformInfo& platform_info, + bool has_ref_vars); + // Returns created options for XLA compiler when TFRT-TPU is used. XlaCompiler::Options GenerateCompilerOptionsForTfrtTpu( - const tsl::DeviceType& device_type, + const DeviceCompiler& + xla_device_compiler, const FunctionLibraryRuntime& function_library); // Returns created options for XLA compiler when PjRt (Device API) is used for diff --git a/tensorflow/compiler/jit/xla_compiler_options_util_test.cc b/tensorflow/compiler/jit/xla_compiler_options_util_test.cc index 539acab667b198..0ea5d3222204fb 100644 --- a/tensorflow/compiler/jit/xla_compiler_options_util_test.cc +++ b/tensorflow/compiler/jit/xla_compiler_options_util_test.cc @@ -28,15 +28,19 @@ limitations under the License. #include "tensorflow/compiler/jit/pjrt_device_compiler_client.h" #include "tensorflow/compiler/jit/test_util.h" #include "tensorflow/compiler/jit/xla_device.h" +#include "tensorflow/compiler/jit/xla_device_compiler_client.h" #include "tensorflow/compiler/jit/xla_platform_info.h" #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_argument.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "xla/client/client_library.h" +#include "xla/client/local_client.h" #include "xla/pjrt/pjrt_client.h" #include "xla/shape.h" -#include "xla/stream_executor/device_address_allocator.h" +#include "xla/stream_executor/device_memory_allocator.h" +#include "xla/stream_executor/host/host_platform_id.h" #include "xla/stream_executor/platform.h" #include "xla/tsl/platform/statusor.h" #include "xla/xla_data.pb.h" @@ -50,11 +54,25 @@ limitations under the License. namespace tensorflow { namespace { +using XlaDeviceCompiler = + DeviceCompiler; +using XlaDeviceExecutablePersistor = + DeviceExecutablePersistor; using PjRtDeviceCompiler = DeviceCompiler; using PjRtDeviceExecutablePersistor = DeviceExecutablePersistor; +XlaDeviceCompiler* CreateXlaDeviceCompiler(DeviceType device_type, + xla::LocalClient* local_client) { + auto persistor = std::make_unique( + XlaDeviceExecutablePersistor::Config(), device_type); + auto compiler_client = + std::make_unique(local_client); + return new XlaDeviceCompiler(std::move(persistor), + std::move(compiler_client)); +} + PjRtDeviceCompiler* CreatePjRtDeviceCompiler(DeviceType device_type, xla::PjRtClient* pjrt_client) { auto persistor = std::make_unique( @@ -211,13 +229,103 @@ TEST_F(XlaCompilerOptionsTest, PjRtOptionsNonXlaDevice) { tensorflow::XlaLayoutPreference::kNoPreference); } +TEST_F(XlaCompilerOptionsTest, XlaOptions) { + device_setup_.AddDevicesAndSetUp({DEVICE_XLA_GPU}); + Device* device = device_setup_.GetDevice(DEVICE_XLA_GPU); + + xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); + DeviceType device_type = DeviceType(DEVICE_XLA_GPU); + DeviceType compilation_device_type = DeviceType(DEVICE_GPU_XLA_JIT); + + auto xla_device_compiler = + CreateXlaDeviceCompiler(compilation_device_type, client); + core::ScopedUnref xla_device_compiler_ref(xla_device_compiler); + + se::Platform::Id platform_id = se::host::kHostPlatformId; + auto xla_device_metadata = CreateXlaDeviceMetadata(compilation_device_type); + std::shared_ptr custom_allocator; + XlaPlatformInfo platform_info( + device_type, platform_id, xla_device_metadata.get(), + /*pjrt_device_metadata=*/nullptr, custom_allocator); + + XlaCompiler::Options options = + GenerateCompilerOptions(*xla_device_compiler, *device_setup_.flr(), + device, nullptr, platform_info, false); + + EXPECT_EQ(options.device_type, compilation_device_type); + EXPECT_NE(options.flib_def, nullptr); + EXPECT_EQ(options.graph_def_version, TF_GRAPH_DEF_VERSION); + EXPECT_TRUE(options.allow_cpu_custom_calls); + EXPECT_NE(options.device_allocator, nullptr); + EXPECT_FALSE(options.alias_passthrough_params); + // Check if options have the supplied shape determination functions set. + TF_ASSERT_OK_AND_ASSIGN( + auto shape, options.shape_determination_fns.shape_representation_fn( + TensorShape(), DT_FLOAT, false, + tensorflow::XlaLayoutPreference::kTpuPreferLinearLayout)); + EXPECT_EQ(shape, xla::Shape()); + EXPECT_EQ(options.shape_determination_fns.layout_preference_fn( + TensorShape(), DT_FLOAT, std::nullopt), + tensorflow::XlaLayoutPreference::kTpuPreferLinearLayout); +} + +TEST_F(XlaCompilerOptionsTest, XlaOptionsHasRefVarsNoXlaDeviceMetadata) { + device_setup_.AddDevicesAndSetUp({DEVICE_CPU}); + Device* device = device_setup_.GetDevice(DEVICE_CPU); + + xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); + DeviceType device_type = DeviceType(DEVICE_CPU); + DeviceType compilation_device_type = DeviceType(DEVICE_CPU_XLA_JIT); + + auto xla_device_compiler = + CreateXlaDeviceCompiler(compilation_device_type, client); + core::ScopedUnref xla_device_compiler_ref(xla_device_compiler); + + se::Platform::Id platform_id = se::host::kHostPlatformId; + std::shared_ptr custom_allocator; + XlaPlatformInfo platform_info( + device_type, platform_id, /*xla_device_metadata=*/nullptr, + /*pjrt_device_metadata=*/nullptr, custom_allocator); + + XlaCompiler::Options options = + GenerateCompilerOptions(*xla_device_compiler, *device_setup_.flr(), + device, nullptr, platform_info, false); + + EXPECT_EQ(options.device_type, compilation_device_type); + EXPECT_NE(options.flib_def, nullptr); + EXPECT_EQ(options.graph_def_version, TF_GRAPH_DEF_VERSION); + EXPECT_TRUE(options.allow_cpu_custom_calls); + EXPECT_NE(options.device_allocator, nullptr); + EXPECT_TRUE(options.alias_passthrough_params); + // Check whether options have default shape determination functions set. + TF_ASSERT_OK_AND_ASSIGN( + auto shape, options.shape_determination_fns.shape_representation_fn( + TensorShape(), DT_FLOAT, false, + tensorflow::XlaLayoutPreference::kNoPreference)); + xla::ShapeProto shape_proto; + shape_proto.set_element_type(xla::PrimitiveType::F32); + shape_proto.mutable_layout(); + TF_ASSERT_OK_AND_ASSIGN(auto expected_shape, + xla::Shape::FromProto(shape_proto)); + EXPECT_EQ(shape, expected_shape); + EXPECT_EQ(options.shape_determination_fns.layout_preference_fn( + TensorShape(), DT_FLOAT, std::nullopt), + tensorflow::XlaLayoutPreference::kNoPreference); +} + TEST_F(XlaCompilerOptionsTest, TfRtTpuOptions) { device_setup_.AddDevicesAndSetUp({DEVICE_TPU_NODE}); + // Just use the default local client for testing purposes. + xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); DeviceType compilation_device_type = DeviceType(DEVICE_TPU_XLA_JIT); + auto xla_device_compiler = + CreateXlaDeviceCompiler(compilation_device_type, client); + core::ScopedUnref xla_device_compiler_ref(xla_device_compiler); + XlaCompiler::Options options = GenerateCompilerOptionsForTfrtTpu( - compilation_device_type, *device_setup_.flr()); + *xla_device_compiler, *device_setup_.flr()); EXPECT_EQ(options.device_type, compilation_device_type); EXPECT_NE(options.flib_def, nullptr); diff --git a/tensorflow/compiler/jit/xla_device.cc b/tensorflow/compiler/jit/xla_device.cc index 45488c44577684..0c8b36e49d1ef4 100644 --- a/tensorflow/compiler/jit/xla_device.cc +++ b/tensorflow/compiler/jit/xla_device.cc @@ -17,48 +17,48 @@ limitations under the License. #include -#include -#include #include #include #include -#include #include "absl/base/call_once.h" -#include "absl/log/log.h" -#include "absl/status/status.h" -#include "absl/status/statusor.h" +#include "absl/memory/memory.h" #include "absl/strings/match.h" -#include "absl/strings/str_cat.h" -#include "absl/strings/string_view.h" +#include "tensorflow/compiler/jit/defs.h" #include "tensorflow/compiler/jit/pjrt_device_context.h" #include "tensorflow/compiler/jit/xla_compile_on_demand_op.h" -#include "tensorflow/compiler/jit/xla_tensor.h" -#include "tensorflow/compiler/tf2xla/layout_util.h" +#include "tensorflow/compiler/jit/xla_compile_util.h" +#include "tensorflow/compiler/jit/xla_device_context.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/client/client_library.h" -#include "xla/client/local_client.h" -#include "xla/service/shaped_buffer.h" -#include "xla/shape.h" -#include "xla/stream_executor/platform.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" +#include "xla/service/stream_pool.h" +#include "tensorflow/core/common_runtime/device.h" +#include "tensorflow/core/common_runtime/device_factory.h" #include "tensorflow/core/common_runtime/dma_helper.h" -#include "tensorflow/core/common_runtime/local_device.h" +#include "tensorflow/core/common_runtime/function.h" +#include "tensorflow/core/common_runtime/graph_constructor.h" +#include "tensorflow/core/common_runtime/renamed_device.h" #include "tensorflow/core/framework/allocator.h" #include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/kernel_def.pb.h" +#include "tensorflow/core/framework/node_def_builder.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor.pb.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/graph/types.h" -#include "tensorflow/core/platform/mutex.h" -#include "tensorflow/core/platform/threadpool.h" +#include "tensorflow/core/lib/core/notification.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" +#include "tensorflow/core/profiler/lib/traceme.h" #include "tensorflow/core/public/session_options.h" +#include "tensorflow/core/public/version.h" #include "tensorflow/core/tfrt/common/async_value_tensor.h" +#include "tensorflow/core/util/device_name_utils.h" +#include "tensorflow/core/util/dump_graph.h" +#include "tensorflow/core/util/stream_executor_util.h" namespace tensorflow { @@ -77,6 +77,59 @@ absl::Status DefaultPaddedShapeFn(const Tensor& tensor, xla::Shape* shape) { return absl::OkStatus(); } +// Caches a XlaDeviceAllocator per pair. A +// XlaDeviceAllocator is created on demand and is associated with a +// XlaDevice. It outlives the device itself (for instance, the buffer +// backing a tensor holds a pointer to the allocator for book-keeping, +// and this buffer can outlast the device). +class XlaDeviceAllocatorState { + public: + // Creates or returns a cached XlaDeviceAllocator for a given + // backend and device_ordinal. + static XlaDeviceAllocator* GetOrCreateXlaDeviceAllocator( + const xla::Backend* backend, int device_ordinal); + + private: + // Returns the singleton instance of XlaDeviceAllocatorState. + static XlaDeviceAllocatorState& Singleton(); + XlaDeviceAllocatorState(); + ~XlaDeviceAllocatorState(); + + mutex allocator_mutex_; // Guards the singleton allocator state. + std::unordered_map, + std::unique_ptr, + hash>> + allocators_ TF_GUARDED_BY(allocator_mutex_); + + XlaDeviceAllocatorState(const XlaDeviceAllocatorState&) = delete; + void operator=(const XlaDeviceAllocatorState&) = delete; +}; + +/* static */ XlaDeviceAllocatorState& XlaDeviceAllocatorState::Singleton() { + static auto a = new XlaDeviceAllocatorState; + return *a; +} + +XlaDeviceAllocatorState::XlaDeviceAllocatorState() = default; +XlaDeviceAllocatorState::~XlaDeviceAllocatorState() = default; + +XlaDeviceAllocator* XlaDeviceAllocatorState::GetOrCreateXlaDeviceAllocator( + const xla::Backend* backend, int device_ordinal) { + XlaDeviceAllocatorState& state = Singleton(); + mutex_lock lock(state.allocator_mutex_); + + auto it = state.allocators_.find({backend, device_ordinal}); + if (it != state.allocators_.end()) { + return it->second.get(); + } + + std::unique_ptr alloc = + std::make_unique( + backend->stream_executors()[device_ordinal]); + XlaDeviceAllocator* alloc_ptr = alloc.get(); + state.allocators_[{backend, device_ordinal}] = std::move(alloc); + return alloc_ptr; +} namespace { @@ -141,6 +194,11 @@ const DeviceType& XlaDevice::Metadata::jit_device_type() const { return GetMetadataFromDevice(ctx->device(), metadata); } +/* static */ mutex XlaDevice::global_mu_(LINKER_INITIALIZED); +/* static */ std::vector>* + XlaDevice::global_compute_streams_ = + new std::vector>; + XlaDevice::XlaDevice(const SessionOptions& session_options, const Options& options) : LocalDevice(session_options, @@ -159,15 +217,27 @@ XlaDevice::XlaDevice(const SessionOptions& session_options, platform_(options.platform), intra_op_parallelism_threads_( session_options.config.intra_op_parallelism_threads()), + use_multiple_streams_(options.use_multiple_streams), shape_determination_fns_(options.shape_determination_fns), - allowed_devices_(options.allowed_devices) { + allowed_devices_(options.allowed_devices), + use_global_compute_stream_(options.use_global_compute_stream) { if (options.shape_determination_fns.empty()) { LOG(ERROR) << "shape_representation_fns must be non-empty."; } VLOG(1) << "Created XLA device " << options.compilation_device_name << " " << options.device_ordinal << " " << this; + VLOG(1) << "XlaDevice options: use_multiple_streams: " + << options.use_multiple_streams << " use_global_compute_stream: " + << options.use_global_compute_stream; thread_pool_.reset(new thread::ThreadPool(session_options.env, "xla_device", /*num_threads=*/1)); + + // We have multiple device to device streams to allow for some concurrency + // between transfers. The particular value of '4' is chosen fairly + // arbitrarily. It may be necessary to make this tunable via + // XlaDevice::Options. + static constexpr int kNumDeviceToDeviceStreams = 4; + device_to_device_streams_.resize(kNumDeviceToDeviceStreams); } XlaDevice::~XlaDevice() { @@ -178,6 +248,19 @@ XlaDevice::~XlaDevice() { } } +absl::StatusOr XlaDevice::GetOrCreateClient() const { + // We lazily create the client because the platform commits to the + // details of the host hardware when the client is created, so we + // don't want to do it until we get a chance to hook the platform up + // to a simulator. + + xla::LocalClientOptions options; + options.set_platform(platform_) + .set_allowed_devices(allowed_devices_) + .set_intra_op_parallelism_threads(intra_op_parallelism_threads_); + return xla::ClientLibrary::GetOrCreateLocalClient(options); +} + Allocator* XlaDevice::GetAllocator(AllocatorAttributes attr) { mutex_lock lock(mu_); return GetAllocatorLocked(attr); @@ -189,9 +272,17 @@ Allocator* XlaDevice::GetAllocatorLocked(AllocatorAttributes attr) { } if (xla_allocator_ == nullptr) { - VLOG(1) << "XlaDevice " << this << " uses AsyncValueAllocator"; - pjrt_allocator_ = std::make_unique(); - xla_allocator_ = pjrt_allocator_.get(); + if (UsePjRtForSingleDeviceCompilation(device_name_)) { + VLOG(1) << "XlaDevice " << this << " uses AsyncValueAllocator"; + pjrt_allocator_ = std::make_unique(); + xla_allocator_ = pjrt_allocator_.get(); + } else { + // TODO(b/78468222): This can fail, at least when the backend is GPU and + // there is no GPU on the host. + xla::Backend* backend = GetOrCreateClient().value()->mutable_backend(); + xla_allocator_ = XlaDeviceAllocatorState::GetOrCreateXlaDeviceAllocator( + backend, device_ordinal_); + } } return xla_allocator_; } @@ -201,24 +292,136 @@ absl::Status XlaDevice::EnsureDeviceContextOk() { return GetDeviceContextLocked().status(); } +absl::Status XlaDevice::EnsureStreamOkLocked( + xla::Backend* backend, const std::string& name, + std::shared_ptr* stream, bool* stream_was_changed) { + if (!(*stream) || !(*stream)->ok()) { + xla::StreamPool::Ptr ptr; + TF_ASSIGN_OR_RETURN(ptr, backend->BorrowStream(device_ordinal_)); + *stream = std::shared_ptr(std::move(ptr)); + VLOG(1) << "XlaDevice " << this << " new " << name + << " stream=" << (*stream); + *stream_was_changed = true; + } + return absl::OkStatus(); +} + absl::StatusOr> XlaDevice::GetDeviceContextLocked() { - if (device_contexts_.empty()) { - for (const auto& iter : shape_determination_fns_) { - auto device_context = new PjRtDeviceContext(iter); - VLOG(1) << "XlaDevice " << this << " new PjRtDeviceContext " - << device_context; - device_contexts_.emplace_back(device_context); + if (UsePjRtForSingleDeviceCompilation(device_name_)) { + if (device_contexts_.empty()) { + for (const auto& iter : shape_determination_fns_) { + auto device_context = new PjRtDeviceContext(iter); + VLOG(1) << "XlaDevice " << this << " new PjRtDeviceContext " + << device_context; + device_contexts_.emplace_back(device_context); + } + if (use_accelerator_device_info_) { + auto accelerator_device_info = + std::make_unique(); + accelerator_device_info->default_context = device_contexts_.at(0); + set_tensorflow_accelerator_device_info(accelerator_device_info.get()); + accelerator_device_info_ = std::move(accelerator_device_info); + VLOG(1) << "XlaDevice " << this << " new AcceleratorDeviceInfo " + << accelerator_device_info_.get(); + } + } + + return device_contexts_; + } + + TF_ASSIGN_OR_RETURN(xla::LocalClient * client, GetOrCreateClient()); + xla::Backend* backend = client->mutable_backend(); + + // Ensure all our streams are valid, borrowing new streams if necessary. + bool need_new_device_context = device_contexts_.empty(); + if (use_global_compute_stream_) { + mutex_lock lock(global_mu_); + if (global_compute_streams_->size() <= device_ordinal_) { + global_compute_streams_->resize(device_ordinal_ + 1, nullptr); } - if (use_accelerator_device_info_) { - auto accelerator_device_info = - std::make_unique(); - accelerator_device_info->default_context = device_contexts_.at(0); - set_tensorflow_accelerator_device_info(accelerator_device_info.get()); - accelerator_device_info_ = std::move(accelerator_device_info); - VLOG(1) << "XlaDevice " << this << " new AcceleratorDeviceInfo " - << accelerator_device_info_.get(); + + auto& global_stream = global_compute_streams_->at(device_ordinal_); + if (global_stream != nullptr && global_stream->ok()) { + stream_ = global_stream; + } else { + // Directly create the stream here instead of borrowing from the stream + // pool to avoid potential lifetime issues. + TF_ASSIGN_OR_RETURN( + stream_, + backend->stream_executors()[device_ordinal_]->CreateStream()); + TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "stream", &stream_, + &need_new_device_context)); + (*global_compute_streams_)[device_ordinal_] = stream_; + } + } else { + TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "stream", &stream_, + &need_new_device_context)); + } + + std::shared_ptr host_to_device_stream; + std::shared_ptr device_to_host_stream; + std::vector> device_to_device_streams; + if (use_multiple_streams_) { + TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "host_to_device_stream", + &host_to_device_stream_, + &need_new_device_context)); + for (std::shared_ptr& stream : device_to_device_streams_) { + TF_RETURN_IF_ERROR( + EnsureStreamOkLocked(backend, "device_to_device_stream", &stream, + &need_new_device_context)); } + host_to_device_stream = host_to_device_stream_; + device_to_device_streams = device_to_device_streams_; + // The data transfer requests from device to host could arrive out of order, + // so a single stream would cause deadlock. For this case, + // xla_device_context would borrow a stream for each transfer request. + device_to_host_stream = nullptr; + } else { + host_to_device_stream = stream_; + device_to_host_stream = stream_; + device_to_device_streams = {stream_}; + } + + if (!need_new_device_context) { + return device_contexts_; + } + + // At this point we know we need a new device context. + // Call GetAllocator for the side-effect of ensuring the allocator is created. + GetAllocatorLocked({}); + for (const auto& iter : device_contexts_) { + iter->Unref(); + } + // The XlaDeviceContext keeps a reference count to the streams, and the + // XlaDeviceContext remains live for the duration of a Executor run. This + // ensures that the streams remain live for the duration of a run, even if + // an error is encountered and the streams are replaced with new ones. + for (const auto& iter : shape_determination_fns_) { + auto device_context = new XlaDeviceContext( + stream_, host_to_device_stream, device_to_host_stream, + device_to_device_streams, client, iter, thread_pool_.get()); + VLOG(1) << "XlaDevice " << this << " new XlaDeviceContext " + << device_context; + device_contexts_.emplace_back(device_context); + } + + // Create and set a new AcceleratorDeviceInfo, if necessary. + // + // TODO(b/78232898): This isn't thread-safe; there is a race between the call + // to set_tensorflow_accelerator_device_info() with ops that call the getter + // tensorflow_accelerator_device_info(). This isn't trivially fixed by adding + // locking to those methods; see the bug for details. Our only saving grace at + // the moment is that this race doesn't seem to occur in practice. + if (use_accelerator_device_info_) { + auto accelerator_device_info = + std::make_unique(); + accelerator_device_info->stream = stream_.get(); + accelerator_device_info->default_context = device_contexts_.at(0); + set_tensorflow_accelerator_device_info(accelerator_device_info.get()); + accelerator_device_info_ = std::move(accelerator_device_info); + VLOG(1) << "XlaDevice " << this << " new AcceleratorDeviceInfo " + << accelerator_device_info_.get(); } return device_contexts_; @@ -280,6 +483,21 @@ void XlaDevice::ComputeAsync(AsyncOpKernel* op_kernel, OpKernelContext* context, absl::Status XlaDevice::Sync() { VLOG(1) << "XlaDevice::Sync"; + tsl::profiler::TraceMe activity("XlaDevice::Sync", + tsl::profiler::TraceMeLevel::kInfo); + std::shared_ptr stream; + { + mutex_lock lock(mu_); + stream = stream_; + } + if (!stream) return absl::OkStatus(); + + absl::Status status = stream->BlockHostUntilDone(); + TF_RETURN_IF_ERROR(status); + if (!stream->ok()) { + return absl::InternalError("XlaDevice::Sync() failed."); + } + VLOG(1) << "XlaDevice::Sync completed"; return absl::OkStatus(); } @@ -347,7 +565,24 @@ absl::Status XlaDevice::HandleDeviceError() { return absl::OkStatus(); } -absl::Status XlaDevice::RefreshStatus() { return absl::OkStatus(); } +absl::Status XlaDevice::RefreshStatus() { + std::shared_ptr stream; + { + mutex_lock lock(mu_); + stream = stream_; + } + if (!stream) { + return absl::OkStatus(); + } + absl::Status status = stream->RefreshStatus(); + if (!status.ok()) { + // Ignore errors from HandleDeviceError, since by definition the status is + // already non-ok, so there's nothing extra to report if HandleDeviceError + // itself returns an error. + HandleDeviceError().IgnoreError(); + } + return status; +} XlaDeviceOpRegistrations* RegisterXlaDeviceKernels( const char* device, const char* jit_device, diff --git a/tensorflow/compiler/jit/xla_device.h b/tensorflow/compiler/jit/xla_device.h index 30396631d4a5e7..bedbf87a3cd2e8 100644 --- a/tensorflow/compiler/jit/xla_device.h +++ b/tensorflow/compiler/jit/xla_device.h @@ -205,8 +205,15 @@ class XlaDevice : public LocalDevice { absl::Status RefreshStatus() override TF_LOCKS_EXCLUDED(mu_); + private: + absl::StatusOr GetOrCreateClient() const; Allocator* GetAllocatorLocked(AllocatorAttributes attr) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_); + absl::Status EnsureStreamOkLocked(xla::Backend* backend, + const std::string& name, + std::shared_ptr* stream, + bool* stream_was_changed) + TF_EXCLUSIVE_LOCKS_REQUIRED(mu_); // Return a vector of device context, ordered by the sequence in the given // shape_representation_fns. @@ -233,6 +240,24 @@ class XlaDevice : public LocalDevice { Allocator* xla_allocator_ TF_GUARDED_BY(mu_) = nullptr; // Not owned. std::unique_ptr pjrt_allocator_ TF_GUARDED_BY(mu_); + // Stream associated with this device. Operations enqueued on this + // stream are executed on the device. Operations include data + // copying back and forth between CPU and the device, and + // computations enqueued by XLA. + std::shared_ptr stream_ TF_GUARDED_BY(mu_); + // If false, only stream_ is valid and all computation and transfers use + // stream_. If true, computation is performed by stream_ and transfers are + // performed by host_to_device/device_to_device stream or borrowing a stream + // for each device to host transfer. + const bool use_multiple_streams_; + // If use_multiple_streams_, host to device transfers are performed using this + // stream. + std::shared_ptr host_to_device_stream_ TF_GUARDED_BY(mu_); + // If use_multiple_streams_, transfers between different devices are performed + // using these streams. + std::vector> device_to_device_streams_ + TF_GUARDED_BY(mu_); + // See comments in options. std::vector shape_determination_fns_; @@ -263,6 +288,15 @@ class XlaDevice : public LocalDevice { // platform will have resources allocated. For GPUs this will be // filled from visible_gpu_devices list from session configuration. std::optional> allowed_devices_; + + const bool use_global_compute_stream_; + + // A static vector with device_ordinal as its index, describing the global + // compute streams used in each XLA device. It is only used if + // `use_global_compute_stream` in `XlaDevice::Options` is set to true. + static mutex global_mu_; + static std::vector>* global_compute_streams_ + TF_GUARDED_BY(global_mu_); }; // Builds OpKernel registrations on 'device' for the JIT operators diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc index 223d3c5e2b9bd4..12c9f7d884ad66 100644 --- a/tensorflow/compiler/jit/xla_launch_util.cc +++ b/tensorflow/compiler/jit/xla_launch_util.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/jit/xla_launch_util.h" +#include #include #include #include @@ -23,7 +24,6 @@ limitations under the License. #include #include "absl/algorithm/container.h" -#include "absl/base/casts.h" #include "absl/cleanup/cleanup.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" @@ -38,29 +38,40 @@ limitations under the License. #include "tensorflow/compiler/jit/pjrt_tensor_buffer_util.h" #include "tensorflow/compiler/jit/variable_info.h" #include "tensorflow/compiler/jit/variable_info_util.h" +#include "tensorflow/compiler/jit/xla_tensor.h" #include "tensorflow/compiler/tf2xla/const_analysis.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/compiler/tf2xla/xla_helpers.h" #include "tensorflow/compiler/tf2xla/xla_resource.h" +#include "xla/client/local_client.h" #include "xla/future.h" +#include "xla/hlo/ir/hlo_input_output_alias_config.h" #include "xla/pjrt/pjrt_client.h" -#include "xla/pjrt/pjrt_compiler.h" +#include "xla/pjrt/pjrt_common.h" #include "xla/pjrt/pjrt_executable.h" -#include "xla/pjrt/plugin/xla_cpu/cpu_execute_options.h" -#include "xla/runtime/device_id.h" +#include "xla/service/executable.h" +#include "xla/service/maybe_owning_device_memory.h" +#include "xla/service/shaped_buffer.h" +#include "xla/service/transfer_manager.h" #include "xla/shape.h" +#include "xla/shape_util.h" #include "xla/status_macros.h" -#include "xla/stream_executor/device_address_allocator.h" +#include "xla/stream_executor/device_memory.h" +#include "xla/stream_executor/device_memory_allocator.h" +#include "xla/stream_executor/event.h" +#include "xla/stream_executor/host/host_platform_id.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/platform_manager.h" #include "xla/tsl/framework/device_id_utils.h" #include "xla/tsl/framework/serving_device_selector_policies.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/statusor.h" #include "xla/util.h" -#include "tensorflow/core/common_runtime/composite_device.h" #include "tensorflow/core/common_runtime/dma_helper.h" #include "tensorflow/core/common_runtime/gpu/gpu_serving_device_selector.h" #include "tensorflow/core/common_runtime/gpu_device_context.h" +#include "tensorflow/core/framework/allocator.h" #include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/resource_handle.h" @@ -68,12 +79,28 @@ limitations under the License. #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.h" +#include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/refcount.h" +#include "tensorflow/core/platform/errors.h" #include "tensorflow/core/platform/status.h" #include "tensorflow/core/tfrt/common/async_value_tensor.h" +#include "tsl/platform/casts.h" namespace tensorflow { namespace { +using xla::ScopedShapedBuffer; +using xla::ShapedBuffer; + +// Fetch the platform Id from device. +se::Platform::Id XlaPlatformInfoFromDevice(DeviceBase* device_base) { + auto device = static_cast(device_base); + se::Platform::Id platform_id = nullptr; + if (device->device_type() == DEVICE_CPU) { + platform_id = se::host::kHostPlatformId; + } + + return platform_id; +} absl::flat_hash_map CreateVariableLookup( const std::vector& variables) { @@ -109,6 +136,177 @@ absl::StatusOr> GetConstantInputIndicesFromContext( return constant_input_indices; } +XlaComputationLaunchContext::XlaComputationLaunchContext( + xla::LocalClient* client, + stream_executor::DeviceAddressAllocator* xla_allocator, int device_ordinal, + bool allocate_xla_tensors, bool use_multiple_streams) + : client_(client), + xla_allocator_(xla_allocator), + allocate_xla_tensors_(allocate_xla_tensors), + use_multiple_streams_(use_multiple_streams), + device_ordinal_(device_ordinal) { + if (use_multiple_streams_) { + CHECK(allocate_xla_tensors_) << "To use multiple streams correctly we must " + "be allocating XLA tensors!"; + } +} + +// Fills in `execution_input` with `buffer` for `index`. +static void PopulateExecutionInputBuffer( + xla::ExecutionInput& execution_input, const xla::ShapeIndex& index, + stream_executor::DeviceAddressBase buffer, bool donate_buffer, + int device_ordinal, stream_executor::DeviceAddressAllocator* allocator) { + xla::MaybeOwningDeviceAddress* in_buffer = + execution_input.MutableBuffer(index); + if (donate_buffer) { + // Here we pass ownership of the buffer to execution_input without releasing + // ownership from the caller of PopulateExecutionInputBuffer. If execution + // succeeds, we'll take back that duplicate ownership in + // GetOrCreateTensorForOutput. If execution fails, the ExecutionInput will + // release that duplicate ownership automatically. + *in_buffer = stream_executor::ScopedDeviceAddress( + buffer, device_ordinal, allocator); + } else { + *in_buffer = buffer; + } +} + +absl::StatusOr> +XlaComputationLaunchContext::PopulateInputs( + OpKernelContext* ctx, + const XlaCompiler::CompilationResult* compilation_result, + const absl::flat_hash_map& resource_vars, + int missing_ctx_input_prefix, + const xla::HloInputOutputAliasConfig& input_output_alias) { + std::vector arguments; + arguments.reserve(compilation_result->xla_input_shapes.size()); + + xla::ShapeIndex root_index = {}; + + for (int i = 0; i < compilation_result->xla_input_shapes.size(); ++i) { + int arg_num = compilation_result->input_mapping[i]; + CHECK_GE(arg_num, missing_ctx_input_prefix); + const xla::Shape& device_shape = compilation_result->xla_input_shapes[i]; + + auto resource_var_it = resource_vars.find(arg_num); + bool is_resource_variable = resource_var_it != resource_vars.end(); + bool is_updated_resource_variable = + is_resource_variable && + absl::c_any_of(compilation_result->resource_updates, + [&](const XlaCompiler::ResourceUpdate& update) { + // XlaCompiler records `arg_num` (instead of kernel + // parameters) in `resource_updates`. + return update.input_index == arg_num && + update.modified; + }); + + const Tensor* t = is_resource_variable + ? resource_var_it->second + : &(ctx->input(arg_num - missing_ctx_input_prefix)); + CHECK(t); + bool donate_buffer = t->RefCountIsOne() && is_updated_resource_variable && + input_output_alias.ParameterHasAlias(i, root_index); + VLOG(3) << "Processing input: " << i + << "; is_resource_variable=" << is_resource_variable + << "; is_updated_resource_variable=" << is_updated_resource_variable + << "; donate_buffer=" << donate_buffer; + + if (use_multiple_streams_) { + CHECK(ctx->op_device_context() && ctx->op_device_context()->stream()) + << "Must have a stream available when using XLA tensors!"; + XlaTensor* xla_tensor = XlaTensor::FromTensor(t); + CHECK(xla_tensor); + xla_tensor->WaitForDefinitionEventOnStream( + ctx->op_device_context()->stream()); + } + + arguments.emplace_back(&device_shape); + xla::ExecutionInput& execution_input = arguments.back(); + stream_executor::DeviceAddressBase dmem = + XlaTensor::DeviceMemoryFromTensor(*t); + PopulateExecutionInputBuffer(execution_input, root_index, dmem, + donate_buffer, device_ordinal_, + xla_allocator_); + } + return std::move(arguments); +} + +// Construct the tensor for the given type and buffer. +static Tensor MakeTensor(DataType dtype, const TensorShape& shape, + stream_executor::DeviceAddressBase buffer, + Allocator* allocator) { + size_t expected_size = shape.num_elements() * DataTypeSize(dtype); + auto* tensor_buffer = new XlaTensorBuffer(buffer.opaque(), expected_size, + buffer.size(), allocator); + Tensor t(dtype, shape, tensor_buffer); + tensor_buffer->Unref(); + return t; +} + +// Get aliased tensor from output, or make a new one for the corresponding +// output operation. Transfers ownership of the buffer from output to the +// returned tensor. +static absl::StatusOr GetOrCreateTensorForOutput( + xla::ScopedShapedBuffer& output, int output_num, OpKernelContext* ctx, + int missing_ctx_input_prefix, + const xla::HloInputOutputAliasConfig& input_output_alias, + absl::Span input_mapping, + const absl::flat_hash_map& resource_vars_snapshots, + DataType output_dtype, const TensorShape& output_shape, + Allocator* output_allocator, bool allocate_xla_tensors, se::Stream* stream, + bool use_multiple_streams, std::shared_ptr definition_event) { + xla::ShapeIndex output_index = input_output_alias.shape().IsTuple() + ? xla::ShapeIndex({output_num}) + : xla::ShapeIndex({}); + CHECK(input_output_alias.shape().IsTuple() || output_num == 0); + if (std::optional alias = + input_output_alias.GetAliasedParameter(output_index)) { + VLOG(3) << "Found alias: " << alias->ToString(); + int tf_param = + input_mapping[alias->parameter_number] - missing_ctx_input_prefix; + const Tensor input_tensor = + ctx->input(tf_param).dtype() != DT_RESOURCE + ? ctx->input(tf_param) + : *resource_vars_snapshots.at(missing_ctx_input_prefix + tf_param); + stream_executor::DeviceAddressBase input_buffer = + XlaTensor::DeviceMemoryFromTensor(input_tensor); + stream_executor::DeviceAddressBase output_buffer = + output.buffer({output_num}); + if (input_buffer.opaque() == output_buffer.opaque()) { + // In the case of a donated buffer, both input_tensor and output think + // they have ownership of the buffer (see comment in + // PopulateExecutionInputBuffer). Release ownership from output to avoid + // double free. + output.set_buffer(stream_executor::ScopedDeviceAddress(), + {output_num}); + return input_tensor; + } + } + + if (allocate_xla_tensors) { + Tensor output_tensor; + TF_RETURN_IF_ERROR( + ctx->allocate_temp(output_dtype, output_shape, &output_tensor)); + if (output_tensor.TotalBytes() > 0) { + XlaTensor* xla_tensor = XlaTensor::FromTensor(&output_tensor); + TF_RET_CHECK(xla_tensor); + xla_tensor->set_shaped_buffer(output.TakeSubTree({output_num})); + if (use_multiple_streams) { + xla_tensor->ResetDefinitionEvent(definition_event, stream); + } + } + return output_tensor; + } + + stream_executor::DeviceAddressBase output_buffer = + output.buffer({output_num}); + Tensor output_tensor = + MakeTensor(output_dtype, output_shape, output_buffer, output_allocator); + output.set_buffer(stream_executor::ScopedDeviceAddress(), + {output_num}); + return output_tensor; +} + // Sets output `output_num` for `ctx` provided it is known at a compile time. absl::Status SetOutputForConstant( OpKernelContext* ctx, bool requires_copy_to_device, @@ -188,6 +386,157 @@ absl::StatusOr> GatherVariableInfo( return std::move(out); } +absl::Status XlaComputationLaunchContext::PopulateOutputs( + OpKernelContext* ctx, + const XlaCompiler::CompilationResult* compilation_result, + ScopedShapedBuffer output, int missing_ctx_input_prefix, + absl::Span variable_infos, + const xla::HloInputOutputAliasConfig& input_output_alias, + const absl::flat_hash_map& resource_vars) { + se::Stream* stream = + ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr; + Allocator* allocator = ctx->device()->GetAllocator({}); + + // Computation output should always be a tuple. + VLOG(2) << "Result tuple shape: " << output.on_host_shape().ToString(); + VLOG(2) << "Result tuple shape (on device): " + << output.on_device_shape().ToString(); + CHECK_EQ(ctx->num_outputs(), compilation_result->outputs.size()); + + // If the on-host-shape isn't a tuple, create a new single-element tuple + // buffer with a nullptr root index table. This allows the code below to treat + // output as a tuple unconditionally. + if (!output.on_host_shape().IsTuple()) { + ShapedBuffer nontuple_buffer = output.release(); + ShapedBuffer buffer( + xla::ShapeUtil::MakeTupleShape({nontuple_buffer.on_device_shape()}), + output.device_ordinal()); + buffer.buffers().CopySubtreeFrom(nontuple_buffer.buffers(), + /*src_index=*/{}, + /*dst_index=*/{0}); + output = ScopedShapedBuffer(std::move(buffer), output.memory_allocator()); + } + + std::shared_ptr definition_event; + if (use_multiple_streams_ && stream) { + TF_ASSIGN_OR_RETURN(definition_event, stream->parent()->CreateEvent()); + TF_RETURN_IF_ERROR(stream->RecordEvent(definition_event.get())); + } + + for (const XlaOutputDescription& descr : compilation_result->outputs) { + if (descr.type == DT_VARIANT) { + return xla::Unimplemented( + "Support for TensorList crossing the XLA/TF boundary " + "is not implemented"); + } + } + + std::vector output_tensor_shapes; + output_tensor_shapes.reserve(ctx->num_outputs()); + if (output.on_host_shape().is_dynamic()) { + const se::Platform* platform = nullptr; + if (stream != nullptr) { + platform = stream->parent()->GetPlatform(); + } else { + // Stream is not set for the host platform. + TF_ASSIGN_OR_RETURN(platform, + se::PlatformManager::PlatformWithId( + XlaPlatformInfoFromDevice(ctx->device()))); + } + TF_ASSIGN_OR_RETURN(auto transfer_manager, + xla::TransferManager::GetForPlatform(platform)); + + xla::Shape output_device_shape = output.on_device_shape(); + TF_RETURN_IF_ERROR(transfer_manager->ReadDynamicShapes( + stream, &output, &output_device_shape)); + + output.set_shapes(output_device_shape, output_device_shape); + for (int i = 0; i < ctx->num_outputs(); ++i) { + const xla::Shape& subshape = + xla::ShapeUtil::GetSubshape(output_device_shape, {i}); + TensorShape shape; + TF_RETURN_IF_ERROR(XLAShapeToTensorShape(subshape, &shape)); + output_tensor_shapes.push_back(shape); + } + } else { + for (int i = 0; i < ctx->num_outputs(); ++i) { + output_tensor_shapes.push_back(compilation_result->outputs[i].shape); + } + } + + // Copy XLA results to the OpOutputList. + int output_num = 0; + for (int i = 0, end = ctx->num_outputs(); i < end; ++i) { + const TensorShape& shape = output_tensor_shapes[i]; + const DataType& type = compilation_result->outputs[i].type; + VLOG(2) << "Populating output for retval " << i << " shape " + << shape.DebugString() << " type " << DataTypeString(type); + + if (compilation_result->outputs[i].is_constant) { + TF_RETURN_IF_ERROR(SetOutputForConstant( + ctx, /*requires_copy_to_device=*/stream != nullptr, + compilation_result, i)); + } else if (type == DT_RESOURCE) { + int input_index = + compilation_result->outputs[i].input_index - missing_ctx_input_prefix; + TF_RET_CHECK(input_index >= 0 && input_index < ctx->num_inputs()) + << "Invalid input for outputs " << i << ": " << input_index; + ctx->set_output(i, ctx->input(input_index)); + } else { + TF_ASSIGN_OR_RETURN( + Tensor output_tensor, + GetOrCreateTensorForOutput( + output, output_num, ctx, missing_ctx_input_prefix, + input_output_alias, compilation_result->input_mapping, + resource_vars, ctx->expected_output_dtype(i), shape, allocator, + allocate_xla_tensors_, stream, use_multiple_streams_, + definition_event)); + ctx->set_output(i, output_tensor); + ++output_num; + } + } + + // input_index -> index into variable_infos. + absl::flat_hash_map variable_info_lookup; + for (int i = 0; i < variable_infos.size(); i++) { + variable_info_lookup.emplace(variable_infos[i].index(), i); + } + + // Apply variable updates, if any. + for (int i = 0, end = compilation_result->resource_updates.size(); i < end; + ++i) { + const XlaCompiler::ResourceUpdate& write = + compilation_result->resource_updates[i]; + int actual_input_index = write.input_index - missing_ctx_input_prefix; + CHECK_GE(actual_input_index, 0); + CHECK_LT(actual_input_index, ctx->num_inputs()); + Var* var = variable_infos[variable_info_lookup[actual_input_index]].var(); + CHECK(var); + + VLOG(2) << "Updating variable #" << i + << " at input index: " << actual_input_index << " with shape " + << write.shape.DebugString() << "; variable tensor has shape: " + << var->tensor()->shape().DebugString(); + + if (var->is_initialized && var->tensor()->dtype() != write.type) { + return absl::InternalError("Mismatched type in variable write"); + } + + TF_ASSIGN_OR_RETURN( + Tensor output_tensor, + GetOrCreateTensorForOutput(output, output_num, ctx, + missing_ctx_input_prefix, input_output_alias, + compilation_result->input_mapping, + resource_vars, write.type, write.shape, + allocator, allocate_xla_tensors_, stream, + use_multiple_streams_, definition_event)); + var->is_initialized |= write.modified; + *var->tensor() = output_tensor; + ++output_num; + } + return absl::OkStatus(); +} + absl::StatusOr> XlaComputationLaunchContext::BuildXlaCompilerArguments( absl::Span must_be_constant_idxs, @@ -466,8 +815,7 @@ absl::Status PopulateCtxOutputsFromPjRtExecutableOutputs( xla::ExecuteOptions GetPjRtExecuteOptions( const DeviceType& device_type, - absl::flat_hash_set non_donatable_input_indices, - xla::ExecuteContext* execute_context) { + absl::flat_hash_set non_donatable_input_indices) { xla::ExecuteOptions options; // Hardcode run id to always be one: TF distributed strategy // differentiates between subsequent runs using dependency edges. This @@ -486,7 +834,6 @@ xla::ExecuteOptions GetPjRtExecuteOptions( // tensorflow/compiler/xla/pjrt/pjrt_c_api_client.cc;l=923-927;rcl=519286815 options.use_major_to_minor_data_layout_for_callbacks = true; options.non_donatable_input_indices = std::move(non_donatable_input_indices); - options.context = execute_context; return options; } @@ -520,17 +867,14 @@ absl::Status RunPjRtExecutable( OpKernelContext* ctx) { const CompositeDevice::AcceleratorDeviceInfo* accelerator_device_info = ctx->device()->tensorflow_accelerator_device_info(); - const DeviceType& device_type = GetDeviceType(ctx); const bool use_pjrt_tensor_buffer = - device_type == DEVICE_GPU || accelerator_device_info == nullptr || - accelerator_device_info->use_pjrt_tensor_buffer; - int pjrt_device_id = 0; - if (device_type == DEVICE_CPU || device_type == DEVICE_XLA_CPU) { - pjrt_device_id = 0; - } else { - pjrt_device_id = - tsl::GetDeviceIdFromDeviceParsedName(ctx->device()->parsed_name()); - } + (accelerator_device_info == nullptr) + ? true + : accelerator_device_info->use_pjrt_tensor_buffer; + + const DeviceType& device_type = GetDeviceType(ctx); + const int pjrt_device_id = + tsl::GetDeviceIdFromDeviceParsedName(ctx->device()->parsed_name()); TF_ASSIGN_OR_RETURN( xla::PjRtDevice * device, pjrt_client->LookupAddressableDevice(xla::LocalDeviceId(pjrt_device_id))); @@ -554,17 +898,12 @@ absl::Status RunPjRtExecutable( executable->FingerprintExecutable()); device_selector_resource->selector()->Enqueue(pjrt_device_id, fingerprint); } - const Eigen::ThreadPoolDevice* intra_op_thread_pool = nullptr; - if (ctx != nullptr && ctx->device() != nullptr && - ctx->device()->has_eigen_cpu_device()) { - intra_op_thread_pool = ctx->device()->eigen_cpu_device(); - } TF_ASSIGN_OR_RETURN( std::vector> execute_outputs, RunPjRtExecutable(num_missing_prefix_ctx_inputs, inputs, variable_snapshots, updated_variables, device_type, use_pjrt_tensor_buffer, compilation_result, device, - pjrt_client, executable, intra_op_thread_pool)); + pjrt_client, executable)); if (device_selector_resource != nullptr) { device_selector_resource->selector()->Completed(pjrt_device_id, /*had_error=*/false); @@ -583,8 +922,7 @@ absl::StatusOr>> RunPjRtExecutable( const DeviceType& device_type, bool use_pjrt_tensor_buffer, const XlaCompiler::CompilationResult& compilation_result, xla::PjRtDevice* device, xla::PjRtClient* pjrt_client, - xla::PjRtLoadedExecutable* executable, - const Eigen::ThreadPoolDevice* intra_op_thread_pool) { + xla::PjRtLoadedExecutable* executable) { std::vector executable_args; executable_args.reserve(compilation_result.input_mapping.size()); std::vector> owned_executable_args; @@ -597,37 +935,32 @@ absl::StatusOr>> RunPjRtExecutable( std::vector> execute_outputs; std::optional> future; - - std::unique_ptr cpu_execute_context; - if (pjrt_client->platform_id() == xla::CpuId() && - intra_op_thread_pool != nullptr) { - cpu_execute_context = std::make_unique(); - cpu_execute_context->set_intra_op_thread_pool(intra_op_thread_pool); - } - xla::ExecuteOptions execute_options = - GetPjRtExecuteOptions(device_type, std::move(non_donatable_input_indices), - cpu_execute_context.get()); - if (executable->num_replicas() != 1 || executable->num_partitions() != 1) { - TF_ASSIGN_OR_RETURN(execute_outputs, - executable->ExecuteSharded(executable_args, device, - execute_options, future)); + TF_ASSIGN_OR_RETURN( + execute_outputs, + executable->ExecuteSharded( + executable_args, device, + GetPjRtExecuteOptions(device_type, + std::move(non_donatable_input_indices)), + future)); } else { - TF_ASSIGN_OR_RETURN(execute_outputs, - executable->ExecutePortable(executable_args, device, - execute_options, future)); + TF_ASSIGN_OR_RETURN( + execute_outputs, + executable->ExecutePortable( + executable_args, device, + GetPjRtExecuteOptions(device_type, + std::move(non_donatable_input_indices)), + future)); } - // We need to ensure the PjRtBuffers owned by `owned_executable_args` and - // the `cpu_execute_context` live until execution is complete. We do this by - // capturing them by move, so they are owned by the lambda that is executed - // when the future returned by ExecutePortable/ExecuteSharded is ready - // i.e. when the execution is complete. - if ((!owned_executable_args.empty() || cpu_execute_context != nullptr) && - future.has_value()) { - future->OnReady([owned_executable_args = std::move(owned_executable_args), - cpu_execute_context = - std::move(cpu_execute_context)](absl::Status s) {}); + // We need to ensure the PjRtBuffers owned by `owned_executable_args` live + // until execution is complete. We do this by capturing + // `owned_executable_args` by move, so it is owned by the lambda + // that is executed when the future returned by ExecutePortable/ExecuteSharded + // is ready i.e. when the execution is complete. + if (!owned_executable_args.empty() && future.has_value()) { + future->OnReady([owned_executable_args = + std::move(owned_executable_args)](absl::Status s) {}); } return execute_outputs; diff --git a/tensorflow/compiler/jit/xla_launch_util.h b/tensorflow/compiler/jit/xla_launch_util.h index 14ab1968e0b947..57ed70bb565fa8 100644 --- a/tensorflow/compiler/jit/xla_launch_util.h +++ b/tensorflow/compiler/jit/xla_launch_util.h @@ -18,35 +18,24 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_JIT_XLA_LAUNCH_UTIL_H_ #define TENSORFLOW_COMPILER_JIT_XLA_LAUNCH_UTIL_H_ -#include -#include +#include #include +#include #include #include "absl/container/flat_hash_map.h" -#include "absl/container/flat_hash_set.h" -#include "absl/status/status.h" -#include "absl/status/statusor.h" -#include "absl/types/span.h" #include "tensorflow/compiler/jit/variable_info.h" +#include "tensorflow/compiler/jit/xla_tensor.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "xla/client/local_client.h" -#include "xla/hlo/ir/hlo_input_output_alias_config.h" #include "xla/pjrt/pjrt_client.h" -#include "xla/pjrt/pjrt_executable.h" -#include "xla/service/executable.h" +#include "xla/service/maybe_owning_device_memory.h" #include "xla/service/shaped_buffer.h" -#include "xla/stream_executor/device_address_allocator.h" +#include "xla/stream_executor/device_memory_allocator.h" #include "tensorflow/core/framework/allocation_description.pb.h" -#include "tensorflow/core/framework/allocator.h" -#include "tensorflow/core/framework/device.h" -#include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/framework/types.h" - -namespace Eigen { -struct ThreadPoolDevice; -} // namespace Eigen +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/thread_annotations.h" namespace tensorflow { @@ -119,8 +108,7 @@ absl::Status PopulateCtxOutputsFromPjRtExecutableOutputs( // Returns the options used for executing a PjRtLoadedExecutable. xla::ExecuteOptions GetPjRtExecuteOptions( const DeviceType& device_type, - absl::flat_hash_set non_donatable_input_indices, - xla::ExecuteContext* execute_context = nullptr); + absl::flat_hash_set non_donatable_input_indices); // Returns the device ordinal from the parsed name of the device. int GetDeviceOrdinal(const DeviceBase* device); @@ -164,13 +152,25 @@ absl::StatusOr>> RunPjRtExecutable( const DeviceType& device_type, bool use_pjrt_tensor_buffer, const XlaCompiler::CompilationResult& compilation_result, xla::PjRtDevice* device, xla::PjRtClient* pjrt_client, - xla::PjRtLoadedExecutable* executable, - const Eigen::ThreadPoolDevice* intra_op_thread_pool = nullptr); + xla::PjRtLoadedExecutable* executable); // Helper class to perform the marshalling of TensorFlow inputs and outputs to // ShapedBuffers suitable for passing to an XLA computation. class XlaComputationLaunchContext { public: + // Create a new launch context. 'allocate_xla_tensors' is true if allocated + // output tensors and variables are always XlaTensors. If false they are + // assumed to be "normal" device pointers. + // If 'use_multiple_streams' is true, tensors may be defined and used on + // multiple streams and so se::Events must be defined and waited for. If + // 'use_multiple_streams' is true, 'allocate_xla_tensors' must also be true + // because we track inter-stream dependencies through events inside XlaTensor + // objects. + XlaComputationLaunchContext( + xla::LocalClient* client, + stream_executor::DeviceAddressAllocator* xla_allocator, + int device_ordinal, bool allocate_xla_tensors, bool use_multiple_streams); + // Builds a XlaCompiler::Argument vector from the arguments to an XlaLaunch // op. // Precondition: variables in `variable_args` are locked. @@ -179,6 +179,45 @@ class XlaComputationLaunchContext { absl::Span inputs, absl::Span variable_args, Device* device); + + // Add all inputs within `ctx` as XLA arguments (returned by arguments()). + // `variables` is a map from TensorFlow argument number to resource variable. + // + // Assumes that the first `missing_ctx_input_prefix` inputs to the kernel are + // missing and adjusts input indices accordingly. All elements in kernel's + // input_mapping must be greater than or equal to `missing_ctx_input_prefix` + // (in other words, no inputs actually required by the kernel can be missing). + absl::StatusOr> PopulateInputs( + OpKernelContext* ctx, + const XlaCompiler::CompilationResult* compilation_result, + const absl::flat_hash_map& resource_vars, + int missing_ctx_input_prefix, + const xla::HloInputOutputAliasConfig& input_output_alias); + + // Given the XLA output in `output`, populate all outputs of `ctx`. Also + // writes out the resource variable updates. + // + // Updates to all resource variables are written in a single atomic operation. + // This models *->Write dependencies between resource variable operations. + // See jit/resource_operation_safety_analysis for details. + // + // + // Assumes that the first `missing_ctx_input_prefix` inputs to the + // compilation_result are missing and adjusts input indices accordingly. + absl::Status PopulateOutputs( + OpKernelContext* ctx, + const XlaCompiler::CompilationResult* compilation_result, + xla::ScopedShapedBuffer output, int missing_ctx_input_prefix, + absl::Span variable_infos, + const xla::HloInputOutputAliasConfig& input_output_alias, + const absl::flat_hash_map& resource_vars); + + private: + xla::LocalClient* client_; + stream_executor::DeviceAddressAllocator* xla_allocator_; + bool allocate_xla_tensors_; + bool use_multiple_streams_; + int device_ordinal_; }; // A simple TensorBuffer implementation that allows us to create Tensors that diff --git a/tensorflow/compiler/jit/xla_launch_util_gpu_test.cc b/tensorflow/compiler/jit/xla_launch_util_gpu_test.cc index dd7c4f6090ba6a..e3f32f8403379a 100644 --- a/tensorflow/compiler/jit/xla_launch_util_gpu_test.cc +++ b/tensorflow/compiler/jit/xla_launch_util_gpu_test.cc @@ -130,6 +130,7 @@ class PjRtExecutionUtilGpuTest : public OpsTestBase { delete tensor; } tensors_.clear(); + device_context_->Unref(); core::ScopedUnref device_compiler_ref(device_compiler_); core::ScopedUnref profiler_ref(profiler_); } @@ -223,8 +224,8 @@ TEST_F(PjRtExecutionUtilGpuTest, PreparePjRtExecutableArguments) { std::vector input_mapping{3, 4}; std::vector variables; - auto pjrt_device = pjrt_client_->LookupAddressableDevice( - xla::LocalDeviceId(device_->parsed_name().id)); + auto pjrt_device = + pjrt_client_->LookupAddressableDevice(device_->parsed_name().id); TF_EXPECT_OK(pjrt_device.status()); std::vector exec_args; diff --git a/tensorflow/compiler/jit/xla_launch_util_test.cc b/tensorflow/compiler/jit/xla_launch_util_test.cc index d42a0a5f362019..22ce12448ce347 100644 --- a/tensorflow/compiler/jit/xla_launch_util_test.cc +++ b/tensorflow/compiler/jit/xla_launch_util_test.cc @@ -96,6 +96,8 @@ class PjRtExecutionUtilTest : public OpsTestBase { // Add and setup the XLA_CPU device. auto device_type = DeviceType(DEVICE_XLA_CPU); + rollout_config.AllowForDeviceInXlaLaunch(device_type); + rollout_config.AllowForDeviceInXlaCompileOnDemand(device_type); auto jit_device_type = DeviceType(DEVICE_CPU_XLA_JIT); auto device = diff --git a/tensorflow/compiler/mlir/lite/python/slim_model_importer.cc b/tensorflow/compiler/mlir/lite/python/slim_model_importer.cc index 99fa96d4dc99f7..9e55cafa495a13 100644 --- a/tensorflow/compiler/mlir/lite/python/slim_model_importer.cc +++ b/tensorflow/compiler/mlir/lite/python/slim_model_importer.cc @@ -374,8 +374,19 @@ absl::StatusOr> LoadSlimModel( {combined_builder.getStringAttr(res_name)})); } - std::string inputs_str = absl::StrJoin(arg_names, ","); - std::string outputs_str = absl::StrJoin(res_names, ","); + std::vector entry_input_names; + entry_input_names.reserve(arg_names.size()); + for (const auto& arg_name : arg_names) { + entry_input_names.push_back(absl::StrCat(sig_name, "_", arg_name)); + } + std::vector entry_output_names; + entry_output_names.reserve(res_names.size()); + for (const auto& res_name : res_names) { + entry_output_names.push_back( + absl::StrCat(sig_name, "_", res_name, "_output")); + } + std::string inputs_str = absl::StrJoin(entry_input_names, ","); + std::string outputs_str = absl::StrJoin(entry_output_names, ","); llvm::SmallVector entry_function_attrs; entry_function_attrs.push_back(combined_builder.getNamedAttr( "inputs", combined_builder.getStringAttr(inputs_str))); diff --git a/tensorflow/compiler/mlir/lite/stablehlo/tests/tfl_legalize_hlo.mlir b/tensorflow/compiler/mlir/lite/stablehlo/tests/tfl_legalize_hlo.mlir index 92f991f2513686..0d2afcf1e9dfc0 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/tests/tfl_legalize_hlo.mlir +++ b/tensorflow/compiler/mlir/lite/stablehlo/tests/tfl_legalize_hlo.mlir @@ -2451,6 +2451,34 @@ func.func @iota_1d() -> tensor<123xf32> { // ----- +// tfl.range has no bf16 support, so the range is built in f32 and cast back. +// CHECK-LABEL: iota_1d_bf16 +func.func @iota_1d_bf16() -> tensor<123xbf16> { + %0 = "mhlo.iota"() <{ iota_dimension = 0 : i64 }> : () -> tensor<123xbf16> + func.return %0 : tensor<123xbf16> +} + +// CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<0.000000e+00> : tensor +// CHECK-DAG: %[[CST_2:.*]] = arith.constant dense<1.230000e+02> : tensor +// CHECK-DAG: %[[CST_3:.*]] = arith.constant dense<1.000000e+00> : tensor +// CHECK: %[[RANGE:.*]] = "tfl.range"(%[[CST_1]], %[[CST_2]], %[[CST_3]]) : (tensor, tensor, tensor) -> tensor<123xf32> +// CHECK: "tfl.cast"(%[[RANGE]]) : (tensor<123xf32>) -> tensor<123xbf16> + +// ----- + +// CHECK-LABEL: iota_2d_bf16 +func.func @iota_2d_bf16() -> tensor<4x123xbf16> { + %0 = "mhlo.iota"() <{ iota_dimension = 1 : i64 }> : () -> tensor<4x123xbf16> + func.return %0 : tensor<4x123xbf16> +} + +// CHECK: %[[RANGE:.*]] = "tfl.range"({{.*}}) : (tensor, tensor, tensor) -> tensor<123xf32> +// CHECK: %[[RESHAPE:.*]] = "tfl.reshape"(%[[RANGE]], {{.*}}) : (tensor<123xf32>, tensor<2xi32>) -> tensor<1x123xf32> +// CHECK: %[[BCAST:.*]] = "tfl.broadcast_to"(%[[RESHAPE]], {{.*}}) : (tensor<1x123xf32>, tensor<2xi64>) -> tensor<4x123xf32> +// CHECK: "tfl.cast"(%[[BCAST]]) : (tensor<4x123xf32>) -> tensor<4x123xbf16> + +// ----- + // CHECK-LABEL: iota_i64 func.func @iota_i64() -> tensor<123xi64> { %0 = "mhlo.iota"() <{ iota_dimension = 0 : i64 }> : () -> tensor<123xi64> @@ -2524,6 +2552,23 @@ func.func @dynamic_iota_f32_1d(%arg0: tensor<1xi32>) -> tensor { // ----- +// tfl.range has no bf16 support, so the range is built in f32 and cast back. +// CHECK-LABEL: dynamic_iota_bf16_1d +func.func @dynamic_iota_bf16_1d(%arg0: tensor<1xi32>) -> tensor { + %0 = "mhlo.dynamic_iota"(%arg0) <{iota_dimension = 0 : i64}> : (tensor<1xi32>) -> tensor + func.return %0 : tensor +} + +// CHECK-DAG: %[[CST:.*]] = arith.constant dense<0.000000e+00> : tensor +// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<1.000000e+00> : tensor +// CHECK: %0 = "tfl.cast"(%arg0) : (tensor<1xi32>) -> tensor<1xf32> +// CHECK: %[[CST_1:.*]] = arith.constant dense<> : tensor<0xi32> +// CHECK: %1 = "tfl.reshape"(%0, %[[CST_1]]) : (tensor<1xf32>, tensor<0xi32>) -> tensor +// CHECK: %2 = "tfl.range"(%[[CST]], %1, %[[CST_0]]) : (tensor, tensor, tensor) -> tensor +// CHECK: "tfl.cast"(%2) : (tensor) -> tensor + +// ----- + //===----------------------------------------------------------------------===// // mhlo.dynamic_slice //===----------------------------------------------------------------------===// @@ -3073,6 +3118,18 @@ func.func @log1p_dynamic(%arg0: tensor) -> tensor { // ----- +// CHECK-LABEL: log1p_bf16 +func.func @log1p_bf16(%arg0: tensor<2xbf16>) -> tensor<2xbf16> { + %0 = "mhlo.log_plus_one"(%arg0) : (tensor<2xbf16>) -> tensor<2xbf16> + func.return %0 : tensor<2xbf16> +} + +// CHECK: %cst = arith.constant dense<1.000000e+00> : tensor +// CHECK: %0 = tfl.add(%arg0, %cst) <{fused_activation_function = "NONE"}> : (tensor<2xbf16>, tensor) -> tensor<2xbf16> +// CHECK: %1 = "tfl.log"(%0) : (tensor<2xbf16>) -> tensor<2xbf16> + +// ----- + // CHECK-LABEL: neg func.func @neg(%arg0: tensor<2xf32>) -> tensor<2xf32> { %0 = "mhlo.negate"(%arg0) : (tensor<2xf32>) -> tensor<2xf32> diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/iota.cc b/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/iota.cc index b88f8deacade29..22ce8a8ede33bb 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/iota.cc +++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/iota.cc @@ -47,10 +47,16 @@ class LegalizeIota : public OpConversionPattern { ConversionPatternRewriter& rewriter) const final; }; +// tfl.range only supports i32, i64 and f32. Narrower floats are materialized +// in f32 and cast back, which is exact because an iota only holds small +// integral values. Leaving them unconverted emits a STABLEHLO_IOTA builtin +// that the TFLite runtime has no kernel for. +bool IsNarrowFloat(Type e_type) { return e_type.isBF16() || e_type.isF16(); } + bool IsIotaLegal(mhlo::IotaOp op) { auto e_type = llvm::cast(op.getType()).getElementType(); - return !(e_type.isF32() || e_type.isSignlessInteger(32) || - e_type.isSignlessInteger(64)); + return !(e_type.isF32() || IsNarrowFloat(e_type) || + e_type.isSignlessInteger(32) || e_type.isSignlessInteger(64)); } std::tuple @@ -73,52 +79,55 @@ LogicalResult LegalizeIota::matchAndRewrite( mhlo::IotaOp op, OpAdaptor adaptor, ConversionPatternRewriter& rewriter) const { if (IsIotaLegal(op)) { - return rewriter.notifyMatchFailure(op, "Must be i32, i64 or f32"); + return rewriter.notifyMatchFailure(op, "Must be i32, i64 or a float"); } auto type = llvm::cast(op.getType()); auto e_type = type.getElementType(); + const bool needs_cast = IsNarrowFloat(e_type); + const Type range_e_type = needs_cast ? rewriter.getF32Type() : e_type; const int64_t iota_dim_size = type.getDimSize(op.getIotaDimension()); auto [start, limit, delta] = - BuildRangeParams(e_type, iota_dim_size, rewriter); + BuildRangeParams(range_e_type, iota_dim_size, rewriter); auto start_op = arith::ConstantOp::create(rewriter, op->getLoc(), start); auto limit_op = arith::ConstantOp::create(rewriter, op->getLoc(), limit); auto delta_op = arith::ConstantOp::create(rewriter, op->getLoc(), delta); - auto range_type = RankedTensorType::get({iota_dim_size}, e_type); - auto range_op = TFL::RangeOp::create(rewriter, op->getLoc(), range_type, - start_op, limit_op, delta_op); - - if (type.getRank() == 1) { - rewriter.replaceOp(op, range_op); - return success(); + auto range_type = RankedTensorType::get({iota_dim_size}, range_e_type); + Value result = TFL::RangeOp::create(rewriter, op->getLoc(), range_type, + start_op, limit_op, delta_op); + + if (type.getRank() > 1) { + // mhlo.iota allows filling ND tensors iota-style. Reshape and broadcast + // tfl 1D range output. + llvm::SmallVector reshape_shape(type.getRank(), 1); + reshape_shape[op.getIotaDimension()] = iota_dim_size; + Value reshape_shape_cst = arith::ConstantOp::create( + rewriter, op->getLoc(), rewriter.getI64TensorAttr(reshape_shape)); + reshape_shape_cst = + TFL::CastOp::create(rewriter, op->getLoc(), + llvm::cast(reshape_shape_cst.getType()) + .clone(rewriter.getI32Type()), + reshape_shape_cst); + + auto reshape_type = RankedTensorType::get(reshape_shape, range_e_type); + auto reshape_op = TFL::ReshapeOp::create( + rewriter, op->getLoc(), reshape_type, result, reshape_shape_cst); + + auto broad_cast_shape_cst = arith::ConstantOp::create( + rewriter, op->getLoc(), rewriter.getI64TensorAttr(type.getShape())); + result = TFL::BroadcastToOp::create(rewriter, op->getLoc(), + type.clone(range_e_type), reshape_op, + broad_cast_shape_cst); } - // mhlo.iota allows filling ND tensors iota-style. Reshape and broadcast - // tfl 1D range output. - - llvm::SmallVector reshape_shape(type.getRank(), 1); - reshape_shape[op.getIotaDimension()] = iota_dim_size; - Value reshape_shape_cst = arith::ConstantOp::create( - rewriter, op->getLoc(), rewriter.getI64TensorAttr(reshape_shape)); - reshape_shape_cst = - TFL::CastOp::create(rewriter, op->getLoc(), - llvm::cast(reshape_shape_cst.getType()) - .clone(rewriter.getI32Type()), - reshape_shape_cst); - - auto reshape_type = RankedTensorType::get(reshape_shape, e_type); - auto reshape_op = TFL::ReshapeOp::create(rewriter, op->getLoc(), reshape_type, - range_op, reshape_shape_cst); - - auto broad_cast_shape_cst = arith::ConstantOp::create( - rewriter, op->getLoc(), rewriter.getI64TensorAttr(type.getShape())); - - rewriter.replaceOpWithNewOp(op, type, reshape_op, - broad_cast_shape_cst); + if (needs_cast) { + result = TFL::CastOp::create(rewriter, op->getLoc(), type, result); + } + rewriter.replaceOp(op, result); return success(); } @@ -138,7 +147,8 @@ class LegalizeDynamicIotaOp : public OpConversionPattern { bool IsDynamicIotaLegal(mhlo::DynamicIotaOp op) { auto type = llvm::cast(op.getType()); auto element_type = type.getElementType(); - return (!element_type.isF32() && !element_type.isSignlessInteger(32) && + return (!element_type.isF32() && !IsNarrowFloat(element_type) && + !element_type.isSignlessInteger(32) && !element_type.isSignlessInteger(64)) || type.getRank() > 1 || op.getIotaDimension() != 0; } @@ -152,17 +162,19 @@ LogicalResult LegalizeDynamicIotaOp::matchAndRewrite( auto type = llvm::cast(op.getType()); Type element_type = type.getElementType(); + const bool needs_cast = IsNarrowFloat(element_type); + const Type range_e_type = needs_cast ? rewriter.getF32Type() : element_type; auto [start, unused_limit, delta] = - BuildRangeParams(element_type, /*iota_dim_size*/ 0, rewriter); + BuildRangeParams(range_e_type, /*iota_dim_size*/ 0, rewriter); auto start_op = arith::ConstantOp::create(rewriter, op.getLoc(), start); auto delta_op = arith::ConstantOp::create(rewriter, op.getLoc(), delta); auto output_shape = op.getOperand(); - if (mlir::isa(element_type)) { + if (mlir::isa(range_e_type)) { auto cast_type = - mlir::cast(output_shape.getType()).clone(element_type); + mlir::cast(output_shape.getType()).clone(range_e_type); output_shape = TFL::CastOp::create(rewriter, op.getLoc(), cast_type, output_shape); } @@ -173,15 +185,21 @@ LogicalResult LegalizeDynamicIotaOp::matchAndRewrite( auto scalar_shape = arith::ConstantOp::create(rewriter, op.getLoc(), scalar_attr); auto limit_scalar = TFL::ReshapeOp::create( - rewriter, op.getLoc(), RankedTensorType::get({}, element_type), + rewriter, op.getLoc(), RankedTensorType::get({}, range_e_type), output_shape, scalar_shape); const uint64_t dimension = op.getIotaDimension(); auto range_type = - RankedTensorType::get({type.getShape()[dimension]}, element_type); + RankedTensorType::get({type.getShape()[dimension]}, range_e_type); + + Value result = TFL::RangeOp::create(rewriter, op.getLoc(), range_type, + start_op, limit_scalar, delta_op); + + if (needs_cast) { + result = TFL::CastOp::create(rewriter, op.getLoc(), type, result); + } - rewriter.replaceOpWithNewOp(op, range_type, start_op, - limit_scalar, delta_op); + rewriter.replaceOp(op, result); return success(); } diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo.cc b/tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo.cc index b283dea3098232..5ecfdc88638fce 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo.cc +++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo.cc @@ -31,6 +31,7 @@ limitations under the License. #include "mlir/IR/BuiltinTypes.h" // from @llvm-project #include "mlir/IR/MLIRContext.h" // from @llvm-project #include "mlir/IR/PatternMatch.h" // from @llvm-project +#include "mlir/IR/TypeUtilities.h" // from @llvm-project #include "mlir/Pass/PassRegistry.h" // from @llvm-project #include "mlir/Support/LLVM.h" // from @llvm-project #include "mlir/Support/LogicalResult.h" // from @llvm-project @@ -72,6 +73,19 @@ arith::ConstantOp ShapeToConst(PatternRewriter& rewriter, Value value) { return arith::ConstantOp::create(rewriter, value.getLoc(), attr_type, attr); } +// Returns a scalar (rank-0) constant holding `raw_value`, typed with the same +// floating point element type as `value`. Decompositions that introduce a +// literal operand must match the input precision; hardcoding an f32 constant +// restricts the pattern to f32 inputs and leaves e.g. bf16 operands illegal. +arith::ConstantOp ScalarConstantLike(OpBuilder& b, Value value, + double raw_value) { + Type element_type = getElementTypeOrSelf(value.getType()); + auto attr_type = RankedTensorType::get({}, element_type); + auto attr = DenseElementsAttr::get(attr_type, + b.getFloatAttr(element_type, raw_value)); + return arith::ConstantOp::create(b, value.getLoc(), attr_type, attr); +} + // Returns true if broadcast_dimensions obey Tensorflow convention, as in new // dimensions are added as prefix. bool IsTFLStyleBroadcast(DenseIntElementsAttr broadcast_dimensions, diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo_patterns.td b/tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo_patterns.td index db16a99353ba85..d1fe3a889b8451 100644 --- a/tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo_patterns.td +++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/tflite_legalize_hlo_patterns.td @@ -28,6 +28,12 @@ def ShapeToConst : NativeCodeCall<"ShapeToConst($_builder, $0)">; def CreateTFLCastToInt32Op : NativeCodeCall< "CreateCastToInt32($0, $_loc, $_builder)">; +// Builds a rank-0 constant with the same float element type as the given value. +def ScalarZeroLike : NativeCodeCall< + "ScalarConstantLike($_builder, $0, 0.0)">; +def ScalarOneLike : NativeCodeCall< + "ScalarConstantLike($_builder, $0, 1.0)">; + def ConstDefaultResultAccuracyAttr : ConstantAttr; @@ -381,27 +387,24 @@ def : Pat< (TFL_CastOp $input)>; def : Pat< - (MHLO_Expm1Op F32Tensor:$x, ConstDefaultResultAccuracyAttr), + (MHLO_Expm1Op TensorOf<[AnyFloat]>:$x, ConstDefaultResultAccuracyAttr), (TFL_SubOp (TFL_ExpOp $x), - (Arith_ConstantOp - ConstantAttr, "1.0f">), + (ScalarOneLike $x), TFL_AF_None)>; def : Pat< - (MHLO_IsFiniteOp F32Tensor:$x), + (MHLO_IsFiniteOp TensorOf<[AnyFloat]>:$x), (TFL_EqualOp (TFL_SubOp $x, $x, TFL_AF_None), - (Arith_ConstantOp - ConstantAttr, "0.0f">))>; + (ScalarZeroLike $x))>; def : Pat< - (MHLO_Log1pOp F32Tensor:$x, ConstDefaultResultAccuracyAttr), + (MHLO_Log1pOp TensorOf<[AnyFloat]>:$x, ConstDefaultResultAccuracyAttr), (TFL_LogOp (TFL_AddOp $x, - (Arith_ConstantOp - ConstantAttr, "1.0f">), + (ScalarOneLike $x), TFL_AF_None))>; //===----------------------------------------------------------------------===// diff --git a/tensorflow/compiler/mlir/lite/tests/cast_bf16_ops_to_f32.mlir b/tensorflow/compiler/mlir/lite/tests/cast_bf16_ops_to_f32.mlir index ac70df6c99d2cd..a1037a7b81722a 100644 --- a/tensorflow/compiler/mlir/lite/tests/cast_bf16_ops_to_f32.mlir +++ b/tensorflow/compiler/mlir/lite/tests/cast_bf16_ops_to_f32.mlir @@ -80,3 +80,24 @@ func.func @skip_call_and_region_ops(%arg0: tensor, %arg1: tensor<2x2xbf16>) func.func private @callee(%arg0: tensor<2x2xbf16>) -> tensor<2x2xbf16> { return %arg0 : tensor<2x2xbf16> } + +// ----- + +// CHECK-LABEL: func.func @skip_composite_and_impl +func.func @skip_composite_and_impl(%arg0: tensor<2x2xbf16>) -> tensor<2x2xbf16> { + // CHECK: %[[COMP:.*]] = stablehlo.composite "odml.test" %arg0 {decomposition = @test.impl} : (tensor<2x2xbf16>) -> tensor<2x2xbf16> + // CHECK: return %[[COMP]] : tensor<2x2xbf16> + %0 = "stablehlo.composite"(%arg0) <{name = "odml.test"}> {decomposition = @test.impl} : (tensor<2x2xbf16>) -> tensor<2x2xbf16> + return %0 : tensor<2x2xbf16> +} + +// CHECK-LABEL: func.func private @test.impl +func.func private @test.impl(%arg0: tensor<2x2xbf16>) -> tensor<2x2xbf16> { + // CHECK-DAG: %[[CAST0:.*]] = "tfl.cast"(%arg0) : (tensor<2x2xbf16>) -> tensor<2x2xf32> + // CHECK-DAG: %[[CAST1:.*]] = "tfl.cast"(%arg0) : (tensor<2x2xbf16>) -> tensor<2x2xf32> + // CHECK: %[[ADD:.*]] = tfl.add %[[CAST0]], %[[CAST1]] {fused_activation_function = "NONE"} : tensor<2x2xf32> + // CHECK: %[[CAST_OUT:.*]] = "tfl.cast"(%[[ADD]]) : (tensor<2x2xf32>) -> tensor<2x2xbf16> + // CHECK: return %[[CAST_OUT]] : tensor<2x2xbf16> + %0 = "tfl.add"(%arg0, %arg0) <{fused_activation_function = "NONE"}> : (tensor<2x2xbf16>, tensor<2x2xbf16>) -> tensor<2x2xbf16> + return %0 : tensor<2x2xbf16> +} diff --git a/tensorflow/compiler/mlir/lite/transforms/cast_bf16_ops_to_f32_pass.cc b/tensorflow/compiler/mlir/lite/transforms/cast_bf16_ops_to_f32_pass.cc index 25424be72adbfd..77b35466acc608 100644 --- a/tensorflow/compiler/mlir/lite/transforms/cast_bf16_ops_to_f32_pass.cc +++ b/tensorflow/compiler/mlir/lite/transforms/cast_bf16_ops_to_f32_pass.cc @@ -68,11 +68,17 @@ class CastBf16OpsToF32 : public RewritePattern { // Skip cast ops, constants, zero-operand ops, terminators, return-like ops, // region-bearing ops (e.g. tfl.while, tfl.if), and call ops to prevent type // mismatches with nested block arguments or function signatures. + // + // stablehlo.composite is also skipped: its operand/result types must stay + // in sync with its decomposition function's signature, which this pass does + // not rewrite. It is matched by name because this pass intentionally does + // not depend on the StableHLO dialect. if (isa(op) || op.hasTrait() || op.getName().hasTrait() || op.hasTrait() || op.hasTrait() || op.getNumRegions() > 0 || - isa(op)) { + isa(op) || + op.getName().getStringRef() == "stablehlo.composite") { return failure(); } for (Value input : op.getOperands()) { diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/BUILD b/tensorflow/compiler/mlir/quantization/stablehlo/BUILD index 5f0f6ed959ed2d..420737dc98007b 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/BUILD +++ b/tensorflow/compiler/mlir/quantization/stablehlo/BUILD @@ -404,7 +404,9 @@ tf_cc_test( "@llvm-project//llvm:Support", "@llvm-project//mlir:Pass", "@xla//xla:shape_util", - "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/client:client_library", + "@xla//xla/stream_executor:platform", + "@xla//xla/stream_executor:platform_manager", ], ) diff --git a/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/legalize_tf_quant_test.cc b/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/legalize_tf_quant_test.cc index 74ed3547fce3dd..2e9d7a8794bb6f 100644 --- a/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/legalize_tf_quant_test.cc +++ b/tensorflow/compiler/mlir/quantization/stablehlo/passes/bridge/legalize_tf_quant_test.cc @@ -21,8 +21,10 @@ limitations under the License. #include "llvm/ADT/StringRef.h" #include "mlir/Pass/Pass.h" // from @llvm-project #include "tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h" -#include "xla/pjrt/pjrt_compiler.h" +#include "xla/client/client_library.h" #include "xla/shape.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/platform_manager.h" #include "xla/tsl/lib/core/status_test_util.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.pb.h" @@ -44,8 +46,10 @@ class LegalizeTFQuantTest : public Test { mlir_to_hlo_args.rollout_state = tensorflow::ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED; mlir_to_hlo_args.mlir_module = mlir_module_string; - auto compiler = xla::GetDefaultPjRtCompiler(xla::CpuName()); - TF_ASSERT_OK(compiler.status()); + tensorflow::se::Platform* platform = + tensorflow::se::PlatformManager::PlatformWithName("Host").value(); + auto client = + xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform).value(); tensorflow::tpu::TPUCompileMetadataProto metadata_proto; // Set up an arg per arg_shape with the specified type. for (int i = 0; i < arg_shapes.size(); ++i) { @@ -65,7 +69,7 @@ class LegalizeTFQuantTest : public Test { mlir_to_hlo_args, metadata_proto, use_tuple_args, /*device_type=*/"XLA_TPU_JIT", custom_legalization_passes, /*shape_determination_fns=*/{}, arg_shapes, - &arg_core_mapping, &per_core_arg_shapes, *compiler) + &arg_core_mapping, &per_core_arg_shapes, client) .status()); } }; diff --git a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td index 53b78049f72704..f3c809637ce577 100644 --- a/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td +++ b/tensorflow/compiler/mlir/tensorflow/ir/tf_generated_ops.td @@ -9626,6 +9626,12 @@ a tensor shape `[..., M, K]`. If `adjoint` is `False` then each output matrix satisfies `matrix[..., :, :] * output[..., :, :] = rhs[..., :, :]`. If `adjoint` is `True` then each output matrix satisfies `adjoint(matrix[..., :, :]) * output[..., :, :] = rhs[..., :, :]`. + +If a matrix is not invertible there is no guarantee what the op does. It +may detect the condition and raise an exception or it may simply return a +garbage result. In particular, under XLA compilation (e.g. +`tf.function(jit_compile=True)`) the condition is not detected and the +result is implementation-defined. }]; let arguments = (ins diff --git a/tensorflow/compiler/mlir/tf2xla/api/v1/BUILD b/tensorflow/compiler/mlir/tf2xla/api/v1/BUILD index b67e1381abfab4..996123eb6308dc 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v1/BUILD +++ b/tensorflow/compiler/mlir/tf2xla/api/v1/BUILD @@ -170,9 +170,9 @@ cc_library( "@llvm-project//mlir:Support", "@xla//xla:shape_util", "@xla//xla:status_macros", + "@xla//xla/client:compile_only_client", "@xla//xla/hlo/ir:hlo", "@xla//xla/mlir_hlo:hlo_dialect_registration", - "@xla//xla/pjrt:pjrt_compiler", "@xla//xla/pjrt/proto:compile_options_proto_cc", "@xla//xla/service:hlo_proto_cc", ], @@ -213,9 +213,11 @@ tf_cc_test( "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", "@xla//xla:shape_util", + "@xla//xla/client:client_library", "@xla//xla/hlo/translate/mhlo_to_hlo:type_to_shape", "@xla//xla/mlir_hlo:hlo_dialect_registration", - "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/stream_executor:platform", + "@xla//xla/stream_executor:platform_manager", "@xla//xla/tsl/lib/core:status_test_util", "@xla//xla/tsl/lib/monitoring:test_utils", "@xla//xla/tsl/platform:errors", diff --git a/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.cc b/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.cc index c751b4e6a4a02f..82c13f5fd71cd5 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h" +#include #include #include #include @@ -48,9 +49,9 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h" #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" +#include "xla/client/compile_only_client.h" #include "xla/hlo/ir/hlo_input_output_alias_config.h" #include "xla/mlir_hlo/mhlo/IR/register.h" -#include "xla/pjrt/pjrt_compiler.h" #include "xla/service/hlo.pb.h" #include "xla/shape.h" #include "xla/shape_util.h" @@ -217,13 +218,13 @@ absl::Status CompileTFFunctionWithoutMlir( const DeviceType& device_type, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, - xla::PjRtCompiler* compiler, + xla::CompileOnlyClient* client, XlaCompiler::CompilationResult* compilation_result) { absl::Status comp_status = CompileTFFunctionToHlo( *function_computation.flib_def, function_computation.graph_def_version, shape_determination_funcs, arg_shapes, device_type, function_computation.guaranteed_constants, *function_computation.function, - metadata, compiler, arg_core_mapping, per_core_arg_shapes, use_tuple_args, + metadata, client, arg_core_mapping, per_core_arg_shapes, use_tuple_args, compilation_result); if (comp_status.ok()) { phase2_bridge_compilation_status->GetCell(kOldBridgeNoMlirSuccess) @@ -245,7 +246,7 @@ absl::Status CompileMLIRTFFunction( const DeviceType& device_type, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, - xla::PjRtCompiler* compiler, + xla::CompileOnlyClient* client, XlaCompiler::CompilationResult* compilation_result) { if (!mlir::SetTPUInfeedLayout(mlir_module)) return absl::InternalError("Failed to set layouts attribute"); @@ -283,7 +284,7 @@ absl::Status CompileMLIRTFFunction( TF_RETURN_IF_ERROR(CompileTFFunctionToHlo( *flib_def, versions.producer(), shape_determination_funcs, arg_shapes, - device_type, consts, func, metadata, compiler, arg_core_mapping, + device_type, consts, func, metadata, client, arg_core_mapping, per_core_arg_shapes, use_tuple_args, compilation_result)); if (compilation_result != nullptr && @@ -307,13 +308,13 @@ absl::Status CompileMLIRTFFunction( const DeviceType& device_type, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, - xla::PjRtCompiler* compiler, + xla::CompileOnlyClient* client, XlaCompiler::CompilationResult* compilation_result) { if (mlir_computation.mlir_module_op.has_value()) { return CompileMLIRTFFunction( mlir_computation.mlir_module_op.value(), metadata, use_tuple_args, shape_determination_funcs, arg_shapes, device_type, arg_core_mapping, - per_core_arg_shapes, compiler, compilation_result); + per_core_arg_shapes, client, compilation_result); } mlir::DialectRegistry registry; @@ -325,10 +326,10 @@ absl::Status CompileMLIRTFFunction( TF_RETURN_IF_ERROR(DeserializeMlirModule(mlir_computation.mlir_module, &context, &mlir_module)); - return CompileMLIRTFFunction( - *mlir_module, metadata, use_tuple_args, shape_determination_funcs, - arg_shapes, device_type, arg_core_mapping, per_core_arg_shapes, compiler, - compilation_result); + return CompileMLIRTFFunction(*mlir_module, metadata, use_tuple_args, + shape_determination_funcs, arg_shapes, + device_type, arg_core_mapping, + per_core_arg_shapes, client, compilation_result); } } // namespace @@ -342,7 +343,7 @@ absl::Status CompileTensorflowGraphToHlo( tsl::DeviceType device_type, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, - xla::PjRtCompiler* compiler, + xla::CompileOnlyClient* client, XlaCompiler::CompilationResult* compilation_result) { LOG_FIRST_N(INFO, 1) << "Compiling MLIR computation to XLA HLO using the " "old (non-MLIR) tf2xla bridge"; @@ -359,14 +360,14 @@ absl::Status CompileTensorflowGraphToHlo( TF_RETURN_IF_ERROR(CompileMLIRTFFunction( std::get<0>(computation), metadata, use_tuple_args, shape_determination_funcs, arg_shapes, device_type, arg_core_mapping, - per_core_arg_shapes, compiler, compilation_result)); + per_core_arg_shapes, client, compilation_result)); } else { FunctionToHloArgs function_computation = std::get<1>(computation); TF_RETURN_IF_ERROR(CompileTFFunctionWithoutMlir( function_computation, metadata, use_tuple_args, shape_determination_funcs, arg_shapes, device_type, arg_core_mapping, - per_core_arg_shapes, compiler, compilation_result)); + per_core_arg_shapes, client, compilation_result)); } phase2_bridge_compilation_time->GetCell(kBridgePhase2Config) diff --git a/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h b/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h index d70cbad1fdd4c8..44d4bd83e0e31c 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h +++ b/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph.h @@ -23,7 +23,7 @@ limitations under the License. #include "absl/types/variant.h" #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "xla/pjrt/pjrt_compiler.h" +#include "xla/client/compile_only_client.h" #include "xla/pjrt/proto/compile_options.pb.h" #include "xla/shape.h" #include "tensorflow/core/framework/tensor_shape.h" @@ -46,7 +46,7 @@ absl::Status CompileTensorflowGraphToHlo( tsl::DeviceType device_type, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, - xla::PjRtCompiler* compiler, + xla::CompileOnlyClient* client, XlaCompiler::CompilationResult* compilation_result); } // namespace v1 diff --git a/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph_test.cc b/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph_test.cc index aea90d2b38fd4d..7098fb36a7073e 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph_test.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v1/compile_tf_graph_test.cc @@ -35,10 +35,12 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/internal/utils/test_metadata_config.h" #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "xla/client/client_library.h" #include "xla/hlo/translate/mhlo_to_hlo/type_to_shape.h" #include "xla/mlir_hlo/mhlo/IR/register.h" -#include "xla/pjrt/pjrt_compiler.h" #include "xla/shape.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/platform_manager.h" #include "xla/tsl/framework/device_type.h" #include "xla/tsl/lib/core/status_test_util.h" #include "xla/tsl/lib/monitoring/test_utils.h" @@ -73,6 +75,7 @@ static constexpr char kCompilationTimeStreamzName[] = static constexpr char kCompilationStatusStreamzName[] = "/tensorflow/core/tf2xla/api/v1/phase2_compilation_status"; +static constexpr char kPlatformName[] = "Host"; constexpr char kEntryFuncName[] = "main"; static constexpr char kMlirModuleStr[] = R"( @@ -96,8 +99,10 @@ absl::StatusOr CompileWithComputation( bool use_serialized_mlir) { XlaCompilationResult compilation_result; - TF_ASSIGN_OR_RETURN(auto* compiler, - xla::GetDefaultPjRtCompiler(xla::CpuName())); + se::Platform* platform = + se::PlatformManager::PlatformWithName(kPlatformName).value(); + auto client = + xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform).value(); bool use_tuple_args = true; std::vector arg_core_mapping; @@ -130,7 +135,7 @@ absl::StatusOr CompileWithComputation( tensorflow::tf2xla::v1::CompileTensorflowGraphToHlo( computation, metadata_proto, use_tuple_args, shape_determination_fns, arg_shapes, tsl::DeviceType("XLA_TPU_JIT"), &arg_core_mapping, - &per_core_arg_shapes, compiler, &compilation_result); + &per_core_arg_shapes, client, &compilation_result); if (!compilation_status.ok()) return compilation_status; diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD b/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD index 5631953b4b6737..343ab53fbe32d9 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/BUILD @@ -56,12 +56,13 @@ cc_library( "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:variant", "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", "@llvm-project//mlir:Pass", "@tsl//tsl/platform:protobuf", "@xla//xla:shape_util", "@xla//xla:xla_proto_cc", + "@xla//xla/client:compile_only_client", "@xla//xla/hlo/ir:hlo", - "@xla//xla/pjrt:pjrt_compiler", "@xla//xla/pjrt/proto:compile_options_proto_cc", "@xla//xla/tsl/platform:statusor", ], @@ -90,7 +91,9 @@ tf_cc_test( "@com_google_googletest//:gtest", "@llvm-project//mlir:Pass", "@xla//xla:shape_util", - "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/client:client_library", + "@xla//xla/stream_executor:platform", + "@xla//xla/stream_executor:platform_manager", "@xla//xla/tsl/lib/monitoring:test_utils", "@xla//xla/tsl/platform:statusor", ], diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.cc index 2f05202cc3355e..c391db870fac9f 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.cc @@ -35,8 +35,8 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/internal/reproducer.pb.h" #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "xla/client/compile_only_client.h" #include "xla/hlo/ir/hlo_module.h" -#include "xla/pjrt/pjrt_compiler.h" #include "xla/shape.h" #include "xla/tsl/framework/device_type.h" #include "xla/tsl/lib/monitoring/sampler.h" @@ -158,7 +158,7 @@ absl::StatusOr LegalizeMlirToHlo( const std::vector& arg_shapes, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, - xla::PjRtCompiler* compiler) { + xla::CompileOnlyClient* client) { CompilationTimer timer; auto record_time = llvm::scope_exit([&timer] { phase2_bridge_compilation_time->GetCell(kFullBridge) @@ -174,7 +174,7 @@ absl::StatusOr LegalizeMlirToHlo( TF_RETURN_IF_ERROR(tf2xla::v1::CompileTensorflowGraphToHlo( computation, metadata, use_tuple_args, shape_determination_fns, arg_shapes, tsl::DeviceType(device_type.str()), arg_core_mapping, - per_core_arg_shapes, compiler, compilation_result.get())); + per_core_arg_shapes, client, compilation_result.get())); DumpHloCompilationResult("legalize_tf_fallback.hlo", compilation_result.get()) @@ -185,7 +185,7 @@ absl::StatusOr LegalizeMlirToHlo( auto combined_bridge_status = internal::LegalizeTfToHlo( std::get<0>(computation), metadata, use_tuple_args, device_type, shape_determination_fns, arg_shapes, arg_core_mapping, - per_core_arg_shapes, custom_legalization_passes, compiler, + per_core_arg_shapes, custom_legalization_passes, client, compilation_result.get()); if (combined_bridge_status.ok()) { diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h index fc8d4cb8bfae9f..4a33f9760326ca 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h @@ -21,14 +21,14 @@ limitations under the License. #include #include "absl/status/status.h" -#include "llvm/ADT/StringRef.h" +#include "absl/types/variant.h" +#include "mlir/IR/BuiltinOps.h" // from @llvm-project #include "mlir/Pass/Pass.h" // from @llvm-project #include "tensorflow/compiler/mlir/tf2xla/api/v2/device_type.pb.h" -#include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" -#include "xla/pjrt/pjrt_compiler.h" +#include "xla/client/compile_only_client.h" #include "xla/pjrt/proto/compile_options.pb.h" -#include "xla/shape.h" +#include "xla/tsl/platform/statusor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h" @@ -48,7 +48,7 @@ namespace v2 { // arg_shapes - The shapes of the args. // arg_core_mapping - Which args go on which cores. // per_core_arg_shapes - For each core, the shapes for each argument. -// compiler - The Xla PjRtCompiler. +// client - The Xla Compilation client. absl::StatusOr LegalizeMlirToHlo( const std::variant& computation, const tpu::TPUCompileMetadataProto& metadata, bool use_tuple_args, @@ -58,7 +58,7 @@ absl::StatusOr LegalizeMlirToHlo( const std::vector& arg_shapes, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, - xla::PjRtCompiler* compiler = nullptr); + xla::CompileOnlyClient* client); }; // namespace v2 }; // namespace tf2xla diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf_test.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf_test.cc index 7c69394204d2d0..428bf6d0290dcb 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf_test.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf_test.cc @@ -28,8 +28,10 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.h" #include "tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "xla/pjrt/pjrt_compiler.h" +#include "xla/client/client_library.h" #include "xla/shape.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/platform_manager.h" #include "xla/tsl/lib/core/status_test_util.h" #include "xla/tsl/lib/monitoring/test_utils.h" #include "xla/tsl/platform/statusor.h" @@ -252,8 +254,10 @@ TEST(LegalizeTFTest, RecordsStreamzForNoMlirFallback) { /*graph_def_version=*/0, {&guaranteed_constants}}; - auto compiler = xla::GetDefaultPjRtCompiler(xla::CpuName()); - TF_ASSERT_OK(compiler.status()); + se::Platform* cpu_platform = + se::PlatformManager::PlatformWithName("Host").value(); + auto client = + xla::ClientLibrary::GetOrCreateCompileOnlyClient(cpu_platform).value(); std::vector arg_shapes; TPUCompileMetadataProto metadata_proto; @@ -268,7 +272,7 @@ TEST(LegalizeTFTest, RecordsStreamzForNoMlirFallback) { /*device_type=*/"XLA_CPU_JIT", custom_legalization_passes, /*shape_determination_fns=*/{}, arg_shapes, - &arg_core_mapping, &per_core_arg_shapes, *compiler); + &arg_core_mapping, &per_core_arg_shapes, client); EXPECT_FALSE(compile_result.ok()); } diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/testing/BUILD b/tensorflow/compiler/mlir/tf2xla/api/v2/testing/BUILD index 6213d9ad025f90..2fe5b204522837 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/testing/BUILD +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/testing/BUILD @@ -36,6 +36,7 @@ cc_library( "//tensorflow/compiler/tf2xla:xla_helpers", "//tensorflow/core:framework", "//tensorflow/core:lib", + "//tensorflow/core:test", "//tensorflow/core/protobuf:for_core_protos_cc", "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc", "//tensorflow/core/tpu/kernels:tpu_compile_op_support", @@ -43,8 +44,9 @@ cc_library( "@com_google_absl//absl/strings:string_view", "@llvm-project//mlir:Pass", "@xla//xla:shape_util", - "@xla//xla/pjrt:pjrt_compiler", - "@xla//xla/pjrt/plugin/xla_cpu:cpu_static_registration", + "@xla//xla/client:client_library", + "@xla//xla/stream_executor:platform", + "@xla//xla/stream_executor:platform_manager", ], ) diff --git a/tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.cc b/tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.cc index 8470889ba76f30..17d2490cd8c37e 100644 --- a/tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.cc +++ b/tensorflow/compiler/mlir/tf2xla/api/v2/testing/compile_mlir.cc @@ -24,8 +24,10 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h" #include "tensorflow/compiler/mlir/tf2xla/internal/utils/test_metadata_config.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "xla/pjrt/pjrt_compiler.h" +#include "xla/client/client_library.h" #include "xla/shape.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/platform_manager.h" #include "xla/tsl/platform/statusor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/protobuf/config.pb.h" @@ -49,8 +51,10 @@ absl::StatusOr CompileMlirModule( mlir_to_hlo_args.rollout_state = rollout_state; mlir_to_hlo_args.mlir_module = mlir_module_str; - TF_ASSIGN_OR_RETURN(auto* compiler, - xla::GetDefaultPjRtCompiler(xla::CpuName())); + TF_ASSIGN_OR_RETURN(se::Platform * platform, + se::PlatformManager::PlatformWithName("Host")); + TF_ASSIGN_OR_RETURN( + auto client, xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform)); std::vector arg_shapes; TPUCompileMetadataProto metadata_proto; @@ -68,7 +72,7 @@ absl::StatusOr CompileMlirModule( return LegalizeMlirToHlo(mlir_to_hlo_args, metadata_proto, use_tuple_args, device_type, custom_legalization_passes, /*shape_determination_fns=*/{}, arg_shapes, - &arg_core_mapping, &per_core_arg_shapes, compiler); + &arg_core_mapping, &per_core_arg_shapes, client); } } // namespace testing diff --git a/tensorflow/compiler/mlir/tf2xla/internal/BUILD b/tensorflow/compiler/mlir/tf2xla/internal/BUILD index ce091086781921..04ecb10a60da8d 100644 --- a/tensorflow/compiler/mlir/tf2xla/internal/BUILD +++ b/tensorflow/compiler/mlir/tf2xla/internal/BUILD @@ -185,8 +185,8 @@ cc_library( "@llvm-project//mlir:Pass", "@stablehlo//:register", "@xla//xla:shape_util", + "@xla//xla/client:compile_only_client", "@xla//xla/mlir_hlo:hlo_dialect_registration", - "@xla//xla/pjrt:pjrt_compiler", "@xla//xla/tsl/platform:errors", "@xla//xla/tsl/platform:statusor", ], @@ -206,6 +206,7 @@ tf_cc_test( "//tensorflow/core:framework", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test_main", + "//tensorflow/core/lib/monitoring:cell_reader", "//tensorflow/core/protobuf:for_core_protos_cc", "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc", "//tensorflow/core/tpu/kernels:tpu_compile_op_support", @@ -215,11 +216,11 @@ tf_cc_test( "@llvm-project//mlir:Pass", "@stablehlo//:register", "@xla//xla:shape_util", + "@xla//xla/client:client_library", "@xla//xla/mlir_hlo:hlo_dialect_registration", - "@xla//xla/pjrt:pjrt_compiler", - "@xla//xla/tsl/lib/monitoring:cell_reader", + "@xla//xla/stream_executor:platform", + "@xla//xla/stream_executor:platform_manager", "@xla//xla/tsl/platform:errors", - "@xla//xla/tsl/platform:statusor", ], ) diff --git a/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.cc b/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.cc index 5ad143b8104fd6..fcfbc0997cae05 100644 --- a/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.cc +++ b/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.cc @@ -35,8 +35,8 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_mlir.h" #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "xla/client/compile_only_client.h" #include "xla/mlir_hlo/mhlo/IR/register.h" -#include "xla/pjrt/pjrt_compiler.h" #include "xla/shape.h" #include "xla/tsl/framework/device_type.h" #include "xla/tsl/platform/errors.h" @@ -72,7 +72,7 @@ absl::StatusOr LegalizeTfToHlo( std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, std::vector>& custom_legalization_passes, - xla::PjRtCompiler* compiler, XlaCompilationResult* compilation_result) { + xla::CompileOnlyClient* client, XlaCompilationResult* compilation_result) { if (!mlir::SetTPUInfeedLayout(mlir_module_op)) return absl::AbortedError("Failed to set layouts attribute"); @@ -93,7 +93,7 @@ absl::StatusOr LegalizeTfToHlo( absl::Status old_bridge_status = v1::CompileTensorflowGraphToHlo( mlir_to_hlo_args, metadata, use_tuple_args, shape_determination_fns, arg_shapes, tsl::DeviceType(device_type.str()), arg_core_mapping, - per_core_arg_shapes, compiler, compilation_result); + per_core_arg_shapes, client, compilation_result); TF_RETURN_IF_ERROR(CheckAndIncrementCounter( old_bridge_status, MlirBridgeSecondPhaseMetric::kMlirCombinedOldFailure)); @@ -114,7 +114,7 @@ absl::StatusOr LegalizeTfToHlo( std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, std::vector>& custom_legalization_passes, - xla::PjRtCompiler* compiler, XlaCompilationResult* compilation_result) { + xla::CompileOnlyClient* client, XlaCompilationResult* compilation_result) { LOG_FIRST_N(INFO, 1) << "Compiling MLIR computation to XLA HLO using the " "Combined MLIR Tf2Xla Bridge."; @@ -122,7 +122,7 @@ absl::StatusOr LegalizeTfToHlo( return LegalizeTfToHlo(computation.mlir_module_op.value(), metadata, use_tuple_args, device_type, shape_determination_fns, arg_shapes, arg_core_mapping, per_core_arg_shapes, - custom_legalization_passes, compiler, + custom_legalization_passes, client, compilation_result); } @@ -142,7 +142,7 @@ absl::StatusOr LegalizeTfToHlo( return LegalizeTfToHlo(mlir_module.get(), metadata, use_tuple_args, device_type, shape_determination_fns, arg_shapes, arg_core_mapping, per_core_arg_shapes, - custom_legalization_passes, compiler, + custom_legalization_passes, client, compilation_result); } diff --git a/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h b/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h index c28b395ed69725..5c6e0f12f02b08 100644 --- a/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h +++ b/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo.h @@ -20,7 +20,7 @@ limitations under the License. #include "llvm/ADT/StringRef.h" #include "mlir/Pass/Pass.h" // from @llvm-project #include "tensorflow/compiler/tf2xla/xla_helpers.h" -#include "xla/pjrt/pjrt_compiler.h" +#include "xla/client/compile_only_client.h" #include "xla/tsl/platform/statusor.h" #include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h" #include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h" @@ -40,7 +40,7 @@ absl::StatusOr LegalizeTfToHlo( std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, std::vector>& custom_legalization_passes, - xla::PjRtCompiler* compiler, XlaCompilationResult* compilation_result); + xla::CompileOnlyClient* client, XlaCompilationResult* compilation_result); }; // namespace internal }; // namespace tf2xla diff --git a/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo_test.cc b/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo_test.cc index 77323a4ed42e2d..324a7f8d346a6f 100644 --- a/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo_test.cc +++ b/tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_to_hlo_test.cc @@ -33,14 +33,15 @@ limitations under the License. #include "tensorflow/compiler/mlir/tf2xla/internal/test_matchers.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "xla/client/client_library.h" #include "xla/mlir_hlo/mhlo/IR/register.h" -#include "xla/pjrt/pjrt_compiler.h" #include "xla/shape.h" -#include "xla/tsl/lib/monitoring/cell_reader.h" +#include "xla/stream_executor/platform.h" +#include "xla/stream_executor/platform_manager.h" #include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/lib/monitoring/cell_reader.h" #include "tensorflow/core/protobuf/config.pb.h" #include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h" #include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h" @@ -49,10 +50,10 @@ namespace tensorflow { namespace tf2xla { namespace internal { +using ::tensorflow::monitoring::testing::CellReader; using tpu::MlirToHloArgs; using tpu::ShardingAndIndex; using tpu::TPUCompileMetadataProto; -using tsl::monitoring::testing::CellReader; static constexpr char kMlirLegalizeCount[] = "/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_count"; @@ -101,8 +102,10 @@ absl::StatusOr CompileMlirModule( mlir_to_hlo_args.mlir_module_op = mlir_module.get(); } - TF_ASSIGN_OR_RETURN(auto* compiler, - xla::GetDefaultPjRtCompiler(xla::CpuName())); + se::Platform* platform = + se::PlatformManager::PlatformWithName("Host").value(); + auto client = + xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform).value(); std::vector arg_shapes = {{1}}; TPUCompileMetadataProto metadata_proto; @@ -120,7 +123,7 @@ absl::StatusOr CompileMlirModule( /*device_type=*/"XLA_TPU_JIT", /*shape_determination_fns=*/{}, arg_shapes, &arg_core_mapping, &per_core_arg_shapes, - custom_legalization_passes, compiler, + custom_legalization_passes, client, compilation_result.get()); } diff --git a/tensorflow/compiler/mlir/tf2xla/tests/registration/BUILD b/tensorflow/compiler/mlir/tf2xla/tests/registration/BUILD index e217bebb7a98e1..6ef93d44495b82 100644 --- a/tensorflow/compiler/mlir/tf2xla/tests/registration/BUILD +++ b/tensorflow/compiler/mlir/tf2xla/tests/registration/BUILD @@ -28,8 +28,8 @@ cc_library( ], deps = [ "//tensorflow/compiler/mlir/tensorflow", + "//tensorflow/compiler/mlir/tensorflow/translate:mlir_roundtrip_flags", "//tensorflow/compiler/mlir/tensorflow/translate/tools:file_tf_mlir_translate", - "//tensorflow/compiler/mlir/tf2xla/api/v2:mlir_roundtrip_flags", "//tensorflow/compiler/mlir/tf2xla/api/v2:tf_executor_to_graph", "//tensorflow/compiler/mlir/tools:translate_cl_options", "//tensorflow/compiler/tf2xla:xla_compiler", @@ -46,8 +46,13 @@ cc_library( "@llvm-project//mlir:Support", "@llvm-project//mlir:TranslateLib", "@tsl//tsl/platform:protobuf", - "@xla//xla/pjrt:pjrt_compiler", - "@xla//xla/pjrt/plugin/xla_cpu:cpu_static_registration", + "@xla//xla/client:client_library", + "@xla//xla/client:compile_only_client", + "@xla//xla/service/cpu:cpu_compiler", + "@xla//xla/service/cpu:cpu_transfer_manager", + "@xla//xla/stream_executor:platform_manager", + "@xla//xla/stream_executor/host:host_platform", + "@xla//xla/stream_executor/host:host_platform_id", "@xla//xla/tsl/platform:status", ], alwayslink = 1, diff --git a/tensorflow/compiler/mlir/tf2xla/tests/registration/graph_to_tf_executor_registration.cc b/tensorflow/compiler/mlir/tf2xla/tests/registration/graph_to_tf_executor_registration.cc index 53b69438337437..8edb01f90063f6 100644 --- a/tensorflow/compiler/mlir/tf2xla/tests/registration/graph_to_tf_executor_registration.cc +++ b/tensorflow/compiler/mlir/tf2xla/tests/registration/graph_to_tf_executor_registration.cc @@ -27,13 +27,16 @@ limitations under the License. #include "mlir/Support/LogicalResult.h" // from @llvm-project #include "mlir/Tools/mlir-translate/Translation.h" // from @llvm-project #include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h" +#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h" #include "tensorflow/compiler/mlir/tensorflow/translate/tools/file_tf_mlir_translate.h" -#include "tensorflow/compiler/mlir/tf2xla/api/v2/mlir_roundtrip_flags.h" #include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_executor_to_graph.h" #include "tensorflow/compiler/mlir/tools/tf_mlir_translate_cl.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" -#include "xla/pjrt/pjrt_compiler.h" +#include "xla/client/client_library.h" +#include "xla/client/compile_only_client.h" +#include "xla/stream_executor/host/host_platform_id.h" +#include "xla/stream_executor/platform_manager.h" #include "xla/tsl/platform/status.h" #include "tensorflow/core/common_runtime/graph_constructor.h" #include "tensorflow/core/framework/function.h" @@ -55,10 +58,10 @@ static constexpr char kMlirToGraphCompilationCheckName[] = static constexpr char kArbitraryDeviceName[] = "XLA_CPU_JIT"; static Status CompileGraph(tensorflow::Graph* graph, - xla::PjRtCompiler* compiler) { - if (!graph || !compiler) { + xla::CompileOnlyClient* client) { + if (!graph || !client) { return Status(absl::StatusCode::kInvalidArgument, - "Invalid graph or compiler"); + "Invalid graph or client"); } tensorflow::FunctionDefLibrary flib; @@ -67,9 +70,9 @@ static Status CompileGraph(tensorflow::Graph* graph, tensorflow::XlaCompiler::Options options; options.device_type = tensorflow::DeviceType(kArbitraryDeviceName); - options.compiler = compiler; + options.client = client; options.flib_def = flib_def.get(); - tensorflow::XlaCompiler xla_compiler(options); + tensorflow::XlaCompiler compiler(options); std::unique_ptr graph_copy = std::make_unique(tensorflow::OpRegistry::Global()); @@ -77,9 +80,9 @@ static Status CompileGraph(tensorflow::Graph* graph, tensorflow::XlaCompiler::CompileOptions compile_options; tensorflow::XlaCompiler::CompilationResult result; - return xla_compiler.CompileGraph(compile_options, - kMlirToGraphCompilationCheckName, - std::move(graph_copy), {}, &result); + return compiler.CompileGraph(compile_options, + kMlirToGraphCompilationCheckName, + std::move(graph_copy), {}, &result); } static mlir::OwningOpRef GraphdefToMlirTranslateFunction( @@ -120,15 +123,19 @@ static mlir::LogicalResult MlirToGraphTranslateFunction( return mlir::failure(); } - auto compiler = xla::GetDefaultPjRtCompiler(xla::CpuName()); - if (!compiler.ok()) { + // Use Host platform, which should always exist, to make sure graphs compile. + auto platform = stream_executor::PlatformManager::PlatformWithId( + stream_executor::host::kHostPlatformId); + if (!platform.ok()) { return mlir::failure(); } + auto client = + xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform.value()); tensorflow::XlaOpRegistry::RegisterCompilationKernels(); // Verify that the resulting graph can compile. - if (!CompileGraph(graph.get(), compiler.value()).ok()) { + if (client.ok() && !CompileGraph(graph.get(), client.value()).ok()) { return mlir::failure(); } diff --git a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD index 3fc654d38b506d..c2564414993142 100644 --- a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD +++ b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD @@ -162,7 +162,6 @@ cc_library( "//tensorflow/core/tpu/kernels/xla:host_compute_ops", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", - "@com_google_absl//absl/log:vlog_is_on", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -173,15 +172,15 @@ cc_library( "@llvm-project//mlir:Pass", "@llvm-project//mlir:Support", "@tsl//tsl/platform:fingerprint", - "@tsl//tsl/platform:protobuf", "@xla//xla:shape_util", "@xla//xla:xla_data_proto_cc", + "@xla//xla/client:client_library", "@xla//xla/hlo/translate/hlo_to_mhlo:hlo_to_mlir_hlo", "@xla//xla/pjrt:pjrt_compiler", - "@xla//xla/pjrt/plugin/xla_cpu:cpu_static_registration", "@xla//xla/python/ifrt", "@xla//xla/service:device_assignment", "@xla//xla/service:hlo_proto_cc", + "@xla//xla/stream_executor:platform_manager", "@xla//xla/tsl/platform:errors", ], ) diff --git a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.cc b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.cc index 5a810d4d7c5245..b1b8aa73d877fd 100644 --- a/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.cc +++ b/tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.cc @@ -23,7 +23,6 @@ limitations under the License. #include "absl/log/check.h" #include "absl/log/log.h" -#include "absl/log/vlog_is_on.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" @@ -52,12 +51,14 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "xla/client/client_library.h" #include "xla/hlo/translate/hlo_to_mhlo/hlo_to_mlir_hlo.h" #include "xla/pjrt/pjrt_compiler.h" #include "xla/python/ifrt/client.h" #include "xla/python/ifrt/layout.h" #include "xla/service/device_assignment.h" #include "xla/shape.h" +#include "xla/stream_executor/platform_manager.h" #include "xla/tsl/lib/strings/proto_serialization.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/statusor.h" @@ -70,7 +71,6 @@ limitations under the License. #include "tensorflow/core/protobuf/tpu/topology.pb.h" #include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h" #include "tsl/platform/fingerprint.h" -#include "tsl/platform/protobuf.h" namespace tensorflow { namespace ifrt_serving { @@ -244,8 +244,12 @@ absl::StatusOr CompileTfToHlo(const Tf2HloArg& arg) { } VLOG(1) << "device_type: " << device_type; - TF_ASSIGN_OR_RETURN(auto* compiler, - xla::GetDefaultPjRtCompiler(xla::CpuName())); + TF_ASSIGN_OR_RETURN( + auto* platform, + stream_executor::PlatformManager::PlatformWithName("Host")); + TF_ASSIGN_OR_RETURN( + auto* client, xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform)); + std::vector arg_shapes; arg_shapes.reserve(arg.input_dtypes_and_shapes.size()); @@ -282,7 +286,7 @@ absl::StatusOr CompileTfToHlo(const Tf2HloArg& arg) { tensorflow::XlaShapeLayoutHelpers::ShapeDeterminationFns( tensorflow::UseNoPreferenceLayoutFn(), arg.shape_representation_fn), - arg_shapes, &arg_core_mapping, &per_core_arg_shapes, compiler)); + arg_shapes, &arg_core_mapping, &per_core_arg_shapes, client)); for (auto arg_shapes_iter = per_core_arg_shapes.begin() + 1; arg_shapes_iter != per_core_arg_shapes.end(); ++arg_shapes_iter) { diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index 7291da1ad16f9a..56b3fc99746d7d 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -110,6 +110,10 @@ py_test( ], ) +#LINT.IfChange(combined_tests) +# If you add a new tf_xla_py_strict_test please either add the test file to one of the combined test +# targets that matches in all tags and other settings or add a new combined test target. + tf_xla_combined_py_test( name = "combined_ops_test_a", size = "medium", @@ -272,7 +276,9 @@ tf_xla_combined_py_test( # go/keep-sorted end ], ) +#LINT.ThenChange(:individual_tests) +#LINT.IfChange(individual_tests) tf_xla_py_strict_test( name = "adadelta_test", size = "medium", @@ -2583,12 +2589,13 @@ tf_xla_py_strict_test( "//tensorflow/compiler/mlir/stablehlo", "//tensorflow/compiler/tf2xla/ops:gen_xla_ops", "//tensorflow/compiler/tf2xla/python:xla", - "//tensorflow/core/protobuf:for_core_protos_py_proto", + "//tensorflow/core:protos_all_py", "//tensorflow/python/eager:def_function", "//tensorflow/python/framework:dtypes", "//tensorflow/python/framework:errors", "//tensorflow/python/framework:function", "//tensorflow/python/framework:ops", + "//tensorflow/python/ops:array_ops", "//tensorflow/python/platform:client_testlib", "//tensorflow/python/platform:test", "//third_party/py/numpy", @@ -2613,12 +2620,13 @@ tf_xla_py_strict_test( "//tensorflow/compiler/mlir/stablehlo", "//tensorflow/compiler/tf2xla/ops:gen_xla_ops", "//tensorflow/compiler/tf2xla/python:xla", - "//tensorflow/core/protobuf:for_core_protos_py_proto", + "//tensorflow/core:protos_all_py", "//tensorflow/python/eager:def_function", "//tensorflow/python/framework:dtypes", "//tensorflow/python/framework:errors", "//tensorflow/python/framework:function", "//tensorflow/python/framework:ops", + "//tensorflow/python/ops:array_ops", "//tensorflow/python/platform:client_testlib", "//tensorflow/python/platform:test", "//third_party/py/numpy", @@ -2642,12 +2650,13 @@ tf_xla_py_strict_test( "//tensorflow/compiler/mlir/stablehlo", "//tensorflow/compiler/tf2xla/ops:gen_xla_ops", "//tensorflow/compiler/tf2xla/python:xla", - "//tensorflow/core/protobuf:for_core_protos_py_proto", + "//tensorflow/core:protos_all_py", "//tensorflow/python/eager:def_function", "//tensorflow/python/framework:dtypes", "//tensorflow/python/framework:errors", "//tensorflow/python/framework:function", "//tensorflow/python/framework:ops", + "//tensorflow/python/ops:array_ops", "//tensorflow/python/platform:client_testlib", "//tensorflow/python/platform:test", "//third_party/py/numpy", @@ -2738,21 +2747,23 @@ tf_xla_py_strict_test( ) # copybara:uncomment_begin(google-only) -tf_xla_py_strict_test( - name = "xla_dump_to_sponge_test", - size = "medium", - srcs = ["xla_dump_to_sponge_test.py"], - tags = [ - "optonly", - ], - deps = [ - ":xla_test", - "//tensorflow/python/ops:array_ops", - "//tensorflow/python/ops:math_ops", - "//tensorflow/python/platform:test", - "//third_party/py/numpy", - ], -) +# tf_xla_py_strict_test( +# name = "xla_dump_to_sponge_test", +# size = "medium", +# srcs = ["xla_dump_to_sponge_test.py"], +# tags = [ +# "optonly", +# ], +# deps = [ +# ":xla_test", +# "//third_party/py/numpy", +# "//tensorflow/python/ops:array_ops", +# "//tensorflow/python/ops:math_ops", +# "//tensorflow/python/platform:test", +# ], +# ) +# copybara:uncomment_end +#LINT.ThenChange(:combined_tests) cuda_py_strict_test( name = "xla_device_gpu_test", diff --git a/tensorflow/compiler/tests/searchsorted_op_test.py b/tensorflow/compiler/tests/searchsorted_op_test.py index 216b9cb799a5af..d22e329e95c6b0 100644 --- a/tensorflow/compiler/tests/searchsorted_op_test.py +++ b/tensorflow/compiler/tests/searchsorted_op_test.py @@ -47,12 +47,15 @@ def testNanValue(self): sorted_sequence = np.array([[1, 3, 5, 7, 9]], dtype) values = np.array([[np.nan]], dtype) + expected_left = np.array([5], dtype=np.int32) + expected_right = np.array([5], dtype=np.int32) + self._test2DExample( dtype, 'left', sorted_sequence, values, - np.array([[0]], dtype=np.int32), + expected_left.reshape(1, -1).astype(np.int32), ) self._test2DExample( @@ -60,7 +63,7 @@ def testNanValue(self): 'right', sorted_sequence, values, - np.array([[5]], dtype=np.int32), + expected_right.reshape(1, -1).astype(np.int32), ) def _test2DExample(self, dtype, side, sorted_sequence, values, correct_ans): @@ -87,6 +90,82 @@ def testUpperBound2DExample(self): correct_ans = np.array([[1, 2, 4], [0, 2, 5]], dtype) self._test2DExample(dtype, 'right', sorted_sequence, values, correct_ans) + def testNanInSortedSequence(self): + for dtype in self.float_types: + sorted_sequence = np.array([[2.0, 4.0, 8.0, np.nan, np.nan]], dtype) + values = np.array([[3.0, np.nan]], dtype) + + expected_left = np.array([1, 3], dtype=np.int32) + expected_right = np.array([1, 5], dtype=np.int32) + + self._test2DExample( + dtype, + 'left', + sorted_sequence, + values, + expected_left.reshape(1, -1).astype(np.int32), + ) + + self._test2DExample( + dtype, + 'right', + sorted_sequence, + values, + expected_right.reshape(1, -1).astype(np.int32), + ) + + def testNanRegression(self): + for dtype in self.float_types: + sorted_sequence = np.array([[2.0, 4.0, np.nan, np.nan]], dtype) + values = np.array([[3.0, 5.0, np.nan]], dtype) + + expected_left = np.array([1, 2, 2], dtype=np.int32) + expected_right = np.array([1, 2, 4], dtype=np.int32) + + self._test2DExample( + dtype, + 'left', + sorted_sequence, + values, + expected_left.reshape(1, -1).astype(np.int32), + ) + + self._test2DExample( + dtype, + 'right', + sorted_sequence, + values, + expected_right.reshape(1, -1).astype(np.int32), + ) + + def testLowerBoundNaN(self): + sorted_sequence = np.array( + [[2.0, 4.0, 8.0, 16.0, 32.0, 64.0]], dtype=np.float32 + ) + values = np.array([[np.nan, 8.0]], dtype=np.float32) + expected = np.array([6, 2], dtype=np.int32) + with self.session() as session: + with self.test_scope(): + tf_ans = array_ops.searchsorted(sorted_sequence, values, side='left') + tf_out = session.run(tf_ans) + self.assertAllEqual(expected.reshape(1, -1), tf_out) + + def testUpperBoundNaN(self): + sorted_sequence = np.array( + [[2.0, 4.0, 8.0, 16.0, 32.0, 64.0]], dtype=np.float32 + ) + + values = np.array([[np.nan, 8.0]], dtype=np.float32) + + expected = np.array([6, 3], dtype=np.int32) + + with self.session() as session: + with self.test_scope(): + tf_ans = array_ops.searchsorted(sorted_sequence, values, side='right') + tf_out = session.run(tf_ans) + + self.assertAllEqual(expected.reshape(1, -1), tf_out) + if __name__ == '__main__': test.main() diff --git a/tensorflow/compiler/tests/xla_call_module_test.py b/tensorflow/compiler/tests/xla_call_module_test.py index 4e6f56172c9bf7..de4a3ab8018cfc 100644 --- a/tensorflow/compiler/tests/xla_call_module_test.py +++ b/tensorflow/compiler/tests/xla_call_module_test.py @@ -1590,8 +1590,7 @@ def f(x): errors.InvalidArgumentError, 'Cannot deserialize computation: .+', ): - with self.test_scope(): - f(x) + f(x) if __name__ == '__main__': diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index 513459692a6370..59283a2bbe0d59 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -186,8 +186,10 @@ cc_library( hdrs = ["tf2xla.h"], visibility = ["//visibility:public"], deps = [ + ":common", ":graph_compiler_util", ":tf2xla_proto_cc", + ":tf2xla_util", ":xla_compiler", ":xla_op_registry", "//tensorflow/compiler/aot:aot_only_var_handle_op", @@ -199,13 +201,9 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/platform:status", - "@com_google_absl//absl/log", - "@com_google_absl//absl/status", "@com_google_absl//absl/strings", + "@xla//xla/client", "@xla//xla/hlo/builder:xla_computation", - "@xla//xla/pjrt:pjrt_client", - "@xla//xla/pjrt:pjrt_compiler", - "@xla//xla/pjrt/plugin/xla_cpu:cpu_static_registration", ], ) @@ -243,6 +241,7 @@ cc_library( "@llvm-project//mlir:IR", "@tsl//tsl/platform:errors", "@tsl//tsl/platform:statusor", + "@xla//xla/client", "@xla//xla/hlo/builder:xla_computation", "@xla//xla/tsl/platform:statusor", ], @@ -558,37 +557,29 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core/platform:errors", - "@com_google_absl//absl/base", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", - "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/types:span", "@tsl//tsl/platform:casts", - "@xla//xla:executable_run_options", "@xla//xla:shape_util", - "@xla//xla:status_macros", "@xla//xla:xla_data_proto_cc", "@xla//xla/backends/cpu:buffer_allocation_info", "@xla//xla/backends/cpu:buffer_allocation_info_util", "@xla//xla/backends/cpu/codegen:compiled_function_library", "@xla//xla/backends/cpu/runtime/thunk_serdes:all_thunk_serdes", # buildcleaner: keep + "@xla//xla/client:client_library", "@xla//xla/client:executable_build_options", + "@xla//xla/client:local_client", "@xla//xla/hlo/builder:xla_computation", - "@xla//xla/pjrt:common_pjrt_client", - "@xla//xla/pjrt:pjrt_client", - "@xla//xla/pjrt:pjrt_compiler", - "@xla//xla/pjrt:pjrt_executable", - "@xla//xla/pjrt/cpu:cpu_client", - "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options", - "@xla//xla/pjrt/plugin/xla_cpu:cpu_static_registration", - "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client", - "@xla//xla/service:buffer_assignment", + "@xla//xla/service:platform_util", "@xla//xla/service/cpu:cpu_aot_compilation_result", "@xla//xla/service/cpu:executable_proto_cc", + "@xla//xla/stream_executor:platform", ] + if_libtpu( if_false = [ + "@xla//xla/service:cpu_plugin", "@xla//xla/service/cpu:cpu_executable", ], if_true = [], @@ -671,12 +662,8 @@ cc_library( "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:inlined_vector", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", "@com_google_absl//absl/memory", "@com_google_absl//absl/status", - "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/synchronization", @@ -701,16 +688,12 @@ cc_library( "@xla//xla/hlo/builder:xla_computation", "@xla//xla/hlo/ir:hlo", "@xla//xla/hlo/translate/mhlo_to_hlo:layout_util", - "@xla//xla/pjrt:pjrt_client", - "@xla//xla/pjrt:pjrt_compiler", - "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client", "@xla//xla/service:device_assignment", "@xla//xla/service:hlo_proto_cc", "@xla//xla/service/spmd/shardy:constants", "@xla//xla/service/spmd/shardy:utils", "@xla//xla/service/spmd/shardy/stablehlo_round_trip:stablehlo_import", "@xla//xla/stream_executor:device_memory_allocator", - "@xla//xla/tsl/platform:logging", ] + if_libtpu([ ":xla_tpu_backend_registration", ]) + if_static(["@tsl//tsl/platform:tensor_float_32_utils"]), @@ -833,18 +816,11 @@ cc_library( "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings", - "@xla//xla:literal", - "@xla//xla:shape_util", + "@com_google_absl//absl/types:optional", + "@xla//xla/client", "@xla//xla/hlo/builder:value_inference", "@xla//xla/hlo/builder:xla_builder", - "@xla//xla/hlo/builder:xla_computation", - "@xla//xla/pjrt:pjrt_client", - "@xla//xla/pjrt:pjrt_compiler", - "@xla//xla/pjrt:pjrt_executable", ], alwayslink = 1, ) @@ -1195,20 +1171,14 @@ tf_cc_test( "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", - "@com_google_absl//absl/status", - "@com_google_absl//absl/status:status_matchers", "@com_google_absl//absl/status:statusor", - "@com_google_absl//absl/strings:string_view", - "@com_google_googletest//:gtest_main", "@tsl//tsl/platform:tensor_float_32_hdr_lib", "@xla//xla:literal", "@xla//xla:literal_util", + "@xla//xla/client:client_library", + "@xla//xla/client:local_client", "@xla//xla/hlo/builder:xla_computation", - "@xla//xla/pjrt:pjrt_client", - "@xla//xla/pjrt:pjrt_compiler", - "@xla//xla/pjrt:pjrt_executable", - "@xla//xla/pjrt/cpu:cpu_client", - "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options", + "@xla//xla/service:cpu_plugin", ] + if_static(["@tsl//tsl/platform:tensor_float_32_utils"]), ) @@ -1220,22 +1190,20 @@ tf_cc_test( ":xla_compiled_cpu_function_thunks", ":xla_jit_compiled_cpu_function", "//tensorflow/core:framework", + "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status:statusor", - "@tsl//tsl/platform:protobuf", "@xla//xla:shape_util", "@xla//xla:xla_data_proto_cc", "@xla//xla/client:executable_build_options", "@xla//xla/hlo/testlib:test", - "@xla//xla/pjrt/plugin/xla_cpu:cpu_static_registration", "@xla//xla/service:compiler", + "@xla//xla/service:platform_util", "@xla//xla/stream_executor:platform", - "@xla//xla/stream_executor:platform_id", "@xla//xla/stream_executor:platform_manager", - "@xla//xla/tsl/platform:statusor", ], ) @@ -1269,31 +1237,34 @@ tf_cc_test( "//tensorflow/core:testlib", "//tensorflow/core/framework:resource_base", "//tensorflow/core/framework:tensor_testutil", + "//tensorflow/core/platform:errors", + "//tensorflow/core/platform:types", "@com_google_absl//absl/log", + "@com_google_absl//absl/memory", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_googletest//:gtest", + "@tsl//tsl/platform:statusor", "@xla//xla:array", "@xla//xla:debug_options_flags", "@xla//xla:literal", "@xla//xla:literal_util", "@xla//xla:shape_util", + "@xla//xla:status_macros", "@xla//xla:xla_data_proto_cc", + "@xla//xla/client", + "@xla//xla/client:client_library", + "@xla//xla/client:local_client", "@xla//xla/hlo/builder:xla_builder", - "@xla//xla/hlo/builder:xla_computation", "@xla//xla/hlo/ir:hlo", "@xla//xla/hlo/ir:hlo_sharding", "@xla//xla/hlo/testlib:filecheck", - "@xla//xla/pjrt:pjrt_client", - "@xla//xla/pjrt:pjrt_compiler", - "@xla//xla/pjrt:pjrt_executable", - "@xla//xla/pjrt/cpu:cpu_client", - "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options", - "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client", "@xla//xla/service", + "@xla//xla/service:cpu_plugin", "@xla//xla/service:hlo_module_util", "@xla//xla/service:hlo_proto_cc", + "@xla//xla/service:hlo_proto_util", "@xla//xla/tests:literal_test_util", "@xla//xla/tsl/platform:errors", "@xla//xla/tsl/platform:statusor", diff --git a/tensorflow/compiler/tf2xla/graph_compiler.cc b/tensorflow/compiler/tf2xla/graph_compiler.cc index a4fb37e1dcf55c..4d3b87a177a25a 100644 --- a/tensorflow/compiler/tf2xla/graph_compiler.cc +++ b/tensorflow/compiler/tf2xla/graph_compiler.cc @@ -15,43 +15,42 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/graph_compiler.h" -#include -#include -#include -#include +#include +#include #include #include -#include "absl/container/inlined_vector.h" -#include "absl/log/check.h" -#include "absl/log/log.h" -#include "absl/status/status.h" -#include "absl/strings/str_cat.h" #include "tensorflow/compiler/tf2xla/const_analysis.h" +#include "tensorflow/compiler/tf2xla/literal_util.h" +#include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/side_effect_util.h" +#include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_context.h" #include "tensorflow/compiler/tf2xla/xla_expression.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" -#include "tensorflow/compiler/tf2xla/xla_resource.h" +#include "xla/client/client_library.h" #include "xla/hlo/builder/xla_builder.h" -#include "xla/status_macros.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/logging.h" -#include "xla/tsl/platform/statusor.h" -#include "tensorflow/core/common_runtime/function_utils.h" -#include "tensorflow/core/framework/allocator.h" +#include "tensorflow/core/common_runtime/device.h" +#include "tensorflow/core/common_runtime/executor.h" +#include "tensorflow/core/common_runtime/function.h" +#include "tensorflow/core/common_runtime/graph_constructor.h" +#include "tensorflow/core/common_runtime/graph_optimizer.h" #include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/attr_value_util.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/op_kernel.h" -#include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/graph/algorithm.h" -#include "tensorflow/core/graph/graph_node_util.h" +#include "tensorflow/core/graph/node_builder.h" #include "tensorflow/core/graph/validate.h" +#include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/gtl/cleanup.h" +#include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/lib/monitoring/counter.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/public/version.h" +#include "tensorflow/core/util/dump_graph.h" namespace tensorflow { diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD index f89737523dcef9..d2502de9996580 100644 --- a/tensorflow/compiler/tf2xla/kernels/BUILD +++ b/tensorflow/compiler/tf2xla/kernels/BUILD @@ -2647,7 +2647,6 @@ tf_kernel_library( name = "lower_upper_bound_ops", srcs = ["lower_upper_bound_ops.cc"], deps = [ - "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:xla_compilation_device", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla:xla_context", @@ -2656,11 +2655,10 @@ tf_kernel_library( "//tensorflow/compiler/tf2xla:xla_resource", "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/core:framework", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core/platform:errors", "@com_google_absl//absl/status", "@xla//xla:comparison_util", "@xla//xla/hlo/builder:xla_builder", + "@xla//xla/hlo/builder/lib:math", ], ) diff --git a/tensorflow/compiler/tf2xla/kernels/lower_upper_bound_ops.cc b/tensorflow/compiler/tf2xla/kernels/lower_upper_bound_ops.cc index 642438b6014e0e..53b439186e1fc0 100644 --- a/tensorflow/compiler/tf2xla/kernels/lower_upper_bound_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/lower_upper_bound_ops.cc @@ -18,12 +18,9 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/comparison_util.h" +#include "xla/hlo/builder/lib/math.h" #include "xla/hlo/builder/xla_builder.h" -#include "tensorflow/core/framework/op_kernel.h" -#include "tensorflow/core/framework/op_requires.h" -#include "tensorflow/core/framework/tensor_shape.h" -#include "tensorflow/core/framework/types.pb.h" -#include "tensorflow/core/platform/errors.h" +#include "tensorflow/core/framework/types.h" namespace tensorflow { namespace { @@ -33,8 +30,7 @@ namespace { // Note that this is an O(MN) algorithm: all entries in each sorted_inputs row // are considered, and their sorted nature is not fully exploited. void BuildLowerUpperBoundOp(XlaOpKernelContext* ctx, DataType out_dtype, - xla::ComparisonDirection comparison_direction, - bool nan_values_compare_greater = false) { + xla::ComparisonDirection comparison_direction) { const TensorShape sorted_inputs_shape = ctx->InputShape("sorted_inputs"); const TensorShape values_shape = ctx->InputShape("values"); const xla::XlaOp sorted_inputs = ctx->Input("sorted_inputs"); @@ -47,6 +43,10 @@ void BuildLowerUpperBoundOp(XlaOpKernelContext* ctx, DataType out_dtype, OP_REQUIRES(ctx, values_shape.dims() == 2, absl::FailedPreconditionError("values must be 2D")); + bool is_fp = tensorflow::DataTypeIsFloating(ctx->InputType("sorted_inputs")); + + xla::XlaBuilder* builder = ctx->builder(); + // Add a new inner dimension to values, to allow broadcasting along the inner // dimension of sorted_sequence. auto new_values_shape = values_shape; @@ -60,33 +60,36 @@ void BuildLowerUpperBoundOp(XlaOpKernelContext* ctx, DataType out_dtype, auto sorted_inputs_reshaped = xla::Reshape(sorted_inputs, new_sorted_inputs_shape.dim_sizes()); - // We are relying on broadcasting to compare each value against each entry in - // the associated sorted_inputs row. - // The reshapes above leave the tensors with equal rank of 3, so broadcast - // dimensions are not explicitly specified. - auto comparison = xla::Compare(values_reshaped, sorted_inputs_reshaped, {}, - comparison_direction); - if (nan_values_compare_greater) { - // Match eager searchsorted side='right' behavior for NaN search values. - // A NaN search value is placed after all entries in the sorted input row. - auto value_is_nan = xla::Compare(values_reshaped, values_reshaped, {}, - xla::ComparisonDirection::kNe); - comparison = xla::Or(comparison, value_is_nan); - } const DataType accumulation_type = XlaHelpers::SumAccumulationType(out_dtype); + xla::XlaOp zero = XlaHelpers::Zero(builder, accumulation_type); + + xla::XlaOp cmp; + if (is_fp) { + xla::XlaOp is_val_nan = xla::IsNan(values_reshaped); + xla::XlaOp is_seq_nan = xla::IsNan(sorted_inputs_reshaped); + + if (comparison_direction == xla::ComparisonDirection::kGt) { + // NanAwareGt(a, b) = (!is_b_nan && (is_a_nan || (a > b))) + xla::XlaOp standard_gt = xla::Gt(values_reshaped, sorted_inputs_reshaped); + cmp = xla::And(xla::Not(is_seq_nan), xla::Or(is_val_nan, standard_gt)); + } else { + // NanAwareGe(a, b) = is_a_nan || (!is_b_nan && (a >= b)) + xla::XlaOp standard_ge = xla::Ge(values_reshaped, sorted_inputs_reshaped); + cmp = xla::Or(is_val_nan, xla::And(xla::Not(is_seq_nan), standard_ge)); + } + } else { + if (comparison_direction == xla::ComparisonDirection::kGt) { + cmp = xla::Gt(values_reshaped, sorted_inputs_reshaped); + } else { + cmp = xla::Ge(values_reshaped, sorted_inputs_reshaped); + } + } - // Convert boolean comparison results to integers so we can sum them. - auto comparison_int = - XlaHelpers::ConvertElementType(comparison, accumulation_type); - - // Sum the comparison results over the inner dimension to find the index for - // each value. - xla::XlaBuilder* builder = ctx->builder(); - auto reduced = - xla::Reduce(comparison_int, XlaHelpers::Zero(builder, accumulation_type), - *ctx->GetOrCreateAdd(accumulation_type), {2}); + xla::XlaOp cmp_int = XlaHelpers::ConvertElementType(cmp, accumulation_type); + xla::XlaOp result_count = + xla::Reduce(cmp_int, zero, *ctx->GetOrCreateAdd(accumulation_type), {2}); - ctx->SetOutput(0, reduced); + ctx->SetOutput(0, result_count); } class LowerBoundOp : public XlaOpKernel { @@ -112,8 +115,7 @@ class UpperBoundOp : public XlaOpKernel { } void Compile(XlaOpKernelContext* ctx) override { - BuildLowerUpperBoundOp(ctx, out_dtype_, xla::ComparisonDirection::kGe, - /*nan_values_compare_greater=*/true); + BuildLowerUpperBoundOp(ctx, out_dtype_, xla::ComparisonDirection::kGe); } private: diff --git a/tensorflow/compiler/tf2xla/kernels/while_op.cc b/tensorflow/compiler/tf2xla/kernels/while_op.cc index 4abaa081653f8f..4634d2500bd458 100644 --- a/tensorflow/compiler/tf2xla/kernels/while_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/while_op.cc @@ -159,7 +159,7 @@ absl::Status ConvertLoopInvariantsToConst( const NameAttrList& cond_name_attr, std::vector* args, std::vector* compile_time_const_arg_indices, - int* num_compile_time_const_args) { + int* num_compile_time_const_args, xla::Client* client) { std::vector loop_invariants(ctx->num_inputs()); GetLoopInvariants(ctx, body_name_attr, &loop_invariants); @@ -381,7 +381,7 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { OP_REQUIRES_OK(ctx, ConvertLoopInvariantsToConst( ctx, body_name_attr_, cond_name_attr_, &arguments, &compile_time_const_arg_indices, - &num_compile_time_const_args)); + &num_compile_time_const_args, compiler->client())); } VLOG(1) << "Compiling body"; diff --git a/tensorflow/compiler/tf2xla/tf2xla.cc b/tensorflow/compiler/tf2xla/tf2xla.cc index 8acadcc451dcb9..49c0eaad39b1d6 100644 --- a/tensorflow/compiler/tf2xla/tf2xla.cc +++ b/tensorflow/compiler/tf2xla/tf2xla.cc @@ -15,31 +15,38 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/tf2xla.h" -#include +#include #include #include +#include #include #include -#include "absl/log/log.h" -#include "absl/status/status.h" #include "absl/strings/str_cat.h" +#include "absl/strings/str_join.h" #include "tensorflow/compiler/aot/aot_only_var_handle_op.h" #include "tensorflow/compiler/tf2xla/graph_compiler_util.h" +#include "tensorflow/compiler/tf2xla/shape_util.h" +#include "tensorflow/compiler/tf2xla/tf2xla_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/hlo/builder/xla_computation.h" -#include "xla/pjrt/pjrt_client.h" -#include "xla/pjrt/pjrt_compiler.h" -#include "xla/tsl/platform/errors.h" +#include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/framework/graph_def_util.h" #include "tensorflow/core/framework/node_def.pb.h" +#include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/tensor_shape.h" -#include "tensorflow/core/framework/types.h" #include "tensorflow/core/framework/versions.pb.h" +#include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/graph/graph.h" +#include "tensorflow/core/graph/node_builder.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/platform/errors.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/types.h" +#include "tensorflow/core/util/dump_graph.h" namespace tensorflow { @@ -49,9 +56,8 @@ namespace { // graph symbolically, with each op building up the XLA HLO. absl::Status ConvertGraphToXla(std::unique_ptr graph, const tf2xla::Config& config, - xla::PjRtCompiler* compiler, - xla::XlaComputation* computation, - xla::PjRtClient* client) { + xla::Client* client, + xla::XlaComputation* computation) { XlaOpRegistry::RegisterCompilationKernels(); for (Node* node : graph->nodes()) { node->set_assigned_device_name( @@ -63,20 +69,19 @@ absl::Status ConvertGraphToXla(std::unique_ptr graph, PopulateXlaArgs(config, &xla_args); // Compile the graph into an XLA computation. XlaCompiler::Options compiler_options; - compiler_options.compiler = compiler; compiler_options.client = client; compiler_options.device_type = DeviceType(DEVICE_CPU_XLA_JIT); compiler_options.flib_def = &graph->flib_def(); compiler_options.graph_def_version = graph->versions().producer(); compiler_options.allow_cpu_custom_calls = true; - XlaCompiler xla_compiler(compiler_options); + XlaCompiler compiler(compiler_options); XlaCompiler::CompilationResult result; XlaCompiler::CompileOptions options; options.alias_resource_update = true; - TF_RETURN_IF_ERROR(xla_compiler.CompileGraph( + TF_RETURN_IF_ERROR(compiler.CompileGraph( options, "tfcompile", std::move(graph), xla_args, &result)); *computation = std::move(*result.computation); @@ -154,14 +159,13 @@ absl::Status ConvertVarHandlesToAotVarHandles(GraphDef* graph_def) { absl::Status ConvertGraphDefToXla(GraphDef graph_def, const tf2xla::Config& config, - xla::PjRtCompiler* compiler, - xla::XlaComputation* computation, - xla::PjRtClient* client) { + xla::Client* client, + xla::XlaComputation* computation) { std::unique_ptr graph; TF_RETURN_IF_ERROR(ConvertVarHandlesToAotVarHandles(&graph_def)); TF_RETURN_IF_ERROR(InitGraph(graph_def, config, &graph)); - TF_RETURN_IF_ERROR(ConvertGraphToXla(std::move(graph), config, compiler, - computation, client)); + TF_RETURN_IF_ERROR( + ConvertGraphToXla(std::move(graph), config, client, computation)); return absl::OkStatus(); } diff --git a/tensorflow/compiler/tf2xla/tf2xla.h b/tensorflow/compiler/tf2xla/tf2xla.h index d38b563cc7a2ac..095ad49afc6b66 100644 --- a/tensorflow/compiler/tf2xla/tf2xla.h +++ b/tensorflow/compiler/tf2xla/tf2xla.h @@ -18,15 +18,11 @@ limitations under the License. #include "absl/strings/string_view.h" #include "tensorflow/compiler/tf2xla/tf2xla.pb.h" +#include "xla/client/client.h" #include "xla/hlo/builder/xla_computation.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/platform/status.h" -namespace xla { -class PjRtClient; -class PjRtCompiler; -} // namespace xla - namespace tensorflow { // Converts a tensorflow::GraphDef into an xla::XlaComputation. The given @@ -34,14 +30,12 @@ namespace tensorflow { // fetches. Each feed is a positional input argument for the generated // computation, while each fetch is a positional output argument. // -// The computation is built in the context of the given `compiler` (and optional -// `client`), which may subsequently be used to compile or execute the -// computation. +// The computation is built in the context of the given `client`, which may +// subsequently be used to compile or execute the computation. absl::Status ConvertGraphDefToXla(GraphDef graph_def, const tf2xla::Config& config, - xla::PjRtCompiler* compiler, - xla::XlaComputation* computation, - xla::PjRtClient* client = nullptr); + xla::Client* client, + xla::XlaComputation* computation); // Similar to ConvertGraphDefToXla, but uses MLIR and handle debug information. // diff --git a/tensorflow/compiler/tf2xla/tf2xla_test.cc b/tensorflow/compiler/tf2xla/tf2xla_test.cc index b609deb78d3eb5..72bd28f2b47a8c 100644 --- a/tensorflow/compiler/tf2xla/tf2xla_test.cc +++ b/tensorflow/compiler/tf2xla/tf2xla_test.cc @@ -15,26 +15,24 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/tf2xla.h" -#include #include -#include "absl/status/status.h" -#include "absl/strings/string_view.h" #include "tensorflow/compiler/tf2xla/tf2xla.pb.h" +#include "xla/client/client_library.h" +#include "xla/client/local_client.h" #include "xla/hlo/builder/xla_computation.h" +#include "xla/literal.h" #include "xla/literal_util.h" -#include "xla/pjrt/cpu/cpu_client.h" -#include "xla/pjrt/pjrt_compiler.h" -#include "xla/pjrt/pjrt_executable.h" -#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h" -#include "xla/tsl/lib/core/status_test_util.h" -#include "xla/tsl/platform/statusor.h" #include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/attr_value_util.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/framework/node_def.pb.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_shape.pb.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/platform/statusor.h" +#include "tensorflow/core/platform/stringpiece.h" #include "tensorflow/core/platform/test.h" #include "tsl/platform/tensor_float_32_utils.h" @@ -111,46 +109,35 @@ tf2xla::Config SumConfig() { return config; } -// NOLINTBEGIN - TEST(ConvertGraphDefToXla, Sum) { GraphDef graph_def = SumGraph(); tf2xla::Config config = SumConfig(); - TF_ASSERT_OK_AND_ASSIGN(auto client, - xla::GetPjRtCpuClient(xla::CpuClientOptions())); - TF_ASSERT_OK_AND_ASSIGN(auto compiler, - xla::GetDefaultPjRtCompiler(xla::CpuName())); - + xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); xla::XlaComputation computation; - TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, compiler, &computation, - client.get())); + TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation)); // Set up arguments. auto x_literal = xla::LiteralUtil::CreateR0(10); auto y_literal = xla::LiteralUtil::CreateR0(32); - TF_ASSERT_OK_AND_ASSIGN( - auto x_buf, - client->BufferFromHostLiteral(x_literal, client->memory_spaces()[0])); - TF_ASSERT_OK_AND_ASSIGN( - auto y_buf, - client->BufferFromHostLiteral(y_literal, client->memory_spaces()[0])); + auto x_global_or = client->TransferToServer(x_literal); + auto y_global_or = client->TransferToServer(y_literal); + TF_EXPECT_OK(x_global_or.status()); + TF_EXPECT_OK(y_global_or.status()); + std::unique_ptr x_global = std::move(x_global_or.value()); + std::unique_ptr y_global = std::move(y_global_or.value()); // Execute and check result. - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_exec, - client->CompileAndLoad(computation, xla::CompileOptions{})); - xla::ExecuteOptions execute_options; - TF_ASSERT_OK_AND_ASSIGN( - auto results, - loaded_exec->Execute({{x_buf.get(), y_buf.get()}}, execute_options)); - TF_ASSERT_OK_AND_ASSIGN(auto result, results[0][0]->ToLiteralSync()); - EXPECT_EQ(42, result->Get({})); + auto result_or = + client->ExecuteAndTransfer(computation, {x_global.get(), y_global.get()}); + TF_EXPECT_OK(result_or.status()); + xla::Literal result = std::move(result_or.value()); + EXPECT_EQ("(\ns32[] 42\n)", result.ToString()); config.mutable_feed(0)->mutable_id()->set_output_index( 123); /* invalid output_index */ - EXPECT_TRUE(absl::IsInvalidArgument(ConvertGraphDefToXla( - graph_def, config, compiler, &computation, client.get()))); + EXPECT_TRUE(absl::IsInvalidArgument( + ConvertGraphDefToXla(graph_def, config, client, &computation))); } GraphDef EinsumGraph(DataType dtype = DT_FLOAT) { @@ -193,10 +180,9 @@ TEST(ConvertGraphDefToXla, EinsumIsConvertedToDotWithDefaultPrecision) { GraphDef graph_def = EinsumGraph(); tf2xla::Config config = EinsumConfig(); - TF_ASSERT_OK_AND_ASSIGN(auto compiler, - xla::GetDefaultPjRtCompiler(xla::CpuName())); + xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); xla::XlaComputation computation; - TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, compiler, &computation)); + TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation)); int num_dots = 0; const xla::HloModuleProto& module_proto = computation.proto(); @@ -223,10 +209,9 @@ TEST_F(ConvertGraphDefToXlaWithTF32Disabled, GraphDef graph_def = EinsumGraph(); tf2xla::Config config = EinsumConfig(); - TF_ASSERT_OK_AND_ASSIGN(auto compiler, - xla::GetDefaultPjRtCompiler(xla::CpuName())); + xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); xla::XlaComputation computation; - TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, compiler, &computation)); + TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation)); int num_dots = 0; const xla::HloModuleProto& module_proto = computation.proto(); @@ -253,10 +238,9 @@ TEST_F(ConvertGraphDefToXlaWithTF32Disabled, GraphDef graph_def = EinsumGraph(DT_BFLOAT16); tf2xla::Config config = EinsumConfig(); - TF_ASSERT_OK_AND_ASSIGN(auto compiler, - xla::GetDefaultPjRtCompiler(xla::CpuName())); + xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); xla::XlaComputation computation; - TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, compiler, &computation)); + TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation)); int num_dots = 0; const xla::HloModuleProto& module_proto = computation.proto(); @@ -316,10 +300,9 @@ TEST(ConvertGraphDefToXla, Conv2DIsConvertedToConvolutionWithDefaultPrecision) { GraphDef graph_def = Conv2DGraph(); tf2xla::Config config = Conv2DConfig(); - TF_ASSERT_OK_AND_ASSIGN(auto compiler, - xla::GetDefaultPjRtCompiler(xla::CpuName())); + xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); xla::XlaComputation computation; - TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, compiler, &computation)); + TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation)); int num_convolutions = 0; const xla::HloModuleProto& module_proto = computation.proto(); @@ -346,10 +329,9 @@ TEST_F(ConvertGraphDefToXlaWithTF32Disabled, GraphDef graph_def = Conv2DGraph(); tf2xla::Config config = Conv2DConfig(); - TF_ASSERT_OK_AND_ASSIGN(auto compiler, - xla::GetDefaultPjRtCompiler(xla::CpuName())); + xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); xla::XlaComputation computation; - TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, compiler, &computation)); + TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation)); int num_convolutions = 0; const xla::HloModuleProto& module_proto = computation.proto(); @@ -380,42 +362,31 @@ TEST(ConvertGraphDefToXla, SumWithUnusedArgument) { (*unused->mutable_attr())["dtype"] = TypeAttrValue(DT_INT32); config.add_feed()->mutable_id()->set_node_name("unused"); - TF_ASSERT_OK_AND_ASSIGN(auto client, - xla::GetPjRtCpuClient(xla::CpuClientOptions())); - TF_ASSERT_OK_AND_ASSIGN(auto compiler, - xla::GetDefaultPjRtCompiler(xla::CpuName())); - + xla::LocalClient* client = xla::ClientLibrary::LocalClientOrDie(); xla::XlaComputation computation; - TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, compiler, &computation, - client.get())); + TF_EXPECT_OK(ConvertGraphDefToXla(graph_def, config, client, &computation)); // Set up arguments. auto x_literal = xla::LiteralUtil::CreateR0(10); auto y_literal = xla::LiteralUtil::CreateR0(32); - TF_ASSERT_OK_AND_ASSIGN( - auto x_buf, - client->BufferFromHostLiteral(x_literal, client->memory_spaces()[0])); - TF_ASSERT_OK_AND_ASSIGN( - auto y_buf, - client->BufferFromHostLiteral(y_literal, client->memory_spaces()[0])); - TF_ASSERT_OK_AND_ASSIGN( - auto unused_buf, - client->BufferFromHostLiteral(y_literal, client->memory_spaces()[0])); + auto x_global_or = client->TransferToServer(x_literal); + auto y_global_or = client->TransferToServer(y_literal); + auto unused_global_or = client->TransferToServer(y_literal); + TF_EXPECT_OK(x_global_or.status()); + TF_EXPECT_OK(y_global_or.status()); + TF_EXPECT_OK(unused_global_or.status()); + std::unique_ptr x_global = std::move(x_global_or.value()); + std::unique_ptr y_global = std::move(y_global_or.value()); + std::unique_ptr unused_global = + std::move(unused_global_or.value()); // Execute and check result. - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_exec, - client->CompileAndLoad(computation, xla::CompileOptions{})); - xla::ExecuteOptions execute_options; - TF_ASSERT_OK_AND_ASSIGN( - auto results, - loaded_exec->Execute({{x_buf.get(), y_buf.get(), unused_buf.get()}}, - execute_options)); - TF_ASSERT_OK_AND_ASSIGN(auto result, results[0][0]->ToLiteralSync()); - EXPECT_EQ(42, result->Get({})); + auto result_or = client->ExecuteAndTransfer( + computation, {x_global.get(), y_global.get(), unused_global.get()}); + TF_EXPECT_OK(result_or.status()); + xla::Literal result = std::move(result_or.value()); + EXPECT_EQ("(\ns32[] 42\n)", result.ToString()); } -// NOLINTEND - } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_compiler.h b/tensorflow/compiler/tf2xla/xla_compiler.h index 446381a14ede76..b9abd5006a958a 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.h +++ b/tensorflow/compiler/tf2xla/xla_compiler.h @@ -46,8 +46,6 @@ limitations under the License. #include "xla/hlo/builder/xla_builder.h" #include "xla/hlo/builder/xla_computation.h" #include "xla/hlo/ir/hlo_sharding.h" -#include "xla/pjrt/pjrt_client.h" -#include "xla/pjrt/pjrt_compiler.h" #include "xla/shape.h" #include "xla/status_macros.h" #include "xla/stream_executor/device_memory_allocator.h" @@ -187,8 +185,7 @@ class XlaCompiler { // -1 indicates the default device should be used. int device_ordinal = -1; - xla::PjRtCompiler* compiler = nullptr; - xla::PjRtClient* client = nullptr; + xla::Client* client = nullptr; // Function library in which to find function definitions. Must be non-null. const FunctionLibraryDefinition* flib_def = nullptr; @@ -345,8 +342,7 @@ class XlaCompiler { const std::string& host_compute_name, xla::XlaOp handle); const Options& options() const { return options_; } - xla::PjRtCompiler* compiler() const { return options_.compiler; } - xla::PjRtClient* client() const { return options_.client; } + xla::Client* client() const { return options_.client; } FunctionLibraryRuntime* flib_runtime() const { return flib_runtime_; } void PushNodeTokenMapping(); diff --git a/tensorflow/compiler/tf2xla/xla_compiler_test.cc b/tensorflow/compiler/tf2xla/xla_compiler_test.cc index c6982f86eab099..32dfd44f5fba1f 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler_test.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler_test.cc @@ -21,7 +21,6 @@ limitations under the License. #include #include #include -#include #include #include @@ -38,9 +37,12 @@ limitations under the License. #include "tensorflow/cc/ops/data_flow_ops.h" #include "tensorflow/cc/ops/function_ops.h" #include "tensorflow/cc/ops/functional_ops.h" +#include "tensorflow/cc/ops/list_ops.h" #include "tensorflow/cc/ops/math_ops.h" #include "tensorflow/cc/ops/resource_variable_ops.h" +#include "tensorflow/cc/ops/standard_ops.h" #include "tensorflow/compiler/tf2xla/layout_util.h" +#include "tensorflow/compiler/tf2xla/literal_util.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/side_effect_util.h" #include "tensorflow/compiler/tf2xla/type_util.h" @@ -49,31 +51,30 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/tf2xla/xla_resource.h" #include "xla/array.h" +#include "xla/client/client.h" +#include "xla/client/client_library.h" +#include "xla/client/local_client.h" #include "xla/debug_options_flags.h" #include "xla/hlo/builder/xla_builder.h" -#include "xla/hlo/builder/xla_computation.h" #include "xla/hlo/ir/hlo_module.h" -#include "xla/hlo/ir/hlo_print_options.h" #include "xla/hlo/ir/hlo_sharding.h" #include "xla/hlo/testlib/filecheck.h" #include "xla/layout_util.h" #include "xla/literal.h" #include "xla/literal_util.h" -#include "xla/pjrt/cpu/cpu_client.h" -#include "xla/pjrt/pjrt_client.h" -#include "xla/pjrt/pjrt_compiler.h" -#include "xla/pjrt/pjrt_executable.h" -#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h" -#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h" #include "xla/service/hlo.pb.h" #include "xla/service/hlo_module_util.h" +#include "xla/service/hlo_proto_util.h" +#include "xla/service/service.h" #include "xla/shape.h" #include "xla/shape_util.h" +#include "xla/status_macros.h" #include "xla/tests/literal_test_util.h" #include "xla/tsl/lib/core/status_test_util.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/statusor.h" #include "xla/xla_data.pb.h" +#include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/common_runtime/graph_constructor.h" #include "tensorflow/core/framework/common_shape_fns.h" #include "tensorflow/core/framework/fake_input.h" @@ -97,17 +98,19 @@ limitations under the License. #include "tensorflow/core/graph/graph.h" #include "tensorflow/core/graph/node_builder.h" #include "tensorflow/core/kernels/ops_testutil.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/platform/errors.h" #include "tensorflow/core/platform/test.h" +#include "tensorflow/core/platform/types.h" +#include "tensorflow/core/public/version.h" +#include "tsl/platform/statusor.h" namespace tensorflow { class XlaCompilerTest : public ::testing::Test { protected: void SetUp() override { - TF_ASSERT_OK_AND_ASSIGN(pjrt_client_, - xla::GetXlaPjrtCpuClient(xla::CpuClientOptions())); - TF_ASSERT_OK_AND_ASSIGN(pjrt_compiler_, - xla::GetDefaultPjRtCompiler(xla::CpuName())); + client_ = xla::ClientLibrary::LocalClientOrDie(); XlaOpRegistry::RegisterCompilationKernels(); @@ -118,8 +121,7 @@ class XlaCompilerTest : public ::testing::Test { XlaCompiler::Options DefaultOptions() { XlaCompiler::Options options; options.device_type = DeviceType(DEVICE_CPU_XLA_JIT); - options.compiler = pjrt_compiler_; - options.client = pjrt_client_.get(); + options.client = client_; options.flib_def = flib_def_.get(); return options; } @@ -128,8 +130,7 @@ class XlaCompilerTest : public ::testing::Test { return compiler->local_flib_def_.get(); } - std::unique_ptr pjrt_client_; - xla::PjRtCompiler* pjrt_compiler_; + xla::Client* client_; std::unique_ptr flib_def_; }; @@ -236,11 +237,7 @@ TEST_F(XlaCompilerTest, EmptyReturnValues) { std::move(graph), /*args=*/{}, &result)); - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_exec, - pjrt_client_->CompileAndLoad(*result.computation, xla::CompileOptions{})); - xla::ExecuteOptions execute_options; - TF_ASSERT_OK(loaded_exec->Execute({{}}, execute_options).status()); + TF_ASSERT_OK(client_->Execute(*result.computation, {}).status()); } // Tests compilation and execution of a graph that adds two tensors. @@ -273,28 +270,20 @@ TEST_F(XlaCompilerTest, Simple) { // Tests that the generated computation works. xla::Literal param0_literal = xla::LiteralUtil::CreateR1({7, 42}); xla::Literal param1_literal = xla::LiteralUtil::CreateR1({-3, 101}); - TF_ASSERT_OK_AND_ASSIGN( - auto param0_buffer, - pjrt_client_->BufferFromHostLiteral(param0_literal, - pjrt_client_->memory_spaces()[0])); - TF_ASSERT_OK_AND_ASSIGN( - auto param1_buffer, - pjrt_client_->BufferFromHostLiteral(param1_literal, - pjrt_client_->memory_spaces()[0])); - - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_exec, - pjrt_client_->CompileAndLoad(*result.computation, xla::CompileOptions{})); - xla::ExecuteOptions execute_options; - TF_ASSERT_OK_AND_ASSIGN( - auto actual, - loaded_exec->Execute({{param0_buffer.get(), param1_buffer.get()}}, - execute_options)); - - TF_ASSERT_OK_AND_ASSIGN(auto actual_literal, actual[0][0]->ToLiteralSync()); + std::unique_ptr param0_data = + client_->TransferToServer(param0_literal).value(); + std::unique_ptr param1_data = + client_->TransferToServer(param1_literal).value(); + + std::unique_ptr actual = + client_ + ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) + .value(); + xla::Literal actual_literal = client_->Transfer(*actual).value(); xla::Literal expected0 = xla::LiteralUtil::CreateR1({4, 143}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected0, *actual_literal)); + xla::Literal expected_literal = xla::LiteralUtil::MakeTuple({&expected0}); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); } absl::StatusOr> LoadModuleFromHloProto( @@ -379,27 +368,18 @@ TEST_F(XlaCompilerTest, OutOfOrderGraph) { // Tests that the generated computation works. xla::Literal param0_literal = xla::LiteralUtil::CreateR1({7, 42}); xla::Literal param1_literal = xla::LiteralUtil::CreateR1({-3, 101}); - TF_ASSERT_OK_AND_ASSIGN( - auto param0_buffer, - pjrt_client_->BufferFromHostLiteral(param0_literal, - pjrt_client_->memory_spaces()[0])); - TF_ASSERT_OK_AND_ASSIGN( - auto param1_buffer, - pjrt_client_->BufferFromHostLiteral(param1_literal, - pjrt_client_->memory_spaces()[0])); - - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_exec, - pjrt_client_->CompileAndLoad(*result.computation, xla::CompileOptions{})); - xla::ExecuteOptions execute_options; - TF_ASSERT_OK_AND_ASSIGN( - auto actual, - loaded_exec->Execute({{param0_buffer.get(), param1_buffer.get()}}, - execute_options)); - - TF_ASSERT_OK_AND_ASSIGN(auto actual_literal, actual[0][0]->ToLiteralSync()); - - EXPECT_TRUE(xla::LiteralTestUtil::Equal(param0_literal, *actual_literal)); + std::unique_ptr param0_data = + client_->TransferToServer(param0_literal).value(); + std::unique_ptr param1_data = + client_->TransferToServer(param1_literal).value(); + + std::unique_ptr actual = + client_ + ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) + .value(); + xla::Literal actual_literal = client_->Transfer(*actual).value(); + + EXPECT_TRUE(xla::LiteralTestUtil::Equal(param0_literal, actual_literal)); } // Tests that the compiler can correctly propagate the layout assigned by @@ -440,7 +420,7 @@ TEST_F(XlaCompilerTest, HonorShapeRepresentationFnForUnwrittenResource) { args, &result)); xla::Shape transposed = xla::ShapeUtil::MakeShapeWithDenseLayout(xla::S32, {2, 3}, {0, 1}); - // Check that the return shapes are correctly transposed. + // Check that the return shapes are correctly tranposed. EXPECT_EQ(result.xla_output_shape, xla::ShapeUtil::MakeTupleShape({transposed})); } @@ -542,9 +522,7 @@ TEST_F(XlaCompilerTest, HonorShapeRepresentationFnForRetVal) { // Check that the return shapes are correctly tranposed. EXPECT_EQ(result.xla_output_shape, xla::ShapeUtil::MakeTupleShape({transposed, transposed})); - TF_ASSERT_OK_AND_ASSIGN(auto program_shape, - result.computation->GetProgramShape()); - EXPECT_EQ(program_shape.result(), + EXPECT_EQ(result.computation->GetProgramShape().value().result(), xla::ShapeUtil::MakeTupleShape({transposed, transposed})); } @@ -584,7 +562,7 @@ TEST_F(XlaCompilerTest, TransposeVariables) { std::move(graph), args, &result)); xla::Shape transposed = xla::ShapeUtil::MakeShapeWithDenseLayout(xla::S32, {2, 3}, {1, 0}); - // Check that the return shapes are correctly transposed. + // Check that the return shapes are correctly tranposed. EXPECT_EQ(result.xla_output_shape, xla::ShapeUtil::MakeTupleShape({transposed, transposed})); } @@ -604,7 +582,7 @@ TEST_F(XlaCompilerTest, UnrankedFakeParam) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "compile", std::move(graph), {}, &result)); - // Check that the return shapes are correctly transposed. + // Check that the return shapes are correctly tranposed. EXPECT_EQ(result.xla_output_shape, xla::ShapeUtil::MakeTupleShape( {xla::ShapeUtil::MakeShape(xla::S32, {0})})); @@ -733,28 +711,18 @@ TEST_F(XlaCompilerTest, ConstantOutputs) { // Tests that the generated computation works. xla::Literal param0_literal = xla::LiteralUtil::CreateR1({7, 42}); - TF_ASSERT_OK_AND_ASSIGN( - auto param0_buffer, - pjrt_client_->BufferFromHostLiteral(param0_literal, - pjrt_client_->memory_spaces()[0])); - - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_exec, pjrt_client_->CompileAndLoad(*result.computation, - xla::CompileOptions{})); - xla::ExecuteOptions execute_options; - TF_ASSERT_OK_AND_ASSIGN( - auto actual, - loaded_exec->Execute({{param0_buffer.get()}}, execute_options)); - - TF_ASSERT_OK_AND_ASSIGN(auto actual_literal0, - actual[0][0]->ToLiteralSync()); - TF_ASSERT_OK_AND_ASSIGN(auto actual_literal1, - actual[0][1]->ToLiteralSync()); + std::unique_ptr param0_data = + client_->TransferToServer(param0_literal).value(); + + std::unique_ptr actual = + client_->Execute(*result.computation, {param0_data.get()}).value(); + xla::Literal actual_literal = client_->Transfer(*actual).value(); xla::Literal expected0 = xla::LiteralUtil::CreateR0(7); xla::Literal expected1 = xla::LiteralUtil::CreateR1({-7, -42}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected0, *actual_literal0)); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected1, *actual_literal1)); + xla::Literal expected = + xla::LiteralUtil::MakeTuple({&expected0, &expected1}); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected, actual_literal)); } } @@ -969,58 +937,25 @@ TEST_F(XlaCompilerTest, CanPassTensorArraysToAndFromComputation) { update.tensor_array_gradients_accessed); // Tests that the generated computation works. - // PjRtCpuClient does not support nested tuples, so wrap the computation to - // flatten the input and output tuples. - xla::XlaBuilder b("flatten_wrapper"); - xla::XlaOp p0 = - xla::Parameter(&b, 0, xla::ShapeUtil::MakeShape(xla::S32, {2}), "base"); - xla::XlaOp p1 = - xla::Parameter(&b, 1, xla::ShapeUtil::MakeShape(xla::S32, {2}), "grad2"); - xla::XlaOp call = - xla::Call(&b, *result.computation, {xla::Tuple(&b, {p0, p1})}); - xla::XlaOp read_op = xla::GetTupleElement(call, 0); - xla::XlaOp resource = xla::GetTupleElement(call, 1); - xla::XlaOp out_base = xla::GetTupleElement(resource, 0); - xla::XlaOp out_grad1 = xla::GetTupleElement(resource, 1); - xla::XlaOp out_grad2 = xla::GetTupleElement(resource, 2); - xla::Tuple(&b, {read_op, out_base, out_grad1, out_grad2}); - TF_ASSERT_OK_AND_ASSIGN(xla::XlaComputation flattened_comp, b.Build()); - xla::Literal input_base = xla::LiteralUtil::CreateR1({7, 42}); xla::Literal input_grad2 = xla::LiteralUtil::CreateR1({-3, 101}); - TF_ASSERT_OK_AND_ASSIGN(auto input_base_buffer, - pjrt_client_->BufferFromHostLiteral( - input_base, pjrt_client_->memory_spaces()[0])); - TF_ASSERT_OK_AND_ASSIGN(auto input_grad2_buffer, - pjrt_client_->BufferFromHostLiteral( - input_grad2, pjrt_client_->memory_spaces()[0])); - - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_exec, - pjrt_client_->CompileAndLoad(flattened_comp, xla::CompileOptions{})); - xla::ExecuteOptions execute_options; - TF_ASSERT_OK_AND_ASSIGN( - auto actual, loaded_exec->Execute( - {{input_base_buffer.get(), input_grad2_buffer.get()}}, - execute_options)); - - TF_ASSERT_OK_AND_ASSIGN(auto actual_output_read, - actual[0][0]->ToLiteralSync()); - TF_ASSERT_OK_AND_ASSIGN(auto actual_output_base, - actual[0][1]->ToLiteralSync()); - TF_ASSERT_OK_AND_ASSIGN(auto actual_output_grad1, - actual[0][2]->ToLiteralSync()); - TF_ASSERT_OK_AND_ASSIGN(auto actual_output_grad2, - actual[0][3]->ToLiteralSync()); + xla::Literal input = xla::LiteralUtil::MakeTuple({&input_base, &input_grad2}); + std::unique_ptr param0_data = + client_->TransferToServer(input).value(); + + std::unique_ptr actual = + client_->Execute(*result.computation, {param0_data.get()}).value(); + xla::Literal actual_literal = client_->Transfer(*actual).value(); xla::Literal output_read = xla::LiteralUtil::CreateR0(42); xla::Literal output_base = xla::LiteralUtil::CreateR1({7, 42}); xla::Literal output_grad1 = xla::LiteralUtil::CreateR1({0, 1}); xla::Literal output_grad2 = xla::LiteralUtil::CreateR1({-3, 101}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(output_read, *actual_output_read)); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(output_base, *actual_output_base)); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(output_grad1, *actual_output_grad1)); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(output_grad2, *actual_output_grad2)); + xla::Literal output_resource = + xla::LiteralUtil::MakeTuple({&output_base, &output_grad1, &output_grad2}); + xla::Literal expected_literal = + xla::LiteralUtil::MakeTuple({&output_read, &output_resource}); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); } // Tests compilation and execution of a graph that adds two tensors. @@ -1120,8 +1055,8 @@ FunctionDef FillFn() { TEST_F(XlaCompilerTest, FunctionCallWithConstants) { // Certain operations in a function, "Fill" for example, requires the // operator's argument to be a compile-time constant instead of a parameter. - // This case tests if XlaCompiler can handle such operators inside function - // calls. + // This testcase tests if XlaCompiler can handle such operators inside + // function calls. XlaCompiler compiler(DefaultOptions()); FunctionDefLibrary flib; @@ -1252,32 +1187,25 @@ TEST_F(XlaCompilerTest, SliceWithDynamicBegins) { } void RunAndCheckVariablesComputation( - xla::PjRtClient* client, const XlaCompiler::CompilationResult& result) { + xla::Client* client, const XlaCompiler::CompilationResult& result) { xla::Literal param0_literal = xla::LiteralUtil::CreateR1({7, 42}); xla::Literal param1_literal = xla::LiteralUtil::CreateR1({-3, 101}); - TF_ASSERT_OK_AND_ASSIGN(auto param0_buffer, - client->BufferFromHostLiteral( - param0_literal, client->memory_spaces()[0])); - TF_ASSERT_OK_AND_ASSIGN(auto param1_buffer, - client->BufferFromHostLiteral( - param1_literal, client->memory_spaces()[0])); - - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_exec, - client->CompileAndLoad(*result.computation, xla::CompileOptions{})); - xla::ExecuteOptions execute_options; - TF_ASSERT_OK_AND_ASSIGN( - auto actual, - loaded_exec->Execute({{param0_buffer.get(), param1_buffer.get()}}, - execute_options)); - - TF_ASSERT_OK_AND_ASSIGN(auto actual_literal0, actual[0][0]->ToLiteralSync()); - TF_ASSERT_OK_AND_ASSIGN(auto actual_literal1, actual[0][1]->ToLiteralSync()); + std::unique_ptr param0_data = + client->TransferToServer(param0_literal).value(); + std::unique_ptr param1_data = + client->TransferToServer(param1_literal).value(); + + std::unique_ptr actual = + client + ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) + .value(); + xla::Literal actual_literal = client->Transfer(*actual).value(); xla::Literal expected0 = xla::LiteralUtil::CreateR1({5, 144}); xla::Literal expected1 = xla::LiteralUtil::CreateR1({4, 143}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected0, *actual_literal0)); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected1, *actual_literal1)); + xla::Literal expected_literal = + xla::LiteralUtil::MakeTuple({&expected0, &expected1}); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); } // Tests a simple graph that reads and writes a variable. @@ -1314,7 +1242,7 @@ TEST_F(XlaCompilerTest, Variables) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", std::move(graph), args, &result)); - RunAndCheckVariablesComputation(pjrt_client_.get(), result); + RunAndCheckVariablesComputation(client_, result); } TEST_F(XlaCompilerTest, ResultLayoutSingle) { @@ -1429,20 +1357,15 @@ TEST_F(XlaCompilerTest, ReturnResourceHandleOnly) { // Tests that the generated computation works. xla::Literal param1_literal = xla::LiteralUtil::CreateR1({-3, 101}); - TF_ASSERT_OK_AND_ASSIGN( - auto param1_buffer, - pjrt_client_->BufferFromHostLiteral(param1_literal, - pjrt_client_->memory_spaces()[0])); - - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_exec, - pjrt_client_->CompileAndLoad(*result.computation, xla::CompileOptions{})); - xla::ExecuteOptions execute_options; - TF_ASSERT_OK_AND_ASSIGN( - auto actual, - loaded_exec->Execute({{param1_buffer.get()}}, execute_options)); - - EXPECT_EQ(actual[0].size(), 0); + std::unique_ptr param1_data = + client_->TransferToServer(param1_literal).value(); + + std::unique_ptr actual = + client_->Execute(*result.computation, {param1_data.get()}).value(); + xla::Literal actual_literal = client_->Transfer(*actual).value(); + + xla::Literal expected_literal = xla::LiteralUtil::MakeTuple({}); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); } TEST_F(XlaCompilerTest, ReturnResourceHandle) { @@ -1480,7 +1403,7 @@ TEST_F(XlaCompilerTest, ReturnResourceHandle) { XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", std::move(graph), args, &result)); - RunAndCheckVariablesComputation(pjrt_client_.get(), result); + RunAndCheckVariablesComputation(client_, result); } absl::StatusOr> BuildTestGraph() { @@ -1535,52 +1458,44 @@ TEST_F(XlaCompilerTest, VariableRepresentationShapeFunction) { TF_ASSERT_OK(compiler.CompileGraph(compile_options, "add", std::move(graph), args, &result)); - TF_ASSERT_OK_AND_ASSIGN(xla::ProgramShape program_shape, - result.computation->GetProgramShape()); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr program_shape, + client_->GetComputationShape(*result.computation)); - ASSERT_EQ(program_shape.parameters_size(), 2); + ASSERT_EQ(program_shape->parameters_size(), 2); EXPECT_TRUE( - xla::ShapeUtil::Compatible(program_shape.parameters(0), + xla::ShapeUtil::Compatible(program_shape->parameters(0), xla::ShapeUtil::MakeShape(xla::S32, {2, 2}))); EXPECT_TRUE(xla::ShapeUtil::Compatible( - program_shape.parameters(1), xla::ShapeUtil::MakeShape(xla::S32, {4}))); + program_shape->parameters(1), xla::ShapeUtil::MakeShape(xla::S32, {4}))); EXPECT_TRUE(xla::ShapeUtil::Compatible( - program_shape.result(), xla::ShapeUtil::MakeTupleShape( - {xla::ShapeUtil::MakeShape(xla::S32, {2, 2}), - xla::ShapeUtil::MakeShape(xla::S32, {4})}))); + program_shape->result(), + xla::ShapeUtil::MakeTupleShape( + {xla::ShapeUtil::MakeShape(xla::S32, {2, 2}), + xla::ShapeUtil::MakeShape(xla::S32, {4})}))); // Tests that the generated computation works. xla::Literal param0_literal = xla::LiteralUtil::CreateR2({{4, 55}, {1, -3}}); xla::Literal param1_literal = xla::LiteralUtil::CreateR1({22, 11, 33, 404}); - TF_ASSERT_OK_AND_ASSIGN( - auto param0_buffer, - pjrt_client_->BufferFromHostLiteral(param0_literal, - pjrt_client_->memory_spaces()[0])); - TF_ASSERT_OK_AND_ASSIGN( - auto param1_buffer, - pjrt_client_->BufferFromHostLiteral(param1_literal, - pjrt_client_->memory_spaces()[0])); - - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_exec, - pjrt_client_->CompileAndLoad(*result.computation, xla::CompileOptions{})); - xla::ExecuteOptions execute_options; - TF_ASSERT_OK_AND_ASSIGN( - auto actual, - loaded_exec->Execute({{param0_buffer.get(), param1_buffer.get()}}, - execute_options)); - - TF_ASSERT_OK_AND_ASSIGN(auto actual_literal0, actual[0][0]->ToLiteralSync()); - TF_ASSERT_OK_AND_ASSIGN(auto actual_literal1, actual[0][1]->ToLiteralSync()); + std::unique_ptr param0_data = + client_->TransferToServer(param0_literal).value(); + std::unique_ptr param1_data = + client_->TransferToServer(param1_literal).value(); + + std::unique_ptr actual = + client_ + ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) + .value(); + xla::Literal actual_literal = client_->Transfer(*actual).value(); xla::Literal expected0 = xla::LiteralUtil::CreateR2({{27, 67}, {35, 402}}); xla::Literal expected1 = xla::LiteralUtil::CreateR1({26, 66, 34, 401}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected0, *actual_literal0)); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected1, *actual_literal1)); + xla::Literal expected_literal = + xla::LiteralUtil::MakeTuple({&expected0, &expected1}); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); } TEST_F(XlaCompilerTest, ArgRetvalShapeRepresentationFunction) { @@ -1617,51 +1532,43 @@ TEST_F(XlaCompilerTest, ArgRetvalShapeRepresentationFunction) { TF_ASSERT_OK(compiler.CompileGraph(compile_options, "add", std::move(graph), args, &result)); - TF_ASSERT_OK_AND_ASSIGN(xla::ProgramShape program_shape, - result.computation->GetProgramShape()); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr program_shape, + client_->GetComputationShape(*result.computation)); - ASSERT_EQ(program_shape.parameters_size(), 2); + ASSERT_EQ(program_shape->parameters_size(), 2); EXPECT_TRUE(xla::ShapeUtil::Compatible( - program_shape.parameters(0), xla::ShapeUtil::MakeShape(xla::S32, {4}))); + program_shape->parameters(0), xla::ShapeUtil::MakeShape(xla::S32, {4}))); EXPECT_TRUE(xla::ShapeUtil::Compatible( - program_shape.parameters(1), xla::ShapeUtil::MakeShape(xla::S32, {4}))); + program_shape->parameters(1), xla::ShapeUtil::MakeShape(xla::S32, {4}))); EXPECT_TRUE(xla::ShapeUtil::Compatible( - program_shape.result(), xla::ShapeUtil::MakeTupleShape( - {xla::ShapeUtil::MakeShape(xla::S32, {4}), - xla::ShapeUtil::MakeShape(xla::S32, {4})}))); + program_shape->result(), + xla::ShapeUtil::MakeTupleShape( + {xla::ShapeUtil::MakeShape(xla::S32, {4}), + xla::ShapeUtil::MakeShape(xla::S32, {4})}))); // Tests that the generated computation works. xla::Literal param0_literal = xla::LiteralUtil::CreateR1({4, 55, 1, -3}); xla::Literal param1_literal = xla::LiteralUtil::CreateR1({22, 11, 33, 404}); - TF_ASSERT_OK_AND_ASSIGN( - auto param0_buffer, - pjrt_client_->BufferFromHostLiteral(param0_literal, - pjrt_client_->memory_spaces()[0])); - TF_ASSERT_OK_AND_ASSIGN( - auto param1_buffer, - pjrt_client_->BufferFromHostLiteral(param1_literal, - pjrt_client_->memory_spaces()[0])); - - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_exec, - pjrt_client_->CompileAndLoad(*result.computation, xla::CompileOptions{})); - xla::ExecuteOptions execute_options; - TF_ASSERT_OK_AND_ASSIGN( - auto actual, - loaded_exec->Execute({{param0_buffer.get(), param1_buffer.get()}}, - execute_options)); - - TF_ASSERT_OK_AND_ASSIGN(auto actual_literal0, actual[0][0]->ToLiteralSync()); - TF_ASSERT_OK_AND_ASSIGN(auto actual_literal1, actual[0][1]->ToLiteralSync()); + std::unique_ptr param0_data = + client_->TransferToServer(param0_literal).value(); + std::unique_ptr param1_data = + client_->TransferToServer(param1_literal).value(); + + std::unique_ptr actual = + client_ + ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) + .value(); + xla::Literal actual_literal = client_->Transfer(*actual).value(); xla::Literal expected0 = xla::LiteralUtil::CreateR1({27, 67, 35, 402}); xla::Literal expected1 = xla::LiteralUtil::CreateR1({26, 66, 34, 401}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected0, *actual_literal0)); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected1, *actual_literal1)); + xla::Literal expected_literal = + xla::LiteralUtil::MakeTuple({&expected0, &expected1}); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); } // Tests a graph which has a function with an invalid op. @@ -1844,9 +1751,9 @@ TEST_F(XlaCompilerTest, OpsWithTensorListInput) { { Scope scope = Scope::NewRootScope().ExitOnError(); std::unique_ptr graph(new Graph(OpRegistry::Global())); - (void)ops::_Arg(scope.WithOpName("arg"), DT_VARIANT, 0); + ops::_Arg(scope.WithOpName("arg"), DT_VARIANT, 0); auto result = ops::Const(scope, {true}, {}); - (void)ops::_Retval(scope.WithOpName("ret"), result, 0); + ops::_Retval(scope.WithOpName("ret"), result, 0); TF_ASSERT_OK(scope.ToGraph(graph.get())); FunctionDef fdef; TF_ASSERT_OK(GraphToFunctionDef(*graph, "cond", &fdef)); @@ -1857,7 +1764,7 @@ TEST_F(XlaCompilerTest, OpsWithTensorListInput) { Scope scope = Scope::NewRootScope().ExitOnError(); std::unique_ptr graph(new Graph(OpRegistry::Global())); auto arg = ops::_Arg(scope.WithOpName("arg"), DT_VARIANT, 0); - (void)ops::_Retval(scope.WithOpName("ret"), arg, 0); + ops::_Retval(scope.WithOpName("ret"), arg, 0); TF_ASSERT_OK(scope.ToGraph(graph.get())); FunctionDef fdef; TF_ASSERT_OK(GraphToFunctionDef(*graph, "body", &fdef)); @@ -1999,35 +1906,26 @@ TEST_F(XlaCompilerTest, WhileWithResources) { xla::Literal literal0 = xla::LiteralUtil::CreateR0(0); xla::Literal literal1 = xla::LiteralUtil::CreateR0(2); xla::Literal literal2 = xla::LiteralUtil::CreateR0(1); - TF_ASSERT_OK_AND_ASSIGN(auto buffer0, - pjrt_client_->BufferFromHostLiteral( - literal0, pjrt_client_->memory_spaces()[0])); - TF_ASSERT_OK_AND_ASSIGN(auto buffer1, - pjrt_client_->BufferFromHostLiteral( - literal1, pjrt_client_->memory_spaces()[0])); - TF_ASSERT_OK_AND_ASSIGN(auto buffer2, - pjrt_client_->BufferFromHostLiteral( - literal2, pjrt_client_->memory_spaces()[0])); - - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_exec, - pjrt_client_->CompileAndLoad(*result.computation, xla::CompileOptions{})); - xla::ExecuteOptions execute_options; - TF_ASSERT_OK_AND_ASSIGN( - auto actual, - loaded_exec->Execute({{buffer0.get(), buffer1.get(), buffer2.get()}}, - execute_options)); - - TF_ASSERT_OK_AND_ASSIGN(auto actual_literal0, actual[0][0]->ToLiteralSync()); - TF_ASSERT_OK_AND_ASSIGN(auto actual_literal1, actual[0][1]->ToLiteralSync()); - TF_ASSERT_OK_AND_ASSIGN(auto actual_literal2, actual[0][2]->ToLiteralSync()); + std::unique_ptr data0 = + client_->TransferToServer(literal0).value(); + std::unique_ptr data1 = + client_->TransferToServer(literal1).value(); + std::unique_ptr data2 = + client_->TransferToServer(literal2).value(); + + std::unique_ptr actual = + client_ + ->Execute(*result.computation, + {data0.get(), data1.get(), data2.get()}) + .value(); + xla::Literal actual_literal = client_->Transfer(*actual).value(); xla::Literal expected0 = xla::LiteralUtil::CreateR0(10); xla::Literal expected1 = xla::LiteralUtil::CreateR0(2); xla::Literal expected2 = xla::LiteralUtil::CreateR0(1); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected0, *actual_literal0)); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected1, *actual_literal1)); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected2, *actual_literal2)); + xla::Literal expected_literal = + xla::LiteralUtil::MakeTuple({&expected0, &expected1, &expected2}); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); } TEST_F(XlaCompilerTest, SetShardingForReturnedTuple) { @@ -2230,6 +2128,7 @@ TEST_F(OpsTestBase, CompileSingleOp) { XlaCompiler::SingleOpCompileArgument single_op_arg(*context_); + xla::Client* client = xla::ClientLibrary::LocalClientOrDie(); XlaOpRegistry::RegisterCompilationKernels(); FunctionDefLibrary flib; std::unique_ptr flib_def( @@ -2237,13 +2136,7 @@ TEST_F(OpsTestBase, CompileSingleOp) { XlaCompiler::Options options; options.device_type = DeviceType(DEVICE_CPU_XLA_JIT); - auto pjrt_client = xla::GetPjRtCpuClient(xla::CpuClientOptions()); - TF_ASSERT_OK(pjrt_client.status()); - options.client = pjrt_client->get(); - auto pjrt_compiler = xla::GetDefaultPjRtCompiler(xla::CpuName()); - if (pjrt_compiler.ok()) { - options.compiler = *pjrt_compiler; - } + options.client = client; options.flib_def = flib_def.get(); XlaCompiler compiler(options); @@ -2260,18 +2153,13 @@ TEST_F(OpsTestBase, CompileSingleOp) { single_op_arg, args, &result)); // Tests that the generated computation works. - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_exec, - (*pjrt_client) - ->CompileAndLoad(*result.computation, xla::CompileOptions{})); - xla::ExecuteOptions execute_options; - TF_ASSERT_OK_AND_ASSIGN(auto actual, - loaded_exec->Execute({{}}, execute_options)); - - TF_ASSERT_OK_AND_ASSIGN(auto actual_literal, actual[0][0]->ToLiteralSync()); + std::unique_ptr actual = + client->Execute(*result.computation, {}).value(); + xla::Literal actual_literal = client->Transfer(*actual).value(); xla::Literal expected0 = xla::LiteralUtil::CreateR2({{6.9, 4.2}}); - EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected0, *actual_literal)); + xla::Literal expected_literal = xla::LiteralUtil::MakeTuple({&expected0}); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(expected_literal, actual_literal)); } TEST_F(XlaCompilerTest, SetShardingForParametersAndReturnValues) { diff --git a/tensorflow/compiler/tf2xla/xla_expression.cc b/tensorflow/compiler/tf2xla/xla_expression.cc index a9baa94c9a8ed2..d26a53e2163e25 100644 --- a/tensorflow/compiler/tf2xla/xla_expression.cc +++ b/tensorflow/compiler/tf2xla/xla_expression.cc @@ -15,36 +15,11 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_expression.h" -#include -#include -#include -#include -#include -#include - -#include "absl/log/check.h" -#include "absl/status/status.h" -#include "absl/status/statusor.h" -#include "absl/strings/str_cat.h" #include "tensorflow/compiler/tf2xla/literal_util.h" #include "tensorflow/compiler/tf2xla/shape_util.h" -#include "tensorflow/compiler/tf2xla/xla_resource.h" #include "xla/hlo/builder/value_inference.h" -#include "xla/hlo/builder/xla_builder.h" -#include "xla/hlo/builder/xla_computation.h" -#include "xla/layout.h" -#include "xla/layout_util.h" -#include "xla/literal.h" -#include "xla/pjrt/pjrt_client.h" -#include "xla/pjrt/pjrt_compiler.h" -#include "xla/pjrt/pjrt_executable.h" -#include "xla/shape.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/macros.h" -#include "xla/tsl/platform/statusor.h" -#include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/lib/core/errors.h" namespace tensorflow { @@ -173,7 +148,7 @@ absl::StatusOr XlaExpression::ResolveDynamism() const { } absl::StatusOr> XlaExpression::ResolveConstant( - xla::PjRtClient* client, bool dynamic_dimension_is_minus_one, + xla::Client* client, bool dynamic_dimension_is_minus_one, xla::ValueInferenceMode mode) const { switch (kind()) { case Kind::kConstant: @@ -216,28 +191,17 @@ absl::StatusOr> XlaExpression::ResolveConstant( return {std::nullopt}; } - if (!client) { + if (!client) return absl::InvalidArgumentError("client is required to resolve constant"); - } TF_ASSIGN_OR_RETURN(xla::XlaComputation constant_graph, handle().builder()->BuildConstantSubGraph( handle(), dynamic_dimension_is_minus_one)); - TF_ASSIGN_OR_RETURN( - std::unique_ptr loaded_exec, - client->CompileAndLoad(constant_graph, xla::CompileOptions())); - xla::ExecuteOptions execute_options; - TF_ASSIGN_OR_RETURN(auto results, - loaded_exec->Execute({{}}, execute_options)); - if (results.empty() || results[0].empty() || results[0][0] == nullptr) { - return absl::InternalError("Constant computation produced no output"); - } - TF_ASSIGN_OR_RETURN(std::shared_ptr literal, - results[0][0]->ToLiteralSync()); + TF_ASSIGN_OR_RETURN(xla::Literal literal, + client->ComputeConstant(constant_graph, &layout)); Tensor tensor; - TF_RETURN_IF_ERROR( - LiteralToHostTensor(literal->Relayout(layout), dtype(), &tensor)); + TF_RETURN_IF_ERROR(LiteralToHostTensor(literal, dtype(), &tensor)); return {tensor}; } diff --git a/tensorflow/compiler/tf2xla/xla_expression.h b/tensorflow/compiler/tf2xla/xla_expression.h index a00cd0682e01d7..ed0041fc9942a0 100644 --- a/tensorflow/compiler/tf2xla/xla_expression.h +++ b/tensorflow/compiler/tf2xla/xla_expression.h @@ -16,19 +16,14 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_TF2XLA_XLA_EXPRESSION_H_ #define TENSORFLOW_COMPILER_TF2XLA_XLA_EXPRESSION_H_ -#include -#include -#include - -#include "absl/status/statusor.h" +#include "absl/types/optional.h" #include "tensorflow/compiler/tf2xla/xla_resource.h" +#include "xla/client/client.h" #include "xla/hlo/builder/value_inference.h" #include "xla/hlo/builder/xla_builder.h" -#include "xla/pjrt/pjrt_client.h" -#include "xla/pjrt/pjrt_compiler.h" -#include "xla/shape.h" #include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/statusor.h" namespace tensorflow { @@ -131,7 +126,7 @@ class XlaExpression { // optional if it cannot be resolved. Returns an error if passed a resource // expression. absl::StatusOr> ResolveConstant( - xla::PjRtClient* client, bool dynamic_dimension_is_minus_one = false, + xla::Client* client, bool dynamic_dimension_is_minus_one = false, xla::ValueInferenceMode mode = xla::ValueInferenceMode::kValue) const; // ResolveDynamism computes where a value inside this op is dynamic or can be diff --git a/tensorflow/compiler/tf2xla/xla_expression_test.cc b/tensorflow/compiler/tf2xla/xla_expression_test.cc index c989e5484a8c8a..797002476aeb1c 100644 --- a/tensorflow/compiler/tf2xla/xla_expression_test.cc +++ b/tensorflow/compiler/tf2xla/xla_expression_test.cc @@ -15,29 +15,19 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_expression.h" -#include #include -#include -#include -#include -#include "absl/log/log.h" -#include "absl/status/statusor.h" +#include "absl/memory/memory.h" #include "tensorflow/compiler/tf2xla/xla_resource.h" +#include "xla/client/client_library.h" +#include "xla/client/local_client.h" #include "xla/hlo/builder/xla_builder.h" -#include "xla/hlo/builder/xla_computation.h" #include "xla/literal.h" -#include "xla/literal_util.h" -#include "xla/pjrt/pjrt_client.h" -#include "xla/pjrt/pjrt_executable.h" -#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h" -#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h" #include "xla/shape_util.h" +#include "xla/status_macros.h" #include "xla/tests/literal_test_util.h" -#include "xla/tsl/platform/statusor.h" -#include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/tensor_testutil.h" +#include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/test.h" namespace tensorflow { @@ -46,8 +36,7 @@ namespace { class XlaExpressionTest : public ::testing::Test { protected: void SetUp() override { - TF_ASSERT_OK_AND_ASSIGN(client_, - xla::GetXlaPjrtCpuClient(xla::CpuClientOptions())); + client_ = xla::ClientLibrary::LocalClientOrDie(); builder_ = std::make_unique("acomputation"); constant_ = test::AsScalar(42); op_ = xla::ConstantR0(builder_.get(), 7); @@ -61,7 +50,7 @@ class XlaExpressionTest : public ::testing::Test { /*tensor_array_multiple_writes_aggregate=*/false); } - std::unique_ptr client_; + xla::Client* client_; std::unique_ptr builder_; Tensor constant_; xla::XlaOp op_; @@ -97,16 +86,10 @@ TEST_F(XlaExpressionTest, AsXlaOp) { XlaExpression::Constant(constant_).AsXlaOp(builder_.get()); TF_ASSERT_OK_AND_ASSIGN(xla::XlaComputation computation, builder_->BuildConstantSubGraph(const_as_op)); - ASSERT_NE(client_, nullptr); - TF_ASSERT_OK_AND_ASSIGN( - auto loaded_exec, - client_->CompileAndLoad(computation, xla::CompileOptions{})); - xla::ExecuteOptions execute_options; - TF_ASSERT_OK_AND_ASSIGN(auto results, - loaded_exec->Execute({{}}, execute_options)); - TF_ASSERT_OK_AND_ASSIGN(auto value, results[0][0]->ToLiteralSync()); + TF_ASSERT_OK_AND_ASSIGN(xla::Literal value, + client_->ComputeConstant(computation)); EXPECT_TRUE(xla::LiteralTestUtil::Equal( - xla::LiteralUtil::CreateR0(42), *value)); + xla::LiteralUtil::CreateR0(42), value)); } TEST_F(XlaExpressionTest, GetShape) { @@ -127,27 +110,27 @@ TEST_F(XlaExpressionTest, GetShape) { } TEST_F(XlaExpressionTest, ResolveConstant) { - EXPECT_FALSE(XlaExpression().ResolveConstant(client_.get()).ok()); - EXPECT_FALSE(XlaExpression::Invalid().ResolveConstant(client_.get()).ok()); + EXPECT_FALSE(XlaExpression().ResolveConstant(client_).ok()); + EXPECT_FALSE(XlaExpression::Invalid().ResolveConstant(client_).ok()); EXPECT_FALSE(XlaExpression::Resource(resource_.get()) - .ResolveConstant(client_.get()) + .ResolveConstant(client_) ->has_value()); TF_ASSERT_OK_AND_ASSIGN( std::optional op_constant, - XlaExpression::XlaOp(op_, DT_INT32).ResolveConstant(client_.get())); + XlaExpression::XlaOp(op_, DT_INT32).ResolveConstant(client_)); ASSERT_TRUE(op_constant.has_value()); test::ExpectTensorEqual(test::AsScalar(7), *op_constant); TF_ASSERT_OK_AND_ASSIGN(std::optional op_nonconstant, XlaExpression::XlaOp(non_constant_op_, DT_FLOAT) - .ResolveConstant(client_.get())); + .ResolveConstant(client_)); EXPECT_FALSE(op_nonconstant.has_value()); TF_ASSERT_OK_AND_ASSIGN( std::optional constant_constant, - XlaExpression::Constant(constant_).ResolveConstant(client_.get())); + XlaExpression::Constant(constant_).ResolveConstant(client_)); ASSERT_TRUE(constant_constant.has_value()); test::ExpectTensorEqual(constant_, *constant_constant); } @@ -155,11 +138,11 @@ TEST_F(XlaExpressionTest, ResolveConstant) { TEST_F(XlaExpressionTest, ResolveConstantOnResource) { XlaExpression constant_resource = XlaExpression::ConstantResource(constant_, resource_.get()); - EXPECT_TRUE(constant_resource.ResolveConstant(client_.get()).ok()); + EXPECT_TRUE(constant_resource.ResolveConstant(client_).ok()); EXPECT_TRUE(resource_->SetZeroValue(builder_.get()).ok()); LOG(ERROR) << "Resource is overwritten: " << resource_->IsOverwritten(); absl::StatusOr> resolved_constant = - constant_resource.ResolveConstant(client_.get()); + constant_resource.ResolveConstant(client_); EXPECT_TRUE(resolved_constant.ok()); EXPECT_FALSE(resolved_constant->has_value()); } diff --git a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc index 40c2247a79858c..de0e7bf45eae79 100644 --- a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc +++ b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc @@ -16,16 +16,12 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h" #include -#include #include #include -#include #include #include -#include "absl/base/casts.h" #include "absl/status/status.h" -#include "absl/status/statusor.h" #include "absl/types/span.h" #include "tensorflow/compiler/tf2xla/tf2xla.h" #include "tensorflow/compiler/tf2xla/tf2xla.pb.h" @@ -33,28 +29,28 @@ limitations under the License. #include "xla/backends/cpu/buffer_allocation_info.h" #include "xla/backends/cpu/buffer_allocation_info_util.h" #include "xla/backends/cpu/codegen/compiled_function_library.h" +#include "xla/client/client_library.h" #include "xla/client/executable_build_options.h" -#include "xla/executable_run_options.h" +#include "xla/client/local_client.h" #include "xla/hlo/builder/xla_computation.h" -#include "xla/pjrt/common_pjrt_client.h" -#include "xla/pjrt/cpu/cpu_client.h" -#include "xla/pjrt/pjrt_client.h" -#include "xla/pjrt/pjrt_compiler.h" -#include "xla/pjrt/pjrt_executable.h" -#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h" -#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h" -#include "xla/service/buffer_assignment.h" #include "xla/service/cpu/cpu_aot_compilation_result.h" #include "xla/service/cpu/cpu_executable.h" -#include "xla/shape.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" +#include "xla/service/platform_util.h" +#include "xla/shape_util.h" +#include "xla/stream_executor/platform.h" #include "xla/xla_data.pb.h" #include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/errors.h" +#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/types.h" +#include "tsl/platform/casts.h" +#include "tsl/platform/statusor.h" namespace tensorflow { namespace { +constexpr char kHostPlatform[] = "Host"; // Returns the index of the result in the temp buffers. absl::StatusOr ComputeResultIndex( @@ -111,21 +107,19 @@ void CollectNames(const T& entries, std::vector* nonempty_names, XlaJitCompiledCpuFunction::Compile( const GraphDef& graph_def, const tf2xla::Config& config, const xla::ExecutableBuildOptions& build_options) { - TF_ASSIGN_OR_RETURN(xla::PjRtCompiler * compiler, - xla::GetDefaultPjRtCompiler(xla::CpuName())); - - TF_ASSIGN_OR_RETURN(std::unique_ptr client, - xla::GetXlaPjrtCpuClient(xla::CpuClientOptions())); - // Convert the graph_def into an xla::XlaComputation. + TF_ASSIGN_OR_RETURN(se::Platform * platform, + xla::PlatformUtil::GetPlatform(kHostPlatform)); + TF_ASSIGN_OR_RETURN(xla::LocalClient * client, + xla::ClientLibrary::GetOrCreateLocalClient(platform)); xla::XlaComputation computation; - TF_RETURN_IF_ERROR(tensorflow::ConvertGraphDefToXla( - graph_def, config, compiler, &computation, client.get())); + TF_RETURN_IF_ERROR(tensorflow::ConvertGraphDefToXla(graph_def, config, client, + &computation)); // Get and verify the program shape. - TF_ASSIGN_OR_RETURN(xla::ProgramShape program_shape, - computation.GetProgramShape()); - if (program_shape.result().element_type() != xla::TUPLE) { + TF_ASSIGN_OR_RETURN(std::unique_ptr program_shape, + client->GetComputationShape(computation)); + if (program_shape->result().element_type() != xla::TUPLE) { // The XlaCompiler we use to build the xla computation always generates a // tuple result, and XlaCompiledCpuFunction relies on this for simpler // calling semantics. @@ -134,35 +128,24 @@ XlaJitCompiledCpuFunction::Compile( } // The parameter names are currently meaningless, and redundant with the rest // of our metadata, so clear them out to avoid confusion and save space. - program_shape.clear_parameter_names(); + program_shape->clear_parameter_names(); // Compute arg shapes, needed to compile the executable. - std::vector arg_shapes; - arg_shapes.reserve(program_shape.parameters_size()); - for (int i = 0; i < program_shape.parameters_size(); ++i) { - arg_shapes.push_back(program_shape.parameters(i)); + std::vector arg_shapes; + arg_shapes.reserve(program_shape->parameters_size()); + for (int i = 0; i < program_shape->parameters_size(); ++i) { + arg_shapes.push_back(&program_shape->parameters(i)); } - xla::CompileOptions compile_options; - compile_options.executable_build_options = build_options; - compile_options.argument_layouts = std::move(arg_shapes); - - // Compile the executable. - TF_ASSIGN_OR_RETURN(const xla::PjRtTopologyDescription* topology, - client->GetTopologyDescription()); - TF_ASSIGN_OR_RETURN(std::unique_ptr pjrt_executable, - compiler->Compile(std::move(compile_options), computation, - *topology, client.get())); - - auto* pjrt_cpu_executable = - absl::down_cast(pjrt_executable.get()); - if (!pjrt_cpu_executable) { - return absl::InternalError( - "Expected PjRtCpuExecutable from CPU compilation."); - } + // Compile the executable. The static_cast to the CpuExecutable subclass is + // necessary since the raw function and buffer assignments are only available + // there. + TF_ASSIGN_OR_RETURN(auto executables, + client->Compile(computation, arg_shapes, build_options)); + TF_RET_CHECK(executables.size() == 1); + std::unique_ptr executable = std::move(executables[0]); xla::cpu::CpuExecutable* cpu_executable = - static_cast( - pjrt_cpu_executable->cpu_executable().get()); + static_cast(executable->executable()); const xla::BufferAssignment& buffer_assignment = cpu_executable->buffer_assignment(); @@ -171,6 +154,10 @@ XlaJitCompiledCpuFunction::Compile( xla::cpu::CreateBufferAllocationInfos(cpu_executable->module(), buffer_assignment); + std::vector buffer_allocation_infos = + xla::cpu::CreateBufferAllocationInfos(cpu_executable->module(), + buffer_assignment); + std::vector arg_index_table = xla::cpu::CreateArgIndexTable(buffer_infos); std::vector result_index_table = @@ -229,13 +216,12 @@ XlaJitCompiledCpuFunction::Compile( compiled_function_library->GetTypelessSymbolsMap(); } - jit->pjrt_client_ = std::move(client); - jit->pjrt_executable_ = std::move(pjrt_executable); + jit->executable_ = std::move(executable); jit->buffer_infos_ = std::move(buffer_infos); jit->arg_index_table_ = std::move(arg_index_table); jit->result_index_table_ = std::move(result_index_table); jit->program_shape_ = - std::make_unique(program_shape.ToProto()); + std::make_unique(program_shape->ToProto()); XlaCompiledCpuFunction::set_static_data_compilation_result_proto( &jit->static_data_, jit->compilation_result_proto_.get()); XlaCompiledCpuFunction::set_static_data_function_library_symbol_map( @@ -283,13 +269,4 @@ XlaJitCompiledCpuFunction::Compile( return std::move(jit_unique_ptr); } -const Eigen::ThreadPoolDevice* -XlaJitCompiledCpuFunction::eigen_intra_op_thread_pool_device() const { - auto* common_client = - absl::down_cast(pjrt_client_.get()); - auto* raw_client = - absl::down_cast(common_client->raw_client()); - return raw_client->eigen_intraop_device(); -} - } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h index 6f175c74d6af4c..6f61f472a2fd5a 100644 --- a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h +++ b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h @@ -16,23 +16,20 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_TF2XLA_XLA_JIT_COMPILED_CPU_FUNCTION_H_ #define TENSORFLOW_COMPILER_TF2XLA_XLA_JIT_COMPILED_CPU_FUNCTION_H_ -#include #include #include #include #include "absl/container/flat_hash_map.h" #include "absl/log/check.h" -#include "absl/status/statusor.h" +#include "tensorflow/compiler/tf2xla/encoded_buffer_allocation_info.h" #include "tensorflow/compiler/tf2xla/tf2xla.pb.h" -#include "tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h" +#include "tensorflow/compiler/tf2xla/xla_compiled_cpu_function_thunks.h" #include "xla/backends/cpu/buffer_allocation_info.h" -#include "xla/client/executable_build_options.h" -#include "xla/executable_run_options.h" -#include "xla/pjrt/pjrt_client.h" -#include "xla/pjrt/pjrt_executable.h" +#include "xla/client/local_client.h" #include "xla/service/cpu/executable.pb.h" #include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/platform/types.h" namespace tensorflow { @@ -65,21 +62,16 @@ class XlaJitCompiledCpuFunction { return static_data_; } - const xla::PjRtExecutable& pjrt_executable() const { - CHECK(pjrt_executable_ != nullptr); // Crash ok - return *pjrt_executable_; + const xla::LocalExecutable& LocalExecutable() const { + CHECK(executable_); // Crash ok + return *executable_; } - const Eigen::ThreadPoolDevice* eigen_intra_op_thread_pool_device() const; - private: XlaJitCompiledCpuFunction() : compilation_result_proto_(nullptr) {} - // The client holds the thread pool used for execution. - std::unique_ptr pjrt_client_; - // The executable holds the underlying function. - std::unique_ptr pjrt_executable_; + std::unique_ptr executable_; // The compilation result proto. std::unique_ptr compilation_result_proto_; diff --git a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function_test.cc b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function_test.cc index 05fd83d967cc30..3864e4fdeca467 100644 --- a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function_test.cc +++ b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function_test.cc @@ -15,7 +15,6 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h" -#include #include #include @@ -26,24 +25,25 @@ limitations under the License. #include "xla/client/executable_build_options.h" #include "xla/hlo/testlib/test.h" #include "xla/service/compiler.h" +#include "xla/service/platform_util.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/stream_executor/platform.h" -#include "xla/stream_executor/platform_id.h" #include "xla/stream_executor/platform_manager.h" #include "xla/tsl/lib/core/status_test_util.h" -#include "xla/tsl/platform/statusor.h" #include "xla/xla_data.pb.h" #include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/attr_value_util.h" #include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/framework/node_def.pb.h" #include "tensorflow/core/platform/test.h" -#include "tsl/platform/protobuf.h" +#include "tensorflow/core/platform/types.h" namespace tensorflow { namespace { +using ::testing::HasSubstr; + PLATFORM_DEFINE_ID(kFakePlatformId, FakePlatform); AttrValue TypeAttrValue(DataType type) { @@ -347,6 +347,9 @@ TEST(XlaJitCompiledCpuFunction, CanCompileWithAdditionalPlatform) { return std::unique_ptr(nullptr); }); + EXPECT_THAT(xla::PlatformUtil::GetDefaultPlatform().status().message(), + HasSubstr(kFakePlatformId->ToName())); + GraphDef graph_def = SumGraph(); tf2xla::Config config = SumConfig(); TF_ASSERT_OK_AND_ASSIGN( diff --git a/tensorflow/core/api_def/base_api/api_def_MatrixSolve.pbtxt b/tensorflow/core/api_def/base_api/api_def_MatrixSolve.pbtxt index d3b1216d404e3d..1a03ead009cc71 100644 --- a/tensorflow/core/api_def/base_api/api_def_MatrixSolve.pbtxt +++ b/tensorflow/core/api_def/base_api/api_def_MatrixSolve.pbtxt @@ -33,5 +33,11 @@ a tensor shape `[..., M, K]`. If `adjoint` is `False` then each output matrix satisfies `matrix[..., :, :] * output[..., :, :] = rhs[..., :, :]`. If `adjoint` is `True` then each output matrix satisfies `adjoint(matrix[..., :, :]) * output[..., :, :] = rhs[..., :, :]`. + +If a matrix is not invertible there is no guarantee what the op does. It +may detect the condition and raise an exception or it may simply return a +garbage result. In particular, under XLA compilation (e.g. +`tf.function(jit_compile=True)`) the condition is not detected and the +result is implementation-defined. END } diff --git a/tensorflow/core/common_runtime/eager/BUILD b/tensorflow/core/common_runtime/eager/BUILD index 5aaddc171f0fb9..8935b2fc2a4b84 100644 --- a/tensorflow/core/common_runtime/eager/BUILD +++ b/tensorflow/core/common_runtime/eager/BUILD @@ -13,6 +13,7 @@ # limitations under the License. # ============================================================================== +load("@bazel_skylib//lib:selects.bzl", "selects") load( "@xla//xla/tsl/mkl:build_defs.bzl", "if_mkl", @@ -20,7 +21,8 @@ load( ) load( "//tensorflow:tensorflow.bzl", - "if_cuda_or_rocm", + "clean_dep", + "if_google", "if_zendnn", "tf_cc_test", "tf_cc_test_mkl", @@ -276,11 +278,30 @@ tf_cuda_library( "@xla//xla/tsl/distributed_runtime/preemption:preemption_notifier", "@xla//xla/tsl/platform:statusor", ], - }) + if_cuda_or_rocm([ - "@xla//xla/pjrt/distributed:key_value_store_interface", - "//tensorflow/core/common_runtime/gpu:gpu_runtime_impl", - "//tensorflow/core/framework:resource_base", - ]), + }) + if_google( + # TODO(b/282068262): PJRT pulls in TFRT components that are incompatible with ARM platform. + # Clean up so that PJRT can run on ARM (and remove "#if defined(PLATFORM_GOOGLE) ..." use + # from gpu_util.cc). + # Also it won't build with WeightWatcher which tracks OSS build binaries. + # TODO(b/290533709): Clean up this build rule. + selects.with_or({ + clean_dep("//tensorflow:linux_x86_64_with_weightwatcher"): [], + ( + clean_dep("//tensorflow:linux_x86_64"), + ): [ + "//tensorflow/core", + "//tensorflow/core/framework:resource_base", + "@xla//xla/pjrt/distributed:key_value_store_interface", + "@xla//xla/pjrt:local_device_state", + "@xla//xla/service:gpu_topology", + "@xla//xla/pjrt:pjrt_client", + "@xla//xla/pjrt:pjrt_compiler", + "@xla//xla/service/gpu:gpu_executable_run_options", + "//tensorflow/core/common_runtime/gpu:gpu_runtime_impl", + ], + "//conditions:default": [], + }), + ), ) tf_cc_test( diff --git a/tensorflow/core/common_runtime/eager/context_distributed_manager.cc b/tensorflow/core/common_runtime/eager/context_distributed_manager.cc index a3a9924112311b..af057e4b39c552 100644 --- a/tensorflow/core/common_runtime/eager/context_distributed_manager.cc +++ b/tensorflow/core/common_runtime/eager/context_distributed_manager.cc @@ -78,19 +78,22 @@ limitations under the License. #include "tensorflow/core/distributed_runtime/worker_interface.h" #endif // !IS_MOBILE_PLATFORM -#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \ - (defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM) +#if (defined(PLATFORM_GOOGLE) && defined(TF_PLATFORM_LINUX_X86_64)) +#define TF_GPU_USE_PJRT #include "xla/pjrt/distributed/key_value_store_interface.h" #include "xla/pjrt/gpu/se_gpu_pjrt_client.h" +#include "xla/pjrt/pjrt_compiler.h" #include "xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h" #include "xla/pjrt/se/local_device_state.h" #include "xla/pjrt/se/pjrt_stream_executor_client.h" +#include "xla/service/gpu/gpu_executable_run_options.h" +#include "xla/service/gpu_topology.h" #include "tensorflow/core/framework/resource_base.h" #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/tfrt/common/global_state.h" #include "tensorflow/core/tfrt/common/pjrt_state.h" #include "tensorflow/core/tfrt/common/pjrt_util.h" -#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM +#endif namespace tensorflow { @@ -110,8 +113,7 @@ namespace { } \ } while (0); -#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \ - (defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM) +#ifdef TF_GPU_USE_PJRT // Provide a KeyValue interface to the coordination service agent for use by // BuildDistributedDevices. class XlaKeyValueStore : public xla::KeyValueStoreInterface { @@ -318,20 +320,19 @@ absl::Status CreateClientOnce( return absl::OkStatus(); } } -#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM +#endif // TF_GPU_USE_PJRT absl::Status CreatePjRtGpuClientWithDistributedDevices( int node_id, int num_nodes, tsl::CoordinationServiceAgent* coordination_service_agent) { -#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \ - (defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM) +#ifdef TF_GPU_USE_PJRT if (num_nodes <= 1) { return absl::OkStatus(); } return CreateClientOnce(node_id, num_nodes, coordination_service_agent); -#else +#else // TF_GPU_USE_PJRT return absl::OkStatus(); -#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM +#endif // TF_GPU_USE_PJRT } bool AreLocalDevicesCompatible(const EagerContext* context, diff --git a/tensorflow/core/common_runtime/gpu/BUILD b/tensorflow/core/common_runtime/gpu/BUILD index dd02102cd077d0..fbfb30d068fe55 100644 --- a/tensorflow/core/common_runtime/gpu/BUILD +++ b/tensorflow/core/common_runtime/gpu/BUILD @@ -13,13 +13,16 @@ # limitations under the License. # ============================================================================== +load("@bazel_skylib//lib:selects.bzl", "selects") load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda") load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm") load("@xla//xla/tsl:tsl.default.bzl", "if_cuda_libs") load( "//tensorflow:tensorflow.bzl", + "clean_dep", "if_cuda_or_rocm", "if_google", + "if_linux_x86_64", "tf_cc_test", "tf_copts", "tf_cuda_library", @@ -190,6 +193,7 @@ tf_cuda_library( "@xla//xla/stream_executor/gpu:gpu_cudamallocasync_allocator", ":gpu_runtime_hermetic_cuda_deps", ], + defines = if_linux_x86_64(["TF_PLATFORM_LINUX_X86_64"]), features = ["-layering_check"], visibility = [ "//tensorflow:internal", @@ -206,22 +210,17 @@ tf_cuda_library( "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", - "//tensorflow/core:session_options", "//tensorflow/core/common_runtime:core_cpu_impl", "//tensorflow/core/common_runtime:device_id_utils", "//tensorflow/core/common_runtime:node_file_writer", "//tensorflow/core/framework:device_attributes_proto_cc", "//tensorflow/core/framework:types_proto_cc", - "//tensorflow/core/platform:hash", "//tensorflow/core/platform:stream_executor", "//tensorflow/core/platform:tensor_float_32_utils", "//tensorflow/core/profiler/lib:annotated_traceme", "//tensorflow/core/profiler/lib:scoped_annotation", "//tensorflow/core/profiler/lib:scoped_memory_debug_annotation", - "@com_google_absl//absl/base", - "@com_google_absl//absl/base:core_headers", - "@com_google_absl//absl/container:btree", - "@com_google_absl//absl/container:flat_hash_map", + "//tensorflow/core/protobuf:for_core_protos_cc", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", @@ -231,39 +230,36 @@ tf_cuda_library( "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/synchronization", - "@eigen_archive//:eigen3", - "@tsl//tsl/platform:mutex", - "@tsl//tsl/platform:numbers", - "@tsl//tsl/platform:strcat", - "@tsl//tsl/profiler/lib:scoped_annotation", - "@tsl//tsl/profiler/lib:traceme", "@xla//xla:shape_util", "@xla//xla/pjrt:host_memory_allocator", - "@xla//xla/stream_executor:device_memory", - "@xla//xla/stream_executor:memory_space", - "@xla//xla/stream_executor:stream_executor_h", "@xla//xla/stream_executor/gpu:gpu_init_impl", "@xla//xla/stream_executor/integrations:stream_executor_allocator", "@xla//xla/tsl/framework:device_id_utils", - "@xla//xla/tsl/framework:scoped_allocation_trace", - "@xla//xla/tsl/platform:logging", - "@xla//xla/tsl/platform:types", - ] + if_cuda_or_rocm([ - "//tensorflow/compiler/tf2xla:common", - "//tensorflow/compiler/tf2xla:layout_util", - "//tensorflow/compiler/jit:flags", - "//tensorflow/compiler/jit:pjrt_device_context", - "//tensorflow/compiler/jit:pjrt_tensor_buffer", - "@xla//xla:future", - "@xla//xla:literal", - "@xla//xla/pjrt:local_device_state", - "@xla//xla/pjrt:pjrt_client", + ] + if_google( + # TODO(b/282068262): PJRT pulls in TFRT components that are incompatible with ARM platform. + # Clean up so that PJRT can run on ARM (and remove "#if defined(PLATFORM_GOOGLE) ..." use + # from gpu_util.cc). + # Also it won't build with WeightWatcher which tracks OSS build binaries. + # TODO(b/290533709): Clean up this build rule. + selects.with_or({ + clean_dep("//tensorflow:linux_x86_64_with_weightwatcher"): [], + ( + clean_dep("//tensorflow:linux_x86_64"), + ): [ + "//tensorflow/compiler/tf2xla:layout_util", + "//tensorflow/compiler/jit:flags", + "//tensorflow/compiler/jit:pjrt_device_context", + "@xla//xla/stream_executor/integrations:tf_allocator_adapter", + "//tensorflow/core/tfrt/common:pjrt_util", + ], + "//conditions:default": [], + }) + if_cuda_or_rocm([ + "@xla//xla/service:gpu_plugin_impl", # for registering cuda compiler. + ]), + ) + if_cuda_or_rocm([ + "@tsl//tsl/platform:dso_loader", "@xla//xla/pjrt/gpu:gpu_helpers", "@xla//xla/pjrt/gpu:se_gpu_pjrt_client", - "@xla//xla/service:gpu_plugin_impl", # for registering cuda compiler. - "@xla//xla/stream_executor/integrations:tf_allocator_adapter", - "//tensorflow/core/tfrt/common:pjrt_util", - "@tsl//tsl/platform:dso_loader", ]) + if_cuda([ "@xla//xla/stream_executor/cuda:all_runtime", ]) + if_rocm([ @@ -278,37 +274,17 @@ tf_cuda_library( linkstatic = 1, visibility = ["//visibility:public"], deps = [ - "//tensorflow/compiler/tf2xla:layout_util", "//tensorflow/core:core_cpu_lib", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", - "//tensorflow/core:gpu_headers_lib", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", - "//tensorflow/core:session_options", "//tensorflow/core/platform:stream_executor", - "@com_google_absl//absl/base", - "@com_google_absl//absl/base:core_headers", - "@com_google_absl//absl/container:btree", - "@com_google_absl//absl/container:flat_hash_map", - "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/status", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/strings:string_view", - "@com_google_absl//absl/synchronization", "@eigen_archive//:eigen3", - "@tsl//tsl/platform:numbers", - "@xla//xla/pjrt:local_device_state", - "@xla//xla/stream_executor:stream_executor_h", "@xla//xla/stream_executor/gpu:gpu_init", - "@xla//xla/stream_executor/integrations:tf_allocator_adapter", - "@xla//xla/tsl/framework:scoped_allocation_trace", - "@xla//xla/tsl/platform:logging", - "@xla//xla/tsl/platform:types", - ] + if_static([":gpu_runtime_impl"]) + if_cuda_or_rocm([ - "//tensorflow/compiler/jit:pjrt_device_context", - ]), + ] + if_static([":gpu_runtime_impl"]), ) # This is redundant with the "gpu_runtime_*" targets above. It's useful for @@ -386,25 +362,25 @@ tf_cuda_cc_test( "gpu_device_test.cc", ], features = ["-layering_check"], - tags = tf_cuda_tests_tags() + [ - "nobuilder", - ], + tags = tf_cuda_tests_tags(), deps = [ + ":gpu_id", ":gpu_runtime", "//tensorflow/cc:cc_ops", + "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", + "//tensorflow/core:testlib", + "//tensorflow/core/common_runtime:core_cpu", "//tensorflow/core/common_runtime:core_cpu_internal", "//tensorflow/core/common_runtime:direct_session_internal", - "//tensorflow/core/tfrt/common:pjrt_util", + "//tensorflow/core/kernels:ops_util", "@com_google_absl//absl/synchronization", - "@local_config_cuda//cuda:cuda_headers", - "@xla//xla/pjrt:pjrt_client", "@xla//xla/stream_executor/gpu:gpu_cudamallocasync_allocator", - "@xla//xla/stream_executor/gpu:gpu_init_impl", "@xla//xla/tsl/framework:device_id", ], ) @@ -450,25 +426,24 @@ tf_cuda_cc_test( tags = tf_cuda_tests_tags() + [ "guitar", # "multi_gpu", # TODO(b/287692888): re-enable once the 2gpu test passes. - "nobuilder", ], deps = [ - ":gpu_runtime", + ":gpu_id", "//tensorflow/cc:cc_ops", + "//tensorflow/core:framework", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", + "//tensorflow/core:testlib", + "//tensorflow/core/common_runtime:core_cpu", "//tensorflow/core/common_runtime:core_cpu_internal", "//tensorflow/core/common_runtime:direct_session_internal", - "//tensorflow/core/tfrt/common:pjrt_util", + "//tensorflow/core/kernels:ops_util", "@com_google_absl//absl/synchronization", - "@local_config_cuda//cuda:cuda_headers", - "@xla//xla/pjrt:pjrt_client", "@xla//xla/stream_executor/gpu:gpu_cudamallocasync_allocator", - "@xla//xla/stream_executor/gpu:gpu_init_impl", - "@xla//xla/tsl/framework:device_id", ], ) diff --git a/tensorflow/core/common_runtime/gpu/gpu_device.cc b/tensorflow/core/common_runtime/gpu/gpu_device.cc index 096de4ba7325f3..a55378a9db4806 100644 --- a/tensorflow/core/common_runtime/gpu/gpu_device.cc +++ b/tensorflow/core/common_runtime/gpu/gpu_device.cc @@ -15,9 +15,15 @@ limitations under the License. // TODO(opensource): Use a more generic sounding preprocessor name than // GOOGLE_CUDA +#include "xla/pjrt/host_memory_allocator.h" #if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \ (defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM) +// TODO(b/282059652): Merge google internal and open-source code path once TF +// dependency issue is resolved. +#if (defined(PLATFORM_GOOGLE) && defined(TF_PLATFORM_LINUX_X86_64)) +#define TF_GPU_USE_PJRT +#endif // PLATFORM_GOOGLE && TF_PLATFORM_LINUX_X86_64 #if TENSORFLOW_USE_ROCM #include "rocm/include/hip/hip_runtime.h" @@ -72,11 +78,12 @@ limitations under the License. #elif TENSORFLOW_USE_ROCM #include "tensorflow/core/platform/rocm.h" #endif +#ifdef TF_GPU_USE_PJRT #include "tensorflow/compiler/jit/flags.h" #include "xla/pjrt/gpu/gpu_helpers.h" #include "xla/pjrt/gpu/se_gpu_pjrt_client.h" -#include "xla/pjrt/host_memory_allocator.h" #include "xla/pjrt/pjrt_client.h" +#endif // TF_GPU_USE_PJRT #include "tensorflow/core/framework/log_memory.h" #include "tensorflow/core/platform/fingerprint.h" #include "tensorflow/core/platform/logging.h" @@ -86,11 +93,13 @@ limitations under the License. #include "tensorflow/core/profiler/lib/scoped_annotation.h" #include "tensorflow/core/profiler/lib/scoped_memory_debug_annotation.h" #include "tensorflow/core/public/session_options.h" +#include "tsl/platform/dso_loader.h" +#ifdef TF_GPU_USE_PJRT #include "tensorflow/core/tfrt/common/pjrt_util.h" +#endif // TF_GPU_USE_PJRT #include "tensorflow/core/util/device_name_utils.h" #include "tensorflow/core/util/env_var.h" #include "tensorflow/core/util/stream_executor_util.h" -#include "tsl/platform/dso_loader.h" #if !defined(PLATFORM_GOOGLE) #if GOOGLE_CUDA @@ -285,6 +294,64 @@ class BaseGPUDevice::StreamGroupFactory { return std::make_pair(&insert_result.first->second, insert_result.second); } + // Returns the unique stream group for use with the stream defined by + // {tf_device_id, stream_group_within_gpu}, creating it if it does not yet + // exist. + // This function is thread safe. + BaseGPUDevice::StreamGroup* GetOrCreate(tsl::TfDeviceId tf_device_id, + int stream_group_within_gpu, + se::StreamExecutor* executor, + const GPUOptions& options) { + mutex_lock guard(lock_); + StreamGroup* group = + &streams_[key_type(tf_device_id.value(), stream_group_within_gpu)]; + if (!group->compute) { + int priority = GetPriority(tf_device_id.value(), options); + group->priority = priority; + group->compute = GetInitializedStream(executor, priority); + VLOG(2) << "Created stream[" << stream_group_within_gpu + << "] = " << group->compute << " with priority: " << priority; + +#if TENSORFLOW_USE_ROCM + // ROCm streams are lightweight and will not necessarily trigger device + // queue init until they are first used. For optimal performance, + // compute and nccl streams must be immediate siblings. + group->nccl = GetInitializedStream(executor, priority); + VLOG(2) << "Created nccl_stream[" << stream_group_within_gpu + << "] = " << group->nccl; + + // Force underlying resource creation now. + group->compute->WaitFor(group->nccl).IgnoreError(); + group->nccl->WaitFor(group->compute).IgnoreError(); +#endif + + group->host_to_device = GetInitializedStream(executor, priority); + VLOG(2) << "Created host_to_device_stream[" << stream_group_within_gpu + << "] = " << group->host_to_device; + + group->device_to_host = GetInitializedStream(executor, priority); + VLOG(2) << "Created device_to_host_stream[" << stream_group_within_gpu + << "] = " << group->device_to_host; + + int num_d2d_streams = + options.experimental().num_dev_to_dev_copy_streams(); + if (num_d2d_streams == 0) num_d2d_streams = 1; + if (num_d2d_streams < 1 || num_d2d_streams > 4) { + LOG(ERROR) + << "Illegal GPUOptions.experimental.num_dev_to_dev_copy_streams=" + << num_d2d_streams << " set to 1 instead."; + num_d2d_streams = 1; + } + for (int i = 0; i < num_d2d_streams; ++i) { + se::Stream* stream = GetInitializedStream(executor, priority); + group->device_to_device.push_back(stream); + VLOG(2) << "Created device_to_device_stream[" << stream_group_within_gpu + << "] = " << group->device_to_device.back(); + } + } + return group; + } + // Returns a reference to the StreamGroupFactory singleton. Note that this is // never destroyed, so the objects it owns are never deleted. static StreamGroupFactory& Global() { @@ -295,24 +362,46 @@ class BaseGPUDevice::StreamGroupFactory { // Helper method for unit tests to reset the streams. Never use in production. void TestOnlyReset() { mutex_lock guard(lock_); - // When PJRT is used, streams are managed by PjRtClient. for (auto& item : streams_) { auto& stream = item.second; if (stream.compute) { +#ifndef TF_GPU_USE_PJRT // When PJRT is used, streams are managed by + // PjRtClient. + delete stream.compute; +#endif stream.compute = nullptr; } #if TENSORFLOW_USE_ROCM if (stream.nccl) { +#ifndef TF_GPU_USE_PJRT // When PJRT is used, streams are managed by + // PjRtClient. + delete stream.nccl; +#endif stream.nccl = nullptr; } #endif if (stream.host_to_device) { +#ifndef TF_GPU_USE_PJRT // When PJRT is used, streams are managed by + // PjRtClient. + delete stream.host_to_device; +#endif stream.host_to_device = nullptr; } if (stream.device_to_host) { +#ifndef TF_GPU_USE_PJRT // When PJRT is used, streams are managed by + // PjRtClient. + delete stream.device_to_host; +#endif stream.device_to_host = nullptr; } while (!stream.device_to_device.empty()) { + auto back = stream.device_to_device.back(); + if (back) { +#ifndef TF_GPU_USE_PJRT // When PJRT is used, streams are managed by + // PjRtClient. + delete back; +#endif + } stream.device_to_device.pop_back(); } } @@ -329,9 +418,22 @@ class BaseGPUDevice::StreamGroupFactory { } private: + // Returns a Stream with the underlying GPUStream with the given priority. + se::Stream* GetInitializedStream(se::StreamExecutor* executor, int priority) { + auto stream_or_status = executor->CreateStream(priority); + if (!stream_or_status.ok()) { + LOG(ERROR) << "Failed to create stream: " << stream_or_status.status(); + return nullptr; + } + auto stream_ptr = stream_or_status->get(); + allocated_streams_.emplace_back(std::move(stream_or_status.value())); + return stream_ptr; + } + mutex lock_; using key_type = std::tuple; std::map streams_; + std::vector> allocated_streams_; // StreamGroupFactory cannot be created directly; Call // StreamGroupFactory::Global() to get the global instance. @@ -361,12 +463,14 @@ BaseGPUDevice::BaseGPUDevice(const SessionOptions& options, // names (which include a replica index even for multi-client). set_xla_global_id(Fingerprint32(name) % std::numeric_limits::max()); +#ifdef TF_GPU_USE_PJRT // Note: ShapeDeterminationFns is not used in GPU. XlaShapeLayoutHelpers::ShapeDeterminationFns shape_fns{ UseNoPreferenceLayoutFn(), IdentityShapeRepresentationFn()}; pjrt_device_context_ = core::RefCountPtr( new PjRtDeviceContext(shape_fns, /*use_pjrt_tensor_buffer=*/true)); +#endif // TF_GPU_USE_PJRT GPUProcessState::singleton()->EnableGPUDevice(); } @@ -400,9 +504,13 @@ absl::Status BaseGPUDevice::InitScratchBuffers() { return absl::OkStatus(); } +#ifdef TF_GPU_USE_PJRT absl::Status BaseGPUDevice::Init( const SessionOptions& options, xla::LocalDeviceState* xla_local_device_state) { +#else +Status BaseGPUDevice::Init(const SessionOptions& options) { +#endif // TF_GPU_USE_PJRT auto executor_status = DeviceIdUtil::ExecutorForTfDeviceId( DEVICE_GPU, se::GPUMachineManager(), tf_device_id_); if (!executor_status.status().ok()) { @@ -412,6 +520,7 @@ absl::Status BaseGPUDevice::Init( executor_ = executor_status.value(); +#ifdef TF_GPU_USE_PJRT CHECK(xla_local_device_state != nullptr); // Crash OK. // Construct a StreamGroup and put it inside the global factory. // TODO(tensorflow-team): set up nccl stream when TENSORFLOW_USE_ROCM is set. @@ -435,6 +544,10 @@ absl::Status BaseGPUDevice::Init( << " already exists. This usually only happens in unit tests."; } stream_ = emplace_result.first; +#else + stream_ = StreamGroupFactory::Global().GetOrCreate( + tf_device_id_, 0, executor_, options.config.gpu_options()); +#endif // TF_GPU_USE_PJRT // Get an allocator that allocates pinned memory on host. AllocatorAttributes attr; @@ -481,10 +594,14 @@ absl::Status BaseGPUDevice::Init( accelerator_device_info_ = new DeviceBase::AcceleratorDeviceInfo; accelerator_device_info_->stream = stream_->compute; accelerator_device_info_->default_context = device_context_; +#ifdef TF_GPU_USE_PJRT accelerator_device_info_->pjrt_context = pjrt_device_context_.get(); + bool use_pjrt = + GetXlaOpsCommonFlags()->tf_xla_use_device_api.IsEnabledForGpu(); accelerator_device_info_->use_pjrt_tensor_buffer = - static_cast(pjrt_device_context_.get()) - ->use_pjrt_tensor_buffer(); + use_pjrt && static_cast(pjrt_device_context_.get()) + ->use_pjrt_tensor_buffer(); +#endif // TF_GPU_USE_PJRT accelerator_device_info_->event_mgr = em_; tsl::PlatformDeviceId platform_device_id; TF_RETURN_IF_ERROR( @@ -1323,6 +1440,7 @@ absl::Status BaseGPUDeviceFactory::CreateDevices( bool populate_pjrt_gpu_client_creation_info = gpu_options.experimental().populate_pjrt_gpu_client_creation_info(); +#ifdef TF_GPU_USE_PJRT absl::StatusOr obtained_info = GetPjRtGpuClientCreationInfo(); if (obtained_info.ok() && obtained_info.value() != nullptr) { @@ -1333,6 +1451,7 @@ absl::Status BaseGPUDeviceFactory::CreateDevices( VLOG(3) << "Previous GetPjRtGpuClientCreationInfo does not exist. Will create."; } +#endif std::vector visible_gpu_order; std::vector valid_platform_device_ids; @@ -1604,6 +1723,7 @@ absl::Status BaseGPUDeviceFactory::CreateDevices( TF_RETURN_IF_ERROR(GetDeviceLocalities( tf_device_specs.size(), interconnect_maps, &device_localities)); +#ifdef TF_GPU_USE_PJRT // After the GPU device creation loop, allocator_id_stream_tuples will be // populated. std::vector allocator_id_stream_tuples; @@ -1647,6 +1767,7 @@ absl::Status BaseGPUDeviceFactory::CreateDevices( "create a new PjRt GPU Client."; } } +#endif // TF_GPU_USE_PJRT GPUProcessState* process_state = GPUProcessState::singleton(); @@ -1680,6 +1801,7 @@ absl::Status BaseGPUDeviceFactory::CreateDevices( tf_device_id.value())); } +#ifdef TF_GPU_USE_PJRT // Create xla::LocalDeviceState. const auto executor_status = DeviceIdUtil::ExecutorForTfDeviceId( DEVICE_GPU, gpu_manager, tf_device_id); @@ -1749,14 +1871,10 @@ absl::Status BaseGPUDeviceFactory::CreateDevices( if (should_create_new_pjrt_client || populate_pjrt_gpu_client_creation_info) { - allocator_id_stream_tuples.push_back( - se::MultiDeviceAdapter::AllocatorInfo{ - .allocator = - std::make_unique(gpu_allocator), - .stream = local_device_state->compute_stream(), - .memory_space = 0, - .device_ordinal = tf_device_id.value(), - }); + auto gpu_allocator_ptr = std::unique_ptr(gpu_allocator); + allocator_id_stream_tuples.emplace_back( + std::move(gpu_allocator_ptr), local_device_state->compute_stream(), 0, + tf_device_id.value()); } } @@ -1765,11 +1883,6 @@ absl::Status BaseGPUDeviceFactory::CreateDevices( return absl::OkStatus(); } - auto& pjrt_rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api; - pjrt_rollout_config.AllowForDeviceInXlaLaunch(DEVICE_GPU); - pjrt_rollout_config.AllowForDeviceInXlaCompileOnDemand(DEVICE_GPU); - pjrt_rollout_config.AllowForDeviceInXlaCompileAndRun(DEVICE_GPU); - if (should_create_new_pjrt_client || populate_pjrt_gpu_client_creation_info) { VLOG(3) << "should_create_new_pjrt_client=" << should_create_new_pjrt_client << ", populate_pjrt_gpu_client_creation_info=" @@ -1782,7 +1895,7 @@ absl::Status BaseGPUDeviceFactory::CreateDevices( auto pjrt_gpu_host_allocator = std::make_unique( - std::make_unique( + std::unique_ptr( process_state->GetGpuHostAllocator(/*options=*/{}, numa_node))); if (populate_pjrt_gpu_client_creation_info && @@ -1823,12 +1936,16 @@ absl::Status BaseGPUDeviceFactory::CreateDevices( int node_id = gpu_options.experimental().node_id(); xla::GpuClientOptions options; options.node_id = node_id; - TF_ASSIGN_OR_RETURN( - auto pjrt_client, - xla::GetSharedStreamExecutorGpuClient( - options, xla_client, std::move(local_device_states), - std::move(allocator_adapter), - std::move(pjrt_gpu_host_allocator))); + ASSIGN_OR_RETURN(auto pjrt_client, + xla::GetSharedStreamExecutorGpuClient( + options, xla_client, std::move(local_device_states), + std::move(allocator_adapter), + std::move(pjrt_gpu_host_allocator))); + + auto& pjrt_rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api; + pjrt_rollout_config.AllowForDeviceInXlaLaunch(DEVICE_GPU); + pjrt_rollout_config.AllowForDeviceInXlaCompileOnDemand(DEVICE_GPU); + pjrt_rollout_config.AllowForDeviceInXlaCompileAndRun(DEVICE_GPU); return SetPjRtClientInTFGlobalResourceManager(DeviceType(DEVICE_GPU), std::move(pjrt_client)); @@ -1845,6 +1962,13 @@ absl::Status BaseGPUDeviceFactory::CreateDevices( } else { return obtained_pjrt_client.status(); } +#else + TF_RETURN_IF_ERROR(CreateGPUDevice(options, name_prefix, tf_device_id, + /*dev_locality=*/it->second, + gpu_allocator, devices)); + } + return OkStatus(); +#endif // TF_GPU_USE_PJRT } static std::string GetShortDeviceDescription( @@ -1864,11 +1988,18 @@ static std::string GetShortDeviceDescription( #endif } +#ifdef TF_GPU_USE_PJRT absl::Status BaseGPUDeviceFactory::CreateGPUDevice( const SessionOptions& options, const std::string& name_prefix, tsl::TfDeviceId tf_device_id, const DeviceLocality& dev_locality, xla::LocalDeviceState* xla_local_device_state, Allocator* gpu_allocator, std::vector>* devices) { +#else +Status BaseGPUDeviceFactory::CreateGPUDevice( + const SessionOptions& options, const string& name_prefix, + tsl::TfDeviceId tf_device_id, const DeviceLocality& dev_locality, + Allocator* gpu_allocator, std::vector>* devices) { +#endif // TF_GPU_USE_PJRT CHECK_GE(tf_device_id.value(), 0); const std::string device_name = absl::StrCat(name_prefix, "/device:GPU:", tf_device_id.value()); @@ -1905,7 +2036,11 @@ absl::Status BaseGPUDeviceFactory::CreateGPUDevice( LOG(INFO) << "Created device " << device_name << " with " << (bytes_limit >> 20) << " MB memory: " << " -> " << GetShortDeviceDescription(platform_device_id, *desc); +#ifdef TF_GPU_USE_PJRT TF_RETURN_IF_ERROR(gpu_device->Init(options, xla_local_device_state)); +#else + TF_RETURN_IF_ERROR(gpu_device->Init(options)); +#endif // TF_GPU_USE_PJRT gpu_allocator->SetStreamAndPreallocateMemory( gpu_device->compute_stream()->platform_specific_handle().stream); diff --git a/tensorflow/core/common_runtime/gpu/gpu_device.h b/tensorflow/core/common_runtime/gpu/gpu_device.h index b0d9ca81861c67..441715bd2d22cb 100644 --- a/tensorflow/core/common_runtime/gpu/gpu_device.h +++ b/tensorflow/core/common_runtime/gpu/gpu_device.h @@ -20,6 +20,12 @@ limitations under the License. #ifndef TENSORFLOW_CORE_COMMON_RUNTIME_GPU_GPU_DEVICE_H_ #define TENSORFLOW_CORE_COMMON_RUNTIME_GPU_GPU_DEVICE_H_ +// TODO(b/282059652): Merge google internal and open-source code path once TF +// dependency issue is resolved. +#if (defined(PLATFORM_GOOGLE) && defined(TF_PLATFORM_LINUX_X86_64)) +#define TF_GPU_USE_PJRT +#endif // PLATFORM_GOOGLE && TF_PLATFORM_LINUX_X86_64 + #include #include #include @@ -30,10 +36,12 @@ limitations under the License. #include #include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive +#ifdef TF_GPU_USE_PJRT #include "tensorflow/compiler/jit/pjrt_device_context.h" #include "tensorflow/compiler/tf2xla/layout_util.h" #include "xla/pjrt/local_device_state.h" #include "xla/stream_executor/integrations/tf_allocator_adapter.h" +#endif // TF_GPU_USE_PJRT #include "xla/tsl/framework/device_id.h" #include "tensorflow/core/common_runtime/device_factory.h" #include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h" @@ -102,8 +110,12 @@ class BaseGPUDevice : public LocalDevice { }; // Initialize the device and return the status of initialization. +#ifdef TF_GPU_USE_PJRT absl::Status Init(const SessionOptions& options, xla::LocalDeviceState* xla_local_device_state); +#else + absl::Status Init(const SessionOptions& options); +#endif // TF_GPU_USE_PJRT void Compute(OpKernel* op_kernel, OpKernelContext* context) override; @@ -410,6 +422,7 @@ class BaseGPUDeviceFactory : public DeviceFactory { // Creates a BaseGPUDevice associated with 'tf_device_id', and adds it to the // 'devices' vector. The 'gpu_allocator' is created by the caller and usually // preallocates a set amount of GPU memory. +#ifdef TF_GPU_USE_PJRT absl::Status CreateGPUDevice(const SessionOptions& options, const std::string& name_prefix, tsl::TfDeviceId tf_device_id, @@ -417,6 +430,14 @@ class BaseGPUDeviceFactory : public DeviceFactory { xla::LocalDeviceState* xla_local_device_state, Allocator* gpu_allocator, std::vector>* devices); +#else + absl::Status CreateGPUDevice(const SessionOptions& options, + const std::string& name_prefix, + tsl::TfDeviceId tf_device_id, + const DeviceLocality& dev_locality, + Allocator* gpu_allocator, + std::vector>* devices); +#endif // TF_GPU_USE_PJRT virtual std::unique_ptr CreateGPUDevice( const SessionOptions& options, const std::string& name, diff --git a/tensorflow/core/common_runtime/gpu/gpu_device_test.cc b/tensorflow/core/common_runtime/gpu/gpu_device_test.cc index fd6c180f4ca5db..f191c8f32b082f 100644 --- a/tensorflow/core/common_runtime/gpu/gpu_device_test.cc +++ b/tensorflow/core/common_runtime/gpu/gpu_device_test.cc @@ -15,10 +15,13 @@ limitations under the License. #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM +#if (defined(PLATFORM_GOOGLE) && defined(TF_PLATFORM_LINUX_X86_64)) +#define TF_GPU_USE_PJRT +#endif // PLATFORM_GOOGLE && TF_PLATFORM_LINUX_X86_64 + #include "tensorflow/core/common_runtime/gpu/gpu_device.h" #include "absl/synchronization/notification.h" -#include "xla/pjrt/pjrt_client.h" #include "xla/stream_executor/gpu/gpu_cudamallocasync_allocator.h" #include "xla/stream_executor/gpu/gpu_init.h" #include "xla/tsl/framework/device_id.h" @@ -29,7 +32,11 @@ limitations under the License. #include "tensorflow/core/platform/random.h" #include "tensorflow/core/platform/status.h" #include "tensorflow/core/platform/test.h" + +#ifdef TF_GPU_USE_PJRT +#include "xla/pjrt/pjrt_client.h" #include "tensorflow/core/tfrt/common/pjrt_util.h" +#endif // TF_GPU_USE_PJRT #if GOOGLE_CUDA // Needed for CUDA_VERSION preprocessor directive @@ -310,6 +317,7 @@ TEST_F(GPUDeviceTest, VirtualDeviceConfigConflictsWithVisibleDeviceList) { "list."); } +#ifdef TF_GPU_USE_PJRT TEST_F(GPUDeviceTest, GpuDeviceWithPjrt) { SessionOptions opts = MakeSessionOptions("0"); std::vector> devices; @@ -319,8 +327,9 @@ TEST_F(GPUDeviceTest, GpuDeviceWithPjrt) { EXPECT_GE(devices[0]->attributes().memory_limit(), 0); EXPECT_EQ(static_cast(devices[0].get())->priority(), 0); auto pjrt_client = GetPjRtClient(DeviceType(DEVICE_GPU)); - TF_EXPECT_OK(pjrt_client.status()); + EXPECT_OK(pjrt_client.status()); } +#endif // TF_GPU_USE_PJRT TEST_F(GPUDeviceTest, EmptyVirtualDeviceConfig) { // It'll create single virtual device when the virtual device config is empty. diff --git a/tensorflow/core/common_runtime/gpu/gpu_process_state.cc b/tensorflow/core/common_runtime/gpu/gpu_process_state.cc index 2e32288b24ca12..15fd92a873bea0 100644 --- a/tensorflow/core/common_runtime/gpu/gpu_process_state.cc +++ b/tensorflow/core/common_runtime/gpu/gpu_process_state.cc @@ -13,9 +13,14 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +// TODO(b/282059652): Merge google internal and open-source code path once TF +// dependency issue is resolved. +#if (defined(PLATFORM_GOOGLE) && defined(TF_PLATFORM_LINUX_X86_64)) +#define TF_GPU_USE_PJRT +#endif // PLATFORM_GOOGLE && TF_PLATFORM_LINUX_X86_64 + #include "tensorflow/core/common_runtime/gpu/gpu_process_state.h" -#include #include #include #include @@ -23,11 +28,9 @@ limitations under the License. #include #include "absl/container/flat_hash_set.h" -#include "absl/memory/memory.h" #include "xla/stream_executor/gpu/gpu_init.h" #include "xla/stream_executor/integrations/device_mem_allocator.h" #include "xla/stream_executor/integrations/stream_executor_allocator.h" -#include "xla/stream_executor/memory_space.h" #include "xla/stream_executor/stream_executor.h" #include "xla/tsl/framework/allocator.h" #include "xla/tsl/framework/bfc_allocator.h" @@ -42,14 +45,9 @@ limitations under the License. #include "tensorflow/core/common_runtime/gpu/gpu_debug_allocator.h" #include "tensorflow/core/common_runtime/gpu/gpu_id_manager.h" #include "tensorflow/core/common_runtime/pool_allocator.h" -#include "tensorflow/core/common_runtime/process_state.h" #include "tensorflow/core/common_runtime/shared_counter.h" -#include "tensorflow/core/framework/allocator.h" #include "tensorflow/core/framework/log_memory.h" #include "tensorflow/core/framework/tracking_allocator.h" -#include "tensorflow/core/framework/types.h" -#include "tensorflow/core/platform/mutex.h" -#include "tensorflow/core/platform/numa.h" #include "tsl/platform/mutex.h" #include "tsl/platform/strcat.h" @@ -57,11 +55,6 @@ limitations under the License. #include "xla/stream_executor/gpu/gpu_cudamallocasync_allocator.h" #endif // GOOGLE_CUDA -#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \ - (defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM) -#include "tensorflow/core/tfrt/common/pjrt_util.h" -#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM - namespace tensorflow { // NOLINTNEXTLINE(clang-diagnostic-unused-function) @@ -235,16 +228,31 @@ Allocator* GPUProcessState::GetGPUAllocator( recording_allocator = new internal::RecordingAllocator( &process_state_->mem_desc_map_, gpu_allocator, md, &mu_); } - allocator_parts.allocator.reset(gpu_allocator); +#ifdef TF_GPU_USE_PJRT + // Owning allocator is not set if `allocator_not_owned` is set. + allocator_parts.allocator_not_owned = gpu_allocator; allocator_parts.counter.reset(timing_counter); allocator_parts.bfc_allocator = gpu_bfc_allocator.release(); allocator_parts.sub_allocator = sub_allocator_ptr; allocator_parts.recording_allocator.reset(recording_allocator); +#else + allocator_parts = { + std::unique_ptr(gpu_allocator), + std::unique_ptr(timing_counter), + gpu_bfc_allocator.release(), + sub_allocator_ptr, + std::unique_ptr(recording_allocator), + }; +#endif // TF_GPU_USE_PJRT } if (process_state_->ProcessState::FLAGS_brain_gpu_record_mem_types) { return allocator_parts.recording_allocator.get(); } else { +#ifdef TF_GPU_USE_PJRT + return allocator_parts.allocator_not_owned; +#else return allocator_parts.allocator.get(); +#endif // TF_GPU_USE_PJRT } #else LOG(FATAL) << "GPUAllocator unavailable. Not compiled with --config=cuda or " @@ -305,7 +313,11 @@ Allocator* GPUProcessState::GetGpuHostAllocator(const GPUOptions& options, return gpu_host_allocators_[0].recording_allocator.get(); } if (static_cast(gpu_host_allocators_.size()) > numa_node) { +#ifdef TF_GPU_USE_PJRT + return gpu_host_allocators_[0].allocator_not_owned; +#else return gpu_host_allocators_[0].allocator.get(); +#endif // TF_GPU_USE_PJRT } } @@ -319,7 +331,11 @@ Allocator* GPUProcessState::GetGpuHostAllocator(const GPUOptions& options, // it knows is valid. se::StreamExecutor* se = nullptr; for (int i = 0; i < static_cast(gpu_allocators_.size()); ++i) { +#ifdef TF_GPU_USE_PJRT + if (gpu_allocators_[i].allocator_not_owned != nullptr) { +#else if (gpu_allocators_[i].allocator != nullptr) { +#endif // TF_GPU_USE_PJRT se = DeviceIdUtil::ExecutorForTfDeviceId( DEVICE_GPU, se::GPUMachineManager(), tsl::TfDeviceId(i)) .value(); @@ -368,13 +384,22 @@ Allocator* GPUProcessState::GetGpuHostAllocator(const GPUOptions& options, // at the cost of performance. allocator = new TrackingAllocator(allocator, true); } +#ifdef TF_GPU_USE_PJRT + // Ownership of the GPU host allocator will be transferred to PJRT. AllocatorParts gpu_host_allocator({ - /*allocator=*/std::unique_ptr(allocator), + /*allocator=*/nullptr, std::unique_ptr(nullptr), /*bfc_allocator=*/nullptr, sub_allocator, /*recording_allocator=*/nullptr, + /*allocator_not_owned=*/allocator, }); +#else + AllocatorParts gpu_host_allocator({std::unique_ptr(allocator), + std::unique_ptr(nullptr), + /*bfc_allocator=*/nullptr, sub_allocator, + /*recording_allocator=*/nullptr}); +#endif // TF_GPU_USE_PJRT gpu_host_allocators_.push_back(std::move(gpu_host_allocator)); AllocatorParts& allocator_parts = gpu_host_allocators_.back(); if (process_state_->ProcessState::FLAGS_brain_gpu_record_mem_types) { @@ -385,13 +410,18 @@ Allocator* GPUProcessState::GetGpuHostAllocator(const GPUOptions& options, md.nic_registered = false; allocator_parts.recording_allocator = std::make_unique( - &process_state_->mem_desc_map_, allocator, md, &mu_); + &process_state_->mem_desc_map_, allocator_parts.allocator.get(), + md, &mu_); } } if (process_state_->ProcessState::FLAGS_brain_gpu_record_mem_types) { return gpu_host_allocators_[0].recording_allocator.get(); } else { +#ifdef TF_GPU_USE_PJRT + return gpu_host_allocators_[0].allocator_not_owned; +#else return gpu_host_allocators_[0].allocator.get(); +#endif // TF_GPU_USE_PJRT } } @@ -442,11 +472,6 @@ void GPUProcessState::AddGpuHostFreeVisitor( } void GPUProcessState::TestOnlyReset() { -#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \ - (defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM) - TestOnlyResetPjRtClientInTFGlobalResourceManager(DeviceType(DEVICE_GPU)) - .IgnoreError(); -#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM if (process_state_) { process_state_->ProcessState::TestOnlyReset(); } diff --git a/tensorflow/core/common_runtime/gpu/gpu_process_state.h b/tensorflow/core/common_runtime/gpu/gpu_process_state.h index aff7a6b3ff3555..19f8448ed5e67b 100644 --- a/tensorflow/core/common_runtime/gpu/gpu_process_state.h +++ b/tensorflow/core/common_runtime/gpu/gpu_process_state.h @@ -16,8 +16,16 @@ limitations under the License. #ifndef TENSORFLOW_CORE_COMMON_RUNTIME_GPU_GPU_PROCESS_STATE_H_ #define TENSORFLOW_CORE_COMMON_RUNTIME_GPU_GPU_PROCESS_STATE_H_ -#include +// TODO(b/282059652): Merge google internal and open-source code path once TF +// dependency issue is resolved. +#if (defined(PLATFORM_GOOGLE) && defined(TF_PLATFORM_LINUX_X86_64)) +#define TF_GPU_USE_PJRT +#endif // PLATFORM_GOOGLE && TF_PLATFORM_LINUX_X86_64 + +#include +#include #include +#include #include #include "xla/tsl/framework/device_id.h" @@ -26,6 +34,7 @@ limitations under the License. #include "tensorflow/core/framework/allocator.h" #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/thread_annotations.h" +#include "tensorflow/core/platform/types.h" #include "tensorflow/core/protobuf/config.pb.h" namespace tensorflow { @@ -39,7 +48,7 @@ class GPUProcessState { // If ps == nullptr, returns pointer to the single instance of this class to // be used within this process. // - // If ps != nullptr, accepts a value to be returned by all subsequent calls. + // If ps != nullptrs, accepts a value to be returned by all subsequent calls. // A non-null ps may ONLY be provided during program static storage // initialization. Must not be called more than once with a non-null ps. // @@ -152,9 +161,16 @@ class GPUProcessState { struct AllocatorParts { std::unique_ptr allocator; std::unique_ptr counter; - GPUBFCAllocator* bfc_allocator = nullptr; - SubAllocator* sub_allocator = nullptr; // owned by allocator + GPUBFCAllocator* bfc_allocator; + SubAllocator* sub_allocator; // owned by allocator std::unique_ptr recording_allocator; + +#ifdef TF_GPU_USE_PJRT + // Not owning GPU allocator. The allocator is owned by PJRT. If + // `allocator_not_owned` is set, `allocator` owned by AllocatorParts won't + // be set. + Allocator* allocator_not_owned; +#endif // TF_GPU_USE_PJRT }; std::vector gpu_allocators_ TF_GUARDED_BY(mu_); std::vector> gpu_visitors_ diff --git a/tensorflow/core/common_runtime/gpu/gpu_util.cc b/tensorflow/core/common_runtime/gpu/gpu_util.cc index 986b0969735b65..8ab059105c08fb 100644 --- a/tensorflow/core/common_runtime/gpu/gpu_util.cc +++ b/tensorflow/core/common_runtime/gpu/gpu_util.cc @@ -19,7 +19,9 @@ limitations under the License. #include #include #include +#include #include +#include #include "absl/log/check.h" #include "absl/log/log.h" @@ -27,37 +29,46 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/synchronization/notification.h" -#include "xla/stream_executor/device_memory.h" +#include "xla/layout_util.h" +#include "xla/shape.h" +#include "tensorflow/core/framework/types.pb.h" + +// TODO(b/282059652): Merge google internal and open-source code path once TF +// dependency issue is resolved. +#if (defined(PLATFORM_GOOGLE) && defined(TF_PLATFORM_LINUX_X86_64)) +#define TF_GPU_USE_PJRT +#endif // PLATFORM_GOOGLE && TF_PLATFORM_LINUX_X86_64 + +#ifdef TF_GPU_USE_PJRT +#include "tensorflow/compiler/jit/pjrt_tensor_buffer.h" +#include "tensorflow/compiler/tf2xla/literal_util.h" +#include "xla/future.h" +#include "xla/literal.h" +#endif // TF_GPU_USE_PJRT + #include "tensorflow/core/common_runtime/copy_tensor.h" +#include "tensorflow/core/common_runtime/device.h" #include "tensorflow/core/common_runtime/device/device_event_mgr.h" #include "tensorflow/core/common_runtime/dma_helper.h" #include "tensorflow/core/common_runtime/gpu/gpu_process_state.h" #include "tensorflow/core/common_runtime/gpu_device_context.h" -#include "tensorflow/core/framework/allocator.h" -#include "tensorflow/core/framework/device.h" -#include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/log_memory.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor.pb.h" #include "tensorflow/core/framework/tensor_reference.h" #include "tensorflow/core/framework/types.h" -#include "tensorflow/core/framework/types.pb.h" -#include "tensorflow/core/platform/hash.h" -#include "tensorflow/core/platform/status.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/refcount.h" +#include "tensorflow/core/lib/hash/hash.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/lib/strings/stringprintf.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/stream_executor.h" #include "tensorflow/core/platform/tensor_coding.h" +#include "tensorflow/core/profiler/lib/scoped_annotation.h" #include "tensorflow/core/util/util.h" -#include "tsl/profiler/lib/scoped_annotation.h" #include "tsl/profiler/lib/traceme.h" -#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \ - (defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM) -#include "tensorflow/compiler/jit/pjrt_tensor_buffer.h" -#include "tensorflow/compiler/tf2xla/literal_util.h" -#include "xla/future.h" -#include "xla/layout_util.h" -#include "xla/literal.h" -#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM - // IMPLEMENTATION NOTE: // // 1. Within this module, we intentionally LOG(FATAL) if any stream @@ -350,13 +361,12 @@ void GPUUtil::CopyGPUTensorToCPU(Device* gpu_device, } } -#if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \ - (defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM) +#ifdef TF_GPU_USE_PJRT // The above `WaitFor(send_stream)` for the PjRt case eliminates race // conditions caused by either non-XLA ops that have not finished or a case in // the PJRT client implementation where an event on the buffer is not sent // properly. A possible future improvement is to specifically handle the - // relevant case(s) and move this codeblock to the start of + // relevant case(s) and move this TF_GPU_USE_PJRT codeblock to the start of // this function (avoiding the need for `WaitFor(send_stream)` for the // PjRt case). const PjRtTensorBuffer* pjrt_tensor_buffer = @@ -368,10 +378,6 @@ void GPUUtil::CopyGPUTensorToCPU(Device* gpu_device, literal = std::make_unique(); auto status = tensorflow::HostTensorToMutableBorrowingLiteral( cpu_tensor, literal.get()); - if (!status.ok()) { - done(status); - return; - } } else { xla::Shape shape = pjrt_tensor_buffer->pjrt_buffer()->on_device_shape(); *shape.mutable_layout() = @@ -385,7 +391,7 @@ void GPUUtil::CopyGPUTensorToCPU(Device* gpu_device, done](const absl::Status& status) { done(status); }); return; } -#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM +#endif // TF_GPU_USE_PJRT VLOG(1) << "CopyGPUTensorToCPU using AcceleratorDeviceInfo"; const int64_t total_bytes = gpu_tensor->TotalBytes(); diff --git a/tensorflow/core/config/BUILD b/tensorflow/core/config/BUILD index 7be8f9e203f17a..a0959f5073a19a 100644 --- a/tensorflow/core/config/BUILD +++ b/tensorflow/core/config/BUILD @@ -47,7 +47,7 @@ filegroup( "flag_defs.h", "flags.h", ], - visibility = ["//tensorflow:internal"], + visibility = ["//visibility:private"], ) filegroup( diff --git a/tensorflow/core/grappler/costs/graph_memory.cc b/tensorflow/core/grappler/costs/graph_memory.cc index a8ed2c8fd4229e..18530d5a5b005f 100644 --- a/tensorflow/core/grappler/costs/graph_memory.cc +++ b/tensorflow/core/grappler/costs/graph_memory.cc @@ -173,7 +173,8 @@ void GraphMemory::InferFromTrace(const StepStats& timeline) { } for (const auto& dev_stats : timeline.dev_stats()) { const std::string& device_name = dev_stats.device(); - const bool is_gpu = (device_name.find("GPU:") || device_name.find("gpu:")); + const bool is_gpu = device_name.find("GPU:") != std::string::npos || + device_name.find("gpu:") != std::string::npos; std::deque& device_tensors = live_tensors_per_device[dev_stats.device()]; for (const auto& node_stats : dev_stats.node_stats()) { diff --git a/tensorflow/core/grappler/costs/graph_memory_test.cc b/tensorflow/core/grappler/costs/graph_memory_test.cc index 82c6e06df0f33f..1e99bc545b949e 100644 --- a/tensorflow/core/grappler/costs/graph_memory_test.cc +++ b/tensorflow/core/grappler/costs/graph_memory_test.cc @@ -14,9 +14,11 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/core/grappler/costs/graph_memory.h" + #include "tensorflow/cc/ops/standard_ops.h" #include "tensorflow/core/grappler/grappler_item.h" #include "tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.h" +#include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/test.h" namespace tensorflow { @@ -190,6 +192,43 @@ TEST_F(GraphMemoryTest, GpuSwapping) { } } +TEST_F(GraphMemoryTest, CpuIgnoresSwapToHost) { + // _swap_to_host only applies to GPU devices: setting it on a CPU node must + // not change the CPU memory estimate. + TrivialTestGraphInputYielder fake_input(4, 2, 1024 * 1024, false, {"/CPU:0"}); + GrapplerItem item; + CHECK(fake_input.NextItem(&item)); + item.feed.clear(); + + GraphMemory memory(item); + TF_ASSERT_OK(memory.InferStatically(devices_)); + const GraphMemory::MemoryUsage& cpu_mem = memory.GetPeakMemoryUsage("/CPU:0"); + std::set cpu_tensors; + for (const auto& t : cpu_mem.live_tensors) { + cpu_tensors.insert(absl::StrCat(t.node, ":", t.output_id)); + } + + bool node_found = false; + for (auto& node : *item.graph.mutable_node()) { + if (node.name() == "AddN_1") { + (*node.mutable_attr())["_swap_to_host"].mutable_list()->add_i(0); + node_found = true; + } + } + ASSERT_TRUE(node_found); + GraphMemory swapped_memory(item); + TF_ASSERT_OK(swapped_memory.InferStatically(devices_)); + const GraphMemory::MemoryUsage& swapped_cpu_mem = + swapped_memory.GetPeakMemoryUsage("/CPU:0"); + std::set swapped_cpu_tensors; + for (const auto& t : swapped_cpu_mem.live_tensors) { + swapped_cpu_tensors.insert(absl::StrCat(t.node, ":", t.output_id)); + } + + EXPECT_EQ(cpu_mem.used_memory, swapped_cpu_mem.used_memory); + EXPECT_EQ(cpu_tensors, swapped_cpu_tensors); +} + TEST_F(GraphMemoryTest, CtrlDependencies) { // Build a simple graph with a control dependency. Scope s = Scope::NewRootScope(); diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD index 7389958ff6094b..dcb26690b0099b 100644 --- a/tensorflow/core/kernels/BUILD +++ b/tensorflow/core/kernels/BUILD @@ -2801,7 +2801,7 @@ tf_kernel_library( tf_kernel_library( name = "tensor_array_ops", prefix = "tensor_array_ops", - deps = DATA_FLOW_DEPS, + deps = DATA_FLOW_DEPS + ["//tensorflow/core/util:overflow"], ) DYNAMIC_DEPS = [ diff --git a/tensorflow/core/kernels/conv_ops_fused_image_transform.cc b/tensorflow/core/kernels/conv_ops_fused_image_transform.cc index 1df29d70b320d6..4ba8566addd364 100644 --- a/tensorflow/core/kernels/conv_ops_fused_image_transform.cc +++ b/tensorflow/core/kernels/conv_ops_fused_image_transform.cc @@ -649,6 +649,10 @@ class FusedResizeConv2DUsingGemmOp : public OpKernel { // Input tensor is of the following dimensions: // [ batch, in_rows, in_cols, in_depth ] const Tensor& input = context->input(0); + OP_REQUIRES( + context, input.dims() == 4, + absl::InvalidArgumentError(absl::StrCat("input must be 4-dimensional: ", + input.shape().DebugString()))); OP_REQUIRES(context, (input.shape().num_elements() > 0), absl::InvalidArgumentError("Input tensor can't be empty")); @@ -764,10 +768,6 @@ class FusedResizeConv2DUsingGemmOp : public OpKernel { const Tensor& filter = context->input(filter_index); // For 2D convolution, there should be 4 dimensions. - OP_REQUIRES( - context, padded_shape.dims() == 4, - absl::InvalidArgumentError(absl::StrCat("input must be 4-dimensional", - padded_shape.DebugString()))); OP_REQUIRES( context, filter.dims() == 4, absl::InvalidArgumentError(absl::StrCat( diff --git a/tensorflow/core/kernels/conv_ops_test.cc b/tensorflow/core/kernels/conv_ops_test.cc index caff583b570092..28713763cba8a3 100644 --- a/tensorflow/core/kernels/conv_ops_test.cc +++ b/tensorflow/core/kernels/conv_ops_test.cc @@ -367,6 +367,27 @@ TEST_F(FusedResizePadConvOpTest, NoResizePadOnlySymmetricComparative) { "SAME", DT_FLOAT); } +TEST_F(FusedResizePadConvOpTest, InvalidInputRank) { + TF_ASSERT_OK(NodeDefBuilder("fused_pad_conv_op", "FusedPadConv2D") + .Input(FakeInput(DT_FLOAT)) + .Input(FakeInput(DT_INT32)) + .Input(FakeInput(DT_FLOAT)) + .Attr("T", DT_FLOAT) + .Attr("mode", "REFLECT") + .Attr("strides", {1, 1, 1, 1}) + .Attr("padding", "VALID") + .Finalize(node_def())); + TF_ASSERT_OK(InitOp()); + AddInputFromArray(TensorShape({3, 2}), {0, 0, 0, 0, 0, 0}); + AddInputFromArray(TensorShape({4, 2}), {0, 0, 1, 1, 1, 1, 0, 0}); + AddInputFromArray(TensorShape({2, 2, 1, 1}), {0, 0, 0, 0}); + + const absl::Status status = RunOpKernel(); + EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument); + EXPECT_NE(status.message().find("input must be 4-dimensional"), + std::string::npos); +} + class ConvOpTest : public OpsTestBase { protected: void HandwrittenConv() { diff --git a/tensorflow/core/kernels/cwise_ops.h b/tensorflow/core/kernels/cwise_ops.h index 7c0c3f88fe4a74..2a138de8473323 100644 --- a/tensorflow/core/kernels/cwise_ops.h +++ b/tensorflow/core/kernels/cwise_ops.h @@ -653,7 +653,11 @@ struct xlogy_op { scalar_log_op log_op; Packet log_y = log_op.packetOp(y); Packet x_log_y = pmul(x, log_y); - return pselect(mask, x, x_log_y); + // Select zeros rather than x. An x can compare equal to zero without being + // +0: -0 does, and so does a subnormal when denormals are flushed, as they + // are in TensorFlow's kernels. Returning x would leak either one, while + // the scalar path above returns +0 for both. + return pselect(mask, zeros, x_log_y); } }; @@ -683,7 +687,8 @@ struct xlog1py_op { scalar_log1p_op log1p_op; Packet log1p_y = log1p_op.packetOp(y); Packet x_log1p_y = pmul(x, log1p_y); - return pselect(mask, x, x_log1p_y); + // Select zeros rather than x; see xlogy_op. + return pselect(mask, zeros, x_log1p_y); } }; @@ -715,7 +720,8 @@ struct xdivy_op { Packet zeros = pzero(x); Packet mask = pcmp_eq(x, zeros); Packet x_div_y = pdiv(x, y); - return pselect(mask, x, x_div_y); + // Select zeros rather than x; see xlogy_op. + return pselect(mask, zeros, x_div_y); } }; @@ -876,6 +882,42 @@ struct functor_traits> { }; }; +// igamma(a, x) = P(a, x) is defined only for a > 0, x >= 0. Eigen's +// scalar_igamma_op short-circuits to 0 when x == 0 before the domain check +// fires, so a <= 0 with x == 0 silently returns 0 instead of NaN. The +// wrapper restores the mathematically correct NaN for out-of-domain inputs. +template +struct igamma_op : binary_op_base { + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar + operator()(const Scalar& a, const Scalar& x) const { + if (x == Scalar(0) && !(a > Scalar(0))) { + return Eigen::NumTraits::quiet_NaN(); + } + return Eigen::internal::scalar_igamma_op()(a, x); + } + template + EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet packetOp(const Packet& a, + const Packet& x) const { + Packet zeros = pzero(x); + Packet x_is_zero = pcmp_eq(x, zeros); + Packet a_gt_zero = pcmp_lt(zeros, a); + Packet domain_error = pandnot(x_is_zero, a_gt_zero); + Packet nan = pset1(Eigen::NumTraits::quiet_NaN()); + Packet igamma_val = + Eigen::internal::scalar_igamma_op().packetOp(a, x); + return pselect(domain_error, nan, igamma_val); + } +}; + +template +struct functor_traits> { + enum { + Cost = functor_traits>::Cost + + Eigen::NumTraits::AddCost, + PacketAccess = functor_traits>::PacketAccess, + }; +}; + } // end namespace internal } // end namespace Eigen @@ -1209,7 +1251,7 @@ struct minimum : base> {}; template -struct igamma : base> {}; +struct igamma : base> {}; template struct random_gamma_grad diff --git a/tensorflow/core/kernels/depthwise_conv_grad_op.cc b/tensorflow/core/kernels/depthwise_conv_grad_op.cc index 34a326e685f625..6028a38a625cd1 100644 --- a/tensorflow/core/kernels/depthwise_conv_grad_op.cc +++ b/tensorflow/core/kernels/depthwise_conv_grad_op.cc @@ -1147,6 +1147,14 @@ class DepthwiseConv2dNativeBackpropFilterOp : public OpKernel { // If there is nothing to compute, return. if (out_backprop.shape().num_elements() == 0) { + // The output is sized from `filter_shape`, which is independent of + // `out_backprop`, so it can be non-empty on this path. Nothing below + // writes it and `forward_input_or_allocate_output` does not initialize + // the buffer, so it would be returned to the caller uninitialized. A + // filter that received no contribution has a gradient of exactly zero. + if (filter_backprop->NumElements() > 0) { + filter_backprop->template flat().setZero(); + } return; } diff --git a/tensorflow/core/kernels/dynamic_partition_op_gpu.cu.cc b/tensorflow/core/kernels/dynamic_partition_op_gpu.cu.cc index 51db545a580bdd..fb88e09e5e3dba 100644 --- a/tensorflow/core/kernels/dynamic_partition_op_gpu.cu.cc +++ b/tensorflow/core/kernels/dynamic_partition_op_gpu.cu.cc @@ -61,7 +61,9 @@ namespace { template __global__ void RangeInitKernel(const T start, const T delta, const int32_t size, T* out) { - GPU_1D_KERNEL_LOOP(i, size) { out[i] = start + i * delta; } + for (int32_t i : GpuGridRangeX(size)) { + out[i] = start + i * delta; + } } __global__ void MoveValuesKernel(const int32_t* keys, const int32_t* values, diff --git a/tensorflow/core/kernels/gpu_prim_helpers.h b/tensorflow/core/kernels/gpu_prim_helpers.h index d5bfdaba2c85fc..cf24deaa48a495 100644 --- a/tensorflow/core/kernels/gpu_prim_helpers.h +++ b/tensorflow/core/kernels/gpu_prim_helpers.h @@ -33,7 +33,9 @@ namespace detail { template __global__ void RangeInitKernel(const T start, const T delta, const T size, T* out) { - GPU_1D_KERNEL_LOOP(i, size) { out[i] = start + i * delta; } + for (T i : GpuGridRangeX(size)) { + out[i] = start + i * delta; + } } // Initialize out with range start, start + delta, start + 2 * delta, ... diff --git a/tensorflow/core/kernels/mlir_generated/build_defs.bzl b/tensorflow/core/kernels/mlir_generated/build_defs.bzl index fd207b936288eb..a2bb4e21b8b42c 100644 --- a/tensorflow/core/kernels/mlir_generated/build_defs.bzl +++ b/tensorflow/core/kernels/mlir_generated/build_defs.bzl @@ -181,9 +181,9 @@ def _gen_kernel_bin_impl(ctx): "--arch=%s" % ",".join(ctx.attr.gpu_archs), "--input=%s" % ctx.file.mlir_op.path, "--output=%s" % gpu_bin.path, - "--enable_ftz=%s" % (ctx.attr.data_type == "f32"), - "--jit_i64_indexed_for_large_tensors=%s" % ctx.attr.jit_i64_indexed_for_large_tensors, - "--jit=%s" % ctx.attr.jit, + "--enable_ftz=%s" % ("true" if ctx.attr.data_type == "f32" else "false"), + "--jit_i64_indexed_for_large_tensors=%s" % ("true" if ctx.attr.jit_i64_indexed_for_large_tensors else "false"), + "--jit=%s" % ("true" if ctx.attr.jit else "false"), ], use_default_shell_env = True, mnemonic = "compile", @@ -406,8 +406,8 @@ def _gen_kernel_library( ), "--arch={}".format(gpu_arch_option), "--tile_sizes=%s" % typed_tile_size, - "--jit_i64_indexed_for_large_tensors=%s" % jit_i64_indexed_for_large_tensors, - "--enable_ftz=%s" % (type == "f32"), + "--jit_i64_indexed_for_large_tensors=%s" % ("true" if jit_i64_indexed_for_large_tensors else "false"), + "--enable_ftz=%s" % ("true" if type == "f32" else "false"), ] if typed_unroll_factors: test_args.append("--unroll_factors=%s" % typed_unroll_factors) diff --git a/tensorflow/core/kernels/ragged_gather_op.cc b/tensorflow/core/kernels/ragged_gather_op.cc index 72e5c840e86ad5..8e12616d26ad23 100644 --- a/tensorflow/core/kernels/ragged_gather_op.cc +++ b/tensorflow/core/kernels/ragged_gather_op.cc @@ -12,6 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include #include #include #include @@ -40,10 +41,10 @@ void WriteValueSlices( const auto& params_dense_values = params_dense_values_in.flat_outer_dims(); auto values = values_out->flat_outer_dims(); - int out_pos = 0; + SPLITS_TYPE out_pos = 0; for (const auto& slice : value_slices) { - for (int i = slice.first; i < slice.second; ++i) { - for (int j = 0; j < value_size; ++j) { + for (SPLITS_TYPE i = slice.first; i < slice.second; ++i) { + for (SPLITS_TYPE j = 0; j < value_size; ++j) { values(out_pos, j) = params_dense_values(i, j); } ++out_pos; @@ -171,17 +172,25 @@ class RaggedGatherOpBase : public OpKernel { // should add a new split point to out_splits that is 4 greater than the // previous split point in out_splits. for (int i = 0; i < indices.size(); ++i) { - int start = indices(i); - int limit = indices(i) + 1; + SPLITS_TYPE start = indices(i); + SPLITS_TYPE limit = indices(i) + 1; // Copy splits. for (int dim = 0; dim < params_nested_splits.size(); ++dim) { const auto& splits = params_nested_splits[dim]; int out_dim = dim + indices_in.dims() - 1; if (out_dim >= 0) { - SPLITS_TYPE delta = out_splits->at(out_dim).back() - splits(start); - for (int j = start; j < limit; ++j) { - out_splits->at(out_dim).push_back(splits(j + 1) + delta); + // Splits are validated to be non-negative and sorted, so `length` is + // in [0, max]. Check before adding so neither int32 nor int64 Tsplits + // can overflow. + const SPLITS_TYPE last = out_splits->at(out_dim).back(); + for (SPLITS_TYPE j = start; j < limit; ++j) { + const SPLITS_TYPE length = splits(j + 1) - splits(start); + if (length > std::numeric_limits::max() - last) { + return absl::InvalidArgumentError( + "Output splits value exceeds limits of Tsplits"); + } + out_splits->at(out_dim).push_back(last + length); } } start = splits(start); diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc index 6b8f50774a0cf1..8a1c33a8f88424 100644 --- a/tensorflow/core/kernels/resource_variable_ops.cc +++ b/tensorflow/core/kernels/resource_variable_ops.cc @@ -1221,6 +1221,18 @@ class ResourceScatterUpdateOp : public OpKernel { "updates.shape ", updates.shape().DebugString(), ", indices.shape ", indices.shape().DebugString(), ", params.shape ", params->shape().DebugString()))); + for (int d = 0; d < updates.dims(); ++d) { + const int64_t expected_dim = + d < indices.dims() ? indices.dim_size(d) + : params->dim_size(d - indices.dims() + 1); + OP_REQUIRES(c, updates.dim_size(d) == expected_dim, + absl::InvalidArgumentError(absl::StrCat( + "Must have updates.shape = indices.shape + " + "params.shape[1:] or updates.shape = [], got ", + "updates.shape ", updates.shape().DebugString(), + ", indices.shape ", indices.shape().DebugString(), + ", params.shape ", params->shape().DebugString()))); + } // Check that we have enough index space const int64_t N_big = indices.NumElements(); diff --git a/tensorflow/core/kernels/scatter_functor_gpu.cu.h b/tensorflow/core/kernels/scatter_functor_gpu.cu.h index e4f43d51b46075..41fe2ed65263f4 100644 --- a/tensorflow/core/kernels/scatter_functor_gpu.cu.h +++ b/tensorflow/core/kernels/scatter_functor_gpu.cu.h @@ -92,12 +92,14 @@ __global__ void ScatterOpCustomKernel(T* __restrict__ params, GPU_1D_KERNEL_LOOP(i, updates_size) { int indices_i = i / update_block; int updates_i = i; - int param_first_index = indices[indices_i]; + Index param_first_index = indices[indices_i]; if (!(param_first_index >= 0 && param_first_index < first_dim_size)) { // Ignore indices that are out of range. continue; } - int64_t params_i = param_first_index * update_block + (i % update_block); + const int64_t params_i = + static_cast(param_first_index) * update_block + + (i % update_block); body(¶ms[params_i], ldg(updates + updates_i)); } } @@ -113,13 +115,15 @@ __global__ void ScatterScalarOpCustomKernel(T* __restrict__ params, ScatterOpKernelBody body; GPU_1D_KERNEL_LOOP(i, synthesized_updates_size) { int indices_i = i / update_block; - int param_first_index = indices[indices_i]; + Index param_first_index = indices[indices_i]; const T update_val = *update; if (!(param_first_index >= 0 && param_first_index < first_dim_size)) { // Ignore indices that are out of range. continue; } - int params_i = param_first_index * update_block + (i % update_block); + const int64_t params_i = + static_cast(param_first_index) * update_block + + (i % update_block); body(¶ms[params_i], update_val); } } diff --git a/tensorflow/core/kernels/scatter_nd_op.cc b/tensorflow/core/kernels/scatter_nd_op.cc index d56134c19f9d11..49d8236dcd0102 100644 --- a/tensorflow/core/kernels/scatter_nd_op.cc +++ b/tensorflow/core/kernels/scatter_nd_op.cc @@ -134,7 +134,7 @@ class ScatterNdOp : public ScatterOpBase { TensorShapeUtils::MakeShape(vec.data(), vec.size(), &shape)); OP_REQUIRES(c, - ValidEmptyOutputShape(shape_input.NumElements(), + ValidEmptyOutputShape(shape.num_elements(), indices.shape().num_elements(), updates.shape().num_elements()), absl::InvalidArgumentError( diff --git a/tensorflow/core/kernels/scatter_nd_op_cpu_impl.h b/tensorflow/core/kernels/scatter_nd_op_cpu_impl.h index c4cc570b97f07a..84f191264d9830 100644 --- a/tensorflow/core/kernels/scatter_nd_op_cpu_impl.h +++ b/tensorflow/core/kernels/scatter_nd_op_cpu_impl.h @@ -20,10 +20,11 @@ limitations under the License. #define EIGEN_USE_THREADS +#include #include +#include #include "unsupported/Eigen/CXX11/Tensor" // from @eigen_archive - #include "tensorflow/core/framework/bounds_check.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/register_types.h" @@ -127,20 +128,62 @@ struct ScatterNdFunctor { batch_strides[dim + 1] * output_shape_prefix[dim + 1]; } - for (Eigen::DenseIndex loc = 0; loc < batch_size; ++loc) { - Index i = 0; - bool out_of_bounds = false; - for (int dim = 0; dim < IXDIM; ++dim) { - const Index ix_d = internal::SubtleMustCopy(Tindices(loc, dim)); - out_of_bounds |= !FastBoundsCheck(ix_d, output_shape_prefix[dim]); - i += ix_d * batch_strides[dim]; + if (slice_size == 1) { + T* out_data = Toutput.data(); + const T* update_data = Tupdates.data(); + for (Eigen::DenseIndex loc = 0; loc < batch_size; ++loc) { + Index i = 0; + bool out_of_bounds = false; + for (int dim = 0; dim < IXDIM; ++dim) { + const Index ix_d = internal::SubtleMustCopy(Tindices(loc, dim)); + out_of_bounds |= !FastBoundsCheck(ix_d, output_shape_prefix[dim]); + i += ix_d * batch_strides[dim]; + } + if (TF_PREDICT_FALSE(out_of_bounds)) { + error_loc = loc; + continue; + } + if constexpr (OP == scatter_nd_op::UpdateOp::ASSIGN) { + out_data[i] = update_data[loc]; + } else if constexpr (OP == scatter_nd_op::UpdateOp::ADD) { + if constexpr (std::is_same_v) { + out_data[i].append(update_data[loc]); + } else if constexpr (std::is_same_v) { + out_data[i] |= update_data[loc]; + } else { + out_data[i] += update_data[loc]; + } + } else if constexpr (OP == scatter_nd_op::UpdateOp::SUB) { + if constexpr (std::is_same_v) { + out_data[i] ^= update_data[loc]; + } else if constexpr (!std::is_same_v) { + out_data[i] -= update_data[loc]; + } + } else if constexpr (OP == scatter_nd_op::UpdateOp::MIN) { + if constexpr (!std::is_same_v) { + out_data[i] = std::min(out_data[i], update_data[loc]); + } + } else if constexpr (OP == scatter_nd_op::UpdateOp::MAX) { + if constexpr (!std::is_same_v) { + out_data[i] = std::max(out_data[i], update_data[loc]); + } + } } - if (TF_PREDICT_FALSE(out_of_bounds)) { - error_loc = loc; - // Don't break the loop here, but continue to update the rest because - // the caller might ignore bad indices. - continue; - } else { + } else { + for (Eigen::DenseIndex loc = 0; loc < batch_size; ++loc) { + Index i = 0; + bool out_of_bounds = false; + for (int dim = 0; dim < IXDIM; ++dim) { + const Index ix_d = internal::SubtleMustCopy(Tindices(loc, dim)); + out_of_bounds |= !FastBoundsCheck(ix_d, output_shape_prefix[dim]); + i += ix_d * batch_strides[dim]; + } + if (TF_PREDICT_FALSE(out_of_bounds)) { + error_loc = loc; + // Don't break the loop here, but continue to update the rest because + // the caller might ignore bad indices. + continue; + } auto input_chip = Toutput.template chip<0>(i); auto output_chip = input_chip; auto update_chip = Tupdates.template chip<0>(loc); diff --git a/tensorflow/core/kernels/searchsorted_op.cc b/tensorflow/core/kernels/searchsorted_op.cc index 78fa6675cc178f..2e6875b80ccebb 100644 --- a/tensorflow/core/kernels/searchsorted_op.cc +++ b/tensorflow/core/kernels/searchsorted_op.cc @@ -33,6 +33,21 @@ typedef Eigen::ThreadPoolDevice CPUDevice; typedef Eigen::GpuDevice GPUDevice; namespace functor { + +// Comparator that establishes a strict total order consistent with NaN being +// the largest value. Without it, IEEE 754 NaN comparisons are all false and +// lower_bound/upper_bound misplace NaN at position 0 instead of the end. +template +struct NanAwareCompare { + bool operator()(const T& a, const T& b) const { + if (a < b) return true; + if constexpr (!Eigen::NumTraits::IsInteger) { + return !Eigen::numext::isnan(a) && Eigen::numext::isnan(b); + } + return false; + } +}; + template struct UpperBoundFunctor { static absl::Status Compute( @@ -48,7 +63,8 @@ struct UpperBoundFunctor { for (int64_t i = first; i < last; ++i) { output_ptr[i] = std::upper_bound(sorted_inputs_ptr, sorted_inputs_ptr + num_inputs, - values(i + b * num_values)) - + values(i + b * num_values), + NanAwareCompare()) - sorted_inputs_ptr; } } @@ -78,7 +94,8 @@ struct LowerBoundFunctor { for (int64_t i = first; i < last; ++i) { output_ptr[i] = std::lower_bound(sorted_inputs_ptr, sorted_inputs_ptr + num_inputs, - values(i + b * num_values)) - + values(i + b * num_values), + NanAwareCompare()) - sorted_inputs_ptr; } } diff --git a/tensorflow/core/kernels/searchsorted_op_gpu.cu.cc b/tensorflow/core/kernels/searchsorted_op_gpu.cu.cc index df60599d90f32f..343c49a25aa520 100644 --- a/tensorflow/core/kernels/searchsorted_op_gpu.cu.cc +++ b/tensorflow/core/kernels/searchsorted_op_gpu.cu.cc @@ -30,6 +30,56 @@ namespace tensorflow { typedef Eigen::GpuDevice GPUDevice; namespace { + +template +struct GpuNanAwareCompare { + __device__ bool operator()(T a, T b) const { + if (a < b) return true; + if constexpr (!Eigen::NumTraits::IsInteger) { + return !Eigen::numext::isnan(a) && Eigen::numext::isnan(b); + } + return false; + } +}; + +template +__device__ OutType gpu_lower_bound(const T* first, OutType count, T val, + Compare comp) { + const T* orig_first = first; + OutType step = 0; + while (count > 0) { + const T* it = first; + step = count / 2; + it += step; + if (comp(*it, val)) { + first = ++it; + count -= step + 1; + } else { + count = step; + } + } + return first - orig_first; +} + +template +__device__ OutType gpu_upper_bound(const T* first, OutType count, T val, + Compare comp) { + const T* orig_first = first; + OutType step = 0; + while (count > 0) { + const T* it = first; + step = count / 2; + it += step; + if (!comp(val, *it)) { + first = ++it; + count -= step + 1; + } else { + count = step; + } + } + return first - orig_first; +} + template __global__ void UpperBoundKernel(const T* __restrict__ sorted_inputs, int64_t batch_size, int64_t sorted_inputs_size, @@ -39,8 +89,9 @@ __global__ void UpperBoundKernel(const T* __restrict__ sorted_inputs, for (int64_t work_unit_id : GpuGridRangeX(values_size * batch_size)) { int64_t bid = work_unit_id / values_size; T value = values[work_unit_id]; - outputs[work_unit_id] = gpu_helper::upper_bound( - sorted_inputs + bid * sorted_inputs_size, sorted_inputs_size, value); + outputs[work_unit_id] = gpu_upper_bound( + sorted_inputs + bid * sorted_inputs_size, sorted_inputs_size, value, + GpuNanAwareCompare()); } } @@ -53,8 +104,9 @@ __global__ void LowerBoundKernel(const T* __restrict__ sorted_inputs, for (int64_t work_unit_id : GpuGridRangeX(values_size * batch_size)) { int64_t bid = work_unit_id / values_size; T value = values[work_unit_id]; - outputs[work_unit_id] = gpu_helper::lower_bound( - sorted_inputs + bid * sorted_inputs_size, sorted_inputs_size, value); + outputs[work_unit_id] = gpu_lower_bound( + sorted_inputs + bid * sorted_inputs_size, sorted_inputs_size, value, + GpuNanAwareCompare()); } } } // namespace diff --git a/tensorflow/core/kernels/tensor_array_ops.cc b/tensorflow/core/kernels/tensor_array_ops.cc index 74c59ceae54559..fa69fb20abb35b 100644 --- a/tensorflow/core/kernels/tensor_array_ops.cc +++ b/tensorflow/core/kernels/tensor_array_ops.cc @@ -48,6 +48,7 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/thread_annotations.h" #include "tensorflow/core/platform/types.h" +#include "tensorflow/core/util/overflow.h" typedef Eigen::ThreadPoolDevice CPUDevice; #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM @@ -1326,7 +1327,17 @@ class TensorArraySplitOp : public OpKernel { cumulative_lengths.reserve(num_tensors); int64_t total_length = 0; for (int i = 0; i < num_tensors; ++i) { - total_length += tensor_lengths_t(i); + const int64_t length = tensor_lengths_t(i); + OP_REQUIRES(ctx, length >= 0, + absl::InvalidArgumentError(absl::StrCat( + "Expected lengths to be non-negative, but found ", length, + " at index ", i))); + const int64_t next_total_length = + AddWithoutOverflow(total_length, length); + OP_REQUIRES( + ctx, next_total_length >= 0, + absl::InvalidArgumentError("TensorArraySplit lengths overflowed")); + total_length = next_total_length; cumulative_lengths.push_back(total_length); } diff --git a/tensorflow/core/runtime_fallback/kernel/BUILD b/tensorflow/core/runtime_fallback/kernel/BUILD index c4f1ee59a0df01..696444e71c4a17 100644 --- a/tensorflow/core/runtime_fallback/kernel/BUILD +++ b/tensorflow/core/runtime_fallback/kernel/BUILD @@ -339,6 +339,7 @@ cc_library( ":kernel_fallback_tensor", ":kernel_fallback_utils", "//tensorflow/core/framework:tensor_proto_cc", + "//tensorflow/core/profiler/lib:traceme", "//tensorflow/core/runtime_fallback/runtime:op_logger", "//tensorflow/core/runtime_fallback/util:attr_util", "//tensorflow/core/runtime_fallback/util:type_util", @@ -352,22 +353,11 @@ cc_library( "//tensorflow/core/tfrt/utils:tensor_util", "@com_google_absl//absl/base", "@com_google_absl//absl/base:core_headers", - "@com_google_absl//absl/container:flat_hash_set", - "@com_google_absl//absl/container:inlined_vector", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", - "@com_google_absl//absl/log:vlog_is_on", - "@com_google_absl//absl/status", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/strings:string_view", "@llvm-project//llvm:Support", "@tf_runtime//:core_runtime", - "@tf_runtime//:dtype", "@tf_runtime//:hostcontext", "@tf_runtime//:support", "@tf_runtime//:tensor_alwayslink", - "@tsl//tsl/profiler/lib:traceme", - "@tsl//tsl/profiler/lib:traceme_encode", ] + select({ "//tensorflow:android": [ "//tensorflow/core:portable_tensorflow_lib_lite", # TODO(annarev): exclude runtime srcs diff --git a/tensorflow/core/runtime_fallback/kernel/kernel_fallback_execute_compat.cc b/tensorflow/core/runtime_fallback/kernel/kernel_fallback_execute_compat.cc index eed265a33a83fa..587e9b1835a033 100644 --- a/tensorflow/core/runtime_fallback/kernel/kernel_fallback_execute_compat.cc +++ b/tensorflow/core/runtime_fallback/kernel/kernel_fallback_execute_compat.cc @@ -15,7 +15,6 @@ limitations under the License. #include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_execute_compat.h" #include -#include #include #include #include @@ -26,30 +25,15 @@ limitations under the License. #include "absl/base/attributes.h" #include "absl/base/casts.h" -#include "absl/container/flat_hash_set.h" -#include "absl/container/inlined_vector.h" -#include "absl/log/check.h" -#include "absl/log/log.h" -#include "absl/log/vlog_is_on.h" -#include "absl/status/status.h" -#include "absl/strings/str_cat.h" -#include "absl/strings/string_view.h" -#include "absl/strings/strip.h" -#include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" -#include "llvm/Support/raw_ostream.h" -#include "xla/tsl/framework/allocator.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/macros.h" -#include "tensorflow/core/framework/allocator.h" #include "tensorflow/core/framework/logging.h" -#include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor.pb.h" -#include "tensorflow/core/framework/tensor_shape.h" -#include "tensorflow/core/framework/types.h" #include "tensorflow/core/lib/gtl/cleanup.h" -#include "tensorflow/core/platform/tstring.h" +#include "tensorflow/core/platform/errors.h" +#include "tensorflow/core/platform/status.h" +#include "tensorflow/core/platform/types.h" +#include "tensorflow/core/profiler/lib/traceme.h" #include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_compat_request_state.h" #include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_tensor.h" #include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_utils.h" @@ -63,9 +47,7 @@ limitations under the License. #include "tensorflow/core/tfrt/utils/fallback_tensor.h" #include "tensorflow/core/tfrt/utils/tensor_util.h" #include "tensorflow/core/tfrt/utils/utils.h" -#include "tsl/profiler/lib/traceme.h" -#include "tsl/profiler/lib/traceme_encode.h" -#include "tfrt/dtype/dtype.h" // from @tf_runtime +#include "tsl/platform/errors.h" #include "tfrt/host_context/async_dispatch.h" // from @tf_runtime #include "tfrt/host_context/async_value_ref.h" // from @tf_runtime #include "tfrt/host_context/attribute_utils.h" // from @tf_runtime @@ -74,13 +56,10 @@ limitations under the License. #include "tfrt/host_context/execution_context.h" // from @tf_runtime #include "tfrt/host_context/function.h" // from @tf_runtime #include "tfrt/host_context/kernel_registry.h" // from @tf_runtime -#include "tfrt/host_context/kernel_utils.h" // from @tf_runtime #include "tfrt/support/error_util.h" // from @tf_runtime #include "tfrt/support/forward_decls.h" // from @tf_runtime -#include "tfrt/support/ref_count.h" // from @tf_runtime #include "tfrt/support/string_util.h" // from @tf_runtime #include "tfrt/tensor/tensor.h" // from @tf_runtime -#include "tfrt/tensor/tensor_shape.h" // from @tf_runtime namespace tensorflow { namespace tfd { @@ -1010,21 +989,6 @@ class TestAllocator : public tensorflow::AllocatorWrapper { fflush(stdout); return wrapped()->AllocateRaw(alignment, num_bytes, allocation_attr); } - - void* AllocateRawAlignedNew(size_t alignment, size_t num_bytes) override { - std::printf("Using TestAllocator\n"); - fflush(stdout); - return wrapped()->AllocateRawAlignedNew(alignment, num_bytes); - } - - void* AllocateRawAlignedNew( - size_t alignment, size_t num_bytes, - const AllocationAttributes& allocation_attr) override { - std::printf("Using TestAllocator\n"); - fflush(stdout); - return wrapped()->AllocateRawAlignedNew(alignment, num_bytes, - allocation_attr); - } }; tensorflow::Allocator* GetTestAllocator() { diff --git a/tensorflow/core/tfrt/common/BUILD b/tensorflow/core/tfrt/common/BUILD index f744c4bf99717d..0b492b5425ba2c 100644 --- a/tensorflow/core/tfrt/common/BUILD +++ b/tensorflow/core/tfrt/common/BUILD @@ -146,12 +146,13 @@ cc_library( ":pjrt_state", "//tensorflow/core:framework", "//tensorflow/core:framework_types_hdr", + "//tensorflow/core/platform:errors", "//tensorflow/core/platform:refcount", "//tensorflow/core/platform:status", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", + "@tsl//tsl/platform:errors", "@xla//xla/pjrt:pjrt_client", - "@xla//xla/tsl/platform:errors", ], ) @@ -207,10 +208,12 @@ tf_cc_test( ":pjrt_util", "//tensorflow/core:framework", "@com_google_googletest//:gtest_main", + "@tsl//tsl/platform:status_matchers", + "@tsl//tsl/platform:statusor", + "@tsl//tsl/platform:test", "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options", "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client", "@xla//xla/tsl/lib/core:status_test_util", - "@xla//xla/tsl/platform:statusor", "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc", ], ) @@ -271,13 +274,12 @@ cc_library( deps = [ ":pjrt_client_factory_options", ":pjrt_client_factory_registry", - "//tensorflow/compiler/tf2xla:xla_op_registry", "//tensorflow/core:framework_types_hdr", "@com_google_absl//absl/status:statusor", + "@tsl//tsl/platform:statusor", "@xla//xla/pjrt:pjrt_client", "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options", "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client", - "@xla//xla/tsl/platform:statusor", ], alwayslink = True, ) diff --git a/tensorflow/core/tfrt/common/pjrt_client_factory_options.h b/tensorflow/core/tfrt/common/pjrt_client_factory_options.h index e27a69bbff5bc5..70e3092c2df654 100644 --- a/tensorflow/core/tfrt/common/pjrt_client_factory_options.h +++ b/tensorflow/core/tfrt/common/pjrt_client_factory_options.h @@ -19,10 +19,6 @@ limitations under the License. #include #include -namespace Eigen { -struct ThreadPoolDevice; -} // namespace Eigen - namespace xla { // PjrtClientFactoryOptions store arguments to create PJRT client. // Caller is responsible to set option value for corresponding PJRT client @@ -37,7 +33,6 @@ struct PjrtClientFactoryOptions { struct CpuClientCreateOptions { bool asynchronous = false; - const Eigen::ThreadPoolDevice* intra_op_device = nullptr; }; GpuClientCreateOptions gpu_options; CpuClientCreateOptions cpu_options; diff --git a/tensorflow/core/tfrt/common/pjrt_cpu_client_registration.cc b/tensorflow/core/tfrt/common/pjrt_cpu_client_registration.cc index 993dc2577305a3..5cb91b4986d9ac 100644 --- a/tensorflow/core/tfrt/common/pjrt_cpu_client_registration.cc +++ b/tensorflow/core/tfrt/common/pjrt_cpu_client_registration.cc @@ -17,14 +17,13 @@ limitations under the License. #include #include "absl/status/statusor.h" -#include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/pjrt/pjrt_client.h" #include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h" #include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h" -#include "xla/tsl/platform/statusor.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/tfrt/common/pjrt_client_factory_options.h" #include "tensorflow/core/tfrt/common/pjrt_client_factory_registry.h" +#include "tsl/platform/statusor.h" namespace xla { @@ -32,7 +31,6 @@ absl::StatusOr> GetCpuClient( const PjrtClientFactoryOptions& option) { xla::CpuClientOptions cpu_options; cpu_options.asynchronous = option.cpu_options.asynchronous; - cpu_options.intra_op_device = option.cpu_options.intra_op_device; TF_ASSIGN_OR_RETURN(std::unique_ptr client, xla::GetXlaPjrtCpuClient(cpu_options)); @@ -42,7 +40,4 @@ absl::StatusOr> GetCpuClient( REGISTER_PJRT_CLIENT_FACTORY(cpu_client, tensorflow::DEVICE_CPU, GetCpuClient); -REGISTER_PJRT_CLIENT_FACTORY(xla_cpu_client, tensorflow::DEVICE_XLA_CPU, - GetCpuClient); - } // namespace xla diff --git a/tensorflow/core/tfrt/common/pjrt_state.cc b/tensorflow/core/tfrt/common/pjrt_state.cc index 6c8600493acac8..2c6403edc48798 100644 --- a/tensorflow/core/tfrt/common/pjrt_state.cc +++ b/tensorflow/core/tfrt/common/pjrt_state.cc @@ -88,18 +88,6 @@ absl::Status PjRtState::MovePjRtClientToUnused(const DeviceType& device_type) { device_type); } -absl::Status PjRtState::TestOnlyResetPjRtClient(const DeviceType& device_type) { - absl::MutexLock lock(mu_); - if (auto it = clients_.find(device_type); it != clients_.end()) { - unused_.push_back(std::move(it->second)); - clients_.erase(it); - } - if (device_type == DeviceType(DEVICE_GPU)) { - pjrt_gpu_client_creation_info_.reset(); - } - return absl::OkStatus(); -} - absl::Status PjRtState::SetPjRtGpuClientCreationInfo( std::unique_ptr info) { absl::MutexLock lock(mu_); diff --git a/tensorflow/core/tfrt/common/pjrt_state.h b/tensorflow/core/tfrt/common/pjrt_state.h index c229588646bbb5..e0e9f8657bb8a8 100644 --- a/tensorflow/core/tfrt/common/pjrt_state.h +++ b/tensorflow/core/tfrt/common/pjrt_state.h @@ -63,7 +63,6 @@ class PjRtState : public ResourceBase { // Moves PJRT client to `unused_`. The PJRT client moved to `unused_` will not // be returned by `GetPjRtClient`. absl::Status MovePjRtClientToUnused(const DeviceType& device_type); - absl::Status TestOnlyResetPjRtClient(const DeviceType& device_type); std::string DebugString() const override; // Saves information needed to create a PJRT client (to enable creating a diff --git a/tensorflow/core/tfrt/common/pjrt_state_test.cc b/tensorflow/core/tfrt/common/pjrt_state_test.cc index 6de048977482b9..68731ea4bffccb 100644 --- a/tensorflow/core/tfrt/common/pjrt_state_test.cc +++ b/tensorflow/core/tfrt/common/pjrt_state_test.cc @@ -129,40 +129,4 @@ TEST_F(PjRtStateTestFixture, GetOrCreatePjRtClientNotExist) { EXPECT_THAT(pjrt_client, testing::NotNull()); } -TEST_F(PjRtStateTestFixture, TestOnlyResetPjRtClientPreservesUnusedClients) { - xla::CpuClientOptions options; - options.asynchronous = true; - options.cpu_device_count = 1; - - TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client_1, - xla::GetXlaPjrtCpuClient(options)); - xla::PjRtClient* pjrt_client_1_ptr = pjrt_client_1.get(); - TF_ASSERT_OK(pjrt_state_->SetPjRtClient(tensorflow::DEVICE_CPU, - std::move(pjrt_client_1))); - TF_ASSERT_OK(pjrt_state_->MovePjRtClientToUnused(tensorflow::DEVICE_CPU)); - - TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client_2, - xla::GetXlaPjrtCpuClient(options)); - xla::PjRtClient* pjrt_client_2_ptr = pjrt_client_2.get(); - TF_ASSERT_OK(pjrt_state_->SetPjRtClient(tensorflow::DEVICE_CPU, - std::move(pjrt_client_2))); - - TF_ASSERT_OK(pjrt_state_->TestOnlyResetPjRtClient(tensorflow::DEVICE_CPU)); - - EXPECT_THAT(pjrt_state_->GetPjRtClient(tensorflow::DEVICE_CPU), - absl_testing::StatusIs( - tensorflow::error::NOT_FOUND, - HasSubstr("PjRt client not found for device type"))); - // Verifies that both the previously retired client and the reset client are - // still alive in `unused_`. - EXPECT_EQ(pjrt_client_1_ptr->platform_name(), "cpu"); - EXPECT_EQ(pjrt_client_2_ptr->platform_name(), "cpu"); - - // Resetting another device type (e.g., GPU) should not destroy unused - // clients. - TF_ASSERT_OK(pjrt_state_->TestOnlyResetPjRtClient(tensorflow::DEVICE_GPU)); - EXPECT_EQ(pjrt_client_1_ptr->platform_name(), "cpu"); - EXPECT_EQ(pjrt_client_2_ptr->platform_name(), "cpu"); -} - } // namespace diff --git a/tensorflow/core/tfrt/common/pjrt_util.cc b/tensorflow/core/tfrt/common/pjrt_util.cc index ba6d87a3375205..e9d9b23135498b 100644 --- a/tensorflow/core/tfrt/common/pjrt_util.cc +++ b/tensorflow/core/tfrt/common/pjrt_util.cc @@ -20,12 +20,14 @@ limitations under the License. #include "absl/status/status.h" #include "absl/status/statusor.h" #include "xla/pjrt/pjrt_client.h" -#include "xla/tsl/platform/errors.h" #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/framework/types.h" +#include "tensorflow/core/platform/errors.h" #include "tensorflow/core/platform/refcount.h" +#include "tensorflow/core/platform/status.h" #include "tensorflow/core/tfrt/common/global_state.h" #include "tensorflow/core/tfrt/common/pjrt_state.h" +#include "tsl/platform/errors.h" namespace tensorflow { @@ -90,18 +92,4 @@ absl::StatusOr GetPjRtGpuClientCreationInfo() { core::ScopedUnref pjrt_state_ref(pjrt_state); return pjrt_state->GetPjRtGpuClientCreationInfo(); } - -absl::Status TestOnlyResetPjRtClientInTFGlobalResourceManager( - const DeviceType& device_type) { - ResourceMgr* rmgr = tfrt_global::GetTFGlobalResourceMgr(); - PjRtState* pjrt_state; - TF_RETURN_IF_ERROR(rmgr->LookupOrCreate( - rmgr->default_container(), kPjRtStateResourceName, &pjrt_state, - [&](PjRtState** ret) { - *ret = PjRtState::Create(); - return absl::OkStatus(); - })); - core::ScopedUnref pjrt_state_ref(pjrt_state); - return pjrt_state->TestOnlyResetPjRtClient(device_type); -} } // namespace tensorflow diff --git a/tensorflow/core/tfrt/common/pjrt_util.h b/tensorflow/core/tfrt/common/pjrt_util.h index 77c8ae6e087986..d47e7714515e2c 100644 --- a/tensorflow/core/tfrt/common/pjrt_util.h +++ b/tensorflow/core/tfrt/common/pjrt_util.h @@ -40,8 +40,6 @@ absl::StatusOr GetPjRtClient(const DeviceType& device_type); absl::Status SetPjRtGpuClientCreationInfoInTFGlobalResourceManager( std::unique_ptr info); absl::StatusOr GetPjRtGpuClientCreationInfo(); -absl::Status TestOnlyResetPjRtClientInTFGlobalResourceManager( - const DeviceType& device_type); } // namespace tensorflow diff --git a/tensorflow/core/tfrt/common/pjrt_util_test.cc b/tensorflow/core/tfrt/common/pjrt_util_test.cc index ee531d158707a6..bfd8e9f8c3b34f 100644 --- a/tensorflow/core/tfrt/common/pjrt_util_test.cc +++ b/tensorflow/core/tfrt/common/pjrt_util_test.cc @@ -22,10 +22,12 @@ limitations under the License. #include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h" #include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h" #include "xla/tsl/lib/core/status_test_util.h" -#include "xla/tsl/platform/statusor.h" #include "xla/tsl/protobuf/error_codes.pb.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/tfrt/common/pjrt_state.h" +#include "tsl/platform/status_matchers.h" +#include "tsl/platform/statusor.h" +#include "tsl/platform/test.h" namespace tensorflow { namespace { @@ -62,26 +64,5 @@ TEST(PjRtGpuClientCreationInfoTest, SetAndGet) { EXPECT_THAT(retrieved_info->allowed_devices, ElementsAre(123)); } -TEST(PjRtUtilTest, TestOnlyResetPjRtClientInTFGlobalResourceManager) { - xla::CpuClientOptions options; - options.asynchronous = true; - options.cpu_device_count = 1; - TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetXlaPjrtCpuClient(options)); - TF_ASSERT_OK(SetPjRtClientInTFGlobalResourceManager(DEVICE_CPU, - std::move(pjrt_client))); - TF_ASSERT_OK(TestOnlyResetPjRtClientInTFGlobalResourceManager(DEVICE_CPU)); - EXPECT_THAT(GetPjRtClient(DEVICE_CPU), - absl_testing::StatusIs(error::NOT_FOUND)); - - auto info = std::make_unique(); - info->allowed_devices.insert(123); - TF_ASSERT_OK( - SetPjRtGpuClientCreationInfoInTFGlobalResourceManager(std::move(info))); - TF_ASSERT_OK(TestOnlyResetPjRtClientInTFGlobalResourceManager(DEVICE_GPU)); - TF_ASSERT_OK_AND_ASSIGN(PjRtGpuClientCreationInfo * retrieved_info, - GetPjRtGpuClientCreationInfo()); - EXPECT_EQ(retrieved_info, nullptr); -} - } // namespace } // namespace tensorflow diff --git a/tensorflow/core/tfrt/mlrt/kernel/BUILD b/tensorflow/core/tfrt/mlrt/kernel/BUILD index cdab3c879e41d1..6457ebfce1455c 100644 --- a/tensorflow/core/tfrt/mlrt/kernel/BUILD +++ b/tensorflow/core/tfrt/mlrt/kernel/BUILD @@ -141,6 +141,7 @@ cc_library( "@tf_runtime//:hostcontext", "@tsl//tsl/profiler/lib:connected_traceme", "@tsl//tsl/profiler/lib:context_types_hdrs", + "@tsl//tsl/profiler/lib:traceme", "@xla//xla/tsl/platform:criticality", ], ) @@ -235,6 +236,7 @@ tf_cc_shared_test( "//tensorflow/core/framework:tensor_testutil", "//tensorflow/core/kernels:math", "//tensorflow/core/ops:math_ops_op_lib", + "//tensorflow/core/profiler/lib:profiler_session", "//tensorflow/core/tfrt/fallback:device_with_custom_allocator", "//tensorflow/core/tfrt/fallback:fallback_state", "//tensorflow/core/tfrt/mlrt/bytecode", @@ -248,8 +250,16 @@ tf_cc_shared_test( "@com_google_googletest//:gtest_main", "@tf_runtime//:hostcontext", "@tsl//tsl/platform:status_matchers", + "@tsl//tsl/profiler/lib:profiler_session", + "@tsl//tsl/profiler/protobuf:xplane_proto_cc", + "@xla//xla/backends/profiler/cpu:host_tracer", # buildcleaner: keep "@xla//xla/tsl/lib/core:status_test_util", "@xla//xla/tsl/platform:status_matchers", + "@xla//xla/tsl/profiler/utils:group_events", + "@xla//xla/tsl/profiler/utils:preprocess_xplane", + "@xla//xla/tsl/profiler/utils:tf_xplane_visitor", + "@xla//xla/tsl/profiler/utils:xplane_schema", + "@xla//xla/tsl/profiler/utils:xplane_visitor", ], ) diff --git a/tensorflow/core/tfrt/mlrt/kernel/batch_kernel.cc b/tensorflow/core/tfrt/mlrt/kernel/batch_kernel.cc index 4a65ce42f0f1e5..65f80468b6530f 100644 --- a/tensorflow/core/tfrt/mlrt/kernel/batch_kernel.cc +++ b/tensorflow/core/tfrt/mlrt/kernel/batch_kernel.cc @@ -44,6 +44,7 @@ limitations under the License. #include "tensorflow/core/tfrt/utils/fallback_tensor.h" #include "tsl/profiler/lib/connected_traceme.h" #include "tsl/profiler/lib/context_types.h" +#include "tsl/profiler/lib/traceme.h" #include "tfrt/concurrency/chain.h" // from @tf_runtime #include "tfrt/host_context/resource_context.h" // from @tf_runtime @@ -345,13 +346,14 @@ void MlrtBatchResource::ProcessFuncBatchImpl( fallback_request_state.set_runtime_config( caller_fallback_request_state.runtime_config()); + const uint64_t batch_activity_id = tsl::profiler::TraceMe::NewActivityId(); tsl::profiler::TraceMeProducer activity( // To TraceMeConsumers in WorkQueue. [step_id] { return tsl::profiler::TraceMeEncode("RunMlrtFunction", {{"id", step_id}, {"_r", 1}}); }, - tsl::profiler::ContextType::kTfrtExecutor, step_id, + tsl::profiler::ContextType::kTfrtExecutor, batch_activity_id, tsl::profiler::TraceMeLevel::kInfo); auto trace_me_context_id = activity.GetContextId(); diff --git a/tensorflow/core/tfrt/mlrt/kernel/batch_kernel_test.cc b/tensorflow/core/tfrt/mlrt/kernel/batch_kernel_test.cc index a785d23d12f972..e0bde738ebad6e 100644 --- a/tensorflow/core/tfrt/mlrt/kernel/batch_kernel_test.cc +++ b/tensorflow/core/tfrt/mlrt/kernel/batch_kernel_test.cc @@ -28,6 +28,11 @@ limitations under the License. #include "absl/synchronization/notification.h" #include "absl/types/span.h" #include "xla/tsl/lib/core/status_test_util.h" +#include "xla/tsl/profiler/utils/group_events.h" +#include "xla/tsl/profiler/utils/preprocess_xplane.h" +#include "xla/tsl/profiler/utils/tf_xplane_visitor.h" +#include "xla/tsl/profiler/utils/xplane_schema.h" +#include "xla/tsl/profiler/utils/xplane_visitor.h" #include "tensorflow/core/framework/allocator.h" #include "tensorflow/core/framework/tensor_testutil.h" #include "tensorflow/core/tfrt/fallback/device_with_custom_allocator.h" @@ -38,6 +43,8 @@ limitations under the License. #include "tensorflow/core/tfrt/mlrt/interpreter/interpreter_testutil.h" #include "tensorflow/core/tfrt/mlrt/kernel/context.h" #include "tensorflow/core/tfrt/mlrt/kernel/kernel.h" +#include "tsl/profiler/lib/profiler_session.h" +#include "tsl/profiler/protobuf/xplane.pb.h" #include "tfrt/host_context/concurrent_work_queue.h" // from @tf_runtime #include "tfrt/host_context/execution_context.h" // from @tf_runtime @@ -500,6 +507,107 @@ TEST(KernelTest, BatchFunctionWithDeviceOp) { EXPECT_EQ(test_device->allocator_uses(), 2); } +// Verifies that two batches sharing the same step_id (as in split-batch +// execution) get separate XProf event groups: +// 1. Run 2 batches (max_batch_size=1) with the same step_id under +// ProfilerSession. +// 2. Run XProf's PreprocessXSpace + GroupTfEvents on the captured XSpace. +// 3. Check kGroupId on the 2 "RunMlrtFunction::Execute" worker events: +// - With step_id as context_id: both get group_id=0 ([0, 0] -> FAILS). +// - With NewActivityId(): they get distinct group_ids ([0, 1] -> PASSES). +TEST(KernelTest, SplitBatchesProduceDisjointXProfEventGroups) { + auto buffer = CreateExecutableForBatchFunctionOp(); + + mlrt::bc::Executable executable(buffer.data()); + + mlrt::KernelRegistry registry; + RegisterTfMlrtKernels(registry); + RegisterTfMlrtBatchKernels(registry); + mlrt::LoadedExecutable loaded_executable(executable, registry); + + auto work_queue = tfrt::CreateMultiThreadedWorkQueue( + /*num_threads=*/4, /*num_blocking_threads=*/4); + + tensorflow::SessionOptions session_options; + tensorflow::FunctionDefLibrary fdef_lib; + TF_ASSERT_OK_AND_ASSIGN(auto fallback_state, tfrt_stub::FallbackState::Create( + session_options, fdef_lib)); + + std::function)> runner = + [](const std::function& f) { f(); }; + + // Both calls share this ResourceContext (and batching queue). With + // max_batch_size=1, each call forms a separate batch. + tfrt::ResourceContext resource_context; + + auto run_one_batch = [&]() { + tfrt_stub::OpKernelRunnerTable runner_table; + tfd::FallbackResourceArray resource_array; + tfd::KernelFallbackCompatRequestState fallback_request_state( + &runner, &fallback_state->device_manager(), /*step_id=*/0, + &runner_table, &resource_array, /*user_intra_op_threadpool=*/nullptr, + /*model_metadata=*/std::nullopt, + &fallback_state->process_function_library_runtime()); + + mlrt::ExecutionContext execution_context(&loaded_executable); + execution_context.set_work_queue(work_queue.get()); + execution_context.AddUserContext( + std::make_unique(&fallback_request_state, &resource_context)); + + tensorflow::Tensor input_tensor(tensorflow::DT_INT32, {1}); + input_tensor.flat()(0) = 100; + mlrt::Value arg(tfrt_stub::FallbackTensor(std::move(input_tensor))); + mlrt::Value result; + + absl::Notification notification; + execution_context.set_exit_handler( + [¬ification]() { notification.Notify(); }); + + std::vector last_uses = {true}; + execution_context.Call(executable.functions()[0], last_uses, + absl::MakeSpan(&arg, 1), absl::MakeSpan(&result, 1)); + mlrt::Execute(execution_context); + + notification.WaitForNotification(); + TF_ASSERT_OK(execution_context.status()); + }; + + std::unique_ptr profiler = + tsl::ProfilerSession::Create(tsl::ProfilerSession::DefaultOptions()); + + run_one_batch(); + run_one_batch(); + + tensorflow::profiler::XSpace space; + TF_ASSERT_OK(profiler->CollectData(&space)); + + tsl::profiler::PreprocessXSpace(&space); + tsl::profiler::GroupTfEvents(&space); + + int num_process_batch_events = 0; + std::vector execute_group_ids; + for (const tensorflow::profiler::XPlane& plane : space.planes()) { + if (plane.name() != tsl::profiler::kHostThreadsPlaneName) continue; + tsl::profiler::XPlaneVisitor plane_visitor = + tsl::profiler::CreateTfXPlaneVisitor(&plane); + plane_visitor.ForEachLine([&](const tsl::profiler::XLineVisitor& line) { + line.ForEachEvent([&](const tsl::profiler::XEventVisitor& event) { + if (event.Name() == "ProcessBatch") ++num_process_batch_events; + if (event.Name() != "RunMlrtFunction::Execute") return; + std::optional group_id = + event.GetStat(tsl::profiler::StatType::kGroupId); + if (group_id.has_value()) { + execute_group_ids.push_back(group_id->IntValue()); + } + }); + }); + } + + ASSERT_EQ(num_process_batch_events, 2); + ASSERT_EQ(execute_group_ids.size(), 2); + EXPECT_NE(execute_group_ids[0], execute_group_ids[1]); +} + } // namespace } // namespace tf_mlrt } // namespace tensorflow diff --git a/tensorflow/core/tpu/BUILD b/tensorflow/core/tpu/BUILD index dcc00dbc8b79cd..e3d661185edf6a 100644 --- a/tensorflow/core/tpu/BUILD +++ b/tensorflow/core/tpu/BUILD @@ -295,8 +295,8 @@ cc_library( "@xla//xla:status_macros", "@xla//xla:xla_data_proto_cc", "@xla//xla:xla_proto_cc", + "@xla//xla/client:compile_only_client", "@xla//xla/hlo/ir:hlo", - "@xla//xla/pjrt:pjrt_compiler", "@xla//xla/service:hlo_module_config", "@xla//xla/tsl/platform:errors", "@xla//xla/tsl/platform:statusor", diff --git a/tensorflow/core/tpu/tpu_compile.cc b/tensorflow/core/tpu/tpu_compile.cc index c7702f9f0d0e14..91cfa36f565594 100644 --- a/tensorflow/core/tpu/tpu_compile.cc +++ b/tensorflow/core/tpu/tpu_compile.cc @@ -46,12 +46,12 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" #include "tensorflow/compiler/tf2xla/xla_resource.h" +#include "xla/client/compile_only_client.h" #include "xla/debug_options_flags.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_sharding.h" -#include "xla/pjrt/pjrt_compiler.h" #include "xla/service/hlo_module_config.h" #include "xla/shape.h" #include "xla/shape_util.h" @@ -578,19 +578,19 @@ absl::Status CompileTFFunctionToHlo( const std::vector& arg_shapes, const DeviceType& device_type, const GuaranteedConsts& guaranteed_constants, const NameAttrList& function, const tpu::TPUCompileMetadataProto& metadata, - xla::PjRtCompiler* pjrt_compiler, + xla::CompileOnlyClient* client, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, bool use_tuple_args, XlaCompiler::CompilationResult* compilation_result) { XlaCompiler::Options compiler_options; FunctionLibraryDefinition flib_definition(flib_def); compiler_options.device_type = device_type; + compiler_options.client = client; compiler_options.flib_def = &flib_definition; compiler_options.allow_cpu_custom_calls = false; compiler_options.graph_def_version = graph_def_version; compiler_options.shape_determination_fns = shape_determination_fns; compiler_options.use_shardy_partitioner = metadata.use_shardy_partitioner(); - compiler_options.compiler = pjrt_compiler; auto compiler = std::make_unique(compiler_options); diff --git a/tensorflow/core/tpu/tpu_compile.h b/tensorflow/core/tpu/tpu_compile.h index 7e1dffe5e9a337..dfef9ae0fc927a 100644 --- a/tensorflow/core/tpu/tpu_compile.h +++ b/tensorflow/core/tpu/tpu_compile.h @@ -22,7 +22,7 @@ limitations under the License. #include "absl/types/span.h" #include "tensorflow/compiler/tf2xla/layout_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "xla/pjrt/pjrt_compiler.h" +#include "xla/client/compile_only_client.h" #include "xla/shape.h" #include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/function.h" @@ -42,7 +42,7 @@ absl::Status CompileTFFunctionToHlo( const std::vector& arg_shapes, const DeviceType& device_type, const GuaranteedConsts& guaranteed_constants, const NameAttrList& function, const tpu::TPUCompileMetadataProto& metadata, - xla::PjRtCompiler* pjrt_compiler, + xla::CompileOnlyClient* client, std::vector* arg_core_mapping, std::vector>* per_core_arg_shapes, bool use_tuple_args, XlaCompiler::CompilationResult* compilation_result); diff --git a/tensorflow/core/util/example_proto_fast_parsing.cc b/tensorflow/core/util/example_proto_fast_parsing.cc index 5776007faf3572..40475158509cb9 100644 --- a/tensorflow/core/util/example_proto_fast_parsing.cc +++ b/tensorflow/core/util/example_proto_fast_parsing.cc @@ -2563,10 +2563,16 @@ absl::Status ParseSequenceDenseFeatures( TensorShape dense_shape, row_shape; DataType dtype = c.dtype; const size_t expected_max_elements = feature.length; - if (!c.shape.AsTensorShape(&row_shape) || - expected_max_elements != - (expected_max_elements / row_shape.num_elements()) * - row_shape.num_elements()) { + const int64_t row_num_elements = + c.shape.AsTensorShape(&row_shape) ? row_shape.num_elements() : -1; + // Guard against division by zero when any shape dimension is 0. When + // row_num_elements==0, the only consistent state is zero total elements; + // any positive count would be impossible to fit into zero-element rows. + if (row_num_elements < 0 || + (row_num_elements == 0 + ? expected_max_elements != 0 + : expected_max_elements % static_cast(row_num_elements) != + 0)) { PartialTensorShape total_shape = row_shape; total_shape.InsertDim(0, -1); return absl::InvalidArgumentError(absl::StrCat( @@ -2576,7 +2582,10 @@ absl::Status ParseSequenceDenseFeatures( " is not consistent with output shape: ", total_shape.DebugString())); } int64_t expected_max_rows = - expected_max_elements / row_shape.num_elements(); + row_num_elements == 0 + ? 0 + : static_cast(expected_max_elements / + static_cast(row_num_elements)); if (is_batch) { dense_shape.AddDim(num_examples); } diff --git a/tensorflow/core/util/gpu_device_functions.h b/tensorflow/core/util/gpu_device_functions.h index 35197fcbdf6960..bdbc8610c23bb5 100644 --- a/tensorflow/core/util/gpu_device_functions.h +++ b/tensorflow/core/util/gpu_device_functions.h @@ -110,10 +110,11 @@ namespace detail { template class GpuGridRange { struct Iterator { - __device__ Iterator(T index, T delta) : index_(index), delta_(delta) {} + __device__ Iterator(T index, T delta, T end) + : index_(index), delta_(delta), end_(end) {} __device__ T operator*() const { return index_; } __device__ Iterator& operator++() { - index_ += delta_; + index_ += (delta_ < end_ - index_) ? delta_ : (end_ - index_); return *this; } __device__ bool operator!=(const Iterator& other) const { @@ -133,14 +134,15 @@ class GpuGridRange { private: T index_; const T delta_; + const T end_; }; public: __device__ GpuGridRange(T begin, T delta, T end) : begin_(begin), delta_(delta), end_(end) {} - __device__ Iterator begin() const { return Iterator{begin_, delta_}; } - __device__ Iterator end() const { return Iterator{end_, 0}; } + __device__ Iterator begin() const { return Iterator{begin_, delta_, end_}; } + __device__ Iterator end() const { return Iterator{end_, 0, end_}; } private: T begin_; diff --git a/tensorflow/lite/CMakeLists.txt b/tensorflow/lite/CMakeLists.txt index 27e25e83750461..39809306861f58 100644 --- a/tensorflow/lite/CMakeLists.txt +++ b/tensorflow/lite/CMakeLists.txt @@ -40,32 +40,6 @@ cmake_policy(SET CMP0042 NEW) project(tensorflow-lite C CXX) set(CMAKE_POSITION_INDEPENDENT_CODE ON) -# Use FetchContent to download TensorFlow if not provided -include(FetchContent) - -set(TENSORFLOW_SOURCE_DIR "" CACHE PATH - "Directory that contains the TensorFlow project" -) -if(NOT TENSORFLOW_SOURCE_DIR) - message(STATUS "Downloading TensorFlow repository...") - FetchContent_Declare( - tensorflow - GIT_REPOSITORY https://github.com/tensorflow/tensorflow.git - # v2.21.0-rc0 has updated schema_generated.h file. - GIT_TAG v2.21.0-rc0 - GIT_SHALLOW TRUE - GIT_PROGRESS TRUE - SOURCE_DIR ${CMAKE_CURRENT_BINARY_DIR}/tensorflow-src - ) - FetchContent_GetProperties(tensorflow) - if(NOT tensorflow_POPULATED) - FetchContent_Populate(tensorflow) - endif() - set(TENSORFLOW_SOURCE_DIR "${tensorflow_SOURCE_DIR}") -endif() -set(TF_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/tensorflow") -set(TSL_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/third_party/xla/third_party/tsl") -set(XLA_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/third_party/xla/") set(TFLITE_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}") set(CMAKE_MODULE_PATH "${TFLITE_SOURCE_DIR}/tools/cmake/modules" @@ -164,14 +138,6 @@ macro(populate_tflite_source_vars RELATIVE_DIR SOURCES_VAR) "${TFLITE_SOURCE_DIR}/${RELATIVE_DIR}" ${SOURCES_VAR} ${ARGN} ) endmacro() -# Simplifies inclusion of non-test sources and headers from a directory -# relative to TF_SOURCE_DIR. See populate_source_vars() for the description of -# arguments including and following SOURCES_VAR. -macro(populate_tf_source_vars RELATIVE_DIR SOURCES_VAR) - populate_source_vars( - "${TF_SOURCE_DIR}/${RELATIVE_DIR}" ${SOURCES_VAR} ${ARGN} - ) -endmacro() # Make sure all repos have licenses. set(OVERRIDABLE_FETCH_CONTENT_LICENSE_CHECK ON) # Additional library dependencies based upon enabled features. @@ -190,11 +156,9 @@ endif() find_package(cpuinfo REQUIRED) #CPUINFO is used by XNNPACK and RUY library find_package(ml_dtypes REQUIRED) find_package(ruy REQUIRED) -# Include TSL, which is in tensorflow/third_party -include_directories( - ${TSL_SOURCE_DIR} - ${XLA_SOURCE_DIR} -) +# Header-only TensorFlow dependencies (xla/tsl/..., third_party/fft2d/...). +find_package(tensorflow_headers REQUIRED) +include_directories(${TENSORFLOW_HEADERS_INCLUDE_DIR}) # Download necessary dependencies. if(TFLITE_ENABLE_XNNPACK) # pthreadpool is used by XNNPACK. @@ -410,7 +374,7 @@ if(TFLITE_ENABLE_GPU) ${TFLITE_DELEGATES_GPU_COMMON_TRANSFORMATIONS_SRCS} ${TFLITE_SOURCE_DIR}/tools/versioning/gpu_compatibility.cc ${TFLITE_SOURCE_DIR}/tools/versioning/op_signature.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/tools/versioning/op_signature.cc + ${TFLITE_SOURCE_DIR}/converter/tools/versioning/op_signature.cc ) include_directories( AFTER @@ -439,7 +403,7 @@ if(TFLITE_ENABLE_GPU) target_include_directories(metal_delegate PUBLIC ${CMAKE_BINARY_DIR}/abseil-cpp ${CMAKE_BINARY_DIR}/flatbuffers/include - PRIVATE ${TENSORFLOW_SOURCE_DIR} + PRIVATE ${TFLITE_SOURCE_DIR}/.. PRIVATE ${PROJECT_BINARY_DIR} ) # @@ -455,7 +419,7 @@ if(TFLITE_ENABLE_GPU) OUTPUT ${METAL_GENERATED_DIR}/inference_context_generated.h COMMAND ${CMAKE_COMMAND} -E make_directory ${METAL_GENERATED_DIR} COMMAND ${FLATC} --scoped-enums - -I ${TENSORFLOW_SOURCE_DIR} + -I ${TFLITE_SOURCE_DIR}/.. -o ${METAL_GENERATED_DIR} -c ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/inference_context.fbs DEPENDS ${FLATC_TARGET} @@ -533,7 +497,7 @@ if(TFLITE_ENABLE_XNNPACK) "${PROJECT_BINARY_DIR}/tensorflow/lite/delegates/xnnpack/weight_cache_schema_generated.h" ) target_include_directories(xnnpack-delegate - PUBLIC $ $ + PUBLIC $ $ PRIVATE "${PROJECT_BINARY_DIR}" PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/.." PRIVATE "${PROJECT_BINARY_DIR}/tensorflow" @@ -683,26 +647,26 @@ set(_ALL_TFLITE_SRCS ${TFLITE_SOURCE_DIR}/schema/conversion_metadata_generated.h ${TFLITE_SOURCE_DIR}/schema/schema_generated.h ${TFLITE_SRCS} - ${TF_SOURCE_DIR}/compiler/mlir/lite/allocation.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/allocation.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/api/error_reporter.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/api/error_reporter.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/api/flatbuffer_conversions.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/api/verifier.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/c/tflite_types.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/macros.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/model_builder_base.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/core/model_builder_base.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/experimental/remat/metadata_util.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/experimental/remat/metadata_util.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/mmap_allocation.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/mmap_allocation_disabled.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/schema/conversion_metadata_generated.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/schema/schema_generated.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/schema/schema_utils.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/utils/control_edges.h - ${TF_SOURCE_DIR}/compiler/mlir/lite/utils/string_utils.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/utils/string_utils.h + ${TFLITE_SOURCE_DIR}/converter/allocation.cc + ${TFLITE_SOURCE_DIR}/converter/allocation.h + ${TFLITE_SOURCE_DIR}/converter/core/api/error_reporter.cc + ${TFLITE_SOURCE_DIR}/converter/core/api/error_reporter.h + ${TFLITE_SOURCE_DIR}/converter/core/api/flatbuffer_conversions.cc + ${TFLITE_SOURCE_DIR}/converter/core/api/verifier.h + ${TFLITE_SOURCE_DIR}/converter/core/c/tflite_types.h + ${TFLITE_SOURCE_DIR}/converter/core/macros.h + ${TFLITE_SOURCE_DIR}/converter/core/model_builder_base.cc + ${TFLITE_SOURCE_DIR}/converter/core/model_builder_base.h + ${TFLITE_SOURCE_DIR}/converter/experimental/remat/metadata_util.cc + ${TFLITE_SOURCE_DIR}/converter/experimental/remat/metadata_util.h + ${TFLITE_SOURCE_DIR}/converter/mmap_allocation.cc + ${TFLITE_SOURCE_DIR}/converter/mmap_allocation_disabled.cc + ${TFLITE_SOURCE_DIR}/converter/schema/conversion_metadata_generated.h + ${TFLITE_SOURCE_DIR}/converter/schema/schema_generated.h + ${TFLITE_SOURCE_DIR}/converter/schema/schema_utils.cc + ${TFLITE_SOURCE_DIR}/converter/utils/control_edges.h + ${TFLITE_SOURCE_DIR}/converter/utils/string_utils.cc + ${TFLITE_SOURCE_DIR}/converter/utils/string_utils.h # go/keep-sorted end ) @@ -719,7 +683,6 @@ set(_ALL_TFLITE_HDRS ${_ALL_TFLITE_SRCS}) list(FILTER _ALL_TFLITE_HDRS INCLUDE REGEX ".*\\.h$") target_include_directories(tensorflow-lite PUBLIC - $ $ $ $ @@ -807,21 +770,10 @@ if(TFLITE_KERNEL_TEST) add_subdirectory(${TFLITE_SOURCE_DIR}/kernels) endif() -# Add the generated headers directory. Required for maintaining the -# tensorflow/lite directory structure for generated headers. -set(TFLITE_GENERATED_HEADERS_DIR ${CMAKE_BINARY_DIR}/tensorflow/lite) - -# Add the profiling proto directory. -add_subdirectory(${TFLITE_SOURCE_DIR}/profiling/proto) - -# Add the benchmark result proto directory. -add_subdirectory(${TFLITE_SOURCE_DIR}/tools/benchmark/proto) - -# Add the tf example directory. -add_subdirectory(${TF_SOURCE_DIR}/core/example ${CMAKE_BINARY_DIR}/example_proto_generated) - # The benchmark tool. -add_subdirectory(${TFLITE_SOURCE_DIR}/tools/benchmark) +if(TFLITE_ENABLE_BENCHMARK_MODEL) + add_subdirectory(${TFLITE_SOURCE_DIR}/tools/benchmark) +endif() # The label_image example. add_subdirectory(${TFLITE_SOURCE_DIR}/examples/label_image) @@ -838,11 +790,6 @@ add_library(_pywrap_tensorflow_interpreter_wrapper SHARED EXCLUDE_FROM_ALL # To remove "lib" prefix. set_target_properties(_pywrap_tensorflow_interpreter_wrapper PROPERTIES PREFIX "") -target_include_directories(_pywrap_tensorflow_interpreter_wrapper - PUBLIC - ${TENSORFLOW_SOURCE_DIR} -) - target_link_libraries(_pywrap_tensorflow_interpreter_wrapper tensorflow-lite ${CMAKE_DL_LIBS} diff --git a/tensorflow/lite/examples/label_image/CMakeLists.txt b/tensorflow/lite/examples/label_image/CMakeLists.txt index 823d90dc044b94..6dda556abe8e67 100644 --- a/tensorflow/lite/examples/label_image/CMakeLists.txt +++ b/tensorflow/lite/examples/label_image/CMakeLists.txt @@ -21,7 +21,7 @@ populate_source_vars("${TFLITE_SOURCE_DIR}/examples/label_image" FILTER "_test\\.cc$" ) list(APPEND TFLITE_LABEL_IMAGE_SRCS - ${XLA_SOURCE_DIR}/xla/tsl/util/stats_calculator.cc + ${TENSORFLOW_HEADERS_INCLUDE_DIR}/xla/tsl/util/stats_calculator.cc ${TFLITE_SOURCE_DIR}/profiling/memory_info.cc ${TFLITE_SOURCE_DIR}/profiling/profile_buffer.cc ${TFLITE_SOURCE_DIR}/profiling/profile_summarizer.cc @@ -82,8 +82,10 @@ target_compile_options(label_image PRIVATE ${TFLITE_LABEL_IMAGE_CC_OPTIONS} ) +target_compile_definitions(label_image + PRIVATE + TFLITE_WITHOUT_PROTO +) target_link_libraries(label_image tensorflow-lite - profiling_info_proto - libprotobuf ) diff --git a/tensorflow/lite/java/src/test/java/org/tensorflow/lite/gpu/GpuDelegateTest.java b/tensorflow/lite/java/src/test/java/org/tensorflow/lite/gpu/GpuDelegateTest.java index e1c8fd7fb5f40c..ee23a362acd639 100644 --- a/tensorflow/lite/java/src/test/java/org/tensorflow/lite/gpu/GpuDelegateTest.java +++ b/tensorflow/lite/java/src/test/java/org/tensorflow/lite/gpu/GpuDelegateTest.java @@ -214,7 +214,7 @@ private static ArrayList getTopKLabels(float[][] labels, int k) { @Override public int compare(Map.Entry o1, Map.Entry o2) { // Intentionally reversed to put high confidence at the head of the queue. - return o1.getValue().compareTo(o2.getValue()) * -1; + return o2.getValue().compareTo(o1.getValue()); } }); diff --git a/tensorflow/lite/kernels/BUILD b/tensorflow/lite/kernels/BUILD index d693d8f2aa11f9..d43aaa4fcbbdd6 100644 --- a/tensorflow/lite/kernels/BUILD +++ b/tensorflow/lite/kernels/BUILD @@ -174,8 +174,10 @@ cc_library( "@com_google_absl//absl/base", "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/base:no_destructor", + "@com_google_absl//absl/log", "@com_google_absl//absl/log:absl_check", "@com_google_absl//absl/log:absl_log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/synchronization", @@ -183,7 +185,6 @@ cc_library( "@com_google_googletest//:gtest", "@eigen_archive//:eigen3", "@flatbuffers", - "@tsl//tsl/platform:logging", ], ) diff --git a/tensorflow/lite/kernels/CMakeLists.txt b/tensorflow/lite/kernels/CMakeLists.txt index b9872323ac698d..ca9f48a5b3412f 100644 --- a/tensorflow/lite/kernels/CMakeLists.txt +++ b/tensorflow/lite/kernels/CMakeLists.txt @@ -18,10 +18,8 @@ find_package(google_benchmark REQUIRED) find_package(re2 REQUIRED) # Generate the mutable schema_generated.h header for tests. -set(TF_SOURCE_DIR "${TENSORFLOW_SOURCE_DIR}/tensorflow") -set(SCHEMA_GENERATED_ROOT ${TF_SOURCE_DIR}/compiler/mlir/lite/schema) -set(SCHEMA_FILE ${SCHEMA_GENERATED_ROOT}/schema.fbs) -set(SCHEMA_GENERATED_PATH ${SCHEMA_GENERATED_ROOT}/mutable) +set(SCHEMA_FILE ${TFLITE_SOURCE_DIR}/converter/schema/schema.fbs) +set(SCHEMA_GENERATED_PATH ${PROJECT_BINARY_DIR}/tflite/converter/schema/mutable) set(SCHEMA_GENERATED_FILE ${SCHEMA_GENERATED_PATH}/schema_generated.h) # Use the util function in flatbuffer to generate the schema header. @@ -84,18 +82,15 @@ set(TEST_FRAMEWORK_SRC ${TFLITE_SOURCE_DIR}/delegates/nnapi/acceleration_test_list.cc ${TFLITE_SOURCE_DIR}/delegates/nnapi/acceleration_test_util.cc ${TFLITE_SOURCE_DIR}/profiling/memory_info.cc - ${TFLITE_SOURCE_DIR}/../compiler/mlir/lite/schema/schema_conversion_utils.cc + ${TFLITE_SOURCE_DIR}/converter/schema/schema_conversion_utils.cc ${TFLITE_SOURCE_DIR}/tools/command_line_flags.cc ${DELEGATE_PROVIDERS} ${TFLITE_SOURCE_DIR}/tools/optimize/model_utils.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/tools/optimize/operator_property.cc - ${TF_SOURCE_DIR}/compiler/mlir/tools/optimize/quantization_utils.cc - ${TF_SOURCE_DIR}/compiler/mlir/lite/kernels/internal/runtime_shape.cc + ${TFLITE_SOURCE_DIR}/converter/tools/optimize/operator_property.cc + ${TFLITE_SOURCE_DIR}/converter/kernels/internal/runtime_shape.cc ${TFLITE_SOURCE_DIR}/tools/optimize/quantization_utils.cc ${TFLITE_SOURCE_DIR}/tools/tool_params.cc - ${TFLITE_SOURCE_DIR}/tools/versioning/op_version.cc - ${XLA_SOURCE_DIR}/xla/tsl/platform/default/env_time.cc - ${XLA_SOURCE_DIR}/xla/tsl/platform/default/logging.cc + ${TFLITE_SOURCE_DIR}/converter/tools/versioning/op_version.cc internal/test_util.cc acceleration_test_util.cc acceleration_test_util_internal.cc @@ -111,10 +106,13 @@ endif() if(NOT TFLITE_ENABLE_GPU) list(APPEND TEST_FRAMEWORK_SRC ${TFLITE_SOURCE_DIR}/tools/versioning/op_signature.cc + ${TFLITE_SOURCE_DIR}/converter/tools/versioning/op_signature.cc ) endif() -set(TEST_FRAMEWORK_OPTIONS "") +# Model dumping (--dump_tflite_model_dir) needs tools/serialization, which is not +# built with CMake. +set(TEST_FRAMEWORK_OPTIONS "-DTFLITE_WITHOUT_MODEL_DUMP") if(TFLITE_ENABLE_XNNPACK) list(APPEND TEST_FRAMEWORK_SRC ${TFLITE_SOURCE_DIR}/tools/delegates/xnnpack_delegate_provider.cc @@ -130,9 +128,12 @@ target_link_libraries(tensorflow-lite-test-base re2 tensorflow-lite benchmark + absl::check + absl::log + absl::random_random ) add_dependencies(tensorflow-lite-test-base mutable_schema_file) -target_include_directories(tensorflow-lite-test-base PUBLIC ${SCHEMA_GENERATED_ROOT}) +target_include_directories(tensorflow-lite-test-base PUBLIC ${PROJECT_BINARY_DIR}) target_compile_options(tensorflow-lite-test-base PUBLIC ${TEST_FRAMEWORK_OPTIONS}) add_library(tensorflow-lite-test-gtest-main INTERFACE) @@ -163,7 +164,7 @@ macro(add_kernel_test TEST_SRC TEST_LIB) ${TEST_LIB} ) add_test(NAME ${TEST_NAME} COMMAND ${TEST_NAME} - WORKING_DIRECTORY ${TENSORFLOW_SOURCE_DIR} + WORKING_DIRECTORY ${TFLITE_SOURCE_DIR}/.. ) set_tests_properties(${TEST_NAME} PROPERTIES LABELS "plain") diff --git a/tensorflow/lite/kernels/test_util.cc b/tensorflow/lite/kernels/test_util.cc index 8596eae3639ac6..f66ea6ff0a49db 100644 --- a/tensorflow/lite/kernels/test_util.cc +++ b/tensorflow/lite/kernels/test_util.cc @@ -37,6 +37,8 @@ limitations under the License. #include "absl/base/const_init.h" #include "absl/base/no_destructor.h" #include "absl/base/thread_annotations.h" +#include "absl/log/check.h" +#include "absl/log/log.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/strings/str_replace.h" @@ -65,11 +67,12 @@ limitations under the License. #include "tensorflow/lite/string_type.h" #include "tensorflow/lite/string_util.h" #include "tensorflow/lite/tools/logging.h" +#ifndef TFLITE_WITHOUT_MODEL_DUMP #include "tensorflow/lite/tools/serialization/writer_lib.h" +#endif // TFLITE_WITHOUT_MODEL_DUMP #include "tensorflow/lite/tools/versioning/op_version.h" #include "tensorflow/lite/types/fp16.h" // IWYU pragma: keep #include "tensorflow/lite/version.h" -#include "tsl/platform/logging.h" namespace tflite { @@ -164,6 +167,7 @@ MATCHER(Fp16Eq, "") { return AlmostEquals(actual, expected, fp16_ulps_in_fp32); } +#ifndef TFLITE_WITHOUT_MODEL_DUMP // Returns the name of the dumped model. The name is in the format of // DTS---.tflite. The model serial // number is used to distinguish different models dumped in the same test. @@ -281,6 +285,7 @@ std::unique_ptr ModifyDumpedModel( std::make_unique(std::move(data), fbb.GetSize()); return FlatBufferModel::VerifyAndBuildFromAllocation(std::move(allocation)); } +#endif // TFLITE_WITHOUT_MODEL_DUMP } // namespace @@ -703,6 +708,7 @@ int SingleOpModel::CountNumberOfDelegatedPartitions() const { } void SingleOpModel::MaybeDumpModel() { +#ifndef TFLITE_WITHOUT_MODEL_DUMP std::string dump_directory( tflite::KernelTestDelegateProviders::Get() ->ConstParams() @@ -750,6 +756,7 @@ void SingleOpModel::MaybeDumpModel() { output_file.write(reinterpret_cast(allocation->base()), allocation->bytes()); output_file.close(); +#endif // TFLITE_WITHOUT_MODEL_DUMP } SingleOpModel::~SingleOpModel() { diff --git a/tensorflow/lite/kernels/test_util.h b/tensorflow/lite/kernels/test_util.h index ce2ce190554a55..3b9121c24125bc 100644 --- a/tensorflow/lite/kernels/test_util.h +++ b/tensorflow/lite/kernels/test_util.h @@ -42,6 +42,8 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/log/absl_check.h" #include "absl/log/absl_log.h" +#include "absl/log/check.h" +#include "absl/log/log.h" #include "absl/types/span.h" #include "Eigen/Core" // from @eigen_archive #include "flatbuffers/buffer.h" // from @flatbuffers @@ -63,7 +65,6 @@ limitations under the License. #include "tensorflow/lite/types/fp16.h" #include "tensorflow/lite/types/half.h" #include "tensorflow/lite/util.h" -#include "tsl/platform/logging.h" namespace tflite { diff --git a/tensorflow/lite/nnapi/BUILD b/tensorflow/lite/nnapi/BUILD index 4fff5610d9ab79..58b6dabe2fc25e 100644 --- a/tensorflow/lite/nnapi/BUILD +++ b/tensorflow/lite/nnapi/BUILD @@ -101,7 +101,6 @@ cc_library( deps = [ ":nnapi_implementation", ":nnapi_lib", - "//tensorflow/core/platform:logging", "//tensorflow/lite:framework", "@com_google_absl//absl/log:check", ], diff --git a/tensorflow/lite/nnapi/nnapi_handler.h b/tensorflow/lite/nnapi/nnapi_handler.h index 57ab9d7dfae339..31e6916dd65e25 100644 --- a/tensorflow/lite/nnapi/nnapi_handler.h +++ b/tensorflow/lite/nnapi/nnapi_handler.h @@ -20,7 +20,6 @@ limitations under the License. #include #include "absl/log/check.h" -#include "tensorflow/core/platform/logging.h" #include "tensorflow/lite/nnapi/NeuralNetworksTypes.h" #include "tensorflow/lite/nnapi/nnapi_implementation.h" diff --git a/tensorflow/lite/profiling/profile_summary_formatter.cc b/tensorflow/lite/profiling/profile_summary_formatter.cc index afc9dc2db36fd2..20c3eb53a1bbf5 100644 --- a/tensorflow/lite/profiling/profile_summary_formatter.cc +++ b/tensorflow/lite/profiling/profile_summary_formatter.cc @@ -29,7 +29,9 @@ limitations under the License. #include "tensorflow/core/util/stat_summarizer_options.h" #include "tensorflow/core/util/stats_calculator.h" +#ifndef TFLITE_WITHOUT_PROTO #include "tensorflow/lite/profiling/proto/profiling_info.pb.h" +#endif // TFLITE_WITHOUT_PROTO #include "tensorflow/lite/tools/logging.h" namespace tflite { @@ -139,6 +141,7 @@ ProfileSummaryCSVFormatter::GetStatSummarizerOptions() const { return options; } +#ifndef TFLITE_WITHOUT_PROTO std::vector ProfileSummaryProtoFormatter::GetDetailsSortedByRunOrder( const tensorflow::StatsCalculator* stats_calculator) const { @@ -305,6 +308,7 @@ void ProfileSummaryProtoFormatter::HandleOutput( benchmark_profiling_data.SerializeToOstream(output_stream); } } +#endif // TFLITE_WITHOUT_PROTO } // namespace profiling } // namespace tflite diff --git a/tensorflow/lite/profiling/profile_summary_formatter.h b/tensorflow/lite/profiling/profile_summary_formatter.h index 62514eafecda93..5eab2f1539d572 100644 --- a/tensorflow/lite/profiling/profile_summary_formatter.h +++ b/tensorflow/lite/profiling/profile_summary_formatter.h @@ -28,7 +28,9 @@ limitations under the License. #include "tensorflow/core/util/stat_summarizer_options.h" #include "tensorflow/core/util/stats_calculator.h" +#ifndef TFLITE_WITHOUT_PROTO #include "tensorflow/lite/profiling/proto/profiling_info.pb.h" +#endif // TFLITE_WITHOUT_PROTO namespace tflite { namespace profiling { @@ -98,6 +100,7 @@ class ProfileSummaryCSVFormatter : public ProfileSummaryDefaultFormatter { tensorflow::StatSummarizerOptions GetStatSummarizerOptions() const override; }; +#ifndef TFLITE_WITHOUT_PROTO class ProfileSummaryProtoFormatter : public ProfileSummaryFormatter { public: std::string GetOutputString( @@ -139,6 +142,7 @@ class ProfileSummaryProtoFormatter : public ProfileSummaryFormatter { std::vector GetDetailsSortedByRunOrder( const tensorflow::StatsCalculator* stats_calculator) const; }; +#endif // TFLITE_WITHOUT_PROTO } // namespace profiling } // namespace tflite diff --git a/tensorflow/lite/profiling/proto/CMakeLists.txt b/tensorflow/lite/profiling/proto/CMakeLists.txt deleted file mode 100644 index 5738c992fc2839..00000000000000 --- a/tensorflow/lite/profiling/proto/CMakeLists.txt +++ /dev/null @@ -1,57 +0,0 @@ -# -# Copyright 2024 The TensorFlow Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -find_package(Protobuf REQUIRED) - -add_library(profiling_info_proto profiling_info.proto) - -list(APPEND profiling_info_generated_files - ${CMAKE_BINARY_DIR}/tensorflow/lite/profiling/proto/profiling_info.pb.cc - ${CMAKE_BINARY_DIR}/tensorflow/lite/profiling/proto/profiling_info.pb.h) - -# Generate profiling_info.pb.cc and profiling_info.pb.h from -# profiling_info.proto using protoc. Once the protobuf package version is -# upgraded, we can use protobuf_generate_cpp/protobuf_generate here directly. -add_custom_command( - OUTPUT ${profiling_info_generated_files} - COMMAND ${Protobuf_PROTOC_EXECUTABLE} - ARGS --cpp_out=${CMAKE_BINARY_DIR} --proto_path=${TENSORFLOW_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/profiling_info.proto - DEPENDS ${Protobuf_PROTOC_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/profiling_info.proto -) - -set_source_files_properties(${profiling_info_generated_files} PROPERTIES GENERATED TRUE) -target_sources(profiling_info_proto PRIVATE ${profiling_info_generated_files}) -target_link_libraries(profiling_info_proto protobuf::libprotobuf) -target_include_directories(profiling_info_proto PUBLIC ${CMAKE_BINARY_DIR}) - -add_library(model_runtime_info_proto model_runtime_info.proto) -list(APPEND model_runtime_info_generated_files - ${CMAKE_BINARY_DIR}/tensorflow/lite/profiling/proto/model_runtime_info.pb.cc - ${CMAKE_BINARY_DIR}/tensorflow/lite/profiling/proto/model_runtime_info.pb.h -) - -# Generate model_runtime_info.pb.cc and model_runtime_info.pb.h from -# model_runtime_info.proto using protoc. Once the protobuf package version is -# upgraded, we can use protobuf_generate_cpp/protobuf_generate here directly. -add_custom_command( - OUTPUT ${model_runtime_info_generated_files} - COMMAND ${Protobuf_PROTOC_EXECUTABLE} - ARGS --cpp_out=${CMAKE_BINARY_DIR} --proto_path=${TENSORFLOW_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/model_runtime_info.proto - DEPENDS ${Protobuf_PROTOC_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/model_runtime_info.proto ${profiling_info_generated_files} -) - -set_source_files_properties(${model_runtime_info_generated_files} PROPERTIES GENERATED TRUE) -target_sources(model_runtime_info_proto PRIVATE ${model_runtime_info_generated_files}) -target_link_libraries(model_runtime_info_proto protobuf::libprotobuf) -target_include_directories(model_runtime_info_proto PUBLIC ${CMAKE_BINARY_DIR}) diff --git a/tensorflow/lite/toco/tflite/BUILD b/tensorflow/lite/toco/tflite/BUILD index ab265a0264ef4f..b665f95ada7714 100644 --- a/tensorflow/lite/toco/tflite/BUILD +++ b/tensorflow/lite/toco/tflite/BUILD @@ -142,6 +142,7 @@ cc_library( "//tensorflow/lite/toco:toco_port", "//tensorflow/lite/toco:tooling_util", "@com_google_absl//absl/log", + "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@flatbuffers", ], diff --git a/tensorflow/lite/toco/tflite/export.cc b/tensorflow/lite/toco/tflite/export.cc index d04d598dc62a6d..63889cb572c91b 100644 --- a/tensorflow/lite/toco/tflite/export.cc +++ b/tensorflow/lite/toco/tflite/export.cc @@ -23,6 +23,7 @@ limitations under the License. #include "flatbuffers/flexbuffers.h" #include "absl/log/log.h" +#include "absl/status/status.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "flatbuffers/buffer.h" // from @flatbuffers diff --git a/tensorflow/lite/tools/benchmark/CMakeLists.txt b/tensorflow/lite/tools/benchmark/CMakeLists.txt index aee420c6a52b9f..087d73ab41ce60 100644 --- a/tensorflow/lite/tools/benchmark/CMakeLists.txt +++ b/tensorflow/lite/tools/benchmark/CMakeLists.txt @@ -21,11 +21,10 @@ populate_source_vars("${TFLITE_SOURCE_DIR}/tools/benchmark" FILTER "(_test|_plus_flex_main|_performance_options.*)\\.cc$" ) list(APPEND TFLITE_BENCHMARK_SRCS - ${XLA_SOURCE_DIR}/xla/tsl/util/stats_calculator.cc + ${TENSORFLOW_HEADERS_INCLUDE_DIR}/xla/tsl/util/stats_calculator.cc ${TFLITE_SOURCE_DIR}/kernels/internal/utils/sparsity_format_converter.cc ${TFLITE_SOURCE_DIR}/profiling/memory_info.cc ${TFLITE_SOURCE_DIR}/profiling/memory_usage_monitor.cc - ${TFLITE_SOURCE_DIR}/profiling/model_runtime_info.cc ${TFLITE_SOURCE_DIR}/profiling/profile_buffer.cc ${TFLITE_SOURCE_DIR}/profiling/profile_summarizer.cc ${TFLITE_SOURCE_DIR}/profiling/profile_summary_formatter.cc @@ -44,15 +43,6 @@ list(APPEND TFLITE_BENCHMARK_SRCS list(APPEND TFLITE_BENCHMARK_LIBS tensorflow-lite -) - -list(APPEND TFLITE_BENCHMARK_LIBS - profiling_info_proto - feature_proto - example_proto - model_runtime_info_proto - benchmark_result_proto - protobuf::libprotobuf absl::log ) @@ -110,9 +100,11 @@ target_compile_options(benchmark_model PRIVATE ${TFLITE_BENCHMARK_CC_OPTIONS} ) -target_include_directories(benchmark_model - PUBLIC - ${CMAKE_BINARY_DIR} +# Proto-based outputs (e.g. --result_file_path, --export_model_runtime_info and +# the 'proto' op profiling output mode) are not supported in the CMake build. +target_compile_definitions(benchmark_model + PRIVATE + TFLITE_WITHOUT_PROTO ) target_link_libraries(benchmark_model ${TFLITE_BENCHMARK_LIBS} diff --git a/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc b/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc index f6ccb18a963143..cfc7c7a65bc19c 100644 --- a/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc +++ b/tensorflow/lite/tools/benchmark/benchmark_tflite_model.cc @@ -39,8 +39,6 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "ruy/profiler/profiler.h" // from @ruy -#include "tensorflow/core/example/example.pb.h" -#include "tensorflow/core/example/feature.pb.h" #include "tensorflow/lite/core/c/c_api_types.h" #include "tensorflow/lite/core/c/common.h" #include "tensorflow/lite/core/kernels/register.h" @@ -52,18 +50,25 @@ limitations under the License. #include "tensorflow/lite/kernels/cpu_backend_context.h" #include "tensorflow/lite/op_resolver.h" #include "tensorflow/lite/optional_debug_tools.h" -#include "tensorflow/lite/profiling/model_runtime_info.h" #include "tensorflow/lite/profiling/profile_summary_formatter.h" #include "tensorflow/lite/string_util.h" #include "tensorflow/lite/tools/benchmark/benchmark_params.h" #include "tensorflow/lite/tools/benchmark/benchmark_utils.h" #include "tensorflow/lite/tools/benchmark/profiling_listener.h" -#include "tensorflow/lite/tools/benchmark/proto/benchmark_result.pb.h" #include "tensorflow/lite/tools/delegates/delegate_provider.h" #include "tensorflow/lite/tools/logging.h" #include "tensorflow/lite/tools/model_loader.h" #include "tensorflow/lite/tools/utils.h" +// Proto-based outputs are compiled out when TFLITE_WITHOUT_PROTO is defined +// (e.g. the CMake build), which avoids a dependency on protobuf and TensorFlow. +#ifndef TFLITE_WITHOUT_PROTO +#include "tensorflow/core/example/example.pb.h" +#include "tensorflow/core/example/feature.pb.h" +#include "tensorflow/lite/profiling/model_runtime_info.h" +#include "tensorflow/lite/tools/benchmark/proto/benchmark_result.pb.h" +#endif // TFLITE_WITHOUT_PROTO + void RegisterSelectedOps(::tflite::MutableOpResolver* resolver); // Version with Weak linker attribute doing nothing: if someone links this @@ -75,7 +80,9 @@ RegisterSelectedOps(::tflite::MutableOpResolver* resolver) {} namespace tflite { namespace benchmark { namespace { +#ifndef TFLITE_WITHOUT_PROTO using ::tflite::tools::benchmark::BenchmarkResult; +#endif // TFLITE_WITHOUT_PROTO using utils::InputTensorData; using utils::VoidUniquePtr; @@ -89,12 +96,19 @@ constexpr bool kOpProfilingEnabledDefault = false; // Op profiling output modes. constexpr char kOpProfilingOutputModeStdout[] = "stdout"; constexpr char kOpProfilingOutputModeCsv[] = "csv"; +#ifndef TFLITE_WITHOUT_PROTO constexpr char kOpProfilingOutputModeProto[] = "proto"; +#endif // TFLITE_WITHOUT_PROTO + +const char* kOpProfilingOutputModes[] = { + kOpProfilingOutputModeStdout, + kOpProfilingOutputModeCsv, +#ifndef TFLITE_WITHOUT_PROTO + kOpProfilingOutputModeProto, +#endif // TFLITE_WITHOUT_PROTO +}; -const char* kOpProfilingOutputModes[] = {kOpProfilingOutputModeStdout, - kOpProfilingOutputModeCsv, - kOpProfilingOutputModeProto}; - +#ifndef TFLITE_WITHOUT_PROTO // Sets feature values in the tensorflow::Example proto from the tflite tensor. // Returns an error if the tensor type is not supported or the tensor dime is a // nullptr. @@ -137,6 +151,7 @@ TfLiteStatus MaybeSetFeatureValuesFromTensor(const TfLiteTensor& tensor, return kTfLiteError; } } +#endif // TFLITE_WITHOUT_PROTO // Dumps ruy profiling events if the ruy profiler is enabled. class RuyProfileListener : public BenchmarkListener { @@ -218,6 +233,7 @@ class OutputSaver : public BenchmarkListener { } } +#ifndef TFLITE_WITHOUT_PROTO // If the output_proto_filepath is specified, save the output tensors as // tensorflow::Example proto and serialize it to the file. const std::string output_proto_path = @@ -236,6 +252,12 @@ class OutputSaver : public BenchmarkListener { ofs.close(); } } +#else + if (!params_->Get("output_proto_filepath").empty()) { + TFLITE_LOG(WARN) + << "--output_proto_filepath is not supported in this build."; + } +#endif // TFLITE_WITHOUT_PROTO } private: @@ -243,6 +265,7 @@ class OutputSaver : public BenchmarkListener { const BenchmarkParams* params_ = nullptr; }; +#ifndef TFLITE_WITHOUT_PROTO // Dumps the Model Runtime Info if enabled when export_model_runtime_info is // set to true. class ModelRuntimeInfoListener : public BenchmarkListener { @@ -337,6 +360,7 @@ class ProtoBenchmarkReporter : public BenchmarkListener { private: std::string result_file_path_; }; +#endif // TFLITE_WITHOUT_PROTO std::vector Split(const std::string& str, const char delim) { if (str.empty()) { @@ -488,8 +512,10 @@ std::shared_ptr CreateProfileSummaryFormatter(const std::string& output_mode) { if (output_mode == kOpProfilingOutputModeCsv) { return std::make_shared(); +#ifndef TFLITE_WITHOUT_PROTO } else if (output_mode == kOpProfilingOutputModeProto) { return std::make_shared(); +#endif // TFLITE_WITHOUT_PROTO } else { return std::make_shared(); } @@ -1205,8 +1231,13 @@ TfLiteStatus BenchmarkTfLiteModel::Init() { new InterpreterStatePrinter(interpreter_.get()))); if (params_.Get("export_model_runtime_info")) { +#ifndef TFLITE_WITHOUT_PROTO AddOwnedListener(std::unique_ptr( new ModelRuntimeInfoListener(interpreter_.get()))); +#else + TFLITE_LOG(WARN) + << "--export_model_runtime_info is not supported in this build."; +#endif // TFLITE_WITHOUT_PROTO } interpreter_->SetAllowFp16PrecisionForFp32(params_.Get("allow_fp16")); @@ -1375,8 +1406,14 @@ TfLiteStatus BenchmarkTfLiteModel::Init() { AddOwnedListener( std::unique_ptr(new RuyProfileListener())); +#ifndef TFLITE_WITHOUT_PROTO AddOwnedListener( std::unique_ptr(new ProtoBenchmarkReporter())); +#else + if (!params_.Get("result_file_path").empty()) { + TFLITE_LOG(WARN) << "--result_file_path is not supported in this build."; + } +#endif // TFLITE_WITHOUT_PROTO AddOwnedListener(std::unique_ptr( new OutputSaver(interpreter_runner_.get()))); diff --git a/tensorflow/lite/tools/benchmark/proto/CMakeLists.txt b/tensorflow/lite/tools/benchmark/proto/CMakeLists.txt deleted file mode 100644 index 6a39f06e03c32b..00000000000000 --- a/tensorflow/lite/tools/benchmark/proto/CMakeLists.txt +++ /dev/null @@ -1,36 +0,0 @@ -# -# Copyright 2025 The TensorFlow Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -find_package(Protobuf REQUIRED) - -add_library(benchmark_result_proto benchmark_result.proto) - -list(APPEND benchmark_result_generated_files - ${CMAKE_BINARY_DIR}/tensorflow/lite/tools/benchmark/proto/benchmark_result.pb.cc - ${CMAKE_BINARY_DIR}/tensorflow/lite/tools/benchmark/proto/benchmark_result.pb.h) - -# Generate benchmark_result.pb.cc and benchmark_result.pb.h from -# benchmark_result.proto using protoc. Once the protobuf package version is -# upgraded, we can use protobuf_generate_cpp/protobuf_generate here directly. -add_custom_command( - OUTPUT ${benchmark_result_generated_files} - COMMAND ${Protobuf_PROTOC_EXECUTABLE} - ARGS --cpp_out=${CMAKE_BINARY_DIR} --proto_path=${TENSORFLOW_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/benchmark_result.proto - DEPENDS ${Protobuf_PROTOC_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/benchmark_result.proto -) - -set_source_files_properties(${benchmark_result_generated_files} PROPERTIES GENERATED TRUE) -target_sources(benchmark_result_proto PRIVATE ${benchmark_result_generated_files}) -target_link_libraries(benchmark_result_proto protobuf::libprotobuf) -target_include_directories(benchmark_result_proto PUBLIC ${CMAKE_BINARY_DIR}) diff --git a/tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake b/tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake new file mode 100644 index 00000000000000..7fef6f8c81c55e --- /dev/null +++ b/tensorflow/lite/tools/cmake/modules/Findtensorflow_headers.cmake @@ -0,0 +1,77 @@ +# +# Copyright 2026 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Downloads the few files from the TensorFlow repository that the TFLite runtime +# and the benchmark_model tool use (xla/tsl/..., tensorflow/core/util/... and +# third_party/fft2d/...), instead of cloning the whole TensorFlow repository. +# +# Sets TENSORFLOW_HEADERS_INCLUDE_DIR. + +if(TENSORFLOW_HEADERS_INCLUDE_DIR) + set(tensorflow_headers_FOUND TRUE) + return() +endif() + +set(_tensorflow_headers_base_url + "https://raw.githubusercontent.com/tensorflow/tensorflow/v2.21.0-rc0" +) +set(_tensorflow_headers_dir "${CMAKE_BINARY_DIR}/tensorflow_headers") +# || +set(_tensorflow_headers_files + "third_party/fft2d/fft.h|third_party/fft2d/fft.h|2db045d17dfd4b4fa5201e86a1653f0c0b7741e14c927f0c426007127109e825" + "third_party/fft2d/fft2d.h|third_party/fft2d/fft2d.h|b24c63e77d5daf3affd7386085c41c046f3ad0c8dcab56a621fd8e84a5af5c1a" + "third_party/xla/xla/tsl/framework/convolution/eigen_convolution_helpers.h|xla/tsl/framework/convolution/eigen_convolution_helpers.h|3fd52ebb0f14b9c4f3b3e225d1e7b29c712c3046d34e07936bfec0e0c1f41152" + "third_party/xla/xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h|xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h|c0bec189723d52b4495dec62e197e2a4ffe7725bc313676984384e8b8fbb7a13" + "third_party/xla/xla/tsl/lib/random/philox_random.h|xla/tsl/lib/random/philox_random.h|7a7659f95c59419373261af305736311bfd839b281259703b23f4df9b391a43e" + "third_party/xla/xla/tsl/lib/random/random_distributions_utils.h|xla/tsl/lib/random/random_distributions_utils.h|1d65158a878510a1ec5835c5f26962b0ebe45bd56ecaec7e38314a8ac5e7a517" + # Used by the benchmark_model tool. + "tensorflow/core/util/stat_summarizer_options.h|tensorflow/core/util/stat_summarizer_options.h|e91f94b2fbb7e3a8d8b99664a1ff2443088ca1433b4372b1fcd2477b78492b3b" + "tensorflow/core/util/stats_calculator.h|tensorflow/core/util/stats_calculator.h|8d80ca00444162a55a288b3b7369f314e5cf65a3258387aef4af1df39db08e0a" + "third_party/xla/xla/tsl/util/stat_summarizer_options.h|xla/tsl/util/stat_summarizer_options.h|9f7d7cc5de38ae6e97a8982dc386a012ed70144c768ef2525fd25054de41c2b2" + "third_party/xla/xla/tsl/util/stats_calculator.cc|xla/tsl/util/stats_calculator.cc|704c6d22240a521a7cbee591cd932121f0f6dcbaf179037758a7f3bd4d851e82" + "third_party/xla/xla/tsl/util/stats_calculator.h|xla/tsl/util/stats_calculator.h|f5628ba1fbf39e7c4d2b280ff78c8daba40d9d462a97f8aa2abc0c5371b9f241" +) + +foreach(_entry IN LISTS _tensorflow_headers_files) + string(REPLACE "|" ";" _entry "${_entry}") + list(LENGTH _entry _entry_len) + if(NOT _entry_len EQUAL 3) + message(FATAL_ERROR "Malformed entry in _tensorflow_headers_files: ${_entry}") + endif() + list(GET _entry 0 _src) + list(GET _entry 1 _path) + list(GET _entry 2 _sha256) + set(_dst "${_tensorflow_headers_dir}/${_path}") + if(EXISTS "${_dst}") + file(SHA256 "${_dst}" _actual) + if(_actual STREQUAL _sha256) + continue() + endif() + endif() + message(STATUS "Downloading ${_src}") + file(DOWNLOAD "${_tensorflow_headers_base_url}/${_src}" "${_dst}" + EXPECTED_HASH SHA256=${_sha256} + STATUS _status + ) + list(GET _status 0 _code) + if(NOT _code EQUAL 0) + message(FATAL_ERROR "Failed to download ${_src}: ${_status}") + endif() +endforeach() + +set(TENSORFLOW_HEADERS_INCLUDE_DIR "${_tensorflow_headers_dir}" CACHE PATH + "Include directory with the TensorFlow headers used by TFLite" +) +set(tensorflow_headers_FOUND TRUE) diff --git a/tensorflow/python/eager/polymorphic_function/polymorphic_function_test_cpu_only.py b/tensorflow/python/eager/polymorphic_function/polymorphic_function_test_cpu_only.py index f1aad51a95b6d2..da01ed73c49113 100644 --- a/tensorflow/python/eager/polymorphic_function/polymorphic_function_test_cpu_only.py +++ b/tensorflow/python/eager/polymorphic_function/polymorphic_function_test_cpu_only.py @@ -33,9 +33,8 @@ def testJitCompileRaisesExceptionWhenXlaIsUnsupported(self): if test.is_built_with_rocm() or test_util.is_xla_enabled(): return - with self.assertRaisesRegex( - errors.NotFoundError, 'The PJRT client factory' - ): + with self.assertRaisesRegex(errors.UnimplementedError, + 'support for that platform linked in'): @polymorphic_function.function(jit_compile=True) def fn(x): diff --git a/tensorflow/python/eager/pywrap_gradient_exclusions.cc b/tensorflow/python/eager/pywrap_gradient_exclusions.cc index 10f18cdf75d991..3ff8a8ce366a1c 100644 --- a/tensorflow/python/eager/pywrap_gradient_exclusions.cc +++ b/tensorflow/python/eager/pywrap_gradient_exclusions.cc @@ -428,7 +428,7 @@ absl::optional> OpGradientUnusedInputIndices( absl::optional> OpGradientUnusedOutputIndices( const tensorflow::string &op_name) { - static std::array a = {{ + static std::array a = {{ {"Abs"}, {"AccumulateNV2"}, {"Acos"}, @@ -507,7 +507,6 @@ absl::optional> OpGradientUnusedOutputIndices( {"CrossReplicaSum"}, {"Cumprod"}, {"Cumsum"}, - {"CumulativeLogsumexp"}, {"DecodeBase64"}, {"DecodePaddedRaw"}, {"DecodeRaw"}, diff --git a/tensorflow/python/eager/record.py b/tensorflow/python/eager/record.py index ea6c4ae4f8338e..848d08d4b17f09 100644 --- a/tensorflow/python/eager/record.py +++ b/tensorflow/python/eager/record.py @@ -30,7 +30,7 @@ class VariableWatcher(object): with VariableWatcher() as variable_watcher: var.assign_add(1.0) - assert variable_watcher.watched_variables == [var] + assert variable_watcher.watched_variables() == [var] """ __slots__ = ["_variable_watcher"] diff --git a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py index 9665e7917c50bf..47f3001eff1aa2 100644 --- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py +++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py @@ -160,6 +160,7 @@ def CheckVersusNumpy(self, ndims_mask, arr_shape, make_mask=None, axis=None): mask = make_mask(arr_shape[:ndims_mask]) if axis is not None: mask = make_mask(arr_shape[axis:ndims_mask + axis]) + masked_arr = None if axis is None or axis == 0: masked_arr = arr[mask] elif axis == 1: @@ -2356,6 +2357,60 @@ def g(): _ = g.get_concrete_function() + def testNaN(self): + """NaN should always be placed at the end of a sorted sequence.""" + # Test NaN in values + sorted_sequence = np.array( + [2.0, 4.0, 8.0, 16.0, 32.0, 64.0], dtype=np.float32 + ) + values = np.array([np.nan, 8.0], dtype=np.float32) + for side in ("left", "right"): + with self.subTest(side=side, case="nan_in_values"): + expected = np.searchsorted(sorted_sequence, values, side=side) + actual = self.evaluate( + array_ops.searchsorted(sorted_sequence, values, side=side) + ) + self.assertAllEqual(expected, actual) + + # Test NaN in sorted_sequence + sorted_sequence_with_nan = np.array( + [2.0, 4.0, 8.0, np.nan, np.nan], dtype=np.float32 + ) + values_to_search = np.array([3.0, np.nan], dtype=np.float32) + for side in ("left", "right"): + with self.subTest(side=side, case="nan_in_sequence"): + expected = np.searchsorted( + sorted_sequence_with_nan, values_to_search, side=side + ) + actual = self.evaluate( + array_ops.searchsorted( + sorted_sequence_with_nan, values_to_search, side=side + ) + ) + self.assertAllEqual(expected, actual) + + # Regression test: NaN in sorted_sequence with multiple NaNs, searching + # for finite values that fall before, between, and after non-NaN elements. + sorted_sequence_with_nans = np.array( + [2.0, 4.0, np.nan, np.nan], dtype=np.float32 + ) + for val in [3.0, 5.0, np.nan]: + for side in ("left", "right"): + with self.subTest(side=side, val=val, case="nan_in_sorted_regression"): + expected = np.searchsorted( + sorted_sequence_with_nans, + np.array([val], dtype=np.float32), + side=side, + ) + actual = self.evaluate( + array_ops.searchsorted( + sorted_sequence_with_nans, + np.array([val], dtype=np.float32), + side=side, + ) + ) + self.assertAllEqual(expected, actual) + def testInvalidValuesLowerBound(self): arg_0_tensor = random_ops.random_uniform([3, 3], dtype=dtypes.float32) arg_0 = array_ops.identity(arg_0_tensor) diff --git a/tensorflow/python/kernel_tests/array_ops/huge_slice_op_test.py b/tensorflow/python/kernel_tests/array_ops/huge_slice_op_test.py index 3df288faf21d46..bb702cca84e902 100644 --- a/tensorflow/python/kernel_tests/array_ops/huge_slice_op_test.py +++ b/tensorflow/python/kernel_tests/array_ops/huge_slice_op_test.py @@ -33,7 +33,11 @@ def testInt64Slicing(self): self.assertAllEqual([True, False, True], slice_val) slice_t = array_ops.slice( - a_large, constant_op.constant([long(2)**32 + 3], dtype=dtypes.int64), - [3]) + a_large, constant_op.constant([2**32 + 3], dtype=dtypes.int64), [3] + ) slice_val = self.evaluate(slice_t) self.assertAllEqual([True, False, True], slice_val) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/python/kernel_tests/data_structures/tensor_array_ops_test.py b/tensorflow/python/kernel_tests/data_structures/tensor_array_ops_test.py index 7fd15a50c18e75..8550256de304d9 100644 --- a/tensorflow/python/kernel_tests/data_structures/tensor_array_ops_test.py +++ b/tensorflow/python/kernel_tests/data_structures/tensor_array_ops_test.py @@ -642,6 +642,46 @@ def testTensorArraySplitIncompatibleShapesFails(self): r"dynamically resizeable"): self.evaluate(ta.split([1.0], [1]).flow) + def testTensorArraySplitRejectsNegativeLengths(self): + if context.executing_eagerly(): + self.skipTest("TensorArraySplitOp requires graph mode") + + for lengths, index in [([-1, 1], 0), ([1, -1], 1)]: + with self.subTest(lengths=lengths), self.session(): + handle, flow = gen_data_flow_ops.tensor_array_v3( + size=2, dtype=dtypes.float32, tensor_array_name="negative_lengths" + ) + with self.assertRaisesOpError( + "Expected lengths to be non-negative, but found -1 at index " + f"{index}" + ): + self.evaluate( + gen_data_flow_ops.tensor_array_split_v3( + handle=handle, + value=constant_op.constant([], dtype=dtypes.float32), + lengths=lengths, + flow_in=flow, + ) + ) + + def testTensorArraySplitRejectsOverflowingLengths(self): + if context.executing_eagerly(): + self.skipTest("TensorArraySplitOp requires graph mode") + + with self.session(): + handle, flow = gen_data_flow_ops.tensor_array_v3( + size=2, dtype=dtypes.float32, tensor_array_name="overflowing_lengths" + ) + with self.assertRaisesOpError("TensorArraySplit lengths overflowed"): + self.evaluate( + gen_data_flow_ops.tensor_array_split_v3( + handle=handle, + value=constant_op.constant([], dtype=dtypes.float32), + lengths=[np.iinfo(np.int64).max, 1], + flow_in=flow, + ) + ) + def _testTensorArrayWriteGradientAddMultipleAdds(self, dtype): with self.cached_session(): ta = tensor_array_ops.TensorArray( diff --git a/tensorflow/python/kernel_tests/io_ops/parsing_ops_test.py b/tensorflow/python/kernel_tests/io_ops/parsing_ops_test.py index 55d61d3c451cf1..38145d21e08993 100644 --- a/tensorflow/python/kernel_tests/io_ops/parsing_ops_test.py +++ b/tensorflow/python/kernel_tests/io_ops/parsing_ops_test.py @@ -2436,6 +2436,17 @@ def testSerializedContainingMisalignedNestedRaggedFeature(self): # Message for batch=false in eager mode: "|Incompatible shapes|required broadcastable shapes")) + def testFixedLenSequenceFeatureZeroDimDoesNotCrash(self): + # Regression test for https://github.com/tensorflow/tensorflow/issues/123476 + # ParseSequenceDenseFeatures divided by row_shape.num_elements() without + # guarding against the zero case, causing a SIGFPE when shape=[0]. + # The public Python API now rejects shape=[0] with a recoverable error + # rather than crashing the process; verify the error is catchable. + with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)): + parsing_ops.FixedLenSequenceFeature( + shape=[0], dtype=dtypes.float32, allow_missing=True + ) + @test_util.run_all_in_graph_and_eager_modes class DecodeRawTest(test.TestCase): diff --git a/tensorflow/python/kernel_tests/linalg/linalg_ops_test.py b/tensorflow/python/kernel_tests/linalg/linalg_ops_test.py index 755ece0f3e1f84..fcb0b0d667207f 100644 --- a/tensorflow/python/kernel_tests/linalg/linalg_ops_test.py +++ b/tensorflow/python/kernel_tests/linalg/linalg_ops_test.py @@ -60,7 +60,12 @@ def test_works_with_five_different_random_pos_def_matrices(self): with self.subTest(n=n, np_type=np_type, atol=atol, k=k): rhs = self.rng.randn(2, n, k).astype(np_type) x = linalg_ops.cholesky_solve(chol, rhs) - self.assertAllClose(rhs, math_ops.matmul(array, x), atol=atol) + rhs_pred = ( + test_util.matmul_without_tf32(array, x) + if test.is_gpu_available() + else math_ops.matmul(array, x) + ) + self.assertAllClose(rhs, rhs_pred, atol=atol) class LogdetTest(test.TestCase): @@ -435,7 +440,10 @@ def test_symmetric(self): expected_a_pinv_ = self.expected_pinv(a_, rcond) a_pinv = linalg.pinv(a, rcond, validate_args=True) a_pinv_ = self.evaluate(a_pinv) - self.assertAllClose(expected_a_pinv_, a_pinv_, atol=2e-5, rtol=2e-5) + use_tf32_tol = self.dtype == np.float32 and test.is_gpu_available() + atol = 3e-3 if use_tf32_tol else 2e-5 + rtol = 1e-3 if use_tf32_tol else 2e-5 + self.assertAllClose(expected_a_pinv_, a_pinv_, atol=atol, rtol=rtol) if not self.use_static_shape: return self.assertAllEqual(expected_a_pinv_.shape, a_pinv.shape) @@ -453,7 +461,10 @@ def test_nonsquare(self): expected_a_pinv_ = self.expected_pinv(a_, rcond) a_pinv = linalg.pinv(a, rcond, validate_args=True) a_pinv_ = self.evaluate(a_pinv) - self.assertAllClose(expected_a_pinv_, a_pinv_, atol=1e-5, rtol=1e-4) + use_tf32_tol = self.dtype == np.float32 and test.is_gpu_available() + atol = 1e-3 if use_tf32_tol else 1e-5 + rtol = 1e-3 if use_tf32_tol else 1e-4 + self.assertAllClose(expected_a_pinv_, a_pinv_, atol=atol, rtol=rtol) if not self.use_static_shape: return self.assertAllEqual(expected_a_pinv_.shape, a_pinv.shape) diff --git a/tensorflow/python/kernel_tests/math_ops/approx_topk_test.py b/tensorflow/python/kernel_tests/math_ops/approx_topk_test.py index a182f442df0690..c9b2aeca400620 100644 --- a/tensorflow/python/kernel_tests/math_ops/approx_topk_test.py +++ b/tensorflow/python/kernel_tests/math_ops/approx_topk_test.py @@ -163,7 +163,14 @@ def ann(qy, db, db_half_norm_sq, k): math_ops.matmul(qy_op, db_op, transpose_b=True)) gt = np.argsort(scores)[:, :k] ann_recall = self.compute_recall(idx, gt) - self.assertGreaterEqual(ann_recall, 0.95) + use_tf32_tol = ( + test_util.is_gpu_available() + and dtype == dtypes.float32 + and k == 1 + and feature_dim == 2 + ) + min_recall = 0.80 if use_tf32_tol else 0.95 + self.assertGreaterEqual(ann_recall, min_recall) def test_highdim(self): db = self._rng.random([2, 10, 200, 3], dtype=np.float32) diff --git a/tensorflow/python/kernel_tests/math_ops/batch_matmul_op_test.py b/tensorflow/python/kernel_tests/math_ops/batch_matmul_op_test.py index c3182d259b1d59..7f1812e71a3674 100644 --- a/tensorflow/python/kernel_tests/math_ops/batch_matmul_op_test.py +++ b/tensorflow/python/kernel_tests/math_ops/batch_matmul_op_test.py @@ -168,7 +168,8 @@ def testRank4xRank2(self): w = np.random.normal(size=[4, 16]).astype(np.float32) expected = np.matmul(p, w) result = math_ops.matmul(p, w) - self.assertAllClose(self.evaluate(result), expected, rtol=1e-5, atol=1e-5) + tol = 1e-3 if test.is_gpu_available() else 1e-5 + self.assertAllClose(self.evaluate(result), expected, rtol=tol, atol=tol) def testRank4xRank2BatchDim(self): np.random.seed(42) @@ -176,7 +177,8 @@ def testRank4xRank2BatchDim(self): w = np.random.normal(size=[5, 7]).astype(np.float32) expected = np.matmul(p, w) result = math_ops.matmul(p, w) - self.assertAllClose(self.evaluate(result), expected, rtol=1e-5, atol=1e-5) + tol = 1e-3 if test.is_gpu_available() else 1e-5 + self.assertAllClose(self.evaluate(result), expected, rtol=tol, atol=tol) def testRank3xRank2(self): np.random.seed(42) @@ -184,7 +186,8 @@ def testRank3xRank2(self): w = np.random.normal(size=[4, 6]).astype(np.float32) expected = np.matmul(p, w) result = math_ops.matmul(p, w) - self.assertAllClose(self.evaluate(result), expected, rtol=1e-5, atol=1e-5) + tol = 1e-3 if test.is_gpu_available() else 1e-5 + self.assertAllClose(self.evaluate(result), expected, rtol=tol, atol=tol) class BatchMatmulGradientTest(test.TestCase): diff --git a/tensorflow/python/kernel_tests/math_ops/cumulative_logsumexp_test.py b/tensorflow/python/kernel_tests/math_ops/cumulative_logsumexp_test.py index 888e8afa7fb052..c192d69e866b32 100644 --- a/tensorflow/python/kernel_tests/math_ops/cumulative_logsumexp_test.py +++ b/tensorflow/python/kernel_tests/math_ops/cumulative_logsumexp_test.py @@ -16,13 +16,9 @@ import numpy as np -from tensorflow.python.eager import backprop -from tensorflow.python.eager import forwardprop -from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops -from tensorflow.python.ops import array_ops_stack from tensorflow.python.ops import gradient_checker_v2 from tensorflow.python.ops import map_fn from tensorflow.python.ops import math_ops @@ -148,113 +144,6 @@ def testPlusInfinityAllInf(self): msg=f'Expected +inf outputs for all-inf input, got {result}', ) - def testSecondDerivativeNestedForwardAccumulator(self): - # Regression test for GitHub issue #127243: - # cumulative_logsumexp produces NaN second derivative under nested - # forward-mode autodiff. - def target(t): - x = array_ops.reshape( - array_ops_stack.stack([t, -t + 1, t / 2 - 2, -2 * t, t + 3, t / 4]), - [2, 3], - ) - y = math_ops.cumulative_logsumexp( - x, axis=-1, exclusive=False, reverse=True - ) - return ( - y[0, 0] + 2 * y[0, 1] - y[0, 2] + 3 * y[1, 0] + y[1, 1] - 2 * y[1, 2] - ) - - t = constant_op.constant(-40.0, dtype=dtypes.float64) - dt = constant_op.constant(1.0, dtype=dtypes.float64) - expected = 1.0572349592992632e-12 - - # Verify forward and 1st derivative - self.assertAllClose(self.evaluate(target(t)), 395.0000000000019) - - # 1. Reverse-over-Reverse - with backprop.GradientTape() as t2: - t2.watch(t) - with backprop.GradientTape() as t1: - t1.watch(t) - y = target(t) - g1 = t1.gradient(y, t) - rr = t2.gradient(g1, t) - self.assertAllClose(self.evaluate(g1), -9.75, atol=1e-10) - self.assertAllClose(self.evaluate(rr), expected, atol=1e-12) - - # 2. Forward-over-Reverse - with forwardprop.ForwardAccumulator(t, dt) as acc: - with backprop.GradientTape() as t1: - t1.watch(t) - y = target(t) - g1 = t1.gradient(y, t) - fr = acc.jvp(g1) - self.assertAllClose(self.evaluate(fr), expected, atol=1e-12) - - # 3. Reverse-over-Forward - with backprop.GradientTape() as t1: - t1.watch(t) - with forwardprop.ForwardAccumulator(t, dt) as acc: - y = target(t) - j1 = acc.jvp(y) - rf = t1.gradient(j1, t) - self.assertAllClose(self.evaluate(rf), expected, atol=1e-12) - - # 4. Forward-over-Forward (nested ForwardAccumulator) - with forwardprop.ForwardAccumulator(t, dt) as acc2: - with forwardprop.ForwardAccumulator(t, dt) as acc1: - y = target(t) - j1 = acc1.jvp(y) - ff = acc2.jvp(j1) - self.assertAllClose(self.evaluate(j1), -9.75, atol=1e-10) - self.assertAllClose(self.evaluate(ff), expected, atol=1e-12) - - # Verify positive, negative, and zero cotangents across supported dtypes - tol_map = { - dtypes.float64: 1e-12, - dtypes.float32: 1e-5, - dtypes.float16: 1e-2, - dtypes.bfloat16: 2e-1, - } - for reverse in (True, False): - for exclusive in (True, False): - x64 = constant_op.constant([1.0, -2.0, 0.5, 3.0], dtype=dtypes.float64) - dx64 = constant_op.constant([1.0, -1.0, 0.0, 0.5], dtype=dtypes.float64) - w64 = constant_op.constant([2.0, -3.0, 0.0, 1.0], dtype=dtypes.float64) - with backprop.GradientTape() as t2: - t2.watch(x64) - with backprop.GradientTape() as t1: - t1.watch(x64) - lse64 = math_ops.cumulative_logsumexp( - x64, reverse=reverse, exclusive=exclusive - ) - out64 = math_ops.reduce_sum( - (math_ops.exp(lse64) if exclusive else lse64) * w64 - ) - g1 = t1.gradient(out64, x64) - dir_g1 = math_ops.reduce_sum(g1 * dx64) - g2 = t2.gradient(dir_g1, x64) - ref_jvp1 = self.evaluate(dir_g1) - ref_jvp2 = self.evaluate(math_ops.reduce_sum(g2 * dx64)) - - for dtype in self.valid_dtypes: - x_val = constant_op.constant([1.0, -2.0, 0.5, 3.0], dtype=dtype) - dx_val = constant_op.constant([1.0, -1.0, 0.0, 0.5], dtype=dtype) - weights = constant_op.constant([2.0, -3.0, 0.0, 1.0], dtype=dtype) - with forwardprop.ForwardAccumulator(x_val, dx_val) as acc2: - with forwardprop.ForwardAccumulator(x_val, dx_val) as acc1: - lse = math_ops.cumulative_logsumexp( - x_val, reverse=reverse, exclusive=exclusive - ) - out = math_ops.reduce_sum( - (math_ops.exp(lse) if exclusive else lse) * weights - ) - jvp1 = acc1.jvp(out) - jvp2 = acc2.jvp(jvp1) - tol = tol_map[dtype] - self.assertAllClose(self.evaluate(jvp1), ref_jvp1, rtol=tol, atol=tol) - self.assertAllClose(self.evaluate(jvp2), ref_jvp2, rtol=tol, atol=tol) - if __name__ == '__main__': test.main() diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_binary_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_binary_test.py index b0cf1d0058c1a5..4d6bf5c3e77d93 100644 --- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_binary_test.py +++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_binary_test.py @@ -306,6 +306,31 @@ def testDoubleBasic(self): except ImportError as e: tf_logging.warn("Cannot test special functions: %s" % str(e)) + @test_util.run_deprecated_v1 + def testIgammaDomainEdgeCases(self): + # P(a, x) is undefined for a <= 0; x == 0 short-circuits before the domain + # check in Eigen, so the kernel must return NaN for those inputs explicitly. + for dtype in [np.float32, np.float64]: + a_vals = np.array([-0.1, -1.0, 0.0, np.nan], dtype=dtype) + x_vals = np.array([0.0, 0.0, 0.0, 0.0], dtype=dtype) + with self.cached_session(): + result = math_ops.igamma( + constant_op.constant(a_vals), constant_op.constant(x_vals) + ) + result_np = self.evaluate(result) + self.assertTrue( + np.all(np.isnan(result_np)), + "Expected NaN for out-of-domain (a<=0, x==0), got %s" % result_np, + ) + # P(a, 0) == 0 for a > 0; the fix must not disturb this identity. + a_pos = np.array([0.5, 1.0, 2.0], dtype=dtype) + x_zero = np.array([0.0, 0.0, 0.0], dtype=dtype) + with self.cached_session(): + result_pos = math_ops.igamma( + constant_op.constant(a_pos), constant_op.constant(x_zero) + ) + self.assertAllEqual(self.evaluate(result_pos), np.zeros(3, dtype=dtype)) + def testBfloat16Basic(self): bf16_np = dtypes_lib.bfloat16.as_numpy_dtype x = np.linspace(-5, 20, 15).reshape(1, 3, 5).astype(bf16_np) # pylint: disable=too-many-function-args diff --git a/tensorflow/python/kernel_tests/nn_ops/BUILD b/tensorflow/python/kernel_tests/nn_ops/BUILD index ca0b76ac6bbd9a..045ebe9ebbdee5 100644 --- a/tensorflow/python/kernel_tests/nn_ops/BUILD +++ b/tensorflow/python/kernel_tests/nn_ops/BUILD @@ -436,9 +436,11 @@ py_library( srcs = ["depthwise_conv_op_base.py"], strict_deps = True, deps = [ + "//tensorflow/python/eager:def_function", "//tensorflow/python/framework:constant_op", "//tensorflow/python/framework:errors", "//tensorflow/python/framework:for_generated_wrappers", + "//tensorflow/python/framework:tensor_spec", "//tensorflow/python/framework:test_lib", "//tensorflow/python/ops:array_ops", "//tensorflow/python/ops:gradient_checker", diff --git a/tensorflow/python/kernel_tests/nn_ops/atrous_convolution_test.py b/tensorflow/python/kernel_tests/nn_ops/atrous_convolution_test.py index 23665706b06c08..422b2ac0602c39 100644 --- a/tensorflow/python/kernel_tests/nn_ops/atrous_convolution_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/atrous_convolution_test.py @@ -19,8 +19,11 @@ import numpy as np from tensorflow.python.eager import context +from tensorflow.python.eager import def_function from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes +from tensorflow.python.framework import errors +from tensorflow.python.framework import tensor_spec from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import gradient_checker @@ -278,6 +281,47 @@ def testInvalidInputRank(self): with self.assertRaisesRegex(ValueError, "rank at least 3"): nn_ops.convolution(value, filters, padding="VALID") + def testDilatedConvolutionFailsWhenFilterExceedsInput(self): + """Eager and `tf.function` must fail the same way (#113319).""" + x = array_ops.zeros([2, 5, 5, 3]) + f = array_ops.zeros([3, 3, 3, 4]) # spans 7 > 5 once dilated by 3 + for padding in ("VALID", [[0, 0], [0, 0], [0, 0], [0, 0]]): + with self.assertRaisesRegex(ValueError, "no valid window"): + nn_ops.convolution(x, f, padding=padding, dilation_rate=[3, 3]) + + @def_function.function + def fn(a, b, padding=padding): + return nn_ops.convolution(a, b, padding=padding, dilation_rate=[3, 3]) + + with self.assertRaisesRegex(ValueError, "no valid window"): + fn(x, f) + + def testDilatedConvolutionFailsAtRuntimeForUnknownShapes(self): + """Shapes unknown at trace time are caught by the run-time assertion.""" + fn = def_function.function( + lambda a, b: nn_ops.convolution( + a, b, padding="VALID", dilation_rate=[3, 3] + ) + ) + concrete = fn.get_concrete_function( + tensor_spec.TensorSpec([None, None, None, 3], dtypes.float32), + tensor_spec.TensorSpec([3, 3, 3, 4], dtypes.float32), + ) + # `tf.function` converts the `InvalidArgumentError` raised by the + # run-time assertion into a Python `ValueError`, so accept both spellings + # of the same failure. + with self.assertRaises((errors.InvalidArgumentError, ValueError)): + concrete(array_ops.zeros([2, 5, 5, 3]), array_ops.zeros([3, 3, 3, 4])) + + def testDilatedConvolutionFitsWithExplicitPadding(self): + """Explicit padding that admits a window is not rejected.""" + x = array_ops.ones([1, 4, 4, 2]) + f = array_ops.ones([2, 2, 2, 3]) + y = nn_ops.convolution( + x, f, padding=[[0, 0], [1, 1], [1, 1], [0, 0]], dilation_rate=[2, 2] + ) + self.assertEqual(y.shape.as_list(), [1, 4, 4, 3]) + if __name__ == "__main__": test.main() diff --git a/tensorflow/python/kernel_tests/nn_ops/conv3d_backprop_filter_v2_grad_test.py b/tensorflow/python/kernel_tests/nn_ops/conv3d_backprop_filter_v2_grad_test.py index 9d0387da61b5bd..fe06f42e532577 100644 --- a/tensorflow/python/kernel_tests/nn_ops/conv3d_backprop_filter_v2_grad_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/conv3d_backprop_filter_v2_grad_test.py @@ -56,7 +56,7 @@ def testGradient(self): [in_val, out_backprop_val], [in_shape, out_backprop_shape], output, filter_shape) print("conv3d_backprop_filter gradient err = %g " % err) - err_tolerance = 1e-3 + err_tolerance = 0.25 if test.is_gpu_available() else 1e-3 self.assertLess(err, err_tolerance) def testBadFilterShape(self): diff --git a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py index 7d3aa8fdd59c88..57cdfff5889c0e 100644 --- a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py +++ b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py @@ -1705,8 +1705,9 @@ def _GetVal(data_format, use_gpu): values = [] for (data_format, use_gpu) in GetTestConfigs(): values.append(_GetVal(data_format, use_gpu)) + tol = 1e-3 if test.is_gpu_available() else 2e-4 for i in range(1, len(values)): - self.assertAllClose(values[0], values[i], rtol=2e-4, atol=2e-4) + self.assertAllClose(values[0], values[i], rtol=tol, atol=tol) @test_util.run_in_graph_and_eager_modes def testConv2D2x2Depth1ValidBackpropFilter(self): @@ -3795,7 +3796,7 @@ def Test(self): strides=[stride, stride], dilations=[2, 2], padding=padding, - rtol=5e-4) + rtol=1e-3 if test.is_gpu_available() else 5e-4) return Test diff --git a/tensorflow/python/kernel_tests/nn_ops/depthwise_conv_op_base.py b/tensorflow/python/kernel_tests/nn_ops/depthwise_conv_op_base.py index 4e466a0a1c876a..5be78b3fc32b39 100644 --- a/tensorflow/python/kernel_tests/nn_ops/depthwise_conv_op_base.py +++ b/tensorflow/python/kernel_tests/nn_ops/depthwise_conv_op_base.py @@ -16,10 +16,12 @@ import numpy as np +from tensorflow.python.eager import def_function from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_spec from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import gradient_checker @@ -1170,3 +1172,62 @@ def testDepthwiseConv2DForwardExplicitCompare(self): self._CompareForward( input_size, filter_size, output_size, stride, padding, "float64" ) + + def testDilatedDepthwiseConv2DFailsWhenFilterExceedsInput(self): + """A dilated filter larger than the input fails in every mode (#113320). + + With unit strides the dilated filter can still exceed the input, leaving no + valid window. Eager silently returned a result of the wrong shape while + `tf.function` failed; both must now fail, including for unknown spatial + shapes that only the run-time assertion can catch. + """ + x = array_ops.ones([2, 4, 4, 3]) + # Dilation 4 expands the 2x2 filter to 5x5, which does not fit in 4x4. + f = array_ops.ones([2, 2, 3, 1]) + with self.assertRaisesRegex(ValueError, "no valid window"): + nn_impl.depthwise_conv2d(x, f, [1, 1, 1, 1], "VALID", dilations=[4, 4]) + + @def_function.function + def dilated(a, b): + return nn_impl.depthwise_conv2d( + a, b, [1, 1, 1, 1], "VALID", dilations=[4, 4] + ) + + with self.assertRaises(ValueError): + dilated(x, f) + + @def_function.function( + input_signature=[ + tensor_spec.TensorSpec([None, None, None, 3], dtypes.float32), + tensor_spec.TensorSpec([2, 2, 3, 1], dtypes.float32), + ] + ) + def dilated_dynamic(a, b): + return nn_impl.depthwise_conv2d( + a, b, [1, 1, 1, 1], "VALID", dilations=[4, 4] + ) + + with self.assertRaises((errors.InvalidArgumentError, ValueError)): + dilated_dynamic(x, f) + + def testDilatedDepthwiseConv2DWithUnitStrides(self): + """Dilated depthwise conv with unit strides still matches numpy.""" + x_np = (np.arange(1 * 9 * 9 * 2, dtype=np.float32) % 5 + 1).reshape( + [1, 9, 9, 2] + ) + f_np = (np.arange(3 * 3 * 2 * 1, dtype=np.float32) % 3 + 1).reshape( + [3, 3, 2, 1] + ) + for rate in (2, 3): + for padding in ("VALID", "SAME"): + expected = _DepthwiseConv2dNumpy( + x_np, f_np, [1, 1, 1, 1], padding, "NHWC", [rate, rate] + ) + actual = nn_impl.depthwise_conv2d( + constant_op.constant(x_np), + constant_op.constant(f_np), + [1, 1, 1, 1], + padding, + dilations=[rate, rate], + ) + self.assertAllClose(expected, actual, rtol=1e-5, atol=1e-5) diff --git a/tensorflow/python/kernel_tests/random/stateless_random_ops_test.py b/tensorflow/python/kernel_tests/random/stateless_random_ops_test.py index fad01517b67d2c..5d59970e3a0f9a 100644 --- a/tensorflow/python/kernel_tests/random/stateless_random_ops_test.py +++ b/tensorflow/python/kernel_tests/random/stateless_random_ops_test.py @@ -171,9 +171,7 @@ def multinomial_cases(): def wrap(op: ops.Operation, logits, logits_dtype, output_dtype, seed): device_type = get_device().device_type # Some dtypes are not supported on some devices - if (logits_dtype == dtypes.bfloat16 and device_type == 'GPU' and - not test_util.is_gpu_available( - cuda_only=True, min_cuda_compute_capability=(8, 0))): + if logits_dtype == dtypes.bfloat16 and device_type == 'GPU': logits_dtype = dtypes.float32 return op(seed=seed, logits=constant_op.constant(logits, dtype=logits_dtype), diff --git a/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py b/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py index 93640351eb1fe7..21e4af862b730a 100644 --- a/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py +++ b/tensorflow/python/kernel_tests/variables/resource_variable_ops_test.py @@ -1522,9 +1522,37 @@ def testScatterUpdateInvalidArgs(self): # eager execution (where the error is realized during kernel execution), # and XLA auto-clustering execution (where the error is realized in the xla # op kernel) which is triggered when running in eager op as function mode. - with self.assertRaisesRegex(Exception, r"shape.*2.*3|RET_CHECK failure"): + with self.assertRaisesRegex( + Exception, r"Must have updates\.shape|shape.*2.*3|RET_CHECK failure" + ): state_ops.scatter_update(v, [0, 1], [0, 1, 2]) + @test_util.run_in_graph_and_eager_modes + def testScatterUpdateInvalidInnerShape(self): + v = resource_variable_ops.ResourceVariable( + array_ops.zeros([10, 0]), shape=tensor_shape.TensorShape(None) + ) + self.evaluate(variables.global_variables_initializer()) + with self.assertRaisesRegex( + errors.InvalidArgumentError, r"Must have updates\.shape" + ): + self.evaluate( + resource_variable_ops.resource_scatter_update( + v.handle, + constant_op.constant(-1, dtype=dtypes.int64), + array_ops.zeros([5]), + ) + ) + + @test_util.run_gpu_only + def testScatterUpdateDoesNotNarrowInt64Index(self): + with context.eager_mode(), ops.device("gpu:0"): + index = constant_op.constant([np.iinfo(np.int64).min], dtype=dtypes.int64) + for updates in (constant_op.constant(2.0), constant_op.constant([2.0])): + v = resource_variable_ops.ResourceVariable([1.0]) + resource_variable_ops.resource_scatter_update(v.handle, index, updates) + self.assertAllEqual([1.0], v.numpy()) + @test_util.disable_xla("b/208334252") # XLA doesn't have a deterministic impl def testScatterAddDeterministic(self): with context.eager_mode(), test_util.deterministic_ops(): diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD index 37f2a0af409414..2bcda6a1843e58 100644 --- a/tensorflow/python/ops/BUILD +++ b/tensorflow/python/ops/BUILD @@ -4376,6 +4376,7 @@ py_library( ":array_ops", ":array_ops_stack", ":candidate_sampling_ops", + ":check_ops", ":cond", ":ctc_ops", ":custom_gradient", diff --git a/tensorflow/python/ops/control_flow_ops_test.py b/tensorflow/python/ops/control_flow_ops_test.py index 5fb2757295cd70..fdc431faabc1cb 100644 --- a/tensorflow/python/ops/control_flow_ops_test.py +++ b/tensorflow/python/ops/control_flow_ops_test.py @@ -1663,22 +1663,57 @@ def testWhileLoopSameReturnShape_TrueSingleLoopVar(self): @test_util.run_v2_only def testEagerWhileLoopSingleLoopVarBareTensorBodyPreservesShape(self): - x = constant_op.constant([[5.0]], shape=[1, 1]) + x = array_ops.ones([1, 28, 28, 1]) shapes_seen = [] - def cond(x): - shapes_seen.append(x.shape.as_list()) - return math_ops.greater(x[0, 0], 3) + def cond(a): + shapes_seen.append(a.shape.as_list()) + return math_ops.logical_and( + array_ops.shape(a)[0] < 10, math_ops.reduce_max(a) < 4.0 + ) - def body(x): - return x - 1 + def body(a): + return a + 1.0 r = while_loop.while_loop(cond, body, [x], return_same_structure=True) - - self.assertAllEqual(shapes_seen, [[1, 1]] * len(shapes_seen)) - self.assertLen(shapes_seen, 3) + self.assertAllEqual(shapes_seen, [[1, 28, 28, 1]] * 4) self.assertIsInstance(r, list) - self.assertAllClose(self.evaluate(r), [[[3.0]]]) + self.assertLen(r, 1) + self.assertEqual(r[0].shape.as_list(), [1, 28, 28, 1]) + self.assertAllClose( + self.evaluate(r[0]), array_ops.ones([1, 28, 28, 1]) * 4.0 + ) + + shapes_seen.clear() + r_max = while_loop.while_loop( + cond, body, [x], return_same_structure=True, maximum_iterations=3 + ) + self.assertAllEqual(shapes_seen, [[1, 28, 28, 1]] * 4) + self.assertIsInstance(r_max, list) + self.assertLen(r_max, 1) + self.assertEqual(r_max[0].shape.as_list(), [1, 28, 28, 1]) + self.assertAllClose(self.evaluate(r_max[0]), self.evaluate(r[0])) + + @test_util.run_v2_only + def testEagerWhileLoopMultiVarBareTensorBodyRaisesValueError(self): + x = array_ops.ones([1, 28, 28, 1]) + for max_iters in (None, 2): + with self.assertRaises(ValueError): + while_loop.while_loop( + lambda a, b: array_ops.shape(a)[0] < 10, + lambda a, b: array_ops.ones([2, 28, 28, 1]), + [x, x], + return_same_structure=True, + maximum_iterations=max_iters, + ) + with self.assertRaises(ValueError): + while_loop.while_loop( + lambda d: d["a"] < 3, + lambda d: {"a": d["a"] + 1}, + [{"a": constant_op.constant(0)}], + return_same_structure=True, + maximum_iterations=max_iters, + ) @test_util.enable_control_flow_v2 @test_util.run_in_graph_and_eager_modes diff --git a/tensorflow/python/ops/math_grad.py b/tensorflow/python/ops/math_grad.py index 37fe239e5c8a15..a1200cf1f72e22 100644 --- a/tensorflow/python/ops/math_grad.py +++ b/tensorflow/python/ops/math_grad.py @@ -2042,49 +2042,32 @@ def _CumprodGrad(op: ops.Operation, grad): def _CumulativeLogsumexpGrad(op: ops.Operation, grad): x = op.inputs[0] axis = op.inputs[1] + cumulative_logsumexp = op.outputs[0] exclusive = op.get_attr("exclusive") reverse = op.get_attr("reverse") - cumulative_logsumexp = math_ops.cumulative_logsumexp( - x, axis=axis, exclusive=exclusive, reverse=reverse - ) # Split the incoming gradient into positive and negative part # in order to take logs. This is required for stable results. - finite_lse = math_ops.is_finite(cumulative_logsumexp) - pos_mask = math_ops.logical_and(math_ops.greater(grad, 0), finite_lse) - safe_pos_grad = array_ops.where_v2( - pos_mask, grad, math_ops.cast(1.0, grad.dtype) - ) log_grad_positive = array_ops.where_v2( - pos_mask, - math_ops.log(safe_pos_grad) - cumulative_logsumexp, - grad.dtype.min, - ) + math_ops.greater(grad, 0), + math_ops.log(grad), + grad.dtype.min) - neg_mask = math_ops.logical_and(math_ops.less(grad, 0), finite_lse) - safe_neg_grad = array_ops.where_v2( - neg_mask, -grad, math_ops.cast(1.0, grad.dtype) - ) log_grad_negative = array_ops.where_v2( - neg_mask, - math_ops.log(safe_neg_grad) - cumulative_logsumexp, - grad.dtype.min, - ) + math_ops.less(grad, 0), + math_ops.log(-grad), + grad.dtype.min) output_pos = math_ops.exp( math_ops.cumulative_logsumexp( - log_grad_positive, axis=axis, reverse=not reverse, exclusive=exclusive - ) - + x - ) + log_grad_positive - cumulative_logsumexp, + axis=axis, reverse=not reverse, exclusive=exclusive) + x) output_neg = math_ops.exp( math_ops.cumulative_logsumexp( - log_grad_negative, axis=axis, reverse=not reverse, exclusive=exclusive - ) - + x - ) + log_grad_negative - cumulative_logsumexp, + axis=axis, reverse=not reverse, exclusive=exclusive) + x) return [output_pos - output_neg, None] diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py index 40923e38d00ff0..261a0d1bcf4862 100644 --- a/tensorflow/python/ops/math_ops.py +++ b/tensorflow/python/ops/math_ops.py @@ -575,6 +575,10 @@ def _mul(x, y, name=None): @dispatch.register_binary_elementwise_api @dispatch.add_dispatch_support def subtract(x, y, name=None): + if not tensor_util.is_tf_type(x) and tensor_util.is_tf_type(y): + x = ops.convert_to_tensor(x, dtype=y.dtype.base_dtype) + elif tensor_util.is_tf_type(x) and not tensor_util.is_tf_type(y): + y = ops.convert_to_tensor(y, dtype=x.dtype.base_dtype) return gen_math_ops.sub(x, y, name) diff --git a/tensorflow/python/ops/math_ops_test.py b/tensorflow/python/ops/math_ops_test.py index 84b81cb1151e0f..1fbc1795a0b986 100644 --- a/tensorflow/python/ops/math_ops_test.py +++ b/tensorflow/python/ops/math_ops_test.py @@ -1136,6 +1136,19 @@ def testXlogyWithZeroBroadcast(self): self.assertAllClose(zeros_np, xlogy_tf_np[0]) self.assertAllClose(xtimes_logy, xlogy_tf_np[1]) + def testXlogyWithNegativeZero(self): + # A zero x gives +0 at every size. Tensors of at least one SIMD packet used + # to return x itself, which kept the sign of -0. This covers the Eigen CPU + # kernels; the GPU kernels are generated from MLIR separately. + for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]: + for size in [1, 3, 4, 16, 1024]: + x = constant_op.constant(np.full((size,), -0.0), dtype=dtype) + y = constant_op.constant(np.full((size,), 2.0), dtype=dtype) + with test_util.force_cpu(): + xlogy = self.evaluate(math_ops.xlogy(x, y)) + self.assertAllEqual(xlogy, np.zeros(size)) + self.assertFalse(np.signbit(xlogy).any()) + @test_util.run_all_in_graph_and_eager_modes class Xlog1pyTest(test_util.TensorFlowTestCase): @@ -1170,6 +1183,17 @@ def testXlog1pyWithZeroBroadcast(self): self.assertAllClose(zeros_np, xlog1py_tf_np[0]) self.assertAllClose(xtimes_log1py, xlog1py_tf_np[1]) + def testXlog1pyWithNegativeZero(self): + # See XlogyTest.testXlogyWithNegativeZero. + for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]: + for size in [1, 3, 4, 16, 1024]: + x = constant_op.constant(np.full((size,), -0.0), dtype=dtype) + y = constant_op.constant(np.full((size,), 1.0), dtype=dtype) + with test_util.force_cpu(): + xlog1py = self.evaluate(math_ops.xlog1py(x, y)) + self.assertAllEqual(xlog1py, np.zeros(size)) + self.assertFalse(np.signbit(xlog1py).any()) + @test_util.run_all_in_graph_and_eager_modes class XdivyTest(test_util.TensorFlowTestCase): @@ -1203,6 +1227,17 @@ def testXdivyWithZeroBroadcast(self): self.assertAllClose(zeros_np, xdivy_tf_np[0]) self.assertAllClose(x_over_y, xdivy_tf_np[1]) + def testXdivyWithNegativeZero(self): + # See XlogyTest.testXlogyWithNegativeZero. + for dtype in [dtypes.float16, dtypes.float32, dtypes.float64]: + for size in [1, 3, 4, 16, 1024]: + x = constant_op.constant(np.full((size,), -0.0), dtype=dtype) + y = constant_op.constant(np.full((size,), 3.0), dtype=dtype) + with test_util.force_cpu(): + xdivy = self.evaluate(math_ops.xdivy(x, y)) + self.assertAllEqual(xdivy, np.zeros(size)) + self.assertFalse(np.signbit(xdivy).any()) + @test_util.run_all_in_graph_and_eager_modes class NextAfterTest(test_util.TensorFlowTestCase): @@ -1563,6 +1598,26 @@ def test_fn(): @test_util.run_all_in_graph_and_eager_modes class ScalarCoercionTest(test_util.TensorFlowTestCase): + def testSubtractScalarCoercion(self): + x = constant_op.constant([1.0, 2.0, 3.0], dtype=dtypes.float32) + + def subtract_fn(t): + return math_ops.subtract(5, t) + + def subtract_reverse_fn(t): + return math_ops.subtract(t, 5) + + for fn, expected in ( + (subtract_fn, [4.0, 3.0, 2.0]), + (subtract_reverse_fn, [-4.0, -3.0, -2.0]), + ): + self.assertAllClose(fn(x), expected) + self.assertAllClose(def_function.function(fn)(x), expected) + if test_util.is_xla_enabled(): + self.assertAllClose( + def_function.function(fn, jit_compile=True)(x), expected + ) + def testMultiplyAndPowScalarCoercion(self): x1 = constant_op.constant([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32) x2 = constant_op.constant([4.0, 5.0, 6.0, 7.0], dtype=dtypes.bfloat16) diff --git a/tensorflow/python/ops/nn_grad_test.py b/tensorflow/python/ops/nn_grad_test.py index 1a2fc541c02ce1..36d6b66f3e79df 100644 --- a/tensorflow/python/ops/nn_grad_test.py +++ b/tensorflow/python/ops/nn_grad_test.py @@ -81,7 +81,8 @@ def run_test(self, x, y): x.get_shape().as_list(), y, y.get_shape().as_list()) - self.assertLess(error, 2e-3) + tol = 0.06 if test.is_gpu_available() else 2e-3 + self.assertLess(error, tol) @test_util.run_deprecated_v1 def testConv2dGradWRTInput(self): diff --git a/tensorflow/python/ops/nn_impl.py b/tensorflow/python/ops/nn_impl.py index 8a9d491a4be260..0edfc8c542704c 100644 --- a/tensorflow/python/ops/nn_impl.py +++ b/tensorflow/python/ops/nn_impl.py @@ -27,6 +27,7 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import array_ops_stack from tensorflow.python.ops import candidate_sampling_ops +from tensorflow.python.ops import check_ops from tensorflow.python.ops import cond as tf_cond from tensorflow.python.ops import ctc_ops # pylint: disable=unused-import from tensorflow.python.ops import custom_gradient @@ -682,6 +683,131 @@ def zero_fraction(value, name=None): return array_ops.identity(zero_fraction_float32, "fraction") +def _check_dilated_depthwise_shapes(input, filter, rate, padding, data_format): + """Raises a descriptive error if a dilated depthwise conv has no window. + + `with_space_to_batch` evaluates the dilated convolution densely over + space-to-batch blocks. When the input (plus any explicit padding) is smaller + than the dilated filter, no valid window exists: `tf.function` then fails + with an opaque negative-dimension error while eager silently returns a tensor + of the wrong shape. Validating the shapes up front gives both execution modes + the same descriptive error. + + Shapes that are not statically known are left for the op to report at run + time. "SAME" padding always admits an output and is not checked. + + Args: + input: The `input` argument of `depthwise_conv2d`. + filter: The `filter` argument of `depthwise_conv2d`. + rate: The dilation rate, a list of two positive ints. + padding: The `padding` argument of `depthwise_conv2d`. + data_format: The `data_format` argument of `depthwise_conv2d`. + + Raises: + ValueError: If a dilated filter does not fit in the padded input. + """ + if all(r == 1 for r in rate): + return + + channels_first = (data_format or "").startswith("NC") + if isinstance(padding, str): + if padding == "SAME": + # `with_space_to_batch` pads enough for an output to always exist. + return + spatial_pads = [(0, 0), (0, 0)] + else: + # Drop the batch and channel dimensions, keep the two spatial ones. + spatial_pads = [list(p) for p in padding] + spatial_pads = spatial_pads[2:] if channels_first else spatial_pads[1:-1] + + if input.shape.rank != 4 or filter.shape.rank != 4: + return + in_shape = input.shape.as_list() + f_shape = filter.shape.as_list() + if None in in_shape or None in f_shape: + return + + in_spatial = in_shape[2:4] if channels_first else in_shape[1:3] + # The filter layout is `[height, width, in_channels, channel_multiplier]` + # regardless of `data_format`. + f_spatial = f_shape[0:2] + + for i, dim_name in enumerate(("height", "width")): + available = in_spatial[i] + spatial_pads[i][0] + spatial_pads[i][1] + dilated = (f_spatial[i] - 1) * rate[i] + 1 + if available < dilated: + raise ValueError( + f"Input {dim_name} ({in_spatial[i]}) plus padding " + f"({spatial_pads[i][0] + spatial_pads[i][1]}) is too small for " + f"dilation rate {rate[i]} with filter size {f_spatial[i]}: the " + f"dilated filter spans {dilated} values but only {available} are " + "available, so the convolution has no valid window." + ) + + +def _assert_dilated_depthwise_fits(input, filter, rate, padding, data_format): + """Guards a dilated depthwise conv whose spatial shapes are dynamic. + + `_check_dilated_depthwise_shapes` cannot see shapes that are unknown at + tracing time (a `tf.function` accepting a `TensorSpec` with `None` + dimensions). Without a run-time check, an input too small for the dilated + filter makes the underlying op silently return a wrongly shaped result + instead of failing, which is the `tf.function` half of #113320. + + Args: + input: The `input` argument of `depthwise_conv2d`. + filter: The `filter` argument of `depthwise_conv2d`. + rate: The dilation rate, a list of two positive ints. + padding: The `padding` argument of `depthwise_conv2d`. + data_format: The `data_format` argument of `depthwise_conv2d`. + + Returns: + `input`, with a control dependency on the shape assertion. + """ + # A rank other than 4 is invalid for the op itself; leave it to the op to + # report, rather than slicing `array_ops.shape` below on a rank we cannot + # interpret. + if (input.shape.rank is not None and input.shape.rank != 4) or ( + filter.shape.rank is not None and filter.shape.rank != 4 + ): + return input + + channels_first = (data_format or "").startswith("NC") + if isinstance(padding, str): + pads = [0, 0] + else: + # Drop the batch and channel dimensions, keep the two spatial ones. + spatial_pads = [list(p) for p in padding] + spatial_pads = spatial_pads[2:] if channels_first else spatial_pads[1:-1] + pads = [ + spatial_pads[0][0] + spatial_pads[0][1], + spatial_pads[1][0] + spatial_pads[1][1], + ] + + in_shape = array_ops.shape(input) + f_shape = array_ops.shape(filter) + in_spatial = in_shape[2:4] if channels_first else in_shape[1:3] + # The filter layout is `[height, width, in_channels, channel_multiplier]` + # regardless of `data_format`. + f_spatial = f_shape[0:2] + + pads_tensor = ops.convert_to_tensor(pads, dtype=in_spatial.dtype) + available = in_spatial + pads_tensor + dilated = (f_spatial - 1) * ops.convert_to_tensor( + rate, dtype=f_spatial.dtype + ) + 1 + check = check_ops.assert_greater_equal( + available, + dilated, + message=( + "Dilated depthwise convolution: the input (plus padding) is " + "smaller than the dilated filter, so there is no valid window." + ), + ) + with ops.control_dependencies([check]): + return array_ops.identity(input) + + # pylint: disable=redefined-builtin @tf_export(v1=["nn.depthwise_conv2d"]) @dispatch.add_dispatch_support @@ -787,6 +913,21 @@ def depthwise_conv2d(input, f"`strides` {strides}" ) + _check_dilated_depthwise_shapes(input, filter, rate, padding, data_format) + + # The check above only sees static shapes. When tracing a `tf.function` + # whose spatial dimensions are unknown, fall back to a run-time assertion + # so that an oversized dilated filter fails instead of silently producing a + # wrongly shaped result. "SAME" padding always admits an output. + if any(_ > 1 for _ in rate) and not ( + isinstance(padding, str) and padding == "SAME" + ): + in_shape = input.shape.as_list() if input.shape.rank == 4 else None + if in_shape is None or None in in_shape: + input = _assert_dilated_depthwise_fits( + input, filter, rate, padding, data_format + ) + # Use depthwise_conv2d_native if executing on TPU. if device_context.enclosing_tpu_context() is not None: if data_format == "NCHW": @@ -850,7 +991,8 @@ def depthwise_conv2d_v2(input, same horizontal and vertical strides, `strides = [1, stride, stride, 1]`. If any value in `dilations` is greater than 1, we perform atrous depthwise convolution, in which case all values in the `strides` tensor must be equal - to 1. + to 1; a `ValueError` is raised otherwise. A dilated depthwise convolution is + evaluated with space-to-batch, which cannot express a stride. Usage Example: diff --git a/tensorflow/python/ops/nn_ops.py b/tensorflow/python/ops/nn_ops.py index 149668d9766a3e..58c30af22bec4b 100644 --- a/tensorflow/python/ops/nn_ops.py +++ b/tensorflow/python/ops/nn_ops.py @@ -1195,6 +1195,160 @@ def convolution_v2( # pylint: disable=missing-docstring "filter", "filters") +def _check_dilated_convolution_shapes( + input_tensor, + filters, + dilations, + padding, + num_batch_dims, + channels_first, + num_spatial_dims, +): + """Raises a descriptive error if a dilated convolution has no window. + + `tf.nn.convolution` evaluates dilated convolutions with + `with_space_to_batch`, which subsamples the input into blocks and convolves + each block densely. When the dilated filter is larger than the input, no + valid window exists. `tf.function` then fails with an opaque + negative-dimension error, while eager mode silently returns a tensor of the + wrong shape, because the inner convolution returns a tensor shaped like its + input instead of the empty output (#113319). Validating the shapes up front + gives both execution modes the same descriptive error. + + Shapes that are not statically known are left for the run-time assertion in + `_assert_dilated_convolution_fits`. "SAME" padding always admits an output + and is not checked. A malformed `input`, `filters` or `padding` (unknown or + too small a rank, or a padding list that is too short) is left for the core + convolution code, which already reports it with a better message than an + `IndexError` from here would. + + Args: + input_tensor: The `input` argument of `convolution_internal`. + filters: The `filters` argument of `convolution_internal`. + dilations: The dilation rate, a list of `num_spatial_dims` positive ints. + padding: The `padding` argument of `convolution_internal`. + num_batch_dims: The number of leading batch dimensions of `input`. + channels_first: Whether the channel dimension precedes the spatial ones. + num_spatial_dims: The number of spatial dimensions, 1, 2 or 3. + + Raises: + ValueError: If the dilated filter does not fit in the padded input. + """ + if all(d == 1 for d in dilations): + return + + spatial_start = num_batch_dims + 1 if channels_first else num_batch_dims + if isinstance(padding, str): + if padding == "SAME": + # `with_space_to_batch` pads enough for an output to always exist. + return + if padding != "VALID": + return + spatial_pads = [(0, 0)] * num_spatial_dims + else: + # Drop the batch and channel dimensions, keep the spatial ones. + spatial_pads = [ + list(p) + for p in padding[spatial_start : spatial_start + num_spatial_dims] + ] + if len(spatial_pads) < num_spatial_dims: + return + + if input_tensor.shape.rank is None or filters.shape.rank is None: + return + if ( + input_tensor.shape.rank < num_spatial_dims + num_batch_dims + 1 + or filters.shape.rank < num_spatial_dims + ): + # Let the core convolution code report the rank error it already raises. + return + in_shape = input_tensor.shape.as_list() + f_shape = filters.shape.as_list() + if None in in_shape or None in f_shape: + return + + in_spatial = in_shape[spatial_start : spatial_start + num_spatial_dims] + # The spatial dimensions of `filters` come first regardless of the data + # format. + f_spatial = f_shape[:num_spatial_dims] + + for i in range(num_spatial_dims): + available = in_spatial[i] + spatial_pads[i][0] + spatial_pads[i][1] + dilated = (f_spatial[i] - 1) * dilations[i] + 1 + if available < dilated: + raise ValueError( + f"Input spatial dimension {i} ({in_spatial[i]}) plus padding " + f"({spatial_pads[i][0] + spatial_pads[i][1]}) is too small for " + f"dilation rate {dilations[i]} with filter size {f_spatial[i]}: the " + f"dilated filter spans {dilated} values but only {available} are " + "available, so the convolution has no valid window." + ) + + +def _assert_dilated_convolution_fits( + input_tensor, + filters, + dilations, + padding, + num_batch_dims, + channels_first, + num_spatial_dims, +): + """Guards a dilated convolution whose spatial shapes are dynamic. + + `_check_dilated_convolution_shapes` cannot see shapes that are unknown at + tracing time (a `tf.function` accepting a `TensorSpec` with `None` + dimensions). Without a run-time check, an input too small for the dilated + filter makes the convolution silently return a wrongly shaped result instead + of failing, which is the eager half of #113319. + + Args: + input_tensor: The `input` argument of `convolution_internal`. + filters: The `filters` argument of `convolution_internal`. + dilations: The dilation rate, a list of `num_spatial_dims` positive ints. + padding: The `padding` argument of `convolution_internal`. + num_batch_dims: The number of leading batch dimensions of `input`. + channels_first: Whether the channel dimension precedes the spatial ones. + num_spatial_dims: The number of spatial dimensions, 1, 2 or 3. + + Returns: + `input`, with a control dependency on the shape assertion. + """ + spatial_start = num_batch_dims + 1 if channels_first else num_batch_dims + if isinstance(padding, str): + pads = [0] * num_spatial_dims + else: + spatial_pads = [ + list(p) + for p in padding[spatial_start : spatial_start + num_spatial_dims] + ] + if len(spatial_pads) < num_spatial_dims: + # The core convolution code reports the malformed padding already. + return input + pads = [p[0] + p[1] for p in spatial_pads] + + in_shape = array_ops.shape(input_tensor) + f_shape = array_ops.shape(filters) + in_spatial = in_shape[spatial_start : spatial_start + num_spatial_dims] + f_spatial = f_shape[:num_spatial_dims] + + pads_tensor = ops.convert_to_tensor(pads, dtype=in_spatial.dtype) + available = in_spatial + pads_tensor + dilated = (f_spatial - 1) * ops.convert_to_tensor( + dilations, dtype=f_spatial.dtype + ) + 1 + check = check_ops.assert_greater_equal( + available, + dilated, + message=( + "Dilated convolution: the input (plus padding) is smaller than " + "the dilated filter, so there is no valid window." + ), + ) + with ops.control_dependencies([check]): + return array_ops.identity(input_tensor) + + def convolution_internal( input, # pylint: disable=redefined-builtin filters, @@ -1334,6 +1488,37 @@ def convolution_internal( strides = strides[1:-1] dilations = dilations[1:-1] + channels_first = channel_index == num_batch_dims + _check_dilated_convolution_shapes( + input, + filters, + dilations, + padding, + num_batch_dims, + channels_first, + num_spatial_dims, + ) + + # The check above only sees static shapes. When tracing a `tf.function` + # whose spatial dimensions are unknown, fall back to a run-time + # assertion so that an oversized dilated filter fails instead of + # silently producing a wrongly shaped result. "SAME" padding always + # admits an output. + if not (isinstance(padding, str) and padding == "SAME"): + in_shape = ( + input.shape.as_list() if input.shape.rank is not None else None + ) + if in_shape is None or None in in_shape: + input = _assert_dilated_convolution_fits( + input, + filters, + dilations, + padding, + num_batch_dims, + channels_first, + num_spatial_dims, + ) + op = Convolution( tensor_shape.as_shape(input.shape), tensor_shape.as_shape(filters.shape), diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py index 0bca1cbca5b211..9b0490eb60b35c 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py @@ -536,12 +536,20 @@ def f(x1, x2): float_dtype = np_utils.result_type(float) x1 = math_ops.cast(x1, float_dtype) x2 = math_ops.cast(x2, float_dtype) - amax = maximum(x1, x2) delta = x1 - x2 + # `maximum` sends the whole gradient to `x1` where `x1 == x2`, so the two + # partial derivatives (analytically 0.5 each) come out as 1.0 and 0.0. + # Selecting the operands through `where` instead keeps the forward values + # bit-for-bit identical while giving each argument its own differentiable + # expression, so both gradients are 0.5 at `x1 == x2`. Selecting them once + # also avoids evaluating `exp` and `log1p` for both branches. The exponent + # is always <= 0, hence this cannot overflow. + max_val = np_array_ops.where(x1 > x2, x1, x2) + min_val = np_array_ops.where(x1 > x2, x2, x1) return np_array_ops.where( isnan(delta), x1 + x2, # NaNs or infinities of the same sign. - amax + log1p(exp(-abs(delta))), + max_val + log1p(exp(min_val - max_val)), ) return _bin_op(f, x1, x2) diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py index 59f07e0f2db1ab..84869bc94586e8 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py @@ -215,6 +215,19 @@ def testHypotSpecialCases(self): expected = np.hypot(x, y) np.testing.assert_equal(actual.tolist(), expected.tolist()) + def testLogaddexpEqualInputsGradient(self): + # Both partial derivatives of logaddexp(x, x) are 0.5. `maximum` used to + # send the whole gradient to x1 there, giving 1.0 and 0.0. + for dtype in (dtypes.float32, dtypes.float64): + x1 = constant_op.constant([0.0, 1.0, -2.0, 100.0], dtype=dtype) + x2 = constant_op.constant([0.0, 1.0, -2.0, 100.0], dtype=dtype) + with backprop.GradientTape(persistent=True) as tape: + tape.watch(x1) + tape.watch(x2) + y = math_ops.reduce_sum(np_math_ops.logaddexp(x1, x2)) + for grad in tape.gradient(y, (x1, x2)): + self.assertAllClose(grad, [0.5, 0.5, 0.5, 0.5]) + def testLogaddexp(self): self._testBinaryOp(np_math_ops.logaddexp, np.logaddexp, 'logaddexp') self._testBinaryOp(np_math_ops.logaddexp2, np.logaddexp2, 'logaddexp2') diff --git a/tensorflow/python/ops/parallel_for/math_test.py b/tensorflow/python/ops/parallel_for/math_test.py index 2dc84ba913c8ee..133c374c6e7db0 100644 --- a/tensorflow/python/ops/parallel_for/math_test.py +++ b/tensorflow/python/ops/parallel_for/math_test.py @@ -329,7 +329,9 @@ def loop_fn(i): # pylint: enable=cell-var-from-loop - self._test_loop_fn(loop_fn, 2) + rtol = 1e-3 if test.is_gpu_available() else 1e-4 + atol = 1e-3 if test.is_gpu_available() else 1e-5 + self._test_loop_fn(loop_fn, 2, rtol=rtol, atol=atol) def test_batch_matmul_broadcast(self): for broadcast_a in (True, False): @@ -350,7 +352,9 @@ def loop_fn(i): return math_ops.matmul(a, b) # pylint: enable=cell-var-from-loop - self._test_loop_fn(loop_fn, 2) + rtol = 1e-3 if test.is_gpu_available() else 1e-4 + atol = 1e-3 if test.is_gpu_available() else 1e-5 + self._test_loop_fn(loop_fn, 2, rtol=rtol, atol=atol) def test_reduction(self): x = random_ops.random_uniform([2, 3, 4, 5]) diff --git a/tensorflow/python/ops/ragged/ragged_gather_op_test.py b/tensorflow/python/ops/ragged/ragged_gather_op_test.py index d20cc3744060e1..06bdf06867ff15 100644 --- a/tensorflow/python/ops/ragged/ragged_gather_op_test.py +++ b/tensorflow/python/ops/ragged/ragged_gather_op_test.py @@ -242,6 +242,25 @@ def testValueCountOverflowError(self): ) ) + def testSplitsOverflowErrorInt64(self): + # Zero-width values let int64 splits reach 2**62 without allocating memory; + # gathering the row twice would overflow the int64 output splits. + with self.assertRaisesRegex( + errors.InvalidArgumentError, 'exceeds limits of Tsplits' + ): + self.evaluate( + gen_ragged_array_ops.ragged_gather( + params_nested_splits=[ + constant_op.constant([0, 2**62], dtype=dtypes.int64) + ], + params_dense_values=array_ops.zeros( + [2**62, 0], dtype=dtypes.uint8 + ), + indices=array_ops.zeros([2], dtype=dtypes.int32), + OUTPUT_RAGGED_RANK=1, + ) + ) + def testUnknownIndicesRankError(self): if context.executing_eagerly(): return diff --git a/tensorflow/python/ops/while_loop.py b/tensorflow/python/ops/while_loop.py index 690ebd3974c9b0..2c3cf2beb1ef61 100644 --- a/tensorflow/python/ops/while_loop.py +++ b/tensorflow/python/ops/while_loop.py @@ -452,6 +452,25 @@ def while_loop(cond, if not loop_vars: raise ValueError("'loop_vars' must be provided.") try_to_pack = (len(loop_vars) == 1 and not return_same_structure) + if ( + executing_eagerly + and isinstance(loop_vars, (list, tuple)) + and len(loop_vars) == 1 + and not try_to_pack + ): + orig_loop_vars_type = ( + type(loop_vars) if type(loop_vars) in (list, tuple) else list + ) + unpacked_body = body + + def _pack_body(*args): + res = unpacked_body(*args) + if not nest.is_nested(res): + return orig_loop_vars_type((res,)) + return res + + body = _pack_body + if maximum_iterations is not None: maximum_iterations = ops.convert_to_tensor( maximum_iterations, name="maximum_iterations") @@ -482,19 +501,17 @@ def while_loop(cond, if executing_eagerly: packed = False # whether the body result was packed into a 1-item tuple - orig_loop_vars_type = ( - type(loop_vars) if type(loop_vars) in (list, tuple) else list + loop_var_structure = nest.map_structure( + type_spec.type_spec_from_value, loop_vars ) - - loop_var_structure = nest.map_structure(type_spec.type_spec_from_value, - list(loop_vars)) while cond(*loop_vars): loop_vars = body(*loop_vars) - if not isinstance(loop_vars, (list, tuple)): - if try_to_pack: - packed = True - loop_vars = orig_loop_vars_type((loop_vars,)) - nest.assert_same_structure(loop_var_structure, list(loop_vars)) + if try_to_pack and not isinstance(loop_vars, (list, tuple)): + packed = True + loop_vars = (loop_vars,) + nest.assert_same_structure( + loop_var_structure, loop_vars, check_types=False + ) def convert(x): if isinstance(x, tensor_array_ops.TensorArray): diff --git a/tensorflow/python/tools/optimize_for_inference_test.py b/tensorflow/python/tools/optimize_for_inference_test.py index cd2f71fb3e598f..99d96cf40ebec7 100644 --- a/tensorflow/python/tools/optimize_for_inference_test.py +++ b/tensorflow/python/tools/optimize_for_inference_test.py @@ -400,16 +400,19 @@ def testFoldFusedBatchNormsV3(self): original_result = sess.run(["output:0"]) optimized_graph_def = optimize_for_inference_lib.fold_batch_norms( original_graph_def) - with self.cached_session() as sess: - _ = importer.import_graph_def( - optimized_graph_def, input_map={}, name="optimized") - optimized_result = sess.run(["optimized/output:0"]) - self.assertAllClose( - original_result, optimized_result, rtol=1e-04, atol=1e-06) + with self.cached_session() as sess: + _ = importer.import_graph_def( + optimized_graph_def, input_map={}, name="optimized" + ) + optimized_result = sess.run(["optimized/output:0"]) - for node in optimized_graph_def.node: - self.assertNotEqual("FusedBatchNormV3", node.op) + self.assertAllClose( + original_result, optimized_result, rtol=1e-04, atol=1e-06 + ) + + for node in optimized_graph_def.node: + self.assertNotEqual("FusedBatchNormV3", node.op) @test_util.run_deprecated_v1 def testFuseResizePadAndConv(self): diff --git a/tensorflow/workspace2.bzl b/tensorflow/workspace2.bzl index 8e878d43070531..eacb39dbf8a972 100644 --- a/tensorflow/workspace2.bzl +++ b/tensorflow/workspace2.bzl @@ -423,10 +423,10 @@ def _tf_repositories(): tf_http_archive( name = "curl", build_file = "@xla//third_party:curl.BUILD", - sha256 = "d9b327997999045a24cda50f3983e69e51c516bd8be6ef9842fc7f99135e33bb", - strip_prefix = "curl-8.21.0", + sha256 = "d54dd598bf05927a726deb38df31c6a255ba83ff1de57c5d1464dac3ed8f44a1", + strip_prefix = "curl-8.22.0", system_build_file = "//third_party/systemlibs:curl.BUILD", - urls = tf_mirror_urls("https://curl.se/download/curl-8.21.0.tar.gz"), + urls = tf_mirror_urls("https://curl.se/download/curl-8.22.0.tar.gz"), ) tf_http_archive( diff --git a/tensorflow/workspace3.bzl b/tensorflow/workspace3.bzl index bf557b54aa9769..b231388b486f4b 100644 --- a/tensorflow/workspace3.bzl +++ b/tensorflow/workspace3.bzl @@ -80,10 +80,10 @@ def workspace(): # Details: https://github.com/google-ml-infra/rules_ml_toolchain tf_http_archive( name = "rules_ml_toolchain", - sha256 = "3b05687842427041c65d1bc2f4aeda3a3079557120f5be8c34690087a88c5de5", - strip_prefix = "rules_ml_toolchain-2eddbc595cc0bbe650c2640204f66b14f015f1a8", + sha256 = "677f0785f14bf75339a3187cba7fa95e1a1f888a18568edc2e5fac87d513fc7c", + strip_prefix = "rules_ml_toolchain-f56a19573fec3543050a5f4bcd9a07a1de08e8ee", urls = tf_mirror_urls( - "https://github.com/google-ml-infra/rules_ml_toolchain/archive/2eddbc595cc0bbe650c2640204f66b14f015f1a8.tar.gz", + "https://github.com/google-ml-infra/rules_ml_toolchain/archive/f56a19573fec3543050a5f4bcd9a07a1de08e8ee.tar.gz", ), ) diff --git a/third_party/xla/.github/workflows/ci.yml b/third_party/xla/.github/workflows/ci.yml index 39bc138adeb4a3..b31aebfa2f8a23 100644 --- a/third_party/xla/.github/workflows/ci.yml +++ b/third_party/xla/.github/workflows/ci.yml @@ -114,6 +114,7 @@ jobs: uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: path: "openxla/xla" + fetch-depth: 0 persist-credentials: false - name: Checking out ${{ matrix.job_info.repo }} if: ${{ matrix.job_info.repo != 'openxla/xla' }} @@ -135,6 +136,11 @@ jobs: - name: "Run build.py" working-directory: ${{ matrix.job_info.repo }} timeout-minutes: 60 + env: + PYTHONPATH: ${{ github.workspace }}/openxla/xla + XLA_CI_BAZEL_DIFF_BASE_SHA: ${{ github.event.pull_request.base.sha }} + XLA_CI_BAZEL_DIFF_HEAD_SHA: ${{ github.event.pull_request.head.sha || github.sha }} + XLA_CI_DISABLE_BAZEL_DIFF: ${{ contains(github.event.pull_request.labels.*.name, 'ci:disable-bazel-diff') }} run: | # Shellcheck doesn't understand GitHub Actions expressions and incorrectly # thinks this comparison will always fail. @@ -142,7 +148,7 @@ jobs: if [[ "${{ matrix.job_info.pool }}" == *"windows"* ]]; then python "$GITHUB_WORKSPACE"\\openxla\\xla\\build_tools\\ci\\build.py --build="${{ matrix.job_info.name }}_github_actions" else - "$GITHUB_WORKSPACE"/openxla/xla/build_tools/ci/build.py --build="${{ matrix.job_info.name }}_github_actions" + python3 "$GITHUB_WORKSPACE"/openxla/xla/build_tools/ci/build.py --build="${{ matrix.job_info.name }}_github_actions" fi diff --git a/third_party/xla/.github/workflows/clang_tidy.yml b/third_party/xla/.github/workflows/clang_tidy.yml index 99dae1139c8fd6..1cfa192e5a7f6b 100644 --- a/third_party/xla/.github/workflows/clang_tidy.yml +++ b/third_party/xla/.github/workflows/clang_tidy.yml @@ -27,9 +27,6 @@ jobs: shell: bash timeout-minutes: 30 continue-on-error: true - if: | - github.event.sender.type == 'User' || - contains(github.event.pull_request.body, 'FORCE_TEST_ACTIONS') steps: - name: "Checking out repository" uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 diff --git a/third_party/xla/MODULE.bazel b/third_party/xla/MODULE.bazel index 6a7be06dcfc0ce..d858789006e9fb 100644 --- a/third_party/xla/MODULE.bazel +++ b/third_party/xla/MODULE.bazel @@ -6,10 +6,11 @@ module(name = "xla") # go/keep-sorted start bazel_dep(name = "abseil-cpp", version = "20260526.0", repo_name = "com_google_absl") bazel_dep(name = "abseil-py", version = "2.1.0", repo_name = "absl_py") +bazel_dep(name = "bazel_clang_tidy", dev_dependency = True) bazel_dep(name = "bazel_features", version = "1.43.0") bazel_dep(name = "bazel_skylib", version = "1.9.0") bazel_dep(name = "boringssl", version = "0.20250818.0") -bazel_dep(name = "curl", version = "8.12.0.bcr.1") +bazel_dep(name = "curl", version = "8.22.0") bazel_dep(name = "google_benchmark", version = "1.9.4", repo_name = "com_google_benchmark") bazel_dep(name = "googletest", version = "1.17.0", repo_name = "com_google_googletest") bazel_dep(name = "grpc", version = "1.81.0") @@ -97,6 +98,15 @@ archive_override( urls = ["https://github.com/google/googletest/archive/28e9d1f26771c6517c3b4be10254887673c94018.zip"], ) +archive_override( + module_name = "bazel_clang_tidy", + integrity = "sha256-ltpuk1zMkQRc+SjbxX8iUIonKcUff7P1YXgBew3rmzw=", + strip_prefix = "bazel_clang_tidy-c4d35e0d0b838309358e57a2efed831780f85cd0", + urls = [ + "https://github.com/erenon/bazel_clang_tidy/archive/c4d35e0d0b838309358e57a2efed831780f85cd0.tar.gz", + ], +) + ############################################################## # C++ dependencies diff --git a/third_party/xla/build_tools/ci/BUILD b/third_party/xla/build_tools/ci/BUILD index 3bd92b827859fb..2cb0a7e2b7e89a 100644 --- a/third_party/xla/build_tools/ci/BUILD +++ b/third_party/xla/build_tools/ci/BUILD @@ -33,9 +33,27 @@ package( licenses = ["notice"], ) +pytype_strict_library( + name = "bazel_diff", + srcs = ["bazel_diff.py"], +) + +py_test( + name = "bazel_diff_test", + srcs = ["bazel_diff_test.py"], + deps = [ + ":bazel_diff", + "@absl_py//absl/testing:absltest", + "@absl_py//absl/testing:parameterized", + ], +) + pytype_strict_binary( name = "build", srcs = ["build.py"], + deps = [ + ":bazel_diff", + ], ) pytype_strict_library( @@ -72,7 +90,7 @@ py_test( genrule( name = "generated_build_commands", outs = ["generated_commands.txt"], - cmd = "$(location //build_tools/ci:build) --dump_commands > $(OUTS)", + cmd = "$(execpath //build_tools/ci:build) --dump_commands > $(OUTS)", tags = ["not_run:arm"], tools = [":build"], ) diff --git a/third_party/xla/build_tools/ci/bazel_diff.py b/third_party/xla/build_tools/ci/bazel_diff.py new file mode 100644 index 00000000000000..ad82dd0567bebe --- /dev/null +++ b/third_party/xla/build_tools/ci/bazel_diff.py @@ -0,0 +1,747 @@ +# Copyright 2026 The OpenXLA Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""bazel-diff integration module for XLA CI. + +Analyzes changed files between a base SHA and head commit, determines impacted +Bazel targets using bazel-diff (with cquery by default), filters and normalizes +labels for XLA, and produces a decision (FULL, SKIP, or IMPACTED targets). +""" + +from collections.abc import Iterable, Mapping, Sequence +import dataclasses +import enum +import hashlib +import logging +import os +import re +import subprocess +import sys +import tempfile +import time +from typing import Any, List, Optional, Protocol, Tuple + +BAZEL_DIFF_VERSION = "v46.1.0" +BAZEL_DIFF_JAR_URL = "https://github.com/Tinder/bazel-diff/releases/download/v46.1.0/bazel-diff_deploy.jar" +BAZEL_DIFF_JAR_SHA256 = ( + "1026c66304d262e26065fc0ccd0281766c4c8b63abd18ea1fc4366fe05453c1e" +) + +# Global configuration files whose modification invalidates incremental diffs +# and necessitates a full test suite run. +GLOBAL_BAZEL_CONFIG_PATTERNS: Tuple[str, ...] = ( + r"(^|/)MODULE\.bazel$", + r"(^|/)REPO\.bazel$", + r"(^|/)WORKSPACE(\.bzlmod)?$", + r"\.bazelrc$", + r"\.bazelversion$", +) + +# b/519689071: keywords in BUILD/.bzl changes that signal unconfigured hashing +# under-selection (label_flag, config_setting, alias). +FULL_RUN_BUILD_KEYWORDS: Tuple[str, ...] = ( + "label_flag(", + "config_setting(", + "alias(", +) + +# Regex matching files that only contain documentation or repository metadata. +DOCS_OR_METADATA_PATTERN = re.compile( + r"(\.md$|^docs/|(^|/)OWNERS$|^LICENSE|^\.clang|^\.gitignore|^\.vscode/|\.png$|\.jpg$|\.jpeg$|\.svg$|\.webp$|\.gif$)" +) + +# Target prefixes to keep in XLA's CI. +KEEP_TARGET_PREFIXES: Tuple[str, ...] = ( + "//xla/", + "//xla:", + "//build_tools/", + "@tsl//tsl/", +) + + +class BuildConfig(Protocol): + """Protocol describing the Build attributes used by bazel-diff. + + Attributes: + target_patterns: Sequence of Bazel target patterns to build or test. + configs: Sequence of Bazel --config names. + repo_env: Mapping of --repo_env variables. + options: Mapping of additional Bazel flags. + startup_options: Mapping of Bazel startup options. + bazel_diff_use_cquery: Whether bazel-diff should use cquery (--useCquery). + """ + + target_patterns: Sequence[str] + configs: Sequence[str] + repo_env: Mapping[str, Any] + options: Mapping[str, Any] + startup_options: Mapping[str, Any] + bazel_diff_use_cquery: bool + + +class BazelDiffDecisionType(enum.Enum): + FULL = "FULL" + SKIP = "SKIP" + IMPACTED = "IMPACTED" + + +@dataclasses.dataclass(frozen=True) +class BazelDiffDecision: + """Decision result from bazel-diff analysis. + + Attributes: + decision: The decision type (FULL, SKIP, or IMPACTED). + impacted_targets: Tuple of normalized impacted Bazel target labels. + impacted_targets_file: Path to the file containing target patterns for + Bazel's --target_pattern_file flag, if decision is IMPACTED. + reason: Human-readable explanation for the decision. + changed_files_count: Number of git changed files between base and head. + elapsed_seconds: Wall-clock duration of the bazel-diff analysis in seconds. + """ + + decision: BazelDiffDecisionType + impacted_targets: Tuple[str, ...] = () + impacted_targets_file: Optional[str] = None + reason: str = "" + changed_files_count: int = 0 + elapsed_seconds: float = 0.0 + + +def matches_any_pattern(path: str, patterns: Sequence[str]) -> bool: + """Returns True if path matches any regex pattern in patterns.""" + return any(re.search(pattern, path) for pattern in patterns) + + +def is_global_config_changed(changed_files: Sequence[str]) -> bool: + """Returns True if any changed file matches global bazel config patterns.""" + return any( + matches_any_pattern(file_path, GLOBAL_BAZEL_CONFIG_PATTERNS) + for file_path in changed_files + ) + + +def is_docs_or_metadata_only(changed_files: Sequence[str]) -> bool: + """Returns True if all changed files are docs or non-build metadata.""" + if not changed_files: + return False + return all( + DOCS_OR_METADATA_PATTERN.search(file_path) for file_path in changed_files + ) + + +def _run_git_best_effort( + args: Sequence[str], cwd: str = "." +) -> Optional[subprocess.CompletedProcess[str]]: + """Runs a git command best-effort, returning None on OS/subprocess error.""" + try: + return subprocess.run( + ["git", *args], + cwd=cwd, + capture_output=True, + text=True, + check=False, + ) + except (OSError, subprocess.SubprocessError): + return None + + +def get_merge_base( + base_sha: str, head_sha: str = "HEAD", cwd: str = "." +) -> Optional[str]: + """Computes git merge-base between base_sha and head_sha. + + Args: + base_sha: Base git commit SHA. + head_sha: Head git commit SHA or ref. + cwd: Directory where the git command should be executed. + + Returns: + The merge base commit SHA as a string, or None if not found. + """ + result = _run_git_best_effort(["merge-base", base_sha, head_sha], cwd=cwd) + if result and result.returncode == 0 and result.stdout.strip(): + return result.stdout.strip() + return None + + +def get_diff_base(base_sha: str, head_sha: str = "HEAD", cwd: str = ".") -> str: + """Returns the base commit to diff against head_sha. + + Uses git merge-base when available so that changes that landed on the base + branch after the feature branch diverged are not falsely attributed to the PR. + + Args: + base_sha: Base git commit SHA. + head_sha: Head git commit SHA or ref. + cwd: Directory where the git command should be executed. + + Returns: + The merge-base SHA if available, otherwise base_sha. + """ + merge_base = get_merge_base(base_sha, head_sha, cwd=cwd) + return merge_base if merge_base else base_sha + + +def get_changed_files( + base_sha: str, head_sha: str = "HEAD", cwd: str = "." +) -> List[str]: + """Returns list of changed filepaths between base_sha and head_sha. + + Args: + base_sha: Base git commit SHA. + head_sha: Head git commit SHA or ref. + cwd: Directory where the git command should be executed. + + Returns: + List of relative paths of changed files. + """ + diff_base = get_diff_base(base_sha, head_sha, cwd=cwd) + command = ["git", "diff", "--name-only", diff_base, head_sha] + result = subprocess.run( + command, cwd=cwd, capture_output=True, text=True, check=True + ) + return [line.strip() for line in result.stdout.splitlines() if line.strip()] + + +def touches_build_keywords( + changed_files: Sequence[str], + base_sha: str, + head_sha: str, + cwd: str = ".", +) -> bool: + """Returns True if any BUILD or .bzl diff contains sensitive keywords.""" + build_or_bzl_files = [ + file_path + for file_path in changed_files + if file_path.endswith((".bzl", "BUILD", "BUILD.bazel")) + ] + if not build_or_bzl_files: + return False + + try: + diff_base = get_diff_base(base_sha, head_sha, cwd=cwd) + command = [ + "git", + "diff", + "-U0", + diff_base, + head_sha, + "--", + *build_or_bzl_files, + ] + result = subprocess.run( + command, cwd=cwd, capture_output=True, text=True, check=True + ) + diff_text = result.stdout + for keyword in FULL_RUN_BUILD_KEYWORDS: + if keyword in diff_text: + logging.info( + "Found build keyword '%s' in diff; forcing full run", keyword + ) + return True + except (OSError, subprocess.SubprocessError) as e: + logging.warning("Failed to check build keywords: %s", e) + return True + return False + + +def normalize_target_label(label: str) -> str: + """Normalizes Bazel target labels to standard repo forms.""" + label = label.strip() + if not label: + return "" + + # Root workspace bzlmod syntax: @@//... or @//... -> //... + if label.startswith("@@//"): + return label[2:] + if label.startswith("@//"): + return label[1:] + + # External repo canonical labels: @@+tsl+tsl//tsl/... or @@tsl//tsl/... + # -> @tsl//tsl/... + if label.startswith("@@") and "//" in label: + repo_part, target_part = label[2:].split("//", 1) + if "tsl" in repo_part: + return f"@tsl//{target_part}" + return f"@{repo_part}//{target_part}" + + return label + + +def target_matches_pattern(target: str, pattern: str) -> bool: + """Checks if target label matches a Bazel target pattern.""" + if pattern.endswith("/..."): + base = pattern[:-4] + return ( + target == base + or target.startswith(base + ":") + or target.startswith(base + "/") + ) + if pattern.endswith(":all") or pattern.endswith(":*"): + package = pattern.rsplit(":", 1)[0] + return target.startswith(package + ":") + return target == pattern + + +def target_allowed_by_patterns(target: str, patterns: Sequence[str]) -> bool: + """Checks if target is permitted by positive and negative patterns.""" + if not patterns: + return True + positive_patterns = [ + pattern for pattern in patterns if not pattern.startswith("-") + ] + negative_patterns = [ + pattern[1:] for pattern in patterns if pattern.startswith("-") + ] + if positive_patterns and not any( + target_matches_pattern(target, pattern) for pattern in positive_patterns + ): + return False + if any( + target_matches_pattern(target, pattern) for pattern in negative_patterns + ): + return False + return True + + +def filter_and_normalize_targets( + targets: Iterable[str], + keep_prefixes: Sequence[str] = KEEP_TARGET_PREFIXES, + target_patterns: Sequence[str] = (), +) -> List[str]: + """Normalizes and filters impacted targets against patterns.""" + filtered: List[str] = [] + seen = set() + + for raw_target in targets: + line = raw_target.strip() + if not line or line.startswith("#"): + continue + # Drop external synthetic query targets that are not directly buildable + if line.startswith(("//external:", "@external:")): + continue + + normalized = normalize_target_label(line) + if not normalized: + continue + + if any(normalized.startswith(prefix) for prefix in keep_prefixes): + if target_patterns and not target_allowed_by_patterns( + normalized, target_patterns + ): + continue + if normalized not in seen: + seen.add(normalized) + filtered.append(normalized) + + return sorted(filtered) + + +def verify_sha256(file_path: str, expected_sha256: str) -> bool: + """Verifies the SHA256 checksum of a file.""" + if not os.path.isfile(file_path): + return False + hasher = hashlib.sha256() + with open(file_path, "rb") as f: + while chunk := f.read(65536): + hasher.update(chunk) + return hasher.hexdigest().lower() == expected_sha256.lower() + + +def download_and_verify_jar( + dest_path: Optional[str] = None, + url: str = BAZEL_DIFF_JAR_URL, + expected_sha256: str = BAZEL_DIFF_JAR_SHA256, +) -> str: + """Downloads bazel-diff deploy jar with retries and verifies its SHA-256.""" + if dest_path is None: + dest_path = os.path.join(tempfile.gettempdir(), "bazel-diff.jar") + if verify_sha256(dest_path, expected_sha256): + logging.info("Reusing existing verified bazel-diff jar at %s", dest_path) + return dest_path + + os.makedirs(os.path.dirname(os.path.abspath(dest_path)), exist_ok=True) + temp_dest = f"{dest_path}.tmp.{os.getpid()}" + + command = [ + "curl", + "-fLo", + temp_dest, + "--retry", + "5", + "--retry-connrefused", + "--retry-all-errors", + url, + ] + logging.info("Downloading bazel-diff jar via curl: %s", " ".join(command)) + subprocess.run(command, check=True) + if not verify_sha256(temp_dest, expected_sha256): + if os.path.exists(temp_dest): + os.remove(temp_dest) + raise ValueError( + f"SHA-256 mismatch for downloaded bazel-diff jar from {url}" + ) + + os.replace(temp_dest, dest_path) + return dest_path + + +def run_generate_hashes( + jar_path: str, + workspace_dir: str, + output_json: str, + cquery_options: Optional[Sequence[str]] = None, + use_cquery: bool = True, + startup_options: Optional[Sequence[str]] = None, + cwd: Optional[str] = None, +) -> None: + """Executes bazel-diff generate-hashes for the current workspace state.""" + fine_grained_repos = "@tsl,@@+tsl_extension+tsl" if use_cquery else "@tsl" + args = [ + "java", + "-jar", + jar_path, + "generate-hashes", + "-k", + "-w", + workspace_dir, + output_json, + f"--fineGrainedHashExternalRepos={fine_grained_repos}", + ] + if use_cquery: + args.append("--useCquery") + args.append( + "--cqueryExpression=deps(//xla/...:all-targets +" + " //build_tools/...:all-targets + @tsl//tsl/...:all-targets)" + ) + if cquery_options: + args.append(f"--cqueryCommandOptions={' '.join(cquery_options)}") + if startup_options: + args.append(f"--bazelStartupOptions={' '.join(startup_options)}") + + logging.info("Running bazel-diff generate-hashes: %s", " ".join(args)) + subprocess.run(args, cwd=cwd or workspace_dir, check=True) + + +def run_get_impacted_targets( + jar_path: str, + workspace_dir: str, + starting_hashes: str, + final_hashes: str, + output_txt: str, + cwd: Optional[str] = None, +) -> None: + """Executes bazel-diff get-impacted-targets.""" + args = [ + "java", + "-jar", + jar_path, + "get-impacted-targets", + "-sh", + starting_hashes, + "-fh", + final_hashes, + "-w", + workspace_dir, + "-o", + output_txt, + ] + logging.info("Running bazel-diff get-impacted-targets: %s", " ".join(args)) + subprocess.run(args, cwd=cwd or workspace_dir, check=True) + + +def report_decision(decision: BazelDiffDecision, build_name: str) -> None: + """Emits decision summary to stdout and appends to $GITHUB_STEP_SUMMARY.""" + summary_line = ( + f"bazel-diff Build='{build_name}': " + f"Decision={decision.decision.value}, " + f"ChangedFiles={decision.changed_files_count}, " + f"ImpactedTargets={len(decision.impacted_targets)}, " + f"Reason='{decision.reason}', " + f"Duration={decision.elapsed_seconds:.1f}s" + ) + separator = "=" * 55 + sys.stdout.write(f"\n{separator}\n") + sys.stdout.write(f"{summary_line}\n") + if decision.decision == BazelDiffDecisionType.IMPACTED: + sys.stdout.write( + "Sample impacted targets (up to 10 of" + f" {len(decision.impacted_targets)}):\n" + ) + for target in decision.impacted_targets[:10]: + sys.stdout.write(f" {target}\n") + sys.stdout.write(f"{separator}\n\n") + sys.stdout.flush() + + step_summary_file = os.environ.get("GITHUB_STEP_SUMMARY") + if step_summary_file: + try: + with open(step_summary_file, "a", encoding="utf-8") as f: + f.write("### bazel-diff Impact Analysis\n") + f.write( + "| Parameter | Value |\n" + "|---|---|\n" + f"| **Build** | `{build_name}` |\n" + f"| **Decision** | **`{decision.decision.value}`** |\n" + f"| **Changed files** | {decision.changed_files_count} |\n" + f"| **Impacted targets** | {len(decision.impacted_targets)} |\n" + f"| **Reason** | {decision.reason} |\n" + f"| **Duration** | {decision.elapsed_seconds:.1f}s |\n\n" + ) + if decision.decision == BazelDiffDecisionType.IMPACTED: + f.write("
Impacted Targets List\n\n```\n") + for target in decision.impacted_targets: + f.write(f"{target}\n") + f.write("```\n
\n\n") + except OSError as e: + logging.warning("Failed to write to GITHUB_STEP_SUMMARY: %s", e) + + +def dict_to_cli_options(options_dict: Any) -> List[str]: + """Converts a dict of options to Bazel CLI flags.""" + if not isinstance(options_dict, dict): + return [] + options = [] + for key, value in options_dict.items(): + if isinstance(value, bool) and value: + options.append(f"--{key}") + elif isinstance(value, (list, tuple)): + options.extend(f"--{key}={item}" for item in value) + else: + options.append(f"--{key}={value}") + return options + + +def get_cquery_command_options(build: BuildConfig) -> List[str]: + """Extracts --config, --repo_env, and Starlark build flags from a Build.""" + options = [f"--config={config}" for config in build.configs] + for key, value in build.repo_env.items(): + options.append(f"--repo_env={key}={value}") + for key, value in build.options.items(): + if key.startswith(("//", "@")): + if isinstance(value, bool) and value: + options.append(f"--{key}") + else: + options.append(f"--{key}={value}") + return options + + +def compute_impacted_targets( + build: BuildConfig, + base_sha: str, + head_sha: str = "HEAD", + workspace_dir: str = ".", + dest_jar_path: Optional[str] = None, +) -> BazelDiffDecision: + """Computes impacted targets for a given build between base_sha and head_sha. + + Guaranteed fail-open: any unexpected error returns FULL decision so CI never + fails due to a failure in bazel-diff. + + Args: + build: Build object defining the target CI job configuration. + base_sha: Base git SHA to compare against. + head_sha: Head git SHA. + workspace_dir: Absolute or relative path to the workspace root directory. + dest_jar_path: Path where the bazel-diff jar will be cached. + + Returns: + A BazelDiffDecision specifying whether to run FULL, SKIP, or IMPACTED tests. + """ + start_time = time.time() + workspace_dir = os.path.abspath(workspace_dir) + if dest_jar_path is None: + dest_jar_path = os.path.join(tempfile.gettempdir(), "bazel-diff.jar") + + # Record original ref currently checked out so we can faithfully restore it + original_head = head_sha + result = _run_git_best_effort(["rev-parse", "HEAD"], cwd=workspace_dir) + if result and result.returncode == 0 and result.stdout.strip(): + original_head = result.stdout.strip() + + if not head_sha or head_sha == "HEAD": + head_sha = original_head + + # Fetch base SHA if shallow + _run_git_best_effort( + ["fetch", "--depth=1", "origin", base_sha], cwd=workspace_dir + ) + + if not get_merge_base(base_sha, head_sha, cwd=workspace_dir): + _run_git_best_effort(["fetch", "--unshallow", "origin"], cwd=workspace_dir) + + # 1. Changed files analysis + try: + changed_files = get_changed_files( + base_sha, head_sha, cwd=workspace_dir + ) + except (OSError, subprocess.SubprocessError) as e: + return BazelDiffDecision( + decision=BazelDiffDecisionType.FULL, + reason=f"Failed to get git changed files: {e}", + elapsed_seconds=time.time() - start_time, + ) + + changed_count = len(changed_files) + if is_global_config_changed(changed_files): + return BazelDiffDecision( + decision=BazelDiffDecisionType.FULL, + reason="Global Bazel configuration modified", + changed_files_count=changed_count, + elapsed_seconds=time.time() - start_time, + ) + + if touches_build_keywords( + changed_files, base_sha, head_sha, cwd=workspace_dir + ): + return BazelDiffDecision( + decision=BazelDiffDecisionType.FULL, + reason=( + "Diff touches label_flag, config_setting, or alias (b/519689071)" + ), + changed_files_count=changed_count, + elapsed_seconds=time.time() - start_time, + ) + + if is_docs_or_metadata_only(changed_files): + return BazelDiffDecision( + decision=BazelDiffDecisionType.SKIP, + reason="Only documentation or repository metadata modified", + changed_files_count=changed_count, + elapsed_seconds=time.time() - start_time, + ) + + # 2. Download bazel-diff + try: + jar_path = download_and_verify_jar(dest_jar_path) + except (OSError, subprocess.SubprocessError, ValueError) as e: + return BazelDiffDecision( + decision=BazelDiffDecisionType.FULL, + reason=f"Failed to download/verify bazel-diff jar: {e}", + changed_files_count=changed_count, + elapsed_seconds=time.time() - start_time, + ) + + # 3. Generate hashes with git checkout guarantees + temp_dir = tempfile.mkdtemp(prefix="bazel_diff_") + base_hashes_json = os.path.join(temp_dir, "base_hashes.json") + head_hashes_json = os.path.join(temp_dir, "head_hashes.json") + impacted_txt = os.path.join(temp_dir, "impacted_targets.txt") + + cquery_options = get_cquery_command_options(build) + use_cquery = build.bazel_diff_use_cquery + startup_options = None + if build.startup_options: + startup_options = dict_to_cli_options(build.startup_options) + + diff_base = get_diff_base(base_sha, head_sha, cwd=workspace_dir) + try: + # Generate base hashes + subprocess.run( + ["git", "checkout", "-f", diff_base], + cwd=workspace_dir, + check=True, + capture_output=True, + ) + run_generate_hashes( + jar_path, + workspace_dir, + base_hashes_json, + cquery_options=cquery_options, + use_cquery=use_cquery, + startup_options=startup_options, + ) + + # Generate head hashes + subprocess.run( + ["git", "checkout", "-f", head_sha], + cwd=workspace_dir, + check=True, + capture_output=True, + ) + run_generate_hashes( + jar_path, + workspace_dir, + head_hashes_json, + cquery_options=cquery_options, + use_cquery=use_cquery, + startup_options=startup_options, + ) + + # Determine impacted targets + run_get_impacted_targets( + jar_path, + workspace_dir, + base_hashes_json, + head_hashes_json, + impacted_txt, + ) + + with open(impacted_txt, "r", encoding="utf-8") as f: + raw_targets = f.readlines() + except (OSError, subprocess.SubprocessError, ValueError) as e: + logging.warning("bazel-diff hashing/analysis failed: %s", e) + return BazelDiffDecision( + decision=BazelDiffDecisionType.FULL, + reason=f"Exception during bazel-diff run: {e}", + changed_files_count=changed_count, + elapsed_seconds=time.time() - start_time, + ) + finally: + # Guaranteed restoration to original HEAD + _run_git_best_effort(["checkout", "-f", original_head], cwd=workspace_dir) + + target_patterns = build.target_patterns + filtered_targets = filter_and_normalize_targets( + raw_targets, target_patterns=target_patterns + ) + + if not filtered_targets: + return BazelDiffDecision( + decision=BazelDiffDecisionType.SKIP, + impacted_targets=(), + reason="No targets impacted after filtering", + changed_files_count=changed_count, + elapsed_seconds=time.time() - start_time, + ) + + # Rewrite the impacted targets file with the filtered list + filtered_impacted_file = os.path.join( + temp_dir, "filtered_impacted_targets.txt" + ) + try: + with open(filtered_impacted_file, "w", encoding="utf-8") as f: + for target in filtered_targets: + f.write(f"{target}\n") + if target_patterns: + for pattern in target_patterns: + if pattern.startswith("-"): + f.write(f"{pattern}\n") + except OSError as e: + logging.warning("Failed to write filtered impacted targets file: %s", e) + return BazelDiffDecision( + decision=BazelDiffDecisionType.FULL, + reason=f"Exception writing impacted targets file: {e}", + changed_files_count=changed_count, + elapsed_seconds=time.time() - start_time, + ) + + return BazelDiffDecision( + decision=BazelDiffDecisionType.IMPACTED, + impacted_targets=tuple(filtered_targets), + impacted_targets_file=filtered_impacted_file, + reason=f"Found {len(filtered_targets)} impacted targets", + changed_files_count=changed_count, + elapsed_seconds=time.time() - start_time, + ) diff --git a/third_party/xla/build_tools/ci/bazel_diff_test.py b/third_party/xla/build_tools/ci/bazel_diff_test.py new file mode 100644 index 00000000000000..f663d36e0e943d --- /dev/null +++ b/third_party/xla/build_tools/ci/bazel_diff_test.py @@ -0,0 +1,391 @@ +# Copyright 2026 The OpenXLA Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""Tests for bazel_diff module.""" + +import hashlib +import os +import tempfile +from unittest import mock + +from absl.testing import absltest +from absl.testing import parameterized + +from build_tools.ci import bazel_diff + + +class BazelDiffTest(parameterized.TestCase): + + @parameterized.parameters( + (["MODULE.bazel"], True), + (["REPO.bazel"], True), + (["WORKSPACE"], True), + (["WORKSPACE.bzlmod"], True), + ([".bazelrc"], True), + ([".bazelversion"], True), + (["xla/flags.bzl"], False), + (["third_party/tsl/tsl/platform/default/rules.bzl"], False), + (["third_party/tsl/REPO.bazel"], True), + (["third_party/tsl/tsl/BUILD"], False), + (["third_party/tsl/tsl/platform/denormal.cc"], False), + (["xla/service/cpu/cpu_compiler.cc"], False), + (["docs/overview.md", "xla/BUILD"], False), + ) + def test_is_global_config_changed(self, changed_files, expected): + self.assertEqual( + bazel_diff.is_global_config_changed(changed_files), expected + ) + + @parameterized.parameters( + (["README.md", "docs/index.md"], True), + (["OWNERS", "LICENSE"], True), + ([".gitignore", ".vscode/settings.json"], True), + (["logo.png", "diagram.svg"], True), + (["README.md", "xla/service/hlo_parser.cc"], False), + ([], False), + ) + def test_is_docs_or_metadata_only(self, changed_files, expected): + self.assertEqual( + bazel_diff.is_docs_or_metadata_only(changed_files), expected + ) + + @parameterized.parameters( + ("//xla/service:cpu_compiler", "//xla/service:cpu_compiler"), + ("@@//xla/python:types", "//xla/python:types"), + ("@//xla/python:types", "//xla/python:types"), + ("@@+tsl+tsl//tsl/platform:logging", "@tsl//tsl/platform:logging"), + ("@@tsl//tsl/platform:logging", "@tsl//tsl/platform:logging"), + ("@tsl//tsl/platform:logging", "@tsl//tsl/platform:logging"), + ("@@llvm-project//llvm:Support", "@llvm-project//llvm:Support"), + ("", ""), + ) + def test_normalize_target_label(self, input_label, expected): + self.assertEqual(bazel_diff.normalize_target_label(input_label), expected) + + def test_filter_and_normalize_targets(self): + raw_targets = [ + "//xla/service:hlo_parser", + "@@//xla/tools:runner", + "@@+tsl+tsl//tsl/platform:logging", + "//build_tools/ci:build", + "//external:tsl", + "@external:llvm", + "@unrelated//foo:bar", + "//xla/service:hlo_parser", # duplicate + "# comment line", + " ", + ] + expected = [ + "//build_tools/ci:build", + "//xla/service:hlo_parser", + "//xla/tools:runner", + "@tsl//tsl/platform:logging", + ] + self.assertEqual( + bazel_diff.filter_and_normalize_targets(raw_targets), expected + ) + + @parameterized.parameters( + ("//xla/service/gpu:gpu_compiler", "//xla/service/gpu/...", True), + ("//xla/service/gpu/sub:target", "//xla/service/gpu/...", True), + ("//xla/service/cpu:cpu_compiler", "//xla/service/gpu/...", False), + ("//xla/service/gpu_other:target", "//xla/service/gpu/...", False), + ("//xla/backends/cpu:test", "//xla/backends/cpu:all", True), + ("//xla/backends/cpu:test", "//xla/backends/cpu:*", True), + ("//xla/backends/cpu:test", "//xla/backends/cpu:other", False), + ("//xla/backends/cpu:test", "//xla/backends/cpu:test", True), + ) + def test_target_matches_pattern(self, target, pattern, expected): + self.assertEqual( + bazel_diff.target_matches_pattern(target, pattern), expected + ) + + def test_target_allowed_by_patterns(self): + patterns = ( + "//xla/...", + "@tsl//tsl/...", + "-//xla/backends/gpu/...", + "-//xla/service/gpu/...", + "-//xla/tpu/...", + "-//xla/backends/cpu/collectives:gloo_collectives_test", + ) + self.assertTrue( + bazel_diff.target_allowed_by_patterns( + "//xla/service/cpu:cpu_compiler", patterns + ) + ) + self.assertTrue( + bazel_diff.target_allowed_by_patterns( + "@tsl//tsl/platform:default/denormal.cc", patterns + ) + ) + self.assertFalse( + bazel_diff.target_allowed_by_patterns( + "//xla/service/gpu:gpu_compiler", patterns + ) + ) + self.assertFalse( + bazel_diff.target_allowed_by_patterns( + "//xla/backends/gpu/tests:gpu_atomic_test", patterns + ) + ) + self.assertFalse( + bazel_diff.target_allowed_by_patterns( + "//xla/backends/cpu/collectives:gloo_collectives_test", patterns + ) + ) + self.assertTrue( + bazel_diff.target_allowed_by_patterns( + "//xla/backends/cpu/collectives:mpi_test", patterns + ) + ) + self.assertFalse( + bazel_diff.target_allowed_by_patterns( + "@unrelated//pkg:target", patterns + ) + ) + + def test_filter_and_normalize_targets_with_patterns(self): + raw_targets = [ + "//xla/service/gpu:gpu_compiler", + "//xla/service/cpu:cpu_compiler", + "//xla/backends/gpu/tests:gpu_atomic_test", + "@tsl//tsl/platform:default/denormal.cc", + ] + patterns = ( + "//xla/...", + "@tsl//tsl/...", + "-//xla/service/gpu/...", + "-//xla/backends/gpu/...", + ) + expected = [ + "//xla/service/cpu:cpu_compiler", + "@tsl//tsl/platform:default/denormal.cc", + ] + self.assertEqual( + bazel_diff.filter_and_normalize_targets( + raw_targets, target_patterns=patterns + ), + expected, + ) + + def test_verify_sha256(self): + with tempfile.NamedTemporaryFile("wb", delete=False) as f: + f.write(b"hello world") + temp_path = f.name + + try: + expected_hash = hashlib.sha256(b"hello world").hexdigest() + self.assertTrue(bazel_diff.verify_sha256(temp_path, expected_hash)) + self.assertFalse(bazel_diff.verify_sha256(temp_path, "0" * 64)) + self.assertFalse( + bazel_diff.verify_sha256("/non/existent/file", expected_hash) + ) + finally: + os.unlink(temp_path) + + def test_report_decision(self): + with tempfile.NamedTemporaryFile("w", delete=False) as f: + summary_file = f.name + + try: + with mock.patch.dict(os.environ, {"GITHUB_STEP_SUMMARY": summary_file}): + decision = bazel_diff.BazelDiffDecision( + decision=bazel_diff.BazelDiffDecisionType.IMPACTED, + impacted_targets=("//xla/service:hlo_parser",), + reason="Found 1 impacted target", + changed_files_count=2, + elapsed_seconds=12.3, + ) + bazel_diff.report_decision(decision, "xla_linux_x86_gpu_l4") + + with open(summary_file, "r") as f: + content = f.read() + + self.assertIn("bazel-diff Impact Analysis", content) + self.assertIn("xla_linux_x86_gpu_l4", content) + self.assertIn("IMPACTED", content) + self.assertIn("//xla/service:hlo_parser", content) + finally: + os.unlink(summary_file) + + @mock.patch.object(bazel_diff, "get_changed_files") + def test_compute_impacted_targets_global_config_fallback(self, mock_changed): + mock_changed.return_value = ["MODULE.bazel", "xla/service/cpu_compiler.cc"] + build_mock = mock.MagicMock() + decision = bazel_diff.compute_impacted_targets( + build_mock, "base_sha", "head_sha" + ) + self.assertEqual(decision.decision, bazel_diff.BazelDiffDecisionType.FULL) + self.assertIn("Global Bazel configuration", decision.reason) + + @mock.patch.object(bazel_diff.subprocess, "run") + @mock.patch.object(bazel_diff, "get_diff_base") + def test_touches_build_keywords_sensitive(self, mock_base, mock_run): + mock_base.return_value = "base_sha" + mock_run.return_value = mock.MagicMock( + returncode=0, + stdout=( + "+ label_flag(name = 'custom_flag', default = '//xla/pkg:target')\n" + ), + ) + self.assertTrue( + bazel_diff.touches_build_keywords( + ["xla/flags.bzl"], "base_sha", "head_sha" + ) + ) + + @mock.patch.object(bazel_diff.subprocess, "run") + @mock.patch.object(bazel_diff, "get_diff_base") + def test_touches_build_keywords_benign_bzl(self, mock_base, mock_run): + mock_base.return_value = "base_sha" + mock_run.return_value = mock.MagicMock( + returncode=0, + stdout=( + "+def my_macro(name, deps=[]):\n" + "+ native.cc_library(name='hlo_alias_analysis')\n" + ), + ) + self.assertFalse( + bazel_diff.touches_build_keywords( + ["xla/service/custom.bzl"], "base_sha", "head_sha" + ) + ) + mock_run.assert_called_once_with( + [ + "git", + "diff", + "-U0", + "base_sha", + "head_sha", + "--", + "xla/service/custom.bzl", + ], + cwd=".", + capture_output=True, + text=True, + check=True, + ) + + def test_get_cquery_command_options(self): + build_mock = mock.MagicMock() + build_mock.configs = ( + "warnings", + "rbe_linux_cuda_nvcc", + "hermetic_cuda_umd", + ) + build_mock.repo_env = {"TF_CUDA_COMPUTE_CAPABILITIES": "7.5"} + build_mock.options = { + "//xla/tsl:ci_build": True, + "@local_config_cuda//cuda:override_include_cuda_libs": False, + "flaky_test_attempts": 3, + "test_output": "errors", + } + self.assertEqual( + bazel_diff.get_cquery_command_options(build_mock), + [ + "--config=warnings", + "--config=rbe_linux_cuda_nvcc", + "--config=hermetic_cuda_umd", + "--repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5", + "--//xla/tsl:ci_build", + "--@local_config_cuda//cuda:override_include_cuda_libs=False", + ], + ) + + @mock.patch.object(bazel_diff.subprocess, "run") + def test_run_generate_hashes_cquery_options(self, mock_run): + bazel_diff.run_generate_hashes( + "/tmp/bazel-diff.jar", + "/workspace", + "/tmp/hashes.json", + cquery_options=[ + "--config=warnings", + "--config=rbe_linux_cuda_nvcc", + "--repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5", + "--//xla/tsl:ci_build", + ], + ) + args = mock_run.call_args[0][0] + self.assertIn( + "--fineGrainedHashExternalRepos=@tsl,@@+tsl_extension+tsl", + args, + ) + self.assertIn( + "--cqueryCommandOptions=--config=warnings --config=rbe_linux_cuda_nvcc" + " --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --//xla/tsl:ci_build", + args, + ) + + @mock.patch.object(bazel_diff, "touches_build_keywords") + @mock.patch.object(bazel_diff, "get_changed_files") + def test_compute_impacted_targets_build_keywords_fallback( + self, mock_changed, mock_keywords + ): + mock_changed.return_value = ["xla/flags.bzl"] + mock_keywords.return_value = True + build_mock = mock.MagicMock() + decision = bazel_diff.compute_impacted_targets( + build_mock, "base_sha", "head_sha" + ) + self.assertEqual(decision.decision, bazel_diff.BazelDiffDecisionType.FULL) + self.assertIn("label_flag, config_setting, or alias", decision.reason) + + @mock.patch.object(bazel_diff, "get_changed_files") + def test_compute_impacted_targets_docs_only_skip(self, mock_changed): + mock_changed.return_value = ["README.md", "docs/architecture.md"] + build_mock = mock.MagicMock() + decision = bazel_diff.compute_impacted_targets( + build_mock, "base_sha", "head_sha" + ) + self.assertEqual(decision.decision, bazel_diff.BazelDiffDecisionType.SKIP) + self.assertIn("Only documentation", decision.reason) + + @mock.patch.object(bazel_diff, "get_changed_files") + def test_compute_impacted_targets_git_fail_open(self, mock_changed): + mock_changed.side_effect = OSError("git failure") + build_mock = mock.MagicMock() + decision = bazel_diff.compute_impacted_targets( + build_mock, "base_sha", "head_sha" + ) + self.assertEqual(decision.decision, bazel_diff.BazelDiffDecisionType.FULL) + self.assertIn("Failed to get git changed files", decision.reason) + + @mock.patch.object(bazel_diff.subprocess, "run") + def test_get_merge_base_success(self, mock_run): + mock_run.return_value = mock.MagicMock(returncode=0, stdout="base_commit\n") + self.assertEqual(bazel_diff.get_merge_base("base", "head"), "base_commit") + + @mock.patch.object(bazel_diff.subprocess, "run") + def test_get_merge_base_failure(self, mock_run): + mock_run.return_value = mock.MagicMock(returncode=128, stdout="") + self.assertIsNone(bazel_diff.get_merge_base("base", "head")) + + @mock.patch.object(bazel_diff, "get_merge_base") + def test_get_diff_base_with_merge_base(self, mock_mb): + mock_mb.return_value = "merge_base_sha" + self.assertEqual( + bazel_diff.get_diff_base("base_sha", "head_sha"), "merge_base_sha" + ) + + @mock.patch.object(bazel_diff, "get_merge_base") + def test_get_diff_base_fallback(self, mock_mb): + mock_mb.return_value = None + self.assertEqual( + bazel_diff.get_diff_base("base_sha", "head_sha"), "base_sha" + ) + + +if __name__ == "__main__": + absltest.main() diff --git a/third_party/xla/build_tools/ci/build.py b/third_party/xla/build_tools/ci/build.py index 20ea7038739ac5..47f3b412843569 100755 --- a/third_party/xla/build_tools/ci/build.py +++ b/third_party/xla/build_tools/ci/build.py @@ -1,4 +1,3 @@ -#!/usr/bin/python3 # Copyright 2024 The OpenXLA Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -31,7 +30,9 @@ import os import subprocess import sys -from typing import Any, ClassVar, Dict, List, Tuple +from typing import Any, ClassVar, Dict, List, Optional, Tuple + +from build_tools.ci import bazel_diff # TODO(ddunleavy): move this to the bazelrc _DEFAULT_BAZEL_OPTIONS = dict( @@ -132,7 +133,27 @@ def from_str(cls, s): @dataclasses.dataclass(frozen=True, **_KW_ONLY_IF_PYTHON310) class Build: - """Class representing a build of XLA.""" + """Class representing a build of XLA. + + Attributes: + type_: The BuildType enum value identifying this build. + repo: The repository under test (e.g. "openxla/xla"). + target_patterns: Tuple of Bazel target patterns to build or test. + subcommand: The Bazel subcommand to execute ("test" by default). + configs: Tuple of Bazel --config names. + build_tag_filters: Tuple of tags for --build_tag_filters. + test_tag_filters: Tuple of tags for --test_tag_filters. + action_env: Dictionary of --action_env variables. + test_env: Dictionary of --test_env variables. + repo_env: Dictionary of --repo_env variables. + override_repository: Dictionary of --override_repository mappings. + override_module: Dictionary of --override_module mappings. + options: Dictionary of additional Bazel flags. + startup_options: Dictionary of Bazel startup options (before subcommand). + extra_setup_commands: Tuple of shell commands to run before Bazel. + use_bazel_diff: Whether to enable bazel-diff target filtering on presubmit. + bazel_diff_use_cquery: Whether bazel-diff should use cquery (--useCquery). + """ _builds: ClassVar[Dict[BuildType, "Build"]] = {} @@ -151,6 +172,8 @@ class Build: options: Dict[str, Any] = dataclasses.field(default_factory=dict) startup_options: Dict[str, Any] = dataclasses.field(default_factory=dict) extra_setup_commands: Tuple[List[str], ...] = () + use_bazel_diff: bool = False + bazel_diff_use_cquery: bool = True def __post_init__(self): # pylint: disable=protected-access @@ -169,13 +192,17 @@ def all_builds(cls): return cls._builds def bazel_command( - self, subcommand: str = "test", extra_options: Tuple[str, ...] = () + self, + subcommand: str = "test", + extra_options: Tuple[str, ...] = (), + target_pattern_file: Optional[str] = None, ) -> List[str]: """Returns a bazel test command for this build. Args: subcommand: The subcommand to give to bazel. `test` by default. extra_options: Extra options. For now just used to pass in `--nobuild`. + target_pattern_file: Optional file containing explicit target patterns. Returns: List of command line arguments """ @@ -209,6 +236,18 @@ def bazel_command( + options + list(extra_options) ) + if target_pattern_file: + extra = ["--skip_incompatible_explicit_targets"] + if self.subcommand == "test" and "--build_tests_only" not in all_options: + extra.append("--build_tests_only") + return [ + "bazel", + *startup_options, + subcommand, + *all_options, + *extra, + f"--target_pattern_file={target_pattern_file}", + ] return [ "bazel", *startup_options, @@ -218,7 +257,9 @@ def bazel_command( *self.target_patterns, ] - def commands(self) -> List[List[str]]: + def commands( + self, target_pattern_file: Optional[str] = None + ) -> List[List[str]]: """Returns list of commands for a build.""" cmds = [] @@ -241,11 +282,18 @@ def commands(self) -> List[List[str]]: cmds.append( retry( self.bazel_command( - subcommand="build", extra_options=("--nobuild",) + subcommand="build", + extra_options=("--nobuild",), + target_pattern_file=target_pattern_file, ) ) ) - cmds.append(self.bazel_command(subcommand=self.subcommand)) + cmds.append( + self.bazel_command( + subcommand=self.subcommand, + target_pattern_file=target_pattern_file, + ) + ) return cmds @@ -303,6 +351,8 @@ def nvidia_gpu_build_with_compute_capability( configs: Tuple[str, ...], compute_capability: int, multi_gpu: bool = False, + use_bazel_diff: bool = False, + bazel_diff_use_cquery: bool = True, ) -> Build: """Returns a build for a Nvidia GPU build with the given compute capability.""" repo_env = {"TF_CUDA_COMPUTE_CAPABILITIES": f"{compute_capability/10}"} @@ -349,6 +399,8 @@ def nvidia_gpu_build_with_compute_capability( options=options, repo_env=repo_env, extra_setup_commands=(["nvidia-smi"],) if multi_gpu else (), + use_bazel_diff=use_bazel_diff, + bazel_diff_use_cquery=bazel_diff_use_cquery, ) @@ -367,6 +419,8 @@ def nvidia_gpu_build_with_compute_capability( build_tag_filters=cpu_x86_tag_filter, test_tag_filters=cpu_x86_tag_filter, options={**_DEFAULT_BAZEL_OPTIONS, "//xla/tsl:ci_build": True}, + use_bazel_diff=True, + bazel_diff_use_cquery=True, ) windows_x86_tag_filter = ( @@ -433,6 +487,8 @@ def nvidia_gpu_build_with_compute_capability( build_tag_filters=cpu_x86_tag_filter, test_tag_filters=cpu_x86_tag_filter, options={**_DEFAULT_BAZEL_OPTIONS, "//xla/tsl:ci_build": True}, + use_bazel_diff=True, + bazel_diff_use_cquery=True, ) cpu_arm_tag_filter = ( @@ -455,6 +511,8 @@ def nvidia_gpu_build_with_compute_capability( }, build_tag_filters=cpu_arm_tag_filter, test_tag_filters=cpu_arm_tag_filter, + use_bazel_diff=True, + bazel_diff_use_cquery=True, ) nvidia_gpu_build_with_compute_capability( @@ -462,6 +520,8 @@ def nvidia_gpu_build_with_compute_capability( configs=("warnings", "rbe_linux_cuda_nvcc", "hermetic_cuda_umd"), compute_capability=75, multi_gpu=False, + use_bazel_diff=True, + bazel_diff_use_cquery=True, ) nvidia_gpu_build_with_compute_capability( @@ -469,6 +529,8 @@ def nvidia_gpu_build_with_compute_capability( configs=("warnings", "rbe_linux_cuda_nvcc", "hermetic_cuda_umd"), compute_capability=90, multi_gpu=True, + use_bazel_diff=True, + bazel_diff_use_cquery=True, ) oneapi_build_tag_filter = ( @@ -882,8 +944,19 @@ def dump_all_build_commands(): # Awkward workaround b/c Build instances are not hashable for build in sorted(Build.all_builds().values(), key=lambda b: str(b.type_)): sys.stdout.write(f"# BEGIN {build.type_}\n") - for cmd in build.commands(): - sys.stdout.write(" ".join(cmd) + "\n") + for command in build.commands(target_pattern_file=None): + sys.stdout.write(" ".join(command) + "\n") + if build.use_bazel_diff: + sys.stdout.write("# With bazel-diff (presubmit):\n") + if build.bazel_diff_use_cquery: + cquery_options = " ".join(bazel_diff.get_cquery_command_options(build)) + sys.stdout.write(f"# bazel-diff cquery_options: {cquery_options}\n") + else: + sys.stdout.write("# bazel-diff use_cquery: False\n") + for command in build.commands( + target_pattern_file="/tmp/target_pattern_file" + ): + sys.stdout.write(" ".join(command) + "\n") sys.stdout.write(f"# END {build.type_}\n") @@ -913,9 +986,43 @@ def main(): if args.dump_commands: dump_all_build_commands() return - else: - for cmd in Build.all_builds()[args.build].commands(): - sh(cmd) + + build = Build.all_builds()[args.build] + github_event = os.environ.get("GITHUB_EVENT_NAME", "").strip() + is_presubmit = github_event in ("", "pull_request", "pull_request_target") + base_sha = os.environ.get("XLA_CI_BAZEL_DIFF_BASE_SHA", "").strip() + head_sha = os.environ.get("XLA_CI_BAZEL_DIFF_HEAD_SHA", "HEAD").strip() + disabled = os.environ.get("XLA_CI_DISABLE_BAZEL_DIFF", "false").lower() in ( + "true", + "1", + "yes", + ) + + target_pattern_file = None + if build.use_bazel_diff and is_presubmit and base_sha and not disabled: + logging.info( + "Running bazel-diff analysis (base=%s, head=%s)", + base_sha, + head_sha, + ) + decision = bazel_diff.compute_impacted_targets(build, base_sha, head_sha) + bazel_diff.report_decision(decision, str(build.type_)) + + if decision.decision == bazel_diff.BazelDiffDecisionType.SKIP: + logging.info("bazel-diff: 0 impacted targets, skipping Bazel commands!") + return + elif decision.decision == bazel_diff.BazelDiffDecisionType.IMPACTED: + target_pattern_file = decision.impacted_targets_file + + for command in build.commands(target_pattern_file=target_pattern_file): + result = sh(command, check=False) + if result.returncode == 4 and target_pattern_file: + logging.info( + "Bazel returned exit code 4 (no tests found), treating as success." + ) + continue + if result.returncode != 0: + sys.exit(result.returncode) if __name__ == "__main__": diff --git a/third_party/xla/build_tools/ci/clang_apply_replacements_wrapper.sh.tpl b/third_party/xla/build_tools/ci/clang_apply_replacements_wrapper.sh.tpl index 20a55ae3b2f85c..30eaec28d147ca 100644 --- a/third_party/xla/build_tools/ci/clang_apply_replacements_wrapper.sh.tpl +++ b/third_party/xla/build_tools/ci/clang_apply_replacements_wrapper.sh.tpl @@ -23,7 +23,10 @@ echoerr() { printf "${RED}ERROR:${NOCOLOR} %s\n" "$*" >&2 } -REAL_BIN="$PWD/external/%LLVM_REPO_NAME%/bin/clang-apply-replacements" +REAL_BIN="$PWD/external/rules_ml_toolchain++toolchain_ext+%LLVM_REPO_NAME%/bin/clang-apply-replacements" +if [ ! -f "$REAL_BIN" ]; then + REAL_BIN="$PWD/external/%LLVM_REPO_NAME%/bin/clang-apply-replacements" +fi if [ ! -f "$REAL_BIN" ]; then echoerr "Failed to locate clang-apply-replacements binary at: $REAL_BIN" exit 1 diff --git a/third_party/xla/build_tools/ci/clang_tidy_diff.py b/third_party/xla/build_tools/ci/clang_tidy_diff.py index 402fd19202b530..732d2bd7c8b56c 100644 --- a/third_party/xla/build_tools/ci/clang_tidy_diff.py +++ b/third_party/xla/build_tools/ci/clang_tidy_diff.py @@ -218,15 +218,17 @@ def _resolve_clang_apply_replacements( runfiles_dir = os.environ.get("RUNFILES_DIR") if not runfiles_dir: return None - runfiles_path = ( - pathlib.Path(runfiles_dir) - / "xla/build_tools/ci/clang_apply_replacements_bin" - ) - if runfiles_path.exists(): - return runfiles_path + for repo_dir in ("_main", "xla"): + runfiles_path = ( + pathlib.Path(runfiles_dir) + / repo_dir + / "build_tools/ci/clang_apply_replacements_bin" + ) + if runfiles_path.exists(): + return runfiles_path _logger().warning( "clang-apply-replacements binary not found in RUNFILES_DIR: %s", - runfiles_path, + runfiles_dir, ) return None diff --git a/third_party/xla/build_tools/ci/clang_tidy_diff_test.py b/third_party/xla/build_tools/ci/clang_tidy_diff_test.py index e185c831d51383..d8de69fc31cc90 100644 --- a/third_party/xla/build_tools/ci/clang_tidy_diff_test.py +++ b/third_party/xla/build_tools/ci/clang_tidy_diff_test.py @@ -644,6 +644,20 @@ def fake_apply_fixes(temp_dir: pathlib.Path) -> None: self.assertLen(staged_dirs, 1) + @parameterized.parameters("_main", "xla") + def test_resolve_clang_apply_replacements_runfiles(self, repo_dir): + tmpdir = self.create_tempdir().full_path + bin_file = ( + pathlib.Path(tmpdir) + / repo_dir + / "build_tools/ci/clang_apply_replacements_bin" + ) + self.create_tempfile(bin_file.as_posix(), content="") + with absltest.mock.patch.dict("os.environ", {"RUNFILES_DIR": tmpdir}): + self.assertEqual( + clang_tidy_diff._resolve_clang_apply_replacements(None), bin_file + ) + if __name__ == "__main__": absltest.main() diff --git a/third_party/xla/build_tools/ci/clang_tidy_wrapper.sh.tpl b/third_party/xla/build_tools/ci/clang_tidy_wrapper.sh.tpl index c33ca7a301d10c..59ff8ffd9db3ad 100755 --- a/third_party/xla/build_tools/ci/clang_tidy_wrapper.sh.tpl +++ b/third_party/xla/build_tools/ci/clang_tidy_wrapper.sh.tpl @@ -23,7 +23,10 @@ echoerr() { printf "${RED}ERROR:${NOCOLOR} %s\n" "$*" >&2 } -REAL_BIN="$PWD/external/%LLVM_REPO_NAME%/bin/clang-tidy" +REAL_BIN="$PWD/external/rules_ml_toolchain++toolchain_ext+%LLVM_REPO_NAME%/bin/clang-tidy" +if [ ! -f "$REAL_BIN" ]; then + REAL_BIN="$PWD/external/%LLVM_REPO_NAME%/bin/clang-tidy" +fi if [ ! -f "$REAL_BIN" ]; then echoerr "Failed to locate clang-tidy binary at: $REAL_BIN" exit 1 diff --git a/third_party/xla/build_tools/ci/golden_commands.txt b/third_party/xla/build_tools/ci/golden_commands.txt index bffa12b8d1b742..d57f43a389a08e 100644 --- a/third_party/xla/build_tools/ci/golden_commands.txt +++ b/third_party/xla/build_tools/ci/golden_commands.txt @@ -16,6 +16,10 @@ bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu- # BEGIN BuildType.XLA_LINUX_ARM64_CPU_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... +# With bazel-diff (presubmit): +# bazel-diff cquery_options: --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --//xla/tsl:ci_build +parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only --//xla/tsl:ci_build --nobuild --skip_incompatible_explicit_targets --target_pattern_file=/tmp/target_pattern_file +bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only --//xla/tsl:ci_build --skip_incompatible_explicit_targets --target_pattern_file=/tmp/target_pattern_file # END BuildType.XLA_LINUX_ARM64_CPU_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_CPU_128_VCPU_PRESUBMIT_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main //xla/tools:compute_xspace_stats_main @@ -24,15 +28,28 @@ bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu- # BEGIN BuildType.XLA_LINUX_X86_CPU_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... +# With bazel-diff (presubmit): +# bazel-diff cquery_options: --config=warnings --config=nonccl --config=rbe_linux_cpu --//xla/tsl:ci_build +parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild --skip_incompatible_explicit_targets --build_tests_only --target_pattern_file=/tmp/target_pattern_file +bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --skip_incompatible_explicit_targets --build_tests_only --target_pattern_file=/tmp/target_pattern_file # END BuildType.XLA_LINUX_X86_CPU_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_CPU_WORKSPACE_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --config=workspace --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --config=workspace --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... +# With bazel-diff (presubmit): +# bazel-diff cquery_options: --config=warnings --config=nonccl --config=rbe_linux_cpu --config=workspace --//xla/tsl:ci_build +parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --config=workspace --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild --skip_incompatible_explicit_targets --build_tests_only --target_pattern_file=/tmp/target_pattern_file +bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --config=workspace --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --skip_incompatible_explicit_targets --build_tests_only --target_pattern_file=/tmp/target_pattern_file # END BuildType.XLA_LINUX_X86_CPU_WORKSPACE_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_8X_H100_GITHUB_ACTIONS nvidia-smi parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any --test_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=9.0 --repo_env=REMOTE_GPU_TESTING=0 --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --test_sharding_strategy=disabled --test_strategy=exclusive --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any --test_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=9.0 --repo_env=REMOTE_GPU_TESTING=0 --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --test_sharding_strategy=disabled --test_strategy=exclusive -- //xla/... //build_tools/... @tsl//tsl/... +# With bazel-diff (presubmit): +# bazel-diff cquery_options: --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=9.0 --repo_env=REMOTE_GPU_TESTING=0 --//xla/tsl:ci_build +nvidia-smi +parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any --test_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=9.0 --repo_env=REMOTE_GPU_TESTING=0 --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --test_sharding_strategy=disabled --test_strategy=exclusive --nobuild --skip_incompatible_explicit_targets --build_tests_only --target_pattern_file=/tmp/target_pattern_file +bazel test --build_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any --test_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=9.0 --repo_env=REMOTE_GPU_TESTING=0 --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --test_sharding_strategy=disabled --test_strategy=exclusive --skip_incompatible_explicit_targets --build_tests_only --target_pattern_file=/tmp/target_pattern_file # END BuildType.XLA_LINUX_X86_GPU_8X_H100_GITHUB_ACTIONS # BEGIN BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS nvidia-smi @@ -75,6 +92,10 @@ bazel build --build_tag_filters=oneapi-only,requires-gpu-intel,-requires-gpu-amd # BEGIN BuildType.XLA_LINUX_X86_GPU_T4_GITHUB_ACTIONS parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/... //build_tools/... @tsl//tsl/... bazel test --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/... //build_tools/... @tsl//tsl/... +# With bazel-diff (presubmit): +# bazel-diff cquery_options: --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --//xla/tsl:ci_build +parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild --skip_incompatible_explicit_targets --build_tests_only --target_pattern_file=/tmp/target_pattern_file +bazel test --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --skip_incompatible_explicit_targets --build_tests_only --target_pattern_file=/tmp/target_pattern_file # END BuildType.XLA_LINUX_X86_GPU_T4_GITHUB_ACTIONS # BEGIN BuildType.XLA_MACOS_ARM64_CPU_KOKORO df -h diff --git a/third_party/xla/build_tools/ci/run_clang_tidy.sh b/third_party/xla/build_tools/ci/run_clang_tidy.sh index 0d25ff8e5b4273..486dcbd8d6109b 100755 --- a/third_party/xla/build_tools/ci/run_clang_tidy.sh +++ b/third_party/xla/build_tools/ci/run_clang_tidy.sh @@ -57,7 +57,10 @@ get_targets_from_diff() { # Populate global TARGETS variable # We use readarray to handle multi-line output from bazel cquery - readarray -t TARGETS < <($BAZEL_CMD cquery --output=starlark --starlark:expr="target.label" --config="$CONFIG" "$QUERY") + CQUERY_OUT=$($BAZEL_CMD cquery --output=starlark --starlark:expr="target.label" --config="$CONFIG" "$QUERY") + if [ -n "$CQUERY_OUT" ]; then + readarray -t TARGETS <<< "$CQUERY_OUT" + fi if [ ${#TARGETS[@]} -eq 0 ]; then set +x diff --git a/third_party/xla/docs/_toc.yaml b/third_party/xla/docs/_toc.yaml index 313201eeb47f31..7ea289a0b8d40b 100644 --- a/third_party/xla/docs/_toc.yaml +++ b/third_party/xla/docs/_toc.yaml @@ -115,6 +115,8 @@ toc: path: /xla/build_from_source - title: Testing HLO passes path: /xla/test_hlo_passes + - title: "Avoiding #ifdefs" + path: /xla/avoiding_ifdefs - title: Copybara quirks path: /xla/copybara diff --git a/third_party/xla/docs/avoiding_ifdefs.md b/third_party/xla/docs/avoiding_ifdefs.md new file mode 100644 index 00000000000000..5c901aa4360e68 --- /dev/null +++ b/third_party/xla/docs/avoiding_ifdefs.md @@ -0,0 +1,370 @@ +# The Hitchhiker's Guide to fewer `#ifdef`s in XLA + +**TL;DR:** `#ifdef`s are bad for maintainability. Please try to avoid them. +Scroll down for details on how. + +[TOC] + +## Introduction + +Preprocessor conditionals (`#if`, `#ifdef`, `#ifndef`, etc.) are a convenient +way to make existing code work in a different environment than it was originally +intended for. Examples include: + +* Making existing CUDA-specific code work in a ROCm or SYCL environment. +* Adding new library features (for example, from a newer cuDNN or hipDNN + release) while maintaining support for older versions. + +However, they come with a high maintenance cost (more details in the next +section), especially when refactoring existing code. + +This guide explains the downsides of preprocessor conditionals and offers +alternatives for the most common use cases in a cookbook style. It is intended +to help contributors design maintainable, platform-portable changes from the +start and to serve as a shared reference during code reviews. + +## Motivation + +The C preprocessor is the first stage in a C++ compilation pipeline. +Preprocessor conditionals manipulate the stream of preprocessor tokens that is +read from the source code and header files and handed over to the C++ compiler. +The fact that it runs before the compiler and operates on the preprocessor token +level is the main reason why it results in a hard-to-maintain codebase: + +* **Code Complexity:** `#ifdef`s are convenient to use because they can be + inserted almost anywhere in the code. Since they operate on tokens, there + are almost no syntactic or semantic limits. Developers don't need to think + about a proper abstraction; they can just insert the conditionals right + where they are needed. Not being forced to design a proper abstraction makes + the code harder to read and harder to reason about. Having even a simple + abstraction like a free function allows easier testing of this function or + mocking its behavior—either now or later. + + It also renders a subset of compiler warnings meaningless because the + compiler will only ever see one evaluated token stream. A common example is + the need for the `[[maybe_unused]]` attribute on a function parameter that + is only used in one branch of a preprocessor conditional. + +* **Testability (Build):** Since preprocessor conditionals operate on token + streams, all the untaken branches neither need to be syntactically nor + semantically correct and are therefore not checked for correctness. This + severely hinders larger refactorings where a developer applies a + find-and-replace and relies on the compiler to tell them where to make + manual fixups. Incorrect changes in uncompiled preprocessor branches either + go in unnoticed or get detected late in the process (if CI happens to build + that particular preprocessor branch). The latter is also relevant for + day-to-day development (see the next point). + +* **Testability (Coverage):** Having more preprocessor conditionals means + having more build configurations to test, and the number of configurations + increases exponentially with every new condition introduced. It is already + impossible to test all of them. For example, XLA has quite a few + conditionals based on cuDNN version numbers that are not all exercised in + CI. One could argue that this is less relevant as long as we test the + configurations we care about—which is probably true. But unnecessarily + broken code leads to bug reports that someone needs to address, even if it + is just saying that a configuration is not supported. + + More importantly, potentially broken build configurations increase the + developer's iteration dead time. A change might work fine when running + `bazel test` locally, but CI might build a slightly different configuration + and fail. Fixes are often benign—like adding a `[[maybe_unused]]` + attribute—but this additional round trip costs the developer extra time on + every pull request. + + **Therefore, decreasing the number of preprocessor conditionals decreases + the number of build configurations, which in turn decreases the likelihood + of additional CI round trips on a PR.** + +* **Tooling:** C++ parsing is such a complex task that most tools nowadays + rely on a compiler frontend for parsing and semantic analysis and then + operate directly on the AST. Notable examples include `clang-tidy`, + `include-cleaner`, and language servers (`clangd`) that provide code + completion, navigation, and syntax highlighting in your IDE. Another class + of tools relies on compiler instrumentation, including sanitizers and code + coverage tools. + + All of these tools only see a single build configuration, so extensive use + of preprocessor conditionals hinders their usability. `include-cleaner`, for + example, suggests removing `#include`s that are only used in an unevaluated + preprocessor branch. Similarly, your IDE won't show syntax highlighting or + code navigation for ROCm code when it is configured for a CUDA or CPU build, + making it a pain to edit. + +## Mitigations + +### Category I - Skipping test cases in unit tests + +It is very common that a certain test case is only supported: + +* On a certain backend. +* With a certain GPU model. +* When library X is at least of version Y. + +Previously, it was common to skip tests using preprocessor conditionals: + +```cpp +TEST(Foo, Bar) { +#ifdef TENSORFLOW_USE_ROCM + GTEST_SKIP(); +#endif + // ... +} +``` + +#### Alternative: Ask StreamExecutor + +`StreamExecutor` is XLA's hardware abstraction layer, and its +[`stream_executor::DeviceDescription`](https://github.com/openxla/xla/blob/main/xla/stream_executor/device_description.h) +has the information needed to make the same decision at runtime: + +```cpp +TEST_F(FooTest, Bar) { + // `device_description()` is provided by `HloPjRtGpuTestBase`. Other test + // fixtures expose it via `executor->GetDeviceDescription()`. + const se::DeviceDescription& device = device_description(); + + // Skip based on the backend platform. + if (device.gpu_compute_capability().IsRocm()) { + GTEST_SKIP() << "Not supported on ROCm."; + } + + // Skip based on the GPU model / compute capability. + if (const auto* cc = + device.gpu_compute_capability().cuda_compute_capability(); + cc != nullptr && !cc->IsAtLeastHopper()) { + GTEST_SKIP() << "Requires Hopper or newer."; + } + + // Skip based on the runtime or library version. + if (device.runtime_version() < se::SemanticVersion{12, 2, 0}) { + GTEST_SKIP() << "Requires CUDA runtime >= 12.2."; + } + if (device.dnn_version() < se::SemanticVersion{9, 0, 0}) { + GTEST_SKIP() << "Requires cuDNN >= 9.0."; + } + // ... +} +``` + +`DeviceDescription` exposes both compute capabilities and structured version +numbers via +[`se::SemanticVersion`](https://github.com/openxla/xla/blob/main/xla/stream_executor/semantic_version.h): + +* **Backend & architecture:** `device.gpu_compute_capability().IsCuda()`, + `device.gpu_compute_capability().IsRocm()`, + `device.gpu_compute_capability().IsOneAPI()`, and accessors for + `CudaComputeCapability`, `RocmComputeCapability`, and + `OneAPIComputeCapability`. +* **Runtime & driver versions:** `device.runtime_version()`, + `device.driver_version()`, `device.kernel_mode_driver_version()`, + `device.compile_time_toolkit_version()`. +* **Library versions:** `device.dnn_version()`, `device.cub_version()`. + +If a specific version or hardware property is not yet available in +`DeviceDescription`, please add it (or ask the contributor to add it) to +`DeviceDescription` instead of falling back to `#ifdef`. + +#### Long term (author's opinion) + +In the long term, none of our higher-level tests should need to make decisions +based on backends, hardware variants, or runtime/driver versions. Instead, tests +should ask `StreamExecutor` whether a certain feature is available, and +`StreamExecutor` determines that based on all the necessary details. All this +logic should live in one place, though this does not yet have a concrete design. + +### Category II - Could be a runtime `if` statement + +Another class of preprocessor conditionals guards code that would also compile +fine without the conditional: + +```cpp +void Foo::Bar() { +#if TENSORFLOW_USE_ROCM + // Do something that would also compile in CUDA/CPU mode +#endif + // ... +} +``` + +#### Alternative: Use a runtime `if` + +Replace the preprocessor conditional with a runtime conditional: + +```cpp +void Foo::Bar() { + if (stream_executor_.GetDeviceDescription() + .gpu_compute_capability() + .IsRocm()) { + // Do something that compiles fine everywhere + } + // ... +} +``` + +Whether this should all be factored out into a separate function is up to the +reviewer's discretion. At this point it is just "normal" code and normal code +review practices apply. + +**This class of preprocessor conditionals appears most often when checking a +condition that is not readily available in StreamExecutor's +`DeviceDescription`. Everyone is encouraged to ask contributors to add the +relevant pieces of information to `DeviceDescription` instead of accepting the +preprocessor conditional.** + +### Category III - Requires access to the low-level runtime + +The most common (and also the best justified) class of preprocessor conditionals +are those that guard code that wouldn't compile otherwise. Most often this is +because it uses something from a low-level backend-specific header (such as the +CUDA or ROCm API): + +```cpp +#if TENSORFLOW_USE_ROCM +#include +#endif + +void Foo::Bar() { +#if TENSORFLOW_USE_ROCM + // Do something that would *NOT* compile in CUDA/CPU mode +#endif + // ... +} +``` + +The same applies to version guards on library headers (for example, +`#if CUDNN_VERSION >= 90000`) when the code references symbols that do not exist +in older versions of the header. + +#### Alternative I: Let StreamExecutor deal with it + +The first question to ask is whether the changed component should depend on +low-level runtime specifics at all. Ideally, only our hardware abstraction layer +`StreamExecutor` should. So if this code is not part of `StreamExecutor`, the +first option is to see if it can be pushed down into `xla/stream_executor/`. + +That means creating an API (a function, a class, or whatever makes sense) in +`StreamExecutor` that exists for *all* backends. This API can have different +implementations based on the backend. It might also allow querying whether this +particular feature is available, which enables the runtime-`if` solution +described above. + +Of course, that means `StreamExecutor` now has to deal with accessing the +low-level backend APIs. Alternative II describes how this can be modeled without +(too many) `#ifdef`s. + +#### Alternative II: Code in separate targets + +If the usage of backend-specific headers can't be pushed into `StreamExecutor` +or the change is inside `StreamExecutor`, the best option is to split the code +into separate compilation units. In our example, we have one build target with a +single C++ source file containing the ROCm code and another build target with a +single C++ source file containing the non-ROCm code—typically a stub +implementation that returns `absl::UnimplementedError` on all the defined +functions. + +This has the following advantages: + +* Splitting code across multiple files means there needs to be some layer of + abstraction in between—at the very least a free function. This encourages + contributors and reviewers alike to think about the right level of + abstraction, leading to higher-quality code in the long term. +* Both files only contain a single build configuration—giving you full syntax + highlighting and IDE tooling in both of them. +* Both build targets can be built separately in the same build graph. No need + to change your Bazel build flags and evict your build cache. +* Both build targets can have their own tests. For example, if a certain API + only exists for ROCm, we can have unit tests only testing the ROCm + implementation instead of higher-level tests that get skipped in all other + cases, which can speed up testing. (To be fair: in many cases there are good + reasons for having the higher-level tests as well.) + +A concrete design could look like this: + +```cpp +// feature.h +// Do NOT include any platform-specific headers (CUDA/ROCm) here. +#include "absl/status/status.h" + +absl::Status Foo(); +``` + +```cpp +// feature_rocm.cc +#include "feature.h" +#include "something/something/rocm.h" + +absl::Status Foo() { + // Do something ROCm-specific. +} +``` + +```cpp +// feature_stub.cc +#include "feature.h" +#include "absl/status/status.h" + +absl::Status Foo() { + return absl::UnimplementedError("This is a ROCm-only feature."); +} +``` + +```python +# BUILD +load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured") + +cc_library( + name = "feature_rocm", + srcs = [ + "feature.h", + "feature_rocm.cc", + ], + tags = ["manual"], # Exclude this from wildcard builds when ROCm is not enabled + deps = [ + "@com_google_absl//absl/status", + "@local_config_rocm//rocm:rocm_headers", + ], +) + +cc_library( + name = "feature_stub", + srcs = [ + "feature.h", + "feature_stub.cc", + ], + deps = ["@com_google_absl//absl/status"], +) + +cc_library( + name = "feature", + hdrs = ["feature.h"], + deps = if_rocm_is_configured( + [":feature_rocm"], + [":feature_stub"], + ) + [ + "@com_google_absl//absl/status", + ], +) +``` + +Note that `feature.h` is only in the `hdrs` attribute of the `:feature` target +(and in `srcs` of the two implementation targets). This ensures that consumers +including `feature.h` (and automated dependency tools) depend on `:feature` +rather than directly on `:feature_rocm` or `:feature_stub`. + +It is also important that `feature.h` does **not** include any platform-specific +headers (such as CUDA, ROCm, or cuDNN headers); only `.cc` files (and private +headers) should include those. The same build target pattern works for CUDA +(`if_cuda_is_configured` from `//xla/tsl/platform/default:cuda_build_defs.bzl`) +and SYCL. + +#### Long term (author's opinion) + +In the previous example, the targets `:feature_rocm` and `:feature_stub` define +the same symbol `Foo`, so they can't be linked into the same binary without +causing an ODR violation. Long term, XLA could have a plugin infrastructure +where multiple backends can be loaded at the same time. In that case, the design +above needs to be changed slightly so that `:feature_rocm` and `:feature_stub` +define symbols with different names and `:feature` has a dispatch function with +the original name—though it is currently unclear whether this will become a high +enough priority. diff --git a/third_party/xla/docs/contributing.md b/third_party/xla/docs/contributing.md index 66a713ab9480bd..229aea101003ab 100644 --- a/third_party/xla/docs/contributing.md +++ b/third_party/xla/docs/contributing.md @@ -174,6 +174,12 @@ There are two ways to do this: of the flag flip if problems arise, and affected users can temporarily set the flag themselves before a rollback is performed. +* *Preprocessor Conditionals (`#ifdef`)*: Avoid using `#if` / `#ifdef` for + backend-, hardware-, or version-specific logic whenever possible. Prefer + runtime checks via `StreamExecutor`'s `DeviceDescription` or separate build + targets. See [The Hitchhiker's Guide to Fewer `#ifdef`s in XLA](./avoiding_ifdefs.md) + for details and examples. + * When in doubt as to conventions within the code, it is always a good idea to examine pre-existing code and to try to follow the patterns already in place in OpenXLA. diff --git a/third_party/xla/tensorflow.bazelrc b/third_party/xla/tensorflow.bazelrc index 5fade5fffe284a..be9a77b0574489 100644 --- a/third_party/xla/tensorflow.bazelrc +++ b/third_party/xla/tensorflow.bazelrc @@ -142,6 +142,13 @@ common --incompatible_enforce_config_setting_visibility # TODO: also enable this flag after fixing the visibility violations # common --incompatible_config_setting_private_default_visibility +# external_include_paths feature forces passing external repositories headers +# via -isystem alters how Clang resolves their file paths. This feature with +# old versions of rules_cc triggering errors. +# Following lines can be removed after rules_cc will be updated to >= 0.2.22 +common --features=-external_include_paths +common --host_features=-external_include_paths + # Print a stacktrace when a test is killed build --test_env="GTEST_INSTALL_FAILURE_SIGNAL_HANDLER=1" @@ -304,8 +311,7 @@ common:rocm_ci_hermetic --dynamic_mode=fully common:rocm_ci_hermetic --@rules_ml_toolchain//common:enable_xla_test_global_symbol_version=True common:rocm_ci_hermetic --@rules_ml_toolchain//common:enable_hermetic_libraries=True common:rocm_ci_hermetic --config=rocm_clang_hermetic -common:rocm_ci_hermetic --repo_env=TF_ROCM_AMDGPU_TARGETS="gfx908,gfx90a" -common:rocm_ci_hermetic --repo_env=ROCM_DISTRO_VERSION="rocm_7.13.0_gfx90a" +common:rocm_ci_hermetic --repo_env=TF_ROCM_AMDGPU_TARGETS="gfx908,gfx90a,gfx950" common:rocm_ci_hermetic --repo_env=SYSROOT_DIST=linux_glibc_2_31 common:rocm_ci_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic diff --git a/third_party/xla/third_party/curl.BUILD b/third_party/xla/third_party/curl.BUILD index 917bf2b755e172..aa9adbf69a6e0f 100644 --- a/third_party/xla/third_party/curl.BUILD +++ b/third_party/xla/third_party/curl.BUILD @@ -48,17 +48,13 @@ cc_library( "lib/altsvc.h", "lib/amigaos.c", "lib/amigaos.h", + "lib/api.c", + "lib/api.h", "lib/arpa_telnet.h", - "lib/asyn.h", - "lib/asyn-ares.c", - "lib/asyn-base.c", - "lib/asyn-thrdd.c", "lib/bufq.c", "lib/bufq.h", "lib/bufref.c", "lib/bufref.h", - "lib/cf-dns.c", - "lib/cf-dns.h", "lib/cf-h1-proxy.c", "lib/cf-h1-proxy.h", "lib/cf-h2-proxy.c", @@ -79,8 +75,6 @@ cc_library( "lib/cfilters.h", "lib/config-mac.h", "lib/config-os400.h", - "lib/config-riscos.h", - "lib/config-win32.h", "lib/conncache.c", "lib/conncache.h", "lib/connect.c", @@ -96,6 +90,8 @@ cc_library( "lib/curl_addrinfo.c", "lib/curl_addrinfo.h", "lib/curl_ctype.h", + "lib/curl_ed25519.c", + "lib/curl_ed25519.h", "lib/curl_endian.c", "lib/curl_endian.h", "lib/curl_fnmatch.c", @@ -177,10 +173,6 @@ cc_library( "lib/cw-pause.h", "lib/dict.c", "lib/dict.h", - "lib/dnscache.c", - "lib/dnscache.h", - "lib/doh.c", - "lib/doh.h", "lib/dynhds.c", "lib/dynhds.h", "lib/easy.c", @@ -215,10 +207,6 @@ cc_library( "lib/headers.c", "lib/headers.h", "lib/hmac.c", - "lib/hostip.c", - "lib/hostip.h", - "lib/hostip4.c", - "lib/hostip6.c", "lib/hsts.c", "lib/hsts.h", "lib/http.c", @@ -233,14 +221,14 @@ cc_library( "lib/http_chunks.h", "lib/http_digest.c", "lib/http_digest.h", + "lib/http_httpsig.c", + "lib/http_httpsig.h", "lib/http_negotiate.c", "lib/http_negotiate.h", "lib/http_ntlm.c", "lib/http_ntlm.h", "lib/http_proxy.c", "lib/http_proxy.h", - "lib/httpsrr.c", - "lib/httpsrr.h", "lib/idn.c", "lib/idn.h", "lib/if2ip.c", @@ -340,6 +328,8 @@ cc_library( "lib/uint-bset.h", "lib/uint-hash.c", "lib/uint-hash.h", + "lib/uint-hashset.c", + "lib/uint-hashset.h", "lib/uint-spbset.c", "lib/uint-spbset.h", "lib/uint-table.c", @@ -364,6 +354,22 @@ cc_library( "lib/vauth/spnego_sspi.c", "lib/vauth/vauth.c", "lib/vauth/vauth.h", + "lib/vdns/asyn.h", + "lib/vdns/asyn-ares.c", + "lib/vdns/asyn-base.c", + "lib/vdns/asyn-thrdd.c", + "lib/vdns/cf-dns.c", + "lib/vdns/cf-dns.h", + "lib/vdns/dnscache.c", + "lib/vdns/dnscache.h", + "lib/vdns/doh.c", + "lib/vdns/doh.h", + "lib/vdns/hostip.c", + "lib/vdns/hostip.h", + "lib/vdns/hostip4.c", + "lib/vdns/hostip6.c", + "lib/vdns/httpsrr.c", + "lib/vdns/httpsrr.h", "lib/version.c", "lib/vquic/capsule.c", "lib/vquic/capsule.h", @@ -485,6 +491,7 @@ cc_library( "-DEFAULTLIB:advapi32.lib", "-DEFAULTLIB:bcrypt.lib", "-DEFAULTLIB:crypt32.lib", + "-DEFAULTLIB:iphlpapi.lib", "-DEFAULTLIB:secur32.lib", "-DEFAULTLIB:Normaliz.lib", ], @@ -514,7 +521,7 @@ CURL_BIN_WIN_COPTS = [ cc_binary( name = "curl_bin", srcs = [ - "lib/config-win32.h", + "include/curl_config.h", ] + [ "src/config2setopts.c", "src/config2setopts.h", @@ -601,6 +608,8 @@ cc_binary( "src/tool_writeout_json.h", "src/tool_xattr.c", "src/tool_xattr.h", + "src/toolx/tool_time.c", + "src/toolx/tool_time.h", "src/var.c", "src/var.h", ], @@ -633,7 +642,85 @@ genrule( "#endif", "", "#if defined(_WIN32)", - "# include \"lib/config-win32.h\"", + # Starting in curl 8.22.0, lib/config-win32.h has a preprocessor #error + # rejecting _MSC_VER > 1800 (restricting it to legacy MSVC 2010-2013 + # IDE builds). We inline the Windows configuration macros from + # lib/config-win32.h here instead of including it directly. + "# define HAVE_FCNTL_H 1", + "# define HAVE_IO_H 1", + "# define HAVE_LOCALE_H 1", + "# define HAVE_STDBOOL_H 1", + "# define HAVE_BOOL_T 1", + "# define HAVE_SYS_TYPES_H 1", + "# define HAVE_SYS_UTIME_H 1", + "# define HAVE_CLOSESOCKET 1", + "# define HAVE_FREEADDRINFO 1", + "# define HAVE_GETADDRINFO 1", + "# define HAVE_GETADDRINFO_THREADSAFE 1", + "# define HAVE_GETHOSTNAME 1", + "# define HAVE_GETPEERNAME 1", + "# define HAVE_GETSOCKNAME 1", + "# define HAVE_IOCTLSOCKET 1", + "# define HAVE_IOCTLSOCKET_FIONBIO 1", + "# define HAVE_SETLOCALE 1", + "# define HAVE_SOCKET 1", + "# define HAVE_UTIME 1", + "# define HAVE_RECV 1", + "# define RECV_TYPE_ARG1 SOCKET", + "# define RECV_TYPE_ARG2 char *", + "# define RECV_TYPE_ARG3 int", + "# define RECV_TYPE_ARG4 int", + "# define RECV_TYPE_RETV int", + "# define HAVE_SEND 1", + "# define SEND_TYPE_ARG1 SOCKET", + "# define SEND_TYPE_ARG2 char *", + "# define SEND_TYPE_ARG3 int", + "# define SEND_TYPE_ARG4 int", + "# define SEND_TYPE_RETV int", + "# define HAVE_SIGNAL 1", + "# define SIZEOF_INT 4", + "# define SIZEOF_LONG 4", + "# ifdef _WIN64", + "# define SIZEOF_SIZE_T 8", + "# ifndef _SSIZE_T_DEFINED", + "# define _SSIZE_T_DEFINED", + "# define ssize_t __int64", + "# endif", + "# else", + "# define SIZEOF_SIZE_T 4", + "# ifndef _SSIZE_T_DEFINED", + "# define _SSIZE_T_DEFINED", + "# define ssize_t int", + "# endif", + "# endif", + "# define SIZEOF_CURL_OFF_T 8", + "# ifndef _USE_32BIT_TIME_T", + "# define SIZEOF_TIME_T 8", + "# else", + "# define SIZEOF_TIME_T 4", + "# endif", + "# define SIZEOF_OFF_T 4", + "# define HAVE_STRUCT_SOCKADDR_STORAGE 1", + "# define HAVE_STRUCT_TIMEVAL 1", + "# define HAVE_SOCKADDR_IN6_SIN6_SCOPE_ID 1", + "# if !defined(USE_RESOLV_THREADED) && !defined(USE_SYNC_DNS)", + "# define USE_RESOLV_THREADED 1", + "# endif", + "# define HAVE_LDAP_SSL 1", + "# define USE_WIN32_LDAP 1", + "# define USE_WIN32_CRYPTO 1", + "# define USE_UNIX_SOCKETS 1", + "# ifndef CURL_OS", + "# if defined(_M_IX86) || defined(__i386__)", + "# define CURL_OS \"i386-pc-win32\"", + "# elif defined(_M_X64) || defined(__x86_64__)", + "# define CURL_OS \"x86_64-pc-win32\"", + "# elif defined(_M_ARM64) || defined(__aarch64__)", + "# define CURL_OS \"aarch64-pc-win32\"", + "# else", + "# define CURL_OS \"unknown-pc-win32\"", + "# endif", + "# endif", "# define BUILDING_LIBCURL 1", "# define CURL_DISABLE_CRYPTO_AUTH 1", "# define CURL_DISABLE_DICT 1", diff --git a/third_party/xla/third_party/gpus/rocm/rocm_redist.bzl b/third_party/xla/third_party/gpus/rocm/rocm_redist.bzl index 79a49e65260880..f221414111fe20 100644 --- a/third_party/xla/third_party/gpus/rocm/rocm_redist.bzl +++ b/third_party/xla/third_party/gpus/rocm/rocm_redist.bzl @@ -16,31 +16,11 @@ """Rocm distribution list, used as a hermetic rocm dependency.""" rocm_redist = { - "rocm_7.13.0_gfx94X": struct( + "rocm_latest_multiarch": struct( packages = [ { - "url": "https://repo.amd.com/rocm/tarball/therock-dist-linux-gfx94X-dcgpu-7.13.0.tar.gz", - "sha256": "db5543de096fb175ff2ece19dacc28b2a3201df48b38051cc505e508d84e35ab", - }, - ], - required_softlinks = [], - rocm_root = "", - ), - "rocm_7.13.0_gfx908": struct( - packages = [ - { - "url": "https://repo.amd.com/rocm/tarball/therock-dist-linux-gfx908-7.13.0.tar.gz", - "sha256": "5d84753a8d8895ff2f6137a2a922ee8f36ce9c2e01b60a99d3ee776a683bfc34", - }, - ], - required_softlinks = [], - rocm_root = "", - ), - "rocm_7.13.0_gfx90a": struct( - packages = [ - { - "url": "https://repo.amd.com/rocm/tarball/therock-dist-linux-gfx90a-7.13.0.tar.gz", - "sha256": "b2d3c49ef936b3b24b10a25bae3e60df7ccc9c5134095a080bbc721d5062b4c7", + "url": "https://stable.repo.amd.com/rocm/core/tarball/therock-dist-linux-multiarch-10.0.0.tar.gz", + "sha256": "1c5e807875d26a2470ecc7323daa5b5b9009208a55c3290ac255a909cde15fc6", }, ], required_softlinks = [], diff --git a/third_party/xla/third_party/gpus/rocm_configure.bzl b/third_party/xla/third_party/gpus/rocm_configure.bzl index 4acd64749cbf3d..918e2e417464cc 100644 --- a/third_party/xla/third_party/gpus/rocm_configure.bzl +++ b/third_party/xla/third_party/gpus/rocm_configure.bzl @@ -49,7 +49,7 @@ _ROCM_DISTRO_LINKS = "ROCM_DISTRO_LINKS" _TMPDIR = "TMPDIR" # Default hermetic ROCm redistributable version -_DEFAULT_ROCM_DISTRO_VERSION = "rocm_7.13.0_gfx94X" +_DEFAULT_ROCM_DISTRO_VERSION = "rocm_latest_multiarch" _TF_ROCM_RBE_DOCKER_IMAGE = "TF_ROCM_RBE_DOCKER_IMAGE" _TF_ROCM_RBE_POOL = "TF_ROCM_RBE_POOL" diff --git a/third_party/xla/third_party/llvm/generated.patch b/third_party/xla/third_party/llvm/generated.patch index a9462de1d2ba0b..634ed2ef4a9817 100644 --- a/third_party/xla/third_party/llvm/generated.patch +++ b/third_party/xla/third_party/llvm/generated.patch @@ -1,189 +1,664 @@ Auto generated patch. Do not edit or delete it, even if empty. -diff -ruN --strip-trailing-cr a/clang/lib/CodeGen/CGCall.cpp b/clang/lib/CodeGen/CGCall.cpp ---- a/clang/lib/CodeGen/CGCall.cpp -+++ b/clang/lib/CodeGen/CGCall.cpp -@@ -2804,17 +2804,10 @@ - AddAttributesFromFunctionProtoType( - getContext(), FuncAttrs, Fn->getType()->getAs()); - if (AttrOnCallSite && Fn->isReplaceableGlobalAllocationFunction()) { -- // A sane operator new returns a non-aliasing pointer and does not -- // read or write accessible memory. -+ // A sane operator new returns a non-aliasing pointer. - if (getCodeGenOpts().AssumeSaneOperatorNew && -- Fn->getDeclName().isAnyOperatorNew()) { -+ Fn->getDeclName().isAnyOperatorNew()) - RetAttrs.addAttribute(llvm::Attribute::NoAlias); -- // FIXME: inaccessiblemem could cause issues if LTO makes the -- // previously inaccessible memory accessible after linking. -- FuncAttrs.addMemoryAttr( -- llvm::MemoryEffects::inaccessibleOrErrnoMemOnly( -- llvm::ModRefInfo::ModRef, llvm::ModRefInfo::Mod)); -- } - } - const CXXMethodDecl *MD = dyn_cast(Fn); - const bool IsVirtualCall = MD && MD->isVirtual(); -diff -ruN --strip-trailing-cr a/clang/lib/CodeGen/CGExprCXX.cpp b/clang/lib/CodeGen/CGExprCXX.cpp ---- a/clang/lib/CodeGen/CGExprCXX.cpp -+++ b/clang/lib/CodeGen/CGExprCXX.cpp -@@ -1391,9 +1391,19 @@ - /// - /// We model such elidable calls with the 'builtin' attribute. - llvm::Function *Fn = dyn_cast(CalleePtr); -- if (CalleeDecl->isReplaceableGlobalAllocationFunction() && Fn && -- Fn->hasFnAttribute(llvm::Attribute::NoBuiltin)) { -- CallOrInvoke->addFnAttr(llvm::Attribute::Builtin); -+ if (CalleeDecl->isReplaceableGlobalAllocationFunction() && Fn) { -+ if (Fn->hasFnAttribute(llvm::Attribute::NoBuiltin)) -+ CallOrInvoke->addFnAttr(llvm::Attribute::Builtin); +diff -ruN --strip-trailing-cr a/clang/include/clang/AST/TypeBase.h b/clang/include/clang/AST/TypeBase.h +--- a/clang/include/clang/AST/TypeBase.h ++++ b/clang/include/clang/AST/TypeBase.h +@@ -2270,7 +2270,7 @@ + /// increments towards the beginning. + /// Positive non-zero number represents the index + 1. + /// Zero means this is not substituted from an expansion. +- unsigned PackIndex : 15; ++ unsigned PackIndex : 16; + }; + + class SubstPackTypeBitfields { +diff -ruN --strip-trailing-cr a/clang/test/SemaCXX/make_integer_seq.cpp b/clang/test/SemaCXX/make_integer_seq.cpp +--- a/clang/test/SemaCXX/make_integer_seq.cpp ++++ b/clang/test/SemaCXX/make_integer_seq.cpp +@@ -50,3 +50,8 @@ + + __make_integer_seq<__make_integer_seq, int, 10> PR28494; // expected-note{{different template parameters}} + // expected-error@make_integer_seq.cpp:* {{template argument for template template parameter must be a class template or type alias template}} + -+ // A sane operator new does not read or write accessible memory. -+ if (CGF.CGM.getCodeGenOpts().AssumeSaneOperatorNew && -+ CalleeDecl->getDeclName().isAnyOperatorNew()) { -+ // FIXME: inaccessiblemem could cause issues if LTO makes the -+ // previously inaccessible memory accessible after linking. -+ CallOrInvoke->setMemoryEffects( -+ llvm::MemoryEffects::inaccessibleOrErrnoMemOnly( -+ llvm::ModRefInfo::ModRef, llvm::ModRefInfo::Mod)); -+ } ++// The largest type pack whose pack indices fit in SubstTemplateTypeParmType. ++template using Expand = void(Ts...); ++template struct ExpandSeq { using type = Expand; }; ++using LargePack = __make_integer_seq::type; +diff -ruN --strip-trailing-cr a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp +--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp ++++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp +@@ -13901,47 +13901,12 @@ + return SDValue(); + } + +-/// Check if Mask defines a known constant set of enabled lanes, where only the +-/// first N lanes are enabled. N is returned if so. +-static uint64_t calculateConstantLowMaskLanes(SDValue Mask) { +- // We expect masks for masked load/store to be i1 predicates. +- if (Mask.getValueType().getScalarSizeInBits() != 1) +- return 0; +- +- if (Mask.getOpcode() == ISD::BUILD_VECTOR) { +- unsigned NumOnes = 0; +- auto Op = Mask->op_begin(); +- for (; Op != Mask->op_end(); ++Op) { +- if (!isOneConstant(*Op)) +- break; +- ++NumOnes; +- } +- +- if (NumOnes == 0) +- return 0; +- +- for (; Op != Mask->op_end(); ++Op) +- if (!isNullConstant(*Op)) +- return 0; +- return NumOnes; +- } +- +- if (Mask.getOpcode() == ISD::GET_ACTIVE_LANE_MASK && +- isNullConstant(Mask.getOperand(0)) && +- isa(Mask.getOperand(1)) && +- Mask.getOperand(1).getValueType().getSizeInBits() <= 64) +- return Mask.getConstantOperandVal(1); +- +- return 0; +-} +- + SDValue DAGCombiner::visitMSTORE(SDNode *N) { + MaskedStoreSDNode *MST = cast(N); + SDValue Mask = MST->getMask(); + SDValue Chain = MST->getChain(); + SDValue Value = MST->getValue(); + SDValue Ptr = MST->getBasePtr(); +- EVT VT = Value.getValueType(); + + // Zap masked stores with a zero mask. + if (ISD::isConstantSplatVectorAllZeros(Mask.getNode())) +@@ -13964,46 +13929,21 @@ + } } - return RV; -diff -ruN --strip-trailing-cr a/clang/test/CodeGenCXX/new_hot_cold.cpp b/clang/test/CodeGenCXX/new_hot_cold.cpp ---- a/clang/test/CodeGenCXX/new_hot_cold.cpp -+++ b/clang/test/CodeGenCXX/new_hot_cold.cpp -@@ -124,7 +124,7 @@ +- // If this is a masked store with an all ones mask, we can use a unmasked +- // store. ++ // If this is a masked load with an all ones mask, we can use a unmasked load. + // FIXME: Can we do this for indexed, compressing, or truncating stores? +- if (MST->isUnindexed() && !MST->isCompressingStore() && +- !MST->isTruncatingStore()) { +- if (ISD::isConstantSplatVectorAllOnes(Mask.getNode())) +- return DAG.getStore(MST->getChain(), SDLoc(N), MST->getValue(), +- MST->getBasePtr(), MST->getPointerInfo(), +- MST->getBaseAlign(), MST->getMemOperand()->getFlags(), +- MST->getAAInfo()); +- +- // Convert a masked_store with constant getactivelanemask input mask to a +- // standard store. +- if (uint64_t Lanes = calculateConstantLowMaskLanes(Mask)) { +- if (Lanes < VT.getVectorMinNumElements() && isPowerOf2_32(Lanes)) { +- EVT SubVT = EVT::getVectorVT(*DAG.getContext(), +- VT.getVectorElementType(), Lanes); +- unsigned IsFast = 0; +- if ((!LegalTypes || isTypeLegal(SubVT)) && +- TLI.allowsMemoryAccess(*DAG.getContext(), DAG.getDataLayout(), +- SubVT, MST->getAddressSpace(), +- MST->getBaseAlign(), +- MST->getMemOperand()->getFlags(), &IsFast) && +- IsFast) { +- SDLoc DL(N); +- SDValue Ext = DAG.getExtractSubvector(DL, SubVT, MST->getValue(), 0); +- return DAG.getStore(MST->getChain(), DL, Ext, MST->getBasePtr(), +- MST->getPointerInfo(), MST->getBaseAlign(), +- MST->getMemOperand()->getFlags(), +- MST->getAAInfo()); +- } +- } +- } +- } ++ if (ISD::isConstantSplatVectorAllOnes(Mask.getNode()) && MST->isUnindexed() && ++ !MST->isCompressingStore() && !MST->isTruncatingStore()) ++ return DAG.getStore(MST->getChain(), SDLoc(N), MST->getValue(), ++ MST->getBasePtr(), MST->getPointerInfo(), ++ MST->getBaseAlign(), MST->getMemOperand()->getFlags(), ++ MST->getAAInfo()); - // CHECK-DAG: attributes [[ATTR_NOBUILTIN]] = { nobuiltin allocsize(0) {{.*}} } - // CHECK-DAG: attributes [[ATTR_NOBUILTIN_NOTHROW]] = { nobuiltin nounwind allocsize(0) {{.*}} } --// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_CALL]] = { allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } -+// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_CALL]] = { allocsize(0) } - // CHECK-DAG: attributes [[ATTR_BUILTIN_CALL]] = { builtin allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } --// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_NOTHROW_CALL]] = { nounwind allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } -+// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_NOTHROW_CALL]] = { nounwind allocsize(0) } - // CHECK-DAG: attributes [[ATTR_BUILTIN_NOTHROW_CALL]] = { builtin nounwind allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } -diff -ruN --strip-trailing-cr a/clang/test/CodeGenCXX/operator-new.cpp b/clang/test/CodeGenCXX/operator-new.cpp ---- a/clang/test/CodeGenCXX/operator-new.cpp -+++ b/clang/test/CodeGenCXX/operator-new.cpp -@@ -27,5 +27,14 @@ - } + // Try transforming N to an indexed store. + if (CombineToPreIndexedLoadStore(N) || CombineToPostIndexedLoadStore(N)) + return SDValue(N, 0); + +- if (MST->isTruncatingStore() && MST->isUnindexed() && VT.isInteger() && ++ if (MST->isTruncatingStore() && MST->isUnindexed() && ++ Value.getValueType().isInteger() && + (!isa(Value) || + !cast(Value)->isOpaque())) { + APInt TruncDemandedBits = +diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll b/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll +--- a/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll ++++ b/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll +@@ -139,7 +139,14 @@ + define void @store_low2_v4f64(ptr %p, <4 x double> %a) vscale_range(2, 2) { + ; CHECK-LABEL: store_low2_v4f64: + ; CHECK: // %bb.0: +-; CHECK-NEXT: str q0, [x0] ++; CHECK-NEXT: // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1 ++; CHECK-NEXT: ptrue p0.d, vl2 ++; CHECK-NEXT: ptrue p1.h, vl2 ++; CHECK-NEXT: // kill: def $q0 killed $q0 killed $z0_z1 def $z0_z1 ++; CHECK-NEXT: splice z0.d, p0, { z0.d, z1.d } ++; CHECK-NEXT: punpklo p0.h, p1.b ++; CHECK-NEXT: punpklo p0.h, p0.b ++; CHECK-NEXT: st1d { z0.d }, p0, [x0] + ; CHECK-NEXT: ret + %m = tail call <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32 0, i32 2) + tail call void @llvm.masked.store(<4 x double> %a, ptr align 1 %p, <4 x i1> %m) +@@ -149,7 +156,13 @@ + define void @store_low2_v4i32(ptr %p, <4 x i32> %a) vscale_range(1, 1) { + ; CHECK-LABEL: store_low2_v4i32: + ; CHECK: // %bb.0: +-; CHECK-NEXT: str d0, [x0] ++; CHECK-NEXT: ptrue p0.h, vl2 ++; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 ++; CHECK-NEXT: mov z1.h, p0/z, #-1 // =0xffffffffffffffff ++; CHECK-NEXT: ptrue p0.s ++; CHECK-NEXT: sshll v1.4s, v1.4h, #0 ++; CHECK-NEXT: cmpne p1.s, p0/z, z1.s, #0 ++; CHECK-NEXT: st1w { z0.s }, p1, [x0] + ; CHECK-NEXT: ret + %m = tail call <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32 0, i32 2) + tail call void @llvm.masked.store(<4 x i32> %a, ptr align 1 %p, <4 x i1> %m) +@@ -159,7 +172,12 @@ + define void @store_low2c_v4i32(ptr %p, <4 x i32> %a) vscale_range(1, 1) { + ; CHECK-LABEL: store_low2c_v4i32: + ; CHECK: // %bb.0: +-; CHECK-NEXT: str d0, [x0] ++; CHECK-NEXT: adrp x8, .LCPI12_0 ++; CHECK-NEXT: ptrue p0.s ++; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 ++; CHECK-NEXT: ldr q1, [x8, :lo12:.LCPI12_0] ++; CHECK-NEXT: cmpne p1.s, p0/z, z1.s, #0 ++; CHECK-NEXT: st1w { z0.s }, p1, [x0] + ; CHECK-NEXT: ret + tail call void @llvm.masked.store(<4 x i32> %a, ptr align 1 %p, <4 x i1> ) + ret void +@@ -168,7 +186,11 @@ + define void @store_low2_nxv4f64(ptr %p, %a) { + ; CHECK-LABEL: store_low2_nxv4f64: + ; CHECK: // %bb.0: +-; CHECK-NEXT: str q0, [x0] ++; CHECK-NEXT: ptrue p0.s, vl2 ++; CHECK-NEXT: punpkhi p1.h, p0.b ++; CHECK-NEXT: punpklo p0.h, p0.b ++; CHECK-NEXT: st1d { z1.d }, p1, [x0, #1, mul vl] ++; CHECK-NEXT: st1d { z0.d }, p0, [x0] + ; CHECK-NEXT: ret + %m = tail call @llvm.get.active.lane.mask.nxv4i1.i32(i32 0, i32 2) + tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) +@@ -203,10 +225,13 @@ + define void @store_low2_v8i16(ptr %p, <8 x i16> %a) vscale_range(2, 2) { + ; CHECK-LABEL: store_low2_v8i16: + ; CHECK: // %bb.0: ++; CHECK-NEXT: ptrue p0.b, vl2 + ; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 +-; CHECK-NEXT: ptrue p0.s, vl2 +-; CHECK-NEXT: uunpklo z0.s, z0.h +-; CHECK-NEXT: st1h { z0.s }, p0, [x0] ++; CHECK-NEXT: mov z1.b, p0/z, #-1 // =0xffffffffffffffff ++; CHECK-NEXT: ptrue p0.h, vl8 ++; CHECK-NEXT: sshll v1.8h, v1.8b, #0 ++; CHECK-NEXT: cmpne p1.h, p0/z, z1.h, #0 ++; CHECK-NEXT: st1h { z0.h }, p1, [x0] + ; CHECK-NEXT: ret + %m = tail call <8 x i1> @llvm.get.active.lane.mask.v8i1.i32(i32 0, i32 2) + tail call void @llvm.masked.store(<8 x i16> %a, ptr align 1 %p, <8 x i1> %m) +@@ -216,9 +241,8 @@ + define void @store_low2_nxv8i16(ptr %p, %a) { + ; CHECK-LABEL: store_low2_nxv8i16: + ; CHECK: // %bb.0: +-; CHECK-NEXT: mov z1.h, z0.h[1] +-; CHECK-NEXT: str h0, [x0] +-; CHECK-NEXT: str h1, [x0, #2] ++; CHECK-NEXT: ptrue p0.h, vl2 ++; CHECK-NEXT: st1h { z0.h }, p0, [x0] + ; CHECK-NEXT: ret + %m = tail call @llvm.get.active.lane.mask.nxv8i1.i32(i32 0, i32 2) + tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) +@@ -239,7 +263,8 @@ + define void @store_low4_nxv8i16(ptr %p, %a) { + ; CHECK-LABEL: store_low4_nxv8i16: + ; CHECK: // %bb.0: +-; CHECK-NEXT: str d0, [x0] ++; CHECK-NEXT: ptrue p0.h, vl4 ++; CHECK-NEXT: st1h { z0.h }, p0, [x0] + ; CHECK-NEXT: ret + %m = tail call @llvm.get.active.lane.mask.nxv8i1.i32(i32 0, i32 4) + tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) +diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll b/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll +--- a/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll ++++ b/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll +@@ -6,43 +6,40 @@ + ; CHECK-LABEL: f0: + ; CHECK: // %bb.0: // %b0 + ; CHECK-NEXT: { +-; CHECK-NEXT: r0 = add(r1,r0) ++; CHECK-NEXT: [[RADD:[r0-9]+]] = add([[R1:[r0-9]+]],[[R0:[r0-9]+]]) + ; CHECK-NEXT: } + ; CHECK-NEXT: { +-; CHECK-NEXT: allocframe(r29,#128):raw ++; CHECK-NEXT: [[RLCPI:[r0-9]+]] = ##.LCPI0_0 + ; CHECK-NEXT: } + ; CHECK-NEXT: { +-; CHECK-NEXT: r2 = add(r0,#136) ++; CHECK-NEXT: [[RNEG:[r0-9]+]] = #-1 + ; CHECK-NEXT: } + ; CHECK-NEXT: { +-; CHECK-NEXT: r29 = and(r29,#-128) ++; CHECK-NEXT: [[RADD136:[r0-9]+]] = add([[RADD]],#136) + ; CHECK-NEXT: } + ; CHECK-NEXT: { +-; CHECK-NEXT: v1 = vmem(r0+#1) ++; CHECK-NEXT: [[V1:[v0-9]+]] = vmem([[RLCPI]]+#0) + ; CHECK-NEXT: } + ; CHECK-NEXT: { +-; CHECK-NEXT: v0 = vmemu(r2+#0) ++; CHECK-NEXT: [[RADD128:[r0-9]+]] = add([[RADD]],#128) + ; CHECK-NEXT: } + ; CHECK-NEXT: { +-; CHECK-NEXT: r4 = add(r29,#0) ++; CHECK-NEXT: [[Q0:[q0-9]+]] = vand([[V1]],[[RNEG]]) + ; CHECK-NEXT: } + ; CHECK-NEXT: { +-; CHECK-NEXT: v0.w = vadd(v1.w,v0.w) ++; CHECK-NEXT: [[V2:[v0-9]+]] = vmemu([[RADD136]]+#0) + ; CHECK-NEXT: } + ; CHECK-NEXT: { +-; CHECK-NEXT: vmem(r4+#0) = v0 ++; CHECK-NEXT: [[V3:[v0-9]+]] = vmem([[RADD]]+#1) + ; CHECK-NEXT: } + ; CHECK-NEXT: { +-; CHECK-NEXT: r2 = memw(r4+#0) ++; CHECK-NEXT: [[V0:[v0-9]+]].w = vadd([[V3]].w,[[V2]].w) + ; CHECK-NEXT: } + ; CHECK-NEXT: { +-; CHECK-NEXT: r3 = memw(r4+#4) ++; CHECK-NEXT: if ([[Q0]]) vmem([[RADD128]]+#0) = [[V0]] + ; CHECK-NEXT: } + ; CHECK-NEXT: { +-; CHECK-NEXT: memd(r0+#128) = r3:2 +-; CHECK-NEXT: } +-; CHECK-NEXT: { +-; CHECK-NEXT: r31:30 = dealloc_return(r30):raw ++; CHECK-NEXT: jumpr r31 + ; CHECK-NEXT: } + b0: + %v0 = add i32 %a1, 128 +diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll b/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll +--- a/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll ++++ b/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll +@@ -29,12 +29,20 @@ + define void @kmovkr_2() { + ; AVX512-LABEL: kmovkr_2: + ; AVX512: # %bb.0: # %alloca_21 +-; AVX512-NEXT: movq $0, 0 # encoding: [0x48,0xc7,0x04,0x25,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00] ++; AVX512-NEXT: vxorps %xmm0, %xmm0, %xmm0 # encoding: [0xc5,0xf8,0x57,0xc0] ++; AVX512-NEXT: movw $3, %ax # encoding: [0x66,0xb8,0x03,0x00] ++; AVX512-NEXT: kmovw %eax, %k1 # EVEX TO VEX Compression encoding: [0xc5,0xf8,0x92,0xc8] ++; AVX512-NEXT: vmovups %zmm0, 0 {%k1} # encoding: [0x62,0xf1,0x7c,0x49,0x11,0x04,0x25,0x00,0x00,0x00,0x00] ++; AVX512-NEXT: vzeroupper # encoding: [0xc5,0xf8,0x77] + ; AVX512-NEXT: retq # encoding: [0xc3] + ; + ; AVX512BW-LABEL: kmovkr_2: + ; AVX512BW: # %bb.0: # %alloca_21 +-; AVX512BW-NEXT: movq $0, 0 # encoding: [0x48,0xc7,0x04,0x25,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00] ++; AVX512BW-NEXT: vxorps %xmm0, %xmm0, %xmm0 # encoding: [0xc5,0xf8,0x57,0xc0] ++; AVX512BW-NEXT: movw $3, %ax # encoding: [0x66,0xb8,0x03,0x00] ++; AVX512BW-NEXT: kmovd %eax, %k1 # EVEX TO VEX Compression encoding: [0xc5,0xfb,0x92,0xc8] ++; AVX512BW-NEXT: vmovups %zmm0, 0 {%k1} # encoding: [0x62,0xf1,0x7c,0x49,0x11,0x04,0x25,0x00,0x00,0x00,0x00] ++; AVX512BW-NEXT: vzeroupper # encoding: [0xc5,0xf8,0x77] + ; AVX512BW-NEXT: retq # encoding: [0xc3] + alloca_21: + call void @llvm.masked.store.v4f32.p0(<4 x float> zeroinitializer, ptr null, i32 1, <4 x i1> ) +diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/masked-store-illegal-type.ll b/llvm/test/CodeGen/X86/masked-store-illegal-type.ll +--- a/llvm/test/CodeGen/X86/masked-store-illegal-type.ll ++++ b/llvm/test/CodeGen/X86/masked-store-illegal-type.ll +@@ -1,21 +0,0 @@ +-; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py +-; RUN: llc < %s -mtriple=x86_64-linux-gnu -mattr=+avx,+avx2,+avx512bw,+avx512cd,+avx512dq,+avx512f,+avx512vl | FileCheck %s +- +-define fastcc void @render_event(<12 x double> %0) { +-; CHECK-LABEL: render_event: +-; CHECK: # %bb.0: +-; CHECK-NEXT: vmovlhps {{.*#+}} xmm6 = xmm6[0],xmm7[0] +-; CHECK-NEXT: vmovlhps {{.*#+}} xmm4 = xmm4[0],xmm5[0] +-; CHECK-NEXT: vinsertf128 $1, %xmm6, %ymm4, %ymm4 +-; CHECK-NEXT: vmovlhps {{.*#+}} xmm2 = xmm2[0],xmm3[0] +-; CHECK-NEXT: vmovlhps {{.*#+}} xmm0 = xmm0[0],xmm1[0] +-; CHECK-NEXT: vinsertf128 $1, %xmm2, %ymm0, %ymm0 +-; CHECK-NEXT: vinsertf64x4 $1, %ymm4, %zmm0, %zmm0 +-; CHECK-NEXT: vmovsd {{.*#+}} xmm1 = mem[0],zero +-; CHECK-NEXT: vmovsd %xmm1, 64 +-; CHECK-NEXT: vmovups %zmm0, 0 +-; CHECK-NEXT: vzeroupper +-; CHECK-NEXT: retq +- call void @llvm.masked.store.v12f64.p0(<12 x double> %0, ptr null, <12 x i1> ) +- ret void +-} +diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/pr179489.ll b/llvm/test/CodeGen/X86/pr179489.ll +--- a/llvm/test/CodeGen/X86/pr179489.ll ++++ b/llvm/test/CodeGen/X86/pr179489.ll +@@ -9,12 +9,14 @@ + ; X86-LABEL: foo: + ; X86: # %bb.0: # %entry + ; X86-NEXT: movl {{[0-9]+}}(%esp), %eax +-; X86-NEXT: vpextrb $1, %xmm0, (%eax) ++; X86-NEXT: vmovw %xmm0, %ecx ++; X86-NEXT: movb %ch, (%eax) + ; X86-NEXT: retl + ; + ; X64-LABEL: foo: + ; X64: # %bb.0: # %entry +-; X64-NEXT: vpextrb $1, %xmm0, (%rdi) ++; X64-NEXT: vmovw %xmm0, %eax ++; X64-NEXT: movb %ah, (%rdi) + ; X64-NEXT: retq + entry: + %i3 = shufflevector <8 x i16> %arg, <8 x i16> , <8 x i32> +diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/pr46532.ll b/llvm/test/CodeGen/X86/pr46532.ll +--- a/llvm/test/CodeGen/X86/pr46532.ll ++++ b/llvm/test/CodeGen/X86/pr46532.ll +@@ -7,6 +7,8 @@ + ; CHECK-NEXT: movq (%rax), %rax + ; CHECK-NEXT: vxorps %xmm0, %xmm0, %xmm0 + ; CHECK-NEXT: vmovaps %xmm0, 32(%rax) ++; CHECK-NEXT: vmovsd {{.*#+}} xmm0 = [4294967295,4294967295,0,0] ++; CHECK-NEXT: vmaskmovps %ymm0, %ymm0, (%rax) + while.1.body.preheader: + %0 = load ptr, ptr undef, align 8, !invariant.load !0, !dereferenceable !1, !align !2 + %1 = getelementptr inbounds i8, ptr %0, i64 32 +diff -ruN --strip-trailing-cr a/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td b/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td +--- a/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td ++++ b/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td +@@ -1,71 +0,0 @@ +-// RUN: split-file %s %t +-// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/valid.td \ +-// RUN: | FileCheck %t/valid.td +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-bug.td 2>&1 \ +-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=user-sgpr-init16-bug --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-priority.td 2>&1 \ +-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=required-export-priority --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-timing.td 2>&1 \ +-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=gfx125x-lowest-rate-wmma --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-banks.td 2>&1 \ +-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=ldsbankcount32 --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-lds-size.td 2>&1 \ +-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=addressablelocalmemorysize65536 --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/capability.td 2>&1 \ +-// RUN: | FileCheck %t/capability.td -DFILE=%t/capability.td --implicit-check-not="error:" +- +-//--- common.td +-include "AMDGPU.td" +- +-def EmptyGPU : AMDGPUProcessorModel<"gfx999", NoSchedModel, [], [9, 9, 9]>; +- +-//--- valid.td +-include "common.td" +-// Check the real generic targets, including gfx11's combined workarounds, +-// gfx12-5's fixup and WMMA timing, and gfx9-4's LDS size and bank count. +-// CHECK: GK_GFX9_4_GENERIC +-// CHECK: GK_GFX11_GENERIC +-// CHECK: GK_GFX12_5_GENERIC +- +-//--- missing-bug.td +-include "common.td" +-defvar CheckedFeature = FeatureUserSGPRInit16Bug; +-include "missing-generic.td" +- +-//--- missing-priority.td +-include "common.td" +-defvar CheckedFeature = FeatureRequiredExportPriority; +-include "missing-generic.td" +- +-//--- missing-timing.td +-include "common.td" +-defvar CheckedFeature = FeatureGFX125xLowestRateWMMA; +-include "missing-generic.td" +- +-//--- missing-banks.td +-include "common.td" +-defvar CheckedFeature = FeatureLDSBankCount32; +-include "missing-generic.td" +- +-//--- missing-lds-size.td +-include "common.td" +-defvar CheckedFeature = FeatureAddressableLocalMemorySize65536; +-include "missing-generic.td" +- +-//--- missing-generic.td +-// Classified backend-only features still need support from a covered GPU. +-// Included files use the native path separator after the include directory. +-// CHECK: [[DIR]]{{[/\\]}}missing-generic.td:[[#@LINE+1]]:1: error: generic target 'gfx9-9-generic' exposes feature '[[FEATURE]]' not supported by any covered GPU +-def : AMDGPUProcessorModel<"gfx9-9-generic", NoSchedModel, [CheckedFeature], [9, 9, 0]> { +- let CoveredGPUs = [EmptyGPU]; +-} +- +-//--- capability.td +-include "common.td" +-def WMMAGPU : AMDGPUProcessorModel<"gfx998", NoSchedModel, +- [FeatureWMMA256bInsts], [9, 9, 8]>; +-// Instruction support keeps the default all-covered-GPUs rule. +-// CHECK: [[FILE]]:[[#@LINE+1]]:1: error: generic target 'gfx9-9-generic' exposes feature 'wmma-256b-insts' not supported by covered GPU 'gfx999' +-def : AMDGPUProcessorModel<"gfx9-9-generic", NoSchedModel, [FeatureWMMA256bInsts], [9, 9, 0]> { +- let CoveredGPUs = [WMMAGPU, EmptyGPU]; +-} +diff -ruN --strip-trailing-cr a/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td b/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td +--- a/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td ++++ b/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td +@@ -1,33 +1,66 @@ + // RUN: split-file %s %t +-// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/valid.td \ ++// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/valid.td \ + // RUN: | FileCheck %t/valid.td +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-any.td 2>&1 \ ++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-any.td 2>&1 \ + // RUN: | FileCheck %t/missing-any.td -DFILE=%t/missing-any.td --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-all.td 2>&1 \ ++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-all.td 2>&1 \ + // RUN: | FileCheck %t/missing-all.td -DFILE=%t/missing-all.td --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/implied-all.td 2>&1 \ ++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/implied-all.td 2>&1 \ + // RUN: | FileCheck %t/implied-all.td -DFILE=%t/implied-all.td --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-private.td 2>&1 \ ++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-private.td 2>&1 \ + // RUN: | FileCheck %t/missing-private.td -DFILE=%t/missing-private.td --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/numeric-mismatch.td 2>&1 \ ++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/implied-private.td 2>&1 \ ++// RUN: | FileCheck %t/implied-private.td -DFILE=%t/implied-private.td --implicit-check-not="error:" ++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/numeric-mismatch.td 2>&1 \ + // RUN: | FileCheck %t/numeric-mismatch.td -DFILE=%t/numeric-mismatch.td --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/empty-members.td 2>&1 \ ++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/empty-members.td 2>&1 \ + // RUN: | FileCheck %t/empty-members.td -DFILE=%t/empty-members.td --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/self-member.td 2>&1 \ ++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/self-member.td 2>&1 \ + // RUN: | FileCheck %t/self-member.td -DFILE=%t/self-member.td --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/generic-member.td 2>&1 \ ++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/generic-member.td 2>&1 \ + // RUN: | FileCheck %t/generic-member.td -DFILE=%t/generic-member.td --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/pseudo-member.td 2>&1 \ ++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/pseudo-member.td 2>&1 \ + // RUN: | FileCheck %t/pseudo-member.td -DFILE=%t/pseudo-member.td --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/unregistered-member.td 2>&1 \ ++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/unregistered-member.td 2>&1 \ + // RUN: | FileCheck %t/unregistered-member.td -DFILE=%t/unregistered-member.td --implicit-check-not="error:" +-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/r600-member.td 2>&1 \ ++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/r600-member.td 2>&1 \ + // RUN: | FileCheck %t/r600-member.td -DFILE=%t/r600-member.td --implicit-check-not="error:" - // ALL: declare noundef nonnull ptr @_Znaj( + //--- common.td +-include "llvm/Target/Target.td" ++// Exercise the validation rules with small synthetic targets. The real targets ++// are validated when AMDGPUTargetParserDef.inc is generated during the build. ++// Stand-ins for the Target.td classes used by AMDGPUTargetParser.td, limited to ++// the fields the backend reads. ++class InlineBehavior; ++def InlineIgnore : InlineBehavior; ++def InlineDefault : InlineBehavior; ++class SubtargetFeature i = [], ++ InlineBehavior ib = InlineDefault> { ++ string Name = n; ++ string FieldName = f; ++ string Value = v; ++ string Desc = d; ++ list Implies = i; ++ InlineBehavior InlineBehavior = ib; ++} ++ ++class SchedMachineModel; ++def NoSchedModel : SchedMachineModel; ++class ProcessorItineraries; ++def NoItineraries : ProcessorItineraries; + -+void *f3(unsigned long N) { -+ // SANE: call noalias noundef nonnull ptr @_Znwj(i32 noundef {{.*}}) [[ATTR2:#[0-9]+]] -+ // SANENOT: call noundef nonnull ptr @_Znwj(i32 noundef {{.*}}) [[ATTR2:#[0-9]+]] -+ return ::operator new(N); ++class Processor f> { ++ string Name = n; ++ ProcessorItineraries ProcItin = pi; ++ list Features = f; ++} ++class ProcessorModel f> ++ : Processor { ++ SchedMachineModel SchedModel = m; +} + - // SANE: attributes [[ATTR]] = { builtin allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } -+// SANE: attributes [[ATTR2]] = { allocsize(0) } - // SANENOT: attributes [[ATTR]] = { builtin allocsize(0) } -+// SANENOT: attributes [[ATTR2]] = { allocsize(0) } -diff -ruN --strip-trailing-cr a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp ---- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp -+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp -@@ -16145,18 +16145,21 @@ + include "AMDGPUTargetParser.td" +-def MyTarget : Target; - InstructionCost VecCost = VectorCost(CommonCost); - // Check if the current node must be resized, if the parent node is not -- // resized. -+ // resized. Nodes without a parent (roots of the splat gather -+ // subtrees) are resized back to their original type when the reusing -+ // gather nodes are emitted. - if (It != MinBWs.end() && !UnaryInstruction::isCast(E->getOpcode()) && - E->Idx != 0 && - (E->getOpcode() != Instruction::Load || E->UserTreeIndex)) { - const EdgeInfo &EI = E->UserTreeIndex; -- if (!EI.UserTE->hasState() || -+ if (!EI.UserTE || !EI.UserTE->hasState() || - EI.UserTE->getOpcode() != Instruction::Select || - EI.EdgeIdx != 0) { - auto UserBWIt = MinBWs.find(EI.UserTE); - Type *UserScalarTy = -- (EI.UserTE->isGather() || -- EI.UserTE->State == TreeEntry::SplitVectorize) -+ !EI.UserTE ? OrigScalarTy -+ : (EI.UserTE->isGather() || -+ EI.UserTE->State == TreeEntry::SplitVectorize) - ? EI.UserTE->Scalars.front()->getType() - : EI.UserTE->getOperand(EI.EdgeIdx).front()->getType(); - if (UserBWIt != MinBWs.end()) -diff -ruN --strip-trailing-cr a/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll b/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll ---- a/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll -+++ b/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll -@@ -0,0 +1,75 @@ -+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py -+; RUN: opt -passes=slp-vectorizer -S -mtriple=aarch64-linux-gnu < %s | FileCheck %s -+ -+define double @test() { -+; CHECK-LABEL: @test( -+; CHECK-NEXT: bbl: -+; CHECK-NEXT: [[SELECT:%.*]] = select i1 false, i64 0, i64 0 -+; CHECK-NEXT: [[AND:%.*]] = and i64 [[SELECT]], 0 -+; CHECK-NEXT: [[ICMP:%.*]] = icmp eq i64 [[AND]], 0 -+; CHECK-NEXT: [[TMP0:%.*]] = insertelement <2 x i1> poison, i1 true, i64 0 -+; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <2 x i1> [[TMP0]], <2 x i1> poison, <2 x i32> zeroinitializer -+; CHECK-NEXT: [[TMP2:%.*]] = select <2 x i1> [[TMP1]], <2 x double> zeroinitializer, <2 x double> zeroinitializer -+; CHECK-NEXT: [[TMP3:%.*]] = bitcast <2 x double> [[TMP2]] to <2 x i64> -+; CHECK-NEXT: [[TMP4:%.*]] = select <2 x i1> [[TMP1]], <2 x double> zeroinitializer, <2 x double> zeroinitializer -+; CHECK-NEXT: [[TMP5:%.*]] = bitcast <2 x double> [[TMP4]] to <2 x i64> -+; CHECK-NEXT: [[TMP6:%.*]] = insertelement <2 x i1> poison, i1 [[ICMP]], i64 0 -+; CHECK-NEXT: [[TMP7:%.*]] = shufflevector <2 x i1> [[TMP6]], <2 x i1> poison, <2 x i32> zeroinitializer -+; CHECK-NEXT: [[TMP8:%.*]] = select <2 x i1> [[TMP7]], <2 x double> zeroinitializer, <2 x double> zeroinitializer -+; CHECK-NEXT: [[TMP9:%.*]] = bitcast <2 x double> [[TMP8]] to <2 x i64> -+; CHECK-NEXT: [[TMP10:%.*]] = insertelement <2 x i64> poison, i64 [[SELECT]], i64 0 -+; CHECK-NEXT: [[TMP11:%.*]] = shufflevector <2 x i64> [[TMP10]], <2 x i64> poison, <2 x i32> zeroinitializer -+; CHECK-NEXT: [[TMP12:%.*]] = xor <2 x i64> [[TMP11]], [[TMP9]] -+; CHECK-NEXT: [[TMP13:%.*]] = bitcast <2 x i64> [[TMP12]] to <2 x double> -+; CHECK-NEXT: [[TMP14:%.*]] = insertelement <2 x i64> poison, i64 0, i64 0 -+; CHECK-NEXT: [[TMP15:%.*]] = shufflevector <2 x i64> [[TMP14]], <2 x i64> poison, <2 x i32> zeroinitializer -+; CHECK-NEXT: [[TMP16:%.*]] = xor <2 x i64> [[TMP15]], [[TMP3]] -+; CHECK-NEXT: [[TMP17:%.*]] = bitcast <2 x i64> [[TMP16]] to <2 x double> -+; CHECK-NEXT: [[TMP18:%.*]] = xor <2 x i64> [[TMP15]], [[TMP5]] -+; CHECK-NEXT: [[TMP19:%.*]] = fmul <2 x double> [[TMP13]], [[TMP17]] -+; CHECK-NEXT: [[TMP20:%.*]] = bitcast <2 x i64> [[TMP18]] to <2 x double> -+; CHECK-NEXT: [[TMP21:%.*]] = fmul <2 x double> [[TMP19]], [[TMP20]] -+; CHECK-NEXT: [[TMP22:%.*]] = extractelement <2 x double> [[TMP21]], i64 0 -+; CHECK-NEXT: [[TMP23:%.*]] = extractelement <2 x double> [[TMP21]], i64 1 -+; CHECK-NEXT: [[FADD:%.*]] = fadd double [[TMP22]], [[TMP23]] -+; CHECK-NEXT: ret double [[FADD]] -+; -+bbl: -+ %select = select i1 false, i64 0, i64 0 -+ %and = and i64 %select, 0 -+ %icmp = icmp eq i64 %and, 0 -+ %select1 = select i1 %icmp, double 0.000000e+00, double 0.000000e+00 -+ %bitcast = bitcast double %select1 to i64 -+ %xor = xor i64 %select, %bitcast -+ %bitcast2 = bitcast i64 %xor to double -+ %select4 = select i1 false, i64 0, i64 0 -+ %icmp5 = icmp eq i64 %select4, 0 -+ %select6 = select i1 %icmp5, double 0.000000e+00, double 0.000000e+00 -+ %bitcast7 = bitcast double %select6 to i64 -+ %xor8 = xor i64 %select4, %bitcast7 -+ %bitcast9 = bitcast i64 %xor8 to double -+ %select10 = select i1 %icmp5, double 0.000000e+00, double 0.000000e+00 -+ %bitcast11 = bitcast double %select10 to i64 -+ %select12 = select i1 false, i64 0, i64 0 -+ %icmp13 = icmp eq i64 %select12, 0 -+ %select14 = select i1 %icmp13, double 0.000000e+00, double 0.000000e+00 -+ %bitcast15 = bitcast double %select14 to i64 -+ %xor16 = xor i64 %select12, %bitcast15 -+ %select17 = select i1 %icmp13, double 0.000000e+00, double 0.000000e+00 -+ %bitcast18 = bitcast double %select17 to i64 -+ %select3 = select i1 %icmp, double 0.000000e+00, double 0.000000e+00 -+ %bitcast19 = bitcast double %select3 to i64 -+ %xor20 = xor i64 %select, %bitcast19 -+ %bitcast21 = bitcast i64 %xor20 to double -+ %xor22 = xor i64 %select4, %bitcast11 -+ %bitcast23 = bitcast i64 %xor22 to double -+ %fmul = fmul double %bitcast21, %bitcast23 -+ %xor24 = xor i64 %select12, %bitcast18 -+ %bitcast25 = bitcast i64 %xor24 to double -+ %fmul26 = fmul double %fmul, %bitcast25 -+ %fmul27 = fmul double %bitcast2, %bitcast9 -+ %bitcast28 = bitcast i64 %xor16 to double -+ %fmul29 = fmul double %fmul27, %bitcast28 -+ %fadd = fadd double %fmul26, %fmul29 -+ ret double %fadd + def FeatureAll : SubtargetFeature<"all", "HasAll", "true", "">; + def FeatureAnyA : AMDGPUGenericAnyFeature<"any-a", "HasAnyA", "true", "">; +@@ -37,7 +70,13 @@ + def FeatureAnySmall : NumericFeature<4>; + def FeatureAnyB : NumericFeature<8>; + def FeatureAnyLarge : NumericFeature<16>; +-def FeaturePrivate : AMDGPUGenericAnyFeature<"private", "HasPrivate", "true", "">; ++// Classification inherited through a subclass also applies to backend-only ++// features, which must participate in validation without being emitted. ++class PrivateFeature : AMDGPUGenericAnyFeature< ++ "private", "HasPrivate", "true", "">; ++def FeaturePrivate : PrivateFeature; ++def FeatureImpliesPrivate : SubtargetFeature< ++ "implies-private", "HasImpliesPrivate", "true", "", [FeaturePrivate]>; + def FeatureImpliesA : SubtargetFeature< + "implies-a", "HasImpliesA", "true", "", [FeatureAnyA]>; + def FeatureImpliesB : SubtargetFeature< +@@ -55,7 +94,7 @@ + def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, + [FeatureAll, FeatureAnyA, FeatureAnySmall], [9, 0, 0]>; + def GFX901 : AMDGPUProcessorModel<"gfx901", NoSchedModel, +- [FeatureAll, FeatureImpliesB, FeaturePrivate], [9, 0, 1]>; ++ [FeatureAll, FeatureImpliesB, FeatureImpliesPrivate], [9, 0, 1]>; + // Different members supply different marked features. Follow Implies on both + // the generic and the members, and check marked backend-only features too. + def : AMDGPUProcessorModel<"gfx9-generic", NoSchedModel, +@@ -103,6 +142,15 @@ + let CoveredGPUs = [GFX900]; + } + ++//--- implied-private.td ++include "common.td" ++def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, [], [9, 0, 0]>; ++// An unclassified backend-only parent must not hide a classified dependency. ++// CHECK: [[FILE]]:[[#@LINE+1]]:1: error: generic target 'gfx9-generic' exposes feature 'private' not supported by any covered GPU ++def : AMDGPUProcessorModel<"gfx9-generic", NoSchedModel, [FeatureImpliesPrivate], [9, 0, 0]> { ++ let CoveredGPUs = [GFX900]; +} ++ + //--- numeric-mismatch.td + include "common.td" + def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, [FeatureAnySmall], [9, 0, 0]>; +diff -ruN --strip-trailing-cr a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp +--- a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp ++++ b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp +@@ -14,11 +14,12 @@ + using namespace mlir; + using namespace mlir::tosa; + +-// Building this ~5,000-line generated initializer map under MemorySanitizer +-// leads to extreme compile times (>3.5 minutes) during InstCombine and greedy +-// register allocation due to shadow/origin instrumentation overhead. Disable +-// optimization under MSan to keep build times manageable. +-#if LLVM_MEMORY_SANITIZER_BUILD ++// Building this ~5,000-line generated initializer map under HWAddressSanitizer ++// or MemorySanitizer leads to extreme compile times (>3.5 minutes) during ++// InstCombine and greedy register allocation due to instrumentation ++// overhead. Disable optimization under MSan/HWAsan to keep build times ++// manageable. ++#if LLVM_MEMORY_SANITIZER_BUILD || LLVM_HWADDRESS_SANITIZER_BUILD + __attribute__((optnone)) + #else + LLVM_ATTRIBUTE_MINSIZE +diff -ruN --strip-trailing-cr a/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel b/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel +--- a/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel ++++ b/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel +@@ -66,7 +66,7 @@ + + bool_flag( + name = "llvm_enable_lzma", +- build_setting_default = False, ++ build_setting_default = True, + ) + + config_setting( +@@ -74,11 +74,7 @@ + flag_values = {":llvm_enable_lzma": "true"}, + ) + +-# Unlike zlib and zstd there is no liblzma module in the Bazel registry, so this +-# links the system library rather than building one from source. That is why it +-# is a plain cc_library with linkopts instead of a cc_library_wrapper, and why +-# it defaults to off. +-cc_library( ++cc_library_wrapper( + name = "lzma", + defines = select({ + ":llvm_lzma_enabled": [ +@@ -88,9 +84,9 @@ + "LLVM_ENABLE_LZMA=0", + ], + }), +- linkopts = select({ ++ deps = select({ + ":llvm_lzma_enabled": [ +- "-llzma", ++ "@llvm_xz//:lzma", + ], + "//conditions:default": [], + }), +diff -ruN --strip-trailing-cr a/utils/bazel/MODULE.bazel b/utils/bazel/MODULE.bazel +--- a/utils/bazel/MODULE.bazel ++++ b/utils/bazel/MODULE.bazel +@@ -29,6 +29,7 @@ + bazel_dep(name = "nanobind", version = "2.12.0.bcr.1") + bazel_dep(name = "libpfm", version = "4.13.0", repo_name = "pfm") + bazel_dep(name = "vulkan_headers", version = "1.4.349") ++bazel_dep(name = "xz", version = "5.4.5.bcr.9", repo_name = "llvm_xz") + + bazel_dep(name = "emsdk", version = "6.0.2", dev_dependency = True) + bazel_dep(name = "llvm", version = "0.8.20", dev_dependency = True) +diff -ruN --strip-trailing-cr a/utils/bazel/MODULE.bazel.lock b/utils/bazel/MODULE.bazel.lock +--- a/utils/bazel/MODULE.bazel.lock ++++ b/utils/bazel/MODULE.bazel.lock +@@ -55,7 +55,8 @@ + "https://bcr.bazel.build/modules/bazel_lib/3.0.0/MODULE.bazel": "22b70b80ac89ad3f3772526cd9feee2fa412c2b01933fea7ed13238a448d370d", + "https://bcr.bazel.build/modules/bazel_lib/3.1.0/MODULE.bazel": "6809765c14e3c766a9b9286c7b0ec56ed87a73326e48fe01749f0c0fdcfe3287", + "https://bcr.bazel.build/modules/bazel_lib/3.2.2/MODULE.bazel": "e2c890c8a515d6bca9c66d47718aa9e44b458fde64ec7204b8030bf2d349058c", +- "https://bcr.bazel.build/modules/bazel_lib/3.2.2/source.json": "9e84e115c20e14652c5c21401ae85ff4daa8702e265b5c0b3bf89353f17aa212", ++ "https://bcr.bazel.build/modules/bazel_lib/3.7.2/MODULE.bazel": "cbc920fff00a55b2fd663d40850560a6fb17fd5ef3a9f3d6cbf0d32908dd8a40", ++ "https://bcr.bazel.build/modules/bazel_lib/3.7.2/source.json": "7a9c2b06a113c4962d329bb68befdb90d0447788a19aabaa32b476a0458e7d28", + "https://bcr.bazel.build/modules/bazel_skylib/1.0.3/MODULE.bazel": "bcb0fd896384802d1ad283b4e4eb4d718eebd8cb820b0a2c3a347fb971afd9d8", + "https://bcr.bazel.build/modules/bazel_skylib/1.1.1/MODULE.bazel": "1add3e7d93ff2e6998f9e118022c84d163917d912f5afafb3058e3d2f1545b5e", + "https://bcr.bazel.build/modules/bazel_skylib/1.2.0/MODULE.bazel": "44fe84260e454ed94ad326352a698422dbe372b21a1ac9f3eab76eb531223686", +@@ -313,6 +314,8 @@ + "https://bcr.bazel.build/modules/windows_support/0.4.1/source.json": "bdb72e6717b00135e0a4a8512b3f6b626df0baea0bde74bf7312c7484f791287", + "https://bcr.bazel.build/modules/with_cfg.bzl/0.12.0/MODULE.bazel": "b573395fe63aef4299ba095173e2f62ccfee5ad9bbf7acaa95dba73af9fc2b38", + "https://bcr.bazel.build/modules/with_cfg.bzl/0.12.0/source.json": "3f3fbaeafecaf629877ad152a2c9def21f8d330d91aa94c5dc75bbb98c10b8b8", ++ "https://bcr.bazel.build/modules/xz/5.4.5.bcr.9/MODULE.bazel": "ff7cf8a5334a96b7b9ffcbe73418d45929ce125efebac39c1ee7d14bd2ae7a9b", ++ "https://bcr.bazel.build/modules/xz/5.4.5.bcr.9/source.json": "abf61f3c1630eab7b81722f710060b6f31e7c160511bcbfaad74bb4385883051", + "https://bcr.bazel.build/modules/zlib-ng/2.3.3.bcr.2/MODULE.bazel": "86130924115bf8dfc4aba286fb49a2ec8d32430a21771db9ccf5381acf0d10cc", + "https://bcr.bazel.build/modules/zlib-ng/2.3.3.bcr.2/source.json": "7aa4112eb2fbb99757e97cf0b7a4e2bdb3464b2b3cfb351bc807ffcb34bd349f", + "https://bcr.bazel.build/modules/zlib-ng/2.3.3/MODULE.bazel": "0afa781a5f354b4fd811a5c9a086e111daf50e09d8e98ab1cf7c46eeacc33aac", diff --git a/third_party/xla/third_party/llvm/toolchains.patch b/third_party/xla/third_party/llvm/toolchains.patch index 1170c7610ab573..a8c15b85d9235a 100644 --- a/third_party/xla/third_party/llvm/toolchains.patch +++ b/third_party/xla/third_party/llvm/toolchains.patch @@ -13,13 +13,13 @@ # limitations under the License. # ============================================================================== diff --git a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel +index b683f86ca386..70819c622402 100644 --- a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel -@@ -57,6 +57,36 @@ - # This one is needed for building and vendoring out lldb from off tree. +@@ -59,6 +59,36 @@ exports_files([ "utils/lldbDataFormatters.py", ]) -+ + +config_setting( + name = "macos_arm64", + constraint_values = [ @@ -49,9 +49,20 @@ diff --git a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel b/utils/bazel/llv + "@platforms//cpu:x86_64", + ], +) - ++ config_setting( name = "darwin_arm64", + constraint_values = [ +@@ -533,7 +563,7 @@ cc_library( + # We unconditionally depend on the custom LLVM lzma wrapper. This only + # provides the configuration define unless lzma is enabled, in which + # case it also links the system liblzma. +- "//third-party:lzma", ++ # "//third-party:lzma", + "//libc:shared_math_headers_for_apfloat", + ], + ) + diff --git a/utils/bazel/llvm-project-overlay/llvm/config.bzl b/utils/bazel/llvm-project-overlay/llvm/config.bzl --- a/utils/bazel/llvm-project-overlay/llvm/config.bzl +++ b/utils/bazel/llvm-project-overlay/llvm/config.bzl diff --git a/third_party/xla/third_party/llvm/workspace.bzl b/third_party/xla/third_party/llvm/workspace.bzl index d929da090db313..88052c28512ca5 100644 --- a/third_party/xla/third_party/llvm/workspace.bzl +++ b/third_party/xla/third_party/llvm/workspace.bzl @@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive") def repo(name): """Imports LLVM.""" - LLVM_COMMIT = "a22291abb68744cb0b70978eb506a7c5592a3309" - LLVM_SHA256 = "cf19b82f0dfbef1e62afe520ba4ee6de4f4823524d7804490e88d815f822c8e9" + LLVM_COMMIT = "e56c2cefc3e7978de4a2d799fa3abec77d75af21" + LLVM_SHA256 = "affcad31c9455e95faca9d4de38e4a49b73c23bbf48b405d264ff10a6bfd3baf" tf_http_archive( name = name, diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch index 95a2a325ed502a..0d8bf613295b57 100644 --- a/third_party/xla/third_party/shardy/temporary.patch +++ b/third_party/xla/third_party/shardy/temporary.patch @@ -1,194 +1,1215 @@ +diff --git a/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline.mlir b/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline.mlir +index 5697a287..6ef5c695 100644 +--- a/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline.mlir ++++ b/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline.mlir +@@ -1,4 +1,4 @@ +-// RUN: mpmd_opt %s -mpmd-infer-mesh-pipeline='infer-transfers=true infer-cross-mesh-reductions=True' 2>&1 | FileCheck %s ++// RUN: mpmd_opt %s -mpmd-infer-mesh-pipeline='infer-transfers=true infer-cross-mesh-reductions=true' 2>&1 | FileCheck %s + + !mesh_1_tensor_ui32 = !mpmd.mesh_tensor<"m1", tensor> + !mesh_1_tensor_1_ui32 = !mpmd.mesh_tensor<"m1", tensor<1xui32>> +diff --git a/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline_send_recv.mlir b/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline_send_recv.mlir +index d914bd27..54d0ca3a 100644 +--- a/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline_send_recv.mlir ++++ b/shardy/dialect/mpmd/transforms/import/test/infer_mesh_pipeline_send_recv.mlir +@@ -1,4 +1,4 @@ +-// RUN: mpmd_opt %s -mpmd-infer-mesh-pipeline='infer-transfers=true infer-cross-mesh-reductions=True' 2>&1 | FileCheck %s ++// RUN: mpmd_opt %s -mpmd-infer-mesh-pipeline='infer-transfers=true infer-cross-mesh-reductions=true' 2>&1 | FileCheck %s + + // Test that stablehlo.send and stablehlo.create_token ops (which produce + // !stablehlo.token types) are correctly partitioned into fragments by the diff --git a/third_party/llvm/generated.patch b/third_party/llvm/generated.patch -index 509398da..a9462de1 100644 +index a9462de1..634ed2ef 100644 --- a/third_party/llvm/generated.patch +++ b/third_party/llvm/generated.patch -@@ -1 +1,189 @@ +@@ -1,189 +1,664 @@ Auto generated patch. Do not edit or delete it, even if empty. -+diff -ruN --strip-trailing-cr a/clang/lib/CodeGen/CGCall.cpp b/clang/lib/CodeGen/CGCall.cpp -+--- a/clang/lib/CodeGen/CGCall.cpp -++++ b/clang/lib/CodeGen/CGCall.cpp -+@@ -2804,17 +2804,10 @@ -+ AddAttributesFromFunctionProtoType( -+ getContext(), FuncAttrs, Fn->getType()->getAs()); -+ if (AttrOnCallSite && Fn->isReplaceableGlobalAllocationFunction()) { -+- // A sane operator new returns a non-aliasing pointer and does not -+- // read or write accessible memory. -++ // A sane operator new returns a non-aliasing pointer. -+ if (getCodeGenOpts().AssumeSaneOperatorNew && -+- Fn->getDeclName().isAnyOperatorNew()) { -++ Fn->getDeclName().isAnyOperatorNew()) -+ RetAttrs.addAttribute(llvm::Attribute::NoAlias); -+- // FIXME: inaccessiblemem could cause issues if LTO makes the -+- // previously inaccessible memory accessible after linking. -+- FuncAttrs.addMemoryAttr( -+- llvm::MemoryEffects::inaccessibleOrErrnoMemOnly( -+- llvm::ModRefInfo::ModRef, llvm::ModRefInfo::Mod)); -+- } -+ } -+ const CXXMethodDecl *MD = dyn_cast(Fn); -+ const bool IsVirtualCall = MD && MD->isVirtual(); -+diff -ruN --strip-trailing-cr a/clang/lib/CodeGen/CGExprCXX.cpp b/clang/lib/CodeGen/CGExprCXX.cpp -+--- a/clang/lib/CodeGen/CGExprCXX.cpp -++++ b/clang/lib/CodeGen/CGExprCXX.cpp -+@@ -1391,9 +1391,19 @@ -+ /// -+ /// We model such elidable calls with the 'builtin' attribute. -+ llvm::Function *Fn = dyn_cast(CalleePtr); -+- if (CalleeDecl->isReplaceableGlobalAllocationFunction() && Fn && -+- Fn->hasFnAttribute(llvm::Attribute::NoBuiltin)) { -+- CallOrInvoke->addFnAttr(llvm::Attribute::Builtin); -++ if (CalleeDecl->isReplaceableGlobalAllocationFunction() && Fn) { -++ if (Fn->hasFnAttribute(llvm::Attribute::NoBuiltin)) -++ CallOrInvoke->addFnAttr(llvm::Attribute::Builtin); -++ -++ // A sane operator new does not read or write accessible memory. -++ if (CGF.CGM.getCodeGenOpts().AssumeSaneOperatorNew && -++ CalleeDecl->getDeclName().isAnyOperatorNew()) { -++ // FIXME: inaccessiblemem could cause issues if LTO makes the -++ // previously inaccessible memory accessible after linking. -++ CallOrInvoke->setMemoryEffects( -++ llvm::MemoryEffects::inaccessibleOrErrnoMemOnly( -++ llvm::ModRefInfo::ModRef, llvm::ModRefInfo::Mod)); -++ } -+ } -+ -+ return RV; -+diff -ruN --strip-trailing-cr a/clang/test/CodeGenCXX/new_hot_cold.cpp b/clang/test/CodeGenCXX/new_hot_cold.cpp -+--- a/clang/test/CodeGenCXX/new_hot_cold.cpp -++++ b/clang/test/CodeGenCXX/new_hot_cold.cpp -+@@ -124,7 +124,7 @@ -+ -+ // CHECK-DAG: attributes [[ATTR_NOBUILTIN]] = { nobuiltin allocsize(0) {{.*}} } -+ // CHECK-DAG: attributes [[ATTR_NOBUILTIN_NOTHROW]] = { nobuiltin nounwind allocsize(0) {{.*}} } -+-// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_CALL]] = { allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } -++// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_CALL]] = { allocsize(0) } -+ // CHECK-DAG: attributes [[ATTR_BUILTIN_CALL]] = { builtin allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } -+-// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_NOTHROW_CALL]] = { nounwind allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } -++// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_NOTHROW_CALL]] = { nounwind allocsize(0) } -+ // CHECK-DAG: attributes [[ATTR_BUILTIN_NOTHROW_CALL]] = { builtin nounwind allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } -+diff -ruN --strip-trailing-cr a/clang/test/CodeGenCXX/operator-new.cpp b/clang/test/CodeGenCXX/operator-new.cpp -+--- a/clang/test/CodeGenCXX/operator-new.cpp -++++ b/clang/test/CodeGenCXX/operator-new.cpp -+@@ -27,5 +27,14 @@ +-diff -ruN --strip-trailing-cr a/clang/lib/CodeGen/CGCall.cpp b/clang/lib/CodeGen/CGCall.cpp +---- a/clang/lib/CodeGen/CGCall.cpp +-+++ b/clang/lib/CodeGen/CGCall.cpp +-@@ -2804,17 +2804,10 @@ +- AddAttributesFromFunctionProtoType( +- getContext(), FuncAttrs, Fn->getType()->getAs()); +- if (AttrOnCallSite && Fn->isReplaceableGlobalAllocationFunction()) { +-- // A sane operator new returns a non-aliasing pointer and does not +-- // read or write accessible memory. +-+ // A sane operator new returns a non-aliasing pointer. +- if (getCodeGenOpts().AssumeSaneOperatorNew && +-- Fn->getDeclName().isAnyOperatorNew()) { +-+ Fn->getDeclName().isAnyOperatorNew()) +- RetAttrs.addAttribute(llvm::Attribute::NoAlias); +-- // FIXME: inaccessiblemem could cause issues if LTO makes the +-- // previously inaccessible memory accessible after linking. +-- FuncAttrs.addMemoryAttr( +-- llvm::MemoryEffects::inaccessibleOrErrnoMemOnly( +-- llvm::ModRefInfo::ModRef, llvm::ModRefInfo::Mod)); +-- } +- } +- const CXXMethodDecl *MD = dyn_cast(Fn); +- const bool IsVirtualCall = MD && MD->isVirtual(); +-diff -ruN --strip-trailing-cr a/clang/lib/CodeGen/CGExprCXX.cpp b/clang/lib/CodeGen/CGExprCXX.cpp +---- a/clang/lib/CodeGen/CGExprCXX.cpp +-+++ b/clang/lib/CodeGen/CGExprCXX.cpp +-@@ -1391,9 +1391,19 @@ +- /// +- /// We model such elidable calls with the 'builtin' attribute. +- llvm::Function *Fn = dyn_cast(CalleePtr); +-- if (CalleeDecl->isReplaceableGlobalAllocationFunction() && Fn && +-- Fn->hasFnAttribute(llvm::Attribute::NoBuiltin)) { +-- CallOrInvoke->addFnAttr(llvm::Attribute::Builtin); +-+ if (CalleeDecl->isReplaceableGlobalAllocationFunction() && Fn) { +-+ if (Fn->hasFnAttribute(llvm::Attribute::NoBuiltin)) +-+ CallOrInvoke->addFnAttr(llvm::Attribute::Builtin); ++diff -ruN --strip-trailing-cr a/clang/include/clang/AST/TypeBase.h b/clang/include/clang/AST/TypeBase.h ++--- a/clang/include/clang/AST/TypeBase.h +++++ b/clang/include/clang/AST/TypeBase.h ++@@ -2270,7 +2270,7 @@ ++ /// increments towards the beginning. ++ /// Positive non-zero number represents the index + 1. ++ /// Zero means this is not substituted from an expansion. ++- unsigned PackIndex : 15; +++ unsigned PackIndex : 16; ++ }; ++ ++ class SubstPackTypeBitfields { ++diff -ruN --strip-trailing-cr a/clang/test/SemaCXX/make_integer_seq.cpp b/clang/test/SemaCXX/make_integer_seq.cpp ++--- a/clang/test/SemaCXX/make_integer_seq.cpp +++++ b/clang/test/SemaCXX/make_integer_seq.cpp ++@@ -50,3 +50,8 @@ ++ ++ __make_integer_seq<__make_integer_seq, int, 10> PR28494; // expected-note{{different template parameters}} ++ // expected-error@make_integer_seq.cpp:* {{template argument for template template parameter must be a class template or type alias template}} + + +-+ // A sane operator new does not read or write accessible memory. +-+ if (CGF.CGM.getCodeGenOpts().AssumeSaneOperatorNew && +-+ CalleeDecl->getDeclName().isAnyOperatorNew()) { +-+ // FIXME: inaccessiblemem could cause issues if LTO makes the +-+ // previously inaccessible memory accessible after linking. +-+ CallOrInvoke->setMemoryEffects( +-+ llvm::MemoryEffects::inaccessibleOrErrnoMemOnly( +-+ llvm::ModRefInfo::ModRef, llvm::ModRefInfo::Mod)); +-+ } +++// The largest type pack whose pack indices fit in SubstTemplateTypeParmType. +++template using Expand = void(Ts...); +++template struct ExpandSeq { using type = Expand; }; +++using LargePack = __make_integer_seq::type; ++diff -ruN --strip-trailing-cr a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp ++--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp +++++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp ++@@ -13901,47 +13901,12 @@ ++ return SDValue(); + } + -+ // ALL: declare noundef nonnull ptr @_Znaj( -++ -++void *f3(unsigned long N) { -++ // SANE: call noalias noundef nonnull ptr @_Znwj(i32 noundef {{.*}}) [[ATTR2:#[0-9]+]] -++ // SANENOT: call noundef nonnull ptr @_Znwj(i32 noundef {{.*}}) [[ATTR2:#[0-9]+]] -++ return ::operator new(N); ++-/// Check if Mask defines a known constant set of enabled lanes, where only the ++-/// first N lanes are enabled. N is returned if so. ++-static uint64_t calculateConstantLowMaskLanes(SDValue Mask) { ++- // We expect masks for masked load/store to be i1 predicates. ++- if (Mask.getValueType().getScalarSizeInBits() != 1) ++- return 0; ++- ++- if (Mask.getOpcode() == ISD::BUILD_VECTOR) { ++- unsigned NumOnes = 0; ++- auto Op = Mask->op_begin(); ++- for (; Op != Mask->op_end(); ++Op) { ++- if (!isOneConstant(*Op)) ++- break; ++- ++NumOnes; ++- } ++- ++- if (NumOnes == 0) ++- return 0; ++- ++- for (; Op != Mask->op_end(); ++Op) ++- if (!isNullConstant(*Op)) ++- return 0; ++- return NumOnes; ++- } ++- ++- if (Mask.getOpcode() == ISD::GET_ACTIVE_LANE_MASK && ++- isNullConstant(Mask.getOperand(0)) && ++- isa(Mask.getOperand(1)) && ++- Mask.getOperand(1).getValueType().getSizeInBits() <= 64) ++- return Mask.getConstantOperandVal(1); ++- ++- return 0; ++-} ++- ++ SDValue DAGCombiner::visitMSTORE(SDNode *N) { ++ MaskedStoreSDNode *MST = cast(N); ++ SDValue Mask = MST->getMask(); ++ SDValue Chain = MST->getChain(); ++ SDValue Value = MST->getValue(); ++ SDValue Ptr = MST->getBasePtr(); ++- EVT VT = Value.getValueType(); ++ ++ // Zap masked stores with a zero mask. ++ if (ISD::isConstantSplatVectorAllZeros(Mask.getNode())) ++@@ -13964,46 +13929,21 @@ ++ } + } + +- return RV; +-diff -ruN --strip-trailing-cr a/clang/test/CodeGenCXX/new_hot_cold.cpp b/clang/test/CodeGenCXX/new_hot_cold.cpp +---- a/clang/test/CodeGenCXX/new_hot_cold.cpp +-+++ b/clang/test/CodeGenCXX/new_hot_cold.cpp +-@@ -124,7 +124,7 @@ ++- // If this is a masked store with an all ones mask, we can use a unmasked ++- // store. +++ // If this is a masked load with an all ones mask, we can use a unmasked load. ++ // FIXME: Can we do this for indexed, compressing, or truncating stores? ++- if (MST->isUnindexed() && !MST->isCompressingStore() && ++- !MST->isTruncatingStore()) { ++- if (ISD::isConstantSplatVectorAllOnes(Mask.getNode())) ++- return DAG.getStore(MST->getChain(), SDLoc(N), MST->getValue(), ++- MST->getBasePtr(), MST->getPointerInfo(), ++- MST->getBaseAlign(), MST->getMemOperand()->getFlags(), ++- MST->getAAInfo()); ++- ++- // Convert a masked_store with constant getactivelanemask input mask to a ++- // standard store. ++- if (uint64_t Lanes = calculateConstantLowMaskLanes(Mask)) { ++- if (Lanes < VT.getVectorMinNumElements() && isPowerOf2_32(Lanes)) { ++- EVT SubVT = EVT::getVectorVT(*DAG.getContext(), ++- VT.getVectorElementType(), Lanes); ++- unsigned IsFast = 0; ++- if ((!LegalTypes || isTypeLegal(SubVT)) && ++- TLI.allowsMemoryAccess(*DAG.getContext(), DAG.getDataLayout(), ++- SubVT, MST->getAddressSpace(), ++- MST->getBaseAlign(), ++- MST->getMemOperand()->getFlags(), &IsFast) && ++- IsFast) { ++- SDLoc DL(N); ++- SDValue Ext = DAG.getExtractSubvector(DL, SubVT, MST->getValue(), 0); ++- return DAG.getStore(MST->getChain(), DL, Ext, MST->getBasePtr(), ++- MST->getPointerInfo(), MST->getBaseAlign(), ++- MST->getMemOperand()->getFlags(), ++- MST->getAAInfo()); ++- } ++- } ++- } ++- } +++ if (ISD::isConstantSplatVectorAllOnes(Mask.getNode()) && MST->isUnindexed() && +++ !MST->isCompressingStore() && !MST->isTruncatingStore()) +++ return DAG.getStore(MST->getChain(), SDLoc(N), MST->getValue(), +++ MST->getBasePtr(), MST->getPointerInfo(), +++ MST->getBaseAlign(), MST->getMemOperand()->getFlags(), +++ MST->getAAInfo()); + +- // CHECK-DAG: attributes [[ATTR_NOBUILTIN]] = { nobuiltin allocsize(0) {{.*}} } +- // CHECK-DAG: attributes [[ATTR_NOBUILTIN_NOTHROW]] = { nobuiltin nounwind allocsize(0) {{.*}} } +--// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_CALL]] = { allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } +-+// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_CALL]] = { allocsize(0) } +- // CHECK-DAG: attributes [[ATTR_BUILTIN_CALL]] = { builtin allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } +--// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_NOTHROW_CALL]] = { nounwind allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } +-+// CHECK-DAG: attributes [[ATTR_NO_BUILTIN_NOTHROW_CALL]] = { nounwind allocsize(0) } +- // CHECK-DAG: attributes [[ATTR_BUILTIN_NOTHROW_CALL]] = { builtin nounwind allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } +-diff -ruN --strip-trailing-cr a/clang/test/CodeGenCXX/operator-new.cpp b/clang/test/CodeGenCXX/operator-new.cpp +---- a/clang/test/CodeGenCXX/operator-new.cpp +-+++ b/clang/test/CodeGenCXX/operator-new.cpp +-@@ -27,5 +27,14 @@ +- } ++ // Try transforming N to an indexed store. ++ if (CombineToPreIndexedLoadStore(N) || CombineToPostIndexedLoadStore(N)) ++ return SDValue(N, 0); ++ ++- if (MST->isTruncatingStore() && MST->isUnindexed() && VT.isInteger() && +++ if (MST->isTruncatingStore() && MST->isUnindexed() && +++ Value.getValueType().isInteger() && ++ (!isa(Value) || ++ !cast(Value)->isOpaque())) { ++ APInt TruncDemandedBits = ++diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll b/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll ++--- a/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll +++++ b/llvm/test/CodeGen/AArch64/lowmaskedlanes.ll ++@@ -139,7 +139,14 @@ ++ define void @store_low2_v4f64(ptr %p, <4 x double> %a) vscale_range(2, 2) { ++ ; CHECK-LABEL: store_low2_v4f64: ++ ; CHECK: // %bb.0: ++-; CHECK-NEXT: str q0, [x0] +++; CHECK-NEXT: // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1 +++; CHECK-NEXT: ptrue p0.d, vl2 +++; CHECK-NEXT: ptrue p1.h, vl2 +++; CHECK-NEXT: // kill: def $q0 killed $q0 killed $z0_z1 def $z0_z1 +++; CHECK-NEXT: splice z0.d, p0, { z0.d, z1.d } +++; CHECK-NEXT: punpklo p0.h, p1.b +++; CHECK-NEXT: punpklo p0.h, p0.b +++; CHECK-NEXT: st1d { z0.d }, p0, [x0] ++ ; CHECK-NEXT: ret ++ %m = tail call <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32 0, i32 2) ++ tail call void @llvm.masked.store(<4 x double> %a, ptr align 1 %p, <4 x i1> %m) ++@@ -149,7 +156,13 @@ ++ define void @store_low2_v4i32(ptr %p, <4 x i32> %a) vscale_range(1, 1) { ++ ; CHECK-LABEL: store_low2_v4i32: ++ ; CHECK: // %bb.0: ++-; CHECK-NEXT: str d0, [x0] +++; CHECK-NEXT: ptrue p0.h, vl2 +++; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 +++; CHECK-NEXT: mov z1.h, p0/z, #-1 // =0xffffffffffffffff +++; CHECK-NEXT: ptrue p0.s +++; CHECK-NEXT: sshll v1.4s, v1.4h, #0 +++; CHECK-NEXT: cmpne p1.s, p0/z, z1.s, #0 +++; CHECK-NEXT: st1w { z0.s }, p1, [x0] ++ ; CHECK-NEXT: ret ++ %m = tail call <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32 0, i32 2) ++ tail call void @llvm.masked.store(<4 x i32> %a, ptr align 1 %p, <4 x i1> %m) ++@@ -159,7 +172,12 @@ ++ define void @store_low2c_v4i32(ptr %p, <4 x i32> %a) vscale_range(1, 1) { ++ ; CHECK-LABEL: store_low2c_v4i32: ++ ; CHECK: // %bb.0: ++-; CHECK-NEXT: str d0, [x0] +++; CHECK-NEXT: adrp x8, .LCPI12_0 +++; CHECK-NEXT: ptrue p0.s +++; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 +++; CHECK-NEXT: ldr q1, [x8, :lo12:.LCPI12_0] +++; CHECK-NEXT: cmpne p1.s, p0/z, z1.s, #0 +++; CHECK-NEXT: st1w { z0.s }, p1, [x0] ++ ; CHECK-NEXT: ret ++ tail call void @llvm.masked.store(<4 x i32> %a, ptr align 1 %p, <4 x i1> ) ++ ret void ++@@ -168,7 +186,11 @@ ++ define void @store_low2_nxv4f64(ptr %p, %a) { ++ ; CHECK-LABEL: store_low2_nxv4f64: ++ ; CHECK: // %bb.0: ++-; CHECK-NEXT: str q0, [x0] +++; CHECK-NEXT: ptrue p0.s, vl2 +++; CHECK-NEXT: punpkhi p1.h, p0.b +++; CHECK-NEXT: punpklo p0.h, p0.b +++; CHECK-NEXT: st1d { z1.d }, p1, [x0, #1, mul vl] +++; CHECK-NEXT: st1d { z0.d }, p0, [x0] ++ ; CHECK-NEXT: ret ++ %m = tail call @llvm.get.active.lane.mask.nxv4i1.i32(i32 0, i32 2) ++ tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) ++@@ -203,10 +225,13 @@ ++ define void @store_low2_v8i16(ptr %p, <8 x i16> %a) vscale_range(2, 2) { ++ ; CHECK-LABEL: store_low2_v8i16: ++ ; CHECK: // %bb.0: +++; CHECK-NEXT: ptrue p0.b, vl2 ++ ; CHECK-NEXT: // kill: def $q0 killed $q0 def $z0 ++-; CHECK-NEXT: ptrue p0.s, vl2 ++-; CHECK-NEXT: uunpklo z0.s, z0.h ++-; CHECK-NEXT: st1h { z0.s }, p0, [x0] +++; CHECK-NEXT: mov z1.b, p0/z, #-1 // =0xffffffffffffffff +++; CHECK-NEXT: ptrue p0.h, vl8 +++; CHECK-NEXT: sshll v1.8h, v1.8b, #0 +++; CHECK-NEXT: cmpne p1.h, p0/z, z1.h, #0 +++; CHECK-NEXT: st1h { z0.h }, p1, [x0] ++ ; CHECK-NEXT: ret ++ %m = tail call <8 x i1> @llvm.get.active.lane.mask.v8i1.i32(i32 0, i32 2) ++ tail call void @llvm.masked.store(<8 x i16> %a, ptr align 1 %p, <8 x i1> %m) ++@@ -216,9 +241,8 @@ ++ define void @store_low2_nxv8i16(ptr %p, %a) { ++ ; CHECK-LABEL: store_low2_nxv8i16: ++ ; CHECK: // %bb.0: ++-; CHECK-NEXT: mov z1.h, z0.h[1] ++-; CHECK-NEXT: str h0, [x0] ++-; CHECK-NEXT: str h1, [x0, #2] +++; CHECK-NEXT: ptrue p0.h, vl2 +++; CHECK-NEXT: st1h { z0.h }, p0, [x0] ++ ; CHECK-NEXT: ret ++ %m = tail call @llvm.get.active.lane.mask.nxv8i1.i32(i32 0, i32 2) ++ tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) ++@@ -239,7 +263,8 @@ ++ define void @store_low4_nxv8i16(ptr %p, %a) { ++ ; CHECK-LABEL: store_low4_nxv8i16: ++ ; CHECK: // %bb.0: ++-; CHECK-NEXT: str d0, [x0] +++; CHECK-NEXT: ptrue p0.h, vl4 +++; CHECK-NEXT: st1h { z0.h }, p0, [x0] ++ ; CHECK-NEXT: ret ++ %m = tail call @llvm.get.active.lane.mask.nxv8i1.i32(i32 0, i32 4) ++ tail call void @llvm.masked.store( %a, ptr align 1 %p, %m) ++diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll b/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll ++--- a/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll +++++ b/llvm/test/CodeGen/Hexagon/autohvx/vector-align-basic.ll ++@@ -6,43 +6,40 @@ ++ ; CHECK-LABEL: f0: ++ ; CHECK: // %bb.0: // %b0 ++ ; CHECK-NEXT: { ++-; CHECK-NEXT: r0 = add(r1,r0) +++; CHECK-NEXT: [[RADD:[r0-9]+]] = add([[R1:[r0-9]+]],[[R0:[r0-9]+]]) ++ ; CHECK-NEXT: } ++ ; CHECK-NEXT: { ++-; CHECK-NEXT: allocframe(r29,#128):raw +++; CHECK-NEXT: [[RLCPI:[r0-9]+]] = ##.LCPI0_0 ++ ; CHECK-NEXT: } ++ ; CHECK-NEXT: { ++-; CHECK-NEXT: r2 = add(r0,#136) +++; CHECK-NEXT: [[RNEG:[r0-9]+]] = #-1 ++ ; CHECK-NEXT: } ++ ; CHECK-NEXT: { ++-; CHECK-NEXT: r29 = and(r29,#-128) +++; CHECK-NEXT: [[RADD136:[r0-9]+]] = add([[RADD]],#136) ++ ; CHECK-NEXT: } ++ ; CHECK-NEXT: { ++-; CHECK-NEXT: v1 = vmem(r0+#1) +++; CHECK-NEXT: [[V1:[v0-9]+]] = vmem([[RLCPI]]+#0) ++ ; CHECK-NEXT: } ++ ; CHECK-NEXT: { ++-; CHECK-NEXT: v0 = vmemu(r2+#0) +++; CHECK-NEXT: [[RADD128:[r0-9]+]] = add([[RADD]],#128) ++ ; CHECK-NEXT: } ++ ; CHECK-NEXT: { ++-; CHECK-NEXT: r4 = add(r29,#0) +++; CHECK-NEXT: [[Q0:[q0-9]+]] = vand([[V1]],[[RNEG]]) ++ ; CHECK-NEXT: } ++ ; CHECK-NEXT: { ++-; CHECK-NEXT: v0.w = vadd(v1.w,v0.w) +++; CHECK-NEXT: [[V2:[v0-9]+]] = vmemu([[RADD136]]+#0) ++ ; CHECK-NEXT: } ++ ; CHECK-NEXT: { ++-; CHECK-NEXT: vmem(r4+#0) = v0 +++; CHECK-NEXT: [[V3:[v0-9]+]] = vmem([[RADD]]+#1) ++ ; CHECK-NEXT: } ++ ; CHECK-NEXT: { ++-; CHECK-NEXT: r2 = memw(r4+#0) +++; CHECK-NEXT: [[V0:[v0-9]+]].w = vadd([[V3]].w,[[V2]].w) ++ ; CHECK-NEXT: } ++ ; CHECK-NEXT: { ++-; CHECK-NEXT: r3 = memw(r4+#4) +++; CHECK-NEXT: if ([[Q0]]) vmem([[RADD128]]+#0) = [[V0]] ++ ; CHECK-NEXT: } ++ ; CHECK-NEXT: { ++-; CHECK-NEXT: memd(r0+#128) = r3:2 ++-; CHECK-NEXT: } ++-; CHECK-NEXT: { ++-; CHECK-NEXT: r31:30 = dealloc_return(r30):raw +++; CHECK-NEXT: jumpr r31 ++ ; CHECK-NEXT: } ++ b0: ++ %v0 = add i32 %a1, 128 ++diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll b/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll ++--- a/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll +++++ b/llvm/test/CodeGen/X86/apx/kmov-postrapseudos.ll ++@@ -29,12 +29,20 @@ ++ define void @kmovkr_2() { ++ ; AVX512-LABEL: kmovkr_2: ++ ; AVX512: # %bb.0: # %alloca_21 ++-; AVX512-NEXT: movq $0, 0 # encoding: [0x48,0xc7,0x04,0x25,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00] +++; AVX512-NEXT: vxorps %xmm0, %xmm0, %xmm0 # encoding: [0xc5,0xf8,0x57,0xc0] +++; AVX512-NEXT: movw $3, %ax # encoding: [0x66,0xb8,0x03,0x00] +++; AVX512-NEXT: kmovw %eax, %k1 # EVEX TO VEX Compression encoding: [0xc5,0xf8,0x92,0xc8] +++; AVX512-NEXT: vmovups %zmm0, 0 {%k1} # encoding: [0x62,0xf1,0x7c,0x49,0x11,0x04,0x25,0x00,0x00,0x00,0x00] +++; AVX512-NEXT: vzeroupper # encoding: [0xc5,0xf8,0x77] ++ ; AVX512-NEXT: retq # encoding: [0xc3] ++ ; ++ ; AVX512BW-LABEL: kmovkr_2: ++ ; AVX512BW: # %bb.0: # %alloca_21 ++-; AVX512BW-NEXT: movq $0, 0 # encoding: [0x48,0xc7,0x04,0x25,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00] +++; AVX512BW-NEXT: vxorps %xmm0, %xmm0, %xmm0 # encoding: [0xc5,0xf8,0x57,0xc0] +++; AVX512BW-NEXT: movw $3, %ax # encoding: [0x66,0xb8,0x03,0x00] +++; AVX512BW-NEXT: kmovd %eax, %k1 # EVEX TO VEX Compression encoding: [0xc5,0xfb,0x92,0xc8] +++; AVX512BW-NEXT: vmovups %zmm0, 0 {%k1} # encoding: [0x62,0xf1,0x7c,0x49,0x11,0x04,0x25,0x00,0x00,0x00,0x00] +++; AVX512BW-NEXT: vzeroupper # encoding: [0xc5,0xf8,0x77] ++ ; AVX512BW-NEXT: retq # encoding: [0xc3] ++ alloca_21: ++ call void @llvm.masked.store.v4f32.p0(<4 x float> zeroinitializer, ptr null, i32 1, <4 x i1> ) ++diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/masked-store-illegal-type.ll b/llvm/test/CodeGen/X86/masked-store-illegal-type.ll ++--- a/llvm/test/CodeGen/X86/masked-store-illegal-type.ll +++++ b/llvm/test/CodeGen/X86/masked-store-illegal-type.ll ++@@ -1,21 +0,0 @@ ++-; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py ++-; RUN: llc < %s -mtriple=x86_64-linux-gnu -mattr=+avx,+avx2,+avx512bw,+avx512cd,+avx512dq,+avx512f,+avx512vl | FileCheck %s ++- ++-define fastcc void @render_event(<12 x double> %0) { ++-; CHECK-LABEL: render_event: ++-; CHECK: # %bb.0: ++-; CHECK-NEXT: vmovlhps {{.*#+}} xmm6 = xmm6[0],xmm7[0] ++-; CHECK-NEXT: vmovlhps {{.*#+}} xmm4 = xmm4[0],xmm5[0] ++-; CHECK-NEXT: vinsertf128 $1, %xmm6, %ymm4, %ymm4 ++-; CHECK-NEXT: vmovlhps {{.*#+}} xmm2 = xmm2[0],xmm3[0] ++-; CHECK-NEXT: vmovlhps {{.*#+}} xmm0 = xmm0[0],xmm1[0] ++-; CHECK-NEXT: vinsertf128 $1, %xmm2, %ymm0, %ymm0 ++-; CHECK-NEXT: vinsertf64x4 $1, %ymm4, %zmm0, %zmm0 ++-; CHECK-NEXT: vmovsd {{.*#+}} xmm1 = mem[0],zero ++-; CHECK-NEXT: vmovsd %xmm1, 64 ++-; CHECK-NEXT: vmovups %zmm0, 0 ++-; CHECK-NEXT: vzeroupper ++-; CHECK-NEXT: retq ++- call void @llvm.masked.store.v12f64.p0(<12 x double> %0, ptr null, <12 x i1> ) ++- ret void ++-} ++diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/pr179489.ll b/llvm/test/CodeGen/X86/pr179489.ll ++--- a/llvm/test/CodeGen/X86/pr179489.ll +++++ b/llvm/test/CodeGen/X86/pr179489.ll ++@@ -9,12 +9,14 @@ ++ ; X86-LABEL: foo: ++ ; X86: # %bb.0: # %entry ++ ; X86-NEXT: movl {{[0-9]+}}(%esp), %eax ++-; X86-NEXT: vpextrb $1, %xmm0, (%eax) +++; X86-NEXT: vmovw %xmm0, %ecx +++; X86-NEXT: movb %ch, (%eax) ++ ; X86-NEXT: retl ++ ; ++ ; X64-LABEL: foo: ++ ; X64: # %bb.0: # %entry ++-; X64-NEXT: vpextrb $1, %xmm0, (%rdi) +++; X64-NEXT: vmovw %xmm0, %eax +++; X64-NEXT: movb %ah, (%rdi) ++ ; X64-NEXT: retq ++ entry: ++ %i3 = shufflevector <8 x i16> %arg, <8 x i16> , <8 x i32> ++diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/X86/pr46532.ll b/llvm/test/CodeGen/X86/pr46532.ll ++--- a/llvm/test/CodeGen/X86/pr46532.ll +++++ b/llvm/test/CodeGen/X86/pr46532.ll ++@@ -7,6 +7,8 @@ ++ ; CHECK-NEXT: movq (%rax), %rax ++ ; CHECK-NEXT: vxorps %xmm0, %xmm0, %xmm0 ++ ; CHECK-NEXT: vmovaps %xmm0, 32(%rax) +++; CHECK-NEXT: vmovsd {{.*#+}} xmm0 = [4294967295,4294967295,0,0] +++; CHECK-NEXT: vmaskmovps %ymm0, %ymm0, (%rax) ++ while.1.body.preheader: ++ %0 = load ptr, ptr undef, align 8, !invariant.load !0, !dereferenceable !1, !align !2 ++ %1 = getelementptr inbounds i8, ptr %0, i64 32 ++diff -ruN --strip-trailing-cr a/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td b/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td ++--- a/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td +++++ b/llvm/test/TableGen/AMDGPUTargetDefGenericClassification.td ++@@ -1,71 +0,0 @@ ++-// RUN: split-file %s %t ++-// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/valid.td \ ++-// RUN: | FileCheck %t/valid.td ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-bug.td 2>&1 \ ++-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=user-sgpr-init16-bug --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-priority.td 2>&1 \ ++-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=required-export-priority --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-timing.td 2>&1 \ ++-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=gfx125x-lowest-rate-wmma --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-banks.td 2>&1 \ ++-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=ldsbankcount32 --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-lds-size.td 2>&1 \ ++-// RUN: | FileCheck %t/missing-generic.td -DDIR=%t -DFEATURE=addressablelocalmemorysize65536 --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/capability.td 2>&1 \ ++-// RUN: | FileCheck %t/capability.td -DFILE=%t/capability.td --implicit-check-not="error:" ++- ++-//--- common.td ++-include "AMDGPU.td" ++- ++-def EmptyGPU : AMDGPUProcessorModel<"gfx999", NoSchedModel, [], [9, 9, 9]>; ++- ++-//--- valid.td ++-include "common.td" ++-// Check the real generic targets, including gfx11's combined workarounds, ++-// gfx12-5's fixup and WMMA timing, and gfx9-4's LDS size and bank count. ++-// CHECK: GK_GFX9_4_GENERIC ++-// CHECK: GK_GFX11_GENERIC ++-// CHECK: GK_GFX12_5_GENERIC ++- ++-//--- missing-bug.td ++-include "common.td" ++-defvar CheckedFeature = FeatureUserSGPRInit16Bug; ++-include "missing-generic.td" ++- ++-//--- missing-priority.td ++-include "common.td" ++-defvar CheckedFeature = FeatureRequiredExportPriority; ++-include "missing-generic.td" ++- ++-//--- missing-timing.td ++-include "common.td" ++-defvar CheckedFeature = FeatureGFX125xLowestRateWMMA; ++-include "missing-generic.td" ++- ++-//--- missing-banks.td ++-include "common.td" ++-defvar CheckedFeature = FeatureLDSBankCount32; ++-include "missing-generic.td" ++- ++-//--- missing-lds-size.td ++-include "common.td" ++-defvar CheckedFeature = FeatureAddressableLocalMemorySize65536; ++-include "missing-generic.td" ++- ++-//--- missing-generic.td ++-// Classified backend-only features still need support from a covered GPU. ++-// Included files use the native path separator after the include directory. ++-// CHECK: [[DIR]]{{[/\\]}}missing-generic.td:[[#@LINE+1]]:1: error: generic target 'gfx9-9-generic' exposes feature '[[FEATURE]]' not supported by any covered GPU ++-def : AMDGPUProcessorModel<"gfx9-9-generic", NoSchedModel, [CheckedFeature], [9, 9, 0]> { ++- let CoveredGPUs = [EmptyGPU]; ++-} ++- ++-//--- capability.td ++-include "common.td" ++-def WMMAGPU : AMDGPUProcessorModel<"gfx998", NoSchedModel, ++- [FeatureWMMA256bInsts], [9, 9, 8]>; ++-// Instruction support keeps the default all-covered-GPUs rule. ++-// CHECK: [[FILE]]:[[#@LINE+1]]:1: error: generic target 'gfx9-9-generic' exposes feature 'wmma-256b-insts' not supported by covered GPU 'gfx999' ++-def : AMDGPUProcessorModel<"gfx9-9-generic", NoSchedModel, [FeatureWMMA256bInsts], [9, 9, 0]> { ++- let CoveredGPUs = [WMMAGPU, EmptyGPU]; ++-} ++diff -ruN --strip-trailing-cr a/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td b/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td ++--- a/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td +++++ b/llvm/test/TableGen/AMDGPUTargetDefGenericFeatures.td ++@@ -1,33 +1,66 @@ ++ // RUN: split-file %s %t ++-// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/valid.td \ +++// RUN: llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/valid.td \ ++ // RUN: | FileCheck %t/valid.td ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-any.td 2>&1 \ +++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-any.td 2>&1 \ ++ // RUN: | FileCheck %t/missing-any.td -DFILE=%t/missing-any.td --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-all.td 2>&1 \ +++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-all.td 2>&1 \ ++ // RUN: | FileCheck %t/missing-all.td -DFILE=%t/missing-all.td --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/implied-all.td 2>&1 \ +++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/implied-all.td 2>&1 \ ++ // RUN: | FileCheck %t/implied-all.td -DFILE=%t/implied-all.td --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/missing-private.td 2>&1 \ +++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/missing-private.td 2>&1 \ ++ // RUN: | FileCheck %t/missing-private.td -DFILE=%t/missing-private.td --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/numeric-mismatch.td 2>&1 \ +++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/implied-private.td 2>&1 \ +++// RUN: | FileCheck %t/implied-private.td -DFILE=%t/implied-private.td --implicit-check-not="error:" +++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/numeric-mismatch.td 2>&1 \ ++ // RUN: | FileCheck %t/numeric-mismatch.td -DFILE=%t/numeric-mismatch.td --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/empty-members.td 2>&1 \ +++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/empty-members.td 2>&1 \ ++ // RUN: | FileCheck %t/empty-members.td -DFILE=%t/empty-members.td --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/self-member.td 2>&1 \ +++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/self-member.td 2>&1 \ ++ // RUN: | FileCheck %t/self-member.td -DFILE=%t/self-member.td --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/generic-member.td 2>&1 \ +++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/generic-member.td 2>&1 \ ++ // RUN: | FileCheck %t/generic-member.td -DFILE=%t/generic-member.td --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/pseudo-member.td 2>&1 \ +++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/pseudo-member.td 2>&1 \ ++ // RUN: | FileCheck %t/pseudo-member.td -DFILE=%t/pseudo-member.td --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/unregistered-member.td 2>&1 \ +++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/unregistered-member.td 2>&1 \ ++ // RUN: | FileCheck %t/unregistered-member.td -DFILE=%t/unregistered-member.td --implicit-check-not="error:" ++-// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../include -I %p/../../lib/Target/AMDGPU %t/r600-member.td 2>&1 \ +++// RUN: not llvm-tblgen -gen-amdgpu-target-def -I %t -I %p/../../lib/Target/AMDGPU %t/r600-member.td 2>&1 \ ++ // RUN: | FileCheck %t/r600-member.td -DFILE=%t/r600-member.td --implicit-check-not="error:" + +- // ALL: declare noundef nonnull ptr @_Znaj( ++ //--- common.td ++-include "llvm/Target/Target.td" +++// Exercise the validation rules with small synthetic targets. The real targets +++// are validated when AMDGPUTargetParserDef.inc is generated during the build. +++// Stand-ins for the Target.td classes used by AMDGPUTargetParser.td, limited to +++// the fields the backend reads. +++class InlineBehavior; +++def InlineIgnore : InlineBehavior; +++def InlineDefault : InlineBehavior; +++class SubtargetFeature i = [], +++ InlineBehavior ib = InlineDefault> { +++ string Name = n; +++ string FieldName = f; +++ string Value = v; +++ string Desc = d; +++ list Implies = i; +++ InlineBehavior InlineBehavior = ib; ++} ++ -+ // SANE: attributes [[ATTR]] = { builtin allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } -++// SANE: attributes [[ATTR2]] = { allocsize(0) } -+ // SANENOT: attributes [[ATTR]] = { builtin allocsize(0) } -++// SANENOT: attributes [[ATTR2]] = { allocsize(0) } -+diff -ruN --strip-trailing-cr a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp -+--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp -++++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp -+@@ -16145,18 +16145,21 @@ -+ -+ InstructionCost VecCost = VectorCost(CommonCost); -+ // Check if the current node must be resized, if the parent node is not -+- // resized. -++ // resized. Nodes without a parent (roots of the splat gather -++ // subtrees) are resized back to their original type when the reusing -++ // gather nodes are emitted. -+ if (It != MinBWs.end() && !UnaryInstruction::isCast(E->getOpcode()) && -+ E->Idx != 0 && -+ (E->getOpcode() != Instruction::Load || E->UserTreeIndex)) { -+ const EdgeInfo &EI = E->UserTreeIndex; -+- if (!EI.UserTE->hasState() || -++ if (!EI.UserTE || !EI.UserTE->hasState() || -+ EI.UserTE->getOpcode() != Instruction::Select || -+ EI.EdgeIdx != 0) { -+ auto UserBWIt = MinBWs.find(EI.UserTE); -+ Type *UserScalarTy = -+- (EI.UserTE->isGather() || -+- EI.UserTE->State == TreeEntry::SplitVectorize) -++ !EI.UserTE ? OrigScalarTy -++ : (EI.UserTE->isGather() || -++ EI.UserTE->State == TreeEntry::SplitVectorize) -+ ? EI.UserTE->Scalars.front()->getType() -+ : EI.UserTE->getOperand(EI.EdgeIdx).front()->getType(); -+ if (UserBWIt != MinBWs.end()) -+diff -ruN --strip-trailing-cr a/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll b/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll -+--- a/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll -++++ b/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll -+@@ -0,0 +1,75 @@ -++; NOTE: Assertions have been autogenerated by utils/update_test_checks.py -++; RUN: opt -passes=slp-vectorizer -S -mtriple=aarch64-linux-gnu < %s | FileCheck %s -++ -++define double @test() { -++; CHECK-LABEL: @test( -++; CHECK-NEXT: bbl: -++; CHECK-NEXT: [[SELECT:%.*]] = select i1 false, i64 0, i64 0 -++; CHECK-NEXT: [[AND:%.*]] = and i64 [[SELECT]], 0 -++; CHECK-NEXT: [[ICMP:%.*]] = icmp eq i64 [[AND]], 0 -++; CHECK-NEXT: [[TMP0:%.*]] = insertelement <2 x i1> poison, i1 true, i64 0 -++; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <2 x i1> [[TMP0]], <2 x i1> poison, <2 x i32> zeroinitializer -++; CHECK-NEXT: [[TMP2:%.*]] = select <2 x i1> [[TMP1]], <2 x double> zeroinitializer, <2 x double> zeroinitializer -++; CHECK-NEXT: [[TMP3:%.*]] = bitcast <2 x double> [[TMP2]] to <2 x i64> -++; CHECK-NEXT: [[TMP4:%.*]] = select <2 x i1> [[TMP1]], <2 x double> zeroinitializer, <2 x double> zeroinitializer -++; CHECK-NEXT: [[TMP5:%.*]] = bitcast <2 x double> [[TMP4]] to <2 x i64> -++; CHECK-NEXT: [[TMP6:%.*]] = insertelement <2 x i1> poison, i1 [[ICMP]], i64 0 -++; CHECK-NEXT: [[TMP7:%.*]] = shufflevector <2 x i1> [[TMP6]], <2 x i1> poison, <2 x i32> zeroinitializer -++; CHECK-NEXT: [[TMP8:%.*]] = select <2 x i1> [[TMP7]], <2 x double> zeroinitializer, <2 x double> zeroinitializer -++; CHECK-NEXT: [[TMP9:%.*]] = bitcast <2 x double> [[TMP8]] to <2 x i64> -++; CHECK-NEXT: [[TMP10:%.*]] = insertelement <2 x i64> poison, i64 [[SELECT]], i64 0 -++; CHECK-NEXT: [[TMP11:%.*]] = shufflevector <2 x i64> [[TMP10]], <2 x i64> poison, <2 x i32> zeroinitializer -++; CHECK-NEXT: [[TMP12:%.*]] = xor <2 x i64> [[TMP11]], [[TMP9]] -++; CHECK-NEXT: [[TMP13:%.*]] = bitcast <2 x i64> [[TMP12]] to <2 x double> -++; CHECK-NEXT: [[TMP14:%.*]] = insertelement <2 x i64> poison, i64 0, i64 0 -++; CHECK-NEXT: [[TMP15:%.*]] = shufflevector <2 x i64> [[TMP14]], <2 x i64> poison, <2 x i32> zeroinitializer -++; CHECK-NEXT: [[TMP16:%.*]] = xor <2 x i64> [[TMP15]], [[TMP3]] -++; CHECK-NEXT: [[TMP17:%.*]] = bitcast <2 x i64> [[TMP16]] to <2 x double> -++; CHECK-NEXT: [[TMP18:%.*]] = xor <2 x i64> [[TMP15]], [[TMP5]] -++; CHECK-NEXT: [[TMP19:%.*]] = fmul <2 x double> [[TMP13]], [[TMP17]] -++; CHECK-NEXT: [[TMP20:%.*]] = bitcast <2 x i64> [[TMP18]] to <2 x double> -++; CHECK-NEXT: [[TMP21:%.*]] = fmul <2 x double> [[TMP19]], [[TMP20]] -++; CHECK-NEXT: [[TMP22:%.*]] = extractelement <2 x double> [[TMP21]], i64 0 -++; CHECK-NEXT: [[TMP23:%.*]] = extractelement <2 x double> [[TMP21]], i64 1 -++; CHECK-NEXT: [[FADD:%.*]] = fadd double [[TMP22]], [[TMP23]] -++; CHECK-NEXT: ret double [[FADD]] -++; -++bbl: -++ %select = select i1 false, i64 0, i64 0 -++ %and = and i64 %select, 0 -++ %icmp = icmp eq i64 %and, 0 -++ %select1 = select i1 %icmp, double 0.000000e+00, double 0.000000e+00 -++ %bitcast = bitcast double %select1 to i64 -++ %xor = xor i64 %select, %bitcast -++ %bitcast2 = bitcast i64 %xor to double -++ %select4 = select i1 false, i64 0, i64 0 -++ %icmp5 = icmp eq i64 %select4, 0 -++ %select6 = select i1 %icmp5, double 0.000000e+00, double 0.000000e+00 -++ %bitcast7 = bitcast double %select6 to i64 -++ %xor8 = xor i64 %select4, %bitcast7 -++ %bitcast9 = bitcast i64 %xor8 to double -++ %select10 = select i1 %icmp5, double 0.000000e+00, double 0.000000e+00 -++ %bitcast11 = bitcast double %select10 to i64 -++ %select12 = select i1 false, i64 0, i64 0 -++ %icmp13 = icmp eq i64 %select12, 0 -++ %select14 = select i1 %icmp13, double 0.000000e+00, double 0.000000e+00 -++ %bitcast15 = bitcast double %select14 to i64 -++ %xor16 = xor i64 %select12, %bitcast15 -++ %select17 = select i1 %icmp13, double 0.000000e+00, double 0.000000e+00 -++ %bitcast18 = bitcast double %select17 to i64 -++ %select3 = select i1 %icmp, double 0.000000e+00, double 0.000000e+00 -++ %bitcast19 = bitcast double %select3 to i64 -++ %xor20 = xor i64 %select, %bitcast19 -++ %bitcast21 = bitcast i64 %xor20 to double -++ %xor22 = xor i64 %select4, %bitcast11 -++ %bitcast23 = bitcast i64 %xor22 to double -++ %fmul = fmul double %bitcast21, %bitcast23 -++ %xor24 = xor i64 %select12, %bitcast18 -++ %bitcast25 = bitcast i64 %xor24 to double -++ %fmul26 = fmul double %fmul, %bitcast25 -++ %fmul27 = fmul double %bitcast2, %bitcast9 -++ %bitcast28 = bitcast i64 %xor16 to double -++ %fmul29 = fmul double %fmul27, %bitcast28 -++ %fadd = fadd double %fmul26, %fmul29 -++ ret double %fadd +++class SchedMachineModel; +++def NoSchedModel : SchedMachineModel; +++class ProcessorItineraries; +++def NoItineraries : ProcessorItineraries; + + +-+void *f3(unsigned long N) { +-+ // SANE: call noalias noundef nonnull ptr @_Znwj(i32 noundef {{.*}}) [[ATTR2:#[0-9]+]] +-+ // SANENOT: call noundef nonnull ptr @_Znwj(i32 noundef {{.*}}) [[ATTR2:#[0-9]+]] +-+ return ::operator new(N); +++class Processor f> { +++ string Name = n; +++ ProcessorItineraries ProcItin = pi; +++ list Features = f; ++} +++class ProcessorModel f> +++ : Processor { +++ SchedMachineModel SchedModel = m; + +} + + +- // SANE: attributes [[ATTR]] = { builtin allocsize(0) memory(inaccessiblemem: readwrite, errnomem: write) } +-+// SANE: attributes [[ATTR2]] = { allocsize(0) } +- // SANENOT: attributes [[ATTR]] = { builtin allocsize(0) } +-+// SANENOT: attributes [[ATTR2]] = { allocsize(0) } +-diff -ruN --strip-trailing-cr a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp +---- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp +-+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp +-@@ -16145,18 +16145,21 @@ ++ include "AMDGPUTargetParser.td" ++-def MyTarget : Target; + +- InstructionCost VecCost = VectorCost(CommonCost); +- // Check if the current node must be resized, if the parent node is not +-- // resized. +-+ // resized. Nodes without a parent (roots of the splat gather +-+ // subtrees) are resized back to their original type when the reusing +-+ // gather nodes are emitted. +- if (It != MinBWs.end() && !UnaryInstruction::isCast(E->getOpcode()) && +- E->Idx != 0 && +- (E->getOpcode() != Instruction::Load || E->UserTreeIndex)) { +- const EdgeInfo &EI = E->UserTreeIndex; +-- if (!EI.UserTE->hasState() || +-+ if (!EI.UserTE || !EI.UserTE->hasState() || +- EI.UserTE->getOpcode() != Instruction::Select || +- EI.EdgeIdx != 0) { +- auto UserBWIt = MinBWs.find(EI.UserTE); +- Type *UserScalarTy = +-- (EI.UserTE->isGather() || +-- EI.UserTE->State == TreeEntry::SplitVectorize) +-+ !EI.UserTE ? OrigScalarTy +-+ : (EI.UserTE->isGather() || +-+ EI.UserTE->State == TreeEntry::SplitVectorize) +- ? EI.UserTE->Scalars.front()->getType() +- : EI.UserTE->getOperand(EI.EdgeIdx).front()->getType(); +- if (UserBWIt != MinBWs.end()) +-diff -ruN --strip-trailing-cr a/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll b/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll +---- a/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll +-+++ b/llvm/test/Transforms/SLPVectorizer/AArch64/splat-gather-subtree-minbw.ll +-@@ -0,0 +1,75 @@ +-+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py +-+; RUN: opt -passes=slp-vectorizer -S -mtriple=aarch64-linux-gnu < %s | FileCheck %s +-+ +-+define double @test() { +-+; CHECK-LABEL: @test( +-+; CHECK-NEXT: bbl: +-+; CHECK-NEXT: [[SELECT:%.*]] = select i1 false, i64 0, i64 0 +-+; CHECK-NEXT: [[AND:%.*]] = and i64 [[SELECT]], 0 +-+; CHECK-NEXT: [[ICMP:%.*]] = icmp eq i64 [[AND]], 0 +-+; CHECK-NEXT: [[TMP0:%.*]] = insertelement <2 x i1> poison, i1 true, i64 0 +-+; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <2 x i1> [[TMP0]], <2 x i1> poison, <2 x i32> zeroinitializer +-+; CHECK-NEXT: [[TMP2:%.*]] = select <2 x i1> [[TMP1]], <2 x double> zeroinitializer, <2 x double> zeroinitializer +-+; CHECK-NEXT: [[TMP3:%.*]] = bitcast <2 x double> [[TMP2]] to <2 x i64> +-+; CHECK-NEXT: [[TMP4:%.*]] = select <2 x i1> [[TMP1]], <2 x double> zeroinitializer, <2 x double> zeroinitializer +-+; CHECK-NEXT: [[TMP5:%.*]] = bitcast <2 x double> [[TMP4]] to <2 x i64> +-+; CHECK-NEXT: [[TMP6:%.*]] = insertelement <2 x i1> poison, i1 [[ICMP]], i64 0 +-+; CHECK-NEXT: [[TMP7:%.*]] = shufflevector <2 x i1> [[TMP6]], <2 x i1> poison, <2 x i32> zeroinitializer +-+; CHECK-NEXT: [[TMP8:%.*]] = select <2 x i1> [[TMP7]], <2 x double> zeroinitializer, <2 x double> zeroinitializer +-+; CHECK-NEXT: [[TMP9:%.*]] = bitcast <2 x double> [[TMP8]] to <2 x i64> +-+; CHECK-NEXT: [[TMP10:%.*]] = insertelement <2 x i64> poison, i64 [[SELECT]], i64 0 +-+; CHECK-NEXT: [[TMP11:%.*]] = shufflevector <2 x i64> [[TMP10]], <2 x i64> poison, <2 x i32> zeroinitializer +-+; CHECK-NEXT: [[TMP12:%.*]] = xor <2 x i64> [[TMP11]], [[TMP9]] +-+; CHECK-NEXT: [[TMP13:%.*]] = bitcast <2 x i64> [[TMP12]] to <2 x double> +-+; CHECK-NEXT: [[TMP14:%.*]] = insertelement <2 x i64> poison, i64 0, i64 0 +-+; CHECK-NEXT: [[TMP15:%.*]] = shufflevector <2 x i64> [[TMP14]], <2 x i64> poison, <2 x i32> zeroinitializer +-+; CHECK-NEXT: [[TMP16:%.*]] = xor <2 x i64> [[TMP15]], [[TMP3]] +-+; CHECK-NEXT: [[TMP17:%.*]] = bitcast <2 x i64> [[TMP16]] to <2 x double> +-+; CHECK-NEXT: [[TMP18:%.*]] = xor <2 x i64> [[TMP15]], [[TMP5]] +-+; CHECK-NEXT: [[TMP19:%.*]] = fmul <2 x double> [[TMP13]], [[TMP17]] +-+; CHECK-NEXT: [[TMP20:%.*]] = bitcast <2 x i64> [[TMP18]] to <2 x double> +-+; CHECK-NEXT: [[TMP21:%.*]] = fmul <2 x double> [[TMP19]], [[TMP20]] +-+; CHECK-NEXT: [[TMP22:%.*]] = extractelement <2 x double> [[TMP21]], i64 0 +-+; CHECK-NEXT: [[TMP23:%.*]] = extractelement <2 x double> [[TMP21]], i64 1 +-+; CHECK-NEXT: [[FADD:%.*]] = fadd double [[TMP22]], [[TMP23]] +-+; CHECK-NEXT: ret double [[FADD]] +-+; +-+bbl: +-+ %select = select i1 false, i64 0, i64 0 +-+ %and = and i64 %select, 0 +-+ %icmp = icmp eq i64 %and, 0 +-+ %select1 = select i1 %icmp, double 0.000000e+00, double 0.000000e+00 +-+ %bitcast = bitcast double %select1 to i64 +-+ %xor = xor i64 %select, %bitcast +-+ %bitcast2 = bitcast i64 %xor to double +-+ %select4 = select i1 false, i64 0, i64 0 +-+ %icmp5 = icmp eq i64 %select4, 0 +-+ %select6 = select i1 %icmp5, double 0.000000e+00, double 0.000000e+00 +-+ %bitcast7 = bitcast double %select6 to i64 +-+ %xor8 = xor i64 %select4, %bitcast7 +-+ %bitcast9 = bitcast i64 %xor8 to double +-+ %select10 = select i1 %icmp5, double 0.000000e+00, double 0.000000e+00 +-+ %bitcast11 = bitcast double %select10 to i64 +-+ %select12 = select i1 false, i64 0, i64 0 +-+ %icmp13 = icmp eq i64 %select12, 0 +-+ %select14 = select i1 %icmp13, double 0.000000e+00, double 0.000000e+00 +-+ %bitcast15 = bitcast double %select14 to i64 +-+ %xor16 = xor i64 %select12, %bitcast15 +-+ %select17 = select i1 %icmp13, double 0.000000e+00, double 0.000000e+00 +-+ %bitcast18 = bitcast double %select17 to i64 +-+ %select3 = select i1 %icmp, double 0.000000e+00, double 0.000000e+00 +-+ %bitcast19 = bitcast double %select3 to i64 +-+ %xor20 = xor i64 %select, %bitcast19 +-+ %bitcast21 = bitcast i64 %xor20 to double +-+ %xor22 = xor i64 %select4, %bitcast11 +-+ %bitcast23 = bitcast i64 %xor22 to double +-+ %fmul = fmul double %bitcast21, %bitcast23 +-+ %xor24 = xor i64 %select12, %bitcast18 +-+ %bitcast25 = bitcast i64 %xor24 to double +-+ %fmul26 = fmul double %fmul, %bitcast25 +-+ %fmul27 = fmul double %bitcast2, %bitcast9 +-+ %bitcast28 = bitcast i64 %xor16 to double +-+ %fmul29 = fmul double %fmul27, %bitcast28 +-+ %fadd = fadd double %fmul26, %fmul29 +-+ ret double %fadd ++ def FeatureAll : SubtargetFeature<"all", "HasAll", "true", "">; ++ def FeatureAnyA : AMDGPUGenericAnyFeature<"any-a", "HasAnyA", "true", "">; ++@@ -37,7 +70,13 @@ ++ def FeatureAnySmall : NumericFeature<4>; ++ def FeatureAnyB : NumericFeature<8>; ++ def FeatureAnyLarge : NumericFeature<16>; ++-def FeaturePrivate : AMDGPUGenericAnyFeature<"private", "HasPrivate", "true", "">; +++// Classification inherited through a subclass also applies to backend-only +++// features, which must participate in validation without being emitted. +++class PrivateFeature : AMDGPUGenericAnyFeature< +++ "private", "HasPrivate", "true", "">; +++def FeaturePrivate : PrivateFeature; +++def FeatureImpliesPrivate : SubtargetFeature< +++ "implies-private", "HasImpliesPrivate", "true", "", [FeaturePrivate]>; ++ def FeatureImpliesA : SubtargetFeature< ++ "implies-a", "HasImpliesA", "true", "", [FeatureAnyA]>; ++ def FeatureImpliesB : SubtargetFeature< ++@@ -55,7 +94,7 @@ ++ def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, ++ [FeatureAll, FeatureAnyA, FeatureAnySmall], [9, 0, 0]>; ++ def GFX901 : AMDGPUProcessorModel<"gfx901", NoSchedModel, ++- [FeatureAll, FeatureImpliesB, FeaturePrivate], [9, 0, 1]>; +++ [FeatureAll, FeatureImpliesB, FeatureImpliesPrivate], [9, 0, 1]>; ++ // Different members supply different marked features. Follow Implies on both ++ // the generic and the members, and check marked backend-only features too. ++ def : AMDGPUProcessorModel<"gfx9-generic", NoSchedModel, ++@@ -103,6 +142,15 @@ ++ let CoveredGPUs = [GFX900]; ++ } ++ +++//--- implied-private.td +++include "common.td" +++def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, [], [9, 0, 0]>; +++// An unclassified backend-only parent must not hide a classified dependency. +++// CHECK: [[FILE]]:[[#@LINE+1]]:1: error: generic target 'gfx9-generic' exposes feature 'private' not supported by any covered GPU +++def : AMDGPUProcessorModel<"gfx9-generic", NoSchedModel, [FeatureImpliesPrivate], [9, 0, 0]> { +++ let CoveredGPUs = [GFX900]; + +} +++ ++ //--- numeric-mismatch.td ++ include "common.td" ++ def GFX900 : AMDGPUProcessorModel<"gfx900", NoSchedModel, [FeatureAnySmall], [9, 0, 0]>; ++diff -ruN --strip-trailing-cr a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp ++--- a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp +++++ b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp ++@@ -14,11 +14,12 @@ ++ using namespace mlir; ++ using namespace mlir::tosa; ++ ++-// Building this ~5,000-line generated initializer map under MemorySanitizer ++-// leads to extreme compile times (>3.5 minutes) during InstCombine and greedy ++-// register allocation due to shadow/origin instrumentation overhead. Disable ++-// optimization under MSan to keep build times manageable. ++-#if LLVM_MEMORY_SANITIZER_BUILD +++// Building this ~5,000-line generated initializer map under HWAddressSanitizer +++// or MemorySanitizer leads to extreme compile times (>3.5 minutes) during +++// InstCombine and greedy register allocation due to instrumentation +++// overhead. Disable optimization under MSan/HWAsan to keep build times +++// manageable. +++#if LLVM_MEMORY_SANITIZER_BUILD || LLVM_HWADDRESS_SANITIZER_BUILD ++ __attribute__((optnone)) ++ #else ++ LLVM_ATTRIBUTE_MINSIZE ++diff -ruN --strip-trailing-cr a/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel b/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel ++--- a/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel +++++ b/utils/bazel/llvm-project-overlay/third-party/BUILD.bazel ++@@ -66,7 +66,7 @@ ++ ++ bool_flag( ++ name = "llvm_enable_lzma", ++- build_setting_default = False, +++ build_setting_default = True, ++ ) ++ ++ config_setting( ++@@ -74,11 +74,7 @@ ++ flag_values = {":llvm_enable_lzma": "true"}, ++ ) ++ ++-# Unlike zlib and zstd there is no liblzma module in the Bazel registry, so this ++-# links the system library rather than building one from source. That is why it ++-# is a plain cc_library with linkopts instead of a cc_library_wrapper, and why ++-# it defaults to off. ++-cc_library( +++cc_library_wrapper( ++ name = "lzma", ++ defines = select({ ++ ":llvm_lzma_enabled": [ ++@@ -88,9 +84,9 @@ ++ "LLVM_ENABLE_LZMA=0", ++ ], ++ }), ++- linkopts = select({ +++ deps = select({ ++ ":llvm_lzma_enabled": [ ++- "-llzma", +++ "@llvm_xz//:lzma", ++ ], ++ "//conditions:default": [], ++ }), ++diff -ruN --strip-trailing-cr a/utils/bazel/MODULE.bazel b/utils/bazel/MODULE.bazel ++--- a/utils/bazel/MODULE.bazel +++++ b/utils/bazel/MODULE.bazel ++@@ -29,6 +29,7 @@ ++ bazel_dep(name = "nanobind", version = "2.12.0.bcr.1") ++ bazel_dep(name = "libpfm", version = "4.13.0", repo_name = "pfm") ++ bazel_dep(name = "vulkan_headers", version = "1.4.349") +++bazel_dep(name = "xz", version = "5.4.5.bcr.9", repo_name = "llvm_xz") ++ ++ bazel_dep(name = "emsdk", version = "6.0.2", dev_dependency = True) ++ bazel_dep(name = "llvm", version = "0.8.20", dev_dependency = True) ++diff -ruN --strip-trailing-cr a/utils/bazel/MODULE.bazel.lock b/utils/bazel/MODULE.bazel.lock ++--- a/utils/bazel/MODULE.bazel.lock +++++ b/utils/bazel/MODULE.bazel.lock ++@@ -55,7 +55,8 @@ ++ "https://bcr.bazel.build/modules/bazel_lib/3.0.0/MODULE.bazel": "22b70b80ac89ad3f3772526cd9feee2fa412c2b01933fea7ed13238a448d370d", ++ "https://bcr.bazel.build/modules/bazel_lib/3.1.0/MODULE.bazel": "6809765c14e3c766a9b9286c7b0ec56ed87a73326e48fe01749f0c0fdcfe3287", ++ "https://bcr.bazel.build/modules/bazel_lib/3.2.2/MODULE.bazel": "e2c890c8a515d6bca9c66d47718aa9e44b458fde64ec7204b8030bf2d349058c", ++- "https://bcr.bazel.build/modules/bazel_lib/3.2.2/source.json": "9e84e115c20e14652c5c21401ae85ff4daa8702e265b5c0b3bf89353f17aa212", +++ "https://bcr.bazel.build/modules/bazel_lib/3.7.2/MODULE.bazel": "cbc920fff00a55b2fd663d40850560a6fb17fd5ef3a9f3d6cbf0d32908dd8a40", +++ "https://bcr.bazel.build/modules/bazel_lib/3.7.2/source.json": "7a9c2b06a113c4962d329bb68befdb90d0447788a19aabaa32b476a0458e7d28", ++ "https://bcr.bazel.build/modules/bazel_skylib/1.0.3/MODULE.bazel": "bcb0fd896384802d1ad283b4e4eb4d718eebd8cb820b0a2c3a347fb971afd9d8", ++ "https://bcr.bazel.build/modules/bazel_skylib/1.1.1/MODULE.bazel": "1add3e7d93ff2e6998f9e118022c84d163917d912f5afafb3058e3d2f1545b5e", ++ "https://bcr.bazel.build/modules/bazel_skylib/1.2.0/MODULE.bazel": "44fe84260e454ed94ad326352a698422dbe372b21a1ac9f3eab76eb531223686", ++@@ -313,6 +314,8 @@ ++ "https://bcr.bazel.build/modules/windows_support/0.4.1/source.json": "bdb72e6717b00135e0a4a8512b3f6b626df0baea0bde74bf7312c7484f791287", ++ "https://bcr.bazel.build/modules/with_cfg.bzl/0.12.0/MODULE.bazel": "b573395fe63aef4299ba095173e2f62ccfee5ad9bbf7acaa95dba73af9fc2b38", ++ "https://bcr.bazel.build/modules/with_cfg.bzl/0.12.0/source.json": "3f3fbaeafecaf629877ad152a2c9def21f8d330d91aa94c5dc75bbb98c10b8b8", +++ "https://bcr.bazel.build/modules/xz/5.4.5.bcr.9/MODULE.bazel": "ff7cf8a5334a96b7b9ffcbe73418d45929ce125efebac39c1ee7d14bd2ae7a9b", +++ "https://bcr.bazel.build/modules/xz/5.4.5.bcr.9/source.json": "abf61f3c1630eab7b81722f710060b6f31e7c160511bcbfaad74bb4385883051", ++ "https://bcr.bazel.build/modules/zlib-ng/2.3.3.bcr.2/MODULE.bazel": "86130924115bf8dfc4aba286fb49a2ec8d32430a21771db9ccf5381acf0d10cc", ++ "https://bcr.bazel.build/modules/zlib-ng/2.3.3.bcr.2/source.json": "7aa4112eb2fbb99757e97cf0b7a4e2bdb3464b2b3cfb351bc807ffcb34bd349f", ++ "https://bcr.bazel.build/modules/zlib-ng/2.3.3/MODULE.bazel": "0afa781a5f354b4fd811a5c9a086e111daf50e09d8e98ab1cf7c46eeacc33aac", +diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl +index d929da09..88052c28 100644 +--- a/third_party/llvm/workspace.bzl ++++ b/third_party/llvm/workspace.bzl +@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive") + + def repo(name): + """Imports LLVM.""" +- LLVM_COMMIT = "a22291abb68744cb0b70978eb506a7c5592a3309" +- LLVM_SHA256 = "cf19b82f0dfbef1e62afe520ba4ee6de4f4823524d7804490e88d815f822c8e9" ++ LLVM_COMMIT = "e56c2cefc3e7978de4a2d799fa3abec77d75af21" ++ LLVM_SHA256 = "affcad31c9455e95faca9d4de38e4a49b73c23bbf48b405d264ff10a6bfd3baf" + + tf_http_archive( + name = name, +diff --git a/third_party/stablehlo/temporary.patch b/third_party/stablehlo/temporary.patch +index fb55a869..81f3bf2f 100644 +--- a/third_party/stablehlo/temporary.patch ++++ b/third_party/stablehlo/temporary.patch +@@ -123,6 +123,328 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgU + + /// Extracts integer values from the attribute |elements|. + SmallVector extract1DVector(DenseIntElementsAttr elements); ++diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir ++--- stablehlo/stablehlo/conversions/tosa/tests/binary.mlir +++++ stablehlo/stablehlo/conversions/tosa/tests/binary.mlir ++@@ -38,7 +38,7 @@ ++ ++ // CHECK-LABEL: @concatenate ++ func.func @concatenate(%arg0 : tensor<3x3xf32>, %arg1 : tensor<3x3xf32>) -> tensor<6x3xf32> { ++- // CHECK: tosa.concat %arg0, %arg1 {axis = 0 : i32} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32> +++ // CHECK: tosa.concat %arg0, %arg1 axis(0) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32> ++ %0 = "stablehlo.concatenate"(%arg0, %arg1) {dimension = 0 : i64} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32> ++ return %0 : tensor<6x3xf32> ++ } ++@@ -61,10 +61,10 @@ ++ ++ // CHECK-LABEL: @dot_vector_vector ++ func.func @dot_vector_vector(%arg0 : tensor<3xf32>, %arg1 : tensor<3xf32>) -> tensor { ++- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<> : tensor<0xindex>} : () -> !tosa.shape<0> ++- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<> : tensor<0xindex>) : () -> !tosa.shape<0> +++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] ++- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] ++ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] ++ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] ++@@ -74,10 +74,10 @@ ++ ++ // CHECK-LABEL: @dot_vector_matrix ++ func.func @dot_vector_matrix(%arg0 : tensor<2xf32>, %arg1 : tensor<2x3xf32>) -> tensor<3xf32> { ++- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<3> : tensor<1xindex>} : () -> !tosa.shape<1> ++- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 2]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<3> : tensor<1xindex>) : () -> !tosa.shape<1> +++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 2]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] ++- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] ++ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] ++ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] ++@@ -87,10 +87,10 @@ ++ ++ // CHECK-LABEL: @dot_matrix_vector ++ func.func @dot_matrix_vector(%arg0 : tensor<2x3xf32>, %arg1 : tensor<3xf32>) -> tensor<2xf32> { ++- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<2> : tensor<1xindex>} : () -> !tosa.shape<1> ++- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<2> : tensor<1xindex>) : () -> !tosa.shape<1> +++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] ++- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] ++ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] ++ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] ++@@ -100,10 +100,10 @@ ++ ++ // CHECK-LABEL: @dot_matrix_matrix ++ func.func @dot_matrix_matrix(%arg0 : tensor<2x3xf32>, %arg1 : tensor<3x4xf32>) -> tensor<2x4xf32> { ++- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<[2, 4]> : tensor<2xindex>} : () -> !tosa.shape<2> ++- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<[2, 4]> : tensor<2xindex>) : () -> !tosa.shape<2> +++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] ++- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 4]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 4]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] ++ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] ++ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] ++@@ -113,10 +113,10 @@ ++ ++ // CHECK-LABEL: @dot_general_vector_vector ++ func.func @dot_general_vector_vector(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>) -> tensor { ++- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<> : tensor<0xindex>} : () -> !tosa.shape<0> ++- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<> : tensor<0xindex>) : () -> !tosa.shape<0> +++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] ++- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] ++ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] ++ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] ++@@ -126,10 +126,10 @@ ++ ++ // CHECK-LABEL: @dot_general_vector_matrix ++ func.func @dot_general_vector_matrix(%arg0: tensor<2xf32>, %arg1: tensor<2x3xf32>) -> tensor<3xf32> { ++- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<3> : tensor<1xindex>} : () -> !tosa.shape<1> ++- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 2]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<3> : tensor<1xindex>) : () -> !tosa.shape<1> +++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 2]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] ++- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] ++ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] ++ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] ++@@ -139,10 +139,10 @@ ++ ++ // CHECK-LABEL: @dot_general_matrix_vector ++ func.func @dot_general_matrix_vector(%arg0: tensor<2x3xf32>, %arg1: tensor<3xf32>) -> tensor<2xf32> { ++- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<2> : tensor<1xindex>} : () -> !tosa.shape<1> ++- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<2> : tensor<1xindex>) : () -> !tosa.shape<1> +++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] ++- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] ++ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] ++ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] ++@@ -152,10 +152,10 @@ ++ ++ // CHECK-LABEL: @dot_general_matrix_matrix ++ func.func @dot_general_matrix_matrix(%arg0: tensor<2x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<2x4xf32> { ++- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<[2, 4]> : tensor<2xindex>} : () -> !tosa.shape<2> ++- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<[2, 4]> : tensor<2xindex>) : () -> !tosa.shape<2> +++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] ++- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 4]> : tensor<3xindex>} : () -> !tosa.shape<3> +++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 4]> : tensor<3xindex>) : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] ++ // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] ++ // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] ++diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir b/stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir ++--- stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir +++++ stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir ++@@ -11,10 +11,10 @@ ++ // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1431655765> : tensor<1xi32>}> ++ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> ++ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> ++- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) +++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: %[[V2:.+]] = stablehlo.add %[[V0]], %[[V1]] : tensor<2x2xi32> ++- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> ++ %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) ++ -> tensor<2x2x!quant.uniform> ++@@ -32,10 +32,10 @@ ++ // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1431655765> : tensor<1xi32>}> ++ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> ++ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> ++- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) +++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: %[[V2:.+]] = stablehlo.subtract %[[V0]], %[[V1]] : tensor<2x2xi32> ++- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> ++ %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) ++ -> tensor<2x2x!quant.uniform> ++@@ -52,10 +52,10 @@ ++ // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1717986918> : tensor<1xi32>}> ++ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> ++ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> ++- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) +++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: %[[V2:.+]] = stablehlo.multiply %[[V0]], %[[V1]] : tensor<2x2xi32> ++- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_1]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_1]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> ++ %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) ++ -> tensor<2x2x!quant.uniform> ++@@ -74,10 +74,10 @@ ++ // CHECK-DAG: %[[ZP_MINUS_2:.+]] = "tosa.const"() <{values = dense<-2> : tensor<1xi8>}> ++ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> ++ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> ++- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) +++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: %[[V2:.+]] = stablehlo.divide %[[V0]], %[[V1]] : tensor<2x2xi32> ++- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_3]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_3]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> ++ %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) ++ -> tensor<2x2x!quant.uniform> ++@@ -97,10 +97,10 @@ ++ // CHECK-DAG: %[[SHIFT12:.+]] = "tosa.const"() <{values = dense<12> : tensor<1xi8>}> ++ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> ++ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> ++- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) +++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: %[[V2:.+]] = stablehlo.maximum %[[V0]], %[[V1]] : tensor<2x2xi32> ++- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> ++ %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) ++ -> tensor<2x2x!quant.uniform> ++@@ -120,10 +120,10 @@ ++ // CHECK-DAG: %[[SHIFT12:.+]] = "tosa.const"() <{values = dense<12> : tensor<1xi8>}> ++ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> ++ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> ++- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) +++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: %[[V2:.+]] = stablehlo.minimum %[[V0]], %[[V1]] : tensor<2x2xi32> ++- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> ++ %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) ++ -> tensor<2x2x!quant.uniform> ++@@ -140,9 +140,9 @@ ++ // CHECK-DAG: %[[SHIFT30:.+]] = "tosa.const"() <{values = dense<30> : tensor<1xi8>}> ++ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> ++ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> ++- // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: %[[V1:.+]] = stablehlo.abs %[[V0]] : tensor<20x20xi32> ++- // CHECK: %[[V3:.+]] = tosa.rescale %[[V1]], %[[MULTIPLIER_1]], %[[SHIFT33]], %[[ZP_0]], %[[ZP_MINUS_128]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V1]], %[[MULTIPLIER_1]], %[[SHIFT33]], %[[ZP_0]], %[[ZP_MINUS_128]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: return %[[V3]] : tensor<20x20x!quant.uniform> ++ %0 = "stablehlo.abs"(%arg0) : (tensor<20x20x!quant.uniform>) -> tensor<20x20x!quant.uniform> ++ return %0 : tensor<20x20x!quant.uniform> ++@@ -159,8 +159,8 @@ ++ // CHECK-DAG: %[[SHIFT12:.+]] = "tosa.const"() <{values = dense<12> : tensor<1xi8>}> ++ // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> ++ // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> ++- // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++- // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) +++ // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: %[[V2:.+]] = stablehlo.compare GE, %[[V0]], %[[V1]], TOTALORDER : ++ // CHECK: return %[[V2]] ++ %0 = stablehlo.compare GE, %arg0, %arg1, TOTALORDER : (tensor<20x20x!quant.uniform>, tensor<20x20x!quant.uniform>) -> tensor<20x20xi1> ++@@ -177,8 +177,8 @@ ++ // CHECK-DAG: %[[SHIFT15:.+]] = "tosa.const"() <{values = dense<15> : tensor<1xi8>}> ++ // CHECK-DAG: %[[ZP16_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi16>}> ++ // CHECK-DAG: %[[ZP32_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> ++- // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT17]], %[[ZP16_0]], %[[ZP32_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++- // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT15]], %[[ZP16_0]], %[[ZP32_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +++ // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT17]], %[[ZP16_0]], %[[ZP32_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) +++ // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT15]], %[[ZP16_0]], %[[ZP32_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: %[[V2:.+]] = stablehlo.compare LT, %[[V0]], %[[V1]], TOTALORDER : ++ // CHECK: return %[[V2]] ++ %0 = stablehlo.compare LT, %arg0, %arg1, TOTALORDER : (tensor<20x20x!quant.uniform>, tensor<20x20x!quant.uniform>) -> tensor<20x20xi1> ++diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir b/stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir ++--- stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir +++++ stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir ++@@ -7,7 +7,7 @@ ++ %shift = "tosa.const"() {values = dense<13> : tensor<1xi8>} : () -> tensor<1xi8> ++ %input_zp = "tosa.const"() {values = dense<-1> : tensor<1xi8>} : () -> tensor<1xi8> ++ %output_zp = "tosa.const"() {values = dense<0> : tensor<1xi32>} : () -> tensor<1xi32> ++- %0 = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} : +++ %0 = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) : ++ (tensor<2x2x!quant.uniform>, tensor<1xi32>, tensor<1xi8>, tensor<1xi8>, tensor<1xi32>) -> tensor<2x2xi32> ++ ++ // convert input quantized type to storage type ++diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir ++--- stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir +++++ stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir ++@@ -18,7 +18,7 @@ ++ func.func @iota_dimension_0() -> tensor<4x8xf32> { ++ // CHECK-DAG: %[[VAR0:.*]] = "tosa.const"() ++ // CHECK-SAME{LITERAL}: <{values = dense<[[0.000000e+00], [1.000000e+00], [2.000000e+00], [3.000000e+00]]> : tensor<4x1xf32>}> : () -> tensor<4x1xf32> ++- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 8]> : vector<2xindex>} : () -> !tosa.shape<2> +++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 8]> : vector<2xindex>) : () -> !tosa.shape<2> ++ // CHECK-DAG: %[[VAR2:.*]] = tosa.tile %[[VAR0]], %[[VAR1]] ++ %0 = "stablehlo.iota"() {iota_dimension = 0 : i64} : () -> (tensor<4x8xf32>) ++ return %0 : tensor<4x8xf32> ++@@ -28,7 +28,7 @@ ++ func.func @iota_dimension_1() -> tensor<4x8xi32> { ++ // CHECK-DAG: %[[VAR0:.*]] = "tosa.const"() ++ // CHECK-SAME{LITERAL}: <{values = dense<[[0, 1, 2, 3, 4, 5, 6, 7]]> : tensor<1x8xi32>}> : () -> tensor<1x8xi32> ++- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[4, 1]> : vector<2xindex>} : () -> !tosa.shape<2> +++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[4, 1]> : vector<2xindex>) : () -> !tosa.shape<2> ++ // CHECK-DAG: %[[VAR2:.*]] = tosa.tile %[[VAR0]], %[[VAR1]] ++ %0 = "stablehlo.iota"() {iota_dimension = 1 : i64} : () -> (tensor<4x8xi32>) ++ return %0 : tensor<4x8xi32> ++diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir ++--- stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir +++++ stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir ++@@ -2,7 +2,7 @@ ++ ++ // CHECK-LABEL: @concatenate ++ func.func @concatenate(%arg0 : tensor<5x2xf32>, %arg1 : tensor<5x5xf32>, %arg2 : tensor<5x7xf32>) -> tensor<5x14xf32> { ++- // CHECK: tosa.concat %arg0, %arg1, %arg2 {axis = 1 : i32} : (tensor<5x2xf32>, tensor<5x5xf32>, tensor<5x7xf32>) -> tensor<5x14xf32> +++ // CHECK: tosa.concat %arg0, %arg1, %arg2 axis(1) : (tensor<5x2xf32>, tensor<5x5xf32>, tensor<5x7xf32>) -> tensor<5x14xf32> ++ %0 = "stablehlo.concatenate"(%arg0, %arg1, %arg2) {dimension = 1 : i64} : (tensor<5x2xf32>, tensor<5x5xf32>, tensor<5x7xf32>) -> tensor<5x14xf32> ++ return %0 : tensor<5x14xf32> ++ } ++diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/unary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/unary.mlir ++--- stablehlo/stablehlo/conversions/tosa/tests/unary.mlir +++++ stablehlo/stablehlo/conversions/tosa/tests/unary.mlir ++@@ -79,8 +79,8 @@ ++ ++ // CHECK-LABEL: @slice ++ func.func @slice(%arg : tensor<4x3xf32>) -> tensor<2x2xf32> { ++- // CHECK-DAG: %[[SIZE:.*]] = tosa.const_shape {values = dense<[2, 1]> : tensor<2xindex>} : () -> !tosa.shape<2> ++- // CHECK-DAG: %[[START:.*]] = tosa.const_shape {values = dense<2> : tensor<2xindex>} : () -> !tosa.shape<2> +++ // CHECK-DAG: %[[SIZE:.*]] = tosa.const_shape values(dense<[2, 1]> : tensor<2xindex>) : () -> !tosa.shape<2> +++ // CHECK-DAG: %[[START:.*]] = tosa.const_shape values(dense<2> : tensor<2xindex>) : () -> !tosa.shape<2> ++ // CHECK: tosa.slice %arg0, %[[SIZE]], %[[START]] ++ %0 = "stablehlo.slice"(%arg) { ++ start_indices = array, ++@@ -123,7 +123,7 @@ ++ ++ // CHECK-LABEL: @transpose ++ func.func @transpose(%arg0: tensor<1x2x3xf32>) -> tensor<3x2x1xf32> { ++- // CHECK: %[[VAR0:.*]] = tosa.transpose %arg0 {perms = array} +++ // CHECK: %[[VAR0:.*]] = tosa.transpose %arg0 perms([2, 1, 0]) ++ %0 = "stablehlo.transpose"(%arg0) {permutation = array} : (tensor<1x2x3xf32>) -> tensor<3x2x1xf32> ++ return %0 : tensor<3x2x1xf32> ++ } ++@@ -158,7 +158,7 @@ ++ ++ // CHECK-LABEL: @reshape ++ func.func @reshape(%arg0 : tensor<2x3xf32>) -> tensor<6xf32> { ++- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<6> : tensor<1xindex>} : () -> !tosa.shape<1> +++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<6> : tensor<1xindex>) : () -> !tosa.shape<1> ++ // CHECK: tosa.reshape %arg0, %[[VAR0]] ++ %0 = "stablehlo.reshape"(%arg0) : (tensor<2x3xf32>) -> tensor<6xf32> ++ return %0 : tensor<6xf32> + diff --ruN a/stablehlo/stablehlo/dialect/ChloBytecode.cpp b/stablehlo/stablehlo/dialect/ChloBytecode.cpp + --- stablehlo/stablehlo/dialect/ChloBytecode.cpp + +++ stablehlo/stablehlo/dialect/ChloBytecode.cpp diff --git a/third_party/xla/third_party/shardy/workspace.bzl b/third_party/xla/third_party/shardy/workspace.bzl index dd328941030d9f..2e45bac52ee8ca 100644 --- a/third_party/xla/third_party/shardy/workspace.bzl +++ b/third_party/xla/third_party/shardy/workspace.bzl @@ -18,8 +18,8 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") def repo(): - SHARDY_COMMIT = "08e079b1c5f5b0a3a0c696abc5c56b6ac436b835" - SHARDY_SHA256 = "b3de5417edb836ddcda6c1800bb0a5782c58110ff9567675a82ada9562d1033b" + SHARDY_COMMIT = "dc6f5f05c411667e7f2f6a2d3552276c6f4b1a43" + SHARDY_SHA256 = "0130b0bfa450f1fb2a44da6264aec45e81607896790a799d5fc75c6ff08bafa2" tf_http_archive( name = "shardy", diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index fb55a869629448..c0fd93daf55cd0 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -123,6 +123,328 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgU /// Extracts integer values from the attribute |elements|. SmallVector extract1DVector(DenseIntElementsAttr elements); +diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir +--- stablehlo/stablehlo/conversions/tosa/tests/binary.mlir ++++ stablehlo/stablehlo/conversions/tosa/tests/binary.mlir +@@ -38,7 +38,7 @@ + + // CHECK-LABEL: @concatenate + func.func @concatenate(%arg0 : tensor<3x3xf32>, %arg1 : tensor<3x3xf32>) -> tensor<6x3xf32> { +- // CHECK: tosa.concat %arg0, %arg1 {axis = 0 : i32} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32> ++ // CHECK: tosa.concat %arg0, %arg1 axis(0) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32> + %0 = "stablehlo.concatenate"(%arg0, %arg1) {dimension = 0 : i64} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32> + return %0 : tensor<6x3xf32> + } +@@ -61,10 +61,10 @@ + + // CHECK-LABEL: @dot_vector_vector + func.func @dot_vector_vector(%arg0 : tensor<3xf32>, %arg1 : tensor<3xf32>) -> tensor { +- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<> : tensor<0xindex>} : () -> !tosa.shape<0> +- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<> : tensor<0xindex>) : () -> !tosa.shape<0> ++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] +- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] + // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] + // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] +@@ -74,10 +74,10 @@ + + // CHECK-LABEL: @dot_vector_matrix + func.func @dot_vector_matrix(%arg0 : tensor<2xf32>, %arg1 : tensor<2x3xf32>) -> tensor<3xf32> { +- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<3> : tensor<1xindex>} : () -> !tosa.shape<1> +- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 2]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<3> : tensor<1xindex>) : () -> !tosa.shape<1> ++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 2]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] +- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] + // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] + // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] +@@ -87,10 +87,10 @@ + + // CHECK-LABEL: @dot_matrix_vector + func.func @dot_matrix_vector(%arg0 : tensor<2x3xf32>, %arg1 : tensor<3xf32>) -> tensor<2xf32> { +- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<2> : tensor<1xindex>} : () -> !tosa.shape<1> +- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<2> : tensor<1xindex>) : () -> !tosa.shape<1> ++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] +- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] + // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] + // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] +@@ -100,10 +100,10 @@ + + // CHECK-LABEL: @dot_matrix_matrix + func.func @dot_matrix_matrix(%arg0 : tensor<2x3xf32>, %arg1 : tensor<3x4xf32>) -> tensor<2x4xf32> { +- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<[2, 4]> : tensor<2xindex>} : () -> !tosa.shape<2> +- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<[2, 4]> : tensor<2xindex>) : () -> !tosa.shape<2> ++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] +- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 4]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 4]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] + // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] + // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] +@@ -113,10 +113,10 @@ + + // CHECK-LABEL: @dot_general_vector_vector + func.func @dot_general_vector_vector(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>) -> tensor { +- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<> : tensor<0xindex>} : () -> !tosa.shape<0> +- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<> : tensor<0xindex>) : () -> !tosa.shape<0> ++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] +- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] + // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] + // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] +@@ -126,10 +126,10 @@ + + // CHECK-LABEL: @dot_general_vector_matrix + func.func @dot_general_vector_matrix(%arg0: tensor<2xf32>, %arg1: tensor<2x3xf32>) -> tensor<3xf32> { +- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<3> : tensor<1xindex>} : () -> !tosa.shape<1> +- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 2]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<3> : tensor<1xindex>) : () -> !tosa.shape<1> ++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 2]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] +- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] + // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] + // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] +@@ -139,10 +139,10 @@ + + // CHECK-LABEL: @dot_general_matrix_vector + func.func @dot_general_matrix_vector(%arg0: tensor<2x3xf32>, %arg1: tensor<3xf32>) -> tensor<2xf32> { +- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<2> : tensor<1xindex>} : () -> !tosa.shape<1> +- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<2> : tensor<1xindex>) : () -> !tosa.shape<1> ++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] +- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] + // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] + // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] +@@ -152,10 +152,10 @@ + + // CHECK-LABEL: @dot_general_matrix_matrix + func.func @dot_general_matrix_matrix(%arg0: tensor<2x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<2x4xf32> { +- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<[2, 4]> : tensor<2xindex>} : () -> !tosa.shape<2> +- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<[2, 4]> : tensor<2xindex>) : () -> !tosa.shape<2> ++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]] +- // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 4]> : tensor<3xindex>} : () -> !tosa.shape<3> ++ // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 4]> : tensor<3xindex>) : () -> !tosa.shape<3> + // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]] + // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]] + // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]] +diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir b/stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir +--- stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir ++++ stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir +@@ -11,10 +11,10 @@ + // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1431655765> : tensor<1xi32>}> + // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> + // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> +- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: %[[V2:.+]] = stablehlo.add %[[V0]], %[[V1]] : tensor<2x2xi32> +- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> + %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) + -> tensor<2x2x!quant.uniform> +@@ -32,10 +32,10 @@ + // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1431655765> : tensor<1xi32>}> + // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> + // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> +- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: %[[V2:.+]] = stablehlo.subtract %[[V0]], %[[V1]] : tensor<2x2xi32> +- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> + %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) + -> tensor<2x2x!quant.uniform> +@@ -52,10 +52,10 @@ + // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1717986918> : tensor<1xi32>}> + // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> + // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> +- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: %[[V2:.+]] = stablehlo.multiply %[[V0]], %[[V1]] : tensor<2x2xi32> +- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_1]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_1]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> + %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) + -> tensor<2x2x!quant.uniform> +@@ -74,10 +74,10 @@ + // CHECK-DAG: %[[ZP_MINUS_2:.+]] = "tosa.const"() <{values = dense<-2> : tensor<1xi8>}> + // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> + // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> +- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: %[[V2:.+]] = stablehlo.divide %[[V0]], %[[V1]] : tensor<2x2xi32> +- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_3]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_3]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> + %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) + -> tensor<2x2x!quant.uniform> +@@ -97,10 +97,10 @@ + // CHECK-DAG: %[[SHIFT12:.+]] = "tosa.const"() <{values = dense<12> : tensor<1xi8>}> + // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> + // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> +- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: %[[V2:.+]] = stablehlo.maximum %[[V0]], %[[V1]] : tensor<2x2xi32> +- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> + %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) + -> tensor<2x2x!quant.uniform> +@@ -120,10 +120,10 @@ + // CHECK-DAG: %[[SHIFT12:.+]] = "tosa.const"() <{values = dense<12> : tensor<1xi8>}> + // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> + // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> +- // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +- // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: %[[V2:.+]] = stablehlo.minimum %[[V0]], %[[V1]] : tensor<2x2xi32> +- // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform> + %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<2x2x!quant.uniform>, tensor<2x2x!quant.uniform>) + -> tensor<2x2x!quant.uniform> +@@ -140,9 +140,9 @@ + // CHECK-DAG: %[[SHIFT30:.+]] = "tosa.const"() <{values = dense<30> : tensor<1xi8>}> + // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> + // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> +- // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: %[[V1:.+]] = stablehlo.abs %[[V0]] : tensor<20x20xi32> +- // CHECK: %[[V3:.+]] = tosa.rescale %[[V1]], %[[MULTIPLIER_1]], %[[SHIFT33]], %[[ZP_0]], %[[ZP_MINUS_128]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK: %[[V3:.+]] = tosa.rescale %[[V1]], %[[MULTIPLIER_1]], %[[SHIFT33]], %[[ZP_0]], %[[ZP_MINUS_128]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: return %[[V3]] : tensor<20x20x!quant.uniform> + %0 = "stablehlo.abs"(%arg0) : (tensor<20x20x!quant.uniform>) -> tensor<20x20x!quant.uniform> + return %0 : tensor<20x20x!quant.uniform> +@@ -159,8 +159,8 @@ + // CHECK-DAG: %[[SHIFT12:.+]] = "tosa.const"() <{values = dense<12> : tensor<1xi8>}> + // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}> + // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> +- // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +- // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: %[[V2:.+]] = stablehlo.compare GE, %[[V0]], %[[V1]], TOTALORDER : + // CHECK: return %[[V2]] + %0 = stablehlo.compare GE, %arg0, %arg1, TOTALORDER : (tensor<20x20x!quant.uniform>, tensor<20x20x!quant.uniform>) -> tensor<20x20xi1> +@@ -177,8 +177,8 @@ + // CHECK-DAG: %[[SHIFT15:.+]] = "tosa.const"() <{values = dense<15> : tensor<1xi8>}> + // CHECK-DAG: %[[ZP16_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi16>}> + // CHECK-DAG: %[[ZP32_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}> +- // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT17]], %[[ZP16_0]], %[[ZP32_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} +- // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT15]], %[[ZP16_0]], %[[ZP32_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} ++ // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT17]], %[[ZP16_0]], %[[ZP32_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) ++ // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT15]], %[[ZP16_0]], %[[ZP32_0]] scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) + // CHECK: %[[V2:.+]] = stablehlo.compare LT, %[[V0]], %[[V1]], TOTALORDER : + // CHECK: return %[[V2]] + %0 = stablehlo.compare LT, %arg0, %arg1, TOTALORDER : (tensor<20x20x!quant.uniform>, tensor<20x20x!quant.uniform>) -> tensor<20x20xi1> +diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir b/stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir +--- stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir ++++ stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir +@@ -7,7 +7,7 @@ + %shift = "tosa.const"() {values = dense<13> : tensor<1xi8>} : () -> tensor<1xi8> + %input_zp = "tosa.const"() {values = dense<-1> : tensor<1xi8>} : () -> tensor<1xi8> + %output_zp = "tosa.const"() {values = dense<0> : tensor<1xi32>} : () -> tensor<1xi32> +- %0 = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} : ++ %0 = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp scale32(true) rounding_mode per_channel(false) input_unsigned(false) output_unsigned(false) : + (tensor<2x2x!quant.uniform>, tensor<1xi32>, tensor<1xi8>, tensor<1xi8>, tensor<1xi32>) -> tensor<2x2xi32> + + // convert input quantized type to storage type +diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir +--- stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir ++++ stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir +@@ -18,7 +18,7 @@ + func.func @iota_dimension_0() -> tensor<4x8xf32> { + // CHECK-DAG: %[[VAR0:.*]] = "tosa.const"() + // CHECK-SAME{LITERAL}: <{values = dense<[[0.000000e+00], [1.000000e+00], [2.000000e+00], [3.000000e+00]]> : tensor<4x1xf32>}> : () -> tensor<4x1xf32> +- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 8]> : vector<2xindex>} : () -> !tosa.shape<2> ++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 8]> : vector<2xindex>) : () -> !tosa.shape<2> + // CHECK-DAG: %[[VAR2:.*]] = tosa.tile %[[VAR0]], %[[VAR1]] + %0 = "stablehlo.iota"() {iota_dimension = 0 : i64} : () -> (tensor<4x8xf32>) + return %0 : tensor<4x8xf32> +@@ -28,7 +28,7 @@ + func.func @iota_dimension_1() -> tensor<4x8xi32> { + // CHECK-DAG: %[[VAR0:.*]] = "tosa.const"() + // CHECK-SAME{LITERAL}: <{values = dense<[[0, 1, 2, 3, 4, 5, 6, 7]]> : tensor<1x8xi32>}> : () -> tensor<1x8xi32> +- // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[4, 1]> : vector<2xindex>} : () -> !tosa.shape<2> ++ // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[4, 1]> : vector<2xindex>) : () -> !tosa.shape<2> + // CHECK-DAG: %[[VAR2:.*]] = tosa.tile %[[VAR0]], %[[VAR1]] + %0 = "stablehlo.iota"() {iota_dimension = 1 : i64} : () -> (tensor<4x8xi32>) + return %0 : tensor<4x8xi32> +diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir +--- stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir ++++ stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir +@@ -2,7 +2,7 @@ + + // CHECK-LABEL: @concatenate + func.func @concatenate(%arg0 : tensor<5x2xf32>, %arg1 : tensor<5x5xf32>, %arg2 : tensor<5x7xf32>) -> tensor<5x14xf32> { +- // CHECK: tosa.concat %arg0, %arg1, %arg2 {axis = 1 : i32} : (tensor<5x2xf32>, tensor<5x5xf32>, tensor<5x7xf32>) -> tensor<5x14xf32> ++ // CHECK: tosa.concat %arg0, %arg1, %arg2 axis(1) : (tensor<5x2xf32>, tensor<5x5xf32>, tensor<5x7xf32>) -> tensor<5x14xf32> + %0 = "stablehlo.concatenate"(%arg0, %arg1, %arg2) {dimension = 1 : i64} : (tensor<5x2xf32>, tensor<5x5xf32>, tensor<5x7xf32>) -> tensor<5x14xf32> + return %0 : tensor<5x14xf32> + } +diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/unary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/unary.mlir +--- stablehlo/stablehlo/conversions/tosa/tests/unary.mlir ++++ stablehlo/stablehlo/conversions/tosa/tests/unary.mlir +@@ -79,8 +79,8 @@ + + // CHECK-LABEL: @slice + func.func @slice(%arg : tensor<4x3xf32>) -> tensor<2x2xf32> { +- // CHECK-DAG: %[[SIZE:.*]] = tosa.const_shape {values = dense<[2, 1]> : tensor<2xindex>} : () -> !tosa.shape<2> +- // CHECK-DAG: %[[START:.*]] = tosa.const_shape {values = dense<2> : tensor<2xindex>} : () -> !tosa.shape<2> ++ // CHECK-DAG: %[[SIZE:.*]] = tosa.const_shape values(dense<[2, 1]> : tensor<2xindex>) : () -> !tosa.shape<2> ++ // CHECK-DAG: %[[START:.*]] = tosa.const_shape values(dense<2> : tensor<2xindex>) : () -> !tosa.shape<2> + // CHECK: tosa.slice %arg0, %[[SIZE]], %[[START]] + %0 = "stablehlo.slice"(%arg) { + start_indices = array, +@@ -123,7 +123,7 @@ + + // CHECK-LABEL: @transpose + func.func @transpose(%arg0: tensor<1x2x3xf32>) -> tensor<3x2x1xf32> { +- // CHECK: %[[VAR0:.*]] = tosa.transpose %arg0 {perms = array} ++ // CHECK: %[[VAR0:.*]] = tosa.transpose %arg0 perms([2, 1, 0]) + %0 = "stablehlo.transpose"(%arg0) {permutation = array} : (tensor<1x2x3xf32>) -> tensor<3x2x1xf32> + return %0 : tensor<3x2x1xf32> + } +@@ -158,7 +158,7 @@ + + // CHECK-LABEL: @reshape + func.func @reshape(%arg0 : tensor<2x3xf32>) -> tensor<6xf32> { +- // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<6> : tensor<1xindex>} : () -> !tosa.shape<1> ++ // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<6> : tensor<1xindex>) : () -> !tosa.shape<1> + // CHECK: tosa.reshape %arg0, %[[VAR0]] + %0 = "stablehlo.reshape"(%arg0) : (tensor<2x3xf32>) -> tensor<6xf32> + return %0 : tensor<6xf32> diff --ruN a/stablehlo/stablehlo/dialect/ChloBytecode.cpp b/stablehlo/stablehlo/dialect/ChloBytecode.cpp --- stablehlo/stablehlo/dialect/ChloBytecode.cpp +++ stablehlo/stablehlo/dialect/ChloBytecode.cpp @@ -2176,3 +2498,14 @@ diff --ruN a/stablehlo/stablehlo/transforms/optimization/Passes.h b/stablehlo/st } // namespace stablehlo } // namespace mlir +diff --ruN a/stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp b/stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp +--- stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp ++++ stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp +@@ -1766,6 +1766,7 @@ + void populateStablehloHloImportCanonicalizationPatterns( + MLIRContext* context, RewritePatternSet* patterns, + const StablehloAggressiveSimplificationPassOptions& options) { ++ patterns->add(context, options); + patterns->add(context); + patterns->add(context, options); + } diff --git a/third_party/xla/third_party/triton/common/llvm_cl986804242.patch b/third_party/xla/third_party/triton/common/llvm_cl986804242.patch new file mode 100644 index 00000000000000..80bfb47d976837 --- /dev/null +++ b/third_party/xla/third_party/triton/common/llvm_cl986804242.patch @@ -0,0 +1,26 @@ +--- a/test/Conversion/amd/builtin_func_to_llvm.mlir ++++ b/test/Conversion/amd/builtin_func_to_llvm.mlir +@@ -1,7 +1,7 @@ +-// RUN: triton-opt %s --convert-triton-amdgpu-to-llvm="gfx-arch=gfx942 ftz=True" --convert-builtin-func-to-llvm="ftz=True" | FileCheck %s --check-prefix=LLVM_FTZ +-// RUN: triton-opt %s --convert-triton-amdgpu-to-llvm="gfx-arch=gfx950 ftz=True" --convert-builtin-func-to-llvm="ftz=True" | FileCheck %s --check-prefix=LLVM_FTZ +-// RUN: triton-opt %s --convert-triton-amdgpu-to-llvm="gfx-arch=gfx942 ftz=False" --convert-builtin-func-to-llvm="ftz=False" | FileCheck %s --check-prefix=LLVM_NO_FTZ +-// RUN: triton-opt %s --convert-triton-amdgpu-to-llvm="gfx-arch=gfx950 ftz=False" --convert-builtin-func-to-llvm="ftz=False" | FileCheck %s --check-prefix=LLVM_NO_FTZ ++// RUN: triton-opt %s --convert-triton-amdgpu-to-llvm="gfx-arch=gfx942 ftz=true" --convert-builtin-func-to-llvm="ftz=true" | FileCheck %s --check-prefix=LLVM_FTZ ++// RUN: triton-opt %s --convert-triton-amdgpu-to-llvm="gfx-arch=gfx950 ftz=true" --convert-builtin-func-to-llvm="ftz=true" | FileCheck %s --check-prefix=LLVM_FTZ ++// RUN: triton-opt %s --convert-triton-amdgpu-to-llvm="gfx-arch=gfx942 ftz=false" --convert-builtin-func-to-llvm="ftz=false" | FileCheck %s --check-prefix=LLVM_NO_FTZ ++// RUN: triton-opt %s --convert-triton-amdgpu-to-llvm="gfx-arch=gfx950 ftz=false" --convert-builtin-func-to-llvm="ftz=false" | FileCheck %s --check-prefix=LLVM_NO_FTZ + + #blocked = #ttg.blocked<{sizePerThread = [1], threadsPerWarp = [64], warpsPerCTA = [1], order = [0]}> + + +--- a/test/Conversion/amd/math-denorm-handling.mlir ++++ b/test/Conversion/amd/math-denorm-handling.mlir +@@ -1,5 +1,5 @@ +-// RUN: triton-opt %s -split-input-file --convert-triton-amdgpu-to-llvm="gfx-arch=gfx942 ftz=True" | FileCheck %s --check-prefixes=COMMON,LLVM_FTZ +-// RUN: triton-opt %s -split-input-file --convert-triton-amdgpu-to-llvm="gfx-arch=gfx942 ftz=False" | FileCheck %s --check-prefixes=COMMON,LLVM_NO_FTZ ++// RUN: triton-opt %s -split-input-file --convert-triton-amdgpu-to-llvm="gfx-arch=gfx942 ftz=true" | FileCheck %s --check-prefixes=COMMON,LLVM_FTZ ++// RUN: triton-opt %s -split-input-file --convert-triton-amdgpu-to-llvm="gfx-arch=gfx942 ftz=false" | FileCheck %s --check-prefixes=COMMON,LLVM_NO_FTZ + + + #blocked = #ttg.blocked<{sizePerThread = [1], threadsPerWarp = [64], warpsPerCTA = [1], order = [0]}> + diff --git a/third_party/xla/third_party/triton/common/series.bzl b/third_party/xla/third_party/triton/common/series.bzl index 0aca7c09670566..3c729811758214 100644 --- a/third_party/xla/third_party/triton/common/series.bzl +++ b/third_party/xla/third_party/triton/common/series.bzl @@ -63,5 +63,6 @@ common_patch_list = [ "//third_party/triton:common/llvm_cl979179688.patch", "//third_party/triton:common/llvm_cl979526570.patch", "//third_party/triton:common/llvm_cl983115707.patch", + "//third_party/triton:common/llvm_cl986804242.patch", # Add new patches just above this line ] diff --git a/third_party/xla/third_party/tsl/tsl/platform/BUILD b/third_party/xla/third_party/tsl/tsl/platform/BUILD index 9482149821cd3d..c4570495a91101 100644 --- a/third_party/xla/third_party/tsl/tsl/platform/BUILD +++ b/third_party/xla/third_party/tsl/tsl/platform/BUILD @@ -1214,7 +1214,6 @@ tsl_cc_test( ":blocking_counter", ":random", ":unbounded_work_queue", - "@com_google_absl//absl/memory", "@com_google_absl//absl/synchronization", "@com_google_googletest//:gtest_main", "@xla//xla/tsl/platform:env", diff --git a/third_party/xla/third_party/tsl/tsl/platform/unbounded_work_queue_test.cc b/third_party/xla/third_party/tsl/tsl/platform/unbounded_work_queue_test.cc index ea24a27f3f5f5e..32e5e32419ff49 100644 --- a/third_party/xla/third_party/tsl/tsl/platform/unbounded_work_queue_test.cc +++ b/third_party/xla/third_party/tsl/tsl/platform/unbounded_work_queue_test.cc @@ -18,7 +18,6 @@ limitations under the License. #include #include -#include "absl/memory/memory.h" #include "absl/synchronization/mutex.h" #include "xla/tsl/platform/env.h" #include "xla/tsl/platform/test.h" @@ -105,5 +104,26 @@ TEST_F(UnboundedWorkQueueTest, RacyDestructor) { EXPECT_LE(NumClosuresExecuted(), num_closures); } +TEST(UnboundedWorkQueueOptionsTest, CustomStackSize) { + ThreadOptions thread_options; + thread_options.stack_size = 2 * 1024 * 1024; // 2 MB + UnboundedWorkQueue queue(Env::Default(), "custom_stack_test", thread_options); + BlockingCounter counter(10); + for (int i = 0; i < 10; ++i) { + queue.Schedule([&counter]() { counter.DecrementCount(); }); + } + counter.Wait(); +} + +TEST(UnboundedWorkQueueOptionsTest, DefaultThreadOptions) { + UnboundedWorkQueue queue(Env::Default(), "default_options_test", + ThreadOptions{}); + BlockingCounter counter(10); + for (int i = 0; i < 10; ++i) { + queue.Schedule([&counter]() { counter.DecrementCount(); }); + } + counter.Wait(); +} + } // namespace } // namespace tsl diff --git a/third_party/xla/workspace2.bzl b/third_party/xla/workspace2.bzl index cd87c0b18cf762..8a99921e4d611f 100644 --- a/third_party/xla/workspace2.bzl +++ b/third_party/xla/workspace2.bzl @@ -373,9 +373,9 @@ def _tf_repositories(): tf_http_archive( name = "curl", build_file = "//third_party:curl.BUILD", - sha256 = "d9b327997999045a24cda50f3983e69e51c516bd8be6ef9842fc7f99135e33bb", - strip_prefix = "curl-8.21.0", - urls = tf_mirror_urls("https://curl.se/download/curl-8.21.0.tar.gz"), + sha256 = "d54dd598bf05927a726deb38df31c6a255ba83ff1de57c5d1464dac3ed8f44a1", + strip_prefix = "curl-8.22.0", + urls = tf_mirror_urls("https://curl.se/download/curl-8.22.0.tar.gz"), ) tf_http_archive( diff --git a/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc b/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc index fee83871c42993..84507d32cedc7c 100644 --- a/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc +++ b/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc @@ -237,7 +237,8 @@ void AddScalarOptimizationPasses(mlir::OpPassManager& pm, // The final lowering passes common to both scalar and tiled kernels. // These passes are primarily responsible for lowering individual ops to // their LLVM equivalent. -void AddGenericLoweringPasses(mlir::OpPassManager& pm, bool fast_min_max) { +void AddGenericLoweringPasses(mlir::OpPassManager& pm, bool fast_min_max, + absl::string_view cpu_features) { pm.addNestedPass(emitters::createSimplifyArithPass( GetSimplifyArithPassOptions(fast_min_max))); pm.addPass(emitters::createExpandIntegerPowerPass()); @@ -259,7 +260,9 @@ void AddGenericLoweringPasses(mlir::OpPassManager& pm, bool fast_min_max) { pm.addPass(emitters::createEraseDeadFunctionsPass()); pm.addPass(mlir::createLowerAffinePass()); pm.addPass(mlir::createSCFToControlFlowPass()); - pm.addPass(emitters::createLowerXlaIntrinsicLibPass()); + emitters::LowerXlaIntrinsicLibPassOptions intrinsic_lib_options; + intrinsic_lib_options.cpu_features_ = std::string(cpu_features); + pm.addPass(emitters::createLowerXlaIntrinsicLibPass(intrinsic_lib_options)); pm.addNestedPass(CreateConvertMathToLLVMPass()); pm.addPass(mlir::createConvertMathToLibmPass()); pm.addPass(emitters::createLowerToLLVMCPUPass()); @@ -273,7 +276,8 @@ void AddGenericLoweringPasses(mlir::OpPassManager& pm, bool fast_min_max) { // on scalar instructions extracted/inserted from tensor types. // The resulting IR can then be translated to native LLVM. static void AddScalarLoweringPasses(mlir::OpPassManager& pm, - int32_t vector_width, bool fast_min_max) { + int32_t vector_width, bool fast_min_max, + absl::string_view cpu_features) { pm.addNestedPass( emitters::createConvertPureCallOpsPass()); pm.addPass(cpu::createLowerToLLVMPass( @@ -289,7 +293,7 @@ static void AddScalarLoweringPasses(mlir::OpPassManager& pm, pm.addPass(mlir::createCanonicalizerPass()); pm.addPass(mlir::createCSEPass()); - AddGenericLoweringPasses(pm, fast_min_max); + AddGenericLoweringPasses(pm, fast_min_max, cpu_features); } void AddBufferizationPasses(mlir::OpPassManager& pm, bool msan_enabled) { @@ -353,12 +357,14 @@ class ModuleCallbackPass // Optimizations passes for the tiled emitter. // This is currently very simple but will grow to include tiled optimizations // such as transpose hoisting and dimension reduction. -void AddXtileToVectorPasses(mlir::OpPassManager& pm, bool msan_enabled) { +void AddXtileToVectorPasses(mlir::OpPassManager& pm, bool msan_enabled, + int32_t vector_width) { pm.addPass(xtile::createVerifyLegalXTileOpsPass()); emitters::RegisterOptimizationPasses(pm); - pm.addPass(cpu::createLowerXTileEntryPass()); + pm.addPass(cpu::createLowerXTileEntryPass( + cpu::LowerXTileEntryPassOptions{/*prefer_vector_width=*/vector_width})); pm.addNestedPass( mlir::stablehlo::createStablehloTargetIndependentOptimizationPass()); @@ -399,7 +405,7 @@ void AddXtileToVectorPasses(mlir::OpPassManager& pm, bool msan_enabled) { } // Optimizations passes for the tiled emitter. -void AddNewXtileToVectorPasses(mlir::OpPassManager& pm) { +void AddNewXtileToVectorPasses(mlir::OpPassManager& pm, int32_t vector_width) { pm.addPass(xtile::createVerifyLegalXTileOpsPass()); emitters::RegisterOptimizationPasses(pm); @@ -410,7 +416,8 @@ void AddNewXtileToVectorPasses(mlir::OpPassManager& pm) { pm.addPass(cpu::createLegalizeNarrowFloatStoragePass()); pm.addPass(mlir::createCanonicalizerPass()); pm.addPass(cpu::createVectorizeXTilePass()); - pm.addPass(cpu::createLowerXTileEntryPass()); + pm.addPass(cpu::createLowerXTileEntryPass( + cpu::LowerXTileEntryPassOptions{/*prefer_vector_width=*/vector_width})); pm.addNestedPass(emitters::createSimplifyArithPass( GetSimplifyArithPassOptions(/*fast_min_max=*/false))); @@ -434,7 +441,8 @@ void AddNewXtileToVectorPasses(mlir::OpPassManager& pm) { // Lowering passes for the tiled emitter. // The input IR is from the xtile dialect which uses tensors that are converted // first to the vector dialect and then to LLVM. -void AddVectorToLLVMPasses(mlir::OpPassManager& pm, bool fast_min_max) { +void AddVectorToLLVMPasses(mlir::OpPassManager& pm, bool fast_min_max, + absl::string_view cpu_features) { pm.addPass(cpu::createVectorToScalarPass()); pm.addPass(cpu::createMemrefCopyToLoopsPass()); pm.addPass(cpu::createLowerToLLVMPass()); @@ -456,14 +464,15 @@ void AddVectorToLLVMPasses(mlir::OpPassManager& pm, bool fast_min_max) { pm.addPass(emitters::createSafeIntegerArithmeticPass()); - AddGenericLoweringPasses(pm, fast_min_max); + AddGenericLoweringPasses(pm, fast_min_max, cpu_features); } // Lowering passes for the new tiled emitter. // The input IR is from the xtile dialect which uses tensors that are converted // first to the vector dialect and then to LLVM. void AddNewVectorToLLVMPasses(mlir::OpPassManager& pm, bool fast_min_max, - int32_t vector_width) { + int32_t vector_width, + absl::string_view cpu_features) { // Get rid of 0d vectors. pm.addPass(cpu::createVectorToScalarPass()); // Get rid of multi-dimensional vectors. @@ -490,7 +499,7 @@ void AddNewVectorToLLVMPasses(mlir::OpPassManager& pm, bool fast_min_max, pm.addPass(mlir::memref::createExpandStridedMetadataPass()); pm.addPass(emitters::createSafeIntegerArithmeticPass()); - AddGenericLoweringPasses(pm, fast_min_max); + AddGenericLoweringPasses(pm, fast_min_max, cpu_features); } FusionCompiler::FusionCompiler(mlir::MLIRContext* context, Options options, @@ -514,16 +523,17 @@ FusionCompiler::FusionCompiler(mlir::MLIRContext* context, Options options, std::make_unique(hlo_module_, "post-optimization")); } AddScalarLoweringPasses(scalar_pass_manager_, options_.vector_width, - options_.fast_min_max); + options_.fast_min_max, options_.cpu_features); scalar_pass_manager_.addInstrumentation( std::make_unique()); // Tiled passes. if (options_.use_new_xtile_lowering) { - AddNewXtileToVectorPasses(tiled_pass_manager_); + AddNewXtileToVectorPasses(tiled_pass_manager_, options_.vector_width); } else { - AddXtileToVectorPasses(tiled_pass_manager_, options_.msan_enabled); + AddXtileToVectorPasses(tiled_pass_manager_, options_.msan_enabled, + options_.vector_width); } if (should_dump_mlir_passes) { tiled_pass_manager_.addPass( @@ -531,9 +541,10 @@ FusionCompiler::FusionCompiler(mlir::MLIRContext* context, Options options, } if (options_.use_new_xtile_lowering) { AddNewVectorToLLVMPasses(tiled_pass_manager_, options_.fast_min_max, - options_.vector_width); + options_.vector_width, options_.cpu_features); } else { - AddVectorToLLVMPasses(tiled_pass_manager_, options_.fast_min_max); + AddVectorToLLVMPasses(tiled_pass_manager_, options_.fast_min_max, + options_.cpu_features); } tiled_pass_manager_.enableVerifier(should_verify); tiled_pass_manager_.addInstrumentation( diff --git a/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.h b/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.h index 59ac34ce92b49a..34a65bf44c08d1 100644 --- a/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.h +++ b/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.h @@ -47,6 +47,10 @@ class FusionCompiler { llvm::FastMathFlags fast_math_flags; bool use_new_xtile_lowering = false; bool msan_enabled = false; + // LLVM target feature string of the target machine that will compile the + // resulting LLVM module. Selects feature-gated xla.* intrinsic variants, so + // it must not claim features the target machine lacks. + std::string cpu_features; }; FusionCompiler(mlir::MLIRContext* context, Options options, @@ -81,11 +85,14 @@ class FusionCompiler { // Xtile CPU pipeline contains two stages, the first is the conversion from // Xtile to the vector dialect, the second is the conversion from the vector // dialect to LLVM. -void AddXtileToVectorPasses(mlir::OpPassManager& pm, bool msan_enabled); -void AddNewXtileToVectorPasses(mlir::OpPassManager& pm); -void AddVectorToLLVMPasses(mlir::OpPassManager& pm, bool fast_min_max); +void AddXtileToVectorPasses(mlir::OpPassManager& pm, bool msan_enabled, + int32_t vector_width); +void AddNewXtileToVectorPasses(mlir::OpPassManager& pm, int32_t vector_width); +void AddVectorToLLVMPasses(mlir::OpPassManager& pm, bool fast_min_max, + absl::string_view cpu_features); void AddNewVectorToLLVMPasses(mlir::OpPassManager& pm, bool fast_min_max, - int32_t vector_width); + int32_t vector_width, + absl::string_view cpu_features); } // namespace xla::cpu diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/broadcast/broadcast.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/broadcast/broadcast.hlo index 6eaa228dec7d8f..95e210540f6942 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/broadcast/broadcast.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/broadcast/broadcast.hlo @@ -29,7 +29,7 @@ ENTRY main { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG0:[^:]+]]: memref<32xf32>, // CHECK-SAME: %[[ARG1:[^:]+]]: memref<8x32xf32>, -// CHECK-SAME: %{{.*}}: index) +// CHECK-SAME: %{{.*}}: index {xla.range = [0 : index, 0 : index]}) // CHECK: scf.for %[[IDX:.*]] = %{{.*}} to %{{.*}} step %{{.*}} { // CHECK: %[[OFF1:.*]] = xla.apply_indexing diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/broadcast/broadcast_0d.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/broadcast/broadcast_0d.hlo index 6eeec1aadd5706..14ef07c8f492cf 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/broadcast/broadcast_0d.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/broadcast/broadcast_0d.hlo @@ -28,7 +28,7 @@ ENTRY main { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG0:[^:]+]]: memref, // CHECK-SAME: %[[ARG1:[^:]+]]: memref<128x64xi32>, -// CHECK-SAME: %{{.*}}: index) +// CHECK-SAME: %{{.*}}: index {xla.range = [0 : index, 0 : index]}) // CHECK: scf.for %[[IDX:.*]] = %{{.*}} to %{{.*}} step %{{.*}} { // CHECK: %[[EXT:.*]] = xtile.extract %[[ARG0]][] [] [] : memref -> tensor @@ -38,7 +38,7 @@ ENTRY main { // NEW-VECTOR: func.func @wrapped_fusion_impl( // NEW-VECTOR-SAME: %[[ARG0:[^:]+]]: memref, // NEW-VECTOR-SAME: %[[ARG1:[^:]+]]: memref<128x64xi32>, -// NEW-VECTOR-SAME: %{{.*}}: index) +// NEW-VECTOR-SAME: %{{.*}}: index {xla.range = [0 : index, 0 : index]}) // NEW-VECTOR: scf.for %[[IDX:.*]] = %{{.*}} to %{{.*}} step %{{.*}} { // NEW-VECTOR: %[[READ:.*]] = vector.transfer_read %[[ARG0]][], %{{.*}} : memref, vector // NEW-VECTOR: %[[BCAST:.*]] = vector.broadcast %[[READ]] : vector to vector<4x4xi32> diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/compare.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/compare.hlo index 4e171177beba93..2999953fb6ded9 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/compare.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/compare.hlo @@ -26,7 +26,7 @@ fusion { // CHECK-SAME: %[[LHS:[^:]+]]: memref<128x256xf32>, // CHECK-SAME: %[[RHS:[^:]+]]: memref<128x256xf32>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xi8>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK-SAME: attributes {num_opaque_args = 0 : i32} // CHECK: %[[C0:.*]] = arith.constant 0 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_abs.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_abs.hlo index b926aef6b25fec..a689cfdb4aed3d 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_abs.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_abs.hlo @@ -17,7 +17,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[IN:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK: %[[T_0:.*]] = xtile.extract %[[IN]] // CHECK-SAME: : memref<128x256xcomplex> -> tensor<1x4xcomplex> @@ -29,7 +29,7 @@ fusion { // NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl( // NEW-VECTOR-SAME: %[[IN:[^:]+]]: memref<128x256x2xf32>, // NEW-VECTOR-SAME: %[[OUT:[^:]+]]: memref<128x256xf32>, -// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index) +// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // NEW-VECTOR: vector.transfer_read %[[IN]] // NEW-VECTOR: math.absf diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_add.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_add.hlo index 25ec44b1f4b7c6..ee63b3d4185685 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_add.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_add.hlo @@ -22,7 +22,7 @@ fusion { // CHECK-SAME: %[[LHS:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[RHS:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xcomplex>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK: %[[T_0:.*]] = xtile.extract %[[LHS]] // CHECK-SAME: : memref<128x256xcomplex> -> tensor<1x4xcomplex> @@ -38,7 +38,7 @@ fusion { // NEW-VECTOR-SAME: %[[LHS:[^:]+]]: memref<128x256x2xf32>, // NEW-VECTOR-SAME: %[[RHS:[^:]+]]: memref<128x256x2xf32>, // NEW-VECTOR-SAME: %[[OUT:[^:]+]]: memref<128x256x2xf32>, -// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index) +// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // NEW-VECTOR: %[[LHS_VEC:.*]] = vector.transfer_read %[[LHS]] // NEW-VECTOR-SAME: : memref<128x256x2xf32>, vector<1x4x2xf32> diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_compare.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_compare.hlo index 9c5247a5544d3b..44f569e2973d88 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_compare.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_compare.hlo @@ -19,7 +19,7 @@ fusion { // CHECK-SAME: %[[LHS:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[RHS:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xi8>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK: %[[T_0:.*]] = xtile.extract %[[LHS]] // CHECK-SAME: : memref<128x256xcomplex> -> tensor<1x4xcomplex> @@ -34,7 +34,7 @@ fusion { // NEW-VECTOR-SAME: %[[LHS:[^:]+]]: memref<128x256x2xf32>, // NEW-VECTOR-SAME: %[[RHS:[^:]+]]: memref<128x256x2xf32>, // NEW-VECTOR-SAME: %[[OUT:[^:]+]]: memref<128x256xi8>, -// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index) +// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // NEW-VECTOR: vector.transfer_read %[[LHS]] // NEW-VECTOR: vector.transfer_read %[[RHS]] diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_complex.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_complex.hlo index 1d922bcf87e41e..f29c34dadc9b98 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_complex.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_complex.hlo @@ -19,7 +19,7 @@ fusion { // CHECK-SAME: %[[REAL:[^:]+]]: memref<128x256xf32>, // CHECK-SAME: %[[IMAG:[^:]+]]: memref<128x256xf32>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xcomplex>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK: %[[T_0:.*]] = xtile.extract %[[REAL]] // CHECK-SAME: : memref<128x256xf32> -> tensor<1x4xf32> @@ -34,7 +34,7 @@ fusion { // NEW-VECTOR-SAME: %[[REAL:[^:]+]]: memref<128x256xf32>, // NEW-VECTOR-SAME: %[[IMAG:[^:]+]]: memref<128x256xf32>, // NEW-VECTOR-SAME: %[[OUT:[^:]+]]: memref<128x256x2xf32>, -// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index) +// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // NEW-VECTOR: vector.transfer_read %[[REAL]] // NEW-VECTOR: vector.transfer_read %[[IMAG]] diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_constant.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_constant.hlo index 63c5c2828574be..19d0deb1c23fad 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_constant.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_constant.hlo @@ -15,7 +15,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[LHS:[^:]+]]: memref<4x8xcomplex>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<4x8xcomplex>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 0 : index]}) // CHECK-DAG: %[[CST:.*]] = arith.constant dense<(1.000000e+00,2.000000e+00)> : tensor> // CHECK: %[[T_0:.*]] = xtile.extract %[[LHS]] @@ -26,7 +26,7 @@ fusion { // NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl( // NEW-VECTOR-SAME: %[[LHS:[^:]+]]: memref<4x8x2xf32>, // NEW-VECTOR-SAME: %[[OUT:[^:]+]]: memref<4x8x2xf32>, -// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index) +// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 0 : index]}) // NEW-VECTOR-DAG: %[[CST:.*]] = arith.constant dense<[1.000000e+00, 2.000000e+00]> : vector<2xf32> // NEW-VECTOR: vector.transfer_read %[[LHS]] diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_divide.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_divide.hlo index 0520a61085ee55..ecb5bef6c7b6d6 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_divide.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_divide.hlo @@ -22,7 +22,7 @@ fusion { // CHECK-SAME: %[[LHS:[^:]+]]: memref<9xcomplex>, // CHECK-SAME: %[[RHS:[^:]+]]: memref<9xcomplex>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<9xcomplex>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 0 : index]}) // CHECK: %[[T_0:.*]] = xtile.extract %[[LHS]] // CHECK-SAME: : memref<9xcomplex> -> tensor<4xcomplex> @@ -37,7 +37,7 @@ fusion { // NEW-VECTOR-SAME: %[[LHS:[^:]+]]: memref<9x2xf32>, // NEW-VECTOR-SAME: %[[RHS:[^:]+]]: memref<9x2xf32>, // NEW-VECTOR-SAME: %[[OUT:[^:]+]]: memref<9x2xf32>, -// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index) +// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 0 : index]}) // NEW-VECTOR: vector.transfer_read %[[LHS]] // NEW-VECTOR: vector.transfer_read %[[RHS]] diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_imag.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_imag.hlo index b3b59b7652c2b2..ddea1e28510dd3 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_imag.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_imag.hlo @@ -17,7 +17,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[IN:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK: %[[T_0:.*]] = xtile.extract %[[IN]] // CHECK-SAME: : memref<128x256xcomplex> -> tensor<1x4xcomplex> @@ -29,7 +29,7 @@ fusion { // NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl( // NEW-VECTOR-SAME: %[[IN:[^:]+]]: memref<128x256x2xf32>, // NEW-VECTOR-SAME: %[[OUT:[^:]+]]: memref<128x256xf32>, -// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index) +// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // NEW-VECTOR: vector.transfer_read %[[IN]] // NEW-VECTOR: vector.extract_strided_slice diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_multiply.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_multiply.hlo index 1aa2d3ee7bc20c..d67374b52be8a5 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_multiply.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_multiply.hlo @@ -19,7 +19,7 @@ fusion { // CHECK-SAME: %[[LHS:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[RHS:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xcomplex>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK: %[[T_0:.*]] = xtile.extract %[[LHS]] // CHECK-SAME: : memref<128x256xcomplex> -> tensor<1x4xcomplex> @@ -34,7 +34,7 @@ fusion { // NEW-VECTOR-SAME: %[[LHS:[^:]+]]: memref<128x256x2xf32>, // NEW-VECTOR-SAME: %[[RHS:[^:]+]]: memref<128x256x2xf32>, // NEW-VECTOR-SAME: %[[OUT:[^:]+]]: memref<128x256x2xf32>, -// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index) +// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // NEW-VECTOR: vector.transfer_read %[[LHS]] // NEW-VECTOR: vector.transfer_read %[[RHS]] diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_negate.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_negate.hlo index 0ee8fb6a457c58..ad5ff0a35011b7 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_negate.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_negate.hlo @@ -17,7 +17,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[IN:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xcomplex>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK: %[[T_0:.*]] = xtile.extract %[[IN]] // CHECK-SAME: : memref<128x256xcomplex> -> tensor<1x4xcomplex> @@ -29,7 +29,7 @@ fusion { // NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl( // NEW-VECTOR-SAME: %[[IN:[^:]+]]: memref<128x256x2xf32>, // NEW-VECTOR-SAME: %[[OUT:[^:]+]]: memref<128x256x2xf32>, -// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index) +// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // NEW-VECTOR: vector.transfer_read %[[IN]] // NEW-VECTOR: arith.negf diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_real.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_real.hlo index 23f1fbb0e1cbd4..960a48e017a4d9 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_real.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_real.hlo @@ -17,7 +17,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[IN:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK: %[[T_0:.*]] = xtile.extract %[[IN]] // CHECK-SAME: : memref<128x256xcomplex> -> tensor<1x4xcomplex> @@ -29,7 +29,7 @@ fusion { // NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl( // NEW-VECTOR-SAME: %[[IN:[^:]+]]: memref<128x256x2xf32>, // NEW-VECTOR-SAME: %[[OUT:[^:]+]]: memref<128x256xf32>, -// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index) +// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // NEW-VECTOR: vector.transfer_read %[[IN]] // NEW-VECTOR: vector.extract_strided_slice diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_select.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_select.hlo index 0dd1de7bd54b85..7420e97ab227b6 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_select.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_select.hlo @@ -21,7 +21,7 @@ fusion { // CHECK-SAME: %[[ON_TRUE:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[ON_FALSE:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xcomplex>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK: %[[T_0:.*]] = xtile.extract %[[COND]] // CHECK-SAME: : memref<128x256xi8> -> tensor<1x4xi8> @@ -39,7 +39,7 @@ fusion { // NEW-VECTOR-SAME: %[[ON_TRUE:[^:]+]]: memref<128x256x2xf32>, // NEW-VECTOR-SAME: %[[ON_FALSE:[^:]+]]: memref<128x256x2xf32>, // NEW-VECTOR-SAME: %[[OUT:[^:]+]]: memref<128x256x2xf32>, -// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index) +// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // NEW-VECTOR: vector.transfer_read %[[COND]] // NEW-VECTOR: vector.transfer_read %[[ON_TRUE]] diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_subtract.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_subtract.hlo index 8f653d54acbb18..b68e9468f564a9 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_subtract.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/complex_subtract.hlo @@ -22,7 +22,7 @@ fusion { // CHECK-SAME: %[[LHS:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[RHS:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xcomplex>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK: %[[T_0:.*]] = xtile.extract %[[LHS]] // CHECK-SAME: : memref<128x256xcomplex> -> tensor<1x4xcomplex> @@ -37,7 +37,7 @@ fusion { // NEW-VECTOR-SAME: %[[LHS:[^:]+]]: memref<128x256x2xf32>, // NEW-VECTOR-SAME: %[[RHS:[^:]+]]: memref<128x256x2xf32>, // NEW-VECTOR-SAME: %[[OUT:[^:]+]]: memref<128x256x2xf32>, -// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index) +// NEW-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // NEW-VECTOR: %[[LHS_VEC:.*]] = vector.transfer_read %[[LHS]] // NEW-VECTOR-SAME: : memref<128x256x2xf32>, vector<1x4x2xf32> diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_f32_bf16.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_f32_bf16.hlo index 0556729980ea28..6facf00d7b6f34 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_f32_bf16.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_f32_bf16.hlo @@ -3,29 +3,40 @@ // RUN: test_correctness %s // RUN: fusion_to_xtile %s | fusion_compiler_opt \ -// RUN: --xtile-cpu-new-xtile-to-vector --xtile-cpu-new-vector-to-llvm \ +// RUN: --xtile-cpu-new-xtile-to-vector --xla-simplify-arith \ +// RUN: | FileCheck %s --check-prefix=NEW-VECTOR +// RUN: fusion_to_xtile %s | fusion_compiler_opt \ +// RUN: --xtile-cpu-new-xtile-to-vector="prefer_vector_width=256" \ +// RUN: --xtile-cpu-new-vector-to-llvm \ // RUN: | FileCheck %s --check-prefix=NEW-VECTOR-LLVM // RUN: test_correctness %s --xla_cpu_use_new_xtile_lowering=true fusion { - p0 = f32[1, 16] parameter(0) - ROOT convert = bf16[1, 16] convert(p0) + p0 = f32[2, 16] parameter(0) + ROOT convert = bf16[2, 16] convert(p0) } // CHECK-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl +// NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl +// NEW-VECTOR-NOT: scf.if +// NEW-VECTOR: vector.transfer_read {{.*}} {in_bounds = [true, true]} : memref<2x16xf32>, vector<1x16xf32> +// NEW-VECTOR: vector.transfer_write {{.*}} {in_bounds = [true, true]} : vector<1x16xi16>, memref<2x16xi16> + +// NEW-VECTOR-LLVM: llvm.func @wrapped_fusion(%{{.*}}) -> !llvm.ptr +// NEW-VECTOR-LLVM-SAME: "prefer-vector-width", "256" // NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl // NEW-VECTOR-LLVM: %[[LOAD:.*]] = llvm.load %{{.*}} : !llvm.ptr -> vector<16xf32> -// NEW-VECTOR-LLVM-COUNT-2: llvm.call @xla.fptrunc.v8f32.to.v8bf16(%{{.*}}) : (vector<8xf32>) -> vector<8xbf16> +// NEW-VECTOR-LLVM: llvm.call @xla.fptrunc.v16f32.to.v16bf16(%[[LOAD]]) : (vector<16xf32>) -> vector<16xbf16> // NEW-VECTOR-LLVM: llvm.store %{{.*}}, %{{.*}} : vector<16xi16>, !llvm.ptr ENTRY main { - p0 = f32[1, 16] parameter(0) - ROOT fusion = bf16[1, 16] fusion(p0), kind=kLoop, calls=fusion, + p0 = f32[2, 16] parameter(0) + ROOT fusion = bf16[2, 16] fusion(p0), kind=kLoop, calls=fusion, backend_config={ "fusion_backend_config":{ "block_level_fusion_config":{ "output_tiles":[{"sizes":["1", "16"]}], - "num_ctas":"1" + "num_ctas":"2" } } } diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_s8_s32.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_s8_s32.hlo index 5709a6d143fab6..48638c870189ea 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_s8_s32.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_s8_s32.hlo @@ -24,7 +24,7 @@ fusion { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG:[^:]+]]: memref<128x256xi8>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xi32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK-SAME: attributes {num_opaque_args = 0 : i32} // CHECK: %[[C0:.*]] = arith.constant 0 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/iota_2d.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/iota_2d.hlo index 4ef4f5d4ca75af..5573dc9f998222 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/iota_2d.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/iota_2d.hlo @@ -32,7 +32,7 @@ ENTRY main { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 0 : index]}) // CHECK: %[[IOTA:.*]] = stablehlo.iota dim = 0 : tensor<4xi32> // CHECK: %[[CONV:.*]] = stablehlo.convert %{{.*}} : (tensor<4xi32>) -> tensor<4xf32> // CHECK: %[[BCAST:.*]] = stablehlo.broadcast_in_dim %[[CONV]], dims = [1] : (tensor<4xf32>) -> tensor<8x4xf32> diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/iota_f32.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/iota_f32.hlo index e7ac9cb363dbf2..6c92fba71ada0a 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/iota_f32.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/iota_f32.hlo @@ -32,7 +32,7 @@ ENTRY main { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[OUT:[^:]+]]: memref<128xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 0 : index]}) // CHECK: %[[IOTA:.*]] = stablehlo.iota dim = 0 : tensor<32xi32> // CHECK: %[[CONV:.*]] = stablehlo.convert %{{.*}} : (tensor<32xi32>) -> tensor<32xf32> // CHECK: xtile.insert %{{.*}} into %[[OUT]][%{{.*}}] [32] [1] : tensor<32xf32> -> memref<128xf32> diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/iota_s8.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/iota_s8.hlo index 9a0b7ee18125ac..1184dab099bd76 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/iota_s8.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/iota_s8.hlo @@ -32,7 +32,7 @@ ENTRY main { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[OUT:[^:]+]]: memref<128xi8>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 0 : index]}) // CHECK: %[[IOTA:.*]] = stablehlo.iota dim = 0 : tensor<32xi32> // CHECK: %[[CONV:.*]] = stablehlo.convert %{{.*}} : (tensor<32xi32>) -> tensor<32xi8> // CHECK: xtile.insert %{{.*}} into %[[OUT]][%{{.*}}] [32] [1] : tensor<32xi8> -> memref<128xi8> diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite.hlo index 05965d7a9060e5..afba2925296a0f 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite.hlo @@ -15,7 +15,7 @@ // CHECK-LABEL: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG:[^:]+]]: memref<1x9xf16>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<1x9xi8>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 0 : index]}) // CHECK: %[[T_0:.*]] = xtile.extract %[[ARG]] // CHECK-SAME: : memref<1x9xf16> -> tensor<1x16xf16> @@ -29,7 +29,7 @@ // CHECK-VECTOR-LABEL: func.func @wrapped_fusion_impl( // CHECK-VECTOR-SAME: %[[ARG:[^:]+]]: memref<1x9xf16>, // CHECK-VECTOR-SAME: %[[OUT:[^:]+]]: memref<1x9xi8>, -// CHECK-VECTOR-SAME: %[[PID:[^:]+]]: index) +// CHECK-VECTOR-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 0 : index]}) // CHECK-VECTOR: %[[ALLOCA_ARG:.*]] = memref.alloca() : memref<1x16xf16> // CHECK-VECTOR: %[[ALLOCA_OUT:.*]] = memref.alloca() alignment = 64 : memref<1x16xi8> // CHECK-VECTOR: memref.copy %[[ARG]], diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite_2d_tile.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite_2d_tile.hlo index 03262024000283..eebc7c268c0501 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite_2d_tile.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/is_finite_2d_tile.hlo @@ -33,7 +33,7 @@ fusion { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG:[^:]+]]: memref<128x256xf32>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xi8>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 1 : index]}) // CHECK-SAME: attributes {num_opaque_args = 0 : i32} // CHECK: %[[C0:.*]] = arith.constant 0 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum.hlo index 8b98b3e33edfc7..87aace4418d9f5 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum.hlo @@ -26,7 +26,7 @@ fusion { // CHECK-SAME: %[[LHS:[^:]+]]: memref<128x256xf32>, // CHECK-SAME: %[[RHS:[^:]+]]: memref<128x256xf32>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK-SAME: attributes {num_opaque_args = 0 : i32} // CHECK: %[[C0:.*]] = arith.constant 0 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_2d_tile.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_2d_tile.hlo index 4db7f30c217a42..44bb8b50bb6a91 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_2d_tile.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_2d_tile.hlo @@ -26,7 +26,7 @@ fusion { // CHECK-SAME: %[[LHS:[^:]+]]: memref<128x256xf32>, // CHECK-SAME: %[[RHS:[^:]+]]: memref<128x256xf32>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 1 : index]}) // CHECK-SAME: attributes {num_opaque_args = 0 : i32} // CHECK: %[[C0:.*]] = arith.constant 0 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_mask.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_mask.hlo index 8f0e0a6673a99d..a03af7a167a74d 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_mask.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_mask.hlo @@ -26,7 +26,7 @@ fusion { // CHECK-SAME: %[[LHS:[^:]+]]: memref<100x100xf32>, // CHECK-SAME: %[[RHS:[^:]+]]: memref<100x100xf32>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<100x100xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 0 : index]}) // CHECK-SAME: attributes {num_opaque_args = 0 : i32} // CHECK: %[[C0:.*]] = arith.constant 0 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_ui32.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_ui32.hlo index 597a2198c745c8..56ddb425715092 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_ui32.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_ui32.hlo @@ -26,7 +26,7 @@ fusion { // CHECK-SAME: %[[LHS:[^:]+]]: memref<128x256xui32>, // CHECK-SAME: %[[RHS:[^:]+]]: memref<128x256xui32>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x256xui32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK-SAME: attributes {num_opaque_args = 0 : i32} // CHECK: %[[C0:.*]] = arith.constant 0 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/reduce_precision.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/reduce_precision.hlo index b59d9502b3b75d..1d4acd983d489a 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/reduce_precision.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/reduce_precision.hlo @@ -25,7 +25,7 @@ fusion { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG:[^:]+]]: memref<128x32xf32>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<128x32xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK-SAME: attributes {num_opaque_args = 0 : i32} // CHECK: %[[C0:.*]] = arith.constant 0 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/strided_masked.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/strided_masked.hlo index 85b1697bef1d61..5d41203d7ef8c2 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/strided_masked.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/strided_masked.hlo @@ -31,7 +31,7 @@ ENTRY main { // CHECK-LABEL: xtile.entry_func @wrapped_fusion // CHECK-SAME: %[[ARG0:[^:]+]]: memref<10xf32>, // CHECK-SAME: %[[ARG1:[^:]+]]: memref<3xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 0 : index]}) // CHECK: %[[C1:.*]] = arith.constant 1 : index // CHECK: %[[EXT:.*]] = xtile.extract %[[ARG0]][%[[C1]]] [16] [2] : memref<10xf32> -> tensor<16xf32> // CHECK: %[[C0:.*]] = arith.constant 0 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo index 82c69b890bd617..bf668ec4edf2c2 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo @@ -29,7 +29,7 @@ fusion { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG0:[^:]+]]: memref<1x256x56x56xbf16>, // CHECK-SAME: %[[ARG1:[^:]+]]: memref<1x56x56x256xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 11 : index]}) // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[C1:.*]] = arith.constant 1 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo index cca9185f4e72bc..ce27287f921bc3 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo @@ -29,7 +29,7 @@ fusion { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG0:[^:]+]]: memref<1x56x56x256xf32>, // CHECK-SAME: %[[ARG1:[^:]+]]: memref<1x256x56x56xbf16>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 11 : index]}) // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[C1:.*]] = arith.constant 1 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape.hlo index 6bdc00e6644f98..dccc1439a19a6a 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape.hlo @@ -25,7 +25,7 @@ fusion { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG:[^:]+]]: memref<128x256xf32>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<32768xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK-SAME: attributes {num_opaque_args = 0 : i32} // CHECK: %[[C0:.*]] = arith.constant 0 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape_c64.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape_c64.hlo index 45c6f654275375..75a0fcddbabf55 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape_c64.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape_c64.hlo @@ -25,7 +25,7 @@ fusion { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG:[^:]+]]: memref<128x256xcomplex>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<32768xcomplex>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 3 : index]}) // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[C1:.*]] = arith.constant 1 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose.hlo index 2750a1515d771e..afe7f0e4ae7587 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose.hlo @@ -33,7 +33,7 @@ ENTRY main { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG0:[^:]+]]: memref<16x40x40x64xf32>, // CHECK-SAME: %[[ARG1:[^:]+]]: memref<64x40x40x16xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 7 : index]}) // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[C1:.*]] = arith.constant 1 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_102.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_102.hlo index b4ea00c697948d..7e6be98f5855a4 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_102.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_102.hlo @@ -34,7 +34,7 @@ ENTRY main { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG0:[^:]+]]: memref<100x3x70xf32>, // CHECK-SAME: %[[ARG1:[^:]+]]: memref<3x100x70xf32>, -// CHECK-SAME: %{{.*}}: index) +// CHECK-SAME: %{{.*}}: index {xla.range = [0 : index, 0 : index]}) // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: scf.for %[[IDX:.*]] = %{{.*}} to %{{.*}} step %{{.*}} { diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_120.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_120.hlo index a50e56170ce968..3c57a8a6867c12 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_120.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_120.hlo @@ -35,7 +35,7 @@ ENTRY main { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG0:[^:]+]]: memref<128x32x64xf32>, // CHECK-SAME: %[[ARG1:[^:]+]]: memref<32x64x128xf32>, -// CHECK-SAME: %{{.*}}: index) +// CHECK-SAME: %{{.*}}: index {xla.range = [0 : index, 0 : index]}) // CHECK: scf.for %[[IDX:.*]] = %{{.*}} to %{{.*}} step %{{.*}} { // CHECK: %[[OFF0:.*]] = xla.apply_indexing #indexing_map(%[[IDX]]) diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_201.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_201.hlo index 8daf49587df355..d05da8c9680960 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_201.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_201.hlo @@ -35,7 +35,7 @@ ENTRY main { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG0:[^:]+]]: memref<128x32x64xf32>, // CHECK-SAME: %[[ARG1:[^:]+]]: memref<64x128x32xf32>, -// CHECK-SAME: %{{.*}}: index) +// CHECK-SAME: %{{.*}}: index {xla.range = [0 : index, 0 : index]}) // CHECK: scf.for %[[IDX:.*]] = %{{.*}} to %{{.*}} step %{{.*}} { // CHECK: %[[OFF0:.*]] = xla.apply_indexing #indexing_map(%[[IDX]]) diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_c64.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_c64.hlo index cbb82ace8c6af2..30c1ee05acd5dc 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_c64.hlo +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_c64.hlo @@ -25,7 +25,7 @@ fusion { // CHECK: xtile.entry_func @wrapped_fusion( // CHECK-SAME: %[[ARG:[^:]+]]: memref<16x40x40x64xcomplex>, // CHECK-SAME: %[[OUT:[^:]+]]: memref<64x40x40x16xcomplex>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index {xla.range = [0 : index, 7 : index]}) // CHECK-SAME: attributes {num_opaque_args = 0 : i32} // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/lower_xtile_entry.cc b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/lower_xtile_entry.cc index 410a16f31e5df0..a4788634fab093 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/lower_xtile_entry.cc +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/lower_xtile_entry.cc @@ -47,7 +47,6 @@ limitations under the License. namespace xla::cpu { -#define GEN_PASS_DECL_LOWERXTILEENTRYPASS #define GEN_PASS_DEF_LOWERXTILEENTRYPASS #include "xla/backends/cpu/codegen/tiled/transforms/passes.h.inc" @@ -132,6 +131,13 @@ class LowerXTileEntryPass auto kernel_func = FuncOp::create( builder, kernel_name, builder.getFunctionType({call_frame_type}, {error_type})); + if (prefer_vector_width_ > 0) { + kernel_func->setAttr( + "llvm.passthrough", + builder.getArrayAttr({builder.getStrArrayAttr( + {"prefer-vector-width", + absl::StrCat(prefer_vector_width_.getValue())})})); + } builder.setInsertionPointToStart(kernel_func.addEntryBlock()); diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/passes.td b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/passes.td index 1545a80a5a6007..e3e6647bd4aadd 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/passes.td +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/passes.td @@ -25,6 +25,13 @@ def LowerXTileEntryPass : Pass<"xtile-cpu-lower-xtile-entry", "mlir::ModuleOp"> "xla::cpu::XlaCpuDialect", "xla::xtile::XTileDialect" ]; + + let options = [ + Option<"prefer_vector_width_", "prefer_vector_width", "int32_t", + /*default=*/"0", + "prefer-vector-width value to set on the kernel function. " + "Not set if 0.">, + ]; } def ShloToVectorPass : Pass<"xtile-cpu-shlo-to-vector", "mlir::ModuleOp"> { diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/lower_xtile_entry.mlir b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/lower_xtile_entry.mlir index d60c825471e71d..27d3b854cc7330 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/lower_xtile_entry.mlir +++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/lower_xtile_entry.mlir @@ -1,5 +1,8 @@ -// RUN: fusion_compiler_opt %s --xtile-cpu-lower-xtile-entry -split-input-file -//| FileCheck %s +// RUN: fusion_compiler_opt %s --xtile-cpu-lower-xtile-entry -split-input-file \ +// RUN: | FileCheck %s +// RUN: fusion_compiler_opt %s -split-input-file \ +// RUN: --xtile-cpu-lower-xtile-entry="prefer_vector_width=256" \ +// RUN: | FileCheck %s --check-prefix=PREFER_VECTOR xtile.entry_func @simple_wrap(%input: memref<1024xf32> {xla.some_attr = 1}, %output: memref<32xf64>, @@ -7,25 +10,28 @@ xtile.entry_func @simple_wrap(%input: memref<1024xf32> {xla.some_attr = 1}, xtile.return } -// CHECK-DAG: #[[MAP:.*]] = #xla.indexing_map<"(workgroup_id) -> (min(workgroup_id * 64, 1012), min(workgroup_id * 64 + 64, 1012)), domain: workgroup_id in [0, 15]"> - // CHECK: func.func @simple_wrap(%[[CALL_FRAME:.*]]: !xla_cpu.call_frame) -> !xla_cpu.error { // CHECK-DAG: %[[INPUT:.*]] = xla_cpu.load %[[CALL_FRAME]], 0 : memref<1024xf32> // CHECK-DAG: %[[OUTPUT:.*]] = xla_cpu.load %[[CALL_FRAME]], 1 : memref<32xf64> // CHECK-DAG: %[[WORKGROUP_ID:.*]] = xla_cpu.extract_workgroup_id %[[CALL_FRAME]], x -// CHECK: func.call @[[IMPL_FUNC:.*]](%[[INPUT]], %[[OUTPUT]], %[[IDX]]) : (memref<1024xf32>, memref<32xf64>, index) -> () +// CHECK: call @simple_wrap_impl(%[[INPUT]], %[[OUTPUT]], %[[WORKGROUP_ID]]) : (memref<1024xf32>, memref<32xf64>, index) -> () // CHECK: %[[SUCCESS:.*]] = xla_cpu.success : !xla_cpu.error // CHECK: return %[[SUCCESS]] : !xla_cpu.error -// CHECK: func.func @[[IMPL_FUNC]]( +// CHECK: func.func @simple_wrap_impl( // CHECK-SAME: %{{.*}}: memref<1024xf32> {xla.some_attr = 1 : i64}, // CHECK-SAME: %{{.*}}: memref<32xf64>, // CHECK-SAME: %{{.*}}: index) // CHECK-SAME: attributes {llvm.always_inline, llvm.linkage = #llvm.linkage // CHECK: return +// PREFER_VECTOR: func.func @simple_wrap(%{{.*}}: !xla_cpu.call_frame) -> !xla_cpu.error +// PREFER_VECTOR-SAME: attributes {llvm.passthrough = +// PREFER_VECTOR-SAME: ["prefer-vector-width", "256"]]} { +// PREFER_VECTOR: func.func @simple_wrap_impl( +// PREFER_VECTOR-NOT: llvm.passthrough // ----- diff --git a/third_party/xla/xla/backends/cpu/codegen/tools/fusion_compiler_opt.cc b/third_party/xla/xla/backends/cpu/codegen/tools/fusion_compiler_opt.cc index a1c0da1d71d184..5cf1388adf8030 100644 --- a/third_party/xla/xla/backends/cpu/codegen/tools/fusion_compiler_opt.cc +++ b/third_party/xla/xla/backends/cpu/codegen/tools/fusion_compiler_opt.cc @@ -13,6 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include +#include + #include "llvm/ADT/StringRef.h" #include "llvm/Support/CommandLine.h" #include "mlir/Dialect/Func/Extensions/AllExtensions.h" @@ -36,6 +39,20 @@ struct XtileCpuPassOptions *this, "fast_min_max", llvm::cl::desc("Whether to enable fast min/max operations."), llvm::cl::init(false)}; + Option cpu_features{ + *this, "cpu_features", + llvm::cl::desc("LLVM target feature string (e.g. \"+avx,+avx2\") used " + "to select feature-gated intrinsic variants."), + llvm::cl::init("")}; +}; + +struct XtileToVectorPassOptions + : public mlir::PassPipelineOptions { + Option prefer_vector_width{ + *this, "prefer_vector_width", + llvm::cl::desc("prefer-vector-width value to set on the kernel " + "function. Not set if 0."), + llvm::cl::init(0)}; }; int main(int argc, char** argv) { @@ -49,28 +66,33 @@ int main(int argc, char** argv) { xla::xtile::registerXTileTransformsPasses(); mlir::stablehlo::registerStablehloLinalgTransformsPasses(); - mlir::PassPipelineRegistration( + mlir::PassPipelineRegistration( "xtile-cpu-xtile-to-vector", "Run the conversion from XTile to Vector dialect.", - [](mlir::OpPassManager& pm) { - xla::cpu::AddXtileToVectorPasses(pm, /*msan_enabled=*/false); + [](mlir::OpPassManager& pm, const XtileToVectorPassOptions& options) { + xla::cpu::AddXtileToVectorPasses(pm, /*msan_enabled=*/false, + options.prefer_vector_width); }); - mlir::PassPipelineRegistration( + mlir::PassPipelineRegistration( "xtile-cpu-new-xtile-to-vector", "Run the conversion from XTile to Vector dialect.", - [](mlir::OpPassManager& pm) { xla::cpu::AddNewXtileToVectorPasses(pm); }); + [](mlir::OpPassManager& pm, const XtileToVectorPassOptions& options) { + xla::cpu::AddNewXtileToVectorPasses(pm, options.prefer_vector_width); + }); mlir::PassPipelineRegistration( "xtile-cpu-vector-to-llvm", "Run the conversion from Vector to LLVM dialect.", [](mlir::OpPassManager& pm, const XtileCpuPassOptions& options) { - xla::cpu::AddVectorToLLVMPasses(pm, options.fast_min_max); + xla::cpu::AddVectorToLLVMPasses(pm, options.fast_min_max, + options.cpu_features.getValue()); }); mlir::PassPipelineRegistration( "xtile-cpu-new-vector-to-llvm", "Run the conversion from Vector to LLVM dialect.", [](mlir::OpPassManager& pm, const XtileCpuPassOptions& options) { xla::cpu::AddNewVectorToLLVMPasses(pm, options.fast_min_max, - /*vector_width=*/256); + /*vector_width=*/256, + options.cpu_features.getValue()); }); return mlir::failed(MlirOptMain( argc, argv, "XLA:CPU Fusion compiler pass driver\n", registry)); diff --git a/third_party/xla/xla/backends/cpu/runtime/buffer_allocations.h b/third_party/xla/xla/backends/cpu/runtime/buffer_allocations.h index a0d11f61c53d42..2fea17cd6a9b36 100644 --- a/third_party/xla/xla/backends/cpu/runtime/buffer_allocations.h +++ b/third_party/xla/xla/backends/cpu/runtime/buffer_allocations.h @@ -106,7 +106,7 @@ BufferAllocations::GetDeviceAddress(BufferAllocation::Slice slice) const { // Handle empty slices explicitly and return a null pointer device memory to // guarantee that we do not accidentally write through the empty slice which // would hide a real bug in the code. - if (ABSL_PREDICT_FALSE(slice.size() == 0)) { + if (ABSL_PREDICT_FALSE(slice.empty())) { return se::DeviceAddressBase(nullptr, 0); } diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton/BUILD b/third_party/xla/xla/backends/gpu/autotuner/triton/BUILD index 39f669c3275b75..6b90a4b087eff2 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/triton/BUILD +++ b/third_party/xla/xla/backends/gpu/autotuner/triton/BUILD @@ -42,10 +42,12 @@ cc_library( deps = [ "//xla:shape_util", "//xla:util", + "//xla:xla_data_proto_cc", "//xla/backends/gpu/codegen/triton:tma_utils", "//xla/hlo/ir:hlo", "//xla/hlo/utils:hlo_traversal", "//xla/service/gpu:matmul_utils", + "//xla/service/gpu/model:gpu_performance_model_base", "//xla/stream_executor:device_description", "//xla/stream_executor/gpu:tma_metadata", "//xla/tsl/lib/core:bits", @@ -69,6 +71,7 @@ xla_cc_test( "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/hlo/testlib:verified_hlo_module", + "//xla/service/gpu:gpu_device_info_for_tests", "//xla/service/gpu:matmul_utils", "//xla/stream_executor:device_description", "//xla/stream_executor:device_description_proto_cc", diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space.cc b/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space.cc index 67777810bcda26..24ebe2acc6646b 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space.cc @@ -36,6 +36,7 @@ limitations under the License. #include "xla/hlo/utils/hlo_traversal.h" #include "xla/primitive_util.h" #include "xla/service/gpu/matmul_utils.h" +#include "xla/service/gpu/model/gpu_performance_model_base.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/stream_executor/device_description.h" @@ -101,6 +102,8 @@ TritonDotFusionSearchSpace::TritonDotFusionSearchSpace( operand_bitwidth_( // The bitwitdth of both operands is the same. primitive_util::BitWidth(dot->operand(0)->shape().element_type())), compute_bitwidth_(primitive_util::BitWidth(dot->shape().element_type())), + operand_type_(dot->operand(0)->shape().element_type()), + compute_type_(dot->shape().element_type()), // Figure out some basic limitations on tiling based on the above. lhs_has_expensive_op_(HasExpensiveTransitiveParent(dot->operand(0))), rhs_has_expensive_op_(HasExpensiveTransitiveParent(dot->operand(1))), @@ -242,11 +245,54 @@ bool TritonDotFusionSearchSpace::HasExpensiveTransitiveParent( }); } +double TritonDotFusionSearchSpace::GetArithmeticIntensity() const { + const double m = static_cast(lhs_parallel_size_); + const double n = static_cast(rhs_parallel_size_); + const double k = static_cast(contracting_size_); + const double flops = 2.0 * m * n * k; + const double lhs_bytes = m * k * (operand_bitwidth_ / 8.0); + const double rhs_bytes = k * n * (operand_bitwidth_ / 8.0); + const double out_bytes = m * n * (compute_bitwidth_ / 8.0); + const double total_bytes = lhs_bytes + rhs_bytes + out_bytes; + return total_bytes > 0 ? flops / total_bytes : 0.0; +} + +double TritonDotFusionSearchSpace::GetRidgePoint() const { + const int64_t peak_matrix_ops_per_ns = + GpuPerformanceModelBase::CalculatePeakMatrixOpsPerNs(device_description_, + operand_type_); + const int64_t memory_bandwidth = device_description_.memory_bandwidth(); + if (peak_matrix_ops_per_ns > 0 && memory_bandwidth > 0) { + // Ridge point = (peak FLOPs / s) / (memory bandwidth in bytes / s) + // = (peak_matrix_ops_per_ns * 1e9) / memory_bandwidth. + return (static_cast(peak_matrix_ops_per_ns) * 1e9) / + static_cast(memory_bandwidth); + } + // Fallback if matrix unit description or bandwidth is unavailable: + // ~300 FLOP/B for 16-bit Tensor Cores on Hopper/Blackwell, scaling inversely + // with operand bitwidth. + return operand_bitwidth_ > 0 ? 4800.0 / operand_bitwidth_ : 300.0; +} + int TritonDotFusionSearchSpace::GetDesiredTotalWarps() const { constexpr int kSchedulersPerCore = 4; - constexpr int kDesiredWarpsPerCore = - kMaxWarpsPerScheduler * kSchedulersPerCore; - return kDesiredWarpsPerCore * device_description_.core_count(); + const double arithmetic_intensity = GetArithmeticIntensity(); + const double ridge_point = GetRidgePoint(); + const double compute_intensity_ratio = + ridge_point > 0.0 + ? std::clamp(arithmetic_intensity / ridge_point, 0.0, 1.0) + : 1.0; + + // Memory-bound dots benefit from higher warp oversubscription (up to + // kMaxWarpsPerScheduler) to overlap memory latency, while compute-bound dots + // scale down to kMinWarpsPerScheduler to avoid thrashing and contention. + const int desired_warps_per_scheduler = static_cast( + std::round(kMaxWarpsPerScheduler - + compute_intensity_ratio * + (kMaxWarpsPerScheduler - kMinWarpsPerScheduler))); + const int desired_warps_per_core = + desired_warps_per_scheduler * kSchedulersPerCore; + return desired_warps_per_core * device_description_.core_count(); } TritonDotFusionSearchSpace::OutputTile @@ -593,7 +639,9 @@ void TritonDotFusionSearchSpace::EliminateLowOccupancyConfigs( return; } - constexpr int kMinConfigsForOccupancyOptimization = 24; + // Output tilings alone typically produce 12-20 configs (without split-K + // multipliers). We only skip optimization for very small search spaces. + constexpr int kMinConfigsForOccupancyOptimization = 15; if (configs.size() < kMinConfigsForOccupancyOptimization) { VLOG(10) << "Skipping occupancy optimization for small search spaces. " "Configs size: " @@ -601,7 +649,7 @@ void TritonDotFusionSearchSpace::EliminateLowOccupancyConfigs( return; } - ConfigWithNotes last_config = configs.back(); // Largest split. + ConfigWithNotes last_config = configs.back(); auto has_too_few_tiles = [](const ConfigWithNotes& config) { if (config.not_enough_tiles) { VLOG(10) << "Skipping due to fewer tiles than cores, config = " @@ -613,8 +661,7 @@ void TritonDotFusionSearchSpace::EliminateLowOccupancyConfigs( configs.erase(llvm::remove_if(configs, has_too_few_tiles), configs.end()); if (configs.empty()) { // We can get no configs if the problem is small enough to not even occupy - // all cores. In that case, we just use the largest split and smallest - // tiling. + // all cores. In that case, we just use the smallest tiling. last_config.config.block_m = min_out_tile_.lhs_dim; last_config.config.block_n = min_out_tile_.rhs_dim; VLOG(10) << "No configs with sufficient occupancy, using config = " diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space.h b/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space.h index 28b9530a753bcf..738c443adb518c 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space.h +++ b/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space.h @@ -26,6 +26,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_instructions.h" #include "xla/service/gpu/matmul_utils.h" #include "xla/stream_executor/device_description.h" +#include "xla/xla_data.pb.h" namespace xla::gpu { @@ -100,10 +101,12 @@ class TritonDotFusionSearchSpace { // group" (4 warps) within a CTA to cooperate on a single instruction. /// https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions static constexpr int kMinWarpsPerCtaForWgmma = 4; - // Approximation on the maximum number of warps we would want to oversubscribe - // the SMs with to overlap different GPU pipes (memory, tensor core, ALU, - // special function unit, etc.) - // TODO: b/408114338 - Figure out a better model for this. + // Range of warps per scheduler we want to oversubscribe the SMs with to + // overlap different GPU pipes (memory, tensor core, ALU, special function + // unit, etc.). Memory-bound dots use up to kMaxWarpsPerScheduler to hide + // memory latency, while compute-bound dots scale down to + // kMinWarpsPerScheduler to avoid resource contention and thrashing. + static constexpr int kMinWarpsPerScheduler = 2; static constexpr int kMaxWarpsPerScheduler = 5; // Callback type for `ExtendConfigs`. The method should append zero or more @@ -121,20 +124,28 @@ class TritonDotFusionSearchSpace { bool HasExpensiveTransitiveParent(const HloInstruction* operand) const; // Computes the maximum number of total warps we should have to sufficiently - // saturate the GPU. + // saturate the GPU based on hardware properties and the dot's arithmetic + // intensity. // // We're counting warps instead of blocks here, since we already need this // value as a consideration to decide how large the blocks should be (which // then impacts how many of them we should have). int GetDesiredTotalWarps() const; + // Computes the arithmetic intensity (FLOPs / byte) of the dot operation. + double GetArithmeticIntensity() const; + + // Computes the hardware ridge-point arithmetic intensity (FLOPs / byte) + // where the dot transitions from memory-bound to compute-bound, based on + // device peak matrix ops and memory bandwidth. + double GetRidgePoint() const; + // Computes the maximum sensible size of the output tile (block_m, block_n) // based on the dot shape and element type, and the available registers on // the core. OutputTile GetMaxOutputTile() const; - // Computes the number of result tiles we would have without - // splitting the contracting dimension for a given output tile. + // Computes the number of result tiles we would have for a given output tile. int64_t GetNumResultTiles(OutputTile output_tile) const; // Decides if the problem is small enough so it makes sense to trade off @@ -169,10 +180,9 @@ class TritonDotFusionSearchSpace { // input tilings. int GetMaxNumStages(OutputTile output_tile, int contracting_tile_size) const; - // Finds all promising output shape tilings (block_m, block_n), based on - // `config` with already determined contracting split value and appends them - // to `updated_configs`. Each config in the input list might yield zero or - // more configs in the output. + // Finds all promising output shape tilings (block_m, block_n) based on + // `config` and appends them to `updated_configs`. Each config in the input + // list might yield zero or more configs in the output. void AddOutputTilings(const ConfigWithNotes& config, std::vector& updated_configs) const; @@ -184,17 +194,17 @@ class TritonDotFusionSearchSpace { std::vector& updated_configs) const; // Finds all promising values for the contracting dimension tile size - // (block_k), based on `config` with already determined contracting split and - // output tiling, and appends them to `updated_configs`. Each config in the - // input list might yield zero or more configs in the output. + // (block_k), based on `config` with already determined output tiling, and + // appends them to `updated_configs`. Each config in the input list might + // yield zero or more configs in the output. void AddContractingTiling( const ConfigWithNotes& config, std::vector& updated_configs) const; // Finds all promising values for the pipelining parameter, based on - // `config` with already determined contracting split, output tiling, and - // contracting tile size, and appends them to `updated_configs`. Each config - // in the input list might yield zero or more configs in the output. + // `config` with already determined output tiling and contracting tile size, + // and appends them to `updated_configs`. Each config in the input list might + // yield zero or more configs in the output. void AddPipeliningParameter( const ConfigWithNotes& config, std::vector& updated_configs) const; @@ -228,6 +238,8 @@ class TritonDotFusionSearchSpace { int64_t rhs_parallel_size_; int operand_bitwidth_; int compute_bitwidth_; + PrimitiveType operand_type_; + PrimitiveType compute_type_; bool lhs_has_expensive_op_; bool rhs_has_expensive_op_; int desired_total_warps_; diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space_test.cc b/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space_test.cc index 62340d81b85372..43fad93a057a59 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space_test.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space_test.cc @@ -29,6 +29,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/hlo/testlib/verified_hlo_module.h" +#include "xla/service/gpu/gpu_device_info_for_tests.h" #include "xla/service/gpu/matmul_utils.h" #include "xla/stream_executor/cuda/cuda_compute_capability.h" #include "xla/stream_executor/device_description.h" @@ -198,10 +199,88 @@ TEST_F(DotSearchSpaceTest, SerializesSearchSpace) { EXPECT_EQ(search_space.ToString(), "problem_size_BxMxNxKxE: 1x1024x1024x1024x(16->16) " "tile_range_MxNxK: [16-256]x[8-512]x[16-?] " - "desired_total_warps: 2640 occupancy_optimization: 1 " + "desired_total_warps: 1056 occupancy_optimization: 1 " "warps_per_cta: [2-?]"); } +TEST_F(DotSearchSpaceTest, ScalesDesiredWarpsWithArithmeticIntensity) { + // Memory-bound (128x128x128, AI ~ 42.7 FLOP/B): 5 warps/scheduler -> 2640. + ASSERT_OK_AND_ASSIGN( + std::unique_ptr mem_bound_module, + GetDefaultDotModule(/*lhs_parallel_dim=*/128, /*rhs_parallel_dim=*/128, + /*contracting_dim=*/128)); + EXPECT_THAT(MakeSearchSpace(mem_bound_module.get()).ToString(), + ::testing::HasSubstr("desired_total_warps: 2640")); + + // Intermediate (512x512x512, AI ~ 170.7 FLOP/B): 3 warps/scheduler -> 1584. + ASSERT_OK_AND_ASSIGN( + std::unique_ptr mid_module, + GetDefaultDotModule(/*lhs_parallel_dim=*/512, /*rhs_parallel_dim=*/512, + /*contracting_dim=*/512)); + EXPECT_THAT(MakeSearchSpace(mid_module.get()).ToString(), + ::testing::HasSubstr("desired_total_warps: 1584")); + + // Compute-bound (1024x1024x1024, AI ~ 341.3 FLOP/B): 2 warps/scheduler -> + // 1056. + ASSERT_OK_AND_ASSIGN( + std::unique_ptr compute_bound_module, + GetDefaultDotModule(/*lhs_parallel_dim=*/1024, /*rhs_parallel_dim=*/1024, + /*contracting_dim=*/1024)); + EXPECT_THAT(MakeSearchSpace(compute_bound_module.get()).ToString(), + ::testing::HasSubstr("desired_total_warps: 1056")); +} + +TEST_F(DotSearchSpaceTest, ScalesDesiredWarpsWithHardwareDerivedRidgePoint) { + // Use a fully-populated H100 device description containing matrix unit and + // memory bandwidth specifications. + se::DeviceDescription h100_device = TestGpuDeviceInfo::H100SXMDeviceInfo(); + + // Memory-bound (128x128x128, AI ~ 42.7 FLOP/B): 5 warps/scheduler -> 2640. + ASSERT_OK_AND_ASSIGN( + std::unique_ptr mem_bound_module, + GetDefaultDotModule(/*lhs_parallel_dim=*/128, /*rhs_parallel_dim=*/128, + /*contracting_dim=*/128)); + TritonDotFusionSearchSpace mem_space(h100_device, + GetDot(mem_bound_module.get())); + EXPECT_THAT(mem_space.ToString(), + ::testing::HasSubstr("desired_total_warps: 2640")); + + // Intermediate (512x512x512, AI ~ 170.7 FLOP/B): 3 warps/scheduler -> 1584. + ASSERT_OK_AND_ASSIGN( + std::unique_ptr mid_module, + GetDefaultDotModule(/*lhs_parallel_dim=*/512, /*rhs_parallel_dim=*/512, + /*contracting_dim=*/512)); + TritonDotFusionSearchSpace mid_space(h100_device, GetDot(mid_module.get())); + EXPECT_THAT(mid_space.ToString(), + ::testing::HasSubstr("desired_total_warps: 1584")); + + // Compute-bound (1024x1024x1024, AI ~ 341.3 FLOP/B): 2 warps/scheduler -> + // 1056. + ASSERT_OK_AND_ASSIGN( + std::unique_ptr compute_bound_module, + GetDefaultDotModule(/*lhs_parallel_dim=*/1024, /*rhs_parallel_dim=*/1024, + /*contracting_dim=*/1024)); + TritonDotFusionSearchSpace compute_space(h100_device, + GetDot(compute_bound_module.get())); + EXPECT_THAT(compute_space.ToString(), + ::testing::HasSubstr("desired_total_warps: 1056")); +} + +TEST_F(DotSearchSpaceTest, EliminatesLowOccupancyTilesForMediumProblem) { + ASSERT_OK_AND_ASSIGN( + std::unique_ptr module, + GetDefaultDotModule(/*lhs_parallel_dim=*/1024, /*rhs_parallel_dim=*/1024, + /*contracting_dim=*/1024)); + TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); + + // On 1024x1024 with 132 cores, tiles >= 128x128 yield <= 64 result tiles + // (< 132 cores) and should be eliminated in non-exhaustive mode. + EXPECT_THAT( + search_space.GenerateConfigs(), + AllOf(Not(IsEmpty()), + Not(Contains(AllOf(BlockMIs(Ge(128)), BlockNIs(Ge(128))))))); +} + TEST_F(DotSearchSpaceTest, ReturnsValidConfigList) { ASSERT_OK_AND_ASSIGN(std::unique_ptr module, GetDefaultDotModule()); diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc index 9503fafab09fc6..7e26284faaecd0 100644 --- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc +++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc @@ -146,13 +146,13 @@ bool IsSupportedReductionOpcode(HloOpcode opcode) { CodegenDecision CheckCompareInstruction(const HloInstruction& instr) { auto compare = Cast(&instr); - if (compare->order() == ComparisonOrder::kTotal && - primitive_util::IsFloatingPointType( - compare->operand(0)->shape().element_type())) { - return CodegenDecision::Forbid( - "Total order comparison is not supported for floating-point types"); + if (!primitive_util::IsFloatingPointType( + compare->operand(0)->shape().element_type()) || + compare->order() == ComparisonOrder::kPartial) { + return CodegenDecision::Allow(); } - return CodegenDecision::Allow(); + return CodegenDecision::Forbid( + "Only partial order comparison is supported for floating-point types"); } CodegenDecision IsSupportedReductionInstruction(const HloInstruction& instr) { diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_dynamic_slice.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_dynamic_slice.hlo index ebc93e938627a0..af6abaa23ad8c9 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_dynamic_slice.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_dynamic_slice.hlo @@ -19,7 +19,7 @@ fusion (p0: f32[1024, 16], offset: s32[]) -> f32[8192] { // CHECK-SAME: %[[ARG0:.*]]: memref<1024x16xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref, // CHECK-SAME: %[[ARG2:.*]]: memref<8192xf32>, -// CHECK-SAME: %[[ARG3:.*]]: index) attributes {num_opaque_args = 0 : i32} { +// CHECK-SAME: %[[ARG3:.*]]: index{{( \{xla.range = \[0 : index, 1023 : index\]\})?}}) attributes {num_opaque_args = 0 : i32} { // CHECK: %[[CONSTANT_0:.*]] = arith.constant 512 : i32 // CHECK: %[[CONSTANT_1:.*]] = arith.constant 0 : i32 // CHECK: %[[EXTRACT_0:.*]] = xtile.extract %[[ARG1]][] [] [] : memref -> tensor diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_1.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_1.hlo index fcb4d7224cb118..ef302c3620a85a 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_1.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_1.hlo @@ -9,7 +9,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<16xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<1x16xf32>, -// CHECK-SAME: %[[ARG2:.*]]: index) +// CHECK-SAME: %[[ARG2:.*]]: index {xla.range = [0 : index, 0 : index]}) // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C0]]] [16] [1] : memref<16xf32> -> tensor<16xf32> // CHECK: %[[EXPAND:.*]] = tt.expand_dims %[[EXTRACT]] {axis = 0 : i32} : tensor<16xf32> -> tensor<1x16xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_4.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_4.hlo index bf47db610c2b1a..79ddc336308315 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_4.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_4.hlo @@ -9,7 +9,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<16xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<1x16xf32>, -// CHECK-SAME: %[[ARG2:.*]]: index) +// CHECK-SAME: %[[ARG2:.*]]: index {xla.range = [0 : index, 0 : index]}) // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C0]]] [16] [1] : memref<16xf32> -> tensor<16xf32> // CHECK: %[[EXPAND:.*]] = tt.expand_dims %[[EXTRACT]] {axis = 0 : i32} : tensor<16xf32> -> tensor<1x16xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_transposed.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_transposed.hlo index 7d038bc4bbf6d7..ac17b1d58d0a01 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_transposed.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_transposed.hlo @@ -12,7 +12,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<64x16xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<1x16x64xf32, #xtile.layout<[1, 0, 2]>>, -// CHECK-SAME: %[[ARG2:.*]]: index) +// CHECK-SAME: %[[ARG2:.*]]: index {xla.range = [0 : index, 63 : index]}) // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[INDEX0:.*]] = xla.apply_indexing #indexing_map(%[[ARG2]]) // CHECK: %[[INDEX1:.*]] = xla.apply_indexing #indexing_map1(%[[ARG2]]) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_normalized.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_normalized.hlo index 0b4ad323b3898f..63204afc8d231c 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_normalized.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_normalized.hlo @@ -22,7 +22,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<8x48xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<8x3x16xf32>, -// CHECK-SAME: %[[PID:.*]]: index) attributes {num_opaque_args = 0 : i32} { +// CHECK-SAME: %[[PID:.*]]: index{{( \{xla.range = \[0 : index, 23 : index\]\})?}}) attributes {num_opaque_args = 0 : i32} { // CHECK: %[[OFFSET0:.*]] = xla.apply_indexing #[[$ATTR_0]](%[[PID]]) // CHECK: %[[OFFSET1:.*]] = xla.apply_indexing #[[$ATTR_1]](%[[PID]]) // CHECK: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[OFFSET0]], %[[OFFSET1]]] [2, 8] [1, 1] : memref<8x48xf32> -> tensor<2x8xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_collapse.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_collapse.hlo index 47f4abb6f64ef5..67b119916a61b7 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_collapse.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_collapse.hlo @@ -15,7 +15,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<4096x1x768xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<32x128x768xf32>, -// CHECK-SAME: %[[ARG2:.*]]: index) attributes {num_opaque_args = 0 : i32} { +// CHECK-SAME: %[[ARG2:.*]]: index{{( \{xla.range = \[0 : index, 255 : index\]\})?}}) attributes {num_opaque_args = 0 : i32} { // CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index // CHECK: %[[APPLY_INDEXING_0:.*]] = xla.apply_indexing #[[$ATTR_0]](%[[ARG2]]) // CHECK-OLD: %[[EXTRACT_0:.*]] = xtile.extract %[[ARG0]][%[[APPLY_INDEXING_0]], %[[CONSTANT_0]], %[[CONSTANT_0]]] [16, 1, 1024] [128, 0, 1] : memref<4096x1x768xf32> -> tensor<16x1x1024xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_grouped.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_grouped.hlo index c04da10988b2d4..d98a85e4b4bb21 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_grouped.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_grouped.hlo @@ -12,7 +12,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<64xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<32x1x2xf32>, -// CHECK-SAME: %[[ARG2:.*]]: index) +// CHECK-SAME: %[[ARG2:.*]]: index {xla.range = [0 : index, 15 : index]}) // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[INDEX0:.*]] = xla.apply_indexing #indexing_map(%[[ARG2]]) // CHECK: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[INDEX0]]] [4] [1] : memref<64xf32> -> tensor<4xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_1.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_1.hlo index d5c98d466e358a..17ccd97254425f 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_1.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_1.hlo @@ -12,7 +12,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<16xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<1x4x4xf32>, -// CHECK-SAME: %[[ARG2:.*]]: index) +// CHECK-SAME: %[[ARG2:.*]]: index {xla.range = [0 : index, 1 : index]}) // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[INDEX0:.*]] = xla.apply_indexing #indexing_map(%[[ARG2]]) // CHECK: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[INDEX0]]] [8] [1] : memref<16xf32> -> tensor<8xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_2.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_2.hlo index 91a3db850d72f5..e69acf73893a44 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_2.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_2.hlo @@ -12,7 +12,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<16xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<1x4x4xf32>, -// CHECK-SAME: %[[ARG2:.*]]: index) +// CHECK-SAME: %[[ARG2:.*]]: index {xla.range = [0 : index, 1 : index]}) // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[INDEX0:.*]] = xla.apply_indexing #indexing_map(%[[ARG2]]) // CHECK: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[INDEX0]]] [8] [1] : memref<16xf32> -> tensor<8xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_collapse.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_collapse.hlo index a188a1412e2175..b2929845c73bd4 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_collapse.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_collapse.hlo @@ -15,7 +15,7 @@ fusion { // CHECK: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<32xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<4x1xf32>, -// CHECK-SAME: %[[PID:.*]]: index) +// CHECK-SAME: %[[PID:.*]]: index{{( \{xla.range = \[0 : index, 1 : index\]\})?}}) // CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index // CHECK: %[[OFFSET_IN:.*]] = xla.apply_indexing #[[$MAP_0]](%[[PID]]) // CHECK: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[OFFSET_IN]]] [2] [8] : memref<32xf32> -> tensor<2xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_expand.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_expand.hlo index 1eaac915c64527..bfa7cbbe1098e2 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_expand.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_expand.hlo @@ -14,7 +14,7 @@ fusion { // CHECK: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<32x4xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<32xf32>, -// CHECK-SAME: %[[PID:.*]]: index) +// CHECK-SAME: %[[PID:.*]]: index{{( \{xla.range = \[0 : index, 1 : index\]\})?}}) // CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index // CHECK: %[[OFFSET:.*]] = xla.apply_indexing #[[$MAP_0]](%[[PID]]) // CHECK-OLD: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[OFFSET]], %[[CONSTANT_0]]] [16, 1] [1, 0] : memref<32x4xf32> -> tensor<16x1xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_only.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_only.hlo index a9f37112e6041b..b54b9cb33b3b47 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_only.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_only.hlo @@ -20,7 +20,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<64xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<8x8xf32>, -// CHECK-SAME: %[[ARG2:.*]]: index) attributes {num_opaque_args = 0 : i32} { +// CHECK-SAME: %[[ARG2:.*]]: index{{( \{xla.range = \[0 : index, 15 : index\]\})?}}) attributes {num_opaque_args = 0 : i32} { // CHECK: %[[APPLY_INDEXING_0:.*]] = xla.apply_indexing #[[$ATTR_0]](%[[ARG2]]) // CHECK: %[[EXTRACT_0:.*]] = xtile.extract %[[ARG0]]{{\[}}%[[APPLY_INDEXING_0]]] [4] [1] : memref<64xf32> -> tensor<4xf32> // CHECK: %[[RESHAPE_0:.*]] = tt.reshape %[[EXTRACT_0]] : tensor<4xf32> -> tensor<1x4xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d.hlo index 1ae15c79368b3b..3c153a7c2ac64c 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d.hlo @@ -11,7 +11,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<1x1x1xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<1xf32>, -// CHECK-SAME: %[[PID:.*]]: index) attributes {num_opaque_args = 0 : i32} { +// CHECK-SAME: %[[PID:.*]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}}) attributes {num_opaque_args = 0 : i32} { // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK-OLD: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [1, 1, 1] [0, 0, 0] : memref<1x1x1xf32> -> tensor<1x1x1xf32> // CHECK-NEW: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [1, 1, 1] [1, 1, 1] : memref<1x1x1xf32> -> tensor<1x1x1xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d_tiled.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d_tiled.hlo index 79cf372bf91ae0..c2d0c87737e750 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d_tiled.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d_tiled.hlo @@ -11,7 +11,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<1x1x1xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<1xf32>, -// CHECK-SAME: %[[PID:.*]]: index) attributes {num_opaque_args = 0 : i32} { +// CHECK-SAME: %[[PID:.*]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}}) attributes {num_opaque_args = 0 : i32} { // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [1, 1, 1] {{.*}} : memref<1x1x1xf32> -> tensor<1x1x1xf32> // CHECK: %[[UNSPLAT:.*]] = tt.unsplat %[[EXTRACT]] : tensor<1x1x1xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_scalar.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_scalar.hlo index 2580a5dd63169b..f3cbbafea04887 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_scalar.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_scalar.hlo @@ -11,7 +11,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<1x1x1xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref, -// CHECK-SAME: %[[PID:.*]]: index) attributes {num_opaque_args = 0 : i32} { +// CHECK-SAME: %[[PID:.*]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}}) attributes {num_opaque_args = 0 : i32} { // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK-OLD: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [1, 1, 1] [0, 0, 0] : memref<1x1x1xf32> -> tensor<1x1x1xf32> // CHECK-NEW: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [1, 1, 1] [1, 1, 1] : memref<1x1x1xf32> -> tensor<1x1x1xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_only.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_only.hlo index 86e9221178191c..d37186f7a65376 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_only.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_only.hlo @@ -17,7 +17,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<42x5xi8, #xtile.layout<[0, 1]>>, // CHECK-SAME: %[[ARG1:.*]]: memref<5x42xi8>, -// CHECK-SAME: %[[PID:.*]]: index) attributes {num_opaque_args = 0 : i32} { +// CHECK-SAME: %[[PID:.*]]: index{{( \{xla.range = \[0 : index, 83 : index\]\})?}}) attributes {num_opaque_args = 0 : i32} { // CHECK: %[[OFFSET0:.*]] = xla.apply_indexing #[[$ATTR_0]](%[[PID]]) // CHECK: %[[OFFSET1:.*]] = xla.apply_indexing #[[$ATTR_1]](%[[PID]]) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_permutation_inversion.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_permutation_inversion.hlo index cc87ec0120fbdd..9b7c572cc8c6bf 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_permutation_inversion.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_permutation_inversion.hlo @@ -12,7 +12,7 @@ fusion { // CHECK: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:[^:]+]]: memref<8x16x16x4xbf16, #xtile.layout<[2, 0, 3, 1]>>, // CHECK-SAME: %[[ARG1:[^:]+]]: memref<16x4x8x16xbf16>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}}) // CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index // CHECK: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[CONSTANT_0]], %[[CONSTANT_0]], %[[CONSTANT_0]], %[[CONSTANT_0]]] [8, 16, 16, 4] [1, 1, 1, 1] : memref<8x16x16x4xbf16, #xtile.layout<[2, 0, 3, 1]>> -> tensor<8x16x16x4xbf16> // CHECK: %[[TRANSPOSE:.*]] = tt.trans %[[EXTRACT]] {order = array} : tensor<8x16x16x4xbf16> -> tensor<16x4x8x16xbf16> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo index f7ced8b68f0753..39b3bc64e1b787 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo @@ -36,7 +36,7 @@ ENTRY entry { // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: i32, // CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: i32, // CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: !tt.ptr, -// CHECK-SAME: %[[PID:[a-zA-Z0-9_]+]]: index) +// CHECK-SAME: %[[PID:[a-zA-Z0-9_]+]]: index {xla.range = [0 : index, 127 : index]}) // CHECK: %[[IDX0:.*]] = xla.apply_indexing #[[MAP0]](%[[PID]]) // CHECK: %[[IDX1:.*]] = xla.apply_indexing #[[MAP1]](%[[PID]]) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo index 1855eb967a4034..66986fdad786db 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo @@ -19,7 +19,7 @@ fusion{ // CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: i32, // CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: !tt.ptr, // CHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: !tt.ptr, -// CHECK-SAME: %[[PID:[a-zA-Z0-9_]+]]: index) +// CHECK-SAME: %[[PID:[a-zA-Z0-9_]+]]: index{{( \{xla.range = \[0 : index, 15 : index\]\})?}}) // CHECK-SAME: attributes {num_opaque_args = 4 : i32} { // CHECK: %[[EXTRACT_0:.*]] = xtile.extract diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot.hlo index f3b4996da427f1..35d0f38d580a77 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot.hlo @@ -23,7 +23,7 @@ fusion { // CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]: memref<32x16x256xf32, #xtile.layout<[2, 0, 1]>>, // CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]: memref<256x16x512xf32, #xtile.layout<[0, 1, 2]>>, // CHECK-SAME: %[[ARG2:[A-Za-z0-9]+]]: memref<16x32x512xf32, #xtile.layout<[2, 0, 1]>>, -// CHECK-SAME: %[[ARG3:[A-Za-z0-9]+]]: index) attributes {num_opaque_args = 0 : i32} { +// CHECK-SAME: %[[ARG3:[A-Za-z0-9]+]]: index{{( \{xla.range = \[0 : index, 255 : index\]\})?}}) attributes {num_opaque_args = 0 : i32} { // CHECK-DAG: %[[TC0:.*]] = arith.constant dense<0.000000e+00> : tensor<16x64xf32> // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_broadcast_sequential_dim_mismatch.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_broadcast_sequential_dim_mismatch.hlo index 2972d8832d7d67..247321c43ca601 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_broadcast_sequential_dim_mismatch.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_broadcast_sequential_dim_mismatch.hlo @@ -17,7 +17,7 @@ fusion { // CHECK-SAME: %[[ARG0:[^:]+]]: memref<16x32xf32>, // CHECK-SAME: %[[ARG1:[^:]+]]: memref<32x16xf32>, // CHECK-SAME: %[[ARG2:[^:]+]]: memref<2x16x16xf32>, -// CHECK-SAME: %[[PID:[^:]+]]: index) +// CHECK-SAME: %[[PID:[^:]+]]: index{{( \{xla.range = \[0 : index, 1 : index\]\})?}}) // CHECK-DAG: %[[CST:.*]] = arith.constant dense<0.000000e+00> : tensor<16x16xf32> // CHECK-DAG: %[[CONSTANT_0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[CONSTANT_1:.*]] = arith.constant 1 : index diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/duplicate_dot_operand.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/duplicate_dot_operand.hlo index be87f719908f51..33a6a4320a82b5 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/duplicate_dot_operand.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/duplicate_dot_operand.hlo @@ -11,7 +11,7 @@ fusion { // CHECK-SAME: %[[ARG0:.*]]: memref<4x8xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<8x16xf32>, // CHECK-SAME: %[[ARG2:.*]]: memref<4x16xf32>, -// CHECK-SAME: %[[ARG3:.*]]: index) +// CHECK-SAME: %[[ARG3:.*]]: index {xla.range = [0 : index, 31 : index]}) // CHECK-DAG: %[[CST:.*]] = arith.constant dense<0.000000e+00> : tensor<1x2xf32> // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[APPLY_IDX0:.*]] = xla.apply_indexing #indexing_map(%[[ARG3]]) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_multiple_dimensions.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_multiple_dimensions.hlo index f646d96b5cd889..f1cb0233879b06 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_multiple_dimensions.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_multiple_dimensions.hlo @@ -10,7 +10,7 @@ // CHECK-SAME: %[[ARG0:.*]]: memref<64x32x16xi32>, // CHECK-SAME: %[[ARG1:.*]]: memref, // CHECK-SAME: %[[ARG2:.*]]: memref<10x10x8xi32>, -// CHECK-SAME: %[[ARG3:.*]]: index) +// CHECK-SAME: %[[ARG3:.*]]: index{{( \{xla.range = \[0 : index, 1 : index\]\})?}}) // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C50:.*]] = arith.constant 50 : index // CHECK-DAG: %[[C22:.*]] = arith.constant 22 : i32 diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_add.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_add.hlo index a0db3183e293e0..5a7f06d1dd45c2 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_add.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_add.hlo @@ -14,7 +14,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<64xi32>, // CHECK-SAME: %[[ARG1:.*]]: memref<10xi32>, -// CHECK-SAME: %[[ARG2:.*]]: index) +// CHECK-SAME: %[[ARG2:.*]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}}) // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C20:.*]] = arith.constant 20 : index // CHECK: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C20]]] [32] [1] : memref<64xi32> -> tensor<32xi32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice.hlo index 246df4f03ead2c..42844dfa63c9f2 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice.hlo @@ -20,7 +20,7 @@ fusion { // CHECK-NEW: xtile.entry_func @triton_fn( // CHECK-NEW-SAME: %[[ARG0:.*]]: memref<64xf32>, // CHECK-NEW-SAME: %[[ARG1:.*]]: memref<10xi32>, -// CHECK-NEW-SAME: %[[ARG2:.*]]: index) +// CHECK-NEW-SAME: %[[ARG2:.*]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}}) // CHECK-NEW-DAG: %[[C10:.*]] = arith.constant 10 : i32 // CHECK-NEW-DAG: %[[C44:.*]] = arith.constant 44 : i32 // CHECK-NEW-DAG: %[[C0_I32:.*]] = arith.constant 0 : i32 diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice_no_convert.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice_no_convert.hlo index 47de9a9b17eb2d..865050e7a74914 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice_no_convert.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice_no_convert.hlo @@ -18,7 +18,7 @@ fusion { // CHECK-NEW: xtile.entry_func @triton_fn( // CHECK-NEW-SAME: %[[ARG0:.*]]: memref<64xi32>, // CHECK-NEW-SAME: %[[ARG1:.*]]: memref<10xi32>, -// CHECK-NEW-SAME: %[[ARG2:.*]]: index) +// CHECK-NEW-SAME: %[[ARG2:.*]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}}) // CHECK-NEW-DAG: %[[C10:.*]] = arith.constant 10 : i32 // CHECK-NEW-DAG: %[[C44:.*]] = arith.constant 44 : i32 // CHECK-NEW-DAG: %[[C0_I32:.*]] = arith.constant 0 : i32 diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_offset_data_same_input.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_offset_data_same_input.hlo index 8342fa8208623b..cdeda3bc59a8ec 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_offset_data_same_input.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_offset_data_same_input.hlo @@ -13,7 +13,7 @@ fusion { // CHECK: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:[A-Za-z0-9]+]]: memref<64xi32>, // CHECK-SAME: %[[ARG1:[A-Za-z0-9]+]]: memref<10xi32>, -// CHECK-SAME: %[[PID:[A-Za-z0-9]+]]: index) attributes {num_opaque_args = 0 : i32} +// CHECK-SAME: %[[PID:[A-Za-z0-9]+]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}}) attributes {num_opaque_args = 0 : i32} // CHECK-DAG: %[[C54_I32:.*]] = arith.constant 54 : i32 // CHECK-DAG: %[[C0_I32:.*]] = arith.constant 0 : i32 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_three_tiles_per_block.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_three_tiles_per_block.hlo index 7423871947adf4..6fbe52e11dfa41 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_three_tiles_per_block.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_three_tiles_per_block.hlo @@ -30,7 +30,7 @@ ENTRY main { // CHECK-SAME: %[[LHS:[a-zA-Z0-9]*]]: memref<8x4xf32>, // CHECK-SAME: %[[RHS:[a-zA-Z0-9]*]]: memref<8x4xf32>, // CHECK-SAME: %[[OUT:[a-zA-Z0-9]*]]: memref<8x4xf32>, -// CHECK-SAME: %[[PID:[a-zA-Z0-9]*]]: index) +// CHECK-SAME: %[[PID:[a-zA-Z0-9]*]]: index {xla.range = [0 : index, 2 : index]}) // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_two_tiles_per_block.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_two_tiles_per_block.hlo index 23ad37d9f3ebeb..49698a9ce8017b 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_two_tiles_per_block.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_two_tiles_per_block.hlo @@ -29,7 +29,7 @@ ENTRY main { // CHECK-SAME: %[[LHS:[a-zA-Z0-9]*]]: memref<8x4xf32>, // CHECK-SAME: %[[RHS:[a-zA-Z0-9]*]]: memref<8x4xf32>, // CHECK-SAME: %[[OUT:[a-zA-Z0-9]*]]: memref<8x4xf32>, -// CHECK-SAME: %[[PID:[a-zA-Z0-9]*]]: index) +// CHECK-SAME: %[[PID:[a-zA-Z0-9]*]]: index {xla.range = [0 : index, 3 : index]}) // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/not.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/not.hlo index 13c52cffc7fe85..4523f518110051 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/not.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/not.hlo @@ -11,7 +11,7 @@ fusion { // CHECK: xtile.entry_func // CHECK-SAME: (%[[INPUT:[a-zA-Z0-9]*]]: memref<100xi32>, -// CHECK-SAME: %[[OUTPUT:[a-zA-Z0-9]*]]: memref<100xi32>, %[[PID:.*]]: index) +// CHECK-SAME: %[[OUTPUT:[a-zA-Z0-9]*]]: memref<100xi32>, %[[PID:.*]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}}) // CHECK-SAME: attributes {num_opaque_args = 0 : i32} { // CHECK-DAG: %[[Cm1:.*]] = arith.constant dense<-1> : tensor<128xi32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc index f0d5a10809ef9c..d4c17e5f31c651 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc @@ -530,7 +530,7 @@ ENTRY main { auto xtile_module_and_hlo_module, CreateXTileIrAndFileCheck(std::move(module), "triton_softmax_computation", R"( -CHECK: xtile.entry_func @xtile_dialect_fn(%[[P0:.*]]: {{.*}}, %[[P1:.*]]: {{.*}}, %[[PID:.*]]: index) +CHECK: xtile.entry_func @xtile_dialect_fn(%[[P0:.*]]: {{.*}}, %[[P1:.*]]: {{.*}}, %[[PID:[^:]*]]: index{{( \{xla.range = \[0 : index, 124 : index\]\})?}}) CHECK-DAG: %[[C_0:.*]] = arith.constant 0 : index CHECK: xtile.extract %[[P0]] CHECK-SAME: [%[[PID]], %{{.*}}] [1, 128] [1, 1] @@ -544,7 +544,7 @@ CHECK: } TF_EXPECT_OK(LowerXTileIrToTritonAndFileCheck( xtile_module_and_hlo_module.first.get(), R"( -CHECK: xtile.entry_func @xtile_dialect_fn(%[[P0:.*]]: {{.*}}, %[[P1:.*]]: {{.*}}, %[[PID:.*]]: index) +CHECK: xtile.entry_func @xtile_dialect_fn(%[[P0:.*]]: {{.*}}, %[[P1:.*]]: {{.*}}, %[[PID:[^:]*]]: index{{( \{xla.range = \[0 : index, 124 : index\]\})?}}) CHECK-DAG: %[[C_0:.*]] = arith.constant 0 : index CHECK: xtile.extract %[[P0]] CHECK-SAME: [%[[PID]], %{{.*}}] [1, 128] [1, 1] @@ -609,7 +609,7 @@ CHECK: xtile.entry_func @xtile_dialect_fn( CHECK-SAME: %[[P0:[A-Za-z0-9_]*]]: memref<125x127xf32> CHECK-SAME: %[[P1:[A-Za-z0-9_]*]]: memref<127xf32> CHECK-SAME: %[[P2:[A-Za-z0-9_]*]]: memref<125x127xf32> -CHECK-SAME: %[[TID:[A-Za-z0-9_]*]]: index) +CHECK-SAME: %[[TID:[A-Za-z0-9_]*]]: index{{( \{xla.range = \[0 : index, 124 : index\]\})?}}) CHECK: xtile.extract %[[P0]][%[[TID]], %{{.*}}] [1, 128] [1, 1] : {{.*}} -> tensor<1x128xf32> CHECK: %{{.*}} = arith.constant 0 : index CHECK: xtile.extract %[[P1]][%{{.*}}] [128] [1] : {{.*}} -> tensor<128xf32> @@ -625,7 +625,7 @@ CHECK: xtile.entry_func @xtile_dialect_fn( CHECK-SAME: %[[P0:[A-Za-z0-9_]*]]: memref<125x127xf32> CHECK-SAME: %[[P1:[A-Za-z0-9_]*]]: memref<127xf32> CHECK-SAME: %[[P2:[A-Za-z0-9_]*]]: memref<125x127xf32> -CHECK-SAME: %[[TID:[A-Za-z0-9_]*]]: index) +CHECK-SAME: %[[TID:[A-Za-z0-9_]*]]: index{{( \{xla.range = \[0 : index, 124 : index\]\})?}}) CHECK: xtile.extract %[[P0]][%[[TID]], %{{.*}}] [1, 128] [1, 1] : {{.*}} -> tensor<1x128xf32> CHECK: xtile.extract %[[P1]][%{{.*}}] [128] [1] : {{.*}} -> tensor<128xf32> CHECK: tt.reduce diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/generic_emitter_multi_tiled_dims.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/generic_emitter_multi_tiled_dims.hlo index ac80da89e2b67e..19317f95bdc71d 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/generic_emitter_multi_tiled_dims.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/generic_emitter_multi_tiled_dims.hlo @@ -27,7 +27,7 @@ fusion { // CHECK-NEW: #[[MAP:.*]] = #xla.indexing_map<"(pid) -> (pid / 125), domain: pid in [0, 1249]"> // CHECK-NEW: #[[MAP1:.*]] = #xla.indexing_map<"(pid) -> (pid mod 125), domain: pid in [0, 1249]"> // CHECK: xtile.entry_func @triton_fn -// CHECK-SAME: (%[[P0:.*]]: {{.*}}, %[[P1:.*]]: {{.*}}, %[[P2:.*]]: {{.*}}, %[[P3:.*]]: {{.*}}, %[[TID:.*]]: index) +// CHECK-SAME: (%[[P0:.*]]: {{.*}}, %[[P1:.*]]: {{.*}}, %[[P2:.*]]: {{.*}}, %[[P3:.*]]: {{.*}}, %[[TID:.*]]: index{{( \{xla.range = \[0 : index, 1249 : index\]\})?}}) // CHECK-DAG: %[[C_0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[ROW_INDEX:.*]] = xla.apply_indexing #[[MAP]](%[[TID]] // CHECK-DAG: %[[COL_INDEX:.*]] = xla.apply_indexing #[[MAP1]](%[[TID]] diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_multi_dim_tiled_and_untiled.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_multi_dim_tiled_and_untiled.hlo index 4e8acaae3f1a8b..fc20b5425d488a 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_multi_dim_tiled_and_untiled.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_multi_dim_tiled_and_untiled.hlo @@ -1,7 +1,5 @@ -// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s 2>&1 | FileCheck %s -// RUN: not triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-4 %s - -// CHECK: MapSymbolIdToSequentialDimValue failed for symbolic dimension +// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s +// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-4 %s region { p0 = f32[] parameter(0) @@ -28,3 +26,14 @@ ENTRY main { "num_ctas":"1", "num_stages":"1"}}} } + +// Reduction dim 0 is fully tiled (8/8) and reads at offset 0. Only dims 2 +// (12/4 = 3 iterations) and 4 (16/8 = 2 iterations) are looped over. +// CHECK-LABEL: xtile.entry_func @triton_fn +// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index +// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index +// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index +// CHECK: %[[ACC:.*]] = scf.for {{.*}} = %[[C0]] to %[[C3]] +// CHECK: scf.for {{.*}} = %[[C0]] to %[[C2]] +// CHECK: xtile.extract %arg0[%[[C0]], +// CHECK: "tt.reduce"(%[[ACC]]) diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reshape/reshape_expand_high_rank.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reshape/reshape_expand_high_rank.hlo index d9dbf79b56448f..17f6459afe4049 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reshape/reshape_expand_high_rank.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reshape/reshape_expand_high_rank.hlo @@ -11,7 +11,7 @@ fusion { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<512x16xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<1x1x8x64x16xf32>, -// CHECK-SAME: %[[ARG2:.*]]: index) +// CHECK-SAME: %[[ARG2:.*]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}}) // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C0]], %[[C0]]] [512, 16] [1, 1] : memref<512x16xf32> -> tensor<512x16xf32> // CHECK: %[[RESHAPE:.*]] = tt.reshape %[[EXTRACT]] : tensor<512x16xf32> -> tensor<8x64x16xf32> diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot.hlo index e6b561e1e73042..911ac19c0fc292 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot.hlo @@ -21,7 +21,7 @@ triton_dot { // CHECK-SAME: %[[LHS_SCALE:.*]]: memref<128x4xi8>, // CHECK-SAME: %[[RHS_SCALE:.*]]: memref<4x256xi8>, // CHECK-SAME: %[[OUTPUT:.*]]: memref<128x256xbf16>, -// CHECK-SAME: {{.*}}: index) attributes {num_opaque_args = 0 : i32} { +// CHECK-SAME: {{.*}}: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}}) attributes {num_opaque_args = 0 : i32} { // CHECK-DAG: %[[C0:.*]] = arith.constant dense<0.000000e+00> : tensor<128x256xf32> // CHECK-DAG: %[[C1:.*]] = arith.constant 0 : index // CHECK: %[[LHS_TILE:.*]] = xtile.extract %[[LHS]] diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/boundary_dynamic_slice.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/boundary_dynamic_slice.hlo index a139350d838964..9ad87bf941699a 100644 --- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/boundary_dynamic_slice.hlo +++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/boundary_dynamic_slice.hlo @@ -35,7 +35,7 @@ ENTRY main { // CHECK-LABEL: xtile.entry_func @triton_fn( // CHECK-SAME: %[[ARG0:.*]]: memref<2x97xf32>, // CHECK-SAME: %[[ARG1:.*]]: memref<2xf32>, -// CHECK-SAME: %[[ARG2:.*]]: index) +// CHECK-SAME: %[[ARG2:.*]]: index{{( \{xla.range = \[0 : index, 1 : index\]\})?}}) // CHECK: %[[CST:.*]] = arith.constant 0xFF800000 : f32 // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[ARG2]], %[[C0]]] [1, 128] [1, 1] : memref<2x97xf32> -> tensor<1x128xf32> diff --git a/third_party/xla/xla/backends/gpu/collectives/rccl_types.cc b/third_party/xla/xla/backends/gpu/collectives/rccl_types.cc index 00a37384cd70bc..a0d579b6b6d64d 100644 --- a/third_party/xla/xla/backends/gpu/collectives/rccl_types.cc +++ b/third_party/xla/xla/backends/gpu/collectives/rccl_types.cc @@ -36,13 +36,21 @@ absl::StatusOr ToNcclDataType( stream_executor::RocmComputeCapability rocm_cc) { switch (dtype) { case F8E5M2: - return rocm_cc.has_ocp_fp8_support() ? ncclFloat8e5m2 : ncclInt8; case F8E4M3FN: - return rocm_cc.has_ocp_fp8_support() ? ncclFloat8e4m3 : ncclInt8; - case S8: + if (rocm_cc.has_ocp_fp8_support()) { + return dtype == F8E5M2 ? ncclFloat8e5m2 : ncclFloat8e4m3; + } + [[fallthrough]]; case F8E5M2FNUZ: case F8E4M3FNUZ: case F8E8M0FNU: + if (is_reduction_op) { + return InvalidArgument( + "Unsupported data type for reduction operation: %s", + primitive_util::LowercasePrimitiveTypeName(dtype)); + } + [[fallthrough]]; + case S8: return ncclInt8; case PRED: case U8: diff --git a/third_party/xla/xla/backends/gpu/collectives/rccl_types_test.cc b/third_party/xla/xla/backends/gpu/collectives/rccl_types_test.cc index fbe29ca3890711..5ed0a21b39c7aa 100644 --- a/third_party/xla/xla/backends/gpu/collectives/rccl_types_test.cc +++ b/third_party/xla/xla/backends/gpu/collectives/rccl_types_test.cc @@ -80,6 +80,23 @@ TEST(NcclTypesTest, ToNcclDataTypeNoOcpFp8Support) { se::RocmComputeCapability cc("gfx942"); EXPECT_THAT(ToNcclDataType(F8E5M2, false, cc), IsOkAndHolds(ncclInt8)); EXPECT_THAT(ToNcclDataType(F8E4M3FN, false, cc), IsOkAndHolds(ncclInt8)); + EXPECT_THAT(ToNcclDataType(F8E8M0FNU, false, cc), IsOkAndHolds(ncclInt8)); + EXPECT_THAT(ToNcclDataType(F8E5M2FNUZ, false, cc), IsOkAndHolds(ncclInt8)); + EXPECT_THAT(ToNcclDataType(F8E4M3FNUZ, false, cc), IsOkAndHolds(ncclInt8)); +} + +TEST(NcclTypesTest, ToNcclDataTypeNoOcpFp8SupportReductionFails) { + se::RocmComputeCapability cc("gfx942"); + EXPECT_THAT(ToNcclDataType(F8E5M2, true, cc), + StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_THAT(ToNcclDataType(F8E4M3FN, true, cc), + StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_THAT(ToNcclDataType(F8E8M0FNU, true, cc), + StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_THAT(ToNcclDataType(F8E5M2FNUZ, true, cc), + StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_THAT(ToNcclDataType(F8E4M3FNUZ, true, cc), + StatusIs(absl::StatusCode::kInvalidArgument)); } TEST(NcclTypesTest, ToNcclDataTypeOcpFp8Support) { @@ -87,6 +104,21 @@ TEST(NcclTypesTest, ToNcclDataTypeOcpFp8Support) { EXPECT_THAT(ToNcclDataType(F8E5M2, false, cc), IsOkAndHolds(ncclFloat8e5m2)); EXPECT_THAT(ToNcclDataType(F8E4M3FN, false, cc), IsOkAndHolds(ncclFloat8e4m3)); + EXPECT_THAT(ToNcclDataType(F8E8M0FNU, false, cc), IsOkAndHolds(ncclInt8)); + EXPECT_THAT(ToNcclDataType(F8E5M2FNUZ, false, cc), IsOkAndHolds(ncclInt8)); + EXPECT_THAT(ToNcclDataType(F8E4M3FNUZ, false, cc), IsOkAndHolds(ncclInt8)); +} + +TEST(NcclTypesTest, ToNcclDataTypeOcpFp8SupportReductionFails) { + se::RocmComputeCapability cc("gfx950"); + EXPECT_THAT(ToNcclDataType(F8E5M2, true, cc), IsOkAndHolds(ncclFloat8e5m2)); + EXPECT_THAT(ToNcclDataType(F8E4M3FN, true, cc), IsOkAndHolds(ncclFloat8e4m3)); + EXPECT_THAT(ToNcclDataType(F8E8M0FNU, true, cc), + StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_THAT(ToNcclDataType(F8E5M2FNUZ, true, cc), + StatusIs(absl::StatusCode::kInvalidArgument)); + EXPECT_THAT(ToNcclDataType(F8E4M3FNUZ, true, cc), + StatusIs(absl::StatusCode::kInvalidArgument)); } TEST(NcclTypesTest, ToNcclDataTypeUnsupported) { diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD index 789e090da66f79..31122dd1cba818 100644 --- a/third_party/xla/xla/backends/gpu/runtime/BUILD +++ b/third_party/xla/xla/backends/gpu/runtime/BUILD @@ -422,6 +422,8 @@ xla_test( ":thunk_id", ":thunk_proto_cc", "//xla/service:platform_util", + "//xla/stream_executor:event", + "//xla/stream_executor:mock_stream", "//xla/stream_executor:platform", "//xla/stream_executor:platform_manager", "//xla/stream_executor:stream", @@ -469,10 +471,14 @@ xla_test( ":thunk", ":thunk_executor", ":thunk_id", + "//xla:debug_options_flags", + "//xla:executable_run_options", "//xla/service:buffer_assignment", "//xla/service:platform_util", "//xla/service:service_executable_run_options", + "//xla/service:shaped_buffer", "//xla/service/gpu:buffer_allocations", + "//xla/service/gpu:gpu_executable", "//xla/stream_executor:device_address_allocator", "//xla/stream_executor:platform", "//xla/stream_executor:platform_manager", @@ -482,6 +488,7 @@ xla_test( "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", ], ) diff --git a/third_party/xla/xla/backends/gpu/runtime/async_execution.cc b/third_party/xla/xla/backends/gpu/runtime/async_execution.cc index 26823b190a8800..09e74083c6c58c 100644 --- a/third_party/xla/xla/backends/gpu/runtime/async_execution.cc +++ b/third_party/xla/xla/backends/gpu/runtime/async_execution.cc @@ -70,7 +70,16 @@ AsyncExecution::ExecutionGuard::ExecutionGuard(se::Event* event, se::Stream* async_stream) : event_(event), async_stream_(async_stream) {} +AsyncExecution::ExecutionGuard::ExecutionGuard(ExecutionGuard&& other) noexcept + : event_(std::exchange(other.event_, nullptr)), + async_stream_(std::exchange(other.async_stream_, nullptr)) {} + AsyncExecution::ExecutionGuard::~ExecutionGuard() { + // Nothing to do for moved-from ExecutionGuards. + if (async_stream_ == nullptr) { + return; + } + // If we fail to record completion event on a stream it is unsafe to continue // as the following computations might not see all the updates done by the // async execution. diff --git a/third_party/xla/xla/backends/gpu/runtime/async_execution.h b/third_party/xla/xla/backends/gpu/runtime/async_execution.h index 2f098e83411f96..302dd35e65e8d9 100644 --- a/third_party/xla/xla/backends/gpu/runtime/async_execution.h +++ b/third_party/xla/xla/backends/gpu/runtime/async_execution.h @@ -89,6 +89,10 @@ class AsyncExecution { // when it goes out of scope. class ExecutionGuard { public: + ExecutionGuard(const ExecutionGuard&) = delete; + ExecutionGuard& operator=(const ExecutionGuard&) = delete; + ExecutionGuard(ExecutionGuard&& other) noexcept; + ExecutionGuard& operator=(ExecutionGuard&&) = delete; ~ExecutionGuard(); private: diff --git a/third_party/xla/xla/backends/gpu/runtime/async_execution_test.cc b/third_party/xla/xla/backends/gpu/runtime/async_execution_test.cc index e55ddf5e602f3a..4728793b67ecf6 100644 --- a/third_party/xla/xla/backends/gpu/runtime/async_execution_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/async_execution_test.cc @@ -29,6 +29,8 @@ limitations under the License. #include "xla/backends/gpu/runtime/thunk.pb.h" #include "xla/backends/gpu/runtime/thunk_id.h" #include "xla/service/platform_util.h" +#include "xla/stream_executor/event.h" +#include "xla/stream_executor/mock_stream.h" #include "xla/stream_executor/platform.h" #include "xla/stream_executor/platform_manager.h" #include "xla/stream_executor/stream.h" @@ -85,6 +87,42 @@ TEST(AsyncExecutionTest, InitializeStartDone) { ASSERT_OK(async_execution.Done(&state, stream.get())); } +TEST(AsyncExecutionTest, CompletionEventRecordedOnce) { + // Create an async_execution. + ASSERT_OK_AND_ASSIGN(se::StreamExecutor * executor, CreateExecutor()); + + Thunk::ThunkInfo thunk_info; + thunk_info.thunk_id = ThunkId(1); + thunk_info.profile_annotation = "test-thunk"; + + AsyncExecution async_execution(thunk_info); + Thunk::ExecutionScopedState state; + ASSERT_OK(async_execution.Initialize(&state, executor)); + + // Set up mock streams to verify behavior. + testing::NiceMock stream; + testing::NiceMock async_stream; + + ON_CALL(stream, parent()).WillByDefault(testing::Return(executor)); + ON_CALL(async_stream, WaitFor(&stream)) + .WillByDefault(testing::Return(absl::OkStatus())); + ON_CALL(stream, WaitFor(testing::A())) + .WillByDefault(testing::Return(absl::OkStatus())); + + // The async execution event is recorded only once. + ON_CALL(async_stream, RecordEvent(testing::_)) + .WillByDefault(testing::Return(absl::OkStatus())); + EXPECT_CALL(async_stream, RecordEvent(testing::_)).Times(1); + + // Run the test by calling async_execution.{Start/Done}(). + { + ASSERT_OK_AND_ASSIGN(auto guard, + async_execution.Start(&state, &stream, &async_stream)); + } + + ASSERT_OK(async_execution.Done(&state, &stream)); +} + TEST(AsyncExecutionTest, DoneWithoutStartFails) { ASSERT_OK_AND_ASSIGN(se::StreamExecutor * executor, CreateExecutor()); diff --git a/third_party/xla/xla/backends/gpu/runtime/async_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/async_thunk.cc index de59a8da98c85c..3c81a6a3d5e754 100644 --- a/third_party/xla/xla/backends/gpu/runtime/async_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/async_thunk.cc @@ -91,6 +91,20 @@ absl::Status AsyncStartThunk::ExecuteOnStream(const ExecuteParams& params) { } return params.additional_compute_streams[idx]; } + if (execution_stream_id_.is_memcpy()) { + MemcpyStreamId id = execution_stream_id_.memcpy_id(); + if (id != kMemcpyD2HStreamId && id != kMemcpyH2DStreamId) { + return InvalidArgument("Invalid memcpy stream id: %d", id.value()); + } + bool is_d2h = id == kMemcpyD2HStreamId; + se::Stream* stream = + is_d2h ? params.device_to_host_stream : params.host_to_device_stream; + if (stream == nullptr) { + return Internal("%s stream is not available for async execution", + is_d2h ? "device_to_host" : "host_to_device"); + } + return stream; + } return params.collective_params->async_streams.at( execution_stream_id_.communication_id().value()); }; @@ -178,6 +192,8 @@ absl::StatusOr AsyncStartThunk::ToProto() const { if (execution_stream_id_.is_computation()) { start_proto->set_computation_stream_id( execution_stream_id_.computation_id().value()); + } else if (execution_stream_id_.is_memcpy()) { + start_proto->set_memcpy_stream_id(execution_stream_id_.memcpy_id().value()); } else { start_proto->set_communication_stream_id( execution_stream_id_.communication_id().value()); @@ -200,6 +216,14 @@ absl::StatusOr> AsyncStartThunk::FromProto( return ComputationStreamId(proto.computation_stream_id()); case AsyncStartThunkProto::kCommunicationStreamId: return CommunicationStreamId(proto.communication_stream_id()); + case AsyncStartThunkProto::kMemcpyStreamId: { + MemcpyStreamId id(proto.memcpy_stream_id()); + if (id != kMemcpyD2HStreamId && id != kMemcpyH2DStreamId) { + return InvalidArgument( + "Invalid memcpy stream id in AsyncStartThunk: %d", id.value()); + } + return id; + } default: return Internal("Unknown execution stream id type in AsyncStartThunk"); } diff --git a/third_party/xla/xla/backends/gpu/runtime/async_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/async_thunk_test.cc index 09d263926e786d..f66e7bb55a31f8 100644 --- a/third_party/xla/xla/backends/gpu/runtime/async_thunk_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/async_thunk_test.cc @@ -26,16 +26,21 @@ limitations under the License. #include "absl/status/status_macros.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" +#include "absl/types/span.h" #include "xla/backends/gpu/runtime/command_state.h" #include "xla/backends/gpu/runtime/execution_stream_id.h" #include "xla/backends/gpu/runtime/memset_thunk.h" #include "xla/backends/gpu/runtime/thunk.h" #include "xla/backends/gpu/runtime/thunk_executor.h" #include "xla/backends/gpu/runtime/thunk_id.h" +#include "xla/debug_options_flags.h" +#include "xla/executable_run_options.h" #include "xla/service/buffer_assignment.h" #include "xla/service/gpu/buffer_allocations.h" +#include "xla/service/gpu/gpu_executable.h" #include "xla/service/platform_util.h" #include "xla/service/service_executable_run_options.h" +#include "xla/service/shaped_buffer.h" #include "xla/stream_executor/device_address_allocator.h" #include "xla/stream_executor/platform.h" #include "xla/stream_executor/platform_manager.h" @@ -146,6 +151,135 @@ TEST(AsyncThunkTest, ConcurrentMemsets) { } } +// Test that AsyncStartThunk dispatches to device_to_host_stream when +// execution_stream_id is kMemcpyD2HStreamId, and to host_to_device_stream +// when it is kMemcpyH2DStreamId. We do this by running a no-op thunk on each +// memcpy stream and confirming execution completes without error. +TEST(AsyncThunkTest, MemcpyStreamDispatch) { + ASSERT_OK_AND_ASSIGN(auto executor, CreateExecutor()); + ASSERT_OK_AND_ASSIGN(auto main_stream, executor->CreateStream()); + ASSERT_OK_AND_ASSIGN(auto d2h_stream, executor->CreateStream()); + ASSERT_OK_AND_ASSIGN(auto h2d_stream, executor->CreateStream()); + + Thunk::ThunkInfo d2h_info; + d2h_info.thunk_id = ThunkId(10); + d2h_info.profile_annotation = "d2h_start"; + + Thunk::ThunkInfo h2d_info; + h2d_info.thunk_id = ThunkId(11); + h2d_info.profile_annotation = "h2d_start"; + + ThunkSequence thunks; + thunks.Emplace(d2h_info, kMemcpyD2HStreamId, + ThunkSequence{}); + auto* d2h_start = static_cast(thunks[0].get()); + thunks.Emplace(Thunk::ThunkInfo(), + d2h_start->async_execution()); + + thunks.Emplace(h2d_info, kMemcpyH2DStreamId, + ThunkSequence{}); + auto* h2d_start = static_cast(thunks[2].get()); + thunks.Emplace(Thunk::ThunkInfo(), + h2d_start->async_execution()); + + ThunkExecutor thunk_executor(std::move(thunks)); + + se::StreamExecutorAddressAllocator allocator(executor); + BufferAllocations allocations({}, 0, &allocator); + ServiceExecutableRunOptions run_options; + Thunk::ExecutionScopedState state; + + Thunk::ExecuteParams params = Thunk::ExecuteParams::Create( + run_options, allocations, main_stream.get(), main_stream.get(), + /*collective_params=*/nullptr, /*collective_cliques=*/nullptr, + /*collective_memory=*/nullptr, /*additional_compute_streams=*/{}, &state); + + params.device_to_host_stream = d2h_stream.get(); + params.host_to_device_stream = h2d_stream.get(); + + Thunk::InitializeParams init_params; + init_params.executor = executor; + init_params.execution_scoped_state = &state; + ASSERT_OK(thunk_executor.Initialize(init_params)); + ASSERT_OK(thunk_executor.ExecuteOnStream(params)); + ASSERT_OK(main_stream->BlockHostUntilDone()); +} + +TEST(AsyncThunkTest, MemcpyStreamNullReturnsError) { + ASSERT_OK_AND_ASSIGN(auto executor, CreateExecutor()); + ASSERT_OK_AND_ASSIGN(auto main_stream, executor->CreateStream()); + + AsyncStartThunk d2h_start(Thunk::ThunkInfo(), kMemcpyD2HStreamId, + ThunkSequence{}); + AsyncStartThunk h2d_start(Thunk::ThunkInfo(), kMemcpyH2DStreamId, + ThunkSequence{}); + + se::StreamExecutorAddressAllocator allocator(executor); + BufferAllocations allocations({}, 0, &allocator); + ServiceExecutableRunOptions run_options; + Thunk::ExecutionScopedState state; + + Thunk::ExecuteParams params = Thunk::ExecuteParams::Create( + run_options, allocations, main_stream.get(), main_stream.get(), + /*collective_params=*/nullptr, /*collective_cliques=*/nullptr, + /*collective_memory=*/nullptr, /*additional_compute_streams=*/{}, &state); + ASSERT_EQ(params.device_to_host_stream, nullptr); + ASSERT_EQ(params.host_to_device_stream, nullptr); + + Thunk::InitializeParams init_params; + init_params.executor = executor; + init_params.execution_scoped_state = &state; + ASSERT_OK(d2h_start.Initialize(init_params)); + ASSERT_OK(h2d_start.Initialize(init_params)); + + EXPECT_FALSE(d2h_start.ExecuteOnStream(params).ok()); + EXPECT_FALSE(h2d_start.ExecuteOnStream(params).ok()); +} + +TEST(AsyncThunkTest, GpuExecutableFallsBackToMainStreamForMemcpyStreams) { + ASSERT_OK_AND_ASSIGN(auto executor, CreateExecutor()); + ASSERT_OK_AND_ASSIGN(auto main_stream, executor->CreateStream()); + + ThunkSequence thunks; + thunks.Emplace(Thunk::ThunkInfo(), kMemcpyD2HStreamId, + ThunkSequence{}); + auto* d2h_start = static_cast(thunks[0].get()); + thunks.Emplace(Thunk::ThunkInfo(), + d2h_start->async_execution()); + thunks.Emplace(Thunk::ThunkInfo(), kMemcpyH2DStreamId, + ThunkSequence{}); + auto* h2d_start = static_cast(thunks[2].get()); + thunks.Emplace(Thunk::ThunkInfo(), + h2d_start->async_execution()); + + GpuExecutable::Params params; + params.executable = std::make_unique(std::move(thunks)); + params.debug_options = GetDebugOptionsFromFlags(); + params.buffer_assignment_proto = BufferAssignmentProto(); + params.buffer_allocations_debug_summary = "dummy"; + params.module_name = "memcpy_stream_fallback_test"; + params.device_description = executor->GetDeviceDescription(); + params.enable_debug_info_manager = false; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr executable, + GpuExecutable::Create(std::move(params))); + + se::StreamExecutorAddressAllocator allocator(executor); + ExecutableRunOptions run_options; + run_options.set_stream(main_stream.get()); + run_options.set_allocator(&allocator); + run_options.set_device_ordinal(executor->device_ordinal()); + // No stream_borrower is installed, so GpuExecutable must fall back to + // main_stream for device_to_host_stream/host_to_device_stream. + ServiceExecutableRunOptions service_run_options(run_options); + + ASSERT_OK(executable + ->ExecuteAsyncOnStream(&service_run_options, + absl::Span{}) + .status()); + ASSERT_OK(main_stream->BlockHostUntilDone()); +} + // Test that AsyncDoneThunk::Record() creates an empty command buffer node and // is a no-op on update. TEST(AsyncThunkTest, AsyncDoneRecordCommandBuffer) { @@ -165,7 +299,8 @@ TEST(AsyncThunkTest, AsyncDoneRecordCommandBuffer) { Thunk::ExecuteParams execute_params = Thunk::ExecuteParams::Create( run_options, allocations, stream.get(), stream.get(), /*collective_params=*/nullptr, /*collective_cliques=*/nullptr, - /*collective_memory=*/nullptr, /*additional_streams=*/{}, &thunk_state); + /*collective_memory=*/nullptr, /*additional_compute_streams=*/{}, + &thunk_state); ASSERT_OK_AND_ASSIGN( auto command_buffer, diff --git a/third_party/xla/xla/backends/gpu/runtime/execution_stream_id.h b/third_party/xla/xla/backends/gpu/runtime/execution_stream_id.h index e1512633b2d095..a83fbadf3f632c 100644 --- a/third_party/xla/xla/backends/gpu/runtime/execution_stream_id.h +++ b/third_party/xla/xla/backends/gpu/runtime/execution_stream_id.h @@ -64,11 +64,19 @@ TSL_LIB_GTL_DEFINE_INT_TYPE(ComputationStreamId, uint64_t); // Additional streams for launching communication operations. TSL_LIB_GTL_DEFINE_INT_TYPE(CommunicationStreamId, uint64_t); +// Dedicated streams for host↔device memory copies. There is exactly one D2H +// stream (id value 0) and one H2D stream (id value 1), matching the +// `device_to_host_stream` and `host_to_device_stream` in `ExecuteParams`. +TSL_LIB_GTL_DEFINE_INT_TYPE(MemcpyStreamId, uint64_t); +inline constexpr MemcpyStreamId kMemcpyD2HStreamId(0); +inline constexpr MemcpyStreamId kMemcpyH2DStreamId(1); + // Identifies an additional execution stream used by a thunk. This is either a -// computation stream (for async fusions/calls) or a communication stream (for -// collective operations). Default execution stream is not a computation or -// communication stream, and it's always implicitly available. ExecutionStreamId -// always identifies one of the additional streams used at run time. +// computation stream (for async fusions/calls), a communication stream (for +// collective operations), or a dedicated memcpy stream (for host↔device +// copies). Default execution stream is not a computation, communication, or +// memcpy stream, and it's always implicitly available. ExecutionStreamId always +// identifies one of the additional streams used at run time. class ExecutionStreamId { public: // Implicit conversions from strongly-typed stream ids for convenience. @@ -76,6 +84,8 @@ class ExecutionStreamId { ExecutionStreamId(ComputationStreamId id) : id_(id) {} // NOLINTNEXTLINE(google-explicit-constructor) ExecutionStreamId(CommunicationStreamId id) : id_(id) {} + // NOLINTNEXTLINE(google-explicit-constructor) + ExecutionStreamId(MemcpyStreamId id) : id_(id) {} bool is_computation() const { return std::holds_alternative(id_); @@ -85,6 +95,8 @@ class ExecutionStreamId { return std::holds_alternative(id_); } + bool is_memcpy() const { return std::holds_alternative(id_); } + ComputationStreamId computation_id() const { return std::get(id_); } @@ -93,6 +105,8 @@ class ExecutionStreamId { return std::get(id_); } + MemcpyStreamId memcpy_id() const { return std::get(id_); } + friend bool operator==(const ExecutionStreamId& a, const ExecutionStreamId& b) { return a.id_ == b.id_; @@ -109,7 +123,7 @@ class ExecutionStreamId { } private: - std::variant id_; + std::variant id_; }; template @@ -122,12 +136,20 @@ void AbslStringify(Sink& sink, CommunicationStreamId id) { absl::Format(&sink, "communication_stream[%v]", id.value()); } +template +void AbslStringify(Sink& sink, MemcpyStreamId id) { + sink.Append(id == kMemcpyD2HStreamId ? "memcpy_d2h_stream" + : "memcpy_h2d_stream"); +} + template void AbslStringify(Sink& sink, const ExecutionStreamId& id) { if (id.is_computation()) { absl::Format(&sink, "%v", id.computation_id()); - } else { + } else if (id.is_communication()) { absl::Format(&sink, "%v", id.communication_id()); + } else { + absl::Format(&sink, "%v", id.memcpy_id()); } } diff --git a/third_party/xla/xla/backends/gpu/runtime/execution_stream_id_test.cc b/third_party/xla/xla/backends/gpu/runtime/execution_stream_id_test.cc index 34c0f183f2ab06..940079c8ea0935 100644 --- a/third_party/xla/xla/backends/gpu/runtime/execution_stream_id_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/execution_stream_id_test.cc @@ -62,5 +62,55 @@ TEST(ExecutionStreamIdTest, Equality) { ExecutionStreamId(CommunicationStreamId(1))); } +// ---- MemcpyStreamId tests ------------------------------------------------ + +TEST(MemcpyStreamIdTest, Constants) { + EXPECT_EQ(kMemcpyD2HStreamId.value(), 0u); + EXPECT_EQ(kMemcpyH2DStreamId.value(), 1u); + EXPECT_NE(kMemcpyD2HStreamId, kMemcpyH2DStreamId); +} + +TEST(MemcpyStreamIdTest, AbslStringifyD2H) { + EXPECT_EQ(absl::StrCat(kMemcpyD2HStreamId), "memcpy_d2h_stream"); +} + +TEST(MemcpyStreamIdTest, AbslStringifyH2D) { + EXPECT_EQ(absl::StrCat(kMemcpyH2DStreamId), "memcpy_h2d_stream"); +} + +TEST(ExecutionStreamIdTest, MemcpyAccessors) { + ExecutionStreamId d2h(kMemcpyD2HStreamId); + EXPECT_TRUE(d2h.is_memcpy()); + EXPECT_FALSE(d2h.is_computation()); + EXPECT_FALSE(d2h.is_communication()); + EXPECT_EQ(d2h.memcpy_id(), kMemcpyD2HStreamId); + + ExecutionStreamId h2d(kMemcpyH2DStreamId); + EXPECT_TRUE(h2d.is_memcpy()); + EXPECT_FALSE(h2d.is_computation()); + EXPECT_FALSE(h2d.is_communication()); + EXPECT_EQ(h2d.memcpy_id(), kMemcpyH2DStreamId); +} + +TEST(ExecutionStreamIdTest, AbslStringifyMemcpy) { + EXPECT_EQ(absl::StrCat(ExecutionStreamId(kMemcpyD2HStreamId)), + "memcpy_d2h_stream"); + EXPECT_EQ(absl::StrCat(ExecutionStreamId(kMemcpyH2DStreamId)), + "memcpy_h2d_stream"); +} + +TEST(ExecutionStreamIdTest, EqualityWithMemcpy) { + EXPECT_EQ(ExecutionStreamId(kMemcpyD2HStreamId), + ExecutionStreamId(kMemcpyD2HStreamId)); + EXPECT_EQ(ExecutionStreamId(kMemcpyH2DStreamId), + ExecutionStreamId(kMemcpyH2DStreamId)); + EXPECT_NE(ExecutionStreamId(kMemcpyD2HStreamId), + ExecutionStreamId(kMemcpyH2DStreamId)); + EXPECT_NE(ExecutionStreamId(kMemcpyD2HStreamId), + ExecutionStreamId(ComputationStreamId(0))); + EXPECT_NE(ExecutionStreamId(kMemcpyD2HStreamId), + ExecutionStreamId(CommunicationStreamId(0))); +} + } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/runtime/memset_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/memset_thunk.cc index 402f57808bd875..e7d15b0485ad8d 100644 --- a/third_party/xla/xla/backends/gpu/runtime/memset_thunk.cc +++ b/third_party/xla/xla/backends/gpu/runtime/memset_thunk.cc @@ -54,7 +54,7 @@ absl::StatusOr MemzeroThunk::Record( VLOG(5) << "MemzeroThunk::Record"; VLOG(5) << " dest: " << dest_ << " (" << dest_data.opaque() << ")"; - if (dest_.slice.size() == 0) { + if (dest_.slice.empty()) { VLOG(5) << "Skip recording MemzeroThunk command of 0 bytes"; return nullptr; } diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk.proto b/third_party/xla/xla/backends/gpu/runtime/thunk.proto index b32f8d74310d25..00c062783c010d 100644 --- a/third_party/xla/xla/backends/gpu/runtime/thunk.proto +++ b/third_party/xla/xla/backends/gpu/runtime/thunk.proto @@ -68,6 +68,8 @@ message AsyncStartThunkProto { oneof execution_stream_id { uint64 computation_stream_id = 4; uint64 communication_stream_id = 5; + // 0 = D2H (device_to_host_stream), 1 = H2D (host_to_device_stream). + uint64 memcpy_stream_id = 6; } reserved 2; diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization_test.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization_test.cc index 6bad031e63fea1..8e2b41368773f5 100644 --- a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization_test.cc +++ b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization_test.cc @@ -1341,6 +1341,48 @@ TEST(ThunkProtoDeserializationTest, AsyncStartAndDoneThunk) { EXPECT_THAT(round_trip_done, EqualsProto(done_proto)); } +TEST(ThunkProtoDeserializationTest, AsyncStartThunkMemcpyStreamRoundTrip) { + // Verify that AsyncStartThunk with kMemcpyD2HStreamId and kMemcpyH2DStreamId + // round-trips correctly through proto serialization. + for (MemcpyStreamId stream_id : {kMemcpyD2HStreamId, kMemcpyH2DStreamId}) { + Thunk::ThunkInfo start_info; + start_info.profile_annotation = "memcpy_async_start"; + + AsyncStartThunk start_thunk(start_info, stream_id, ThunkSequence{}); + AsyncDoneThunk done_thunk(Thunk::ThunkInfo(), + start_thunk.async_execution()); + + ASSERT_OK_AND_ASSIGN(ThunkProto start_proto, start_thunk.ToProto()); + ASSERT_OK_AND_ASSIGN(ThunkProto done_proto, done_thunk.ToProto()); + + // The proto must have used the memcpy_stream_id field. + EXPECT_EQ(start_proto.async_start_thunk().execution_stream_id_case(), + AsyncStartThunkProto::kMemcpyStreamId); + EXPECT_EQ(start_proto.async_start_thunk().memcpy_stream_id(), + stream_id.value()); + + ThunkSequenceProto thunk_protos; + *thunk_protos.add_thunks() = start_proto; + *thunk_protos.add_thunks() = done_proto; + ASSERT_OK_AND_ASSIGN( + ThunkSequence sequence, + DeserializeThunkSequenceProto(thunk_protos, /*buffer_allocations=*/{}, + /*hlo_module=*/nullptr, kTestPlatformName, + se::GpuComputeCapability())); + + ASSERT_EQ(sequence.size(), 2); + auto* deserialized_start = + dynamic_cast(sequence[0].get()); + ASSERT_NE(deserialized_start, nullptr); + EXPECT_TRUE(deserialized_start->execution_stream_id().is_memcpy()); + EXPECT_EQ(deserialized_start->execution_stream_id().memcpy_id(), stream_id); + + ASSERT_OK_AND_ASSIGN(ThunkProto round_trip_start, + deserialized_start->ToProto()); + EXPECT_THAT(round_trip_start, EqualsProto(start_proto)); + } +} + TEST(ThunkProtoDeserializationTest, SendThunk) { ThunkProto proto = ParseTextProtoOrDie( R"pb( @@ -1387,13 +1429,13 @@ TEST(ThunkProtoDeserializationTest, SendThunk) { BufferAllocation(/*index=*/0, /*size=*/1024, /*color=*/0), BufferAllocation(/*index=*/1, /*size=*/1024, /*color=*/0)}; - TF_ASSERT_OK_AND_ASSIGN( + ASSERT_OK_AND_ASSIGN( std::unique_ptr thunk, DeserializeThunkProto(proto, buffer_allocations, /*hlo_module=*/nullptr, kTestPlatformName, se::GpuComputeCapability())); auto* send_thunk = dynamic_cast(thunk.get()); ASSERT_NE(send_thunk, nullptr); - TF_ASSERT_OK_AND_ASSIGN(ThunkProto round_trip_proto, send_thunk->ToProto()); + ASSERT_OK_AND_ASSIGN(ThunkProto round_trip_proto, send_thunk->ToProto()); EXPECT_THAT(round_trip_proto, EqualsProto(proto)); } @@ -1564,11 +1606,11 @@ TEST(ThunkProtoDeserializationTest, ConcurrentRegionIdPreserved) { SequentialThunk thunk(thunk_info, ThunkSequence{}); EXPECT_EQ(thunk.concurrent_region_id(), 42); - TF_ASSERT_OK_AND_ASSIGN(ThunkProto proto, thunk.ToProto()); + ASSERT_OK_AND_ASSIGN(ThunkProto proto, thunk.ToProto()); EXPECT_TRUE(proto.thunk_info().has_concurrent_region_id()); EXPECT_EQ(proto.thunk_info().concurrent_region_id(), 42); - TF_ASSERT_OK_AND_ASSIGN( + ASSERT_OK_AND_ASSIGN( std::unique_ptr deserialized, DeserializeThunkProto(proto, /*buffer_allocations=*/{}, /*hlo_module=*/nullptr, kTestPlatformName, diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD index 2013ccd9e9dc89..c78569045a3131 100644 --- a/third_party/xla/xla/backends/gpu/transforms/BUILD +++ b/third_party/xla/xla/backends/gpu/transforms/BUILD @@ -756,6 +756,45 @@ xla_cc_test( ], ) +cc_library( + name = "constant_fill_copy_rewriter", + srcs = ["constant_fill_copy_rewriter.cc"], + hdrs = ["constant_fill_copy_rewriter.h"], + deps = [ + "//xla:shape_util", + "//xla/hlo/ir:hlo", + "//xla/hlo/ir:hlo_instruction_utils", + "//xla/hlo/pass:hlo_pass", + "//xla/hlo/utils:hlo_query", + "//xla/service:decision", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/log", + "@com_google_absl//absl/status:status_macros", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:string_view", + ], +) + +xla_cc_test( + name = "constant_fill_copy_rewriter_test", + srcs = ["constant_fill_copy_rewriter_test.cc"], + deps = [ + ":constant_fill_copy_rewriter", + ":copy_fusion", + "//xla:shape_util", + "//xla/hlo/ir:hlo", + "//xla/hlo/testlib:hlo_hardware_independent_test_base", + "//xla/hlo/transforms/simplifiers:hlo_memory_scheduler", + "//xla/service:buffer_value", + "//xla/service/gpu:alias_info", + "//xla/service/gpu:gpu_device_info_for_tests", + "//xla/tsl/platform:status_matchers", + "@com_google_absl//absl/strings", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "copy_fusion", srcs = ["copy_fusion.cc"], diff --git a/third_party/xla/xla/backends/gpu/transforms/constant_fill_copy_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/constant_fill_copy_rewriter.cc new file mode 100644 index 00000000000000..f8cb54a691ec99 --- /dev/null +++ b/third_party/xla/xla/backends/gpu/transforms/constant_fill_copy_rewriter.cc @@ -0,0 +1,141 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/backends/gpu/transforms/constant_fill_copy_rewriter.h" + +#include + +#include "absl/algorithm/container.h" +#include "absl/log/log.h" +#include "absl/status/status_macros.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_instruction_utils.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/utils/hlo_query.h" +#include "xla/service/decision.h" +#include "xla/shape_util.h" + +namespace xla::gpu { +namespace { + +// Splitting small fills can cost more in launches than it saves in memory. +constexpr int64_t kMinFillBytes = 1024 * 1024; + +// Returns the instruction that ultimately consumes `user`'s value, looking +// through the same single-user bitcast and copy that the rewrite looks through. +const HloInstruction* Consumer(const HloInstruction* user) { + if (user->opcode() == HloOpcode::kBitcast && user->user_count() == 1) { + user = user->users()[0]; + } + if (user->opcode() == HloOpcode::kCopy && user->user_count() == 1) { + user = user->users()[0]; + } + return user; +} + +Decision CanReplaceCopyWithFill(const HloInstruction& copy) { + const HloInstruction* source = copy.operand(0); + if (!copy.shape().IsArray() || copy.shape() != source->shape() || + copy.HasControlDependencies() || copy.has_sharding() || copy.IsDead()) { + return Decision::Forbid( + "Copy changes layout, has ordering or sharding, or is dead"); + } + const HloInstruction* fill = source; + if (source->opcode() == HloOpcode::kBitcast) { + // CopyFusion looks through such bitcasts as well, so leaving the copy in + // place would fuse it back into the fill. + if (source->user_count() != 1 || source->HasControlDependencies() || + source->has_sharding() || + !hlo_instruction_utils::KeepsBitwidth(*source)) { + return Decision::Forbid("Bitcast is shared or changes bitwidth"); + } + fill = source->operand(0); + } + if (fill->opcode() != HloOpcode::kFusion || + fill->fusion_kind() != HloInstruction::FusionKind::kLoop || + fill->operand_count() != 0 || fill->HasControlDependencies() || + fill->has_sharding()) { + return Decision::Forbid("Not an independent loop fill"); + } + if (fill->fused_instructions_computation()->instruction_count() != 2 || + !hlo_query::IsBroadcastOfScalarConstant(*fill->fused_expression_root())) { + return Decision::Forbid("Not a broadcast of a scalar literal"); + } + if (ShapeUtil::ByteSizeOfElements(copy.shape()) < kMinFillBytes) { + return Decision::Forbid("Fill is too small to split"); + } + // When every use of the fill reaches the copy's consumer, all of those + // buffers are live there anyway. A separate fill would only add a launch and + // free the scheduler to pull it away from the consumer, so leave the copy for + // CopyFusion to fold into the fill. + if (copy.user_count() == 1 && + absl::c_all_of(fill->users(), [&](const HloInstruction* user) { + return Consumer(user) == copy.users()[0]; + })) { + return Decision::Forbid("All uses of the fill feed the same consumer"); + } + return Decision::Allow(); +} + +} // namespace + +absl::StatusOr ConstantFillCopyRewriter::RunImpl( + HloModule* module, + const absl::flat_hash_set& execution_threads) { + HloComputation* entry = module->entry_computation(); + if (!HloInstruction::IsThreadIncluded(entry->execution_thread(), + execution_threads)) { + return false; + } + bool changed = false; + // Each decision sees the replacements made so far: once the copies feeding + // other consumers have their own fills, the remaining copies of a fill that + // all feed one consumer stay with it. + for (HloInstruction* copy : entry->MakeInstructionPostOrder()) { + if (copy->opcode() != HloOpcode::kCopy) { + continue; + } + if (Decision decision = CanReplaceCopyWithFill(*copy); + decision.IsForbidden()) { + VLOG(4) << "Not rematerializing " << copy->name() << ": " + << decision.Explain(); + continue; + } + HloInstruction* source = copy->mutable_operand(0); + HloInstruction* fill = source->opcode() == HloOpcode::kBitcast + ? source->mutable_operand(0) + : source; + HloInstruction* replacement = source; + if (fill->user_count() > 1) { + // Other uses keep the original fill alive, so the copy gets a fill of its + // own: still a separate writable value, but one that no longer pins the + // original's buffer until this consumer runs. + replacement = entry->AddInstruction(fill->Clone("rematerialized")); + if (source != fill) { + replacement = entry->AddInstruction( + HloInstruction::CreateBitcast(copy->shape(), replacement)); + } + } + // Otherwise the copy is the fill's only use and the fill stands in for it + // directly. Either way the original fill keeps a user, so its fusion body + // is never orphaned. + ABSL_RETURN_IF_ERROR(entry->ReplaceInstruction(copy, replacement)); + changed = true; + } + return changed; +} + +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/constant_fill_copy_rewriter.h b/third_party/xla/xla/backends/gpu/transforms/constant_fill_copy_rewriter.h new file mode 100644 index 00000000000000..4f837b191eccaa --- /dev/null +++ b/third_party/xla/xla/backends/gpu/transforms/constant_fill_copy_rewriter.h @@ -0,0 +1,47 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef XLA_BACKENDS_GPU_TRANSFORMS_CONSTANT_FILL_COPY_REWRITER_H_ +#define XLA_BACKENDS_GPU_TRANSFORMS_CONSTANT_FILL_COPY_REWRITER_H_ + +#include "absl/container/flat_hash_set.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "xla/hlo/ir/hlo_module.h" +#include "xla/hlo/pass/hlo_pass_interface.h" + +namespace xla::gpu { + +// Replaces entry-computation copies of large scalar-constant fill fusions with +// independent fills, looking through the same single-user bitcasts as +// CopyFusion. Copies whose consumer already receives every other use of the +// fill are left for CopyFusion to fold into it. Run after CopyInsertion and +// before CopyFusion so separate writable buffers can be initialized +// independently by the scheduler. +class ConstantFillCopyRewriter : public HloModulePass { + public: + absl::string_view name() const override { + return "constant-fill-copy-rewriter"; + } + + protected: + absl::StatusOr RunImpl( + HloModule* module, + const absl::flat_hash_set& execution_threads) override; +}; + +} // namespace xla::gpu + +#endif // XLA_BACKENDS_GPU_TRANSFORMS_CONSTANT_FILL_COPY_REWRITER_H_ diff --git a/third_party/xla/xla/backends/gpu/transforms/constant_fill_copy_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/constant_fill_copy_rewriter_test.cc new file mode 100644 index 00000000000000..04d3c1cef8b63a --- /dev/null +++ b/third_party/xla/xla/backends/gpu/transforms/constant_fill_copy_rewriter_test.cc @@ -0,0 +1,361 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/backends/gpu/transforms/constant_fill_copy_rewriter.h" + +#include +#include +#include +#include + +#include +#include "absl/strings/str_replace.h" +#include "absl/strings/string_view.h" +#include "xla/backends/gpu/transforms/copy_fusion.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/ir/hlo_schedule.h" +#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" +#include "xla/hlo/transforms/simplifiers/hlo_memory_scheduler.h" +#include "xla/service/buffer_value.h" +#include "xla/service/gpu/alias_info.h" +#include "xla/service/gpu/gpu_device_info_for_tests.h" +#include "xla/shape_util.h" +#include "xla/tsl/platform/status_matchers.h" + +namespace xla::gpu { +namespace { + +constexpr char kHlo[] = R"( +HloModule test + +fill_body { + value = f32[] constant(7) + ROOT broadcast = f32[512,512]{1,0} broadcast(value), dimensions={} +} + +ENTRY main { + predecessor = f32[] constant(1) + fill = f32[512,512]{1,0} fusion(), kind=kLoop, calls=fill_body + ROOT copy = f32[512,512]{1,0} copy(fill) +} +)"; + +constexpr absl::string_view kRootCopy = + "ROOT copy = f32[512,512]{1,0} copy(fill)"; + +using ConstantFillCopyRewriterTest = HloHardwareIndependentTestBase; + +// Checks that `fill` is a fresh zero-operand fill equivalent to `original`. +void ExpectRematerializedFill(const HloInstruction* fill, + const HloInstruction* original) { + EXPECT_NE(fill, original); + ASSERT_EQ(fill->opcode(), HloOpcode::kFusion); + EXPECT_EQ(fill->operand_count(), 0); + EXPECT_EQ(fill->shape(), original->shape()); + EXPECT_EQ(*fill->fused_instructions_computation(), + *original->fused_instructions_computation()); +} + +TEST_F(ConstantFillCopyRewriterTest, ReplacesSoleCopyWithTheFillItself) { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo)); + HloComputation* entry = module->entry_computation(); + const HloInstruction* fill = entry->GetInstructionWithName("fill"); + ConstantFillCopyRewriter pass; + ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, module.get())); + EXPECT_TRUE(changed); + // The copy was the fill's only use, so the fill stands in for it directly and + // nothing is cloned. + EXPECT_EQ(entry->root_instruction(), fill); + EXPECT_EQ(entry->GetInstructionWithName("copy"), nullptr); + EXPECT_EQ(module->computation_count(), 2); + ASSERT_OK_AND_ASSIGN(changed, RunHloPass(&pass, module.get())); + EXPECT_FALSE(changed); +} + +TEST_F(ConstantFillCopyRewriterTest, RematerializesFillForItsOwnConsumer) { + std::string hlo = absl::StrReplaceAll(kHlo, {{kRootCopy, R"( + copy = f32[512,512]{1,0} copy(fill) + early = f32[512,512]{1,0} negate(copy) + ROOT late = (f32[512,512]{1,0}, f32[512,512]{1,0}) tuple(early, fill))"}}); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + HloComputation* entry = module->entry_computation(); + const HloInstruction* fill = entry->GetInstructionWithName("fill"); + ConstantFillCopyRewriter pass; + ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, module.get())); + EXPECT_TRUE(changed); + ExpectRematerializedFill(entry->GetInstructionWithName("early")->operand(0), + fill); + // The original fill keeps its late consumer and its fusion body. + EXPECT_EQ(entry->root_instruction()->operand(1), fill); + EXPECT_EQ(entry->GetInstructionWithName("copy"), nullptr); + EXPECT_EQ(module->computation_count(), 3); + + // The following pipeline pass must leave the independent fills intact. + const auto device = TestGpuDeviceInfo::RTXA6000DeviceInfo(); + CopyFusion copy_fusion(device); + ASSERT_OK_AND_ASSIGN(changed, RunHloPass(©_fusion, module.get())); + EXPECT_FALSE(changed); +} + +TEST_F(ConstantFillCopyRewriterTest, + KeepsCopyWhoseConsumerReceivesEveryOtherUse) { + // copy0 feeds early on its own, while copy1 and the fill both feed late. + std::string hlo = absl::StrReplaceAll(kHlo, {{kRootCopy, R"( + copy0 = f32[512,512]{1,0} copy(fill) + early = f32[512,512]{1,0} negate(copy0) + copy1 = f32[512,512]{1,0} copy(fill) + ROOT late = (f32[512,512]{1,0}, f32[512,512]{1,0}, f32[512,512]{1,0}) + tuple(early, fill, copy1))"}}); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + HloComputation* entry = module->entry_computation(); + const HloInstruction* fill = entry->GetInstructionWithName("fill"); + ConstantFillCopyRewriter pass; + ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, module.get())); + EXPECT_TRUE(changed); + ExpectRematerializedFill(entry->GetInstructionWithName("early")->operand(0), + fill); + // Once copy0 has its own fill, every remaining use of the fill feeds late, + // so copy1 stays a copy for CopyFusion to fold into the fill. + const HloInstruction* late = entry->root_instruction(); + EXPECT_EQ(late->operand(1), fill); + EXPECT_EQ(late->operand(2), entry->GetInstructionWithName("copy1")); + EXPECT_EQ(late->operand(2)->operand(0), fill); + + const auto device = TestGpuDeviceInfo::RTXA6000DeviceInfo(); + CopyFusion copy_fusion(device); + ASSERT_OK_AND_ASSIGN(changed, RunHloPass(©_fusion, module.get())); + EXPECT_TRUE(changed); + ASSERT_EQ(late->operand(1)->opcode(), HloOpcode::kGetTupleElement); + ASSERT_EQ(late->operand(2)->opcode(), HloOpcode::kGetTupleElement); + EXPECT_EQ(late->operand(1)->operand(0), late->operand(2)->operand(0)); + EXPECT_TRUE(late->operand(1)->operand(0)->IsMultiOutputFusion()); +} + +TEST_F(ConstantFillCopyRewriterTest, LooksThroughBitcastLikeCopyFusion) { + std::string hlo = absl::StrReplaceAll(kHlo, {{kRootCopy, R"( + bitcast = f32[262144]{0} bitcast(fill) + copy = f32[262144]{0} copy(bitcast) + early = f32[262144]{0} negate(copy) + ROOT late = (f32[262144]{0}, f32[512,512]{1,0}) tuple(early, fill))"}}); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + HloComputation* entry = module->entry_computation(); + const HloInstruction* fill = entry->GetInstructionWithName("fill"); + const HloInstruction* bitcast = entry->GetInstructionWithName("bitcast"); + ConstantFillCopyRewriter pass; + ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, module.get())); + EXPECT_TRUE(changed); + // The copy becomes a bitcast of a fresh fill and the old bitcast dies with + // it, leaving predecessor, fill, the new fill and bitcast, early, and late. + const HloInstruction* early = entry->GetInstructionWithName("early"); + const HloInstruction* new_bitcast = early->operand(0); + EXPECT_NE(new_bitcast, bitcast); + ASSERT_EQ(new_bitcast->opcode(), HloOpcode::kBitcast); + EXPECT_EQ(new_bitcast->shape(), early->shape()); + ExpectRematerializedFill(new_bitcast->operand(0), fill); + EXPECT_EQ(entry->root_instruction()->operand(1), fill); + EXPECT_EQ(entry->instruction_count(), 6); + EXPECT_EQ(module->computation_count(), 3); + + const auto device = TestGpuDeviceInfo::RTXA6000DeviceInfo(); + CopyFusion copy_fusion(device); + ASSERT_OK_AND_ASSIGN(changed, RunHloPass(©_fusion, module.get())); + EXPECT_FALSE(changed); +} + +TEST_F(ConstantFillCopyRewriterTest, ReplacesSoleCopyOfBitcastWithTheBitcast) { + std::string hlo = absl::StrReplaceAll(kHlo, {{kRootCopy, R"( + bitcast = f32[262144]{0} bitcast(fill) + ROOT copy = f32[262144]{0} copy(bitcast))"}}); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + HloComputation* entry = module->entry_computation(); + const HloInstruction* bitcast = entry->GetInstructionWithName("bitcast"); + ConstantFillCopyRewriter pass; + ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, module.get())); + EXPECT_TRUE(changed); + EXPECT_EQ(entry->root_instruction(), bitcast); + EXPECT_EQ(bitcast->operand(0), entry->GetInstructionWithName("fill")); + EXPECT_EQ(module->computation_count(), 2); +} + +TEST_F(ConstantFillCopyRewriterTest, RespectsExecutionThreads) { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo)); + ASSERT_OK_AND_ASSIGN(bool changed, + ConstantFillCopyRewriter().Run(module.get(), {"other"})); + EXPECT_FALSE(changed); +} + +TEST_F(ConstantFillCopyRewriterTest, ReducesScheduledPeakMemory) { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"( +HloModule test +sum { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) +} +fill_body { + zero = f32[] constant(0) + ROOT broadcast = f32[512,512]{1,0} broadcast(zero), dimensions={} +} +ENTRY main { + update = f32[1,512]{1,0} parameter(0) + zero = s32[] constant(0) + one = s32[] constant(1) + initial_sum = f32[] constant(0) + fill = f32[512,512]{1,0} fusion(), kind=kLoop, calls=fill_body + copy = f32[512,512]{1,0} copy(fill) + early = f32[512,512]{1,0} dynamic-update-slice(copy, update, zero, zero) + early_sum = f32[] reduce(early, initial_sum), dimensions={0,1}, to_apply=sum + middle = f32[1024,1024]{1,0} broadcast(early_sum), dimensions={} + middle_sum = f32[] reduce(middle, initial_sum), dimensions={0,1}, to_apply=sum + late_update = f32[1,512]{1,0} broadcast(middle_sum), dimensions={} + ROOT late = f32[512,512]{1,0} + dynamic-update-slice(fill, late_update, one, zero) +})")); + auto baseline = module->Clone(); + const auto device = TestGpuDeviceInfo::RTXA6000DeviceInfo(); + GpuAliasInfo alias_info(device); + CopyFusion copy_fusion(device); + ASSERT_OK_AND_ASSIGN(bool baseline_changed, + RunHloPass(©_fusion, baseline.get())); + EXPECT_TRUE(baseline_changed); + auto size = [](const BufferValue& buffer) { + return ShapeUtil::ByteSizeOf(buffer.shape(), /*pointer_size=*/8); + }; + int64_t baseline_peak; + ASSERT_OK_AND_ASSIGN( + auto baseline_schedule, + ScheduleModule(baseline.get(), &alias_info, size, {}, &baseline_peak)); + ASSERT_OK(baseline_schedule.Verify()); + + ConstantFillCopyRewriter pass; + ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, module.get())); + EXPECT_TRUE(changed); + ASSERT_OK_AND_ASSIGN(changed, RunHloPass(©_fusion, module.get())); + EXPECT_FALSE(changed); + int64_t peak; + ASSERT_OK_AND_ASSIGN(auto schedule, ScheduleModule(module.get(), &alias_info, + size, {}, &peak)); + ASSERT_OK(schedule.Verify()); + // The original 1 MiB fill no longer needs to span the 4 MiB middle buffer. + EXPECT_LT(peak, baseline_peak); +} + +TEST_F(ConstantFillCopyRewriterTest, LeavesCopiesInsideWhileBodyAlone) { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"( +HloModule test +fill_body { + zero = f32[] constant(0) + ROOT broadcast = f32[512,512]{1,0} broadcast(zero), dimensions={} +} +condition { + p = f32[512,512]{1,0} parameter(0) + ROOT stop = pred[] constant(false) +} +body { + p = f32[512,512]{1,0} parameter(0) + fill = f32[512,512]{1,0} fusion(), kind=kLoop, calls=fill_body + ROOT copy = f32[512,512]{1,0} copy(fill) +} +ENTRY main { + p = f32[512,512]{1,0} parameter(0) + ROOT loop = f32[512,512]{1,0} while(p), condition=condition, body=body +})")); + ConstantFillCopyRewriter pass; + ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, module.get())); + EXPECT_FALSE(changed); +} + +struct RejectedCopy { + const char* name; + std::vector> replacements; +}; + +class RejectedConstantFillCopyTest + : public ConstantFillCopyRewriterTest, + public ::testing::WithParamInterface {}; + +TEST_P(RejectedConstantFillCopyTest, LeavesCopyUnchanged) { + const RejectedCopy& test = GetParam(); + std::string hlo = absl::StrReplaceAll(kHlo, test.replacements); + ASSERT_NE(hlo, kHlo); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo)); + std::string before = module->ToString(); + ConstantFillCopyRewriter pass; + ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass, module.get())); + EXPECT_FALSE(changed); + EXPECT_EQ(module->ToString(), before); +} + +INSTANTIATE_TEST_SUITE_P( + Safeguards, RejectedConstantFillCopyTest, + ::testing::Values( + RejectedCopy{"SmallFill", {{"512", "256"}}}, + RejectedCopy{"LayoutChange", + {{"ROOT copy = f32[512,512]{1,0}", + "ROOT copy = f32[512,512]{0,1}"}}}, + RejectedCopy{ + "ParameterFill", + {{"value = f32[] constant(7)", "value = f32[] parameter(0)"}, + {"fusion()", "fusion(predecessor)"}}}, + RejectedCopy{"ComputedFill", + {{"ROOT broadcast", + "negated = f32[] negate(value)\n" + " ROOT broadcast"}, + {"broadcast(value)", "broadcast(negated)"}}}, + RejectedCopy{"NonscalarLiteral", + {{"f32[] constant(7)", "f32[1] constant({7})"}, + {"512,512", "1,262144"}, + {"dimensions={}", "dimensions={0}"}}}, + RejectedCopy{ + "CopyControlDependency", + {{"copy(fill)", "copy(fill), control-predecessors={predecessor}"}}}, + RejectedCopy{"FillControlDependency", + {{"calls=fill_body", + "calls=fill_body, control-predecessors={predecessor}"}}}, + RejectedCopy{"CopySharding", + {{"copy(fill)", "copy(fill), sharding={replicated}"}}}, + RejectedCopy{ + "FillSharding", + {{"calls=fill_body", "calls=fill_body, sharding={replicated}"}}}, + RejectedCopy{"DeadCopy", {{kRootCopy, R"( + copy = f32[512,512]{1,0} copy(fill) + ROOT other = f32[512,512]{1,0} negate(fill))"}}}, + // Every use of the fill feeds one consumer, so all of those buffers + // are live there regardless of how they are produced. + RejectedCopy{"SharedConsumer", {{kRootCopy, R"( + copy0 = f32[512,512]{1,0} copy(fill) + copy1 = f32[512,512]{1,0} copy(fill) + ROOT result = (f32[512,512]{1,0}, f32[512,512]{1,0}, f32[512,512]{1,0}) + tuple(fill, copy0, copy1))"}}}, + RejectedCopy{"SharedConsumerOfCopiesOnly", {{kRootCopy, R"( + copy0 = f32[512,512]{1,0} copy(fill) + copy1 = f32[512,512]{1,0} copy(fill) + ROOT result = (f32[512,512]{1,0}, f32[512,512]{1,0}) tuple(copy0, copy1))"}}}, + RejectedCopy{"SharedBitcast", {{kRootCopy, R"( + bitcast = f32[512,512]{1,0} bitcast(fill) + copy = f32[512,512]{1,0} copy(bitcast) + other = f32[512,512]{1,0} negate(bitcast) + ROOT result = (f32[512,512]{1,0}, f32[512,512]{1,0}) tuple(copy, other))"}}}, + RejectedCopy{"BitcastChangesBitwidth", {{kRootCopy, R"( + bitcast = f16[512,1024]{1,0} bitcast(fill) + ROOT copy = f16[512,1024]{1,0} copy(bitcast))"}}}), + [](const ::testing::TestParamInfo& info) { + return info.param.name; + }); + +} // namespace +} // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc index 9d43cde32149e1..909690650944b6 100644 --- a/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc +++ b/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc @@ -855,6 +855,13 @@ class PriorityFusionQueue { *consumer, producer); } + return CanFuseWithNativeEmitter(producer, consumer); + } + + // Checks whether `producer` can be fused into `consumer` using native + // emitters. + FusionDecision CanFuseWithNativeEmitter(HloInstruction* producer, + HloInstruction* consumer) { if (IsFusibleBitcast(*consumer)) { return FusionDecision::Forbid( "not fusing into a single bitcast as consumer"); @@ -1018,6 +1025,21 @@ class PriorityFusionQueue { return FusionDecision::Allow(); } + // Records that `producer` cannot be fused into `user`. + void RecordUnfusableUser(HloInstruction* producer, HloInstruction* user, + const FusionDecision& fusion_decision) { + VLOG(10) << "Cannot fuse " << producer->name() << " with " << user->name() + << ", because: " << fusion_decision.Explain(); + if (dump_fusion_visualization_) { + RegisterFusionState( + *computation_, + absl::StrCat("Cannot fuse producer |", producer->name(), + "| with consumer |", user->name(), + "|: ", fusion_decision.Explain()), + *user, producer); + } + } + FusionDecision CanFuseWithAllNonBitcastUsers(HloInstruction* producer) { if (producer->users().empty()) { return FusionDecision::Forbid("No users to fuse"); @@ -1031,16 +1053,7 @@ class PriorityFusionQueue { has_non_bitcast_user = true; if (auto fusion_decision = CanFuseCached(producer, user); !fusion_decision) { - VLOG(10) << "Cannot fuse " << producer->name() << " with " - << user->name() << ", because: " << fusion_decision.Explain(); - if (dump_fusion_visualization_) { - RegisterFusionState( - *computation_, - absl::StrCat("Cannot fuse producer |", producer->name(), - "| with consumer |", user->name(), - "|: ", fusion_decision.Explain()), - *user, producer); - } + RecordUnfusableUser(producer, user, fusion_decision); return fusion_decision; } } diff --git a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py index 44a34f79ad6d6b..a7066167f4f1e6 100644 --- a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py +++ b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/dense_matmul_lib.py @@ -241,10 +241,23 @@ def _target_fn(lhs, rhs): rhs = memory_utils.with_large_2nd_minor_layout(rhs) return pallas_func(lhs, rhs) + compiler_args: dict[str, Any] = { + "xla_detailed_logging": True, + "xla_tpu_control_large_2nd_minor_layout_for_x16": True, + } + if ( + cfg.lhs_mem == pltpu.VMEM + and cfg.rhs_mem == pltpu.VMEM + and cfg.out_mem == pltpu.VMEM + ): + # For VMEM matmuls, if scoped VMEM is too high then we may hit OOM since we + # need space for the input operands outside of scoped VMEM. So set it to a + # relatively low value. + compiler_args["xla_tpu_scoped_vmem_limit_kib"] = 8 * 1024 return _target_fn.lower( jax.ShapeDtypeStruct((m, k), lhs_dtype), jax.ShapeDtypeStruct((k, n), rhs_dtype), - ).compile({"xla_detailed_logging": True}) + ).compile(compiler_args) def dense_matmul_reference(cfg: DenseMatmulConfig) -> Callable[..., Any]: diff --git a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py index 9bd532696b8e69..d8ce0734a4a17f 100644 --- a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py +++ b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/subchannel_matmul_lib.py @@ -322,7 +322,29 @@ def _target_fn(lhs, rhs, lhs_scales, rhs_scales): rhs = memory_utils.with_large_2nd_minor_layout(rhs) return pallas_func(lhs, rhs, lhs_scales, rhs_scales) - return _target_fn + compiler_args: dict[str, Any] = { + "xla_detailed_logging": True, + "xla_tpu_control_large_2nd_minor_layout_for_x16": True, + } + if ( + cfg.lhs_mem == pltpu.VMEM + and cfg.rhs_mem == pltpu.VMEM + and cfg.out_mem == pltpu.VMEM + ): + # For VMEM matmuls, if scoped VMEM is too high then we may hit OOM since we + # need space for the input operands outside of scoped VMEM. So set it to a + # relatively low value. + compiler_args["xla_tpu_scoped_vmem_limit_kib"] = 8 * 1024 + lower_args = [ + jax.ShapeDtypeStruct((m, k), cfg.lhs_dtype), + jax.ShapeDtypeStruct((k, n), cfg.rhs_quantized_dtype), + None, + jax.ShapeDtypeStruct((subblocks_per_tile, n), cfg.rhs_dtype), + ] + if cfg.pre_quantize_lhs: + lower_args[0] = jax.ShapeDtypeStruct((m, k), cfg.lhs_quantized_dtype) + lower_args[2] = jax.ShapeDtypeStruct((m, subblocks_per_tile), cfg.lhs_dtype) + return _target_fn.lower(*lower_args).compile(compiler_args) def subchannel_matmul_jax( diff --git a/third_party/xla/xla/codegen/emitters/ir/BUILD b/third_party/xla/xla/codegen/emitters/ir/BUILD index 32bde228b16d2f..523ae0c8cf4911 100644 --- a/third_party/xla/xla/codegen/emitters/ir/BUILD +++ b/third_party/xla/xla/codegen/emitters/ir/BUILD @@ -117,6 +117,7 @@ cc_library( "@llvm-project//mlir:ArithDialect", "@llvm-project//mlir:BytecodeOpInterface", "@llvm-project//mlir:FuncDialect", + "@llvm-project//mlir:FunctionInterfaces", "@llvm-project//mlir:IR", "@llvm-project//mlir:InferTypeOpInterface", "@llvm-project//mlir:InliningUtils", diff --git a/third_party/xla/xla/codegen/emitters/ir/xla_ops.cc b/third_party/xla/xla/codegen/emitters/ir/xla_ops.cc index 96ea4544edb8bb..2def844cfcf441 100644 --- a/third_party/xla/xla/codegen/emitters/ir/xla_ops.cc +++ b/third_party/xla/xla/codegen/emitters/ir/xla_ops.cc @@ -51,6 +51,7 @@ limitations under the License. #include "mlir/IR/Types.h" #include "mlir/IR/Value.h" #include "mlir/IR/ValueRange.h" +#include "mlir/Interfaces/FunctionInterfaces.h" #include "mlir/Support/LLVM.h" #include "mlir/Support/LogicalResult.h" #include "xla/codegen/emitters/ir/xla_dialect.cc.inc" @@ -119,7 +120,7 @@ std::optional GetRange(mlir::Value value) { } auto parent = bbarg.getParentBlock()->getParentOp(); - if (auto func_op = mlir::dyn_cast(parent)) { + if (auto func_op = mlir::dyn_cast(parent)) { return attr_to_range(func_op.getArgAttr(bbarg.getArgNumber(), "xla.range")); } return GetIVRange(value); diff --git a/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc b/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc index 5b6ce3e1fd4d41..0eec56a66c6ecb 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc @@ -79,12 +79,6 @@ struct TargetTypes { bool needs_upcast = false; }; -std::string GetCpuFeaturesStr(mlir::ModuleOp module_op) { - mlir::StringAttr features = - module_op->template getAttrOfType("mhlo.cpu_features"); - return !features ? "" : features.getValue().str(); -} - bool NeedsUpcast(mlir::Type type) { mlir::Type elem_type = mlir::getElementTypeOrSelf(type); return elem_type.isBF16() || elem_type.isF16(); @@ -290,56 +284,65 @@ Value CallIntrinsicFunc(mlir::ImplicitLocOpBuilder& b, Op op, } template -mlir::LogicalResult LowerIntrinsicPattern(Op op, - mlir::PatternRewriter& rewriter) { - auto vec_type = mlir::dyn_cast(op.getType()); - if (vec_type && vec_type.getRank() != 1) { - // These will later be converted to loops of 1D vectors but will then miss - // the XLA intrinsic lowering. - op->emitWarning() << "Missed XLA intrinsic lowering as vector rank != 1."; - return rewriter.notifyMatchFailure(op, "Vector rank is not 1."); - } +class LowerIntrinsicPattern : public mlir::OpRewritePattern { + public: + LowerIntrinsicPattern(mlir::MLIRContext* context, + absl::string_view cpu_features) + : mlir::OpRewritePattern(context), cpu_features_(cpu_features) {} + + mlir::LogicalResult matchAndRewrite( + Op op, mlir::PatternRewriter& rewriter) const override { + auto vec_type = mlir::dyn_cast(op.getType()); + if (vec_type && vec_type.getRank() != 1) { + // These will later be converted to loops of 1D vectors but will then + // miss the XLA intrinsic lowering. + op->emitWarning() << "Missed XLA intrinsic lowering as vector rank != 1."; + return rewriter.notifyMatchFailure(op, "Vector rank is not 1."); + } - mlir::ImplicitLocOpBuilder b(op.getLoc(), rewriter); - TargetTypes target_types = GetTargetTypes(op); - auto module_op = op->template getParentOfType(); + mlir::ImplicitLocOpBuilder b(op.getLoc(), rewriter); + TargetTypes target_types = GetTargetTypes(op); - std::string features_str = GetCpuFeaturesStr(module_op); - bool vector_supported = - Intrinsic::IsSupported(features_str, target_types.vector_types); - std::optional subvector_size; - if (vec_type && !vector_supported) { - subvector_size = FindSupportedSubvectorSize( - op, features_str, vec_type.getNumElements()); - } - bool scalar_supported = - Intrinsic::IsSupported(features_str, target_types.scalar_types); - if (!vector_supported && !subvector_size.has_value() && !scalar_supported) { - return rewriter.notifyMatchFailure(op, "unsupported type"); + bool vector_supported = + Intrinsic::IsSupported(cpu_features_, target_types.vector_types); + std::optional subvector_size; + if (vec_type && !vector_supported) { + subvector_size = FindSupportedSubvectorSize( + op, cpu_features_, vec_type.getNumElements()); + } + bool scalar_supported = + Intrinsic::IsSupported(cpu_features_, target_types.scalar_types); + if (!vector_supported && !subvector_size.has_value() && !scalar_supported) { + return rewriter.notifyMatchFailure(op, "unsupported type"); + } + Value result = CallIntrinsicFunc( + b, op, target_types, vec_type, vector_supported, subvector_size); + rewriter.replaceOp(op, result); + return mlir::success(); } - Value result = CallIntrinsicFunc(b, op, target_types, vec_type, - vector_supported, subvector_size); - rewriter.replaceOp(op, result); - return mlir::success(); -} + + private: + std::string cpu_features_; +}; class LowerXlaIntrinsicLibPass : public impl::LowerXlaIntrinsicLibPassBase { public: - LowerXlaIntrinsicLibPass() - : impl::LowerXlaIntrinsicLibPassBase() {} + using LowerXlaIntrinsicLibPassBase::LowerXlaIntrinsicLibPassBase; void runOnOperation() override { mlir::MLIRContext* context = &getContext(); mlir::ModuleOp module_op = getOperation(); mlir::RewritePatternSet patterns(context); - patterns.add(LowerIntrinsicPattern); - patterns.add(LowerIntrinsicPattern); - patterns.add(LowerIntrinsicPattern); - patterns.add(LowerIntrinsicPattern); - patterns.add(LowerIntrinsicPattern); - patterns.add(LowerIntrinsicPattern); - patterns.add(LowerIntrinsicPattern); + absl::string_view cpu_features = cpu_features_.getValue(); + patterns.add, + LowerIntrinsicPattern, + LowerIntrinsicPattern, + LowerIntrinsicPattern, + LowerIntrinsicPattern, + LowerIntrinsicPattern, + LowerIntrinsicPattern>(context, + cpu_features); if (mlir::failed( mlir::applyPatternsGreedily(module_op, std::move(patterns)))) { signalPassFailure(); diff --git a/third_party/xla/xla/codegen/emitters/transforms/passes.td b/third_party/xla/xla/codegen/emitters/transforms/passes.td index f8f05afdf8ef66..9743cdc91eeb79 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/passes.td +++ b/third_party/xla/xla/codegen/emitters/transforms/passes.td @@ -29,6 +29,13 @@ def LowerXlaIntrinsicLibPass : Pass<"xla-lower-xla-intrinsic-lib", "mlir::Module "mlir::func::FuncDialect", "mlir::vector::VectorDialect", ]; + + let options = [ + Option<"cpu_features_", "cpu_features", "std::string", /*default=*/"", + "LLVM target feature string (e.g. '+avx,+avx2') used to select " + "feature-gated intrinsic variants. Must match the features of the " + "target machine that compiles the module.">, + ]; } def ConvertPureCallOpsPass diff --git a/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir b/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir index 1f697977a957c0..7eedc95ced80b2 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir +++ b/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir @@ -13,6 +13,9 @@ // limitations under the License. // ============================================================================== // RUN: emitters_opt %s -split-input-file -xla-lower-xla-intrinsic-lib | FileCheck %s +// RUN: emitters_opt %s -split-input-file \ +// RUN: -xla-lower-xla-intrinsic-lib="cpu_features=+avx,+avx2,+avx512f" \ +// RUN: | FileCheck %s --check-prefix=AVX512 module { func.func @exp_f64(%arg0: f64) -> f64 { @@ -307,13 +310,61 @@ module { // CHECK: %[[S0:.*]] = vector.extract_strided_slice %arg0 offsets = [0], sizes = [16], strides = [1] : vector<32xbf16> to vector<16xbf16> // CHECK: %[[EXT0:.*]] = arith.extf %[[S0]] : vector<16xbf16> to vector<16xf32> // CHECK: %[[C0:.*]] = call @xla.tanh.v16f32(%[[EXT0]]) -// CHECK: %[[TR0:.*]] = call @xla.fptrunc.v8f32.to.v8bf16 +// CHECK: %[[TR0:.*]] = call @xla.fptrunc.v16f32.to.v16bf16(%[[C0]]) // CHECK: %[[S1:.*]] = vector.extract_strided_slice %arg0 offsets = [16], sizes = [16], strides = [1] : vector<32xbf16> to vector<16xbf16> // CHECK: %[[EXT1:.*]] = arith.extf %[[S1]] : vector<16xbf16> to vector<16xf32> // CHECK: %[[C1:.*]] = call @xla.tanh.v16f32(%[[EXT1]]) -// CHECK: %[[TR1:.*]] = call @xla.fptrunc.v8f32.to.v8bf16 +// CHECK: %[[TR1:.*]] = call @xla.fptrunc.v16f32.to.v16bf16(%[[C1]]) // CHECK: return +// ----- + +func.func @log1p_f32_vector_16(%arg0: vector<16xf32>) -> vector<16xf32> { + %ret = math.log1p %arg0 : vector<16xf32> + return %ret : vector<16xf32> +} +// CHECK-LABEL: @log1p_f32_vector_16 +// CHECK-NOT: vector.extract_strided_slice +// CHECK: %[[CALL:.*]] = call @xla.log1p.v16f32(%arg0) : (vector<16xf32>) -> vector<16xf32> +// CHECK: return %[[CALL]] + +// ----- + +func.func @erf_f32_vector_16(%arg0: vector<16xf32>) -> vector<16xf32> { + %ret = math.erf %arg0 : vector<16xf32> + return %ret : vector<16xf32> +} +// CHECK-LABEL: @erf_f32_vector_16 +// CHECK-NOT: vector.extract_strided_slice +// CHECK: %[[CALL:.*]] = call @xla.erf.v16f32(%arg0) : (vector<16xf32>) -> vector<16xf32> +// CHECK: return %[[CALL]] + +// ----- + +func.func @trunc_vector_16(%arg0: vector<16xf32>) -> vector<16xbf16> { + %ret = arith.truncf %arg0 : vector<16xf32> to vector<16xbf16> + return %ret : vector<16xbf16> +} +// CHECK-LABEL: @trunc_vector_16 +// CHECK-NOT: vector.extract_strided_slice +// CHECK: %[[CALL:.*]] = call @xla.fptrunc.v16f32.to.v16bf16(%arg0) : (vector<16xf32>) -> vector<16xbf16> +// CHECK: return %[[CALL]] + +// ----- + +// rsqrt vector variants are gated on CPU features. +func.func @rsqrt_f32_vector_16(%arg0: vector<16xf32>) -> vector<16xf32> { + %ret = math.rsqrt %arg0 : vector<16xf32> + return %ret : vector<16xf32> +} +// CHECK-LABEL: @rsqrt_f32_vector_16 +// CHECK-NOT: @xla.rsqrt.v +// CHECK-COUNT-16: call @xla.rsqrt.f32 + +// AVX512-LABEL: @rsqrt_f32_vector_16 +// AVX512: %[[CALL:.*]] = call @xla.rsqrt.v16f32(%arg0) : (vector<16xf32>) -> vector<16xf32> +// AVX512: return %[[CALL]] + diff --git a/third_party/xla/xla/codegen/intrinsic/erf.h b/third_party/xla/xla/codegen/intrinsic/erf.h index b627e0f697680a..301ae84ebcee0a 100644 --- a/third_party/xla/xla/codegen/intrinsic/erf.h +++ b/third_party/xla/xla/codegen/intrinsic/erf.h @@ -31,11 +31,8 @@ class Erf : public Intrinsic { static constexpr absl::string_view kName = "erf"; static std::vector> SupportedVectorTypes( absl::string_view features = "") { - return {{Type::S(F32)}, - {Type::S(F64)}, - {Type::V(F32, 2)}, - {Type::V(F32, 4)}, - {Type::V(F32, 8)}}; + return {{Type::S(F32)}, {Type::S(F64)}, {Type::V(F32, 2)}, + {Type::V(F32, 4)}, {Type::V(F32, 8)}, {Type::V(F32, 16)}}; } static absl::StatusOr CreateDefinition(llvm::Module* module, diff --git a/third_party/xla/xla/codegen/intrinsic/erf_test.cc b/third_party/xla/xla/codegen/intrinsic/erf_test.cc index 5903a633bbf94d..c3dc6544cf3dc7 100644 --- a/third_party/xla/xla/codegen/intrinsic/erf_test.cc +++ b/third_party/xla/xla/codegen/intrinsic/erf_test.cc @@ -15,7 +15,9 @@ limitations under the License. #include "xla/codegen/intrinsic/erf.h" +#include #include +#include #include #include #include @@ -86,6 +88,28 @@ TEST(ErfTest, F32) { } } +TEST(ErfTest, F32Vector16) { + constexpr size_t kN = 16; + Type type = Type::V(F32, kN); + JitRunner runner = CreateJitRunner(type); + auto fn = runner.GetVectorizedFn(Erf::Name(type)); + + std::vector test_values = GetTestValues(); + std::array vals; + for (size_t i = 0; i < kN; ++i) { + vals[i] = test_values[i % test_values.size()]; + } + std::array results = fn(vals); + for (size_t i = 0; i < kN; ++i) { + float expected = std::erf(vals[i]); + if (std::isnan(expected)) { + EXPECT_TRUE(std::isnan(results[i])); + } else { + EXPECT_THAT(results[i], NearUlps(expected, 1)); + } + } +} + TEST(ErfTest, F64) { Type type = Type::S(F64); JitRunner runner = CreateJitRunner(type); diff --git a/third_party/xla/xla/codegen/intrinsic/fptrunc.h b/third_party/xla/xla/codegen/intrinsic/fptrunc.h index c6753c5f350b24..d02d83b7cecd18 100644 --- a/third_party/xla/xla/codegen/intrinsic/fptrunc.h +++ b/third_party/xla/xla/codegen/intrinsic/fptrunc.h @@ -42,6 +42,7 @@ class FpTrunc : public Intrinsic { {Type::V(xla::F32, 2), Type::V(xla::BF16, 2)}, {Type::V(xla::F32, 4), Type::V(xla::BF16, 4)}, {Type::V(xla::F32, 8), Type::V(xla::BF16, 8)}, + {Type::V(xla::F32, 16), Type::V(xla::BF16, 16)}, {Type::S(F8E5M2), Type::S(F16)}, {Type::V(F8E5M2, 2), Type::V(F16, 2)}, {Type::V(F8E5M2, 4), Type::V(F16, 4)}, diff --git a/third_party/xla/xla/codegen/intrinsic/fptrunc_test.cc b/third_party/xla/xla/codegen/intrinsic/fptrunc_test.cc index 03f74a2ad6cded..888f0ac06e29d3 100644 --- a/third_party/xla/xla/codegen/intrinsic/fptrunc_test.cc +++ b/third_party/xla/xla/codegen/intrinsic/fptrunc_test.cc @@ -16,6 +16,7 @@ limitations under the License. #include "xla/codegen/intrinsic/fptrunc.h" #include +#include #include #include #include @@ -371,6 +372,34 @@ TEST(FpTruncExecutionTest, F32ToF8e4m3fn_Vector8) { EXPECT_EQ(actuals[7], static_cast(0b01110111)); } +TEST(FpTruncExecutionTest, F32ToBf16_Vector16) { + constexpr size_t kN = 16; + const Type from = Type::V(F32, kN); + const Type to = Type::V(BF16, kN); + JitRunner jit = CreateJitRunner(from, to); + auto fptrunc = jit.GetVectorizedFn( + FpTrunc::Name({from, to}) + "_fptoi"); + constexpr float kEps = std::numeric_limits::epsilon(); + std::array vals = { + 1.0f, -2.0f, 0.5f, 0.0f, -0.0f, 3.0f, -0.75f, + // Rounds up to 65536.0f. + 65504.0f, + // Ties round to even: 1 + 2^-8 rounds down, 1 + 3 * 2^-8 rounds up. + 1.00390625f, 1.01171875f, + // Just above a tie rounds up. + 1.00390625f + kEps, + // Rounds up to infinity. + std::numeric_limits::max(), + std::numeric_limits::denorm_min(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + std::numeric_limits::quiet_NaN()}; + std::array expected = { + 0x3F80, 0xC000, 0x3F00, 0x0000, 0x8000, 0x4040, 0xBF40, 0x4780, + 0x3F80, 0x3F82, 0x3F81, 0x7F80, 0x0000, 0x7F80, 0xFF80, 0x7FC0}; + EXPECT_EQ(fptrunc(vals), expected); +} + TEST(FpTruncExecutionTest, F32ToF8e4m3b11fnuz) { JitRunner jit = CreateJitRunner(Type::S(F32), Type::S(F8E4M3B11FNUZ)); auto fptrunc = jit.GetScalarFn( diff --git a/third_party/xla/xla/codegen/intrinsic/log1p.h b/third_party/xla/xla/codegen/intrinsic/log1p.h index 1b3f3913eb66c5..3e3b4753e0a8f3 100644 --- a/third_party/xla/xla/codegen/intrinsic/log1p.h +++ b/third_party/xla/xla/codegen/intrinsic/log1p.h @@ -35,9 +35,11 @@ class Log1p : public Intrinsic { static std::vector> SupportedVectorTypes( absl::string_view features = "") { return { - {Type::S(F16)}, {Type::V(F16, 2)}, {Type::V(F16, 4)}, {Type::V(F16, 8)}, - {Type::S(F32)}, {Type::V(F32, 2)}, {Type::V(F32, 4)}, {Type::V(F32, 8)}, - {Type::S(F64)}, {Type::V(F64, 2)}, {Type::V(F64, 4)}, {Type::V(F64, 8)}, + {Type::S(F16)}, {Type::V(F16, 2)}, {Type::V(F16, 4)}, + {Type::V(F16, 8)}, {Type::S(F32)}, {Type::V(F32, 2)}, + {Type::V(F32, 4)}, {Type::V(F32, 8)}, {Type::V(F32, 16)}, + {Type::S(F64)}, {Type::V(F64, 2)}, {Type::V(F64, 4)}, + {Type::V(F64, 8)}, }; } }; diff --git a/third_party/xla/xla/codegen/intrinsic/log1p_test.cc b/third_party/xla/xla/codegen/intrinsic/log1p_test.cc index 0e54593786e1c1..3c90a3358cd54d 100644 --- a/third_party/xla/xla/codegen/intrinsic/log1p_test.cc +++ b/third_party/xla/xla/codegen/intrinsic/log1p_test.cc @@ -15,7 +15,9 @@ limitations under the License. #include "xla/codegen/intrinsic/log1p.h" +#include #include +#include #include #include #include @@ -100,6 +102,28 @@ TEST(Log1pTest, F32) { } } +TEST(Log1pTest, F32Vector16) { + constexpr size_t kN = 16; + Type type = Type::V(F32, kN); + JitRunner runner = CreateJitRunnerWithLog1p(type); + auto fn = runner.GetVectorizedFn(Log1p::Name(type)); + + std::vector test_values = GetTestValues(); + std::array vals; + for (size_t i = 0; i < kN; ++i) { + vals[i] = test_values[i % test_values.size()]; + } + std::array results = fn(vals); + for (size_t i = 0; i < kN; ++i) { + float expected = std::log1pf(vals[i]); + if (std::isnan(expected)) { + EXPECT_TRUE(std::isnan(results[i])); + } else { + EXPECT_THAT(results[i], NearUlps(expected, 1)); + } + } +} + TEST(Log1pTest, F64) { Type type = Type::S(F64); JitRunner runner = CreateJitRunnerWithLog1p(type); diff --git a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc index bfb664909a058c..0c7f9cfd2256d0 100644 --- a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc +++ b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc @@ -1715,6 +1715,12 @@ absl::StatusOr EmitReduceWithRegion( mlir::ImplicitLocOpBuilder nested_b(loc, loop_builder); for (int i = 0; i < sequential_dim_ids.size(); ++i) { + // Fully tiled dimensions were already bound to 0 by + // EmitFullyTiledSequentialDimensions. Their single-trip loop is + // kept so that masking below is handled uniformly. + if (loop_iteration_counts[i] == 1) { + continue; + } const ge::TilingSpace::DimensionInfo& dim_info = tiled_hlo.tile().tiling_space().GetDimensionInfo( *tiled_hlo.hlo(), sequential_dim_ids[i]); @@ -2192,6 +2198,11 @@ absl::Status EmitGeneric(ImplicitLocOpBuilder& b, b.setInsertionPointToStart(&fn.front()); Value program_id = fn.getProgramId(); Value tile_id = program_id; + // Record the bounds of the program id on the function argument. Otherwise + // they only live in the domains of the apply_indexing ops using it, and get + // lost when identity apply_indexing ops are folded. + fn.setArgAttr(mlir::cast(program_id).getArgNumber(), + "xla.range", b.getIndexArrayAttr({0, schedule.num_pids - 1})); // If there are more than one tile per pid, we need to add a scf.for loop to // iterate through the tiles. diff --git a/third_party/xla/xla/comparison_util.cc b/third_party/xla/xla/comparison_util.cc index aed2072aa2d5d5..c0daede5b065b7 100644 --- a/third_party/xla/xla/comparison_util.cc +++ b/third_party/xla/xla/comparison_util.cc @@ -32,19 +32,6 @@ limitations under the License. namespace xla { namespace { -// Returns the X32 primitive type for each Type. -PrimitiveType DefaultPrimitiveType(Comparison::Type type) { - switch (type) { - case Comparison::Type::kFloat: - case Comparison::Type::kFloatTotalOrder: - return PrimitiveType::F32; - case Comparison::Type::kSigned: - return PrimitiveType::S32; - case Comparison::Type::kUnsigned: - return PrimitiveType::U32; - } -} - // Returns the converse of `direction`. Comparison::Direction Converse(Comparison::Direction direction) { switch (direction) { @@ -102,19 +89,6 @@ std::string ComparisonDirectionToString(Comparison::Direction direction) { } } -std::string ComparisonTypeToString(Comparison::Type type) { - switch (type) { - case Comparison::Type::kFloat: - return "FLOAT"; - case Comparison::Type::kFloatTotalOrder: - return "TOTALORDER"; - case Comparison::Type::kSigned: - return "SIGNED"; - case Comparison::Type::kUnsigned: - return "UNSIGNED"; - } -} - absl::string_view ComparisonPrimitiveTypeToString(PrimitiveType type) { return PrimitiveType_Name(type); } @@ -123,6 +97,8 @@ absl::string_view ComparisonOrderToString(Comparison::Order order) { switch (order) { case Comparison::Order::kPartial: return "PARTIALORDER"; + case Comparison::Order::kWeak: + return "WEAKORDER"; case Comparison::Order::kTotal: return "TOTALORDER"; } @@ -132,6 +108,8 @@ absl::string_view ComparisonOrderToShortString(Comparison::Order order) { switch (order) { case Comparison::Order::kPartial: return "PARTIAL"; + case Comparison::Order::kWeak: + return "WEAK"; case Comparison::Order::kTotal: return "TOTAL"; } @@ -140,7 +118,7 @@ absl::string_view ComparisonOrderToShortString(Comparison::Order order) { absl::StatusOr StringToComparisonDirection( absl::string_view direction) { static auto* const map = - new absl::flat_hash_map({ + new absl::flat_hash_map({ {"EQ", Comparison::Direction::kEq}, {"NE", Comparison::Direction::kNe}, {"GE", Comparison::Direction::kGe}, @@ -158,8 +136,9 @@ absl::StatusOr StringToComparisonDirection( absl::StatusOr ShortStringToComparisonOrder( absl::string_view order) { static auto* const map = - new absl::flat_hash_map({ + new absl::flat_hash_map({ {"TOTAL", Comparison::Order::kTotal}, + {"WEAK", Comparison::Order::kWeak}, {"PARTIAL", Comparison::Order::kPartial}, }); auto it = map->find(order); @@ -169,48 +148,22 @@ absl::StatusOr ShortStringToComparisonOrder( return it->second; } -absl::StatusOr StringToComparisonType( - absl::string_view comparison) { +absl::StatusOr ComparisonTypeToOrder( + absl::string_view comparison_type) { static auto* const map = - new absl::flat_hash_map({ - {"FLOAT", Comparison::Type::kFloat}, - {"TOTALORDER", Comparison::Type::kFloatTotalOrder}, - {"SIGNED", Comparison::Type::kSigned}, - {"UNSIGNED", Comparison::Type::kUnsigned}, + new absl::flat_hash_map({ + {"FLOAT", Comparison::Order::kPartial}, + {"TOTALORDER", Comparison::Order::kTotal}, + {"SIGNED", Comparison::Order::kTotal}, + {"UNSIGNED", Comparison::Order::kTotal}, }); - auto it = map->find(comparison); + auto it = map->find(comparison_type); if (it == map->end()) { - return InvalidArgument("Unknown comparison type: %s", comparison); + return InvalidArgument("Unknown comparison type: %s", comparison_type); } return it->second; } -Comparison::Type Comparison::DefaultComparisonType(PrimitiveType type) { - if (primitive_util::IsFloatingPointType(type) || - primitive_util::IsComplexType(type)) { - return Type::kFloat; - } - if (primitive_util::IsSignedIntegralType(type)) { - return Type::kSigned; - } - if (primitive_util::IsUnsignedIntegralType(type) || type == PRED) { - return Type::kUnsigned; - } - LOG(FATAL) << "Unexpected: " << PrimitiveType_Name(type); -} - -// Returns the default ordering for each Comparison::Type. -Comparison::Order Comparison::DefaultOrdering(Comparison::Type type) { - switch (type) { - case Comparison::Type::kFloat: - return Comparison::Order::kPartial; - case Comparison::Type::kFloatTotalOrder: - case Comparison::Type::kSigned: - case Comparison::Type::kUnsigned: - return Comparison::Order::kTotal; - } -} - // Returns the expected ordering for each primitive type. Comparison::Order Comparison::DefaultOrdering(PrimitiveType type) { if (primitive_util::IsFloatingPointType(type) || @@ -232,20 +185,15 @@ Comparison::Comparison(Direction dir, PrimitiveType type, Order order) Comparison::Comparison(Direction dir, PrimitiveType type) : dir_(dir), primitive_type_(type), order_(DefaultOrdering(type)) {} -Comparison::Comparison(Direction dir, Type type) - : dir_(dir), - primitive_type_(DefaultPrimitiveType(type)), - order_(DefaultOrdering(type)) {} - Comparison Comparison::Converse() const { return Comparison(xla::Converse(dir_), primitive_type_, order_); } std::optional Comparison::Inverse() const { if (IsPartialOrder()) { - // We assume comparisons don't have inverses unless they are total order, - // e.g., a partial order floating point comparison can return true if one - // operand is NaN. + // We assume comparisons don't have inverses unless they are total or weak + // order, e.g., a partial order floating point comparison can return false + // for both (a < b) and (a >= b) if one operand is NaN. return std::nullopt; } if (primitive_util::IsArrayType(primitive_type_)) { @@ -259,7 +207,7 @@ bool Comparison::IsReflexive() const { case Direction::kEq: case Direction::kGe: case Direction::kLe: - return IsTotalOrder(); + return IsTotalOrder() || IsWeakOrder(); case Direction::kNe: case Direction::kGt: case Direction::kLt: @@ -270,7 +218,7 @@ bool Comparison::IsReflexive() const { bool Comparison::IsAntireflexive() const { switch (dir_) { case Direction::kNe: - return IsTotalOrder(); + return IsTotalOrder() || IsWeakOrder(); case Direction::kGt: case Direction::kLt: return true; diff --git a/third_party/xla/xla/comparison_util.h b/third_party/xla/xla/comparison_util.h index a626a8d772dcab..ec5fc68220dd6e 100644 --- a/third_party/xla/xla/comparison_util.h +++ b/third_party/xla/xla/comparison_util.h @@ -61,6 +61,8 @@ class Comparison { enum class Order : uint8_t { // https://en.wikipedia.org/wiki/Total_order kTotal, + // https://en.wikipedia.org/wiki/Weak_ordering + kWeak, // https://en.wikipedia.org/wiki/Partially_ordered_set kPartial, }; @@ -84,16 +86,6 @@ class Comparison { kLt, }; - // (DEPRECATED) Represents the type of comparison. Prefer xla::PrimitiveType - // and Comparison::Order, since there are multiple floating point - // representations that support total ordering. - enum class [[deprecated("Use PrimitiveType and Order")]] Type : uint8_t { - kFloat, - kFloatTotalOrder, - kSigned, - kUnsigned, - }; - Comparison() = delete; // This will default to the expected behavior for Comparison::Order: integers @@ -104,13 +96,6 @@ class Comparison { // targets may support total order floating point type comparisons. explicit Comparison(Direction dir, PrimitiveType type, Order order); - // Returns a comparison with a primitive type matching the Comparison::Type - // and using a default bit width of 32. For example, - // Comparison(Direction::kLt, Type::kFloat).PrimitiveType() /* F32 */ - [[deprecated( - "Use Comparison(Comparison::Direction, " - "PrimitiveType)")]] explicit Comparison(Direction dir, Type type); - inline Direction GetDirection() const { return dir_; } inline PrimitiveType GetPrimitiveType() const { return primitive_type_; } inline Order GetOrder() const { return order_; } @@ -122,6 +107,7 @@ class Comparison { inline bool IsLe() const { return dir_ == Direction::kLe; } inline bool IsLt() const { return dir_ == Direction::kLt; } inline bool IsTotalOrder() const { return order_ == Order::kTotal; } + bool IsWeakOrder() const { return order_ == Order::kWeak; } inline bool IsPartialOrder() const { return order_ == Order::kPartial; } // Returns whether this is a floating point total order comparison. @@ -177,14 +163,27 @@ class Comparison { case Direction::kNe: return +[](T l, T r) { return std::not_equal_to()(l, r); }; case Direction::kGe: - return +[](T l, T r) { return std::greater_equal()(l, r); }; + if constexpr (!is_complex_v) { + return +[](T l, T r) { return std::greater_equal()(l, r); }; + } + break; case Direction::kGt: - return +[](T l, T r) { return std::greater()(l, r); }; + if constexpr (!is_complex_v) { + return +[](T l, T r) { return std::greater()(l, r); }; + } + break; case Direction::kLe: - return +[](T l, T r) { return std::less_equal()(l, r); }; + if constexpr (!is_complex_v) { + return +[](T l, T r) { return std::less_equal()(l, r); }; + } + break; case Direction::kLt: - return +[](T l, T r) { return std::less()(l, r); }; + if constexpr (!is_complex_v) { + return +[](T l, T r) { return std::less()(l, r); }; + } + break; } + LOG(FATAL) << "Unsupported comparison: " << ToString(); } template @@ -203,24 +202,22 @@ class Comparison { return GetComparator()(ToSignMagnitude(a), ToSignMagnitude(b)); } } + if (IsWeakOrder()) { + // -Inf < -Finite < (-0 == +0) < +Finite < +Inf < NaN + bool a_is_nan = Eigen::numext::isnan(a); + bool b_is_nan = Eigen::numext::isnan(b); + if (a_is_nan || b_is_nan) { + return GetComparator()(a_is_nan, b_is_nan); + } + } } // Applies the comparison from this Comparison's direction and ordering. return GetComparator()(a, b); } - // Returns the Comparison::Order corresponding to the deprecated - // Comparison::Type. - static Comparison::Order DefaultOrdering(Type type); - // Returns the expected Comparison::Order for each primitive type. static Comparison::Order DefaultOrdering(PrimitiveType type); - // Returns the Comparison::Type for the given primitive type. This assumes - // that each numerical representation follows the standard behavior, e.g., - // integers are total order and floats are partial order. - [[deprecated("Use PrimitiveType and Order")]] static Comparison::Type - DefaultComparisonType(PrimitiveType type); - private: // The direction of the Comparison, e.g., GT. const Direction dir_; @@ -238,7 +235,6 @@ inline std::ostream& operator<<(std::ostream& os, const Comparison& cmp) { } std::string ComparisonDirectionToString(Comparison::Direction direction); -std::string ComparisonTypeToString(Comparison::Type type); absl::string_view ComparisonPrimitiveTypeToString(PrimitiveType type); absl::string_view ComparisonOrderToString(Comparison::Order order); absl::string_view ComparisonOrderToShortString(Comparison::Order order); @@ -252,8 +248,8 @@ absl::StatusOr StringToComparisonDirection( absl::string_view direction); absl::StatusOr ShortStringToComparisonOrder( absl::string_view order); -absl::StatusOr StringToComparisonType( - absl::string_view comparison); +absl::StatusOr ComparisonTypeToOrder( + absl::string_view comparison_type); // Returns a comparison function using the provided key function on each value, // i.e. `key_fn(a) < key_fn(b)`. diff --git a/third_party/xla/xla/comparison_util_test.cc b/third_party/xla/xla/comparison_util_test.cc index a0f10c173e148c..aca20dd4512a42 100644 --- a/third_party/xla/xla/comparison_util_test.cc +++ b/third_party/xla/xla/comparison_util_test.cc @@ -54,24 +54,16 @@ TEST(Comparison, IntegersDefaultToTotalOrder) { Comparison::Order::kTotal); } -TEST(Comparison, LegacyConstructorDefaultsToX32) { - // We expect the legacy constructor to default to a {S,U,F}32 primitive type. - // This is legacy debt of the previous enumeration Comparison::Type, which - // attempted to combine ordering and numerical data classification in a single - // type. - EXPECT_EQ(Comparison(Comparison::Direction::kGe, Comparison::Type::kFloat) - .GetPrimitiveType(), - xla::PrimitiveType::F32); - EXPECT_EQ( - Comparison(Comparison::Direction::kGe, Comparison::Type::kFloatTotalOrder) - .GetPrimitiveType(), - xla::PrimitiveType::F32); - EXPECT_EQ(Comparison(Comparison::Direction::kGe, Comparison::Type::kSigned) - .GetPrimitiveType(), - xla::PrimitiveType::S32); - EXPECT_EQ(Comparison(Comparison::Direction::kGe, Comparison::Type::kUnsigned) - .GetPrimitiveType(), - xla::PrimitiveType::U32); +TEST(Comparison, ComparisonTypeToOrder) { + EXPECT_THAT(ComparisonTypeToOrder("FLOAT"), + IsOkAndHolds(Comparison::Order::kPartial)); + EXPECT_THAT(ComparisonTypeToOrder("TOTALORDER"), + IsOkAndHolds(Comparison::Order::kTotal)); + EXPECT_THAT(ComparisonTypeToOrder("SIGNED"), + IsOkAndHolds(Comparison::Order::kTotal)); + EXPECT_THAT(ComparisonTypeToOrder("UNSIGNED"), + IsOkAndHolds(Comparison::Order::kTotal)); + EXPECT_FALSE(ComparisonTypeToOrder("INVALID").ok()); } TEST(Comparison, PartialOrderReflexivity) { @@ -180,12 +172,14 @@ TEST(Comparison, ToString) { TEST(Comparison, ComparisonOrderToString) { EXPECT_EQ(ComparisonOrderToString(Comparison::Order::kTotal), "TOTALORDER"); + EXPECT_EQ(ComparisonOrderToString(Comparison::Order::kWeak), "WEAKORDER"); EXPECT_EQ(ComparisonOrderToString(Comparison::Order::kPartial), "PARTIALORDER"); } TEST(Comparison, ComparisonOrderToShortString) { EXPECT_EQ(ComparisonOrderToShortString(Comparison::Order::kTotal), "TOTAL"); + EXPECT_EQ(ComparisonOrderToShortString(Comparison::Order::kWeak), "WEAK"); EXPECT_EQ(ComparisonOrderToShortString(Comparison::Order::kPartial), "PARTIAL"); } @@ -193,13 +187,53 @@ TEST(Comparison, ComparisonOrderToShortString) { TEST(Comparison, ShortStringToComparisonOrder) { EXPECT_THAT(ShortStringToComparisonOrder("TOTAL"), IsOkAndHolds(Comparison::Order::kTotal)); + EXPECT_THAT(ShortStringToComparisonOrder("WEAK"), + IsOkAndHolds(Comparison::Order::kWeak)); EXPECT_THAT(ShortStringToComparisonOrder("PARTIAL"), IsOkAndHolds(Comparison::Order::kPartial)); EXPECT_FALSE(ShortStringToComparisonOrder("TOTALORDER").ok()); + EXPECT_FALSE(ShortStringToComparisonOrder("WEAKORDER").ok()); EXPECT_FALSE(ShortStringToComparisonOrder("PARTIALORDER").ok()); EXPECT_FALSE(ShortStringToComparisonOrder("INVALID").ok()); } +TEST(Comparison, WeakOrderFloatComparison) { + const float qnan = std::numeric_limits::quiet_NaN(); + const float neg_qnan = -std::numeric_limits::quiet_NaN(); + const float inf = std::numeric_limits::infinity(); + const float neg_inf = -std::numeric_limits::infinity(); + + EXPECT_TRUE(Comparison(Comparison::Direction::kEq, PrimitiveType::F32, + Comparison::Order::kWeak) + .Compare(qnan, neg_qnan)); + EXPECT_FALSE(Comparison(Comparison::Direction::kNe, PrimitiveType::F32, + Comparison::Order::kWeak) + .Compare(qnan, neg_qnan)); + + EXPECT_TRUE(Comparison(Comparison::Direction::kEq, PrimitiveType::F32, + Comparison::Order::kWeak) + .Compare(-0.0f, +0.0f)); + EXPECT_FALSE(Comparison(Comparison::Direction::kLt, PrimitiveType::F32, + Comparison::Order::kWeak) + .Compare(-0.0f, +0.0f)); + EXPECT_TRUE(Comparison(Comparison::Direction::kLe, PrimitiveType::F32, + Comparison::Order::kWeak) + .Compare(-0.0f, +0.0f)); + + EXPECT_TRUE(Comparison(Comparison::Direction::kLt, PrimitiveType::F32, + Comparison::Order::kWeak) + .Compare(inf, neg_qnan)); + EXPECT_FALSE(Comparison(Comparison::Direction::kLt, PrimitiveType::F32, + Comparison::Order::kWeak) + .Compare(neg_qnan, neg_inf)); + EXPECT_TRUE(Comparison(Comparison::Direction::kGt, PrimitiveType::F32, + Comparison::Order::kWeak) + .Compare(neg_qnan, inf)); + EXPECT_TRUE(Comparison(Comparison::Direction::kGe, PrimitiveType::F32, + Comparison::Order::kWeak) + .Compare(neg_qnan, qnan)); +} + TEST(Comparison, TotalOrderFloatComparison) { EXPECT_TRUE(Comparison(Comparison::Direction::kEq, PrimitiveType::F32, Comparison::Order::kTotal) diff --git a/third_party/xla/xla/experiments/sm_bandwidth_benchmark/BUILD b/third_party/xla/xla/experiments/sm_bandwidth_benchmark/BUILD index 8919b4b4cf8aac..4aab3e9c88a551 100644 --- a/third_party/xla/xla/experiments/sm_bandwidth_benchmark/BUILD +++ b/third_party/xla/xla/experiments/sm_bandwidth_benchmark/BUILD @@ -49,6 +49,7 @@ xla_cc_test( tags = [ "cuda-only", "gpu", + "nomsan", "requires-gpu-sm90", ], deps = [ diff --git a/third_party/xla/xla/hlo/builder/lib/BUILD b/third_party/xla/xla/hlo/builder/lib/BUILD index ad1f5002716126..5bce2511b0c6a0 100644 --- a/third_party/xla/xla/hlo/builder/lib/BUILD +++ b/third_party/xla/xla/hlo/builder/lib/BUILD @@ -324,6 +324,7 @@ xla_test( "//xla:array4d", "//xla:literal", "//xla:literal_util", + "//xla:shape_util", "//xla:types", "//xla/hlo/builder:xla_builder", "//xla/hlo/testlib:test", diff --git a/third_party/xla/xla/hlo/builder/lib/matrix.cc b/third_party/xla/xla/hlo/builder/lib/matrix.cc index 3587d6a447f396..1024575cf8ca6a 100644 --- a/third_party/xla/xla/hlo/builder/lib/matrix.cc +++ b/third_party/xla/xla/hlo/builder/lib/matrix.cc @@ -457,9 +457,14 @@ xla::XlaOp Einsum(xla::XlaOp x, absl::Span x_config, } else if (x_shape.dimensions(i) == 1) { rhs_outer_dims.push_back(rhs_dim); lhs_delete_dims.push_back(i); - } else { + } else if (y_shape.dimensions(rhs_dim) == 1) { lhs_outer_dims.push_back(i); rhs_delete_dims.push_back(rhs_dim); + } else { + return InvalidArgument( + "Batch dimensions %d of LHS and %d of RHS must match or be 1. " + "Got %d and %d.", + i, rhs_dim, x_shape.dimensions(i), y_shape.dimensions(rhs_dim)); } } else if (is_contracting(dim_name)) { if (x_shape.dimensions(i) == y_shape.dimensions(rhs_dim)) { @@ -468,9 +473,15 @@ xla::XlaOp Einsum(xla::XlaOp x, absl::Span x_config, } else if (x_shape.dimensions(i) == 1) { rhs_outer_dims.push_back(rhs_dim); lhs_delete_dims.push_back(i); - } else { + } else if (y_shape.dimensions(rhs_dim) == 1) { lhs_outer_dims.push_back(i); rhs_delete_dims.push_back(rhs_dim); + } else { + return InvalidArgument( + "Contracting dimensions %d of LHS and %d of RHS must match or be " + "1. " + "Got %d and %d.", + i, rhs_dim, x_shape.dimensions(i), y_shape.dimensions(rhs_dim)); } } else { lhs_outer_dims.push_back(i); diff --git a/third_party/xla/xla/hlo/builder/lib/matrix_test.cc b/third_party/xla/xla/hlo/builder/lib/matrix_test.cc index b3bc3b87f953c5..598b5a3c0d459f 100644 --- a/third_party/xla/xla/hlo/builder/lib/matrix_test.cc +++ b/third_party/xla/xla/hlo/builder/lib/matrix_test.cc @@ -36,6 +36,7 @@ limitations under the License. #include "xla/hlo/testlib/test.h" #include "xla/literal.h" #include "xla/literal_util.h" +#include "xla/shape_util.h" #include "xla/tests/client_library_test_runner_mixin.h" #include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h" #include "xla/tests/hlo_pjrt_test_base.h" @@ -329,5 +330,39 @@ TEST_F(MatrixTest, NormalizeEinsumString) { EXPECT_EQ(NormalizeEinsumString("...ba,ca..."), "...ba,ca...->...bc"); } +TEST_F(MatrixTest, EinsumIncompatibleContractingDims) { + XlaBuilder builder(TestName()); + auto a = ConstantR2(&builder, {{1.0, 2.0, 3.0, 4.0}}); + auto b = + ConstantR2(&builder, {{1.0}, {2.0}, {3.0}, {4.0}, {5.0}, {6.0}}); + Einsum(a, b, "ab,bc->ac"); + EXPECT_FALSE(builder.GetCurrentStatus().ok()); +} + +TEST_F(MatrixTest, BatchDotIncompatibleContractingDims) { + XlaBuilder builder(TestName()); + auto a = ConstantR2(&builder, {{1.0, 2.0, 3.0, 4.0}}); + auto b = + ConstantR2(&builder, {{1.0}, {2.0}, {3.0}, {4.0}, {5.0}, {6.0}}); + BatchDot(a, b); + EXPECT_FALSE(builder.GetCurrentStatus().ok()); +} + +TEST_F(MatrixTest, EinsumIncompatibleBatchDims) { + XlaBuilder builder(TestName()); + auto a = Parameter(&builder, 0, ShapeUtil::MakeShape(F32, {3, 2, 2}), "a"); + auto b = Parameter(&builder, 1, ShapeUtil::MakeShape(F32, {5, 2, 2}), "b"); + Einsum(a, b, "...ij,...jk->...ik"); + EXPECT_FALSE(builder.GetCurrentStatus().ok()); +} + +TEST_F(MatrixTest, EinsumBroadcastingContractingDims) { + XlaBuilder builder(TestName()); + auto a = ConstantR2(&builder, {{1.0}}); + auto b = ConstantR2(&builder, {{1.0}, {2.0}, {3.0}}); + Einsum(a, b, "ab,bc->ac"); + EXPECT_TRUE(builder.GetCurrentStatus().ok()); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/hlo/builder/xla_builder.cc b/third_party/xla/xla/hlo/builder/xla_builder.cc index 7072aaf18ec7a0..fafc71bb7d0909 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder.cc +++ b/third_party/xla/xla/hlo/builder/xla_builder.cc @@ -5801,13 +5801,6 @@ XlaOp Compare(const XlaOp lhs, const XlaOp rhs, broadcast_dimensions, direction, order); } -XlaOp Compare(const XlaOp lhs, const XlaOp rhs, - absl::Span broadcast_dimensions, - ComparisonDirection direction, Comparison::Type compare_type) { - return Compare(lhs, rhs, broadcast_dimensions, direction, - Comparison::DefaultOrdering(compare_type)); -} - XlaOp Compare(const XlaOp lhs, const XlaOp rhs, ComparisonDirection direction) { return Compare(lhs, rhs, {}, direction); } diff --git a/third_party/xla/xla/hlo/builder/xla_builder.h b/third_party/xla/xla/hlo/builder/xla_builder.h index d91e6592ebae4a..7cceabb9adc353 100644 --- a/third_party/xla/xla/hlo/builder/xla_builder.h +++ b/third_party/xla/xla/hlo/builder/xla_builder.h @@ -1528,10 +1528,6 @@ class XlaBuilder { friend XlaOp Compare(XlaOp lhs, XlaOp rhs, absl::Span broadcast_dimensions, ComparisonDirection direction, ComparisonOrder order); - friend XlaOp Compare(XlaOp lhs, XlaOp rhs, - absl::Span broadcast_dimensions, - ComparisonDirection direction, - Comparison::Type compare_type); friend XlaOp Dot(XlaOp lhs, XlaOp rhs, const PrecisionConfig* precision_config, std::optional preferred_element_type); @@ -2588,9 +2584,6 @@ XlaOp LeTotalOrder(XlaOp lhs, XlaOp rhs, XlaOp Compare(XlaOp lhs, XlaOp rhs, absl::Span broadcast_dimensions, ComparisonDirection direction, ComparisonOrder order); -XlaOp Compare(XlaOp lhs, XlaOp rhs, - absl::Span broadcast_dimensions, - ComparisonDirection direction, Comparison::Type compare_type); XlaOp Compare(XlaOp lhs, XlaOp rhs, absl::Span broadcast_dimensions, ComparisonDirection direction); diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc index 69b3cc8cff771c..5637a5f897d5f6 100644 --- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc +++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc @@ -102,70 +102,35 @@ template absl::StatusOr Compare(const Shape& shape, Comparison comparison, LiteralSlice lhs_literal, LiteralSlice rhs_literal) { - auto populate = [&](auto compare_op) -> absl::StatusOr { - Literal result(shape); - - // If layout is the same, we can use linear indexing into the literals. - const Layout& lhs_layout = lhs_literal.shape().layout(); - const Layout& rhs_layout = rhs_literal.shape().layout(); - bool same_layout = LayoutUtil::Equal(lhs_layout, rhs_layout) && - LayoutUtil::Equal(lhs_layout, shape.layout()); - - if (same_layout) { - ABSL_RETURN_IF_ERROR(result.PopulateLinearParallel( - [&](int64_t linear_index, int /*thread_id*/) { - auto lhs = lhs_literal.GetLinear(linear_index); - auto rhs = rhs_literal.GetLinear(linear_index); - if constexpr (is_specialized_floating_point_v) { - if (comparison.IsTotalOrder()) { - return compare_op(ToSignMagnitude(lhs), ToSignMagnitude(rhs)); - } - } - return compare_op(lhs, rhs); - })); - } else { - ABSL_RETURN_IF_ERROR(result.PopulateParallel( - [&](absl::Span multi_index, int /*thread_id*/) { - auto lhs = lhs_literal.Get(multi_index); - auto rhs = rhs_literal.Get(multi_index); - if constexpr (is_specialized_floating_point_v) { - if (comparison.IsTotalOrder()) { - return compare_op(ToSignMagnitude(lhs), ToSignMagnitude(rhs)); - } - } - return compare_op(lhs, rhs); - })); + if constexpr (is_complex_v) { + if (comparison.GetDirection() != ComparisonDirection::kEq && + comparison.GetDirection() != ComparisonDirection::kNe) { + return Unimplemented("Unsupported comparison: %s", comparison.ToString()); } - return result; - }; - switch (comparison.GetDirection()) { - case ComparisonDirection::kEq: - return populate([](auto lhs, auto rhs) { return lhs == rhs; }); - case ComparisonDirection::kNe: - return populate([](auto lhs, auto rhs) { return lhs != rhs; }); - case ComparisonDirection::kGe: - if constexpr (!is_complex_v) { - return populate([](auto lhs, auto rhs) { return lhs >= rhs; }); - } - break; - case ComparisonDirection::kGt: - if constexpr (!is_complex_v) { - return populate([](auto lhs, auto rhs) { return lhs > rhs; }); - } - break; - case ComparisonDirection::kLe: - if constexpr (!is_complex_v) { - return populate([](auto lhs, auto rhs) { return lhs <= rhs; }); - } - break; - case ComparisonDirection::kLt: - if constexpr (!is_complex_v) { - return populate([](auto lhs, auto rhs) { return lhs < rhs; }); - } - break; } + Literal result(shape); - return Unimplemented("Unsupported comparison: %s", comparison.ToString()); + // If layout is the same, we can use linear indexing into the literals. + const Layout& lhs_layout = lhs_literal.shape().layout(); + const Layout& rhs_layout = rhs_literal.shape().layout(); + bool same_layout = LayoutUtil::Equal(lhs_layout, rhs_layout) && + LayoutUtil::Equal(lhs_layout, shape.layout()); + + if (same_layout) { + ABSL_RETURN_IF_ERROR(result.PopulateLinearParallel( + [&](int64_t linear_index, int /*thread_id*/) { + return comparison.Compare( + lhs_literal.GetLinear(linear_index), + rhs_literal.GetLinear(linear_index)); + })); + } else { + ABSL_RETURN_IF_ERROR(result.PopulateParallel( + [&](absl::Span multi_index, int /*thread_id*/) { + return comparison.Compare(lhs_literal.Get(multi_index), + rhs_literal.Get(multi_index)); + })); + } + return result; } std::optional GetInstructionStaticValueAsBool( diff --git a/third_party/xla/xla/hlo/ir/hlo_clone_context.h b/third_party/xla/xla/hlo/ir/hlo_clone_context.h index b1eaa31306c7cc..6cbe42966cf68e 100644 --- a/third_party/xla/xla/hlo/ir/hlo_clone_context.h +++ b/third_party/xla/xla/hlo/ir/hlo_clone_context.h @@ -19,6 +19,7 @@ limitations under the License. #include #include "absl/container/flat_hash_map.h" +#include "absl/strings/string_view.h" #include "xla/map_util.h" namespace xla { @@ -34,7 +35,7 @@ class HloCloneContext { // Creates a new HloCloneContext object to clone HloInstruction and // HloComputation objects to be added to the module specified as argument. // The suffix string will be appended to computation names. - explicit HloCloneContext(HloModule* module, const std::string& suffix = "") + explicit HloCloneContext(HloModule* module, absl::string_view suffix = "") : module_(module), suffix_(suffix) {} HloModule* module() const { return module_; } diff --git a/third_party/xla/xla/hlo/ir/hlo_computation.cc b/third_party/xla/xla/hlo/ir/hlo_computation.cc index 1331309626a758..67e52158dc14ad 100644 --- a/third_party/xla/xla/hlo/ir/hlo_computation.cc +++ b/third_party/xla/xla/hlo/ir/hlo_computation.cc @@ -206,6 +206,12 @@ HloComputation::~HloComputation() { CHECK(FusionInstruction()->fused_instructions_computation() == this); FusionInstruction()->ClearCalledComputations(); } + // Every live instruction dies below, so only the edges that leave the + // computation are unlinked; the call also makes ~HloInstruction skip the + // rest. Instructions in to_be_deleted_ were detached at removal. + for (HloInstruction* instruction : instructions()) { + instruction->DetachFromOperandsAndUsersOutside(this); + } Cleanup(); ClearCalledComputations(); diff --git a/third_party/xla/xla/hlo/ir/hlo_computation_test.cc b/third_party/xla/xla/hlo/ir/hlo_computation_test.cc index 06d5157d2a63a4..39b443ae98f2f5 100644 --- a/third_party/xla/xla/hlo/ir/hlo_computation_test.cc +++ b/third_party/xla/xla/hlo/ir/hlo_computation_test.cc @@ -25,9 +25,11 @@ limitations under the License. #include "absl/strings/match.h" #include "absl/strings/string_view.h" #include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_print_options.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" +#include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tsl/platform/statusor.h" #include "xla/xla_data.pb.h" @@ -333,6 +335,85 @@ ENTRY entry { EXPECT_EQ(get_local_id(comp2, "out"), 4); } +// An instruction outside the destroyed computation that uses one of its +// instructions is left with null operand slots, the state +// DetachFromOperandsAndUsers leaves behind. The user names the instruction +// twice so that every slot is scanned. +TEST_F(HLOComputationTest, DestroyingComputationClearsOperandsOfOutsideUsers) { + const Shape shape = ShapeUtil::MakeShape(F32, {}); + auto builder = HloComputation::Builder("inside"); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + HloInstruction* negate = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, param)); + std::unique_ptr computation = builder.Build(); + + std::unique_ptr outside_user = + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, negate, negate); + ASSERT_EQ(outside_user->parent(), nullptr); + ASSERT_EQ(negate->user_count(), 1); + + computation.reset(); + + EXPECT_EQ(outside_user->operand(0), nullptr); + EXPECT_EQ(outside_user->operand(1), nullptr); +} + +// The reverse direction: an instruction of the destroyed computation uses an +// outside instruction, twice. The outside instruction drops the deleted user +// from its user list. +TEST_F(HLOComputationTest, DestroyingComputationDropsUsersOfOutsideOperands) { + const Shape shape = ShapeUtil::MakeShape(F32, {}); + std::unique_ptr outside_operand = + HloInstruction::CreateParameter(0, shape, "outside"); + auto builder = HloComputation::Builder("inside"); + HloInstruction* add = builder.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kAdd, outside_operand.get(), outside_operand.get())); + std::unique_ptr computation = builder.Build(); + ASSERT_EQ(outside_operand->user_count(), 1); + ASSERT_EQ(outside_operand->users().front(), add); + + computation.reset(); + + EXPECT_EQ(outside_operand->user_count(), 0); +} + +// A computation removed from its module is destroyed by HloModule::Cleanup. +// Unlike the two tests above, the outside instructions have a parent +// computation here, so treating only parentless instructions as outside would +// fail this test. +TEST_F(HLOComputationTest, ModuleCleanupUnlinksEdgesToSurvivingComputations) { + const Shape shape = ShapeUtil::MakeShape(F32, {}); + // A plain HloModule: a VerifiedHloModule would fail on the null operand in + // its destructor. + HloModule module("module", GetModuleConfigForTest()); + + auto removed_builder = HloComputation::Builder("removed"); + HloInstruction* removed_param = removed_builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + HloInstruction* removed_negate = removed_builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, removed_param)); + HloComputation* removed = + module.AddEmbeddedComputation(removed_builder.Build()); + + auto kept_builder = HloComputation::Builder("kept"); + HloInstruction* kept_param = kept_builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + HloInstruction* kept_negate = kept_builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, removed_negate)); + module.AddEntryComputation(kept_builder.Build()); + removed->AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, kept_param)); + ASSERT_EQ(kept_negate->operand(0), removed_negate); + ASSERT_EQ(kept_param->user_count(), 1); + + ASSERT_OK(module.RemoveEmbeddedComputation(removed)); + module.Cleanup(); + + EXPECT_EQ(kept_negate->operand(0), nullptr); + EXPECT_EQ(kept_param->user_count(), 0); +} + TEST_F(HLOComputationTest, PrintWithCompactGTE) { absl::string_view hlo_string = R"( HloModule module diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc index 273ace08d474d4..3563727f001b7a 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.cc @@ -124,7 +124,7 @@ void HloInstruction::Users::Clear() { bool HloInstruction::Users::Contains(const HloInstruction* instruction) const { if (user_map_ == nullptr) { - return std::find(users_.begin(), users_.end(), instruction) != users_.end(); + return absl::c_linear_search(users_, instruction); } return user_map_->contains(instruction); } @@ -150,7 +150,7 @@ void HloInstruction::Users::AddUser(HloInstruction* user) { int64_t HloInstruction::Users::UserId(HloInstruction* user) { if (user_map_ == nullptr) { - auto it = std::find(users_.begin(), users_.end(), user); + auto it = absl::c_find(users_, user); CHECK(it != users_.end()); return it - users_.begin(); } @@ -519,29 +519,17 @@ absl::StatusOr> HloInstruction::CreateFromProto( comparison_direction, StringToComparisonDirection(proto.comparison_direction())); } - auto comparison_order_str = proto.comparison_order(); - if (!comparison_order_str.empty()) { - ABSL_ASSIGN_OR_RETURN(auto comparison_order, - ShortStringToComparisonOrder(comparison_order_str)); - instruction = CreateCompare(shape, operands(0), operands(1), - *comparison_direction, comparison_order); - } else { - auto comparison_type_str = proto.comparison_type(); - if (!comparison_type_str.empty()) { - // If a comparison type is specified, it *must* be valid. - ABSL_ASSIGN_OR_RETURN(auto comparison_type, - StringToComparisonType(comparison_type_str)); - instruction = CreateCompare( - shape, operands(0), operands(1), *comparison_direction, - Comparison::DefaultOrdering(comparison_type)); - } else { - // Allow the specification of comparison type to be optional. - // The comparison type will be determined by the types of the - // operands. - instruction = CreateCompare(shape, operands(0), operands(1), - *comparison_direction); - } + std::optional comparison_order; + if (!proto.comparison_order().empty()) { + ABSL_ASSIGN_OR_RETURN(comparison_order, ShortStringToComparisonOrder( + proto.comparison_order())); + } else if (!proto.comparison_type().empty()) { + // If a comparison type is specified, it *must* be valid. + ABSL_ASSIGN_OR_RETURN(comparison_order, + ComparisonTypeToOrder(proto.comparison_type())); } + instruction = CreateCompare(shape, operands(0), operands(1), + *comparison_direction, comparison_order); break; } case HloOpcode::kTriangularSolve: { @@ -2118,9 +2106,9 @@ HloInstruction::CreateCollectivePermuteStart( /* static */ std::unique_ptr HloInstruction::CreateInfeed( const Shape& infeed_shape, HloInstruction* token_operand, - const std::string& config) { + absl::string_view config) { return std::make_unique(infeed_shape, token_operand, - config); + std::string(config)); } /* static */ std::unique_ptr HloInstruction::CreateOutfeed( @@ -2841,7 +2829,7 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( std::unique_ptr HloInstruction::CloneWithNewOperands( const Shape& shape, absl::Span new_operands, - const std::string& suffix, HloCloneContext* context) const { + absl::string_view suffix, HloCloneContext* context) const { VLOG(3) << "CloneWithNewOperands:\n " << ToString(); VLOG(3) << " new operands:"; for (const HloInstruction* new_operand : new_operands) { @@ -3076,15 +3064,22 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( } void HloInstruction::DetachFromOperandsAndUsers() { + DetachFromOperandsAndUsersOutside(/*computation=*/nullptr); +} + +void HloInstruction::DetachFromOperandsAndUsersOutside( + const HloComputation* computation) { if (cleaned_up_) { return; } + DCHECK(computation == nullptr || computation == parent()); cleaned_up_ = true; // Detach from operands. An instruction may be repeated as an operand. To // avoid calling RemoveUser twice on the same operand, check before remove. for (int64_t operand_num = 0; operand_num < operand_count(); ++operand_num) { HloInstruction* operand = operands_[operand_num]; - if (operand == nullptr) { + if (operand == nullptr || + (computation != nullptr && operand->parent() == computation)) { continue; } operand->users_.MaybeRemoveUser(this); @@ -3092,7 +3087,10 @@ void HloInstruction::DetachFromOperandsAndUsers() { } // Update users. Set `nullptr` to the corresponding operand slot for users. - for (auto& user : this->users()) { + for (HloInstruction* user : users()) { + if (computation != nullptr && user->parent() == computation) { + continue; + } for (int i = 0; i < user->operand_count(); ++i) { if (user->operands_[i] == this) { user->operands_[i] = nullptr; @@ -3102,7 +3100,7 @@ void HloInstruction::DetachFromOperandsAndUsers() { } std::unique_ptr HloInstruction::CloneWithNewShape( - const Shape& shape, const std::string& suffix, + const Shape& shape, absl::string_view suffix, HloCloneContext* context) const { std::unique_ptr clone = CloneWithNewOperands(shape, operands_, context); @@ -3115,7 +3113,7 @@ std::unique_ptr HloInstruction::CloneWithNewShape( } std::unique_ptr HloInstruction::Clone( - const std::string& suffix, HloCloneContext* context) const { + absl::string_view suffix, HloCloneContext* context) const { std::unique_ptr clone = CloneWithNewShape(shape(), suffix, context); return clone; @@ -3416,12 +3414,9 @@ void HloInstruction::RemoveOperandsAtAscendingIndices( } bool HloInstruction::HasConstantOperand() const { - for (const HloInstruction* operand : operands_) { - if (operand->IsConstant()) { - return true; - } - } - return false; + return absl::c_any_of(operands_, [](const HloInstruction* operand) { + return operand->IsConstant(); + }); } bool HloInstruction::IdenticalSlowPath( @@ -3610,8 +3605,7 @@ absl::Status HloInstruction::ReplaceUseWithDifferentShape( RemoveUser(user); TF_RET_CHECK(absl::c_count(user->operands_, this) >= 0); - std::replace(user->operands_.begin(), user->operands_.end(), this, - new_producer); + absl::c_replace(user->operands_, this, new_producer); new_producer->AddUser(user); // Custom fusions may not be able to handle deduplicated operands. if (user->opcode() == HloOpcode::kFusion) { @@ -3827,8 +3821,7 @@ absl::Status HloInstruction::ReplaceAllUsesWithDifferentShape( // graph. new_producer remains the only user of this instruction. new_producer_is_user = true; } else { - std::replace(user->operands_.begin(), user->operands_.end(), this, - new_producer); + absl::c_replace(user->operands_, this, new_producer); new_producer->AddUser(user); if (user->opcode() == HloOpcode::kFusion) { ABSL_RETURN_IF_ERROR( @@ -4019,7 +4012,7 @@ std::string PrintCycle(const HloInstruction* child, DFSStack* dfs_stack, while (!dfs.empty() && result.empty()) { bool found_next_instr = false; auto process_users_or_successors = - [&](const std::vector& users_or_successors) { + [&](absl::Span users_or_successors) { for (const auto& user : users_or_successors) { if (user == child) { dfs.push_back(child); @@ -5881,19 +5874,19 @@ absl::StatusOr StringToEnum(absl::string_view value_name, F enum_to_string, } // namespace absl::StatusOr StringToRandomAlgorithm( - const std::string& name) { + absl::string_view name) { return StringToEnum(name, RandomAlgorithmToString, "algorithm"); } absl::StatusOr StringToRandomDistribution( - const std::string& name) { + absl::string_view name) { return StringToEnum(name, RandomDistributionToString, "distribution"); } absl::StatusOr StringToPrecision( - const std::string& name) { + absl::string_view name) { return StringToEnum(name, PrecisionToString, "precision"); } @@ -5915,7 +5908,7 @@ absl::StatusOr StringToResultAccuracy( } absl::StatusOr StringToAlgorithm( - const std::string& name) { + absl::string_view name) { return StringToEnum(name, AlgorithmToString, "algorithm"); } @@ -6291,8 +6284,9 @@ std::string HloInstruction::infeed_config() const { return Cast(this)->infeed_config(); } -void HloInstruction::set_infeed_config(const std::string& config) { - return Cast(this)->set_infeed_config(config); +void HloInstruction::set_infeed_config(absl::string_view config) { + return Cast(this)->set_infeed_config( + std::string(config)); } const Shape& HloInstruction::outfeed_shape() const { @@ -6307,8 +6301,9 @@ const std::string& HloInstruction::outfeed_config() const { return Cast(this)->outfeed_config(); } -void HloInstruction::set_outfeed_config(const std::string& config) { - return Cast(this)->set_outfeed_config(config); +void HloInstruction::set_outfeed_config(absl::string_view config) { + return Cast(this)->set_outfeed_config( + std::string(config)); } const std::vector& HloInstruction::replica_groups() const { diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.h b/third_party/xla/xla/hlo/ir/hlo_instruction.h index a899406d8b6142..c233a089e30a74 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.h +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.h @@ -875,7 +875,7 @@ class HloInstruction { // is a tuple containing the infeed_shape and the TOKEN. static std::unique_ptr CreateInfeed( const Shape& infeed_shape, HloInstruction* token_operand, - const std::string& config); + absl::string_view config); // Creates an outfeed instruction, which outputs data. outfeed_shape is the // shape of the data being outfed *not* the shape of the outfeed instruction @@ -1817,12 +1817,12 @@ class HloInstruction { // the instruction to form the name of the cloned instruction. // Ignores the control predecessors and successors of this HLO instruction. std::unique_ptr Clone( - const std::string& suffix = "clone", + absl::string_view suffix = "clone", HloCloneContext* context = nullptr) const; // Clones the HLO instruction as above but with new shape. std::unique_ptr CloneWithNewShape( - const Shape& shape, const std::string& suffix = "clone", + const Shape& shape, absl::string_view suffix = "clone", HloCloneContext* context = nullptr) const; // Clones the HLO instruction as above but with new shape and operands. @@ -1833,7 +1833,7 @@ class HloInstruction { // Clones the HLO instruction with new shape, operands and suffix. std::unique_ptr CloneWithNewOperands( const Shape& shape, absl::Span new_operands, - const std::string& suffix, HloCloneContext* context = nullptr) const; + absl::string_view suffix, HloCloneContext* context = nullptr) const; // Implementation for non-common logic of CloneWithNewOperands. // CloneWithNewOperands forwards to this method for some of the intstruction @@ -2011,10 +2011,9 @@ class HloInstruction { // Adds a single attribute only if it not already present in the // HloInstruction. Returns false if the attribute was already present. - bool add_frontend_attribute(const std::string& key, - const std::string& value) { - auto it = - mutable_rare()->frontend_attributes.mutable_map()->insert({key, value}); + bool add_frontend_attribute(absl::string_view key, absl::string_view value) { + auto it = mutable_rare()->frontend_attributes.mutable_map()->insert( + {std::string(key), std::string(value)}); return it.second; } @@ -2422,13 +2421,13 @@ class HloInstruction { std::string infeed_config() const; // Delegates to HloInfeedInstruction::set_infeed_config. - void set_infeed_config(const std::string& config); + void set_infeed_config(absl::string_view config); // Returns the config for the Outfeed instruction. const std::string& outfeed_config() const; // Delegates to HloOutfeedInstruction::set_outfeed_config. - void set_outfeed_config(const std::string& config); + void set_outfeed_config(absl::string_view config); // Returns the shape for the Outfeed instruction. const Shape& outfeed_shape() const; @@ -2752,6 +2751,16 @@ class HloInstruction { // HloInstruction. bool IsMarkedAsDead() const { return marked_as_dead_; } + // Implementation of DetachFromOperandsAndUsers that leaves the edges to + // instructions of the given computation (the parent, or null for none) in + // place and still marks this instruction cleaned up, so ~HloInstruction does + // not unlink them either. ~HloComputation uses it because unlinking an + // instruction scans every operand slot of each of its users, and the edges + // among instructions that die together are never read again. + // REQUIRES: every instruction of the given computation is deleted before any + // of those edges is read again. + void DetachFromOperandsAndUsersOutside(const HloComputation* computation); + // Set the unique id for this instruction to "id". Should only be called by // the instruction's parent computation to set an internal unique id that fits // in an int32_t. @@ -2974,16 +2983,15 @@ std::string SparsityConfigToString(const SparsityConfig& sparsity_config); std::string BlockScalingConfigToString( const BlockScalingConfig& block_scaling_config); -absl::StatusOr StringToRandomAlgorithm( - const std::string& name); +absl::StatusOr StringToRandomAlgorithm(absl::string_view name); absl::StatusOr StringToRandomDistribution( - const std::string& name); + absl::string_view name); absl::StatusOr StringToPrecision( - const std::string& name); + absl::string_view name); absl::StatusOr StringToShuffleMode( absl::string_view mode); absl::StatusOr StringToAlgorithm( - const std::string& name); + absl::string_view name); absl::StatusOr StringToResultAccuracy( absl::string_view name); absl::StatusOr StringToCustomCallSchedule( diff --git a/third_party/xla/xla/hlo/ir/hlo_module.h b/third_party/xla/xla/hlo/ir/hlo_module.h index 263db00361f89f..03d728da2d771f 100644 --- a/third_party/xla/xla/hlo/ir/hlo_module.h +++ b/third_party/xla/xla/hlo/ir/hlo_module.h @@ -724,9 +724,8 @@ class HloModule { CHECK(spmd_parameters_shardings_.has_value()); return *spmd_parameters_shardings_; } - void set_spmd_parameters_shardings( - const std::vector& shardings) { - spmd_parameters_shardings_ = shardings; + void set_spmd_parameters_shardings(absl::Span shardings) { + spmd_parameters_shardings_.emplace(shardings.begin(), shardings.end()); } // Checks if this config has the entry computation output's HLO sharding for @@ -856,8 +855,8 @@ class HloModule { } void set_profile_info( - const std::vector& profile_info) { - profile_info_list_ = profile_info; + absl::Span profile_info) { + profile_info_list_.assign(profile_info.begin(), profile_info.end()); } const std::vector& profile_info() const { diff --git a/third_party/xla/xla/hlo/ir/hlo_module_metadata.cc b/third_party/xla/xla/hlo/ir/hlo_module_metadata.cc index 2ede1b3ef6cf43..328b068c81407d 100644 --- a/third_party/xla/xla/hlo/ir/hlo_module_metadata.cc +++ b/third_party/xla/xla/hlo/ir/hlo_module_metadata.cc @@ -24,6 +24,7 @@ limitations under the License. #include "absl/log/log.h" #include "absl/status/status.h" #include "absl/status/status_macros.h" +#include "absl/strings/string_view.h" #include "xla/service/hlo.pb.h" #include "xla/service/metrics.pb.h" #include "xla/tsl/platform/env.h" @@ -107,7 +108,7 @@ absl::Status HloModuleMetadata::set_custom_metadata( return absl::OkStatus(); } -absl::Status HloModuleMetadata::set_key_value_metric(const std::string& key, +absl::Status HloModuleMetadata::set_key_value_metric(absl::string_view key, int64_t value) { ABSL_ASSIGN_OR_RETURN(HloPassMetadata * pass_metadata, GetCurrentHloPassMetadata()); diff --git a/third_party/xla/xla/hlo/ir/hlo_module_metadata.h b/third_party/xla/xla/hlo/ir/hlo_module_metadata.h index ec4d5be8926412..15944e6e5518f5 100644 --- a/third_party/xla/xla/hlo/ir/hlo_module_metadata.h +++ b/third_party/xla/xla/hlo/ir/hlo_module_metadata.h @@ -125,7 +125,7 @@ class HloModuleMetadata { const HloModuleMetadata& prepartitioning_metadata); // Setters for HloModuleMetadataProto. - void set_module_group_name(const std::string& name) { + void set_module_group_name(absl::string_view name) { module_metadata_.set_module_group_name(name); } void set_canonical_module_id(int64_t id) { @@ -137,7 +137,7 @@ class HloModuleMetadata { absl::Status set_custom_metadata(const ::tsl::protobuf::Message& message); // Adds a (key, value) pair metric if none was already set. Otherwise, it // updates the existing value. - absl::Status set_key_value_metric(const std::string& key, int64_t value); + absl::Status set_key_value_metric(absl::string_view key, int64_t value); absl::StatusOr current_pass_id() { ABSL_ASSIGN_OR_RETURN(HloPassMetadata * pass_metadata, @@ -146,23 +146,21 @@ class HloModuleMetadata { } // Setters for the current HloPassMetadata. - absl::Status set_current_pass_name(const std::string& pass_name) { + absl::Status set_current_pass_name(absl::string_view pass_name) { return MutateCurrentHloPassMetadata( - [&pass_name](HloPassMetadata* pass_metadata) { + [pass_name](HloPassMetadata* pass_metadata) { pass_metadata->set_pass_name(pass_name); }); } - absl::Status set_current_pass_pipeline_name( - const std::string& pipeline_name) { + absl::Status set_current_pass_pipeline_name(absl::string_view pipeline_name) { return MutateCurrentHloPassMetadata( - [&pipeline_name](HloPassMetadata* pass_metadata) { + [pipeline_name](HloPassMetadata* pass_metadata) { pass_metadata->set_pipeline_name(pipeline_name); }); } - absl::Status add_current_pass_dump_filename( - const std::string& dump_filename) { + absl::Status add_current_pass_dump_filename(absl::string_view dump_filename) { return MutateCurrentHloPassMetadata( - [&dump_filename](HloPassMetadata* pass_metadata) { + [dump_filename](HloPassMetadata* pass_metadata) { pass_metadata->add_dump_filenames(dump_filename); }); } diff --git a/third_party/xla/xla/hlo/ir/hlo_original_value.h b/third_party/xla/xla/hlo/ir/hlo_original_value.h index 20f2eadbb8b17e..7769db7599ba08 100644 --- a/third_party/xla/xla/hlo/ir/hlo_original_value.h +++ b/third_party/xla/xla/hlo/ir/hlo_original_value.h @@ -117,9 +117,9 @@ class OriginalValue { } bool IsEmpty() const { - return std::all_of( - tree().leaves().begin(), tree().leaves().end(), - [](const auto& pair) { return !pair.second.has_value(); }); + return absl::c_all_of(tree().leaves(), [](const auto& pair) { + return !pair.second.has_value(); + }); } bool IsCompatibleWith(const Shape& shape) const; diff --git a/third_party/xla/xla/hlo/ir/hlo_schedule.h b/third_party/xla/xla/hlo/ir/hlo_schedule.h index aef51e1e44791e..49548f962a7387 100644 --- a/third_party/xla/xla/hlo/ir/hlo_schedule.h +++ b/third_party/xla/xla/hlo/ir/hlo_schedule.h @@ -73,12 +73,10 @@ class HloInstructionSequence { // Removes the instruction from the sequence. void remove_instruction(HloInstruction* instruction) { - auto instruction_it = std::find(instruction_sequence_.begin(), - instruction_sequence_.end(), instruction); + auto instruction_it = absl::c_find(instruction_sequence_, instruction); if (instruction_it != instruction_sequence_.end() && instruction->parent() != nullptr) { - auto id_it = std::find(id_sequence_.begin(), id_sequence_.end(), - instruction->unique_id()); + auto id_it = absl::c_find(id_sequence_, instruction->unique_id()); instruction_sequence_.erase(instruction_it); id_sequence_.erase(id_it); } @@ -87,11 +85,8 @@ class HloInstructionSequence { // Replaces the old instruction with the new instruction in the sequence. void replace_instruction(HloInstruction* old_instruction, HloInstruction* new_instruction) { - auto instruction_it = - std::find(instruction_sequence_.begin(), instruction_sequence_.end(), - old_instruction); - auto id_it = std::find(id_sequence_.begin(), id_sequence_.end(), - old_instruction->unique_id()); + auto instruction_it = absl::c_find(instruction_sequence_, old_instruction); + auto id_it = absl::c_find(id_sequence_, old_instruction->unique_id()); CHECK(instruction_it != instruction_sequence_.end()) << "Do not find instruction id " << old_instruction->unique_id(); CHECK(id_it != id_sequence_.end()); diff --git a/third_party/xla/xla/hlo/parser/hlo_parser.cc b/third_party/xla/xla/hlo/parser/hlo_parser.cc index 268eedd5fbd1e9..b6e173aeb71a64 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser.cc @@ -634,7 +634,7 @@ class HloParserImpl : public HloParser { bool ParsePrimitiveType(PrimitiveType* result); bool ParseComparisonDirection(ComparisonDirection* result); bool ParseComparisonOrder(Comparison::Order* result); - bool ParseComparisonType(Comparison::Type* result); + bool ParseComparisonType(Comparison::Order* result); bool ParseFusionKind(HloInstruction::FusionKind* result); bool ParseRandomDistribution(RandomDistribution* result); bool ParseConvKind(ConvolutionKind* result); @@ -2870,7 +2870,7 @@ HloInstruction* HloParserImpl::CreateInstruction( // NOLINT case HloOpcode::kCompare: { optional direction; optional order; - optional type; + optional type; attrs["direction"] = {/*required=*/true, AttrTy::kComparisonDirection, &direction}; attrs["order"] = {/*required=*/false, AttrTy::kComparisonOrder, &order}; @@ -2897,8 +2897,7 @@ HloInstruction* HloParserImpl::CreateInstruction( // NOLINT } if (type.has_value()) { return builder->AddInstruction(HloInstruction::CreateCompare( - *shape, operands[0], operands[1], *direction, - Comparison::DefaultOrdering(*type))); + *shape, operands[0], operands[1], *direction, *type)); } return builder->AddInstruction(HloInstruction::CreateCompare( *shape, operands[0], operands[1], *direction)); @@ -6201,11 +6200,12 @@ bool HloParserImpl::ParseAttributeHelper( return true; } case AttrTy::kComparisonType: { - Comparison::Type result; + Comparison::Order result; if (!ParseComparisonType(&result)) { return false; } - static_cast*>(attr_out_ptr)->emplace(result); + static_cast*>(attr_out_ptr) + ->emplace(result); return true; } case AttrTy::kEnum: { @@ -8507,13 +8507,13 @@ bool HloParserImpl::ParseComparisonOrder(Comparison::Order* result) { return true; } -bool HloParserImpl::ParseComparisonType(Comparison::Type* result) { +bool HloParserImpl::ParseComparisonType(Comparison::Order* result) { VLOG(kDebugLevel) << "ParseComparisonType"; if (lexer_.GetKind() != TokKind::kIdent) { return TokenError("expects comparison type"); } std::string val = lexer_.GetStrVal(); - auto status_or_result = StringToComparisonType(val); + auto status_or_result = ComparisonTypeToOrder(val); if (!status_or_result.ok()) { return TokenError(StrFormat("expects comparison type but sees: %s", val)); } diff --git a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc index 6502d1f3d52d87..ddf4306437a436 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc @@ -3771,6 +3771,23 @@ ENTRY %entry(p0: f32[], p1: f32[]) -> pred[] { EXPECT_EQ(compare->order(), ComparisonOrder::kTotal); } +TEST_F(HloParserTest, CompareWithWeakOrder) { + const std::string original = R"(HloModule CompareWithWeakOrder +ENTRY %entry(p0: f32[], p1: f32[]) -> pred[] { + %p0 = f32[] parameter(0) + %p1 = f32[] parameter(1) + ROOT %cmp = pred[] compare(f32[] %p0, f32[] %p1), direction=LT, order=WEAK +})"; + auto result = ParseAndReturnVerifiedModule(original); + ASSERT_OK(result.status()); + const HloInstruction* root = + result.value()->entry_computation()->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kCompare); + const auto* compare = static_cast(root); + EXPECT_EQ(compare->direction(), ComparisonDirection::kLt); + EXPECT_EQ(compare->order(), ComparisonOrder::kWeak); +} + TEST_F(HloParserTest, CompareBothTypeAndOrderFails) { const std::string original = R"(HloModule CompareBothTypeAndOrderFails ENTRY %entry(p0: f32[], p1: f32[]) -> pred[] { diff --git a/third_party/xla/xla/hlo/pass/hlo_pass_interface.h b/third_party/xla/xla/hlo/pass/hlo_pass_interface.h index a9f4d1e65a84b9..1fcdf4058c6c6c 100644 --- a/third_party/xla/xla/hlo/pass/hlo_pass_interface.h +++ b/third_party/xla/xla/hlo/pass/hlo_pass_interface.h @@ -118,7 +118,7 @@ class HloPassInterface { // If an HloPassMetadata has previously been created, it adds a (key, value) // pair metric if none was already set or updates the existing value. // If an HloPassMetadata doesn't exist, it simply returns. - static void SetKVMetric(HloModule* module, const std::string& key, + static void SetKVMetric(HloModule* module, absl::string_view key, int64_t value) { auto status = module->metadata()->set_key_value_metric(key, value); if (!status.ok()) { diff --git a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc index f493932b44acbd..5cf8c5712809cc 100644 --- a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc +++ b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc @@ -48,8 +48,8 @@ namespace xla { namespace { -void RecordPassStartMetadata(HloModule& module, const std::string& pass_name, - const std::string& pipeline_name) { +void RecordPassStartMetadata(HloModule& module, absl::string_view pass_name, + absl::string_view pipeline_name) { module.metadata()->RecordPassStart(); // An HloPassMetadata was just created so absl::Status should always be OK. CHECK_OK(module.metadata()->set_current_pass_name(pass_name)); @@ -57,7 +57,7 @@ void RecordPassStartMetadata(HloModule& module, const std::string& pass_name, } absl::Status AttemptRecordPassEndMetadata(HloModule& module, - const std::string& pass_name, + absl::string_view pass_name, bool module_changed) { // Module id is set here instead of RecordPassStartMetadata because it may // change in the middle of the pass, and we want the final id. @@ -69,7 +69,7 @@ absl::Status AttemptRecordPassEndMetadata(HloModule& module, return absl::OkStatus(); } -void RecordPassEndMetadata(HloModule& module, const std::string& pass_name, +void RecordPassEndMetadata(HloModule& module, absl::string_view pass_name, bool module_changed) { absl::Status status = AttemptRecordPassEndMetadata(module, pass_name, module_changed); diff --git a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.h b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.h index 8d53cd9740a770..aa54edc8a4e69d 100644 --- a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.h +++ b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.h @@ -37,7 +37,7 @@ namespace xla { // Pipeline of HLO passes. class HloPassPipeline : public HloPassInterface { public: - explicit HloPassPipeline(const std::string& name, + explicit HloPassPipeline(absl::string_view name, CompilationStats* compilation_stats = nullptr) : name_(name), compilation_stats_(compilation_stats) { if (compilation_stats == nullptr) { diff --git a/third_party/xla/xla/hlo/transforms/BUILD b/third_party/xla/xla/hlo/transforms/BUILD index f89ee99a223c8d..1be14d085b13c6 100644 --- a/third_party/xla/xla/hlo/transforms/BUILD +++ b/third_party/xla/xla/hlo/transforms/BUILD @@ -340,6 +340,7 @@ cc_library( "//xla/service:host_offload_utils", "//xla/service:memory_annotations_hdr", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/cleanup", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/container:inlined_vector", diff --git a/third_party/xla/xla/hlo/transforms/expanders/BUILD b/third_party/xla/xla/hlo/transforms/expanders/BUILD index 03267ff432e6a9..a3e41f7e6b429d 100644 --- a/third_party/xla/xla/hlo/transforms/expanders/BUILD +++ b/third_party/xla/xla/hlo/transforms/expanders/BUILD @@ -85,6 +85,25 @@ cc_library( ], ) +xla_cc_test( + name = "comparison_expander_test", + srcs = ["comparison_expander_test.cc"], + deps = [ + ":comparison_expander", + "//xla:comparison_util", + "//xla:literal", + "//xla:literal_util", + "//xla/hlo/evaluator:hlo_evaluator", + "//xla/hlo/ir:hlo", + "//xla/hlo/testlib:hlo_hardware_independent_test_base", + "//xla/tests:literal_test_util", + "//xla/tsl/platform:statusor", + "@com_google_absl//absl/strings", + "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "cholesky_expander", srcs = ["cholesky_expander.cc"], diff --git a/third_party/xla/xla/hlo/transforms/expanders/comparison_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/comparison_expander.cc index 7218d076aab0fe..b6c60daec69e5c 100644 --- a/third_party/xla/xla/hlo/transforms/expanders/comparison_expander.cc +++ b/third_party/xla/xla/hlo/transforms/expanders/comparison_expander.cc @@ -78,12 +78,64 @@ HloInstruction* BitcastConvertFloatingPointToIntegral( is_negative, flipped_value, signed_value)); } +HloInstruction* ExpandWeakOrderFloatComparison(HloComputation* computation, + const Shape& pred_shape, + HloInstruction* lhs, + HloInstruction* rhs, + ComparisonDirection direction, + PrimitiveType compare_type) { + auto partial_cmp = [&](HloInstruction* a, HloInstruction* b, + ComparisonDirection dir) { + return computation->AddInstruction(HloInstruction::CreateCompare( + pred_shape, a, b, dir, Comparison::Order::kPartial)); + }; + if (!primitive_util::HasNaN(compare_type)) { + return partial_cmp(lhs, rhs, direction); + } + auto is_nan = [&](HloInstruction* a) { + return partial_cmp(a, a, ComparisonDirection::kNe); + }; + auto pred_or = [&](HloInstruction* a, HloInstruction* b) { + return computation->AddInstruction( + HloInstruction::CreateBinary(pred_shape, HloOpcode::kOr, a, b)); + }; + auto pred_and = [&](HloInstruction* a, HloInstruction* b) { + return computation->AddInstruction( + HloInstruction::CreateBinary(pred_shape, HloOpcode::kAnd, a, b)); + }; + auto pred_not = [&](HloInstruction* a) { + return computation->AddInstruction( + HloInstruction::CreateUnary(pred_shape, HloOpcode::kNot, a)); + }; + + switch (direction) { + case ComparisonDirection::kEq: + return pred_or(partial_cmp(lhs, rhs, ComparisonDirection::kEq), + pred_and(is_nan(lhs), is_nan(rhs))); + case ComparisonDirection::kNe: + return pred_and(partial_cmp(lhs, rhs, ComparisonDirection::kNe), + pred_not(pred_and(is_nan(lhs), is_nan(rhs)))); + case ComparisonDirection::kLe: + return pred_or(partial_cmp(lhs, rhs, ComparisonDirection::kLe), + is_nan(rhs)); + case ComparisonDirection::kGe: + return pred_or(partial_cmp(lhs, rhs, ComparisonDirection::kGe), + is_nan(lhs)); + case ComparisonDirection::kLt: + return pred_not(pred_or(partial_cmp(lhs, rhs, ComparisonDirection::kGe), + is_nan(lhs))); + case ComparisonDirection::kGt: + return pred_not(pred_or(partial_cmp(lhs, rhs, ComparisonDirection::kLe), + is_nan(rhs))); + } +} + bool ComparisonExpander::InstructionMatchesPattern( HloInstruction* instruction) { if (HloCompareInstruction* compare = DynCast(instruction)) { HloInstruction* lhs = instruction->operands()[0]; - if (compare->order() == Comparison::Order::kTotal && + if (compare->order() != Comparison::Order::kPartial && primitive_util::IsFloatingPointType(lhs->shape().element_type())) { return true; } @@ -96,8 +148,11 @@ absl::StatusOr ComparisonExpander::ExpandInstruction( CHECK_EQ(instruction->opcode(), HloOpcode::kCompare); HloCompareInstruction* compare = static_cast(instruction); - CHECK(compare->order() == Comparison::Order::kTotal) - << ComparisonOrderToString(compare->order()); + if (compare->order() != Comparison::Order::kTotal && + compare->order() != Comparison::Order::kWeak) { + return InvalidArgument("Unsupported comparison order: %s", + ComparisonOrderToString(compare->order())); + } HloComputation* computation = instruction->parent(); HloInstruction* lhs = instruction->operands()[0]; HloInstruction* rhs = instruction->operands()[1]; @@ -118,6 +173,15 @@ absl::StatusOr ComparisonExpander::ExpandInstruction( ShapeUtil::ChangeElementType(rhs->shape(), compare_type), rhs)); } + if (compare->order() == Comparison::Order::kWeak) { + HloInstruction* new_compare = + ExpandWeakOrderFloatComparison(computation, instruction->shape(), lhs, + rhs, compare->direction(), compare_type); + VLOG(2) << "New comparison instruction for weak order:" + << new_compare->ToString(); + return new_compare; + } + if (compare_type != F8E8M0FNU) { int64_t bit_width = primitive_util::BitWidth(lhs->shape().element_type()); PrimitiveType signed_type = diff --git a/third_party/xla/xla/hlo/transforms/expanders/comparison_expander_test.cc b/third_party/xla/xla/hlo/transforms/expanders/comparison_expander_test.cc new file mode 100644 index 00000000000000..c901a0d971bed0 --- /dev/null +++ b/third_party/xla/xla/hlo/transforms/expanders/comparison_expander_test.cc @@ -0,0 +1,121 @@ +/* Copyright 2026 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xla/hlo/transforms/expanders/comparison_expander.h" + +#include +#include +#include +#include +#include + +#include +#include "absl/strings/substitute.h" +#include "xla/comparison_util.h" +#include "xla/hlo/evaluator/hlo_evaluator.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" +#include "xla/literal.h" +#include "xla/literal_util.h" +#include "xla/tests/literal_test_util.h" +#include "xla/tsl/platform/statusor.h" + +namespace xla { +namespace { + +using ComparisonExpanderTest = HloHardwareIndependentTestBase; + +TEST_F(ComparisonExpanderTest, ExpandWeakOrderFloatComparisonMatchesEvaluator) { + const float inf = std::numeric_limits::infinity(); + const float qnan = std::numeric_limits::quiet_NaN(); + const float neg_qnan = std::copysign(qnan, -1.0f); + const std::vector values = {-inf, -1.0f, -0.0f, +0.0f, + +1.0f, +inf, neg_qnan, qnan}; + std::vector lhs_vec; + std::vector rhs_vec; + for (float a : values) { + for (float b : values) { + lhs_vec.push_back(a); + rhs_vec.push_back(b); + } + } + Literal lhs_lit = LiteralUtil::CreateR1(lhs_vec); + Literal rhs_lit = LiteralUtil::CreateR1(rhs_vec); + + for (ComparisonDirection dir : + {ComparisonDirection::kEq, ComparisonDirection::kNe, + ComparisonDirection::kLt, ComparisonDirection::kLe, + ComparisonDirection::kGt, ComparisonDirection::kGe}) { + std::string hlo_text = absl::Substitute( + R"( +HloModule test +ENTRY main { + p0 = f32[64] parameter(0) + p1 = f32[64] parameter(1) + ROOT cmp = pred[64] compare(p0, p1), direction=$0, order=WEAK +} +)", + ComparisonDirectionToString(dir)); + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_text)); + + HloEvaluator evaluator; + TF_ASSERT_OK_AND_ASSIGN(Literal unexpanded_result, + evaluator.Evaluate(*module, {&lhs_lit, &rhs_lit})); + + ComparisonExpander expander; + TF_ASSERT_OK_AND_ASSIGN(bool changed, expander.Run(module.get())); + EXPECT_TRUE(changed); + + // Verify no WEAK comparisons remain. + for (const HloInstruction* inst : + module->entry_computation()->instructions()) { + if (inst->opcode() == HloOpcode::kCompare) { + EXPECT_EQ(inst->comparison_order(), ComparisonOrder::kPartial); + } + } + + TF_ASSERT_OK_AND_ASSIGN(Literal expanded_result, + evaluator.Evaluate(*module, {&lhs_lit, &rhs_lit})); + EXPECT_TRUE(LiteralTestUtil::Equal(unexpanded_result, expanded_result)) + << "Mismatch for direction " << ComparisonDirectionToString(dir); + } +} + +TEST_F(ComparisonExpanderTest, ExpandWeakOrderFloatWithoutNaN) { + constexpr absl::string_view kHloText = R"( +HloModule test +ENTRY main { + p0 = f4e2m1fn[8] parameter(0) + p1 = f4e2m1fn[8] parameter(1) + ROOT cmp = pred[8] compare(p0, p1), direction=LT, order=WEAK +} +)"; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHloText)); + ComparisonExpander expander; + TF_ASSERT_OK_AND_ASSIGN(bool changed, expander.Run(module.get())); + EXPECT_TRUE(changed); + + const HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kCompare); + EXPECT_EQ(root->comparison_direction(), ComparisonDirection::kLt); + EXPECT_EQ(root->comparison_order(), ComparisonOrder::kPartial); +} + +} // namespace +} // namespace xla diff --git a/third_party/xla/xla/hlo/transforms/host_offload_legalize.cc b/third_party/xla/xla/hlo/transforms/host_offload_legalize.cc index 4b8ffac96f0b68..2f288a0fecb1a2 100644 --- a/third_party/xla/xla/hlo/transforms/host_offload_legalize.cc +++ b/third_party/xla/xla/hlo/transforms/host_offload_legalize.cc @@ -738,7 +738,8 @@ absl::Status MoveCopyDown( // Returns true if the copy should be moved. A copy can be moved if there is // always a place for it after being moved back to device. -bool ShouldMoveCopyDown(HloInstruction* copy_to_move) { +bool ShouldMoveCopyDown(HloInstruction* copy_to_move, + const CallGraph& call_graph) { std::queue queue; queue.push(host_offload_utils::InstructionAndShapeIndex(copy_to_move)); while (!queue.empty()) { @@ -754,7 +755,7 @@ bool ShouldMoveCopyDown(HloInstruction* copy_to_move) { // Push successors onto the queue to be visited. absl::StatusOr> - successors = host_offload_utils::GetSuccessors(current); + successors = host_offload_utils::GetSuccessors(current, call_graph); if (!successors.ok()) { return false; } @@ -943,7 +944,7 @@ absl::StatusOr ProcessAnnotationForCopyMovement( for (auto it = copies_to_move.rbegin(); it != copies_to_move.rend(); ++it) { InstructionAndIndex& copy_to_move_and_index = *it; HloInstruction* copy_to_move = copy_to_move_and_index.instruction; - if (ShouldMoveCopyDown(copy_to_move)) { + if (ShouldMoveCopyDown(copy_to_move, *call_graph)) { ABSL_RETURN_IF_ERROR(MoveCopyDown(copy_to_move_and_index, call_graph, processed_annotations, to_remove)); changed = true; diff --git a/third_party/xla/xla/hlo/transforms/host_offload_legalize_test.cc b/third_party/xla/xla/hlo/transforms/host_offload_legalize_test.cc index a31a642cc65b2d..b460ad64f3e66c 100644 --- a/third_party/xla/xla/hlo/transforms/host_offload_legalize_test.cc +++ b/third_party/xla/xla/hlo/transforms/host_offload_legalize_test.cc @@ -15,6 +15,8 @@ limitations under the License. #include "xla/hlo/transforms/host_offload_legalize.h" +#include + #include #include #include "absl/status/status.h" diff --git a/third_party/xla/xla/hlo/transforms/host_offloader.cc b/third_party/xla/xla/hlo/transforms/host_offloader.cc index 319a87c6b9945a..9a7f45cf83a487 100644 --- a/third_party/xla/xla/hlo/transforms/host_offloader.cc +++ b/third_party/xla/xla/hlo/transforms/host_offloader.cc @@ -24,6 +24,7 @@ limitations under the License. #include #include "absl/algorithm/container.h" +#include "absl/cleanup/cleanup.h" #include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_set.h" #include "absl/container/inlined_vector.h" @@ -184,6 +185,8 @@ absl::StatusOr HostOffloader::WalkDownHostMemoryOffloadPaths( // std::vector move_to_host_dynamic_update_slices; HloInstruction* starting_instruction = starting_instruction_and_index.instruction; + const CallGraph& call_graph = + GetOrBuildCallGraph(starting_instruction->GetModule()); std::queue queue; absl::flat_hash_map previous; @@ -235,10 +238,8 @@ absl::StatusOr HostOffloader::WalkDownHostMemoryOffloadPaths( // When setting the memory space of a parameter, also set the memory // space of the call site of the computation with this parameter if that // caller is an async-start. - std::unique_ptr call_graph = - CallGraph::Build(instruction->GetModule()); std::vector callers = - call_graph->GetComputationCallers(instruction->parent()); + call_graph.GetComputationCallers(instruction->parent()); for (HloInstruction* caller : callers) { if (caller->opcode() == HloOpcode::kAsyncStart) { ShapeIndex tmp_index = instruction_and_shape_index.shape_index; @@ -371,9 +372,9 @@ absl::StatusOr HostOffloader::WalkDownHostMemoryOffloadPaths( starting_instruction->name()); } // Push successors onto the queue to be visited. - ABSL_ASSIGN_OR_RETURN( - const std::vector successors, - host_offload_utils::GetSuccessors(instruction_and_shape_index)); + ABSL_ASSIGN_OR_RETURN(const std::vector successors, + host_offload_utils::GetSuccessors( + instruction_and_shape_index, call_graph)); for (const InstructionAndShapeIndex& successor : successors) { if (VLOG_IS_ON(1)) { previous.emplace(successor, instruction_and_shape_index); @@ -399,6 +400,7 @@ absl::StatusOr HostOffloader::WalkDownHostMemoryOffloadPaths( if (insert_copy_before) { const std::vector predecessors = host_offload_utils::GetPredecessors(starting_instruction_and_index, + call_graph, operand_index.value_or(0)); CHECK_EQ(predecessors.size(), 1); ABSL_ASSIGN_OR_RETURN( @@ -621,10 +623,9 @@ absl::StatusOr HostOffloader::InsertCopyBetween( // To insert a copy between an instruction and a parameter means we actually // want to insert a copy between the instruction and the call site of the // computation with this parameter. - std::unique_ptr call_graph = - CallGraph::Build(after_instruction->GetModule()); - auto callers = - call_graph->GetComputationCallers(after_instruction->parent()); + const std::vector callers = + GetOrBuildCallGraph(after_instruction->GetModule()) + .GetComputationCallers(after_instruction->parent()); for (HloInstruction* caller : callers) { const auto indices = caller->OperandIndices(before_instruction_and_index.instruction); @@ -700,10 +701,12 @@ HostOffloader::GetStartingInstructions( // 2. Does "normal" memory offloading. std::vector result; std::queue queue; + const CallGraph& call_graph = + GetOrBuildCallGraph(custom_call_instruction->GetModule()); ABSL_ASSIGN_OR_RETURN( const std::vector successors_of_custom_call, host_offload_utils::GetSuccessors( - InstructionAndShapeIndex(custom_call_instruction))); + InstructionAndShapeIndex(custom_call_instruction), call_graph)); for (const InstructionAndShapeIndex& successor : successors_of_custom_call) { queue.push(successor); } @@ -722,8 +725,9 @@ HostOffloader::GetStartingInstructions( continue; } // Is a logical bitcast/reshape, we won't offload this yet. - ABSL_ASSIGN_OR_RETURN(const std::vector successors, - host_offload_utils::GetSuccessors(instruction_and_shape)); + ABSL_ASSIGN_OR_RETURN( + const std::vector successors, + host_offload_utils::GetSuccessors(instruction_and_shape, call_graph)); for (const InstructionAndShapeIndex& successor : successors) { queue.push(successor); } @@ -739,9 +743,11 @@ absl::StatusOr HostOffloader::SliceLeadsToMoveToDeviceCustomCall( slice->opcode() == HloOpcode::kSlice) << "This function must only be called with a slice or dynamic slice."; std::queue queue; + const CallGraph& call_graph = GetOrBuildCallGraph(slice->GetModule()); ABSL_ASSIGN_OR_RETURN( const std::vector successors_of_slice, - host_offload_utils::GetSuccessors(InstructionAndShapeIndex(slice))); + host_offload_utils::GetSuccessors(InstructionAndShapeIndex(slice), + call_graph)); for (const InstructionAndShapeIndex& successor : successors_of_slice) { queue.push(successor); } @@ -764,8 +770,9 @@ absl::StatusOr HostOffloader::SliceLeadsToMoveToDeviceCustomCall( HloOpcodeString(slice->opcode()), slice->name(), slice->name()); return false; } - ABSL_ASSIGN_OR_RETURN(const std::vector successors, - host_offload_utils::GetSuccessors(instruction_and_shape)); + ABSL_ASSIGN_OR_RETURN( + const std::vector successors, + host_offload_utils::GetSuccessors(instruction_and_shape, call_graph)); for (const InstructionAndShapeIndex& successor : successors) { queue.push(successor); } @@ -787,6 +794,8 @@ absl::Status HostOffloader::CreateAllocateBufferForDynamicUpdateSlice( // Walk the graph up. We expect to find a broadcast. Also, while walking up // the graph, set host memory space on everything between the AllocateBuffer // and the DynamicUpdateSlice. + const CallGraph& call_graph = + GetOrBuildCallGraph(dynamic_update_slice->GetModule()); std::queue queue; queue.push(InstructionAndShapeIndex(dynamic_update_slice)); std::optional previous_instruction_and_shape = @@ -830,10 +839,8 @@ absl::Status HostOffloader::CreateAllocateBufferForDynamicUpdateSlice( // If this is a parameter of a while_body, we also need to find the // matching parameter in the while_condition and set the memory spaces // there. - std::unique_ptr call_graph = - CallGraph::Build(instruction->GetModule()); const std::vector callers = - call_graph->GetComputationCallers(instruction->parent()); + call_graph.GetComputationCallers(instruction->parent()); for (HloInstruction* caller : callers) { if (caller->opcode() == HloOpcode::kWhile) { // This parameter belongs to a while. @@ -877,8 +884,8 @@ absl::Status HostOffloader::CreateAllocateBufferForDynamicUpdateSlice( Layout::kHostMemorySpace); ABSL_ASSIGN_OR_RETURN( const std::vector successors, - host_offload_utils::GetSuccessors( - nested_instruction_and_shape)); + host_offload_utils::GetSuccessors(nested_instruction_and_shape, + call_graph)); for (const InstructionAndShapeIndex& successor : successors) { nested_queue.push(successor); } @@ -902,9 +909,10 @@ absl::Status HostOffloader::CreateAllocateBufferForDynamicUpdateSlice( // instruction that we're walking up the graph from. CHECK(previous_instruction_and_shape.has_value()) << "We expect to have a previous instruction at this point."; - ABSL_ASSIGN_OR_RETURN(std::vector successors, - host_offload_utils::GetSuccessors( - InstructionAndShapeIndex(instruction, shape_index))); + ABSL_ASSIGN_OR_RETURN( + std::vector successors, + host_offload_utils::GetSuccessors( + InstructionAndShapeIndex(instruction, shape_index), call_graph)); for (const InstructionAndShapeIndex& successor : successors) { if (ShapeUtil::GetSubshape(successor.instruction->shape(), successor.shape_index) @@ -958,7 +966,7 @@ absl::Status HostOffloader::CreateAllocateBufferForDynamicUpdateSlice( return absl::OkStatus(); } const std::vector predecessors = - host_offload_utils::GetPredecessors(instruction_and_shape); + host_offload_utils::GetPredecessors(instruction_and_shape, call_graph); for (const InstructionAndShapeIndex& predecessor : predecessors) { HloInstruction* predecessor_instruction = predecessor.instruction; if (predecessor_instruction->opcode() == HloOpcode::kBroadcast) { @@ -1257,6 +1265,7 @@ absl::StatusOr HostOffloader::HandleRedundantCopiesBackToHost( const Shape& entry_computation_shape = module->entry_computation_layout().result_layout().shape(); + const CallGraph& call_graph = GetOrBuildCallGraph(module); // We collect all usages per output index, stopping at any non host // instruction. @@ -1320,11 +1329,12 @@ absl::StatusOr HostOffloader::HandleRedundantCopiesBackToHost( } } - ABSL_ASSIGN_OR_RETURN( - std::vector successors, - host_offload_utils::GetSuccessors(InstructionAndShapeIndex( - instruction_and_shape_index.instruction, - instruction_and_shape_index.shape_index))); + ABSL_ASSIGN_OR_RETURN(std::vector successors, + host_offload_utils::GetSuccessors( + InstructionAndShapeIndex( + instruction_and_shape_index.instruction, + instruction_and_shape_index.shape_index), + call_graph)); // Check if any of the successors needs to be on device. for (InstructionAndShapeIndex& successor : successors) { @@ -1509,9 +1519,19 @@ absl::StatusOr HostOffloader::HandlePallasKernels(HloModule* module) { return changed; } +const CallGraph& HostOffloader::GetOrBuildCallGraph(const HloModule* module) { + if (call_graph_ == nullptr) { + call_graph_ = CallGraph::Build(module); + } + return *call_graph_; +} + absl::StatusOr HostOffloader::RunImpl( HloModule* module, const absl::flat_hash_set& execution_threads) { + // Dropped on every exit so that no run sees another run's graph. + absl::Cleanup drop_call_graph = [this] { call_graph_.reset(); }; + // Start by removing all host memory space from all shapes. Host memory space // might have been set by other passes, however, this pass is the one which is // solely responsible for the propagation of host memory space throughout the diff --git a/third_party/xla/xla/hlo/transforms/host_offloader.h b/third_party/xla/xla/hlo/transforms/host_offloader.h index 81c56fa1ba7980..6b8365a472a01a 100644 --- a/third_party/xla/xla/hlo/transforms/host_offloader.h +++ b/third_party/xla/xla/hlo/transforms/host_offloader.h @@ -16,6 +16,7 @@ #define XLA_HLO_TRANSFORMS_HOST_OFFLOADER_H_ #include +#include #include #include #include @@ -28,6 +29,7 @@ #include "xla/hlo/analysis/alias_info.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/pass/hlo_pass_interface.h" +#include "xla/service/call_graph.h" #include "xla/service/host_offload_utils.h" namespace xla { @@ -89,6 +91,13 @@ class HostOffloader : public HloModulePass { already_inserted_copy_before_; const AliasInfo* alias_info_; + // Call graph of the module of the current run, built on first use and + // dropped when RunImpl returns. One graph serves the whole run: until the + // final HloCSE, the pass only adds and removes instructions that call no + // computation, so the call sites never change while the graph is in use. + std::unique_ptr call_graph_; + const CallGraph& GetOrBuildCallGraph(const HloModule* module); + // DynamicUpdateSlices are a bit special because they are the only op which // has multiple operands that host memory offloading supports. As a result, // different memory propagation paths can pass through the same diff --git a/third_party/xla/xla/hlo/transforms/host_offloader_test.cc b/third_party/xla/xla/hlo/transforms/host_offloader_test.cc index aba75f03a9c4ac..fda51e9fa63434 100644 --- a/third_party/xla/xla/hlo/transforms/host_offloader_test.cc +++ b/third_party/xla/xla/hlo/transforms/host_offloader_test.cc @@ -5273,5 +5273,89 @@ ENTRY main { } } +TEST_F(HostOffloaderTest, SecondRunUsesItsOwnCallGraph) { + // Both modules cross a computation boundary, so each run queries the call + // graph the pass keeps for a run; module_a offloads through an annotation, + // module_b through a host memory entry parameter. The second run must see + // its own module's graph, not the graph of the first module. The call graph + // is the only per run state the pass drops between runs, so module_b is + // shaped to read none of the rest: no annotation, no DynamicUpdateSlice, no + // copy insertion. + const std::string& hlo_string_a = R"( +HloModule module_a + +callee { + callee_param = f32[2048] parameter(0) + ROOT callee_root = (f32[2048]) tuple(callee_param) +} + +ENTRY main { + data_param = f32[2048] parameter(0) + offload = f32[2048] custom-call(data_param), custom_call_target="MoveToHost" + call = (f32[2048]) call(offload), to_apply=callee + gte = f32[2048] get-tuple-element(call), index=0 + ROOT load = f32[2048] custom-call(gte), custom_call_target="MoveToDevice" +} +)"; + const std::string& hlo_string_b = R"( +HloModule module_b, entry_computation_layout={(f32[2048]{0:S(5)})->f32[2048]{0}} + +while_condition { + cond_param = (s32[], f32[2048]) parameter(0) + i = s32[] get-tuple-element(cond_param), index=0 + limit = s32[] constant(4) + ROOT lt = pred[] compare(i, limit), direction=LT +} + +while_body { + body_param = (s32[], f32[2048]) parameter(0) + i = s32[] get-tuple-element(body_param), index=0 + data = f32[2048] get-tuple-element(body_param), index=1 + one = s32[] constant(1) + next = s32[] add(i, one) + ROOT body_root = (s32[], f32[2048]) tuple(next, data) +} + +ENTRY main { + param = f32[2048] parameter(0) + zero = s32[] constant(0) + init = (s32[], f32[2048]) tuple(zero, param) + loop = (s32[], f32[2048]) while(init), condition=while_condition, body=while_body + gte = f32[2048] get-tuple-element(loop), index=1 + ROOT load = f32[2048] custom-call(gte), custom_call_target="MoveToDevice" +} +)"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module_a, + ParseAndReturnVerifiedModule(hlo_string_a)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module_b, + ParseAndReturnVerifiedModule(hlo_string_b)); + HostOffloader host_offloader(&alias_info_); + + ASSERT_OK_AND_ASSIGN(bool changed_a, host_offloader.Run(module_a.get())); + EXPECT_TRUE(changed_a); + EXPECT_FALSE(HaveRemainingOffloadAnnotations(module_a.get())); + HloInstruction* copy_to_host; + HloInstruction* call; + ASSERT_THAT(module_a->entry_computation()->root_instruction(), + GmockMatch(m::Copy(m::GetTupleElement( + m::Call(&call, m::Copy(©_to_host, m::Parameter(0))))))); + TestShapeHasMemorySpace(copy_to_host->shape(), Layout::kHostMemorySpace); + TestShapeHasMemorySpace(call->shape().tuple_shapes(0), + Layout::kHostMemorySpace); + + ASSERT_OK_AND_ASSIGN(bool changed_b, host_offloader.Run(module_b.get())); + EXPECT_TRUE(changed_b); + EXPECT_FALSE(HaveRemainingOffloadAnnotations(module_b.get())); + HloInstruction* loop; + ASSERT_THAT(module_b->entry_computation()->root_instruction(), + GmockMatch(m::Copy(m::GetTupleElement(m::While(&loop))))); + TestShapeHasMemorySpace(loop->shape().tuple_shapes(1), + Layout::kHostMemorySpace); + TestShapeHasMemorySpace( + loop->while_body()->parameter_instruction(0)->shape().tuple_shapes(1), + Layout::kHostMemorySpace); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD index 7fd054c9e50849..e111be91e2c82f 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD +++ b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD @@ -438,6 +438,7 @@ cc_library( "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/types:span", "@tsl//tsl/platform:numbers", "@tsl//tsl/profiler/lib:scoped_annotation", ], @@ -454,6 +455,7 @@ xla_cc_test( "//xla:xla_data_proto_cc", "//xla/hlo/analysis:alias_info", "//xla/hlo/analysis:hlo_alias_analysis", + "//xla/hlo/analysis:hlo_dataflow_analysis", "//xla/hlo/analysis:hlo_ordering", "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", @@ -466,8 +468,11 @@ xla_cc_test( "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:nullability", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status:status_matchers", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/strings:string_view", "@com_google_absl//absl/types:span", "@com_google_googletest//:gtest_main", diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc index 303b102cd71dde..2193ddb2c91c6e 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc @@ -535,6 +535,15 @@ int64_t GetReduceFlops(const HloInstruction* reduce) { return ShapeUtil::ElementsIn(reduce->shape()) * (reduce_product - 1); } +// Returns true if any edge padding is negative, i.e. the pad crops its operand. +bool HasNegativePadding(const PaddingConfig& config) { + return absl::c_any_of(config.dimensions(), + [](const PaddingConfig::PaddingConfigDimension& dim) { + return dim.edge_padding_low() < 0 || + dim.edge_padding_high() < 0; + }); +} + } // namespace bool AlgebraicSimplifierVisitor::IsNonNegative( @@ -9648,6 +9657,12 @@ absl::Status AlgebraicSimplifierVisitor::HandleReduceWindow( return absl::OkStatus(); } + // A cropping pad does not compose with the window padding by addition. + if (HasNegativePadding(pad_config)) { + VLOG(10) << "Not folding negative pad into reduce-window."; + return absl::OkStatus(); + } + // If reduce_window already has padding, the pad value of the pad op and the // init value of reduce_window must match to allow folding the pad. const HloInstruction* pad_value = pad->operand(1); @@ -10419,6 +10434,11 @@ absl::StatusOr AlgebraicSimplifierVisitor::FoldConvInputPad( } } + // A cropping pad does not compose with the window padding by addition. + if (HasNegativePadding(padding)) { + return false; + } + // Compute the window which is the result of merging the kPad and the // convolution's existing window. Window new_window = window; diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc index bf5cc9ea7f9b67..8af4270bb1c1ba 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc @@ -6581,6 +6581,29 @@ TEST_P(ConvInputPaddingTest, DoTest) { } } +// Test that a pad with negative (cropping) edge padding is not folded into +// the convolution's window padding: a crop followed by window padding on the +// same side is not the sum of the two amounts. +TEST_F(AlgebraicSimplifierTest, DoNotFoldNegativePadIntoConvolution) { + const std::string& hlo_string = R"( +HloModule test +ENTRY entry { + input = f32[1,1,4,1] parameter(0) + filter = f32[1,2,1,1] parameter(1) + zero = f32[] constant(0) + pad = f32[1,1,3,1] pad(input, zero), padding=0_0x0_0x-1_0x0_0 + ROOT conv = f32[1,1,3,1] convolution(pad, filter), window={size=1x2 pad=0_0x1_0}, dim_labels=b01f_01io->b01f +} +)"; + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); + // Keep the negative pad instead of rewriting it into pad + slice. + AlgebraicSimplifierOptions options = default_options_; + options.set_enable_negative_padding_replacement(false); + AlgebraicSimplifier simplifier(options); + EXPECT_FALSE(RunHloPass(&simplifier, module.get()).value()); +} + // ConvFilterPaddingTest (and its one associated TEST_P) checks that a // computation that does // @@ -7352,6 +7375,33 @@ ENTRY entry { EXPECT_EQ(root->window().dimensions(3).padding_high(), 106); } +// Test that a pad with negative (cropping) edge padding is not folded into +// ReduceWindow: a crop followed by window padding on the same side is not the +// sum of the two amounts. +TEST_F(AlgebraicSimplifierTest, DoNotFoldNegativePadIntoReduceWindow) { + const std::string& hlo_string = R"( +HloModule test +fn { + p0 = f32[] parameter(0) + p1 = f32[] parameter(1) + ROOT add = f32[] add(p0, p1) +} +ENTRY entry { + param = f32[1,2] parameter(0) + const = f32[] constant(0) + pad = f32[1,1] pad(param, const), padding=0_0x-1_0 + ROOT r = f32[1,1] reduce-window(pad, const), to_apply=fn, window={size=1x2 pad=0_0x1_0} +} +)"; + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); + // Keep the negative pad instead of rewriting it into pad + slice. + AlgebraicSimplifierOptions options = default_options_; + options.set_enable_negative_padding_replacement(false); + AlgebraicSimplifier simplifier(options); + EXPECT_FALSE(RunHloPass(&simplifier, module.get()).value()); +} + // Test that ReduceWindow(Convert(Pad(op, x)), y) can simplify to // ReduceWindow(Convert(op), x). TEST_F(AlgebraicSimplifierTest, FoldConvertedPadIntoReduceWindow) { diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc index 7d0f805badb84c..4fcbc0e3ccfb98 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_memory_scheduler.cc @@ -37,6 +37,7 @@ limitations under the License. #include "absl/status/statusor.h" #include "absl/strings/str_format.h" #include "absl/strings/string_view.h" +#include "absl/types/span.h" #include "xla/hlo/analysis/alias_info.h" #include "xla/hlo/analysis/hlo_alias_analysis.h" #include "xla/hlo/analysis/hlo_dataflow_analysis.h" @@ -111,6 +112,10 @@ class ListScheduler { instruction.opcode() == HloOpcode::kConstant; } + // The UseRanges in buffer_uses_ refer to this object's value_index_pool_. + ListScheduler(const ListScheduler&) = delete; + ListScheduler& operator=(const ListScheduler&) = delete; + private: // The scheduling priority of an instruction is first the number of bytes // freed by scheduling the instruction, and second (tie-breaker) by the number @@ -119,6 +124,23 @@ class ListScheduler { // comparison operators. using Priority = std::pair; + // Index of an HloValue in the per computation tables values_ and + // unscheduled_use_count_. + using ValueIndex = int32_t; + + // A count of unscheduled uses never exceeds the instruction count plus one + // (the live out use). + using UseCount = int32_t; + + // The buffers an instruction uses, as a range of value_index_pool_. The + // first accounted of them are the buffers whose memory the scheduling + // heuristic accounts for; the rest belong to parameters and constants. + struct UseRange { + int64_t begin = 0; + int64_t size = 0; + int64_t accounted = 0; + }; + ListScheduler(HloComputation* computation, const HloAliasAnalysis& alias_analysis, const BufferValue::SizeFunction* absl_nonnull size_function) @@ -126,84 +148,139 @@ class ListScheduler { size_function_(ABSL_DIE_IF_NULL(size_function)) { const HloDataflowAnalysis& dataflow_analysis = alias_analysis.dataflow_analysis(); + CHECK_LT(computation->instruction_count(), + std::numeric_limits::max()); + + // Assigns a dense index to each HloValue the first time it is seen. + absl::flat_hash_map value_indices; + value_indices.reserve(computation->instruction_count()); + std::vector ignored; + auto index_of = [&](const HloValue* value) -> ValueIndex { + auto [it, inserted] = value_indices.try_emplace( + value, static_cast(values_.size())); + if (inserted) { + CHECK_LE(values_.size(), + static_cast(std::numeric_limits::max())); + values_.push_back(value); + ignored.push_back(IgnoreBuffer(*value)); + unscheduled_use_count_.push_back(0); + } + return it->second; + }; - // Precompute flattened value sets once per instruction in the computation. - absl::flat_hash_map - flattened_value_sets; - flattened_value_sets.reserve(computation->instruction_count()); - for (const HloInstruction* instruction : computation->instructions()) { - flattened_value_sets.emplace( - instruction, dataflow_analysis.GetFlattenedValueSet(instruction)); - } + // Completes the range appended to the pool from begin on: the accounted + // buffers are moved in front of the ignored ones. + auto finish_range = [&](int64_t begin) -> UseRange { + auto range_begin = value_index_pool_.begin() + begin; + auto ignored_begin = + std::partition(range_begin, value_index_pool_.end(), + [&](ValueIndex index) { return !ignored[index]; }); + const int64_t size = value_index_pool_.end() - range_begin; + const int64_t accounted = ignored_begin - range_begin; + return UseRange{begin, size, accounted}; + }; - auto get_flattened_value_set = - [&](const HloInstruction* inst) -> const HloValueSet& { - auto it = flattened_value_sets.find(inst); - if (it != flattened_value_sets.end()) { - return it->second; + // The flattened value set of an instruction as value indices, appended to + // the pool once per instruction and shared by all of its users. + absl::flat_hash_map flattened_ranges; + flattened_ranges.reserve(computation->instruction_count()); + auto flattened_range = [&](const HloInstruction* inst) -> UseRange { + auto [it, inserted] = flattened_ranges.try_emplace(inst); + if (inserted) { + const HloValueSet value_set = + dataflow_analysis.GetFlattenedValueSet(inst); + const int64_t begin = value_index_pool_.size(); + for (const HloValue* value : value_set.values()) { + value_index_pool_.push_back(index_of(value)); + } + it->second = finish_range(begin); } - return flattened_value_sets - .emplace(inst, dataflow_analysis.GetFlattenedValueSet(inst)) - .first->second; + return it->second; }; - // Create a map containing the HloValue uses for each HLO instruction. + // Create a map containing the HloValue uses for each HLO instruction: the + // union of the flattened value sets of its operands. buffer_uses_.reserve(computation->instruction_count()); + std::vector union_indices; for (auto* instruction : computation->instructions()) { - if (instruction->operands().empty()) { - buffer_uses_[instruction] = {}; - } else if (instruction->unique_operands().size() == 1) { - const HloValueSet& vs = - get_flattened_value_set(instruction->operand(0)); - buffer_uses_[instruction] = vs.values(); - } else { - absl::flat_hash_set instr_uses; - for (auto* operand : instruction->unique_operands()) { - const HloValueSet& vs = get_flattened_value_set(operand); - instr_uses.insert(vs.values().begin(), vs.values().end()); - } - buffer_uses_[instruction] = - std::vector(instr_uses.begin(), instr_uses.end()); + const HloInstruction::InstructionVector& operands = + instruction->operands(); + if (operands.empty()) { + buffer_uses_[instruction] = UseRange(); + continue; + } + // Instructions whose operands are all the same instruction share its + // flattened value list. + if (std::all_of(operands.begin(), operands.end(), + [&](const HloInstruction* operand) { + return operand == operands.front(); + })) { + buffer_uses_[instruction] = flattened_range(operands.front()); + continue; + } + // The union is appended to the pool only after every operand's own + // range is, so that finish_range partitions the union alone. + union_indices.clear(); + for (const HloInstruction* operand : operands) { + absl::Span operand_indices = + AllUses(flattened_range(operand)); + union_indices.insert(union_indices.end(), operand_indices.begin(), + operand_indices.end()); } + std::sort(union_indices.begin(), union_indices.end()); + union_indices.erase( + std::unique(union_indices.begin(), union_indices.end()), + union_indices.end()); + const int64_t begin = value_index_pool_.size(); + value_index_pool_.insert(value_index_pool_.end(), union_indices.begin(), + union_indices.end()); + buffer_uses_[instruction] = finish_range(begin); } - // Initialize unscheduled use counts and precompute bytes defined for each - // instruction. - unscheduled_use_count_.reserve(computation->instruction_count()); + // Precompute bytes defined for each instruction. bytes_defined_.reserve(computation->instruction_count()); for (auto* instruction : computation->instructions()) { int64_t bytes_defined = 0; - bool ignore = IgnoreInstruction(*instruction); - dataflow_analysis.GetInstructionValueSet(instruction) - .ForEachElement( - [&](const ShapeIndex& index, const HloValueSet& value_set) { - if (dataflow_analysis.ValueIsDefinedAt(instruction, index)) { - const HloValue* value = - &dataflow_analysis.GetValueDefinedAt(instruction, index); - unscheduled_use_count_[value] = 0; - if (!ignore) { - bytes_defined += (*size_function_)(*value); - } - } - }); + if (!IgnoreInstruction(*instruction)) { + dataflow_analysis.GetInstructionValueSet(instruction) + .ForEachElement([&](const ShapeIndex& index, + const HloValueSet& value_set) { + if (dataflow_analysis.ValueIsDefinedAt(instruction, index)) { + bytes_defined += (*size_function_)( + dataflow_analysis.GetValueDefinedAt(instruction, index)); + } + }); + } bytes_defined_[instruction] = bytes_defined; } + // Initialize unscheduled use counts. for (auto* instruction : computation->instructions()) { - for (const HloValue* buffer : buffer_uses_.at(instruction)) { - ++unscheduled_use_count_[buffer]; + for (ValueIndex index : AllUses(buffer_uses_.at(instruction))) { + ++unscheduled_use_count_[index]; } } // Buffers live out of the computation have an implicit use at the end of // the computation. - const HloValueSet& root_value_set = - get_flattened_value_set(computation->root_instruction()); - for (const HloValue* live_out_buffer : root_value_set.values()) { - ++unscheduled_use_count_[live_out_buffer]; + const UseRange live_out = flattened_range(computation->root_instruction()); + for (ValueIndex index : AllUses(live_out)) { + ++unscheduled_use_count_[index]; } } + // All value indices of range. + absl::Span AllUses(const UseRange& range) const { + return absl::MakeConstSpan(value_index_pool_) + .subspan(range.begin, range.size); + } + + // The prefix of range whose memory the scheduling heuristic accounts for. + absl::Span AccountedUses(const UseRange& range) const { + return absl::MakeConstSpan(value_index_pool_) + .subspan(range.begin, range.accounted); + } + // Returns whether the memory used by the given buffer should be ignored by // the scheduling heuristic. static bool IgnoreBuffer(const HloValue& buffer) { @@ -219,32 +296,16 @@ class ListScheduler { // The total size of all buffers defined by this instruction. int64_t bytes_defined; - // For each buffer B used by this instruction, we keep a pair (B, U), where - // U is the number of uses of B that have not yet been scheduled. This pair - // is a pointer into the unscheduled_use_count_ map, so it gets updated for - // free when we update counts in the map. - std::vector*> - used_buffer_unscheduled_use_counts; + // The buffers used by this instruction whose memory the heuristic + // accounts for, as indices into unscheduled_use_count_, which is read at + // the time of the query. + absl::Span accounted_uses; }; // Creates a ReadyListEntry for the given instruction. ReadyListEntry MakeReadyListEntry(HloInstruction* instruction) { - ReadyListEntry entry; - entry.instruction = instruction; - entry.bytes_defined = bytes_defined_.at(instruction); - - const auto& uses = buffer_uses_.at(instruction); - entry.used_buffer_unscheduled_use_counts.reserve(uses.size()); - for (const HloValue* buffer : uses) { - if (IgnoreBuffer(*buffer)) { - continue; - } - auto unscheduled_use_count_it = unscheduled_use_count_.find(buffer); - CHECK(unscheduled_use_count_it != unscheduled_use_count_.end()); - entry.used_buffer_unscheduled_use_counts.push_back( - &*unscheduled_use_count_it); - } - return entry; + return ReadyListEntry{instruction, bytes_defined_.at(instruction), + AccountedUses(buffer_uses_.at(instruction))}; } // Returns the number of bytes freed *after* the HLO instruction finishes. @@ -271,11 +332,9 @@ class ListScheduler { } int64_t freed_bytes = 0; - for (const auto& kv : entry.used_buffer_unscheduled_use_counts) { - auto buffer = kv->first; - auto use_count = kv->second; - if (use_count == 1) { - freed_bytes += (*size_function_)(*buffer); + for (ValueIndex index : entry.accounted_uses) { + if (unscheduled_use_count_[index] == 1) { + freed_bytes += (*size_function_)(*values_[index]); } } return freed_bytes - entry.bytes_defined; @@ -346,8 +405,8 @@ class ListScheduler { bool adjust_ready_queue = false; // Update the unscheduled uses of the logical buffers. - for (const HloValue* buffer : buffer_uses_.at(best)) { - int64_t& count = unscheduled_use_count_[buffer]; + for (ValueIndex index : AllUses(buffer_uses_.at(best))) { + UseCount& count = unscheduled_use_count_[index]; CHECK_GT(count, 0); --count; if (count == 1) { @@ -406,13 +465,17 @@ class ListScheduler { HloComputation* computation_; const BufferValue::SizeFunction* absl_nonnull size_function_; - // A map containing the HloValues that each instruction uses. - absl::flat_hash_map> - buffer_uses_; + // Per value tables, indexed by ValueIndex: the HloValue and the count of + // unscheduled HLOs using it. + std::vector values_; + std::vector unscheduled_use_count_; + + // The value index lists that the UseRanges refer to. + std::vector value_index_pool_; - // A map containing the count of unscheduled HLOs which using a particular - // HloValue. - absl::flat_hash_map unscheduled_use_count_; + // A map containing the HloValues that each instruction uses. Instructions + // with a single unique operand share the operand's flattened value list. + absl::flat_hash_map buffer_uses_; // A map containing the total bytes defined by each instruction. absl::flat_hash_map bytes_defined_; diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_memory_scheduler_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_memory_scheduler_test.cc index 8a602741f783db..f983d74a6b7094 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_memory_scheduler_test.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_memory_scheduler_test.cc @@ -15,11 +15,17 @@ limitations under the License. #include "xla/hlo/transforms/simplifiers/hlo_memory_scheduler.h" +#include #include #include #include +#include +#include #include +#include +#include #include +#include #include #include @@ -27,12 +33,17 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/base/nullability.h" #include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" #include "absl/log/check.h" #include "absl/status/status_matchers.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_format.h" +#include "absl/strings/str_join.h" #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/hlo/analysis/alias_info.h" #include "xla/hlo/analysis/hlo_alias_analysis.h" +#include "xla/hlo/analysis/hlo_dataflow_analysis.h" #include "xla/hlo/analysis/hlo_ordering.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" @@ -196,6 +207,368 @@ ENTRY root { peak_memory); } +// Knobs of a generated tuple heavy while body. Every element the loop carries +// has its own size, so the priorities of its consumers differ. +struct TupleHeavyWhileBodySpec { + std::string name; + int tuple_width; + // Results are tupled again and read back through get-tuple-element, so + // consumers use a buffer through an alias of its defining instruction. + bool read_through_aliases; + // An add whose two operands hold the same buffer. + bool overlapping_operands; + // A multiply of a value with itself. + bool duplicated_operands; + // Elements that also combine with constants. Two instructions share one + // constant: one on the main chain and an early userless one that reads the + // element through a tuple, so its priority goes stale until the constant's + // other use refreshes it. A second constant feeds a userless negate alone. + int constants; + // Slices without users, which compete with the root for the last slots. + int dead_instructions; + uint32_t seed; +}; + +void PrintTo(const TupleHeavyWhileBodySpec& spec, std::ostream* os) { + *os << spec.name; +} + +// Builds the HLO text of a while loop whose body follows the spec. The seed +// picks the element sizes, the pairs of elements that mix, the sizes of the +// dead slices and the elements the root passes through unchanged. +std::string TupleHeavyWhileBodyHlo(const TupleHeavyWhileBodySpec& spec) { + CHECK_GE(spec.tuple_width, 2); + CHECK_LE(spec.constants, spec.tuple_width); + std::mt19937 rng(spec.seed); + const int width = spec.tuple_width; + std::vector sizes(width); + for (int i = 0; i < width; ++i) { + sizes[i] = 8 * (i + 1); + std::swap(sizes[i], sizes[rng() % (i + 1)]); + } + auto array = [&](int i) { return absl::StrFormat("f32[%d]", sizes[i]); }; + std::vector element_shapes; + for (int i = 0; i < width; ++i) { + element_shapes.push_back(array(i)); + } + const std::string carried_shape = + absl::StrCat("(", absl::StrJoin(element_shapes, ", "), ", s32[])"); + + std::string hlo = absl::StrFormat( + "HloModule tuple_heavy_while_%s\n\nbody {\n p = %s parameter(0)\n", + spec.name, carried_shape); + auto line = [&](absl::string_view name, absl::string_view shape, + absl::string_view op) { + absl::StrAppend(&hlo, " ", name, " = ", shape, " ", op, "\n"); + }; + for (int i = 0; i < width; ++i) { + line(absl::StrCat("g", i), array(i), + absl::StrFormat("get-tuple-element(p), index=%d", i)); + } + line("counter", "s32[]", + absl::StrFormat("get-tuple-element(p), index=%d", width)); + line("one", "s32[]", "constant(1)"); + line("next", "s32[]", "add(counter, one)"); + for (int i = 0; i < width; ++i) { + line(absl::StrCat("a", i), array(i), absl::StrFormat("negate(g%d)", i)); + } + std::vector reads(width); + for (int i = 0; i < width; ++i) { + reads[i] = absl::StrCat("a", i); + } + if (spec.read_through_aliases) { + std::vector operands; + for (int i = 0; i < width; ++i) { + operands.push_back(reads[i]); + } + line("t", absl::StrCat("(", absl::StrJoin(element_shapes, ", "), ")"), + absl::StrFormat("tuple(%s)", absl::StrJoin(operands, ", "))); + for (int i = 0; i < width; ++i) { + reads[i] = absl::StrCat("ta", i); + line(reads[i], array(i), + absl::StrFormat("get-tuple-element(t), index=%d", i)); + } + } + for (int i = 0; i < width; ++i) { + line(absl::StrCat("b", i), array(i), + spec.overlapping_operands + ? absl::StrFormat("add(%s, a%d)", reads[i], i) + : absl::StrFormat("exponential(%s)", reads[i])); + line(absl::StrCat("c", i), array(i), + spec.duplicated_operands ? absl::StrFormat("multiply(b%d, b%d)", i, i) + : absl::StrFormat("negate(b%d)", i)); + if (i < spec.constants) { + std::vector literal(sizes[i]); + absl::c_iota(literal, 0); + line(absl::StrCat("k", i), array(i), + absl::StrFormat("constant({%s})", absl::StrJoin(literal, ", "))); + line(absl::StrCat("ts", i), absl::StrCat("(", array(i), ")"), + absl::StrFormat("tuple(a%d)", i)); + line(absl::StrCat("gs", i), array(i), + absl::StrFormat("get-tuple-element(ts%d), index=0", i)); + line(absl::StrCat("r", i), array(i), + absl::StrFormat("subtract(gs%d, k%d)", i, i)); + line(absl::StrCat("kk", i), array(i), + absl::StrFormat("constant({%s})", absl::StrJoin(literal, ", "))); + line(absl::StrCat("kd", i), array(i), absl::StrFormat("negate(kk%d)", i)); + line(absl::StrCat("d", i), array(i), + absl::StrFormat("add(c%d, k%d)", i, i)); + } else { + line(absl::StrCat("d", i), array(i), absl::StrFormat("negate(c%d)", i)); + } + } + for (int i = 0; i < width; ++i) { + const int partner = (i + 1 + rng() % (width - 1)) % width; + line(absl::StrCat("cat", i), + absl::StrFormat("f32[%d]", sizes[i] + sizes[partner]), + absl::StrFormat("concatenate(d%d, d%d), dimensions={0}", i, partner)); + line(absl::StrCat("e", i), array(i), + absl::StrFormat("slice(cat%d), slice={[0:%d]}", i, sizes[i])); + } + for (int k = 0; k < spec.dead_instructions; ++k) { + const int i = k % width; + const int size = 1 + rng() % (sizes[i] - 1); + line(absl::StrCat("dead", k), absl::StrFormat("f32[%d]", size), + absl::StrFormat("slice(e%d), slice={[0:%d]}", i, size)); + } + std::vector results; + for (int i = 0; i < width; ++i) { + results.push_back(rng() % 3 == 0 ? absl::StrCat("g", i) + : absl::StrCat("e", i)); + } + results.push_back("next"); + absl::StrAppend(&hlo, " ROOT out = ", carried_shape, " tuple(", + absl::StrJoin(results, ", "), ")\n}\n\n"); + + absl::StrAppend(&hlo, "cond {\n cp = ", carried_shape, " parameter(0)\n"); + absl::StrAppendFormat(&hlo, " ci = s32[] get-tuple-element(cp), index=%d\n", + width); + absl::StrAppend(&hlo, + " limit = s32[] constant(3)\n" + " ROOT lt = pred[] compare(ci, limit), direction=LT\n}\n\n"); + + absl::StrAppend(&hlo, "ENTRY main {\n"); + std::vector parameters; + for (int i = 0; i < width; ++i) { + parameters.push_back(absl::StrCat("p", i)); + absl::StrAppendFormat(&hlo, " p%d = %s parameter(%d)\n", i, array(i), i); + } + parameters.push_back("zero"); + absl::StrAppend(&hlo, " zero = s32[] constant(0)\n init = ", carried_shape, + " tuple(", absl::StrJoin(parameters, ", "), ")\n"); + absl::StrAppend(&hlo, " ROOT loop = ", carried_shape, + " while(init), condition=cond, body=body\n}\n"); + return hlo; +} + +// The list scheduling heuristic written the direct way, as the reference for +// ListMemoryScheduler: the unscheduled use count of every buffer lives in a +// map keyed by HloValue, and an entry's priority is recomputed by scanning the +// buffers it uses whenever the heuristic consults it. +HloInstructionSequence ReferenceListSchedule( + HloComputation* computation, const HloAliasAnalysis& alias_analysis, + const BufferValue::SizeFunction& size_function) { + const HloDataflowAnalysis& dataflow = alias_analysis.dataflow_analysis(); + auto ignored = [](const HloInstruction& instruction) { + return instruction.opcode() == HloOpcode::kParameter || + instruction.opcode() == HloOpcode::kConstant; + }; + + // An instruction uses every value in the flattened value sets of its + // unique operands; the root's values have one more use, the live out. + absl::flat_hash_map> + uses; + absl::flat_hash_map unscheduled_uses; + absl::flat_hash_map bytes_defined; + for (HloInstruction* instruction : computation->instructions()) { + absl::flat_hash_set& used = uses[instruction]; + for (const HloInstruction* operand : instruction->unique_operands()) { + const HloValueSet value_set = dataflow.GetFlattenedValueSet(operand); + used.insert(value_set.values().begin(), value_set.values().end()); + } + for (const HloValue* value : used) { + ++unscheduled_uses[value]; + } + int64_t& defined = bytes_defined[instruction]; + if (!ignored(*instruction)) { + dataflow.GetInstructionValueSet(instruction) + .ForEachElement([&](const ShapeIndex& index, const HloValueSet&) { + if (dataflow.ValueIsDefinedAt(instruction, index)) { + defined += + size_function(dataflow.GetValueDefinedAt(instruction, index)); + } + }); + } + } + const HloValueSet live_out = + dataflow.GetFlattenedValueSet(computation->root_instruction()); + for (const HloValue* value : live_out.values()) { + ++unscheduled_uses[value]; + } + + using Priority = std::pair; + auto priority = [&](const HloInstruction* instruction) -> Priority { + if (ShapeUtil::IsEffectiveScalar(instruction->shape())) { + return {std::numeric_limits::max(), + std::numeric_limits::max()}; + } + int64_t freed = 0; + if (instruction->opcode() == HloOpcode::kOutfeed && + !instruction->outfeed_config().empty()) { + freed = INT_MAX; + } else if (instruction->opcode() == HloOpcode::kInfeed && + !instruction->infeed_config().empty()) { + freed = INT_MIN; + } else { + for (const HloValue* value : uses.at(instruction)) { + if (!ignored(*value->instruction()) && + unscheduled_uses.at(value) == 1) { + freed += size_function(*value); + } + } + freed -= bytes_defined.at(instruction); + } + return {freed, instruction->user_count()}; + }; + + absl::flat_hash_map unscheduled_predecessors; + for (HloInstruction* instruction : computation->instructions()) { + for (HloInstruction* user : instruction->users()) { + ++unscheduled_predecessors[user]; + } + for (HloInstruction* successor : instruction->control_successors()) { + ++unscheduled_predecessors[successor]; + } + } + + // Ready instructions ordered by priority; among equal priorities the most + // recently inserted one is scheduled first. + std::multimap ready; + absl::flat_hash_map::iterator> + ready_entry; + auto make_ready = [&](HloInstruction* instruction) { + ready_entry[instruction] = + ready.emplace(priority(instruction), instruction); + }; + for (HloInstruction* instruction : computation->instructions()) { + if (instruction->operands().empty() && + instruction->control_predecessors().empty()) { + make_ready(instruction); + } + } + + HloInstructionSequence sequence; + while (!ready.empty()) { + auto best_it = std::prev(ready.end()); + HloInstruction* best = best_it->second; + ready.erase(best_it); + ready_entry.erase(best); + sequence.push_back(best); + + bool refresh = false; + for (const HloValue* value : uses.at(best)) { + int64_t& count = unscheduled_uses.at(value); + --count; + if (count == 1) { + refresh = true; + } + } + auto release = [&](HloInstruction* instruction) { + if (--unscheduled_predecessors.at(instruction) == 0) { + make_ready(instruction); + } + }; + for (HloInstruction* user : best->users()) { + release(user); + } + for (HloInstruction* successor : best->control_successors()) { + release(successor); + } + if (!refresh) { + continue; + } + // Only the ready users of the scheduled instruction's operands get a + // fresh priority; a changed one moves behind the entries of equal + // priority. Every other ready instruction keeps its stale priority. + for (HloInstruction* operand : best->operands()) { + for (HloInstruction* user : operand->users()) { + auto it = ready_entry.find(user); + if (it == ready_entry.end()) { + continue; + } + const Priority fresh = priority(user); + if (fresh == it->second->first) { + continue; + } + auto stale = it->second; + it->second = ready.emplace(fresh, user); + ready.erase(stale); + } + } + } + return sequence; +} + +std::string InstructionNames(const HloInstructionSequence& sequence) { + std::vector names; + for (const HloInstruction* instruction : sequence.instructions()) { + names.push_back(instruction->name()); + } + return absl::StrJoin(names, " "); +} + +class ListSchedulerStressTest + : public HloSchedulingTest, + public ::testing::WithParamInterface {}; + +// ListMemoryScheduler must produce the reference's sequence on every +// computation of a generated tuple heavy while loop: same buffer use counts +// at every step, same bytes freed, same refresh points, same tie breaking. +TEST_P(ListSchedulerStressTest, MatchesReference) { + const std::string hlo = TupleHeavyWhileBodyHlo(GetParam()); + SCOPED_TRACE(hlo); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo)); + BufferValue::SizeFunction size_fn = [](const BufferValue& buffer) { + return ShapeUtil::ByteSizeOf(buffer.shape(), /*pointer_size=*/8); + }; + ASSERT_OK_AND_ASSIGN( + HloSchedule schedule, + ScheduleModule(module.get(), + ListMemoryScheduler(&alias_info_, &size_fn))); + ASSERT_OK_AND_ASSIGN(std::unique_ptr alias_analysis, + HloAliasAnalysis::Run(module.get(), &alias_info_)); + for (HloComputation* computation : module->MakeNonfusionComputations()) { + EXPECT_EQ(InstructionNames(schedule.sequence(computation)), + InstructionNames( + ReferenceListSchedule(computation, *alias_analysis, size_fn))) + << "computation " << computation->name(); + } +} + +INSTANTIATE_TEST_SUITE_P( + TupleHeavyWhileBodies, ListSchedulerStressTest, + ::testing::ValuesIn(std::vector{ + {"plain_w2", 2, false, false, false, 0, 0, 1}, + {"aliases_w4", 4, true, false, false, 0, 0, 2}, + {"overlap_w4", 4, true, true, false, 0, 0, 3}, + {"duplicates_w4", 4, true, true, true, 0, 0, 4}, + {"constants_w6", 6, true, true, true, 3, 0, 5}, + {"dead_w6", 6, true, true, true, 3, 4, 6}, + {"all_w8_seed7", 8, true, true, true, 4, 5, 7}, + {"all_w8_seed8", 8, true, true, true, 8, 6, 8}, + {"all_w8_seed9", 8, true, true, true, 8, 6, 9}, + {"all_w8_seed10", 8, true, true, true, 5, 8, 10}, + {"all_w12_seed11", 12, true, true, true, 6, 8, 11}, + {"all_w12_seed12", 12, true, true, true, 12, 10, 12}, + {"all_w16_seed13", 16, true, true, true, 8, 12, 13}, + }), + [](const ::testing::TestParamInfo& info) { + return info.param.name; + }); + TEST_F(HloSchedulingTest, DefaultSchedulerRunsThreeSchedulers) { const char* module_str = R"( HloModule test_aliasing_module diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc index 9890ec597b9e52..9ddd8a59f13a4f 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc @@ -69,6 +69,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/hlo/ir/hlo_original_value.h" #include "xla/hlo/ir/hlo_sharding.h" #include "xla/hlo/ir/hlo_sharding_metadata.h" #include "xla/hlo/translate/hlo_to_mhlo/async_importer.h" @@ -447,38 +448,60 @@ absl::StatusOr HloFunctionImporter::ImportAsFunc( return Internal("Expected %d leaf parameters but got %d", flattened_shardings.size(), leaf_count); } - - for (int i = 0; i < leaf_count; ++i) { - mlir::NamedAttrList argAttrs(function.getArgAttrDict(arg_index)); - if (!flattened_shardings.empty()) { - argAttrs.set(xla::kMhloSharding, - ConvertSharding(flattened_shardings[i], builder_)); - } - if (frontend_attributes) { - if (leaf_count > 1) { - return InvalidArgument( - "A tuple parameter that is being flattened shouldn't have " - "frontend attributes"); - } - argAttrs.set(xla::kMhloFrontendAttributes, frontend_attributes); - } - if (parameter->parameter_replicated_at_leaf_buffers() && - parameter->parameter_replicated_at_leaf_buffers()->at(i)) { - argAttrs.set(xla::kMhloParameterReplication, - builder_->getBoolArrayAttr({true})); - } - // NOTE: since we are flattening args, all arguments will share the same - // location as the tuple parameter instruction. - function.getArgument(arg_index).setLoc( - mlir::hlo::GenerateInstructionLocation(instruction, context_)); - funcArgAttrs[arg_index++] = argAttrs.getDictionary(context_); - } + if (frontend_attributes && leaf_count > 1) { + return InvalidArgument( + "A tuple parameter that is being flattened shouldn't have " + "frontend attributes"); + } + + int i = 0; + ShapeUtil::ForEachLeafShape( + parameter->shape(), + [&](const Shape& subshape, const ShapeIndex& index) { + mlir::NamedAttrList argAttrs(function.getArgAttrDict(arg_index)); + if (!flattened_shardings.empty()) { + argAttrs.set(xla::kMhloSharding, + ConvertSharding(flattened_shardings[i], builder_)); + } + if (frontend_attributes) { + argAttrs.set(xla::kMhloFrontendAttributes, frontend_attributes); + } + if (parameter->parameter_replicated_at_leaf_buffers() && + parameter->parameter_replicated_at_leaf_buffers()->at(i)) { + argAttrs.set(xla::kMhloParameterReplication, + builder_->getBoolArrayAttr({true})); + } + if (parameter->original_value()) { + // Parameter instructions are not call instructions, so they never + // carry a call_hierarchy (or is_synthetic_call). Each flattened + // argument only needs the leaf OriginalArray at `index`. + const auto& orig_array = + parameter->original_value()->original_array(index); + if (orig_array.has_value()) { + OriginalValue leaf_ov(subshape); + *leaf_ov.mutable_original_array({}) = *orig_array; + argAttrs.set(xla::kMhloOriginalValueAttr, + ConvertOriginalValue(leaf_ov, builder_)); + } + } + // NOTE: since we are flattening args, all arguments will share the + // same location as the tuple parameter instruction. + function.getArgument(arg_index).setLoc( + mlir::hlo::GenerateInstructionLocation(instruction, context_)); + funcArgAttrs[arg_index++] = argAttrs.getDictionary(context_); + ++i; + }); } else { mlir::NamedAttrList argAttrs(function.getArgAttrDict(arg_index)); if (parameter->has_sharding()) { argAttrs.set(xla::kMhloSharding, ConvertSharding(parameter->sharding(), builder_)); } + if (parameter->original_value()) { + argAttrs.set( + xla::kMhloOriginalValueAttr, + ConvertOriginalValue(*parameter->original_value(), builder_)); + } if (frontend_attributes) { argAttrs.set( xla::kMhloFrontendAttributes, @@ -590,17 +613,9 @@ absl::StatusOr HloFunctionImporter::ImportInstructionsImpl( // Setup the input parameters. const int num_parameters = computation.num_parameters(); - FuncOp func = llvm::dyn_cast(builder->getBlock()->getParentOp()); for (int i = 0; i < num_parameters; i++) { auto* hlo_parameter = computation.parameter_instruction(i); instruction_value_map_[hlo_parameter] = arguments[i]; - // Only add original value attributes to parameters in functions. Skip - // regions. - if (hlo_parameter->original_value() && func) { - func.setArgAttr( - i, kMhloOriginalValueAttr, - ConvertOriginalValue(*hlo_parameter->original_value(), builder_)); - } } for (auto instruction : computation.MakeInstructionPostOrder()) { @@ -2595,18 +2610,24 @@ mlir::NamedAttribute HloFunctionImporter::ConvertComparisonDirection( mlir::NamedAttribute HloFunctionImporter::ConvertComparisonOrder( ComparisonOrder order) { - mlir::stablehlo::ComparisonType type; switch (order) { case ComparisonOrder::kPartial: - type = mlir::stablehlo::ComparisonType::FLOAT; - break; + return builder_->getNamedAttr( + "compare_type", + mlir::stablehlo::ComparisonTypeAttr::get( + builder_->getContext(), mlir::stablehlo::ComparisonType::FLOAT)); case ComparisonOrder::kTotal: - type = mlir::stablehlo::ComparisonType::TOTALORDER; - break; + return builder_->getNamedAttr( + "compare_type", mlir::stablehlo::ComparisonTypeAttr::get( + builder_->getContext(), + mlir::stablehlo::ComparisonType::TOTALORDER)); + case ComparisonOrder::kWeak: + // TODO(b/565612124): Map to mlir::stablehlo::ComparisonType::WEAKORDER + // once added to StableHLO. + LOG(FATAL) << "Unsupported comparison order: " + << ComparisonOrderToString(order); } - return builder_->getNamedAttr( - "compare_type", - mlir::stablehlo::ComparisonTypeAttr::get(builder_->getContext(), type)); + LOG(FATAL) << "Unhandled comparison order: " << static_cast(order); } mlir::DenseIntElementsAttr HloFunctionImporter::ConvertDimensions( diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import.hlo b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import.hlo index f25477a5686791..0e39f1877ea29e 100644 --- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import.hlo +++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import.hlo @@ -2219,3 +2219,19 @@ add_original_value { // CHECK: }) {mhlo.original_value = #mhlo.original_value>]>} ROOT result = f32[8] all-reduce(input), channel_id=1, replica_groups={{0,1,2,3}, {4,5,6,7}}, to_apply=add_original_value, origin={{"c"}} } + +// Test tuple parameter original value +// CHECK-LABEL: func.func private @test_tuple_original_value +// CHECK-SAME: %arg0: tuple, tensor> {mhlo.original_value = #mhlo.original_value>, <[1], <"p1", []>>]>} +// FLATTEN-CHECK-LABEL: func private @test_tuple_original_value +// FLATTEN-CHECK-SAME: %arg0: tensor {mhlo.original_value = #mhlo.original_value>]>} +// FLATTEN-CHECK-SAME: %arg1: tensor {mhlo.original_value = #mhlo.original_value>]>} +%test_tuple_original_value (param.1: (f32[], f32[])) -> f32[] { + %param.1 = (f32[], f32[]) parameter(0), origin={({"p0"}, {"p1"})} + %gte.0 = f32[] get-tuple-element(%param.1), index=0 + %gte.1 = f32[] get-tuple-element(%param.1), index=1 + ROOT %add = f32[] add(%gte.0, %gte.1) +} + + + diff --git a/third_party/xla/xla/megascale/BUILD b/third_party/xla/xla/megascale/BUILD index 39dde3a7bf946d..4540a271b94c2d 100644 --- a/third_party/xla/xla/megascale/BUILD +++ b/third_party/xla/xla/megascale/BUILD @@ -50,4 +50,7 @@ tf_proto_library( srcs = ["megascale_runtime_error_overlay.proto"], create_grpc_library = True, make_default_target_header_only = True, + protodeps = [ + "//xla/tsl/protobuf:status_proto", + ], ) diff --git a/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.cc b/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.cc index 0168d09e5b957c..e75aa32c2a6bc8 100644 --- a/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.cc +++ b/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.cc @@ -93,7 +93,7 @@ CApiMegascaleErrorAggregator::~CApiMegascaleErrorAggregator() { void CApiMegascaleErrorAggregator::AddError( absl::string_view worker_id, - const runtime::MegaScaleRuntimeErrorOverlay& error) { + const runtime::external::MegaScaleRuntimeError& error) { std::string serialized_error = error.SerializeAsString(); PJRT_Megascale_ErrorAggregator_AddError_Args args{}; args.struct_size = PJRT_Megascale_ErrorAggregator_AddError_Args_STRUCT_SIZE; diff --git a/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.h b/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.h index 1285cab2416e49..960291c636c6b3 100644 --- a/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.h +++ b/third_party/xla/xla/megascale/c_api_client/c_api_megascale_error_aggregator.h @@ -58,7 +58,7 @@ class CApiMegascaleErrorAggregator { delete; void AddError(absl::string_view worker_id, - const runtime::MegaScaleRuntimeErrorOverlay& error); + const runtime::external::MegaScaleRuntimeError& error); ErrorDigest ProcessAndShutdown(); diff --git a/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.cc b/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.cc index 8a26c9c477274a..5ac419154b0054 100644 --- a/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.cc +++ b/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.cc @@ -695,7 +695,7 @@ CreateMegascaleCollectives( } typedef absl::AnyInvocable + const runtime::external::MegaScaleRuntimeError& error)> MegaScaleErrorHandler; static absl::Mutex megascale_error_handlers_mutex(absl::kConstInit); @@ -718,7 +718,7 @@ absl::Status RegisterMegascaleErrorHandler(absl::string_view handler_name, args.handler_name_size = handler_name.size(); args.handler = +[](const char* serialized_error, size_t serialized_error_size, void* user_data) { - runtime::MegaScaleRuntimeErrorOverlay error; + runtime::external::MegaScaleRuntimeError error; error.ParseFromString( absl::string_view(serialized_error, serialized_error_size)); MegaScaleErrorHandler* handler = @@ -815,11 +815,12 @@ GetInterfaceAddressesHelper(absl::string_view megascale_port_name, return addresses; } -absl::StatusOr> +absl::StatusOr> GetOrCreateRuntimeError( - runtime::MegaScaleRuntimeErrorOverlay::ErrorType error_type, + runtime::external::MegaScaleRuntimeError::ErrorType error_type, absl::Time start_time, const absl::Status& status, int32_t launch_id, - std::optional + std::optional< + runtime::external::MegaScaleRuntimeError::UnrecoverableErrorType> unrecoverable_error_type) { ABSL_ASSIGN_OR_RETURN(const PJRT_Api* c_api, pjrt::PjrtApi(kTpuPjrtName)); ABSL_ASSIGN_OR_RETURN(PJRT_Megascale_Extension * extension, @@ -849,11 +850,11 @@ GetOrCreateRuntimeError( }; CHECK_NOTNULL(args.serialized_error); - runtime::MegaScaleRuntimeErrorOverlay error; + runtime::external::MegaScaleRuntimeError error; if (!error.ParseFromString(absl::string_view(args.serialized_error, args.serialized_error_size))) { return absl::InternalError( - "Failed to parse MegaScaleRuntimeErrorOverlay proto from serialized " + "Failed to parse MegaScaleRuntimeError proto from serialized " "error."); } diff --git a/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.h b/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.h index 647d0765910188..8e763ac44c421d 100644 --- a/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.h +++ b/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client.h @@ -84,7 +84,8 @@ CreateMegascaleCollectives( absl::Status RegisterMegascaleErrorHandler( absl::string_view handler_name, - absl::AnyInvocable + absl::AnyInvocable< + void(const runtime::external::MegaScaleRuntimeError& error)> handler); absl::Status UnregisterMegascaleErrorHandler(absl::string_view handler_name); @@ -96,11 +97,12 @@ GetInterfaceAddressesHelper(absl::string_view megascale_port_name, bool use_all_interfaces, bool limit_to_process_numa_local_interfaces); -absl::StatusOr> +absl::StatusOr> GetOrCreateRuntimeError( - runtime::MegaScaleRuntimeErrorOverlay::ErrorType error_type, + runtime::external::MegaScaleRuntimeError::ErrorType error_type, absl::Time start_time, const absl::Status& status, int32_t launch_id, - std::optional + std::optional< + runtime::external::MegaScaleRuntimeError::UnrecoverableErrorType> unrecoverable_error_type = std::nullopt); } // namespace c_api_client diff --git a/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client_test.cc b/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client_test.cc index 968d98768562a9..98be9ba8171804 100644 --- a/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client_test.cc +++ b/third_party/xla/xla/megascale/c_api_client/megascale_c_api_client_test.cc @@ -42,7 +42,7 @@ using ::xla::megascale::c_api_client::RegisterMegascaleErrorHandler; using ::xla::megascale::c_api_client::UnregisterMegascaleErrorHandler; using ::xla::megascale::runtime::DCNTopology; using ::xla::megascale::runtime::EndpointAddresses; -using ::xla::megascale::runtime::MegaScaleRuntimeErrorOverlay; +using ::xla::megascale::runtime::external::MegaScaleRuntimeError; TEST(MegaScaleCApiClientTest, CreateDefaultMegaScaleClientContext) { TF_ASSERT_OK_AND_ASSIGN(auto client_context, @@ -89,7 +89,7 @@ TEST(MegaScaleCApiClientTest, CreateMegascaleErrorAggregator) { TF_ASSERT_OK_AND_ASSIGN(auto aggregator, CreateMegascaleErrorAggregator("test_app")); EXPECT_NE(aggregator, nullptr); - MegaScaleRuntimeErrorOverlay error; + MegaScaleRuntimeError error; aggregator->AddError("worker1", error); EXPECT_EQ(aggregator->size(), 1); EXPECT_TRUE(aggregator->active()); @@ -99,7 +99,7 @@ TEST(MegaScaleCApiClientTest, CreateMegascaleErrorAggregator) { } TEST(MegaScaleCApiClientTest, RegisterAndUnregisterMegascaleErrorHandler) { - auto handler = [](const MegaScaleRuntimeErrorOverlay& error) {}; + auto handler = [](const MegaScaleRuntimeError& error) {}; EXPECT_OK(RegisterMegascaleErrorHandler("test_handler", handler)); EXPECT_OK(UnregisterMegascaleErrorHandler("test_handler")); } @@ -115,16 +115,15 @@ TEST(MegaScaleCApiClientTest, GetInterfaceAddressesHelper) { TEST(MegaScaleCApiClientTest, GetOrCreateRuntimeError) { TF_ASSERT_OK_AND_ASSIGN( - auto result, GetOrCreateRuntimeError( - MegaScaleRuntimeErrorOverlay::UNRECOVERABLE_ERROR, - absl::Now(), absl::OkStatus(), /*launch_id=*/0, - MegaScaleRuntimeErrorOverlay::HOST_TO_DEVICE_ERROR)); + auto result, + GetOrCreateRuntimeError(MegaScaleRuntimeError::UNRECOVERABLE_ERROR, + absl::Now(), absl::OkStatus(), /*launch_id=*/0, + MegaScaleRuntimeError::HOST_TO_DEVICE_ERROR)); const auto& [error, is_new] = result; - EXPECT_EQ(error.error_type(), - MegaScaleRuntimeErrorOverlay::UNRECOVERABLE_ERROR); + EXPECT_EQ(error.error_type(), MegaScaleRuntimeError::UNRECOVERABLE_ERROR); EXPECT_EQ(error.unrecoverable_error_type(), - MegaScaleRuntimeErrorOverlay::HOST_TO_DEVICE_ERROR); + MegaScaleRuntimeError::HOST_TO_DEVICE_ERROR); EXPECT_TRUE(is_new); } diff --git a/third_party/xla/xla/megascale/megascale_runtime_error_overlay.proto b/third_party/xla/xla/megascale/megascale_runtime_error_overlay.proto index be1bc8d66e11fa..7109e4294991ec 100644 --- a/third_party/xla/xla/megascale/megascale_runtime_error_overlay.proto +++ b/third_party/xla/xla/megascale/megascale_runtime_error_overlay.proto @@ -15,13 +15,201 @@ syntax = "proto3"; -package xla.megascale.runtime; +package xla.megascale.runtime.external; + +import "xla/tsl/protobuf/status.proto"; option java_multiple_files = true; option java_outer_classname = "Runtime"; -// Overlay for xla.megascale.runtime.MegaScaleRuntimeError -message MegaScaleRuntimeErrorOverlay { +enum TpuCoreTypeProto { + TPU_CORE_TYPE_INVALID = 0; + TPU_CORE_TYPE_TENSOR_CORE = 1; + TPU_CORE_TYPE_RESERVED = 2; + TPU_CORE_TYPE_SPARSE_CORE = 3; +} + +// LINT.IfChange(megascale_action) +message MegaScaleAction { + message DeviceToHost { + // Buffer produced as result of D2H transfer will be made available + // at done_key. + string done_key = 2; + } + + message HostToDevice {} + + message NetworkSend { + message Peer { + string key = 1; + } + + message Transfer { + Peer peer = 3; + } + + message TransferInfo { + Transfer transfer = 4; + } + + // Transfer buffer over network to destinations specified in transfer_info. + TransferInfo transfer_info = 1; + } + + message NetworkReceive { + // Buffer produced as a result of network received with this key will + // be available at done_key in callback registry. + string done_key = 1; + + // Expected size of the buffer. + int64 buffer_size = 2; + } + + message Compute {} + + // An input can be either a valid buffer to represent data dependency or a + // void buffer to indicate control dependency. + message Input { + // Index of action in MegaScaleActionGraph.actions that produces this + // buffer. + int32 action_index = 1; + } + + oneof action { + DeviceToHost device_to_host = 1; + HostToDevice host_to_device = 2; + NetworkSend network_send = 3; + NetworkReceive network_receive = 4; + Compute compute = 5; + } + + // List of inputs. Action is executed once all its inputs are available and + // any run_after actions are done. + repeated Input inputs = 6; + + // Index of this action in the graph. Mainly for debugging. + int32 index = 8; +} +// LINT.ThenChange() + +// LINT.IfChange(megascale_action_graph) +message MegaScaleActionGraph { + // Unique identifier for this graph. Typically this is set to unique + // identifiers generated by XLA compiler. + string id = 1; + + // All actions that belong to this Graph. Actions are executed once all their + // inputs are available (generated by predecessor actions) and any run_after + // actions are done. See MegaScaleAction for more details. The index of + // actions is referred to by dependent actions. + repeated MegaScaleAction actions = 2; + + // The HLO module name corresponding to the send op that triggered this graph. + string hlo_module = 7; +} +// LINT.ThenChange() + +// LINT.IfChange(executor_telemetry) +message ExecutorTelemetry { + message PendingGraphState { + MegaScaleActionGraph graph = 1; + int32 launch_id = 2; + repeated int32 completed_actions = 3; + repeated int32 pending_actions = 4; + repeated int32 active_actions = 6; + } + + repeated PendingGraphState pending_graphs = 1; +} +// LINT.ThenChange() + +// LINT.IfChange(program_run_info) +message ProgramRunInfo { + // run_id is only relevant within a TPU and shouldn't be compared across TPUs. + uint64 run_id = 1; + // launch_id should be consistent across TPUs running the same execution. + // (within and across slices) + int32 launch_id = 2; + // The name of the HLO module/program. + string name = 3; + // We store the fingerprint (as a hex string) just in case there are different + // modules that happen to have the same module_name, which generally shouldn't + // happen. + string fingerprint = 4; +} +// LINT.ThenChange() + +// LINT.IfChange(tpu_core_hlo) +message TpuCoreHLO { + // Chip ID local to the slice. + int32 chip_id = 1; + // Core index within its chip. + int32 core_idx = 2; + // Full HLO details. + string hlo = 3; + // Fingerprint of the running executable. + string executable_fingerprint = 4; + // Current program counter + optional int64 pc = 6; + // Tag information to be used with PC + optional int64 tag = 7; + string hostname = 8; + // Short name of the HLO. + string hlo_name = 9; + // Module name of the running executable. + string module_name = 10; + // Launch ID of the running executable. + int32 launch_id = 12; + // The type of the TPU core. + TpuCoreTypeProto tpu_core_type = 13; + // Short name of the computation where the HLO resides in. + string computation_name = 15; + // Basic debug info about recently completed programs runs. Sorted in reverse + // chronological order. (the first entry is the most recently completed run) + repeated ProgramRunInfo recent_completions = 16; + // True if the TPU is running a program actively. + optional bool active = 18; + // The error when obtaining the TPU core information. + repeated tensorflow.StatusProto errors = 19; +} +// LINT.ThenChange() + +// LINT.IfChange(rapid_eye_info) +message RapidEyeInfo { + ExecutorTelemetry telemetry = 1; + + repeated TpuCoreHLO tpu_core_hlos = 2; + + // Whether the RapidEyeInfo was truncated due to exceeding the size limit. + bool is_truncated = 3; + + // The number of slices, tasks per slice, local slice ID and task ID that are + // obtained from the workers' last used topology. This is useful for the error + // aggregator to determine if all error messages have been received. + int64 num_slices = 4; + int64 num_tasks_per_slice = 5; + int32 slice_id = 6; + int32 task_id = 7; + + // The incarnation ID for the workload shared across all workers. This will + // reset every time the workload restart. + int64 incarnation_id = 8; + + // Specifies the TPU platform type. + int32 tpu_version = 9; + string chip_config_name = 10; +} +// LINT.ThenChange() + +// LINT.IfChange(runtime_state_summary) +message RuntimeStateSummary { + // RapidEye info. + RapidEyeInfo rapid_eye_info = 4; +} +// LINT.ThenChange() + +// LINT.IfChange(megascale_runtime_error) +message MegaScaleRuntimeError { enum ErrorType { NO_ERROR = 0; HANG_DETECTED = 1; @@ -29,6 +217,18 @@ message MegaScaleRuntimeErrorOverlay { CANCELLED = 3; } + // How the hang error is detected. + enum HangErrorType { + // The default value. This indicates that the hang error type is not + // specified. + HANG_ERROR_UNSPECIFIED = 0; + // The hang error type is caused by a megascale timeout. + TIMEOUT_ERROR = 1; + // The hang error type is caused by no progress being made on all registered + // metrics like rpc, pjrt, etc. + NO_PROGRESS_ERROR = 2; + } + enum UnrecoverableErrorType { UNRECOVERABLE_ERROR_UNCLASSIFIED = 0; HOST_TO_DEVICE_ERROR = 1; @@ -36,9 +236,16 @@ message MegaScaleRuntimeErrorOverlay { } ErrorType error_type = 1; - UnrecoverableErrorType unrecoverable_error_type = 12; string error_message = 2; uint64 estimated_start_time_ms = 3; + RuntimeStateSummary runtime_state = 4; int32 triggering_launch_id = 5; string hostname = 6; + string job_name = 8; + int32 task_id = 9; + fixed64 event_id = 11; + UnrecoverableErrorType unrecoverable_error_type = 12; + fixed64 runtime_error_id = 13; + HangErrorType hang_error_type = 14; } +// LINT.ThenChange() diff --git a/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/stablehlo_canonicalize_from_hlo_import.mlir b/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/stablehlo_canonicalize_from_hlo_import.mlir index 93f8ffe1bfa011..b3343841df7093 100644 --- a/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/stablehlo_canonicalize_from_hlo_import.mlir +++ b/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/stablehlo_canonicalize_from_hlo_import.mlir @@ -215,4 +215,27 @@ func.func @while_op_original_value_index_shifting( : tensor<10xf32>, tensor, tensor<10xf32>, tensor<10xf32> } return %1#0, %1#2, %1#3 : tensor<10xf32>, tensor<10xf32>, tensor<10xf32> -} \ No newline at end of file +} + +// ----- + +// CHECK-LABEL: func.func @concatenate_single_operand +// CHECK-SAME: (%[[ARG0:.*]]: tensor<2x4xf32>) -> tensor<2x4xf32> +func.func @concatenate_single_operand(%arg0: tensor<2x4xf32>) -> tensor<2x4xf32> { + // CHECK-NOT: stablehlo.concatenate + // CHECK: return %[[ARG0]] : tensor<2x4xf32> + %0 = stablehlo.concatenate %arg0, dim = 0 : (tensor<2x4xf32>) -> tensor<2x4xf32> + return %0 : tensor<2x4xf32> +} + +// ----- + +// CHECK-LABEL: func.func @concatenate_remove_empty +// CHECK-SAME: (%[[ARG0:.*]]: tensor<2x4xf32>, %[[ARG1:.*]]: tensor<0x4xf32>) -> tensor<2x4xf32> +func.func @concatenate_remove_empty(%arg0: tensor<2x4xf32>, %arg1: tensor<0x4xf32>) -> tensor<2x4xf32> { + // CHECK-NOT: stablehlo.concatenate + // CHECK: return %[[ARG0]] : tensor<2x4xf32> + %0 = stablehlo.concatenate %arg0, %arg1, dim = 0 : (tensor<2x4xf32>, tensor<0x4xf32>) -> tensor<2x4xf32> + return %0 : tensor<2x4xf32> +} + diff --git a/third_party/xla/xla/pjrt/c/BUILD b/third_party/xla/xla/pjrt/c/BUILD index 0a0b99e4367351..1e3ea0d8096bf8 100644 --- a/third_party/xla/xla/pjrt/c/BUILD +++ b/third_party/xla/xla/pjrt/c/BUILD @@ -452,7 +452,6 @@ cc_library( "//xla/pjrt:pjrt_device_dimensions", "//xla/pjrt:pjrt_executable", "//xla/pjrt:pjrt_layout", - "//xla/pjrt:pjrt_stream_executor_client", "//xla/pjrt:raw_buffer", "//xla/pjrt:scoped_async_tracking_event", "//xla/pjrt/distributed:key_value_store_interface", diff --git a/third_party/xla/xla/pjrt/cpu/BUILD b/third_party/xla/xla/pjrt/cpu/BUILD index bb92234b2867c7..87cb2fc05025cf 100644 --- a/third_party/xla/xla/pjrt/cpu/BUILD +++ b/third_party/xla/xla/pjrt/cpu/BUILD @@ -185,6 +185,8 @@ cc_library( "//xla:executable_run_options", "//xla:future", "//xla:literal", + "//xla:literal_util", + "//xla:shape_layout", "//xla:shape_util", "//xla:util", "//xla:xla_data_proto_cc", @@ -208,7 +210,6 @@ cc_library( "//xla/pjrt:dynamic_shapes", "//xla/pjrt:host_callback", "//xla/pjrt:host_memory_spaces", - "//xla/pjrt:infer_dispatch_info", "//xla/pjrt:layout_mode", "//xla/pjrt:maybe_owning_mlir_module", "//xla/pjrt:mlir_to_hlo", @@ -229,7 +230,6 @@ cc_library( "//xla/pjrt/plugin/xla_cpu:cpu_topology", "//xla/pjrt/plugin/xla_cpu:cpu_topology_description", "//xla/pjrt/proto:compile_options_proto_cc", - "//xla/runtime:chip_id", "//xla/runtime:device_id", "//xla/service:buffer_assignment", "//xla/service:compiled_module", @@ -244,7 +244,7 @@ cc_library( "//xla/service:hlo_proto_cc", "//xla/service:hlo_value", "//xla/service:maybe_owning_device_address", - "//xla/service/cpu:cpu_compiler_pure", + "//xla/service/cpu:cpu_compiler", "//xla/service/cpu:cpu_executable", "//xla/service/cpu:cpu_executable_run_options", "//xla/service/cpu:cpu_xfeed", diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc index 0873c87b8560c9..7ef56465262556 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc @@ -21,6 +21,8 @@ limitations under the License. #include #include #include +#include +#include #include #include #include @@ -65,6 +67,7 @@ limitations under the License. #include "xla/layout.h" #include "xla/layout_util.h" #include "xla/literal.h" +#include "xla/literal_util.h" #include "xla/pjrt/common_pjrt_client.h" #include "xla/pjrt/compiled_memory_stats.h" #include "xla/pjrt/cpu/abstract_cpu_buffer.h" @@ -79,7 +82,7 @@ limitations under the License. #include "xla/pjrt/dynamic_shapes.h" #include "xla/pjrt/host_callback.h" #include "xla/pjrt/host_memory_spaces.h" -#include "xla/pjrt/infer_dispatch_info.h" +#include "xla/pjrt/host_to_device_transfer_manager.h" #include "xla/pjrt/layout_mode.h" #include "xla/pjrt/maybe_owning_mlir_module.h" #include "xla/pjrt/mlir_to_hlo.h" @@ -94,11 +97,10 @@ limitations under the License. #include "xla/pjrt/plugin/xla_cpu/cpu_topology_description.h" #include "xla/pjrt/proto/compile_options.pb.h" #include "xla/pjrt/raw_buffer.h" -#include "xla/pjrt/raw_pjrt_client.h" #include "xla/pjrt/semaphore.h" #include "xla/pjrt/thread_pool_async_work_runner.h" #include "xla/pjrt/utils.h" -#include "xla/runtime/chip_id.h" +#include "xla/primitive_util.h" #include "xla/runtime/device_id.h" #include "xla/service/buffer_assignment.h" #include "xla/service/compiled_module.h" @@ -119,6 +121,7 @@ limitations under the License. #include "xla/service/llvm_ir/llvm_command_line_options.h" #include "xla/service/maybe_owning_device_address.h" #include "xla/shape.h" +#include "xla/shape_layout.h" #include "xla/shape_util.h" #include "xla/stream_executor/device_address.h" #include "xla/tsl/concurrency/async_value.h" @@ -131,6 +134,7 @@ limitations under the License. #include "xla/xla.pb.h" #include "tsl/platform/denormal.h" #include "tsl/platform/fingerprint.h" +#include "tsl/platform/protobuf.h" #include "tsl/platform/setround.h" #include "tsl/profiler/lib/traceme.h" @@ -367,7 +371,7 @@ absl::StatusOr> GetPjRtCpuClient( std::move(allocator), std::move(options.collectives), num_threads, options.asynchronous, options.max_transpose_threads, std::move(options.customize_hlo_module_config), cpu_device_count, - options.max_inflight_computations_per_device, options.intra_op_device); + options.max_inflight_computations_per_device); return CreatePjRtCpuClient(std::move(raw_client), std::move(topology), options.process_id); @@ -436,8 +440,7 @@ PjRtCpuRawClient::PjRtCpuRawClient( std::shared_ptr collectives, size_t num_threads, bool asynchronous, int max_transpose_threads, std::function customize_hlo_module_config, - int cpu_device_count, int max_inflight_computations, - const Eigen::ThreadPoolDevice* intra_op_device) + int cpu_device_count, int max_inflight_computations) : local_device_states_( MakeLocalDeviceStates(cpu_device_count, max_inflight_computations)), allocator_(std::move(allocator)), @@ -447,18 +450,12 @@ PjRtCpuRawClient::PjRtCpuRawClient( last_collective_launch_event_( tsl::MakeAvailableAsyncValueRef()), customize_hlo_module_config_(std::move(customize_hlo_module_config)), - eigen_intraop_pool_(intra_op_device == nullptr - ? std::make_unique( - tsl::Env::Default(), GetThreadOptions(), - "XLAEigen", - std::min(num_threads, kMaxIntraOpThreads)) - : std::unique_ptr()), - eigen_intraop_device_(intra_op_device == nullptr - ? std::make_unique( - eigen_intraop_pool_->AsEigenThreadPool(), - eigen_intraop_pool_->NumThreads()) - : std::unique_ptr()), - custom_intraop_device_(intra_op_device), + eigen_intraop_pool_(new tsl::thread::ThreadPool( + tsl::Env::Default(), GetThreadOptions(), "XLAEigen", + std::min(num_threads, kMaxIntraOpThreads))), + eigen_intraop_device_( + new Eigen::ThreadPoolDevice(eigen_intraop_pool_->AsEigenThreadPool(), + eigen_intraop_pool_->NumThreads())), compile_thread_pool_(std::make_unique( tsl::Env::Default(), GetThreadOptions(), "XLACompile", num_threads)), execute_work_runner_(std::make_unique( @@ -1508,6 +1505,7 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute( run_options.set_run_id(run_id_); // Need to keep device_assignment alive until execution completes. run_options.set_device_assignment(device_assignment_.get()); + run_options.set_intra_op_thread_pool(raw_client->eigen_intraop_device()); run_options.set_rng_seed(options.seed); auto cpu_run_options = std::make_unique(); @@ -1517,14 +1515,6 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute( options.context == nullptr ? nullptr : dynamic_cast(options.context); - - const Eigen::ThreadPoolDevice* intra_op_device = - (cpu_execute_context != nullptr && - cpu_execute_context->intra_op_thread_pool() != nullptr) - ? cpu_execute_context->intra_op_thread_pool() - : raw_client->eigen_intraop_device(); - run_options.set_intra_op_thread_pool(intra_op_device); - if (cpu_execute_context != nullptr && cpu_execute_context->process_index().has_value()) { run_options.set_device_ordinal( @@ -1588,6 +1578,7 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute( } auto execute_thunks = [cpu_executable, buffer_table = std::move(buffer_table), + eigen_device = raw_client->eigen_intraop_device(), run_options = std::move(run_options)]() -> absl::StatusOr> { // Set denormal and rounding behavior to match the default TF diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h index c467ebe46834be..b7539a1a9948fd 100644 --- a/third_party/xla/xla/pjrt/cpu/cpu_client.h +++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h @@ -106,8 +106,7 @@ class PjRtCpuRawClient : public PjRtRawClient { std::shared_ptr collectives, size_t num_threads, bool asynchronous, int max_transpose_threads, std::function customize_hlo_module_config, - int cpu_device_count, int max_inflight_computations, - const Eigen::ThreadPoolDevice* intra_op_device = nullptr); + int cpu_device_count, int max_inflight_computations); ~PjRtCpuRawClient() override; @@ -134,9 +133,8 @@ class PjRtCpuRawClient : public PjRtRawClient { return eigen_intraop_pool_.get(); } - const Eigen::ThreadPoolDevice* eigen_intraop_device() const { - return custom_intraop_device_ != nullptr ? custom_intraop_device_ - : eigen_intraop_device_.get(); + Eigen::ThreadPoolDevice* eigen_intraop_device() const { + return eigen_intraop_device_.get(); } cpu::CpuCollectives* collectives() const { return collectives_.get(); } @@ -299,7 +297,6 @@ class PjRtCpuRawClient : public PjRtRawClient { // the member variables of this class that are already destroyed. std::unique_ptr eigen_intraop_pool_; std::unique_ptr eigen_intraop_device_; - const Eigen::ThreadPoolDevice* custom_intraop_device_ = nullptr; std::unique_ptr compile_thread_pool_; std::unique_ptr execute_work_runner_; std::unique_ptr async_work_runner_; diff --git a/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_client_options.h b/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_client_options.h index ddd3343f606357..3af357c6f86763 100644 --- a/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_client_options.h +++ b/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_client_options.h @@ -27,10 +27,6 @@ limitations under the License. #include "xla/pjrt/plugin/xla_cpu/cpu_topology_description.h" #include "xla/service/hlo_module_config.h" -namespace Eigen { -struct ThreadPoolDevice; -} // namespace Eigen - namespace xla { // Options for creating an XLA:CPU PjRtClient. @@ -74,10 +70,6 @@ struct CpuClientOptions { // detected from the host and will try to use cpu_device_count and process_id // to build the topology. const CpuTopologyDescription* topology = nullptr; - - // Intra-op thread pool device. Optional. If not provided, a default thread - // pool will be created. Does not take ownership. Must outlive the client. - const Eigen::ThreadPoolDevice* intra_op_device = nullptr; }; } // namespace xla diff --git a/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_execute_options.h b/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_execute_options.h index 9cf40d39ac43af..665aefb240c102 100644 --- a/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_execute_options.h +++ b/third_party/xla/xla/pjrt/plugin/xla_cpu/cpu_execute_options.h @@ -21,10 +21,6 @@ limitations under the License. #include "xla/backends/cpu/collectives/cpu_collectives.h" #include "xla/pjrt/pjrt_executable.h" -namespace Eigen { -struct ThreadPoolDevice; -} // namespace Eigen - namespace xla { // ExecuteContext for XLA:CPU PjRtLoadedExecutable::Execute calls. @@ -51,21 +47,9 @@ class CpuExecuteContext : public ExecuteContext { cpu::CpuCollectives*& collectives() { return collectives_; } cpu::CpuCollectives* collectives() const { return collectives_; } - // If specified, override the intra-op thread pool for a particular call of - // PjRtLoadedExecutable::Execute. Must remain valid until the execution - // finishes. - const Eigen::ThreadPoolDevice* intra_op_thread_pool() const { - return intra_op_thread_pool_; - } - void set_intra_op_thread_pool( - const Eigen::ThreadPoolDevice* intra_op_thread_pool) { - intra_op_thread_pool_ = intra_op_thread_pool; - } - private: std::optional process_index_; cpu::CpuCollectives* collectives_ = nullptr; - const Eigen::ThreadPoolDevice* intra_op_thread_pool_ = nullptr; }; } // namespace xla diff --git a/third_party/xla/xla/pjrt/tf_pjrt_client.h b/third_party/xla/xla/pjrt/tf_pjrt_client.h index 61295928948a76..0943e22bddc5a1 100644 --- a/third_party/xla/xla/pjrt/tf_pjrt_client.h +++ b/third_party/xla/xla/pjrt/tf_pjrt_client.h @@ -182,14 +182,6 @@ class TfPjRtExecutable : public PjRtLoadedExecutable { void Delete() override { return wrapped_->Delete(); } bool IsDeleted() const override { return wrapped_->IsDeleted(); } - PjRtExecutable* GetExecutable() const override { - return wrapped_->GetExecutable(); - } - - absl::StatusOr GetCompiledMemoryStats() const override { - return wrapped_->GetCompiledMemoryStats(); - } - absl::StatusOr SerializeExecutable() const override { return wrapped_->SerializeExecutable(); } diff --git a/third_party/xla/xla/python/ifrt/BUILD b/third_party/xla/xla/python/ifrt/BUILD index 54b06a7eaf994c..6a6fd39440985c 100644 --- a/third_party/xla/xla/python/ifrt/BUILD +++ b/third_party/xla/xla/python/ifrt/BUILD @@ -583,8 +583,10 @@ cc_library( visibility = internal_visibility([ ":friends", ":internal", + ":users", ]), deps = [ + ":attribute_map", ":basic_device_list", ":ifrt", ":mock", @@ -783,9 +785,7 @@ cc_library( srcs = ["mock.cc"], hdrs = ["mock.h"], visibility = internal_visibility([ - ":friends", - ":internal", - ":users", + ":mock_users", ]), deps = [ ":attribute_map", diff --git a/third_party/xla/xla/python/ifrt/device_test_util.cc b/third_party/xla/xla/python/ifrt/device_test_util.cc index e9b847b0f64635..ad113cd5ac052a 100644 --- a/third_party/xla/xla/python/ifrt/device_test_util.cc +++ b/third_party/xla/xla/python/ifrt/device_test_util.cc @@ -23,6 +23,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/status/statusor.h" #include "absl/types/span.h" +#include "xla/python/ifrt/attribute_map.h" #include "xla/python/ifrt/basic_device_list.h" #include "xla/python/ifrt/device.h" #include "xla/python/ifrt/device_list.h" @@ -60,6 +61,7 @@ struct DeviceTestClientState { std::vector devices; std::vector addressable_devices; std::vector> device_memories; + std::vector device_attributes; }; // Creates a mock client for device tests. The client will have a specified @@ -67,9 +69,16 @@ struct DeviceTestClientState { // `devices()` and `LookupDevice()`. Device implements `id()`, with an // arbitrary deterministic device ids assigned. Each device has "device" memory // (which is also its default memory), and each memory has a single device. -std::shared_ptr MakeDeviceTestClient(int num_devices, - int num_addressable_devices) { +std::shared_ptr MakeDeviceTestClient(const DeviceTestParam& param) { + const int num_devices = param.num_devices; + const int num_addressable_devices = param.num_addressable_devices; CHECK_GE(num_devices, num_addressable_devices); + if (!param.device_ids.empty()) { + CHECK_EQ(param.device_ids.size(), num_devices); + } + if (!param.device_attributes.empty()) { + CHECK_EQ(param.device_attributes.size(), num_devices); + } auto state = std::make_shared(); state->default_memory_kind = MemoryKind(); @@ -83,6 +92,17 @@ std::shared_ptr MakeDeviceTestClient(int num_devices, state->addressable_devices.reserve(num_addressable_devices); state->device_memories.resize(num_devices); + if (param.device_attributes.empty()) { + state->device_attributes.resize(num_devices, + AttributeMap(AttributeMap::Map{})); + } else { + state->device_attributes.reserve(num_devices); + for (int i = 0; i < num_devices; ++i) { + state->device_attributes.push_back( + AttributeMap(param.device_attributes[i])); + } + } + for (int i = 0; i < num_devices; ++i) { const bool addressable = i < num_addressable_devices; auto memory = std::make_unique(); @@ -94,20 +114,24 @@ std::shared_ptr MakeDeviceTestClient(int num_devices, state->memories.push_back(memory.get()); state->memory_map.insert({MemoryId(i + 10), std::move(memory)}); + DeviceId device_id = + param.device_ids.empty() ? DeviceId(i + 10) : param.device_ids[i]; auto device = std::make_unique(); // client will be filled in at the end of the loop. ON_CALL(*device, client).WillByDefault(ReturnPointee(&state->client)); - ON_CALL(*device, Id).WillByDefault(Return(DeviceId(i + 10))); + ON_CALL(*device, Id).WillByDefault(Return(device_id)); ON_CALL(*device, IsAddressable).WillByDefault(Return(addressable)); ON_CALL(*device, DefaultMemory).WillByDefault(Return(state->memories[i])); // device_memories will be filled in at the end of the loop. ON_CALL(*device, Memories) .WillByDefault(ReturnPointee(&state->device_memories[i])); + ON_CALL(*device, Attributes) + .WillByDefault(ReturnRef(state->device_attributes[i])); state->devices.push_back(device.get()); if (addressable) { state->addressable_devices.push_back(device.get()); } - state->device_map.insert({DeviceId(i + 10), std::move(device)}); + CHECK(state->device_map.insert({device_id, std::move(device)}).second); state->device_memories[i] = absl::MakeConstSpan(&state->memories[i], 1); state->memory_devices[i] = absl::MakeConstSpan(&state->devices[i], 1); @@ -142,8 +166,7 @@ std::shared_ptr MakeDeviceTestClient(int num_devices, } // namespace DeviceTestFixture::DeviceTestFixture(const DeviceTestParam& param) { - const auto& [num_devices, num_addressable_devices] = param; - client_ = MakeDeviceTestClient(num_devices, num_addressable_devices); + client_ = MakeDeviceTestClient(param); } DeviceListRef DeviceTestFixture::GetDevices( diff --git a/third_party/xla/xla/python/ifrt/device_test_util.h b/third_party/xla/xla/python/ifrt/device_test_util.h index fd1933c49501a3..e701c66d6f8478 100644 --- a/third_party/xla/xla/python/ifrt/device_test_util.h +++ b/third_party/xla/xla/python/ifrt/device_test_util.h @@ -17,8 +17,10 @@ limitations under the License. #define XLA_PYTHON_IFRT_DEVICE_TEST_UTIL_H_ #include +#include #include "absl/types/span.h" +#include "xla/python/ifrt/attribute_map.h" #include "xla/python/ifrt/client.h" #include "xla/python/ifrt/device.h" #include "xla/python/ifrt/device_list.h" @@ -33,6 +35,14 @@ namespace test_util { struct DeviceTestParam { int num_devices; int num_addressable_devices; + + // Optional: device IDs. If empty, defaults to DeviceId(i + 10). + // If non-empty, must have size `num_devices`. + std::vector device_ids; + + // Optional: per-device attributes. If non-empty, must have size + // `num_devices`. + std::vector device_attributes; }; // Test fixture for device tests. diff --git a/third_party/xla/xla/python/ifrt/ifrt.bzl b/third_party/xla/xla/python/ifrt/ifrt.bzl index f7217a5df6b3c6..e805d7cf7d7d47 100644 --- a/third_party/xla/xla/python/ifrt/ifrt.bzl +++ b/third_party/xla/xla/python/ifrt/ifrt.bzl @@ -50,3 +50,8 @@ def ifrt_package_groups(name = "ifrt_package_groups"): name = "serdes_week_4_old_version_users", packages = ["//..."], ) + + native.package_group( + name = "mock_users", + packages = ["//..."], + ) diff --git a/third_party/xla/xla/python/ifrt/ir/tests/ifrt_compile_atom_program.mlir b/third_party/xla/xla/python/ifrt/ir/tests/ifrt_compile_atom_program.mlir index 20439d0edcfc82..629d50e974e679 100644 --- a/third_party/xla/xla/python/ifrt/ir/tests/ifrt_compile_atom_program.mlir +++ b/third_party/xla/xla/python/ifrt/ir/tests/ifrt_compile_atom_program.mlir @@ -71,3 +71,38 @@ module @call_hlo_sdy_lowered attributes { } } } + +// ----- + +!array0 = !ifrt.array, + #ifrt.sharding_param<2x1 to [0] on 2>, [0,1]> +!array1 = !ifrt.array, + #ifrt.sharding_param<2x1 to [0] on 2>, [2,3]> +// CHECK: #sp = #ifrt.sharding_param<2x1 to [0] on 2> +// CHECK-LABEL: @call_hlo_multiple_calls_same_target_module +module @call_hlo_multiple_calls_same_target_module { + func.func @main(%arg0: !array0, %arg1: !array1) -> (!array0, !array1) attributes {ifrt.function} { + // CHECK: ifrt.CallLoadedExecutable @{{.*}}(%arg0) + %0, %ctrl_0 = ifrt.Call @add_one::@main(%arg0) on devices [0,1] + {ifrt.module_type = "xla"} : (!array0) -> !array0 + // CHECK: ifrt.CallLoadedExecutable @{{.*}}(%arg1) + %1, %ctrl_1 = ifrt.Call @add_one::@main(%arg1) on devices [2,3] + {ifrt.module_type = "xla"} : (!array1) -> !array1 + return %0, %1 : !array0, !array1 + } + + // CHECK: ifrt.LoadedExecutable @ + // CHECK-SAME: on devices [2, 3] + // CHECK: ifrt.LoadedExecutable @ + // CHECK-SAME: on devices [0, 1] + module @add_one attributes {sym_visibility = "private"} { + func.func @main( + %arg0: tensor<2x2xi32> {mhlo.sharding = "{devices=[2,1]<=[2]}"}) + -> (tensor<2x2xi32> {mhlo.sharding = "{devices=[2,1]<=[2]}"}) { + %0 = mhlo.constant dense<1> : tensor<2x2xi32> + %1 = mhlo.add %arg0, %0 : tensor<2x2xi32> + return %1 : tensor<2x2xi32> + } + } +} + diff --git a/third_party/xla/xla/python/ifrt/ir/transforms/ifrt_compile_atom_program_pass.cc b/third_party/xla/xla/python/ifrt/ir/transforms/ifrt_compile_atom_program_pass.cc index 8636cfa2b1a5d4..65d63b56a982d5 100644 --- a/third_party/xla/xla/python/ifrt/ir/transforms/ifrt_compile_atom_program_pass.cc +++ b/third_party/xla/xla/python/ifrt/ir/transforms/ifrt_compile_atom_program_pass.cc @@ -28,6 +28,7 @@ limitations under the License. #include "absl/status/status_macros.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" +#include "llvm/ADT/DenseMap.h" #include "llvm/ADT/MapVector.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" @@ -130,7 +131,10 @@ class IfrtCompileAtomProgramPass // Depending on the type of module, a MLIR pipeline might be executed before // the compilation is dispatched. absl::StatusOr CompileModule( - CallOp call_op, mlir::ModuleOp module_op); + CallOp call_op, mlir::ModuleOp module_op, + llvm::DenseMap>& + cloned_modules, + const mlir::SymbolUserMap& symbol_users); // Gets the XLA compile options for the given atom program module. absl::StatusOr GetXlaCompileOptions( @@ -145,8 +149,11 @@ class IfrtCompileAtomProgramPass // // Note that the method runs `ifrt-compile-xla-preprocessing-pipeline` // before dispatching compilation. - absl::StatusOr CompileXla(CallOp call_op, - mlir::ModuleOp module_op); + absl::StatusOr CompileXla( + CallOp call_op, mlir::ModuleOp module_op, + llvm::DenseMap>& + cloned_modules, + const mlir::SymbolUserMap& symbol_users); // Returns a future of a AtomProgramCompileResult for the MPMD reshard module. absl::StatusOr CompileMpmdReshard( @@ -244,7 +251,10 @@ std::vector> GetAtomProgramCallbacks( } absl::StatusOr IfrtCompileAtomProgramPass::CompileXla( - CallOp call_op, mlir::ModuleOp module_op) { + CallOp call_op, mlir::ModuleOp module_op, + llvm::DenseMap>& + cloned_modules, + const mlir::SymbolUserMap& symbol_users) { ABSL_ASSIGN_OR_RETURN(XlaCompileOptions * xla_compile_options, GetXlaCompileOptions(call_op, module_op)); @@ -262,8 +272,28 @@ absl::StatusOr IfrtCompileAtomProgramPass::CompileXla( // because MLIR printing takes different paths depending on if a ModuleOp // has a parent or not. Thus, by cloning the module we ensure that the // module's string representation is maintained. - ABSL_ASSIGN_OR_RETURN(mlir::OwningOpRef cloned_module, - CloneModuleIntoContext(module_op, *hlo_program_context_)); + // For modules with a single reference, clone directly into the context to + // avoid any extra clone overhead. For modules with multiple references, + // cache the cloned module in the new context and use ModuleOp::clone() to + // avoid repetitive bytecode roundtrips. + bool is_single_reference = symbol_users.getUsers(module_op).size() <= 1; + + mlir::OwningOpRef cloned_module; + if (is_single_reference) { + ABSL_ASSIGN_OR_RETURN(cloned_module, + CloneModuleIntoContext(module_op, *hlo_program_context_)); + } else { + auto it = cloned_modules.find(module_op); + if (it == cloned_modules.end()) { + ABSL_ASSIGN_OR_RETURN( + mlir::OwningOpRef initial_cloned_module, + CloneModuleIntoContext(module_op, *hlo_program_context_)); + it = cloned_modules + .try_emplace(module_op, std::move(initial_cloned_module)) + .first; + } + cloned_module = it->second.get().clone(); + } auto hlo_program = std::make_unique(hlo_program_context_, std::move(cloned_module)); AtomProgramCompileResult result; @@ -310,8 +340,11 @@ IfrtCompileAtomProgramPass::CompileMpmdReshard(mlir::ModuleOp module_op) { } absl::StatusOr -IfrtCompileAtomProgramPass::CompileModule(CallOp call_op, - mlir::ModuleOp module_op) { +IfrtCompileAtomProgramPass::CompileModule( + CallOp call_op, mlir::ModuleOp module_op, + llvm::DenseMap>& + cloned_modules, + const mlir::SymbolUserMap& symbol_users) { auto module_type = call_op->getAttrOfType(kIfrtModuleTypeAttrName); if (module_type == nullptr) { @@ -319,7 +352,7 @@ IfrtCompileAtomProgramPass::CompileModule(CallOp call_op, "CallOp requires `", kIfrtModuleTypeAttrName.str(), "` to be set")); } if (module_type == kIfrtModuleTypeXla) { - return CompileXla(call_op, module_op); + return CompileXla(call_op, module_op, cloned_modules, symbol_users); } if (module_type == kIfrtModuleTypeMpmdReshard) { return CompileMpmdReshard(module_op); @@ -335,6 +368,7 @@ void IfrtCompileAtomProgramPass::runOnOperation() { llvm::DenseMap call_to_compile_results; mlir::ModuleOp module_op = getOperation(); + mlir::SymbolUserMap symbol_users(symbol_table, module_op); mlir::Attribute sdy_meshes_round_trip_attr = module_op->getAttr(kIfrtSdyMeshesRoundTripAttr); @@ -346,6 +380,10 @@ void IfrtCompileAtomProgramPass::runOnOperation() { // while dispatching a compilation. llvm::MapVector call_op_to_error; + // Map from target module to cloned module in the compilation context. + llvm::DenseMap> + cloned_modules; + // Walk and dispatch the compilations in parallel. module_op.walk([&](CallOp call_op) -> mlir::WalkResult { UserContextScope user_context_scope(user_context_); @@ -379,7 +417,7 @@ void IfrtCompileAtomProgramPass::runOnOperation() { } absl::StatusOr compile_result = - CompileModule(call_op, callee_module); + CompileModule(call_op, callee_module, cloned_modules, symbol_users); if (!compile_result.ok()) { call_op_to_error.try_emplace( call_op, diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc index 1917f0008ccce9..4c14b6041e660f 100644 --- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc +++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc @@ -422,14 +422,14 @@ std::vector GatherHloOutputCallbacks( // Options consumed by the PjRt-IFRT layer itself are not forwarded. xla::CustomOptions::Map ToCustomOptionsMap(const AttributeMap& attribute_map) { xla::CustomOptions::Map custom_options; + custom_options.reserve(attribute_map.size()); attribute_map.ForEach([&](const std::string& name, const AttributeMap::Value& value) { - if (name == "use_output_arena" || - name == PjRtCompatibleLoadedExecutable::kCallLocation) { - return; - } - std::visit([&](const auto& v) { custom_options[name] = v.value; }, value); + std::visit([&](const auto& v) { custom_options.insert({name, v.value}); }, + value); }); + custom_options.erase("use_output_arena"); + custom_options.erase(PjRtCompatibleLoadedExecutable::kCallLocation); return custom_options; } diff --git a/third_party/xla/xla/python/version.h b/third_party/xla/xla/python/version.h index da24e715b82b33..9e7b5e5b37d9e3 100644 --- a/third_party/xla/xla/python/version.h +++ b/third_party/xla/xla/python/version.h @@ -18,7 +18,6 @@ limitations under the License. // An increasing version number to protect jax code against breaking changes. // In JAX, reference this via jax._src.lib.ifrt_version. -#define JAX_IFRT_VERSION_NUMBER \ - 71 // Fix use-after-free in CopyArraysToHostBufferShards +#define JAX_IFRT_VERSION_NUMBER 72 // unreduced + partial manual shard_map #endif // XLA_PYTHON_VERSION_H_ diff --git a/third_party/xla/xla/runtime/buffer_use.cc b/third_party/xla/xla/runtime/buffer_use.cc index 8e08b8afbda1df..d95070da284b41 100644 --- a/third_party/xla/xla/runtime/buffer_use.cc +++ b/third_party/xla/xla/runtime/buffer_use.cc @@ -51,16 +51,9 @@ void BufferUse::ReadWriteSet::AddAll(absl::Span uses) { } bool BufferUse::ReadWriteSet::HasConflicts(const BufferUse& use) const { - // A zero-size slice does not refer to any actual memory, so it can never - // alias another slice, even one with the same allocation and offset (e.g. - // buffer assignment may colocate unrelated zero-size scratch/workspace - // buffers at the same degenerate offset). - if (use.slice().size() == 0) return false; - // Returns true if `use` overlaps with any of the slices in set. auto overlaps = [](const std::vector& set, const BufferUse& use) { return absl::c_any_of(set, [&](const BufferUse& other) { - if (other.slice_.size() == 0) return false; return other.slice_.OverlapsWith(use.slice()) || other.slice_ == use.slice_; }); diff --git a/third_party/xla/xla/runtime/buffer_use_test.cc b/third_party/xla/xla/runtime/buffer_use_test.cc index 62a3ec92e6f171..78c9b7b6e7460d 100644 --- a/third_party/xla/xla/runtime/buffer_use_test.cc +++ b/third_party/xla/xla/runtime/buffer_use_test.cc @@ -107,32 +107,5 @@ TEST(BufferUseTest, ReadWriteSet) { EXPECT_TRUE(rwset.HasConflicts({BufferUse::Write(slice2, slice_shape)})); } -TEST(BufferUseTest, ZeroSizeSlicesDoNotConflict) { - BufferUse::ReadWriteSet rwset; - - BufferAllocation alloc(/*index=*/0, /*size=*/1024, /*color=*/0); - Shape zero_size_shape = ShapeUtil::MakeShape(U8, {0}); - - // Two independent zero-size buffers (e.g. workspace outputs of unrelated - // custom-calls) can be assigned the exact same degenerate {index, offset, - // size=0} slice by buffer assignment. Since they don't refer to any actual - // memory, they must never be reported as conflicting. - BufferAllocation::Slice zero_slice(&alloc, /*offset=*/16, /*size=*/0); - - rwset.Add(BufferUse::Write(zero_slice, zero_size_shape)); - EXPECT_FALSE( - rwset.HasConflicts({BufferUse::Write(zero_slice, zero_size_shape)})); - EXPECT_FALSE( - rwset.HasConflicts({BufferUse::Read(zero_slice, zero_size_shape)})); - - // A non-zero-size slice at the same offset should still be able to - // conflict normally. - Shape slice_shape = ShapeUtil::MakeShape(F32, {2}); - BufferAllocation::Slice non_zero_slice(&alloc, /*offset=*/16, /*size=*/8); - rwset.Add(BufferUse::Write(non_zero_slice, slice_shape)); - EXPECT_TRUE( - rwset.HasConflicts({BufferUse::Write(non_zero_slice, slice_shape)})); -} - } // namespace } // namespace xla diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 45d7e446bcae49..8b3ce74bd0d0da 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -1640,7 +1640,6 @@ cc_library( "//xla/backends/gpu/tests:__pkg__", "//xla/backends/gpu/transforms:__subpackages__", "//third_party/pathways/configuration:__subpackages__", - "//learning/brain/google/xla:__pkg__", ]), deps = [ ":buffer_value", @@ -4659,9 +4658,11 @@ xla_cc_test( name = "host_offload_utils_test", srcs = ["host_offload_utils_test.cc"], deps = [ + ":call_graph", ":host_offload_utils", "//xla:shape_util", "//xla:util", + "//xla/hlo/ir:hlo", "//xla/hlo/testlib:hlo_hardware_independent_test_base", "//xla/tests:xla_internal_test_main", "@com_google_absl//absl/container:flat_hash_set", diff --git a/third_party/xla/xla/service/buffer_assignment.h b/third_party/xla/xla/service/buffer_assignment.h index 14b28c1dfb1ef6..7c161803fa820c 100644 --- a/third_party/xla/xla/service/buffer_assignment.h +++ b/third_party/xla/xla/service/buffer_assignment.h @@ -222,6 +222,7 @@ class BufferAllocation { Index index() const { return allocation_->index(); } int64_t offset() const { return offset_; } int64_t size() const { return size_; } + bool empty() const { return size_ == 0; } PrimitiveType element_type() const { return element_type_; } bool operator==(const Slice& other) const { diff --git a/third_party/xla/xla/service/buffer_assignment_test.cc b/third_party/xla/xla/service/buffer_assignment_test.cc index 68b1f8b58934c4..323b2fee808678 100644 --- a/third_party/xla/xla/service/buffer_assignment_test.cc +++ b/third_party/xla/xla/service/buffer_assignment_test.cc @@ -4213,6 +4213,18 @@ TEST(BufferAllocationSliceProtoTest, FromProtoErrorSliceOutOfRange) { "range for allocation #0 of size 200"))); } +TEST(BufferAllocationSliceTest, Empty) { + BufferAllocation alloc(/*index=*/0, /*size=*/100, /*color=*/0); + BufferAllocation::Slice default_slice; + EXPECT_TRUE(default_slice.empty()); + + BufferAllocation::Slice empty_slice(&alloc, /*offset=*/0, /*size=*/0); + EXPECT_TRUE(empty_slice.empty()); + + BufferAllocation::Slice non_empty_slice(&alloc, /*offset=*/0, /*size=*/10); + EXPECT_FALSE(non_empty_slice.empty()); +} + TEST(BufferAllocationTest, ToProto) { BufferAllocation alloc{42, 64, 3}; alloc.set_constant(true); diff --git a/third_party/xla/xla/service/cpu/tests/BUILD b/third_party/xla/xla/service/cpu/tests/BUILD index c8b04f975f4215..631d56b3679fe8 100644 --- a/third_party/xla/xla/service/cpu/tests/BUILD +++ b/third_party/xla/xla/service/cpu/tests/BUILD @@ -56,7 +56,6 @@ cc_library( "//xla/service:executable", "//xla/service:llvm_compiler", "//xla/service:platform_util", - "//xla/service/cpu:cpu_compiler", "//xla/stream_executor:platform", "//xla/tests:codegen_utils", "//xla/tests:hlo_pjrt_test_base", diff --git a/third_party/xla/xla/service/cpu/thunk_emitter.cc b/third_party/xla/xla/service/cpu/thunk_emitter.cc index 13ae4aa5e6d1e7..23937c970fc2e9 100644 --- a/third_party/xla/xla/service/cpu/thunk_emitter.cc +++ b/third_party/xla/xla/service/cpu/thunk_emitter.cc @@ -149,7 +149,8 @@ absl::StatusOr GetFusionFingerprint( } // namespace static FusionCompiler::Options FusionCompilerOptions( - const HloModuleConfig& config) { + const HloModuleConfig& config, + const TargetMachineFeatures& target_machine_features) { const DebugOptions& debug_options = config.debug_options(); return FusionCompiler::Options{ debug_options.xla_cpu_prefer_vector_width(), @@ -157,12 +158,15 @@ static FusionCompiler::Options FusionCompilerOptions( debug_options.xla_cpu_enable_fast_min_max(), llvm_ir::GetCpuFastMathFlags(config), debug_options.xla_cpu_use_new_xtile_lowering(), - options::IsMsanEnabled(config)}; + options::IsMsanEnabled(config), + target_machine_features.get_target_feature_string()}; } -static FusionCompiler FusionCompilerFactory(mlir::MLIRContext* context, - const HloModule& hlo_module) { - FusionCompiler::Options options = FusionCompilerOptions(hlo_module.config()); +static FusionCompiler FusionCompilerFactory( + mlir::MLIRContext* context, const HloModule& hlo_module, + const TargetMachineFeatures& target_machine_features) { + FusionCompiler::Options options = + FusionCompilerOptions(hlo_module.config(), target_machine_features); return FusionCompiler(context, std::move(options), &hlo_module); } @@ -179,10 +183,12 @@ ThunkEmitter::ThunkEmitter(IrEmitter2& ir_emitter, communicator_resource_( Resource::Create(Resource::kCollectiveCommunicator)), mlir_context_(FusionCompiler::CreateContext()), - fusion_compiler_(FusionCompilerFactory(mlir_context_.get(), hlo_module)), + fusion_compiler_(FusionCompilerFactory(mlir_context_.get(), hlo_module, + target_machine_features)), parallel_fusion_emitter_( - thread_pool, FusionCompilerOptions(hlo_module_config_), &hlo_module, - &buffer_assignment, + thread_pool, + FusionCompilerOptions(hlo_module_config_, target_machine_features), + &hlo_module, &buffer_assignment, hlo_module_config_.debug_options() .xla_cpu_generate_unique_c_style_kernel_entry_points(), options::EnableTiledEmitter(hlo_module_config_)) {} diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD index 394417048f9d21..9a1d99fca60280 100644 --- a/third_party/xla/xla/service/gpu/BUILD +++ b/third_party/xla/xla/service/gpu/BUILD @@ -1742,6 +1742,7 @@ cc_library( ":alias_info", "//xla:xla_proto_cc", "//xla/backends/gpu/transforms:alias_passthrough_params", + "//xla/backends/gpu/transforms:constant_fill_copy_rewriter", "//xla/backends/gpu/transforms:copy_fusion", "//xla/backends/gpu/transforms:sanitize_constant_names", "//xla/hlo/ir:hlo", @@ -2308,7 +2309,6 @@ cc_library( deps = [ ":gpu_executable", ":gpu_executable_proto_cc", - "//xla:debug_options_flags", "//xla:printer", "//xla:xla_proto_cc", "//xla/hlo/ir:hlo", @@ -3629,6 +3629,7 @@ cc_library( hdrs = ["execution_stream_assignment.h"], deps = [ ":gpu_latency_hiding_scheduler", + "//xla:shape_util", "//xla:side_effect_util", "//xla/backends/gpu/runtime:execution_stream_id", "//xla/backends/gpu/transforms/collectives:collective_domain", diff --git a/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc b/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc index 4cb4801c2ad94b..6bb44d0f19cc55 100644 --- a/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc +++ b/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc @@ -97,12 +97,15 @@ CompileModuleResults InitializeResults(const HloModule* hlo_module) { hlo_module, ExecutionStreamAssignment::Options{ kDefaultNumComputeStreams, - /*number_of_collective_execution_streams=*/ + /*number_of_communication_execution_streams=*/ hlo_module->config() .debug_options() .xla_gpu_experimental_enable_collective_multi_streaming() ? kDefaultNumCommunicationStreams : 1, + // Keep disabled by default until the 14-day AOT + // forward-compatibility window has elapsed (9th of October 2026). + /*enable_dedicated_memcpy_streams=*/false, }); return results; } diff --git a/third_party/xla/xla/service/gpu/execution_stream_assignment.cc b/third_party/xla/xla/service/gpu/execution_stream_assignment.cc index 8ed23b98d76d1c..9cadb89ee3db2a 100644 --- a/third_party/xla/xla/service/gpu/execution_stream_assignment.cc +++ b/third_party/xla/xla/service/gpu/execution_stream_assignment.cc @@ -37,6 +37,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_opcode.h" +#include "xla/layout.h" #include "xla/service/collective_ops_utils.h" #include "xla/service/collective_opt_utils.h" #include "xla/service/gpu/gpu_latency_hiding_scheduler.h" @@ -45,10 +46,15 @@ limitations under the License. namespace xla::gpu { namespace { -// There are two kinds of async execution scopes: compute and collective. We -// need just two as our goal is to effectively overlap computation with -// communication. -enum class ExecutionScopeKind { kCompute, kCommunication }; +// Execution scope kinds. kMemcpyD2H and kMemcpyH2D use the dedicated +// device_to_host_stream / host_to_device_stream from ExecuteParams, which +// keeps host↔device copies off the compute and collective streams. +enum class ExecutionScopeKind { + kCompute, + kCommunication, + kMemcpyD2H, + kMemcpyH2D, +}; template void AbslStringify(Sink sink, ExecutionScopeKind kind) { @@ -59,6 +65,12 @@ void AbslStringify(Sink sink, ExecutionScopeKind kind) { case ExecutionScopeKind::kCommunication: sink.Append("communication"); break; + case ExecutionScopeKind::kMemcpyD2H: + sink.Append("memcpy_d2h"); + break; + case ExecutionScopeKind::kMemcpyH2D: + sink.Append("memcpy_h2d"); + break; } } @@ -102,6 +114,10 @@ class ExecutionStreams { case ExecutionScopeKind::kCommunication: { return NextCollective(kUnspecifiedCollectiveDomain); } + case ExecutionScopeKind::kMemcpyD2H: + return kMemcpyD2HStreamId; + case ExecutionScopeKind::kMemcpyH2D: + return kMemcpyH2DStreamId; } } @@ -148,16 +164,79 @@ bool IsWrappedCollective(const HloAsyncInstruction* async) { } } +// Returns true if the shape is in host memory space (S(5)). +bool IsHostShape(const Shape& shape) { + return shape.IsArray() && shape.has_layout() && + shape.layout().memory_space() == Layout::kHostMemorySpace; +} + +// Returns the memcpy direction for a copy-start instruction whose shape is the +// tuple {dst, src, u32[]}. Returns nullopt for D2D copies (neither endpoint +// is in host memory). +std::optional CopyStartDirection( + const HloInstruction* copy_start) { + const Shape& shape = copy_start->shape(); + if (!shape.IsTuple() || shape.tuple_shapes_size() < 2) { + return std::nullopt; + } + bool dst_host = IsHostShape(shape.tuple_shapes(0)); + bool src_host = IsHostShape(shape.tuple_shapes(1)); + if (dst_host) { + return ExecutionScopeKind::kMemcpyD2H; + } + if (src_host) { + return ExecutionScopeKind::kMemcpyH2D; + } + return std::nullopt; +} + +// Returns the memcpy direction for an async-start whose wrapped computation +// contains a host↔device DUS or DS. The inner instruction may have been +// wrapped in a kLoop fusion by StreamAttributeAnnotator. +std::optional AsyncHostMemcpyDirection( + const HloAsyncInstruction* async_start) { + const HloInstruction* inner = async_start->async_wrapped_instruction(); + + // After StreamAttributeAnnotator the DUS/DS is wrapped in a kLoop fusion; + // unwrap to the fused root so the checks below apply uniformly to both the + // fused and non-fused (pre-StreamAttributeAnnotator) cases. + if (inner->opcode() == HloOpcode::kFusion && + inner->fusion_kind() != HloInstruction::FusionKind::kCustom) { + inner = inner->fused_expression_root(); + } + + // DUS writes to its first operand (the base buffer) and the result has the + // same shape. DS reads from its first operand. + if (inner->opcode() == HloOpcode::kDynamicUpdateSlice && + IsHostShape(inner->shape())) { + return ExecutionScopeKind::kMemcpyD2H; + } + if (inner->opcode() == HloOpcode::kDynamicSlice && + !inner->operands().empty() && IsHostShape(inner->operand(0)->shape())) { + return ExecutionScopeKind::kMemcpyH2D; + } + return std::nullopt; +} + // Returns an execution scope kind if operations starts it. std::optional IsExecutionScopeStart( - const HloInstruction* hlo) { + const HloInstruction* hlo, bool enable_dedicated_memcpy_streams) { // Async operation that starts a new execution scope. if (auto* start = DynCast(hlo)) { - return IsWrappedCollective(start) || IsCustomCollectiveOp(start) || - start->frontend_attributes().map().contains( - kCollectiveGroupMarkerAttr) - ? ExecutionScopeKind::kCommunication - : ExecutionScopeKind::kCompute; + // Collective operations run on communication streams. + if (IsWrappedCollective(start) || IsCustomCollectiveOp(start) || + start->frontend_attributes().map().contains( + kCollectiveGroupMarkerAttr)) { + return ExecutionScopeKind::kCommunication; + } + // Host↔device memcpy (DUS/DS wrapping host memory) runs on dedicated + // memcpy streams rather than the general compute pool. + if (enable_dedicated_memcpy_streams) { + if (auto dir = AsyncHostMemcpyDirection(start)) { + return dir; + } + } + return ExecutionScopeKind::kCompute; } // Async-collective operations not yet migrated to async wrappers. @@ -173,8 +252,14 @@ std::optional IsExecutionScopeStart( : std::nullopt; } - // A special case of asynchronous compute operation. + // copy-start: D2H and H2D go on their dedicated memcpy streams; D2D stays + // on a compute stream. if (HloPredicateIsOp(hlo)) { + if (enable_dedicated_memcpy_streams) { + if (auto dir = CopyStartDirection(hlo)) { + return dir; + } + } return ExecutionScopeKind::kCompute; } @@ -195,6 +280,11 @@ std::optional FindAssignedStreamId( return ExecutionStreamId(ComputationStreamId(assigned_stream_id)); case ExecutionScopeKind::kCommunication: return ExecutionStreamId(CommunicationStreamId(assigned_stream_id)); + case ExecutionScopeKind::kMemcpyD2H: + case ExecutionScopeKind::kMemcpyH2D: + // Memcpy streams are fixed singletons; per-instruction annotation + // does not apply. + return std::nullopt; } } return std::nullopt; @@ -242,7 +332,8 @@ ExecutionStreamAssignment::ExecutionStreamAssignment(const HloModule* module, for (const HloInstruction* hlo : instructions) { // Only assign execution stream IDs to scope-start operations. - if (std::optional kind = IsExecutionScopeStart(hlo)) { + if (std::optional kind = IsExecutionScopeStart( + hlo, options.enable_dedicated_memcpy_streams)) { // Prefer an explicitly assigned stream id, then a collective-domain // stream or a dedicated P2P stream for pipelined send/recv. Otherwise, // generate a new stream id for the execution scope. diff --git a/third_party/xla/xla/service/gpu/execution_stream_assignment.h b/third_party/xla/xla/service/gpu/execution_stream_assignment.h index d85d525425eabf..cbb3b6eaa65525 100644 --- a/third_party/xla/xla/service/gpu/execution_stream_assignment.h +++ b/third_party/xla/xla/service/gpu/execution_stream_assignment.h @@ -39,6 +39,9 @@ class ExecutionStreamAssignment { // Number of additional communication execution streams in each // collective-domain pool for round-robin assignment. int number_of_communication_execution_streams = 1; + // Whether to assign dedicated D2H and H2D streams for host<->device memory + // copies instead of compute streams. + bool enable_dedicated_memcpy_streams = true; }; // The `HloModule` must be flat. In other words, there must be a one-to-one diff --git a/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc b/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc index 76826c35af37c4..a1cc751d198f44 100644 --- a/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc +++ b/third_party/xla/xla/service/gpu/execution_stream_assignment_test.cc @@ -24,9 +24,7 @@ limitations under the License. #include "absl/strings/string_view.h" #include "xla/backends/gpu/runtime/execution_stream_id.h" #include "xla/backends/gpu/transforms/collectives/collective_domain.h" -#include "xla/hlo/ir/hlo_casting_utils.h" #include "xla/hlo/ir/hlo_instruction.h" -#include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" #include "xla/service/gpu/backend_configs.pb.h" @@ -119,13 +117,160 @@ TEST_F(ExecutionStreamAssignmentTest, CopyStartStreamIdTest) { ExecutionStreamAssignment assignment(module.get()); - // copy-start is a compute scope start, gets ComputationStreamId(0). + // D2D copy-start (neither endpoint is host memory) runs on a compute stream. EXPECT_THAT( assignment.GetExecutionStreamId( FindInstruction(module.get(), "copy-start")), absl_testing::IsOkAndHolds(ExecutionStreamId(ComputationStreamId(0)))); } +TEST_F(ExecutionStreamAssignmentTest, CopyStartD2HStreamId) { + // Destination is in host memory (S(5)) → device-to-host → kMemcpyD2HStreamId. + const char* const hlo = R"( + HloModule Module + + ENTRY main { + p0 = f32[1024]{0} parameter(0) + copy-start = (f32[1024]{0:S(5)}, f32[1024]{0}, u32[]) copy-start(p0) + ROOT copy-done = f32[1024]{0:S(5)} copy-done(copy-start) + } + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo)); + + ExecutionStreamAssignment assignment(module.get()); + + EXPECT_THAT( + assignment.GetExecutionStreamId( + FindInstruction(module.get(), "copy-start")), + absl_testing::IsOkAndHolds(ExecutionStreamId(kMemcpyD2HStreamId))); +} + +TEST_F(ExecutionStreamAssignmentTest, CopyStartH2DStreamId) { + // Source is in host memory (S(5)) → host-to-device → kMemcpyH2DStreamId. + const char* const hlo = R"( + HloModule Module + + ENTRY main { + p0 = f32[1024]{0:S(5)} parameter(0) + copy-start = (f32[1024]{0}, f32[1024]{0:S(5)}, u32[]) copy-start(p0) + ROOT copy-done = f32[1024]{0} copy-done(copy-start) + } + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo)); + + ExecutionStreamAssignment assignment(module.get()); + + EXPECT_THAT( + assignment.GetExecutionStreamId( + FindInstruction(module.get(), "copy-start")), + absl_testing::IsOkAndHolds(ExecutionStreamId(kMemcpyH2DStreamId))); +} + +TEST_F(ExecutionStreamAssignmentTest, AsyncDusToHostMemory) { + // async-start wrapping a DUS whose output is in host memory → D2H. + const char* const hlo = R"( + HloModule Module + + dus_computation { + base = f32[4]{0:S(5)} parameter(0) + update = f32[1]{0} parameter(1) + idx = s32[] parameter(2) + ROOT result = f32[4]{0:S(5)} dynamic-update-slice(base, update, idx) + } + + ENTRY main { + p0 = f32[4]{0:S(5)} parameter(0) + p1 = f32[1]{0} parameter(1) + p2 = s32[] parameter(2) + start = ((f32[4]{0:S(5)}, f32[1]{0}, s32[]), f32[4]{0:S(5)}, u32[]) async-start(p0, p1, p2), calls=dus_computation + ROOT done = f32[4]{0:S(5)} async-done(start) + } + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo)); + + ExecutionStreamAssignment assignment(module.get()); + + EXPECT_THAT( + assignment.GetExecutionStreamId(FindInstruction(module.get(), "start")), + absl_testing::IsOkAndHolds(ExecutionStreamId(kMemcpyD2HStreamId))); + EXPECT_THAT( + assignment.GetExecutionStreamId(FindInstruction(module.get(), "done")), + absl_testing::StatusIs(absl::StatusCode::kNotFound)); +} + +TEST_F(ExecutionStreamAssignmentTest, AsyncFusedDusToHostMemory) { + // async-start wrapping a kLoop fusion whose output is in host memory (the + // post-StreamAttributeAnnotator form of a D2H host-memory DUS). + const char* const hlo = R"( + HloModule Module + + wrapped_dus_computation { + base = f32[4]{0:S(5)} parameter(0) + update = f32[1]{0} parameter(1) + idx = s32[] parameter(2) + ROOT dus = f32[4]{0:S(5)} dynamic-update-slice(base, update, idx) + } + + fused_dus_async { + p0 = f32[4]{0:S(5)} parameter(0) + p1 = f32[1]{0} parameter(1) + p2 = s32[] parameter(2) + ROOT fused = f32[4]{0:S(5)} fusion(p0, p1, p2), kind=kLoop, + calls=wrapped_dus_computation + } + + ENTRY main { + p0 = f32[4]{0:S(5)} parameter(0) + p1 = f32[1]{0} parameter(1) + p2 = s32[] parameter(2) + start = ((f32[4]{0:S(5)}, f32[1]{0}, s32[]), f32[4]{0:S(5)}, u32[]) async-start(p0, p1, p2), calls=fused_dus_async + ROOT done = f32[4]{0:S(5)} async-done(start) + } + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo)); + + ExecutionStreamAssignment assignment(module.get()); + + EXPECT_THAT( + assignment.GetExecutionStreamId(FindInstruction(module.get(), "start")), + absl_testing::IsOkAndHolds(ExecutionStreamId(kMemcpyD2HStreamId))); +} + +TEST_F(ExecutionStreamAssignmentTest, AsyncDsFromHostMemory) { + // async-start wrapping a DS whose first operand is in host memory → H2D. + const char* const hlo = R"( + HloModule Module + + ds_computation { + src = f32[4]{0:S(5)} parameter(0) + idx = s32[] parameter(1) + ROOT result = f32[1]{0} dynamic-slice(src, idx), dynamic_slice_sizes={1} + } + + ENTRY main { + p0 = f32[4]{0:S(5)} parameter(0) + p1 = s32[] parameter(1) + start = ((f32[4]{0:S(5)}, s32[]), f32[1]{0}, u32[]) async-start(p0, p1), calls=ds_computation + ROOT done = f32[1]{0} async-done(start) + } + )"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo)); + + ExecutionStreamAssignment assignment(module.get()); + + EXPECT_THAT( + assignment.GetExecutionStreamId(FindInstruction(module.get(), "start")), + absl_testing::IsOkAndHolds(ExecutionStreamId(kMemcpyH2DStreamId))); + EXPECT_THAT( + assignment.GetExecutionStreamId(FindInstruction(module.get(), "done")), + absl_testing::StatusIs(absl::StatusCode::kNotFound)); +} + TEST_F(ExecutionStreamAssignmentTest, FusionComputations) { const char* kModuleStr = R"( HloModule m diff --git a/third_party/xla/xla/service/gpu/gpu_aot_compilation_result.cc b/third_party/xla/xla/service/gpu/gpu_aot_compilation_result.cc index f0d11ae64187b7..43868aa43599b4 100644 --- a/third_party/xla/xla/service/gpu/gpu_aot_compilation_result.cc +++ b/third_party/xla/xla/service/gpu/gpu_aot_compilation_result.cc @@ -29,7 +29,6 @@ limitations under the License. #include "absl/strings/string_view.h" #include "google/protobuf/arena.h" #include "riegeli/bytes/string_writer.h" -#include "xla/debug_options_flags.h" #include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/ir/hlo_print_options.h" #include "xla/pjrt/compiled_memory_stats.h" @@ -50,29 +49,39 @@ limitations under the License. namespace xla::gpu { -static absl::StatusOr, tsl::Fprint128>> -ParseHloModuleAndFingerprint(const HloModuleProtoWithConfig& proto) { - ABSL_ASSIGN_OR_RETURN(std::unique_ptr module, - HloModule::CreateFromProtoWithConfig(proto)); +// Fingerprint of the canonical HLO module text. +static tsl::Fprint128 HloModuleFingerprint(const HloModule& module) { HighwayHashPrinter printer; - module->Print(&printer, HloPrintOptions::Canonical() - .set_print_backend_config(true) - .set_sort_backend_config(true)); - return std::make_pair(std::move(module), printer.ToFingerprint128()); + module.Print(&printer, HloPrintOptions::Canonical() + .set_print_backend_config(true) + .set_sort_backend_config(true)); + return printer.ToFingerprint128(); +} + +// Fingerprint of the deterministic serialization of the executable proto. +// +// NOTE: This is expensive as it re-serializes the whole proto (including all +// constants) twice, so it should only be called when its result is actually +// needed (e.g. for logging). +static tsl::Fprint128 ExecutableFingerprint(const GpuExecutableProto& proto) { + // Fingerprint twice with differents seeds to catch non-deterministic + // serializations. + return {tsl::DeterministicProtoHash64(proto), + tsl::DeterministicProtoHash64(proto, /*seed=*/1)}; +} + +// Formats a fingerprint as 32 hex digits, low64 first. +static std::string FprintToHex(const tsl::Fprint128& fp) { + return absl::StrFormat("%016x%016x", fp.low64, fp.high64); } absl::StatusOr> GpuAotCompilationResult::FromProto(GpuExecutableProto executable_proto) { - tsl::Fprint128 executable_fingerprint = { - tsl::DeterministicProtoHash64(executable_proto), - tsl::DeterministicProtoHash64(executable_proto, /*seed=*/1)}; - ABSL_ASSIGN_OR_RETURN( - auto module_and_fingerprint, - ParseHloModuleAndFingerprint(executable_proto.hlo_module_with_config())); - auto& [module, hlo_fingerprint] = module_and_fingerprint; + ABSL_ASSIGN_OR_RETURN(std::unique_ptr module, + HloModule::CreateFromProtoWithConfig( + executable_proto.hlo_module_with_config())); return absl::WrapUnique(new GpuAotCompilationResult( - std::move(executable_proto), std::move(module), hlo_fingerprint, - executable_fingerprint)); + std::move(executable_proto), std::move(module))); } absl::StatusOr> @@ -84,17 +93,13 @@ GpuAotCompilationResult::FromSerialized( ABSL_RETURN_IF_ERROR(ReadSplitProto(std::move(reader), *executable_proto)); - tsl::Fprint128 executable_fingerprint = { - tsl::DeterministicProtoHash64(*executable_proto), - tsl::DeterministicProtoHash64(*executable_proto, /*seed=*/1)}; - ABSL_ASSIGN_OR_RETURN( - auto module_and_fingerprint, - ParseHloModuleAndFingerprint(executable_proto->hlo_module_with_config())); - auto& [module, hlo_fingerprint] = module_and_fingerprint; - return absl::WrapUnique(new GpuAotCompilationResult( - internal::ArenaAllocatedGpuExecutableProto(std::move(arena), - executable_proto), - std::move(module), hlo_fingerprint, executable_fingerprint)); + ABSL_ASSIGN_OR_RETURN(std::unique_ptr module, + HloModule::CreateFromProtoWithConfig( + executable_proto->hlo_module_with_config())); + return absl::WrapUnique( + new GpuAotCompilationResult(internal::ArenaAllocatedGpuExecutableProto( + std::move(arena), executable_proto), + std::move(module))); } absl::StatusOr GpuAotCompilationResult::SerializeAsString() const { @@ -118,15 +123,14 @@ GpuAotCompilationResult::LoadExecutable( VLOG(1) << absl::StrFormat( "GpuAotCompilationResult::LoadExecutable: module=%s " - "num_instructions=%d hlo_fingerprint=%016x%016x " - "executable_fingerprint=%016x%016x", + "num_instructions=%d hlo_fingerprint=%s executable_fingerprint=%s", hlo_module_->name(), hlo_module_->instruction_count(), - hlo_fingerprint_.low64, hlo_fingerprint_.high64, - executable_fingerprint_.low64, executable_fingerprint_.high64); + FprintToHex(HloModuleFingerprint(*hlo_module_)), + FprintToHex(ExecutableFingerprint(GetExecutableProto()))); return GpuExecutable::FromProto(GetExecutableProto(), device_description, platform_id->ToName(), debug_options, - symbol_resolver); + symbol_resolver, std::move(hlo_module_)); } const GpuExecutableProto& GpuAotCompilationResult::GetExecutableProto() const { diff --git a/third_party/xla/xla/service/gpu/gpu_aot_compilation_result.h b/third_party/xla/xla/service/gpu/gpu_aot_compilation_result.h index 4c79f4ac121ffe..cb40993081e845 100644 --- a/third_party/xla/xla/service/gpu/gpu_aot_compilation_result.h +++ b/third_party/xla/xla/service/gpu/gpu_aot_compilation_result.h @@ -35,7 +35,6 @@ limitations under the License. #include "xla/stream_executor/device_description.h" #include "xla/stream_executor/platform.h" #include "xla/xla.pb.h" -#include "tsl/platform/fingerprint.h" namespace xla::gpu { @@ -101,18 +100,13 @@ class GpuAotCompilationResult : public CompiledModule { std::variant gpu_executable_proto, - std::shared_ptr hlo_module, tsl::Fprint128 hlo_fingerprint, - tsl::Fprint128 executable_fingerprint) + std::shared_ptr hlo_module) : gpu_executable_proto_(std::move(gpu_executable_proto)), - hlo_module_(std::move(hlo_module)), - hlo_fingerprint_(hlo_fingerprint), - executable_fingerprint_(executable_fingerprint) {} + hlo_module_(std::move(hlo_module)) {} std::variant gpu_executable_proto_; std::shared_ptr hlo_module_; - tsl::Fprint128 hlo_fingerprint_; - tsl::Fprint128 executable_fingerprint_; }; } // namespace xla::gpu diff --git a/third_party/xla/xla/service/gpu/gpu_aot_compilation_result_test.cc b/third_party/xla/xla/service/gpu/gpu_aot_compilation_result_test.cc index 966880078cbea1..436935074bbb8b 100644 --- a/third_party/xla/xla/service/gpu/gpu_aot_compilation_result_test.cc +++ b/third_party/xla/xla/service/gpu/gpu_aot_compilation_result_test.cc @@ -236,10 +236,14 @@ TEST_F(GpuAotCompilationResultTest, LoadExecutable) { EnsureCudaSymbolIsRegistered(); + std::shared_ptr expected_module = + result->shared_optimized_module(); + ASSERT_OK_AND_ASSIGN( std::unique_ptr executable, std::move(*result).LoadExecutable(platform_.id(), GetDeviceDescription(), DebugOptions())); + EXPECT_EQ(executable->shared_module(), expected_module); { ASSERT_OK_AND_ASSIGN( diff --git a/third_party/xla/xla/service/gpu/gpu_executable.cc b/third_party/xla/xla/service/gpu/gpu_executable.cc index 3f628c38ea9483..4032510567639c 100644 --- a/third_party/xla/xla/service/gpu/gpu_executable.cc +++ b/third_party/xla/xla/service/gpu/gpu_executable.cc @@ -243,6 +243,8 @@ static GpuExecutable::NumAdditionalStreams GetNumAdditionalStreams( // Clamp explicitly requested stream counts to non-negative values. compute = std::max(0, compute); comm = std::max(0, comm); + bool need_d2h_stream = false; + bool need_h2d_stream = false; // Then traverse all thunks to see if anyone requested more streams. This // also sizes sparse collective-domain stream pools from their assigned IDs. @@ -252,14 +254,20 @@ static GpuExecutable::NumAdditionalStreams GetNumAdditionalStreams( ExecutionStreamId id = async_start->execution_stream_id(); if (id.is_computation()) { compute = std::max(compute, id.computation_id().value() + 1); - } else { + } else if (id.is_communication()) { comm = std::max(comm, id.communication_id().value() + 1); + } else if (id.is_memcpy()) { + if (id.memcpy_id() == kMemcpyD2HStreamId) { + need_d2h_stream = true; + } else if (id.memcpy_id() == kMemcpyH2DStreamId) { + need_h2d_stream = true; + } } } }); } - return {compute, comm}; + return {compute, comm, need_d2h_stream, need_h2d_stream}; } GpuExecutable::BorrowedStreams GpuExecutable::BorrowedStreams::Assign( @@ -452,7 +460,7 @@ absl::StatusOr> GpuExecutable::Create( // Implementation note: HLO profiling is always enabled for GPU executables, // since we can use timers around thunks. GpuExecutable::GpuExecutable( - std::unique_ptr debug_module, std::vector binary, + std::shared_ptr debug_module, std::vector binary, BinaryMap dnn_compiled_graphs, se::DeviceDescription device_description, std::unique_ptr executable, std::string module_name, ProgramShape program_shape, std::vector allocations, @@ -792,6 +800,35 @@ absl::Status GpuExecutable::ExecuteThunksImpl( &collective_memory, std::move(compute_streams.streams), &execution_scoped_state, persistent_alloc_indices); + // device_to_host_stream/host_to_device_stream above come from run_options, + // which are only populated when running through PJRT. Fall back to a + // borrowed stream (or main_stream, if no stream borrower is available) so + // that async host<->device copy thunks also work for non-PJRT clients. + StreamPool::Ptr borrowed_device_to_host_stream; + if (num_additional_streams.need_d2h_stream && + execute_params.device_to_host_stream == nullptr) { + if (run_options->HasStreamBorrower()) { + ABSL_ASSIGN_OR_RETURN(borrowed_device_to_host_stream, + run_options->BorrowStream(executor->device_ordinal())); + execute_params.device_to_host_stream = + borrowed_device_to_host_stream.get(); + } else { + execute_params.device_to_host_stream = main_stream; + } + } + StreamPool::Ptr borrowed_host_to_device_stream; + if (num_additional_streams.need_h2d_stream && + execute_params.host_to_device_stream == nullptr) { + if (run_options->HasStreamBorrower()) { + ABSL_ASSIGN_OR_RETURN(borrowed_host_to_device_stream, + run_options->BorrowStream(executor->device_ordinal())); + execute_params.host_to_device_stream = + borrowed_host_to_device_stream.get(); + } else { + execute_params.host_to_device_stream = main_stream; + } + } + XLA_VLOG_DEVICE(1, run_options->device_ordinal()) << "Start GpuExecutable::ExecuteOnStream module: " << module_name; ABSL_RETURN_IF_ERROR(thunk_executor.ExecuteOnStream(execute_params)); @@ -1416,17 +1453,21 @@ absl::StatusOr> GpuExecutable::FromProto( const GpuExecutableProto& proto, const se::DeviceDescription& device_description, absl::string_view platform_name, DebugOptions debug_options, - const std::optional& - symbol_resolver) { + const std::optional& symbol_resolver, + std::shared_ptr debug_module) { Params params; params.debug_options = std::move(debug_options); params.enable_debug_info_manager = params.debug_options.xla_gpu_executable_embed_debug_info(); const std::string& binary = proto.binary(); params.binary.assign(binary.begin(), binary.end()); - if (proto.has_hlo_module_with_config()) { + if (debug_module != nullptr) { + params.debug_module = std::move(debug_module); + } else if (proto.has_hlo_module_with_config()) { ABSL_ASSIGN_OR_RETURN(params.debug_module, HloModule::CreateFromProtoWithConfig( proto.hlo_module_with_config())); + } + if (params.debug_module != nullptr) { // The HLO module deserialized from the proto carries xla_dump_to from the // process that originally compiled it. Override with the current process's // dump path so that runtime dumps (checksum logs, etc.) land in the correct diff --git a/third_party/xla/xla/service/gpu/gpu_executable.h b/third_party/xla/xla/service/gpu/gpu_executable.h index 21200c35860afe..9d4670baef4dbb 100644 --- a/third_party/xla/xla/service/gpu/gpu_executable.h +++ b/third_party/xla/xla/service/gpu/gpu_executable.h @@ -83,6 +83,8 @@ class GpuExecutable : public Executable { struct NumAdditionalStreams { int compute = 0; int communication = 0; + bool need_d2h_stream = false; + bool need_h2d_stream = false; }; using ConstantInfo = GpuModuleGlobals::ConstantInfo; @@ -127,7 +129,7 @@ class GpuExecutable : public Executable { std::unique_ptr alias_info; DebugOptions debug_options; se::DeviceDescription device_description; - std::unique_ptr debug_module = nullptr; + std::shared_ptr debug_module = nullptr; bool enable_debug_info_manager = true; ModuleStats module_stats; se::ExecutableAbiVersion executable_abi_version; @@ -242,12 +244,17 @@ class GpuExecutable : public Executable { absl::StatusOr ResolveConstantGlobals( se::Stream* stream); + // Creates a `GpuExecutable` from its proto representation. + // + // If `debug_module` isn't populated, the HLO module is deserialized from the + // proto, which can be expensive. static absl::StatusOr> FromProto( const GpuExecutableProto&, const se::DeviceDescription& device_description, absl::string_view platform, DebugOptions debug_options, const std::optional& - symbol_resolver = std::nullopt); + symbol_resolver = std::nullopt, + std::shared_ptr debug_module = nullptr); absl::StatusOr ToProto() const; @@ -277,7 +284,7 @@ class GpuExecutable : public Executable { // Use GpuExecutable::Create() to create an instance. explicit GpuExecutable( - std::unique_ptr debug_module, std::vector binary, + std::shared_ptr debug_module, std::vector binary, BinaryMap dnn_compiled_graphs, se::DeviceDescription device_description, std::unique_ptr executable, std::string module_name, ProgramShape program_shape, std::vector allocations, diff --git a/third_party/xla/xla/service/gpu/llvm_gpu_backend/amdgpu_backend.cc b/third_party/xla/xla/service/gpu/llvm_gpu_backend/amdgpu_backend.cc index da4935bc9e9848..3b2da8524957cb 100644 --- a/third_party/xla/xla/service/gpu/llvm_gpu_backend/amdgpu_backend.cc +++ b/third_party/xla/xla/service/gpu/llvm_gpu_backend/amdgpu_backend.cc @@ -629,63 +629,28 @@ absl::Status AMDGPUTargetModuleLinker(llvm::Module* module, module->addModuleFlag(llvm::Module::Error, "amdhsa_code_object_version", kAMDGPUAbiVersion); - return absl::OkStatus(); -} - -// The following routine maps a feature token extracted from the -// hipDeviceProp_t::gcnArchName string, and maps it to a valid feature_str -// to be used for creating the AMDGPUTarget. -// This mapping is currently in a state of flux because TF XLA uses its -// own copy of LLVM, which is different from the LLVM version used by -// hipcc/runtime in the ROCm install. Ordinarily this is not a problem, -// but right now, the LLVM version used by hipcc/runtime has "targetID" -// related changes which have not yet been upstreamed (to the LLVM repo) -// When that upstreaming happens (and TF LLVM pointer moves past the -// upstream commit), the following mapping will need to change -std::string MapGCNArchNameTokenToFeatureStr(const std::string& token, - const std::string& gfx) { - if (token == "sramecc+") { - return "+sramecc"; - } - if (token == "sramecc-") { - if (gfx == "gfx90a" || gfx == "gfx942") { - return ""; + absl::InlinedVector tokens = + absl::StrSplit(compute_capability->gcn_arch_name(), ':'); + CHECK(!tokens.empty()); + for (auto token : absl::MakeSpan(tokens).subspan(1)) { + if (token == "xnack+") { + module->addModuleFlag(llvm::Module::Error, "amdgpu.xnack", 1u); + } else if (token == "xnack-") { + module->addModuleFlag(llvm::Module::Error, "amdgpu.xnack", 0u); + } else if (token == "sramecc+") { + module->addModuleFlag(llvm::Module::Error, "amdgpu.sramecc", 1u); + } else if (token == "sramecc-") { + module->addModuleFlag(llvm::Module::Error, "amdgpu.sramecc", 0u); + } else { + LOG(FATAL) << "Unknown feature'" << token << "' in '" + << compute_capability->gcn_arch_name() << "'"; } - return "-sramecc"; - } - if (token == "xnack+") { - return "+xnack"; } - if (token == "xnack-") { - return "-xnack"; - } - return ""; -} -std::pair GetFeatureStrFromGCNArchName( - const std::string& gcn_arch_name) { - std::string gfx = gcn_arch_name; - // For ROCm versions 4.0 and greater, we need to specify the correct - // feature str, based on the underlying GPU HW to get max performance. - std::vector tokens = absl::StrSplit(gcn_arch_name, ':'); - if (!tokens.empty()) { - gfx = tokens[0]; - } - - std::string mapped_tokens; - for (size_t i = 1; i < tokens.size(); i++) { - // Skip the first token, that is the gfxNNN str - // The rest of the tokens are the feature/targetid strings - auto mapped_token = MapGCNArchNameTokenToFeatureStr(tokens[i], gfx); - if (!mapped_token.empty()) { - if (!mapped_tokens.empty()) mapped_tokens += ","; - mapped_tokens += mapped_token; - } - } - return std::pair{gfx, mapped_tokens}; + return absl::OkStatus(); } -void AMDGPUBackendInit(const DebugOptions& debug_options) { +void AMDGPUBackendInit() { // Initialize the AMDGPU target; it's the only target we link with, so call // its specific initialization functions instead of the catch-all // InitializeAll*. @@ -703,7 +668,7 @@ absl::StatusOr CompileToHsacoInternal( llvm::Module* module, se::GpuComputeCapability gpu_version, const DebugOptions& debug_options, std::string* hsaco_temp_path) { static absl::once_flag backend_init_flag; - absl::call_once(backend_init_flag, AMDGPUBackendInit, debug_options); + absl::call_once(backend_init_flag, AMDGPUBackendInit); auto cc = gpu_version.rocm_compute_capability(); if (!cc) { @@ -711,9 +676,8 @@ absl::StatusOr CompileToHsacoInternal( } llvm::Triple default_target_triple("amdgcn--amdhsa-amdgiz"); // Construct LLVM TargetMachine for AMDGPU. - auto [gfx, feature_str] = GetFeatureStrFromGCNArchName(cc->gcn_arch_name()); - auto target_machine = - GetTargetMachine(default_target_triple, gfx, debug_options, feature_str); + auto target_machine = GetTargetMachine(default_target_triple, + cc->gfx_version(), debug_options, ""); // Link with ROCm-Device-Libs, and optimize the LLVM module. ABSL_RETURN_IF_ERROR(gpu::LinkAndOptimizeModule( diff --git a/third_party/xla/xla/service/gpu/model/BUILD b/third_party/xla/xla/service/gpu/model/BUILD index 96be346dcc1781..d221700bdd30fa 100644 --- a/third_party/xla/xla/service/gpu/model/BUILD +++ b/third_party/xla/xla/service/gpu/model/BUILD @@ -642,6 +642,7 @@ xla_cc_test( "//xla/tests:xla_internal_test_main", "//xla/tsl/platform:env", "//xla/tsl/testing:temporary_directory", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/strings", "@com_google_absl//absl/time", diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc index 00ac8fd9258d73..636570437775f0 100644 --- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc +++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc @@ -810,6 +810,23 @@ int64_t GpuPerformanceModelWithIndexingAnalysis::EstimateNumWarps( return EstimateNumWarpsImpl(tiled_hlo_computation); } +namespace internal { + +absl::flat_hash_map PrecomputeFlopsMap( + const HloFusionAdaptor& fusion_adaptor, + absl::FunctionRef flops_per_element_fn) { + auto post_order = fusion_adaptor.MakeInstructionPostOrder(); + absl::flat_hash_map flops_map; + flops_map.reserve(post_order.size()); + for (const HloInstructionAdaptor& instruction_adaptor : post_order) { + flops_map.emplace(&instruction_adaptor.instruction(), + flops_per_element_fn(&instruction_adaptor.instruction())); + } + return flops_map; +} + +} // namespace internal + absl::StatusOr GpuPerformanceModelWithIndexingAnalysis::TryFindTopKBestTilingsForFusion( const HloFusionAdaptor& fusion_adaptor, int top_k) { @@ -837,6 +854,23 @@ GpuPerformanceModelWithIndexingAnalysis::TryFindTopKBestTilingsForFusion( "TryFindTopKBestTilingsForFusion tiling_space evaluating ", tilings.size(), " tilings."); + const absl::flat_hash_map flops_map = + internal::PrecomputeFlopsMap( + fusion_adaptor, + [this](const HloInstruction* hlo) { return FlopsPerElement(hlo); }); + + auto flops_fn = [&flops_map](const HloInstruction* hlo) -> int64_t { + if (hlo->opcode() == HloOpcode::kParameter || + hlo->opcode() == HloOpcode::kTuple || + hlo->opcode() == HloOpcode::kGetTupleElement) { + return 0; + } + auto it = flops_map.find(hlo); + CHECK(it != flops_map.end()) + << "Missing precomputed FLOPs for " << hlo->name(); + return it->second; + }; + for (const llvm::SmallVector& tiling : tilings) { // Assign padded tile size as Triton emitter will require that. // Symbolic analysis route hides this assumption. @@ -882,12 +916,10 @@ GpuPerformanceModelWithIndexingAnalysis::TryFindTopKBestTilingsForFusion( continue; } - ABSL_ASSIGN_OR_RETURN(std::optional tiled_run_time_data, - EstimateTiledRunTimeDataImpl( - fusion_adaptor, *tiled_computation, *device_info_, - shape_size_, [this](const HloInstruction* hlo) { - return FlopsPerElement(hlo); - })); + ABSL_ASSIGN_OR_RETURN( + std::optional tiled_run_time_data, + EstimateTiledRunTimeDataImpl(fusion_adaptor, *tiled_computation, + *device_info_, shape_size_, flops_fn)); if (tiled_run_time_data.has_value()) { VLOG(2) << "Accepted tile sizes [" diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.h b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.h index 2f08f831e48d10..50cc2989f3cf83 100644 --- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.h +++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.h @@ -20,7 +20,9 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_map.h" #include "absl/container/inlined_vector.h" +#include "absl/functional/function_ref.h" #include "absl/status/statusor.h" #include "absl/strings/str_cat.h" #include "mlir/IR/MLIRContext.h" @@ -146,6 +148,15 @@ class GpuPerformanceModelWithIndexingAnalysis : public GpuPerformanceModelBase { bool enable_same_shape_multi_output_fusion_; }; +namespace internal { + +// Precomputes FLOPs per element for all instructions in the fusion adaptor. +absl::flat_hash_map PrecomputeFlopsMap( + const HloFusionAdaptor& fusion_adaptor, + absl::FunctionRef flops_per_element_fn); + +} // namespace internal + } // namespace gpu } // namespace xla diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc index 0354c406cde3ea..711a6ad045107c 100644 --- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc +++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc @@ -23,6 +23,7 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_map.h" #include "absl/container/inlined_vector.h" #include "absl/strings/string_view.h" #include "absl/time/time.h" @@ -1283,6 +1284,52 @@ ENTRY entry_computation { )"); } +TEST_P(GpuIndexingPerformanceModelTest, PrecomputeFlopsMapWorksCorrectly) { + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"( +HloModule m + +add { + p0 = f32[] parameter(0) + p1 = f32[] parameter(1) + ROOT add = f32[] add(p0, p1) +} + +fused_computation { + p0 = f32[128,256] parameter(0) + p1 = f32[256,64] parameter(1) + c0 = f32[] constant(0) + exp = f32[128,256] exponential(p0) + dot = f32[128,64] dot(exp, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} + reduce = f32[128] reduce(dot, c0), dimensions={1}, to_apply=add + ROOT tuple = (f32[128,64], f32[128]) tuple(dot, reduce) +} + +ENTRY entry { + p0 = f32[128,256] parameter(0) + p1 = f32[256,64] parameter(1) + ROOT fusion = (f32[128,64], f32[128]) fusion(p0, p1), kind=kCustom, + calls=fused_computation +} +)")); + + HloInstruction* fusion = module->entry_computation()->root_instruction(); + auto fusion_adaptor = HloFusionAdaptor::ForInstruction(fusion); + + auto flops_map = internal::PrecomputeFlopsMap( + *fusion_adaptor, [this](const HloInstruction* instr) { + return indexing_cost_model_.FlopsPerElement(instr); + }); + + HloComputation* fused_comp = fusion->fused_instructions_computation(); + const HloInstruction* dot = fused_comp->GetInstructionWithName("dot"); + const HloInstruction* reduce = fused_comp->GetInstructionWithName("reduce"); + const HloInstruction* exp = fused_comp->GetInstructionWithName("exp"); + + EXPECT_EQ(flops_map.at(dot), indexing_cost_model_.FlopsPerElement(dot)); + EXPECT_EQ(flops_map.at(reduce), indexing_cost_model_.FlopsPerElement(reduce)); + EXPECT_EQ(flops_map.at(exp), indexing_cost_model_.FlopsPerElement(exp)); +} + } // namespace } // namespace gpu } // namespace xla diff --git a/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.cc b/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.cc index b8ac06a1546d5b..e275473f9fa60c 100644 --- a/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.cc +++ b/third_party/xla/xla/service/gpu/pre_scheduling_copy_insertion_pipeline.cc @@ -20,6 +20,7 @@ limitations under the License. #include #include "xla/backends/gpu/transforms/alias_passthrough_params.h" +#include "xla/backends/gpu/transforms/constant_fill_copy_rewriter.h" #include "xla/backends/gpu/transforms/copy_fusion.h" #include "xla/backends/gpu/transforms/sanitize_constant_names.h" #include "xla/hlo/ir/hlo_module.h" @@ -77,6 +78,7 @@ HloPassPipeline PreSchedulingCopyInsertionPipeline( pipeline.AddPass(alias_info); } + pipeline.AddPass(); pipeline.AddPass(device_description); pipeline.AddPass(); return pipeline; diff --git a/third_party/xla/xla/service/gpu_topology.cc b/third_party/xla/xla/service/gpu_topology.cc index d9ca3573bcec8c..bf83bb006dc446 100644 --- a/third_party/xla/xla/service/gpu_topology.cc +++ b/third_party/xla/xla/service/gpu_topology.cc @@ -73,7 +73,8 @@ GetHostTargetMachineOptions(absl::string_view platform_version) { "+sha,+sse,+sse2,+sse3,+sse4.1,+sse4.2,+ssse3,+tsxldtrk,+vaes,+" "vpclmulqdq,+wbnoinvd,+xsave,+xsavec,+xsaveopt,+xsaves,-amx-avx512,-" "amx-complex,-amx-fp16,-amx-fp8,-amx-movrs,-avx10.1,-avx10.2,-" - "avx512bmm,-avx512vp2intersect,-avxifma,-avxneconvert,-avxvnniint16,-" + "avx10v2aux,-avx512bmm,-avx512vp2intersect,-avxifma,-avxneconvert,-" + "avxvnniint16,-" "avxvnniint8,-ccmp,-cf,-clzero,-cmpccxadd,-egpr,-enqcmd,-fma4,-hreset,-" "jmpabs,-kl,-lwp,-movrs,-mwaitx,-ndd,-nf,-pconfig,-pku,-ppx,-prefetchi," "-ptwrite,-push2pop2,-raoint,-rdpru,-sgx,-sha512,-shstk,-sm3,-sm4,-" diff --git a/third_party/xla/xla/service/hlo.proto b/third_party/xla/xla/service/hlo.proto index ea708ce2b8dc4c..84b3cfe65e62c9 100644 --- a/third_party/xla/xla/service/hlo.proto +++ b/third_party/xla/xla/service/hlo.proto @@ -354,7 +354,8 @@ message HloInstructionProto { xla.RandomAlgorithm rng_algorithm = 70; // The comparison type used for kCompare. - string comparison_type = 72; + // Deprecated and replaced by comparison_order. + string comparison_type = 72 [deprecated = true]; // Specifies if this is a cross-program-prefetch, used by kCopyStart. // Deprecated and replaced by optional_cross_program_prefetch_index. diff --git a/third_party/xla/xla/service/host_offload_utils.cc b/third_party/xla/xla/service/host_offload_utils.cc index 7df2ed499ce1f6..f03cfa7d9b3a91 100644 --- a/third_party/xla/xla/service/host_offload_utils.cc +++ b/third_party/xla/xla/service/host_offload_utils.cc @@ -17,7 +17,6 @@ limitations under the License. #include #include -#include #include #include #include @@ -73,14 +72,14 @@ bool CustomCallReusesBuffer(const HloInstruction* custom_call, } // namespace absl::StatusOr> GetSuccessors( - const InstructionAndShapeIndex& instruction_and_shape_index) { + const InstructionAndShapeIndex& instruction_and_shape_index, + const CallGraph& call_graph) { std::vector result; HloInstruction* instruction = instruction_and_shape_index.instruction; if (instruction->IsRoot()) { // Successor of the root is the call instruction(s). - std::unique_ptr call_graph = - CallGraph::Build(instruction->GetModule()); - auto callers = call_graph->GetComputationCallers(instruction->parent()); + const std::vector callers = + call_graph.GetComputationCallers(instruction->parent()); for (HloInstruction* caller : callers) { result.push_back({caller, instruction_and_shape_index.shape_index}); } @@ -183,7 +182,7 @@ absl::StatusOr> GetSuccessors( std::vector GetPredecessors( const InstructionAndShapeIndex& instruction_and_shape_index, - std::optional operand_index) { + const CallGraph& call_graph, std::optional operand_index) { std::vector result; HloInstruction* instruction = instruction_and_shape_index.instruction; if (instruction->opcode() == HloOpcode::kGetTupleElement) { @@ -209,10 +208,8 @@ std::vector GetPredecessors( result.push_back({called_computation->root_instruction(), instruction_and_shape_index.shape_index}); } else if (instruction->opcode() == HloOpcode::kParameter) { - std::unique_ptr call_graph = - CallGraph::Build(instruction->GetModule()); const std::vector callers = - call_graph->GetComputationCallers(instruction->parent()); + call_graph.GetComputationCallers(instruction->parent()); absl::flat_hash_set unique_callers(callers.begin(), callers.end()); for (HloInstruction* caller : unique_callers) { diff --git a/third_party/xla/xla/service/host_offload_utils.h b/third_party/xla/xla/service/host_offload_utils.h index 9abb2fd5340e7d..f64b86ccd12df6 100644 --- a/third_party/xla/xla/service/host_offload_utils.h +++ b/third_party/xla/xla/service/host_offload_utils.h @@ -28,6 +28,9 @@ limitations under the License. #include "xla/shape_util.h" namespace xla { + +class CallGraph; + namespace host_offload_utils { struct InstructionAndShapeIndex { @@ -58,16 +61,23 @@ bool operator==(const InstructionAndShapeIndex& lhs, // If an instruction's user is a call, we descend into the call first. // Eventually, a later invocation of this function while walking the graph will // return the call itself as a successor of the ROOT instruction of the -// computation. +// computation. The call sites of the root's computation are read from +// call_graph, the graph of the instruction's module: a graph built before a +// walk stays valid as long as no computation and no instruction with called +// computations is added, removed or reordered. absl::StatusOr> GetSuccessors( - const InstructionAndShapeIndex& instruction_and_shape_index); + const InstructionAndShapeIndex& instruction_and_shape_index, + const CallGraph& call_graph); // If an instruction's operand is a call, return the call now. A follow up call // of this function on that call returns the ROOT. Eventually, once the given // instruction is a parameter, the returned predecessor will be the appropriate -// operand of the call (not the call itself, since we already returned it). +// operand of the call (not the call itself, since we already returned it). The +// call sites of the parameter's computation are read from call_graph, under +// the same condition as for GetSuccessors. std::vector GetPredecessors( const InstructionAndShapeIndex& instruction_and_shape_index, + const CallGraph& call_graph, std::optional operand_index = std::nullopt); // Returns true if the instruction is allowed to be in the diff --git a/third_party/xla/xla/service/host_offload_utils_test.cc b/third_party/xla/xla/service/host_offload_utils_test.cc index afac77a9162262..7a6676c801a016 100644 --- a/third_party/xla/xla/service/host_offload_utils_test.cc +++ b/third_party/xla/xla/service/host_offload_utils_test.cc @@ -23,7 +23,11 @@ limitations under the License. #include #include "absl/container/flat_hash_set.h" #include "absl/strings/string_view.h" +#include "xla/hlo/ir/hlo_computation.h" +#include "xla/hlo/ir/hlo_instruction.h" +#include "xla/hlo/ir/hlo_module.h" #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h" +#include "xla/service/call_graph.h" #include "xla/shape_util.h" #include "xla/util.h" @@ -31,6 +35,9 @@ namespace xla { namespace host_offload_utils { namespace { +using ::testing::ElementsAre; +using ::testing::UnorderedElementsAre; + using HostOffloadUtilsTest = HloHardwareIndependentTestBase; TEST_F(HostOffloadUtilsTest, SimpleGetSuccessorsGetPredecessorsTest) { @@ -50,20 +57,22 @@ TEST_F(HostOffloadUtilsTest, SimpleGetSuccessorsGetPredecessorsTest) { )hlo"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + std::unique_ptr call_graph = CallGraph::Build(module.get()); HloInstruction* data_param = FindInstruction(module.get(), "data_param"); ASSERT_NE(data_param, nullptr); HloInstruction* offload_custom_call = FindInstruction(module.get(), "offload_custom_call"); ASSERT_NE(offload_custom_call, nullptr); - ASSERT_OK_AND_ASSIGN(std::vector succ, - GetSuccessors(InstructionAndShapeIndex(data_param, {}))); + ASSERT_OK_AND_ASSIGN( + std::vector succ, + GetSuccessors(InstructionAndShapeIndex(data_param, {}), *call_graph)); std::vector expected_succ = { InstructionAndShapeIndex(offload_custom_call, {})}; EXPECT_EQ(succ, expected_succ); - std::vector pred = - GetPredecessors(InstructionAndShapeIndex(offload_custom_call, {})); + std::vector pred = GetPredecessors( + InstructionAndShapeIndex(offload_custom_call, {}), *call_graph); std::vector expected_pred = { InstructionAndShapeIndex(data_param, {})}; EXPECT_EQ(pred, expected_pred); @@ -89,6 +98,7 @@ TEST_F(HostOffloadUtilsTest, ComputationGetSuccessorsGetPredecessorsTest) { )hlo"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + std::unique_ptr call_graph = CallGraph::Build(module.get()); HloInstruction* call = FindInstruction(module.get(), "call"); ASSERT_NE(call, nullptr); HloInstruction* gte_0 = FindInstruction(module.get(), "gte_0"); @@ -96,14 +106,15 @@ TEST_F(HostOffloadUtilsTest, ComputationGetSuccessorsGetPredecessorsTest) { HloInstruction* tuple = FindInstruction(module.get(), "tuple"); ASSERT_NE(tuple, nullptr); - ASSERT_OK_AND_ASSIGN(std::vector succ, - GetSuccessors(InstructionAndShapeIndex(call, {0}))); + ASSERT_OK_AND_ASSIGN( + std::vector succ, + GetSuccessors(InstructionAndShapeIndex(call, {0}), *call_graph)); std::vector expected_succ = { InstructionAndShapeIndex(gte_0, {})}; EXPECT_EQ(succ, expected_succ); std::vector pred = - GetPredecessors(InstructionAndShapeIndex(call, {0})); + GetPredecessors(InstructionAndShapeIndex(call, {0}), *call_graph); std::vector expected_pred = { InstructionAndShapeIndex(tuple, {0})}; EXPECT_EQ(pred, expected_pred); @@ -115,12 +126,74 @@ TEST_F(HostOffloadUtilsTest, ComputationGetSuccessorsGetPredecessorsTest) { ASSERT_NE(offload_custom_call, nullptr); std::vector param_pred = - GetPredecessors(InstructionAndShapeIndex(param_0, {})); + GetPredecessors(InstructionAndShapeIndex(param_0, {}), *call_graph); std::vector expected_param_pred = { InstructionAndShapeIndex(offload_custom_call, {})}; EXPECT_EQ(param_pred, expected_param_pred); } +TEST_F(HostOffloadUtilsTest, RootAndParameterStepsUseTheGivenCallGraph) { + absl::string_view hlo_string = R"hlo( + HloModule my_module + callee { + callee_param = f32[2048] parameter(0) + ROOT callee_root = f32[2048] negate(callee_param) + } + ENTRY main { + data_param = f32[2048] parameter(0) + call_a = f32[2048] call(data_param), to_apply=callee + call_b = f32[2048] call(call_a), to_apply=callee + ROOT add = f32[2048] add(call_a, call_b) + } + )hlo"; + + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(hlo_string)); + std::unique_ptr call_graph = CallGraph::Build(module.get()); + HloInstruction* callee_param = FindInstruction(module.get(), "callee_param"); + HloInstruction* callee_root = FindInstruction(module.get(), "callee_root"); + HloInstruction* data_param = FindInstruction(module.get(), "data_param"); + HloInstruction* call_a = FindInstruction(module.get(), "call_a"); + HloInstruction* call_b = FindInstruction(module.get(), "call_b"); + ASSERT_NE(callee_param, nullptr); + ASSERT_NE(callee_root, nullptr); + ASSERT_NE(data_param, nullptr); + ASSERT_NE(call_a, nullptr); + ASSERT_NE(call_b, nullptr); + + // The root steps to every call site, in call site order; the parameter + // steps to the operand of every call site. + ASSERT_OK_AND_ASSIGN( + std::vector succ, + GetSuccessors(InstructionAndShapeIndex(callee_root, {}), *call_graph)); + EXPECT_THAT(succ, ElementsAre(InstructionAndShapeIndex(call_a, {}), + InstructionAndShapeIndex(call_b, {}))); + EXPECT_THAT( + GetPredecessors(InstructionAndShapeIndex(callee_param, {}), *call_graph), + UnorderedElementsAre(InstructionAndShapeIndex(data_param, {}), + InstructionAndShapeIndex(call_a, {}))); + + // A call site added after the graph was built is unknown to that graph and + // known to a fresh one: the answers come from the graph, not from the + // module. HostOffloader relies on this to build its graph once per run, so + // the utilities must never rebuild one themselves. + HloInstruction* call_c = module->entry_computation()->AddInstruction( + HloInstruction::CreateCall(data_param->shape(), {data_param}, + FindComputation(module.get(), "callee"))); + ASSERT_OK_AND_ASSIGN( + succ, + GetSuccessors(InstructionAndShapeIndex(callee_root, {}), *call_graph)); + EXPECT_THAT(succ, ElementsAre(InstructionAndShapeIndex(call_a, {}), + InstructionAndShapeIndex(call_b, {}))); + std::unique_ptr fresh_call_graph = CallGraph::Build(module.get()); + ASSERT_OK_AND_ASSIGN(succ, + GetSuccessors(InstructionAndShapeIndex(callee_root, {}), + *fresh_call_graph)); + EXPECT_THAT(succ, ElementsAre(InstructionAndShapeIndex(call_a, {}), + InstructionAndShapeIndex(call_b, {}), + InstructionAndShapeIndex(call_c, {}))); +} + TEST_F(HostOffloadUtilsTest, IsMoveToHostWithDynamicUpdateSliceTest) { absl::string_view hlo_string = R"hlo( HloModule my_module @@ -239,13 +312,14 @@ TEST_F(HostOffloadUtilsTest, SendGetPredecessorsTest) { )hlo"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + std::unique_ptr call_graph = CallGraph::Build(module.get()); HloInstruction* send = FindInstruction(module.get(), "send"); ASSERT_NE(send, nullptr); HloInstruction* data_param = FindInstruction(module.get(), "data_param"); ASSERT_NE(data_param, nullptr); std::vector pred = - GetPredecessors(InstructionAndShapeIndex(send, {})); + GetPredecessors(InstructionAndShapeIndex(send, {}), *call_graph); std::vector expected_pred = { InstructionAndShapeIndex(data_param, {})}; EXPECT_EQ(pred, expected_pred); @@ -295,6 +369,7 @@ TEST_F(HostOffloadUtilsTest, ConditionalGetPredecessorsTest) { )hlo"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + std::unique_ptr call_graph = CallGraph::Build(module.get()); HloInstruction* conditional = FindInstruction(module.get(), "conditional"); ASSERT_NE(conditional, nullptr); HloInstruction* true_operand = FindInstruction(module.get(), "true_operand"); @@ -307,13 +382,13 @@ TEST_F(HostOffloadUtilsTest, ConditionalGetPredecessorsTest) { ASSERT_NE(f_root, nullptr); std::vector pred_param = - GetPredecessors(InstructionAndShapeIndex(t_param, {})); + GetPredecessors(InstructionAndShapeIndex(t_param, {}), *call_graph); std::vector expected_pred_param = { InstructionAndShapeIndex(true_operand, {})}; EXPECT_EQ(pred_param, expected_pred_param); std::vector pred_cond = - GetPredecessors(InstructionAndShapeIndex(conditional, {})); + GetPredecessors(InstructionAndShapeIndex(conditional, {}), *call_graph); std::vector expected_pred_cond = { InstructionAndShapeIndex(t_root, {}), InstructionAndShapeIndex(f_root, {})}; @@ -337,6 +412,7 @@ TEST_F(HostOffloadUtilsTest, ConditionalReusedComputationPredecessorsTest) { ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(hlo_string)); + std::unique_ptr call_graph = CallGraph::Build(module.get()); HloInstruction* s_param = FindInstruction(module.get(), "s_param"); ASSERT_NE(s_param, nullptr); HloInstruction* true_operand = FindInstruction(module.get(), "true_operand"); @@ -346,7 +422,7 @@ TEST_F(HostOffloadUtilsTest, ConditionalReusedComputationPredecessorsTest) { ASSERT_NE(false_operand, nullptr); std::vector pred_param = - GetPredecessors(InstructionAndShapeIndex(s_param, {})); + GetPredecessors(InstructionAndShapeIndex(s_param, {}), *call_graph); std::vector expected_pred_param = { InstructionAndShapeIndex(true_operand, {}), InstructionAndShapeIndex(false_operand, {})}; @@ -453,6 +529,7 @@ TEST_F(HostOffloadUtilsTest, SortTupleSuccessorsAndPredecessorsTest) { )hlo"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string)); + std::unique_ptr call_graph = CallGraph::Build(module.get()); HloInstruction* main_p0 = FindInstruction(module.get(), "main_p0"); ASSERT_NE(main_p0, nullptr); HloInstruction* main_p1 = FindInstruction(module.get(), "main_p1"); @@ -465,22 +542,24 @@ TEST_F(HostOffloadUtilsTest, SortTupleSuccessorsAndPredecessorsTest) { ASSERT_NE(gte1, nullptr); // Successor of main_p0 (operand 0 of sort) is sort with shape_index {0}. - ASSERT_OK_AND_ASSIGN(std::vector succ_p0, - GetSuccessors(InstructionAndShapeIndex(main_p0, {}))); + ASSERT_OK_AND_ASSIGN( + std::vector succ_p0, + GetSuccessors(InstructionAndShapeIndex(main_p0, {}), *call_graph)); std::vector expected_succ_p0 = { InstructionAndShapeIndex(sort, {0})}; EXPECT_EQ(succ_p0, expected_succ_p0); // Successor of {sort, {0}} is gte0 with shape_index {}, not gte1. - ASSERT_OK_AND_ASSIGN(std::vector succ_sort0, - GetSuccessors(InstructionAndShapeIndex(sort, {0}))); + ASSERT_OK_AND_ASSIGN( + std::vector succ_sort0, + GetSuccessors(InstructionAndShapeIndex(sort, {0}), *call_graph)); std::vector expected_succ_sort0 = { InstructionAndShapeIndex(gte0, {})}; EXPECT_EQ(succ_sort0, expected_succ_sort0); // Predecessor of {sort, {1}} is main_p1 with shape_index {}. std::vector pred_sort1 = - GetPredecessors(InstructionAndShapeIndex(sort, {1})); + GetPredecessors(InstructionAndShapeIndex(sort, {1}), *call_graph); std::vector expected_pred_sort1 = { InstructionAndShapeIndex(main_p1, {})}; EXPECT_EQ(pred_sort1, expected_pred_sort1); diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD index 2ca1c537bfc8d0..8b05f417940548 100644 --- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD +++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD @@ -74,6 +74,8 @@ cc_library( deps = [ "//xla:side_effect_util", "//xla/hlo/ir:hlo_sharding", + "//xla/hlo/ir:mesh_and_axis", + "//xla/hlo/ir:named_sharding", "//xla/hlo/parser:hlo_parser", "//xla/hlo/translate/hlo_to_mhlo:hlo_utils", "//xla/service/spmd/shardy:constants", diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc index 30228198eaa7e2..4b3f5b085e2f92 100644 --- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc +++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc @@ -55,6 +55,8 @@ limitations under the License. #include "shardy/dialect/sdy/ir/utils.h" #include "stablehlo/dialect/StablehloOps.h" #include "xla/hlo/ir/hlo_sharding.h" +#include "xla/hlo/ir/mesh_and_axis.h" +#include "xla/hlo/ir/named_sharding.h" #include "xla/hlo/parser/hlo_parser.h" #include "xla/hlo/translate/hlo_to_mhlo/hlo_utils.h" #include "xla/service/spmd/shardy/constants.h" @@ -82,6 +84,7 @@ using ::mlir::func::FuncOp; using ::mlir::sdy::getSharding; using ::mlir::sdy::kShardingAttr; using ::mlir::sdy::kShardingRuleAttr; +using ::mlir::sdy::ManualAxesAttr; using ::mlir::sdy::MeshAttr; using ::mlir::sdy::OpShardingRuleAttr; using ::mlir::sdy::PropagationBarrierOp; @@ -206,6 +209,37 @@ bool handleFuncResultSharding( return anyChanged; } +ManualAxesAttr extractManualAxes(const xla::HloSharding& hloSharding, + mlir::MLIRContext* context) { + if (hloSharding.IsTuple()) { + // Manual axes are only attached to ranked tensors, so a leading token + // element can legitimately have none. Take the first element that does. + for (const xla::HloSharding& element : hloSharding.tuple_elements()) { + if (ManualAxesAttr manualAxes = extractManualAxes(element, context)) { + return manualAxes; + } + } + return nullptr; + } + if (hloSharding.UseNamedShardingLeaf()) { + const xla::NamedSharding& namedSharding = hloSharding.named_sharding(); + llvm::SmallVector manualAxes; + const xla::Mesh& mesh = namedSharding.mesh(); + for (const auto& axisRef : namedSharding.manual_axes()) { + // Silently dropping an axis here would yield a manual computation with + // fewer manual axes than it actually has, which miscompiles rather than + // fails, so treat an out-of-range index as a bug. + CHECK_LT(axisRef.mesh_axis_index(), mesh.axis_names().size()); + manualAxes.push_back(mlir::StringAttr::get( + context, mesh.axis_names()[axisRef.mesh_axis_index()])); + } + if (!manualAxes.empty()) { + return ManualAxesAttr::get(context, manualAxes); + } + } + return nullptr; +} + // The sharding information is in the `kXlaShardingAttr` attribute. void convertShardyAttrsWithHloShardingV3(FuncOp funcOp) { for (auto [argNum, argType] : llvm::enumerate(funcOp.getArgumentTypes())) { @@ -272,12 +306,38 @@ void convertShardyAttrsWithHloShardingV3(FuncOp funcOp) { StringRef targetName = customCallOp.getCallTargetName(); if (targetName == kShardingCustomCallTargetName || targetName == "X64Combine" || + targetName == kGlobalToLocalShapeCallTargetName || + targetName == kLocalToGlobalShapeCallTargetName || isPythonCallbackCustomCall(customCallOp)) { + HloSharding hloSharding = parseShardingFromString(shardingAttr); if (auto sdySharding = convertToSdySharding( - parseShardingFromString(shardingAttr), - customCallOp->getResultTypes(), customCallOp->getContext())) { + hloSharding, customCallOp->getResultTypes(), + customCallOp->getContext())) { customCallOp->setAttr(kShardingAttr, sdySharding); } + // Only `xla.sdy.GlobalToLocalShape` (when there are operands) and the + // manual computation CallOp (when there are no operands) may carry + // manual axes. + if (targetName == kGlobalToLocalShapeCallTargetName) { + if (auto manualAxesAttr = + extractManualAxes(hloSharding, customCallOp->getContext())) { + customCallOp->setAttr(kManualAxes, manualAxesAttr); + } + } + } + } else if (auto callOp = mlir::dyn_cast(op); + callOp && isManualComputation(callOp)) { + // The body call op carries the local out shardings and the manual axes. + // It is the only carrier of the manual axes when the manual computation + // has no operands, since there is no `xla.sdy.GlobalToLocalShape` then. + HloSharding hloSharding = parseShardingFromString(shardingAttr); + if (auto sdySharding = convertToSdySharding( + hloSharding, callOp->getResultTypes(), callOp->getContext())) { + callOp->setAttr(kShardingAttr, sdySharding); + } + if (auto manualAxesAttr = + extractManualAxes(hloSharding, callOp->getContext())) { + callOp->setAttr(kManualAxes, manualAxesAttr); } } diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc index 6b15a224661c57..cb87f6e39b80e0 100644 --- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc +++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc @@ -17,7 +17,10 @@ limitations under the License. #include +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" +#include "llvm/Support/CommandLine.h" #include "llvm/Support/ErrorHandling.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/IR/AffineMap.h" @@ -40,6 +43,7 @@ limitations under the License. #include "mlir/Support/LLVM.h" #include "mlir/Support/TypeID.h" #include "mlir/Transforms/DialectConversion.h" +#include "shardy/dialect/sdy/ir/constants.h" #include "shardy/dialect/sdy/ir/dialect.h" #include "shardy/dialect/sdy/ir/utils.h" #include "stablehlo/dialect/StablehloOps.h" @@ -59,14 +63,44 @@ using ::mlir::func::FuncOp; namespace stablehlo = ::mlir::stablehlo; namespace sdy = ::mlir::sdy; +// Returns `shardings` with `manualAxes` stripped from each element, i.e. the +// local view of tensors that are sharded over those axes. +// +// A `NamedSharding` may not mention the same axis twice, so an op whose result +// is a local tensor carries the manual axes in its `xla.sdy.manual_axes` +// attribute and the remaining axes in its dimension shardings. +sdy::TensorShardingPerValueAttr localizeShardings( + sdy::TensorShardingPerValueAttr shardings, + mlir::ArrayRef manualAxes, mlir::MLIRContext* context) { + llvm::SmallVector localShardings; + localShardings.reserve(shardings.size()); + for (sdy::TensorShardingAttr sharding : shardings.getShardings()) { + localShardings.push_back(sdy::eraseManualAxes(sharding, manualAxes)); + } + return sdy::TensorShardingPerValueAttr::get(context, localShardings); +} + class SdyRoundTripShardMapExportPass : public mlir::PassWrapper> { public: MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SdyRoundTripShardMapExportPass) - SdyRoundTripShardMapExportPass(bool enableHloShardingV3) - : enableHloShardingV3(enableHloShardingV3) {} + SdyRoundTripShardMapExportPass() = default; + explicit SdyRoundTripShardMapExportPass(bool enableHloShardingV3) { + this->enableHloShardingV3 = enableHloShardingV3; + } + SdyRoundTripShardMapExportPass(const SdyRoundTripShardMapExportPass& other) + : mlir::PassWrapper>(other) { + this->enableHloShardingV3 = other.enableHloShardingV3.getValue(); + } + + Option enableHloShardingV3{ + *this, "enable-hlo-sharding-v3", + llvm::cl::desc("Whether to enable HloShardingV3 which is the mesh and " + "axis based sharding representation."), + llvm::cl::init(false)}; void runOnOperation() final { ModuleOp moduleOp = getOperation(); @@ -90,28 +124,81 @@ class SdyRoundTripShardMapExportPass rewriter.setInsertionPoint(manualComputation); mlir::ValueRange operands = manualComputation->getOperands(); if (!operands.empty()) { + llvm::SmallVector globalOperands(operands.begin(), + operands.end()); + sdy::TensorShardingPerValueAttr inShardings = + manualComputation.getInShardings(); + if (enableHloShardingV3) { + inShardings = sdy::inlineMesh(symbolTable, inShardings); + // `GlobalToLocalShape`'s result is the *local* tensor, so it can only + // carry the local in shardings. Materialize the reshard that + // `ManualComputationOp` performs implicitly on its operands as an + // explicit `@Sharding` custom call. This keeps the exported module + // sharding consistent - every op's sharding describes its own result + // - and lets the import side read the global in shardings straight + // off the operand instead of inferring them. + for (auto [globalOperand, inSharding] : + llvm::zip_equal(globalOperands, inShardings.getShardings())) { + auto shardingOp = stablehlo::CustomCallOp::create( + rewriter, loc, globalOperand.getType(), globalOperand); + shardingOp.setCallTargetName(kShardingCustomCallTargetName); + // Import reads the global in shardings back off these wrappers, so + // they have to survive the HLO round trip intact. Without a side + // effect a wrapper of a zero-element tensor is replaced by a + // constant (see `ZeroSizedHloElimination`, which likewise spares + // the side-effecting `xla.sdy.GlobalToLocalShape`), which would + // lose that operand's sharding. The side effect also stops CSE from + // merging wrappers that carry different shardings. + shardingOp.setHasSideEffect(true); + shardingOp->setAttr(mlir::sdy::kShardingAttr, + sdy::TensorShardingPerValueAttr::get( + moduleOp.getContext(), inSharding)); + globalOperand = shardingOp.getResult(0); + } + } auto globalToLocalShape = stablehlo::CustomCallOp::create( - rewriter, loc, manualCompBodyArgTypes, operands); + rewriter, loc, manualCompBodyArgTypes, globalOperands); globalToLocalShape.setCallTargetName(kGlobalToLocalShapeCallTargetName); // We mark `xla.sdy.GlobalToLocalShape` as side-effecting to avoid // CSE deduping it with another taking the same operands, as it would // ignore the frontend attributes that could be different. globalToLocalShape.setHasSideEffect(true); - sdy::TensorShardingPerValueAttr inShardings = - manualComputation.getInShardings(); if (enableHloShardingV3) { - inShardings = sdy::inlineMesh(symbolTable, inShardings); + globalToLocalShape->setAttr( + mlir::sdy::kShardingAttr, + localizeShardings(inShardings, manualComputation.getManualAxes(), + moduleOp.getContext())); + globalToLocalShape->setAttr(kManualAxes, + manualComputation.getManualAxesAttr()); + } else { + setFrontendAttribute(globalToLocalShape, kInShardings, inShardings); + setFrontendAttribute(globalToLocalShape, kManualAxes, + manualComputation.getManualAxesAttr()); } - setFrontendAttribute(globalToLocalShape, kInShardings, inShardings); - setFrontendAttribute(globalToLocalShape, kManualAxes, - manualComputation.getManualAxesAttr()); operands = globalToLocalShape->getResults(); } + sdy::TensorShardingPerValueAttr outShardings = + manualComputation.getOutShardings(); + if (enableHloShardingV3) { + outShardings = sdy::inlineMesh(symbolTable, outShardings); + } + auto callOp = CallOp::create(rewriter, loc, localResultTypes, funcName, operands); setFrontendAttribute(callOp, kXlaInlineableAttr, rewriter.getStringAttr("xla_late")); + if (enableHloShardingV3 && !outShardings.getShardings().empty()) { + // The call op produces the local tensors, so it carries the local view + // of the out shardings together with the manual axes. This is also the + // only carrier of the manual axes when the manual computation has no + // operands, i.e. when there is no `xla.sdy.GlobalToLocalShape`. + callOp->setAttr( + mlir::sdy::kShardingAttr, + localizeShardings(outShardings, manualComputation.getManualAxes(), + moduleOp.getContext())); + callOp->setAttr(kManualAxes, manualComputation.getManualAxesAttr()); + } mlir::ResultRange results = manualComputation->getResults(); if (!results.empty()) { @@ -122,14 +209,13 @@ class SdyRoundTripShardMapExportPass // We mark `xla.sdy.LocalToGlobalShape` as side-effecting to avoid // CSE removing it if it has no users. localToGlobalShape.setHasSideEffect(true); - sdy::TensorShardingPerValueAttr outShardings = - manualComputation.getOutShardings(); if (enableHloShardingV3) { - outShardings = sdy::inlineMesh(symbolTable, outShardings); + localToGlobalShape->setAttr(mlir::sdy::kShardingAttr, outShardings); + } else { + setFrontendAttribute(localToGlobalShape, kOutShardings, outShardings); + setFrontendAttribute(localToGlobalShape, kManualAxes, + manualComputation.getManualAxesAttr()); } - setFrontendAttribute(localToGlobalShape, kOutShardings, outShardings); - setFrontendAttribute(localToGlobalShape, kManualAxes, - manualComputation.getManualAxesAttr()); results = localToGlobalShape->getResults(); } sdy::inlineRegionAndConvertTerminatorOp( @@ -147,15 +233,16 @@ class SdyRoundTripShardMapExportPass "1. A separate function for the body of the `ManualComputationOp`." "2. A `CallOp` calling the function in #1, marked as not inlinable." "3. A pair of `CustomCallOp`s that change the shape of the " - " arguments/results. They save the in/out shardings and manual " - " axes as frontend attrs."; + " arguments/results." + " Under HloShardingV3 the shardings and manual axes are attached as " + "native attributes, and each operand is additionally wrapped in a " + "`@Sharding` custom call holding the global in sharding so that " + "every op's sharding describes its own result. Otherwise they are " + "saved as frontend attrs on the pair of `CustomCallOp`s."; } void getDependentDialects(mlir::DialectRegistry& registry) const final { registry.insert(); } - - private: - bool enableHloShardingV3 = false; }; } // namespace diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc index 2e2151e8d11f62..29ea30dd68e693 100644 --- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc +++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc @@ -75,6 +75,50 @@ using ::mlir::stablehlo::CustomCallOp; namespace sdy = ::mlir::sdy; +// Recovers the global in shardings of a manual computation from the +// `@Sharding` custom calls that export wraps around each operand of +// `globalToLocalShape`, and removes those wrappers. +// +// Export materializes the reshard that `sdy.manual_computation` performs +// implicitly on entry as an explicit `@Sharding` op, so that every op's +// sharding describes its own result. Undoing it here makes the round trip an +// identity. +// +// Returns an empty vector, and leaves the IR untouched, if any operand is not +// wrapped - the shardings are only meaningful as a complete set. +SmallVector takeGlobalInShardings( + CustomCallOp globalToLocalShape, mlir::IRRewriter& rewriter) { + SmallVector globalInShardings; + SmallVector shardingOps; + for (mlir::Value globalOperand : globalToLocalShape.getOperands()) { + auto shardingOp = globalOperand.getDefiningOp(); + if (!shardingOp || + shardingOp.getCallTargetName() != kShardingCustomCallTargetName) { + return {}; + } + sdy::TensorShardingAttr sharding = + sdy::getSharding(shardingOp.getResult(0)); + if (!sharding) { + return {}; + } + globalInShardings.push_back(sharding); + shardingOps.push_back(shardingOp); + } + + for (auto [i, shardingOp] : llvm::enumerate(shardingOps)) { + globalToLocalShape->setOperand(i, shardingOp.getOperand(0)); + } + // The same wrapper can feed several operands, so only visit it once - after + // the first erase the handle is dangling. + llvm::SmallDenseSet seen; + for (CustomCallOp shardingOp : shardingOps) { + if (seen.insert(shardingOp).second && shardingOp.use_empty()) { + rewriter.eraseOp(shardingOp); + } + } + return globalInShardings; +} + mlir::LogicalResult rewriteManualComputation( CallOp callOp, mlir::IRRewriter& rewriter, const mlir::SymbolTable& symbolTable) { @@ -96,7 +140,7 @@ mlir::LogicalResult rewriteManualComputation( // `ManualComputationOp` differently depending on whether the original had // operands/results. CustomCallOp globalToLocalShape; - mlir::ValueRange operands = callOp.getOperands(); + SmallVector operands(callOp.getOperands()); if (!operands.empty()) { // An input to `sdy.manual_computation` can have a dimension of size 0 // (i.e. 0 num-elements), in which case, the corresponding result of @@ -114,7 +158,8 @@ mlir::LogicalResult rewriteManualComputation( globalToLocalShape = (*customCallResIt).getDefiningOp(); CHECK_EQ(globalToLocalShape.getCallTargetName(), kGlobalToLocalShapeCallTargetName); - operands = globalToLocalShape->getOperands(); + operands.assign(globalToLocalShape->operand_begin(), + globalToLocalShape->operand_end()); } mlir::TypeRange resultTypes = callOp->getResultTypes(); @@ -150,10 +195,13 @@ mlir::LogicalResult rewriteManualComputation( if (!customCallOp) { return; } + if (mlir::DictionaryAttr frontendAttrs = getFrontendAttrs(customCallOp)) { - shardings = parseStringAttr( - frontendAttrs, shardingAttrName); - if (manualAxes.empty()) { + if (hasKey(frontendAttrs, shardingAttrName)) { + shardings = parseStringAttr( + frontendAttrs, shardingAttrName); + } + if (manualAxes.empty() && hasKey(frontendAttrs, kManualAxes)) { manualAxes = parseStringAttr(frontendAttrs, kManualAxes); } @@ -162,6 +210,45 @@ mlir::LogicalResult rewriteManualComputation( setShardingAttrs(globalToLocalShape, inShardings, kInShardings); setShardingAttrs(localToGlobalShape, outShardings, kOutShardings); + + // Under HloShardingV3 the shardings and manual axes round trip as native + // attributes rather than frontend attributes, so read them directly. + if (outShardings.empty() && localToGlobalShape) { + // `LocalToGlobalShape` produces the global tensors, so its own sharding is + // the global out shardings. + if (auto sharding = + localToGlobalShape->getAttrOfType( + sdy::kShardingAttr)) { + outShardings = sharding; + } + } + + if (manualAxes.empty()) { + // A manual computation without operands has no `GlobalToLocalShape`, in + // which case the body call op is the only carrier of the manual axes. + for (Operation* op : + {globalToLocalShape.getOperation(), callOp.getOperation()}) { + auto axes = + op ? op->getAttrOfType(kManualAxes) : nullptr; + if (axes && !axes.empty()) { + manualAxes = axes; + break; + } + } + } + + if (inShardings.empty() && globalToLocalShape) { + SmallVector globalInShardings = + takeGlobalInShardings(globalToLocalShape, rewriter); + if (!globalInShardings.empty()) { + inShardings = + sdy::TensorShardingPerValueAttr::get(context, globalInShardings); + // `takeGlobalInShardings` rewired `globalToLocalShape` past the wrappers. + operands.assign(globalToLocalShape->operand_begin(), + globalToLocalShape->operand_end()); + } + } + auto manualComputationOp = rewriter.replaceOpWithNewOp( callOp, resultTypes, operands, inShardings, outShardings, manualAxes); @@ -276,10 +363,12 @@ class SdyRoundTripShardMapImportPass } StringRef getDescription() const override { - return "converts a CallOp calling a @xla.sdy.manual_computation_body func " - "with in/out shardings and manual axes as frontend attrs, wrapped " - "with a pair of `CustomCallOps` that change the shape of the " - "arguments/results, to a ManualComputationOp"; + return "converts a CallOp calling a @xla.sdy.manual_computation_body func, " + "wrapped with a pair of `CustomCallOps` that change the shape of " + "the arguments/results, to a ManualComputationOp. The in/out " + "shardings and manual axes are read from frontend attrs, or, under " + "HloShardingV3, from native attributes on the boundary ops and on " + "the `@Sharding` custom calls wrapping the operands"; } void getDependentDialects(mlir::DialectRegistry& registry) const final { registry.insert(); diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc index 77e4628feeadd9..e1f91c750f6882 100644 --- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc +++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc @@ -136,6 +136,11 @@ void exportFunc(FuncOp funcOp, const SymbolTable& symbolTable, if (ManualAxesAttr manualAxesAttr = funcOp.getArgAttrOfType(argNum, kManualAxes)) { manualAxes = manualAxesAttr.getValue(); + if (mlir::isa( + funcOp.getArgument(argNum).getType()) && + getMeshAttr(sdySharding).getAxes().size() != manualAxes.size()) { + manualAxes = {}; + } attrs.erase(kManualAxes); } attrs.set(kXlaShardingAttr, @@ -161,6 +166,10 @@ void exportFunc(FuncOp funcOp, const SymbolTable& symbolTable, if (auto manualAxesAttr = mlir::dyn_cast_or_null(attrs.get(kManualAxes))) { manualAxes = manualAxesAttr.getValue(); + if (mlir::isa(funcOp.getResultTypes()[resNum]) && + getMeshAttr(sdySharding).getAxes().size() != manualAxes.size()) { + manualAxes = {}; + } attrs.erase(kManualAxes); } attrs.set(kXlaShardingAttr, @@ -319,24 +328,38 @@ HloSharding getHloShardingForOp( std::function getMeshAttr, ArrayRef manualAxes, bool enableHloShardingV3, bool simplifyReplicatedShardings) { + auto getManualAxesForType = [&](mlir::Type type, + TensorShardingAttr sdySharding) { + if (mlir::isa(type) && + getMeshAttr(sdySharding).getAxes().size() != manualAxes.size()) { + return ArrayRef(); + } + return manualAxes; + }; + bool isNoResultMaximal = op->getNumResults() == 0 && shardings.size() == 1 && (getMeshAttr(shardings.front()).isMaximal() || shardings.front().isFullyReplicated()); CHECK(shardings.size() == op->getNumResults() || isNoResultMaximal); if (op->getNumResults() == 1 || isNoResultMaximal) { - return convertToHloSharding(shardings.front(), getMeshAttr, manualAxes, - enableHloShardingV3, + ArrayRef resultManualAxes = + op->getNumResults() == 1 + ? getManualAxesForType(op->getResultTypes().front(), + shardings.front()) + : manualAxes; + return convertToHloSharding(shardings.front(), getMeshAttr, + resultManualAxes, enableHloShardingV3, simplifyReplicatedShardings); } std::vector newShardings; newShardings.reserve(shardings.size()); - llvm::transform(shardings, std::back_inserter(newShardings), - [&](TensorShardingAttr sdySharding) { - return convertToHloSharding(sdySharding, getMeshAttr, - manualAxes, enableHloShardingV3, - simplifyReplicatedShardings); - }); + for (auto [type, sdySharding] : + llvm::zip_equal(op->getResultTypes(), shardings)) { + newShardings.push_back(convertToHloSharding( + sdySharding, getMeshAttr, getManualAxesForType(type, sdySharding), + enableHloShardingV3, simplifyReplicatedShardings)); + } std::vector shapes; shapes.reserve(op->getNumResults()); @@ -502,11 +525,6 @@ HloSharding convertToHloSharding( if (mesh.getAxes().size() == manualAxes.size()) { return HloSharding::Manual(); } - // TODO(b/438306205): Remove this check once we support both unreduced and - // manual axes in subgroup sharding. - CHECK(sdySharding.getUnreducedAxes().empty() || manualAxes.empty()) - << "Only one of unreduced and manual axes can be present: " - << mlir::sdy::attributeToString(sdySharding); // Iterate the dim shardings. for (auto [index, dimSharding] : diff --git a/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir b/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir index 9a05f216ae8a8e..597d749388fdf6 100644 --- a/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir +++ b/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir @@ -154,3 +154,42 @@ func.func @main( } : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) return %0#0, %0#1 : !stablehlo.token, tensor<2xi64> } + +// ----- + +// ***************** Zero-element operand test ***************** + +// A zero-element operand is replaced by a constant during the HLO round trip, +// which drops any non side-effecting op producing it. Make sure the remaining +// operands keep their in shardings rather than the manual computation coming +// back with none at all. + +// Make sure this temp attr doesn't exist anymore. +// CHECK-NOT: sharding_hlo_string + +// CHECK-V2: sdy.mesh @mesh_1 = <["a"=4]> +// CHECK-V3: sdy.mesh @mesh = <["a"=4]> +sdy.mesh @mesh_1 = <["a"=4]> + +// CHECK-LABEL: func.func @main +func.func @main(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> { + // CHECK-NEXT: %[[MANUAL_COMP:.*]] = sdy.manual_computation(%arg0) + // CHECK-V2-SAME{LITERAL}: in_shardings=[<@mesh_1, [{"a"}, {}]>] out_shardings=[<@mesh_1, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<1x8xi32>) { + // CHECK-V3-SAME{LITERAL}: in_shardings=[<@mesh, [{"a"}, {}]>] out_shardings=[<@mesh, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<1x8xi32>) { + // CHECK-NEXT: sdy.return %arg1 : tensor<1x8xi32> + // CHECK-NEXT: } : (tensor<4x8xi32>) -> tensor<4x8xi32> + // CHECK-NEXT: return %[[MANUAL_COMP]] : tensor<4x8xi32> + %c = stablehlo.constant dense<> : tensor<0xi32> + %0 = sdy.manual_computation(%c, %arg0) + in_shardings=[<@mesh_1, [{}], replicated={"a"}>, <@mesh_1, [{"a"}, {}]>] + out_shardings=[<@mesh_1, [{"a"}, {}]>] + manual_axes={"a"} (%arg1: tensor<0xi32>, %arg2: tensor<1x8xi32>) { + %init = stablehlo.constant dense<0> : tensor + %red = stablehlo.reduce(%arg1 init: %init) applies stablehlo.add across dimensions = [0] : (tensor<0xi32>, tensor) -> tensor + %bcast = stablehlo.broadcast_in_dim %red, dims = [] : (tensor) -> tensor<1x8xi32> + %add = stablehlo.add %arg2, %bcast : tensor<1x8xi32> + sdy.return %add : tensor<1x8xi32> + } : (tensor<0xi32>, tensor<4x8xi32>) -> tensor<4x8xi32> + return %0 : tensor<4x8xi32> +} + diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir index 58b7012ab5b707..1b0f94f6667750 100644 --- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir +++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir @@ -204,6 +204,30 @@ module @module_1 { return %3 : tensor<1xi32> } + // Tests importing shard map custom calls annotated with HloShardingV3 mhlo.sharding. + // CHECK-LABEL: func @manual_computation_hlo_sharding_v3 + // CHECK-SAME: (%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> { + // CHECK-NEXT: %[[MAN_COMP:.*]] = sdy.manual_computation(%arg0) + // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_2, [{"a"}, {}]>] out_shardings=[<@mesh_2, [{"a"}, {}]>] manual_axes={"a"} + // CHECK-SAME: (%arg1: tensor<2x16xf32>) { + // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg1, %arg1 : tensor<2x16xf32> + // CHECK-NEXT: sdy.return %[[ADD]] : tensor<2x16xf32> + // CHECK-NEXT: } : (tensor<8x16xf32>) -> tensor<8x16xf32> + // CHECK-NEXT: return %[[MAN_COMP]] : tensor<8x16xf32> + // CHECK-NEXT: } + func.func @manual_computation_hlo_sharding_v3(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> { + %s0 = stablehlo.custom_call @Sharding(%arg0) {mhlo.sharding = "{mesh['a'=4,'b'=2], [{'a'}, {}]}"} : (tensor<8x16xf32>) -> tensor<8x16xf32> + %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, mhlo.sharding = "{mesh['a'=4,'b'=2], [{}, {}], manual={'a'}}"} : (tensor<8x16xf32>) -> tensor<2x16xf32> + %1 = call @xla.sdy.manual_computation_body_hlo_v3(%0) : (tensor<2x16xf32>) -> tensor<2x16xf32> + %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, mhlo.sharding = "{mesh['a'=4,'b'=2], [{'a'}, {}]}"} : (tensor<2x16xf32>) -> tensor<8x16xf32> + return %2 : tensor<8x16xf32> + } + + func.func private @xla.sdy.manual_computation_body_hlo_v3(%arg0: tensor<2x16xf32>) -> tensor<2x16xf32> { + %0 = stablehlo.add %arg0, %arg0 : tensor<2x16xf32> + return %0 : tensor<2x16xf32> + } + // CHECK-LABEL: func @frontend_attr_not_sharding // CHECK-SAME: %arg0: tensor<16x8xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a", "b"}, {?}]>}, // CHECK-SAME: %arg1: tensor<16x8xf32> {mhlo.frontend_attributes = {baz = 1 : i32, foo = "bar"}}, diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_export_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_export_hlo_sharding_v3.mlir new file mode 100644 index 00000000000000..fc83e872e68335 --- /dev/null +++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_export_hlo_sharding_v3.mlir @@ -0,0 +1,285 @@ +// Copyright 2026 The OpenXLA Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ============================================================================== +// RUN: sdy_opt %s -xla-sdy-round-trip-shard-map-export='enable-hlo-sharding-v3=true' 2>&1 | FileCheck %s + +// Tests exporting `sdy.manual_computation` under HloShardingV3 (`enable-hlo-sharding-v3=true`). +// Verifies that: +// 1. Every op carries a native `sdy.sharding` that describes its own result: +// - a `@Sharding` custom call per operand holds the global in-shardings, +// making the reshard `sdy.manual_computation` performs implicitly on entry +// explicit. It is side-effecting so that the HLO round trip cannot drop it +// (see `ZeroSizedHloElimination` for a zero-element operand) or CSE two +// wrappers carrying different shardings; +// - `GlobalToLocalShape` and the body call op hold the local shardings plus +// typed `xla.sdy.manual_axes`; +// - `LocalToGlobalShape` holds the global out-shardings. +// 2. Legacy string `mhlo.frontend_attributes` (xla.sdy.in_shardings, etc.) are NOT emitted. + +sdy.mesh @mesh_0 = <["a"=4, "b"=2]> +sdy.mesh @mesh_1 = <["a"=2, "b"=2, "c"=2, "d"=2]> + +// Ensure string frontend attributes are completely absent in V3 export: +// CHECK-NOT: xla.sdy.in_shardings +// CHECK-NOT: xla.sdy.out_shardings + +// Tests basic export with multiple inputs, outputs, and all-reduce. +// CHECK-LABEL: func @single_manual_comp +func.func @single_manual_comp(%arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a", ?}, {"b", ?}]>}, %arg1: tensor<16x32xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"b", ?}, {?}]>}) -> (tensor<8x32xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a"}, {}]>}) { + // CHECK-NEXT: %[[SHARDING_0:.*]] = stablehlo.custom_call @Sharding(%arg0) + // CHECK-SAME: has_side_effect = true + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {"b"}]>]> + // CHECK-SAME: (tensor<8x16xf32>) -> tensor<8x16xf32> + // CHECK-NEXT: %[[SHARDING_1:.*]] = stablehlo.custom_call @Sharding(%arg1) + // CHECK-SAME: has_side_effect = true + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"b"}, {}], replicated={"a"}>]> + // CHECK-SAME: (tensor<16x32xf32>) -> tensor<16x32xf32> + // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]]:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_0]], %[[SHARDING_1]]) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>, , [{}, {}]>]> + // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy + // CHECK-SAME: (tensor<8x16xf32>, tensor<16x32xf32>) -> (tensor<2x8xf32>, tensor<8x32xf32>) + // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body(%[[GLOBAL_TO_LOCAL]]#0, %[[GLOBAL_TO_LOCAL]]#1) + // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> + // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy + // CHECK-SAME: (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> + // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}], replicated={"b"}>]> + // CHECK-SAME: (tensor<2x32xf32>) -> tensor<8x32xf32> + // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<8x32xf32> + %0 = sdy.manual_computation(%arg0, %arg1) in_shardings=[<@mesh_0, [{"a"}, {"b"}]>, <@mesh_0, [{"b"}, {}], replicated={"a"}>] out_shardings=[<@mesh_0, [{"a"}, {}], replicated={"b"}>] manual_axes={"a", "b"} (%arg2: tensor<2x8xf32>, %arg3: tensor<8x32xf32>) { + %1 = stablehlo.add %arg2, %arg2 : tensor<2x8xf32> + %2 = stablehlo.dot %1, %arg3 : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> + %3 = "stablehlo.all_reduce"(%2) <{channel_handle = #stablehlo.channel_handle, replica_groups = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi64>, use_global_device_ids}> ({ + ^bb0(%arg4: tensor, %arg5: tensor): + %4 = stablehlo.add %arg4, %arg5 : tensor + stablehlo.return %4 : tensor + }) : (tensor<2x32xf32>) -> tensor<2x32xf32> + sdy.return %3 : tensor<2x32xf32> + } : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32> + return %0 : tensor<8x32xf32> +} + +// Tests chained manual computations where the output of one feeds the input of +// another with a *different* sharding. The `@Sharding` custom call on the +// second boundary is what preserves the reshard between them; without it the +// `in_shardings` of the second manual computation would be lost. +// CHECK-LABEL: func @manual_comp_using_another +func.func @manual_comp_using_another(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a"}, {}]>}) + -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{}, {"b"}]>}) { + // CHECK-NEXT: %[[SHARDING_0:.*]] = stablehlo.custom_call @Sharding(%arg0) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> + // CHECK-NEXT: %[[GLOBAL_TO_LOCAL_0:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_0]]) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> + // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy + // CHECK-SAME: (tensor<8x8xf32>) -> tensor<2x8xf32> + // CHECK-NEXT: %[[SHMAP_0:.*]] = call @xla.sdy.manual_computation_body_0(%[[GLOBAL_TO_LOCAL_0]]) + // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} + // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy + // CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x8xf32> + // CHECK-NEXT: %[[LOCAL_TO_GLOBAL_0:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP_0]]) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> + // CHECK-SAME: (tensor<2x8xf32>) -> tensor<8x8xf32> + // CHECK-NEXT: %[[SHARDING_1:.*]] = stablehlo.custom_call @Sharding(%[[LOCAL_TO_GLOBAL_0]]) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> + // CHECK-NEXT: %[[GLOBAL_TO_LOCAL_1:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_1]]) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> + // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy + // CHECK-SAME: (tensor<8x8xf32>) -> tensor<8x4xf32> + // CHECK-NEXT: %[[SHMAP_1:.*]] = call @xla.sdy.manual_computation_body_1(%[[GLOBAL_TO_LOCAL_1]]) + // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} + // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy + // CHECK-SAME: (tensor<8x4xf32>) -> tensor<8x4xf32> + // CHECK-NEXT: %[[LOCAL_TO_GLOBAL_1:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP_1]]) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> + // CHECK-SAME: (tensor<8x4xf32>) -> tensor<8x8xf32> + // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL_1]] : tensor<8x8xf32> + %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh_0, [{"a"}, {}]>] out_shardings=[<@mesh_0, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<2x8xf32>) { + sdy.return %arg1 : tensor<2x8xf32> + } : (tensor<8x8xf32>) -> tensor<8x8xf32> + + %1 = sdy.manual_computation(%0) in_shardings=[<@mesh_0, [{}, {"b"}]>] out_shardings=[<@mesh_0, [{}, {"b"}]>] manual_axes={"b"} (%arg1: tensor<8x4xf32>) { + sdy.return %arg1 : tensor<8x4xf32> + } : (tensor<8x8xf32>) -> tensor<8x8xf32> + return %1 : tensor<8x8xf32> +} + +// Tests nested manual computations: an inner shard_map inside an outer shard_map body. +// CHECK-LABEL: func @nested_shmaps +func.func @nested_shmaps(%arg0: tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a"}, {"b"}]>}) -> (tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a", ?}, {?}]>}) { + // CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> + // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> + // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy + // CHECK-SAME: (tensor<4x8xf32>) -> tensor<2x8xf32> + // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_3(%[[GLOBAL_TO_LOCAL]]) + // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} + // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy + // CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x8xf32> + // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> + // CHECK-SAME: (tensor<2x8xf32>) -> tensor<4x8xf32> + // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<4x8xf32> + %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh_1, [{"a"}, {}]>] out_shardings=[<@mesh_1, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<2x8xf32>) { + %1 = sdy.manual_computation(%arg1) in_shardings=[<@mesh_1, [{}, {"b"}]>] out_shardings=[<@mesh_1, [{}, {"b"}]>] manual_axes={"b"} (%arg2: tensor<2x4xf32>) { + %2 = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32> + sdy.return %2 : tensor<2x4xf32> + } : (tensor<2x8xf32>) -> tensor<2x8xf32> + sdy.return %1 : tensor<2x8xf32> + } : (tensor<4x8xf32>) -> tensor<4x8xf32> + return %0 : tensor<4x8xf32> +} + +// Tests nested manual computation where the outer body performs additional math on the inner result. +// CHECK-LABEL: func @nested_shmaps_extra_op +func.func @nested_shmaps_extra_op(%arg0: tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a"}, {"b"}]>}) -> (tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a", ?}, {?}]>}) { + // CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> + // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> + // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy + // CHECK-SAME: (tensor<4x8xf32>) -> tensor<2x8xf32> + // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_5(%[[GLOBAL_TO_LOCAL]]) + // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} + // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy + // CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x8xf32> + // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"a"}, {}]>]> + // CHECK-SAME: (tensor<2x8xf32>) -> tensor<4x8xf32> + // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<4x8xf32> + %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh_1, [{"a"}, {}]>] out_shardings=[<@mesh_1, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<2x8xf32>) { + %1 = sdy.manual_computation(%arg1) in_shardings=[<@mesh_1, [{}, {"b"}]>] out_shardings=[<@mesh_1, [{}, {"b"}]>] manual_axes={"b"} (%arg2: tensor<2x4xf32>) { + %2 = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32> + sdy.return %2 : tensor<2x4xf32> + } : (tensor<2x8xf32>) -> tensor<2x8xf32> + %3 = stablehlo.add %1, %1 : tensor<2x8xf32> + sdy.return %3 : tensor<2x8xf32> + } : (tensor<4x8xf32>) -> tensor<4x8xf32> + return %0 : tensor<4x8xf32> +} + +// Edge case: Manual computation producing an output but taking no inputs. There +// is no `GlobalToLocalShape`, so the body call op is the only carrier of the +// manual axes. +// CHECK-LABEL: func @manual_computation_no_inputs +func.func @manual_computation_no_inputs() -> tensor<4xi64> { + // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_6() + // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}]>]> + // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy + // CHECK-SAME: () -> tensor<2xi64> + // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"b"}]>]> + // CHECK-SAME: (tensor<2xi64>) -> tensor<4xi64> + // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<4xi64> + %0 = sdy.manual_computation() in_shardings=[] out_shardings=[<@mesh_0, [{"b"}]>] manual_axes={"b"} () { + %1 = stablehlo.constant dense<[2, 3]> : tensor<2xi64> + sdy.return %1 : tensor<2xi64> + } : () -> tensor<4xi64> + func.return %0 : tensor<4xi64> +} + +// Edge case: Manual computation consuming an input but producing no outputs. +// The body call op has no results, so `GlobalToLocalShape` is the only carrier +// of the manual axes. +// CHECK-LABEL: func @manual_computation_no_outputs +func.func @manual_computation_no_outputs(%arg0: tensor<4xi64>) { + // CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"b"}]>]> + // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}]>]> + // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy + // CHECK-SAME: (tensor<4xi64>) -> tensor<2xi64> + // CHECK-NEXT: call @xla.sdy.manual_computation_body_7(%[[GLOBAL_TO_LOCAL]]) + // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"}} + // CHECK-SAME: (tensor<2xi64>) -> () + // CHECK-NEXT: return + sdy.manual_computation(%arg0) in_shardings=[<@mesh_0, [{"b"}]>] out_shardings=[] manual_axes={"b"} (%arg1: tensor<2xi64>) { + stablehlo.custom_call @sdy_testonly(%arg1) : (tensor<2xi64>) -> () + sdy.return + } : (tensor<4xi64>) -> () + func.return +} + +// Edge case: Degenerate manual computation with no inputs and no outputs. +// CHECK-LABEL: func @manual_computation_no_inputs_no_outputs +func.func @manual_computation_no_inputs_no_outputs() { + // CHECK-NEXT: call @xla.sdy.manual_computation_body_8() {mhlo.frontend_attributes = {inlineable = "xla_late"}} : () -> () + sdy.manual_computation() in_shardings=[] out_shardings=[] manual_axes={} () { + sdy.return + } : () -> () + func.return +} + +// Type preservation: Token arguments should not have manual axes attached in V3. +// CHECK-LABEL: func @manual_computation_tokens +func.func @manual_computation_tokens(%arg0: !stablehlo.token, %arg1: tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) { + // CHECK-NEXT: %[[SHARDING_0:.*]] = stablehlo.custom_call @Sharding(%arg0) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, []>]> + // CHECK-SAME: (!stablehlo.token) -> !stablehlo.token + // CHECK-NEXT: %[[SHARDING_1:.*]] = stablehlo.custom_call @Sharding(%arg1) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{"b"}]>]> + // CHECK-SAME: (tensor<2xi64>) -> tensor<2xi64> + // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]]:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_0]], %[[SHARDING_1]]) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, []>, , [{}]>]> + // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy + // CHECK-SAME: (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<1xi64>) + // CHECK-NEXT: %[[SHMAP:.*]]:2 = call @xla.sdy.manual_computation_body_9(%[[GLOBAL_TO_LOCAL]]#0, %[[GLOBAL_TO_LOCAL]]#1) + // CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} + // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy + // CHECK-SAME: (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<1xi64>) + // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]]:2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]#0, %[[SHMAP]]#1) + // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, []>, , [{"b"}]>]> + // CHECK-SAME: (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<2xi64>) + // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]]#0, %[[LOCAL_TO_GLOBAL]]#1 : !stablehlo.token, tensor<2xi64> + %0:2 = sdy.manual_computation(%arg0, %arg1) in_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>] out_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>] manual_axes={"b"} (%arg2: !stablehlo.token, %arg3: tensor<1xi64>) { + sdy.return %arg2, %arg3 : !stablehlo.token, tensor<1xi64> + } : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) + func.return %0#0, %0#1 : !stablehlo.token, tensor<2xi64> +} + + +// CHECK-LABEL: func @xla.sdy.manual_computation_body_3(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> +// CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) +// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> +// CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) +// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> +// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy +// CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x4xf32> +// CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_2(%[[GLOBAL_TO_LOCAL]]) +// CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} +// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy +// CHECK-SAME: (tensor<2x4xf32>) -> tensor<2x4xf32> +// CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) +// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> +// CHECK-SAME: (tensor<2x4xf32>) -> tensor<2x8xf32> +// CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<2x8xf32> + +// CHECK-LABEL: func @xla.sdy.manual_computation_body_5(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> +// CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0) +// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> +// CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]]) +// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {}]>]> +// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy +// CHECK-SAME: (tensor<2x8xf32>) -> tensor<2x4xf32> +// CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_4(%[[GLOBAL_TO_LOCAL]]) +// CHECK-SAME: {mhlo.frontend_attributes = {inlineable = "xla_late"} +// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy +// CHECK-SAME: (tensor<2x4xf32>) -> tensor<2x4xf32> +// CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]) +// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[, [{}, {"b"}]>]> +// CHECK-SAME: (tensor<2x4xf32>) -> tensor<2x8xf32> +// CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[LOCAL_TO_GLOBAL]], %[[LOCAL_TO_GLOBAL]] : tensor<2x8xf32> +// CHECK-NEXT: return %[[ADD]] : tensor<2x8xf32> diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir index 47d1ba107a8bf5..15b7eec24bdca0 100644 --- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir +++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir @@ -45,3 +45,38 @@ func.func @xla.sdy.manual_computation_body(%arg0: tensor<0x8xf32>) -> tensor<0x8 return %arg0 : tensor<0x8xf32> } +// ----- + +sdy.mesh @mesh = <["a"=2, "b"=2]> + +// Tests import error when manual computation body call operand is not produced by GlobalToLocalShape. +func.func @manual_computation_missing_global_to_local_shape_direct_attrs(%arg0: tensor<0x16xf32>) -> (tensor<0x16xf32>) { + %c = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32> + // expected-error @+2 {{'func.call' op expected at least one operand of @xla.sdy.manual_computation_body to be produced by a xla.sdy.GlobalToLocalShape CustomCallOp}} + // expected-error @+1 {{failed to rewrite func.call to manual computation}} + %0 = call @xla.sdy.manual_computation_body(%c) : (tensor<0x8xf32>) -> tensor<0x8xf32> + %1 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<0x8xf32>) -> (tensor<0x16xf32>) + return %1 : tensor<0x16xf32> +} + +func.func @xla.sdy.manual_computation_body(%arg0: tensor<0x8xf32>) -> tensor<0x8xf32> { + return %arg0 : tensor<0x8xf32> +} + +// ----- + +sdy.mesh @mesh = <["a"=2, "b"=2]> + +// Tests import error when manual computation body call result is not consumed by LocalToGlobalShape. +func.func @manual_computation_missing_local_to_global_shape_direct_attrs(%arg0: tensor<0x16xf32>) -> (tensor<0x8xf32>) { + %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%arg0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<0x16xf32>) -> tensor<0x8xf32> + // expected-error @+2 {{'func.call' op expected the first use of @xla.sdy.manual_computation_body to be a xla.sdy.LocalToGlobalShape CustomCallOp}} + // expected-error @+1 {{failed to rewrite func.call to manual computation}} + %1 = call @xla.sdy.manual_computation_body(%0) : (tensor<0x8xf32>) -> tensor<0x8xf32> + return %1 : tensor<0x8xf32> +} + +func.func @xla.sdy.manual_computation_body(%arg0: tensor<0x8xf32>) -> tensor<0x8xf32> { + return %arg0 : tensor<0x8xf32> +} + diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_hlo_sharding_v3.mlir new file mode 100644 index 00000000000000..6e7060cb48ff30 --- /dev/null +++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_hlo_sharding_v3.mlir @@ -0,0 +1,362 @@ +// Copyright 2026 The OpenXLA Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// ============================================================================== +// RUN: sdy_opt %s --split-input-file -xla-sdy-test-flatten-call-graph -xla-sdy-round-trip-shard-map-import 2>&1 | FileCheck %s + +// Tests importing shard map custom calls back into `sdy.manual_computation` under +// HloShardingV3 (where frontend_attributes are omitted). +// Verifies that: +// 1. Shardings and manual axes are read directly off the native attributes of the +// boundary ops, without any inference: +// - global `in_shardings` come from the `@Sharding` custom call wrapping each +// operand of `GlobalToLocalShape`, which is then unwired and erased; +// - global `out_shardings` come from `LocalToGlobalShape`; +// - `manual_axes` come from `GlobalToLocalShape`, or from the body call op +// when the manual computation has no operands. +// 2. Outlined manual computation call bodies are inlined back into `sdy.manual_computation`. + +sdy.mesh @mesh_0 = <["a"=4, "b"=2]> +sdy.mesh @mesh_1 = <["a"=2, "b"=2, "c"=2, "d"=2]> + +// Tests single manual computation with multiple inputs/outputs, all-reduce, and +// in/out shardings read off the `@Sharding` wrappers and `LocalToGlobalShape`. +// CHECK-LABEL: func @single_manual_comp +func.func @single_manual_comp(%arg0: tensor<8x16xf32>, %arg1: tensor<16x32xf32>) -> (tensor<8x32xf32>) { + // CHECK-NOT: call @xla.sdy.manual_computation_body + // CHECK: %[[MAN_COMP:.*]] = sdy.manual_computation(%arg0, %arg1) + // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{"a"}, {"b"}]>, <@mesh_0, [{"b"}, {}], replicated={"a"}>] + // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{"a"}, {}], replicated={"b"}>] + + // CHECK-SAME{LITERAL}: manual_axes={"a", "b"} + // CHECK-SAME: (%arg2: tensor<2x8xf32>, %arg3: tensor<8x32xf32>) { + // CHECK-NEXT: %[[ADD_0:.*]] = stablehlo.add %arg2, %arg2 : tensor<2x8xf32> + // CHECK-NEXT: %[[DOT:.*]] = stablehlo.dot %[[ADD_0]], %arg3 : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> + // CHECK-NEXT: %[[REDUCE:.*]] = "stablehlo.all_reduce"(%[[DOT]]) + // CHECK-NEXT: ^bb0(%arg4: tensor, %arg5: tensor): + // CHECK-NEXT: %[[ADD_1:.*]] = stablehlo.add %arg4, %arg5 : tensor + // CHECK-NEXT: stablehlo.return %[[ADD_1]] : tensor + // CHECK-NEXT: }) : (tensor<2x32xf32>) -> tensor<2x32xf32> + // CHECK-NEXT: sdy.return %[[REDUCE]] : tensor<2x32xf32> + // CHECK-NEXT: } : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32> + // CHECK-NEXT: return %[[MAN_COMP]] : tensor<8x32xf32> + %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {"b"}]>]>} : (tensor<8x16xf32>) -> tensor<8x16xf32> + %s1 = stablehlo.custom_call @Sharding(%arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}, {}], replicated={"a"}>]>} : (tensor<16x32xf32>) -> tensor<16x32xf32> + %0:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0, %s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>, <@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x16xf32>, tensor<16x32xf32>) -> (tensor<2x8xf32>, tensor<8x32xf32>) + %1 = call @xla.sdy.manual_computation_body(%0#0, %0#1) : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> + %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}], replicated={"b"}>]>} : (tensor<2x32xf32>) -> tensor<8x32xf32> + return %2 : tensor<8x32xf32> +} + +// CHECK-LABEL: func @single_manual_comp_name_is_not_prefix_nor_suffix +func.func @single_manual_comp_name_is_not_prefix_nor_suffix(%arg0: tensor<8x8xf32>) -> (tensor<8x8xf32>) { + // CHECK-NOT: call @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234 + // CHECK: %[[MAN_COMP:.*]] = sdy.manual_computation(%arg0) + // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{"a"}, {}]>] + // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{"a"}, {}]>] + // CHECK-SAME{LITERAL}: manual_axes={"a"} + // CHECK-SAME: (%arg1: tensor<2x8xf32>) { + // CHECK-NEXT: sdy.return %arg1 : tensor<2x8xf32> + // CHECK-NEXT: } : (tensor<8x8xf32>) -> tensor<8x8xf32> + // CHECK-NEXT: return %[[MAN_COMP]] : tensor<8x8xf32> + %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32> + %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x8xf32>) -> tensor<2x8xf32> + %1 = call @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32> + %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<8x8xf32> + return %2 : tensor<8x8xf32> +} + +// Tests chained manual computations: outputs of first shard map feed second shard map. +// CHECK-LABEL: func @manual_comp_using_another +func.func @manual_comp_using_another(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> { + // CHECK-NOT: call @xla.sdy.manual_computation_body_0 + // CHECK: %[[MAN_COMP_0:.*]] = sdy.manual_computation(%arg0) + // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{"a"}, {}]>] + // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{"a"}, {}]>] + // CHECK-SAME{LITERAL}: manual_axes={"a"} + // CHECK-SAME: (%arg1: tensor<2x8xf32>) { + // CHECK-NEXT: sdy.return %arg1 : tensor<2x8xf32> + // CHECK-NEXT: } : (tensor<8x8xf32>) -> tensor<8x8xf32> + // CHECK-NOT: call @xla.sdy.manual_computation_body_1 + // CHECK-NEXT: %[[MAN_COMP_1:.*]] = sdy.manual_computation(%[[MAN_COMP_0]]) + // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{}, {"b"}]>] + // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{}, {"b"}]>] + // CHECK-SAME{LITERAL}: manual_axes={"b"} + // CHECK-SAME: (%arg1: tensor<8x4xf32>) { + // CHECK-NEXT: sdy.return %arg1 : tensor<8x4xf32> + // CHECK-NEXT: } : (tensor<8x8xf32>) -> tensor<8x8xf32> + // CHECK-NEXT: return %[[MAN_COMP_1]] : tensor<8x8xf32> + %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32> + %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x8xf32>) -> tensor<2x8xf32> + %1 = call @xla.sdy.manual_computation_body_0(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32> + %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<8x8xf32> + %s1 = stablehlo.custom_call @Sharding(%2) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {"b"}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32> + %3 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<8x8xf32>) -> tensor<8x4xf32> + %4 = call @xla.sdy.manual_computation_body_1(%3) : (tensor<8x4xf32>) -> tensor<8x4xf32> + %5 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%4) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {"b"}]>]>} : (tensor<8x4xf32>) -> tensor<8x8xf32> + return %5 : tensor<8x8xf32> +} + +// Tests nested manual computations: reconstructing inner and outer sdy.manual_computation regions. +// CHECK-LABEL: func @nested_shmaps +func.func @nested_shmaps(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> { + // CHECK-NOT: call @xla.sdy.manual_computation_body_2 + // CHECK-NOT: call @xla.sdy.manual_computation_body_3 + // CHECK: %[[MAN_COMP_0:.*]] = sdy.manual_computation(%arg0) + // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_1, [{"a"}, {}]>] + // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_1, [{"a"}, {}]>] + // CHECK-SAME{LITERAL}: manual_axes={"a"} + // CHECK-SAME: (%arg1: tensor<2x8xf32>) { + // CHECK-NEXT: %[[MAN_COMP_1:.*]] = sdy.manual_computation(%arg1) + // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_1, [{}, {"b"}]>] + // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_1, [{}, {"b"}]>] + // CHECK-SAME{LITERAL}: manual_axes={"b"} + // CHECK-SAME: (%arg2: tensor<2x4xf32>) { + // CHECK-NEXT: %[[MUL:.*]] = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32> + // CHECK-NEXT: sdy.return %[[MUL]] : tensor<2x4xf32> + // CHECK-NEXT: } : (tensor<2x8xf32>) -> tensor<2x8xf32> + // CHECK-NEXT: sdy.return %[[MAN_COMP_1]] : tensor<2x8xf32> + // CHECK-NEXT: } : (tensor<4x8xf32>) -> tensor<4x8xf32> + // CHECK-NEXT: return %[[MAN_COMP_0]] : tensor<4x8xf32> + %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<4x8xf32>) -> tensor<4x8xf32> + %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<4x8xf32>) -> tensor<2x8xf32> + %1 = call @xla.sdy.manual_computation_body_3(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32> + %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<4x8xf32> + return %2 : tensor<4x8xf32> +} + +// Tests nested manual computation with operations in the outer region around the inner region. +// CHECK-LABEL: func @nested_shmaps_extra_op +func.func @nested_shmaps_extra_op(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> { + // CHECK-NOT: call @xla.sdy.manual_computation_body_4 + // CHECK-NOT: call @xla.sdy.manual_computation_body_5 + // CHECK: %[[MAN_COMP_0:.*]] = sdy.manual_computation(%arg0) + // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_1, [{"a"}, {}]>] + // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_1, [{"a"}, {}]>] + // CHECK-SAME{LITERAL}: manual_axes={"a"} + // CHECK-SAME: (%arg1: tensor<2x8xf32>) { + // CHECK-NEXT: %[[MAN_COMP_1:.*]] = sdy.manual_computation(%arg1) + // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_1, [{}, {"b"}]>] + // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_1, [{}, {"b"}]>] + // CHECK-SAME{LITERAL}: manual_axes={"b"} + // CHECK-SAME: (%arg2: tensor<2x4xf32>) { + // CHECK-NEXT: %[[MUL:.*]] = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32> + // CHECK-NEXT: sdy.return %[[MUL]] : tensor<2x4xf32> + // CHECK-NEXT: } : (tensor<2x8xf32>) -> tensor<2x8xf32> + // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[MAN_COMP_1]], %[[MAN_COMP_1]] : tensor<2x8xf32> + // CHECK-NEXT: sdy.return %[[ADD]] : tensor<2x8xf32> + // CHECK-NEXT: } : (tensor<4x8xf32>) -> tensor<4x8xf32> + // CHECK-NEXT: return %[[MAN_COMP_0]] : tensor<4x8xf32> + %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<4x8xf32>) -> tensor<4x8xf32> + %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<4x8xf32>) -> tensor<2x8xf32> + %1 = call @xla.sdy.manual_computation_body_5(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32> + %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<4x8xf32> + return %2 : tensor<4x8xf32> +} + +// Tests manual computation with no inputs: there is no `GlobalToLocalShape`, so +// the body call op is the only carrier of the manual axes. +// CHECK-LABEL: func @manual_computation_no_inputs +func.func @manual_computation_no_inputs() -> tensor<4xi64> { + + // CHECK-NOT: call @xla.sdy.manual_computation_body_6 + // CHECK: %[[SHMAP:.*]] = sdy.manual_computation() + // CHECK-SAME{LITERAL}: in_shardings=[] + // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{"b"}]>] + // CHECK-SAME{LITERAL}: manual_axes={"b"} + // CHECK-SAME{LITERAL}: () { + // CHECK-NEXT: %[[C:.*]] = stablehlo.constant dense<[2, 3]> : tensor<2xi64> + // CHECK-NEXT: sdy.return %[[C]] : tensor<2xi64> + // CHECK-NEXT: } : () -> tensor<4xi64> + // CHECK-NEXT: return %[[SHMAP]] : tensor<4xi64> + %0 = call @xla.sdy.manual_computation_body_6() {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}]>]>, xla.sdy.manual_axes = #sdy} : () -> tensor<2xi64> + %1 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}]>]>} : (tensor<2xi64>) -> tensor<4xi64> + return %1 : tensor<4xi64> +} + +// Edge case: Manual computation consuming an input but producing no outputs. +// CHECK-LABEL: func @manual_computation_no_outputs +func.func @manual_computation_no_outputs(%arg0: tensor<4xi64>) { + // CHECK-NOT: call @xla.sdy.manual_computation_body_7 + // CHECK: sdy.manual_computation(%arg0) + // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{"b"}]>] + // CHECK-SAME{LITERAL}: out_shardings=[] + // CHECK-SAME{LITERAL}: manual_axes={"b"} + // CHECK-SAME{LITERAL}: (%arg1: tensor<2xi64>) { + // CHECK-NEXT: stablehlo.custom_call @sdy_testonly(%arg1) : (tensor<2xi64>) -> () + // CHECK-NEXT: sdy.return + // CHECK-NEXT: } : (tensor<4xi64>) -> () + // CHECK-NEXT: return + %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}]>]>} : (tensor<4xi64>) -> tensor<4xi64> + %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<4xi64>) -> tensor<2xi64> + call @xla.sdy.manual_computation_body_7(%0) : (tensor<2xi64>) -> () + return +} + +// Edge case: Degenerate manual computation with no inputs and no outputs. +// CHECK-LABEL: func @manual_computation_no_inputs_no_outputs +func.func @manual_computation_no_inputs_no_outputs() { + // CHECK-NEXT: sdy.manual_computation() in_shardings=[] out_shardings=[] manual_axes={} () { + // CHECK-NEXT: sdy.return + // CHECK-NEXT: } : () -> () + // CHECK-NEXT: return + call @xla.sdy.manual_computation_body_8() {mhlo.frontend_attributes = {inlineable = "false"}} : () -> () + return +} + +// Edge case: Manual computation with 0-sized dimension tensors. +// CHECK-LABEL: func @manual_computation_some_zero_dim_inputs_outputs +func.func @manual_computation_some_zero_dim_inputs_outputs(%arg0: tensor<0x16xf32>, %arg1: tensor<16x32xf32>) -> (tensor<0x32xf32>, tensor<16x32xf32>) { + + // CHECK-NOT: call @xla.sdy.manual_computation_body + // CHECK: %[[CONST_0_32:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<0x32xf32> + // CHECK: %[[MAN_COMP:.*]]:2 = sdy.manual_computation(%arg0, %arg1) + // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, [{}, {"b"}]>, <@mesh_0, [{"b"}, {}]>] + // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, [{}, {}], replicated={"b"}>, <@mesh_0, [{"b"}, {}]>] + // CHECK-SAME{LITERAL}: manual_axes={"b"} + // CHECK-SAME: (%arg2: tensor<0x8xf32>, %arg3: tensor<8x32xf32>) { + // CHECK-NEXT: %[[DOT:.*]] = stablehlo.dot %arg2, %arg3 + // CHECK-NEXT: sdy.return %[[DOT]], %arg3 + // CHECK-NEXT: } : (tensor<0x16xf32>, tensor<16x32xf32>) -> (tensor<0x32xf32>, tensor<16x32xf32>) + // CHECK-NEXT: return %[[CONST_0_32]], %[[MAN_COMP]]#1 + %c1 = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32> + %c2 = stablehlo.constant dense<0.000000e+00> : tensor<0x32xf32> + %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {"b"}]>]>} : (tensor<0x16xf32>) -> tensor<0x16xf32> + %s1 = stablehlo.custom_call @Sharding(%arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}, {}]>]>} : (tensor<16x32xf32>) -> tensor<16x32xf32> + %0:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0, %s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>, <@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<0x16xf32>, tensor<16x32xf32>) -> (tensor<0x8xf32>, tensor<8x32xf32>) + %1:2 = call @xla.sdy.manual_computation_body_9(%c1, %0#1) : (tensor<0x8xf32>, tensor<8x32xf32>) -> (tensor<0x32xf32>, tensor<8x32xf32>) + %2:2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%c2, %1#1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}], replicated={"b"}>, <@mesh_0, [{"b"}, {}]>]>} : (tensor<0x32xf32>, tensor<8x32xf32>) -> (tensor<0x32xf32>, tensor<16x32xf32>) + return %c2, %2#1 : tensor<0x32xf32>, tensor<16x32xf32> +} + +// Type preservation: Token arguments are preserved with empty sharding and not assigned manual axes. +// CHECK-LABEL: func @manual_computation_tokens +func.func @manual_computation_tokens(%arg0: !stablehlo.token, %arg1: tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) { + // CHECK-NOT: call @xla.sdy.manual_computation_body_11 + // CHECK: %[[MAN_COMP:.*]]:2 = sdy.manual_computation(%arg0, %arg1) + // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>] + // CHECK-SAME{LITERAL}: out_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>] + // CHECK-SAME{LITERAL}: manual_axes={"b"} + // CHECK-SAME: (%arg2: !stablehlo.token, %arg3: tensor<1xi64>) { + // CHECK-NEXT: sdy.return %arg2, %arg3 : !stablehlo.token, tensor<1xi64> + // CHECK-NEXT: } : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) + // CHECK-NEXT: return %[[MAN_COMP]]#0, %[[MAN_COMP]]#1 : !stablehlo.token, tensor<2xi64> + %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, []>]>} : (!stablehlo.token) -> !stablehlo.token + %s1 = stablehlo.custom_call @Sharding(%arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}]>]>} : (tensor<2xi64>) -> tensor<2xi64> + %0:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0, %s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, []>, <@mesh_0, [{}]>]>, xla.sdy.manual_axes = #sdy} : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<1xi64>) + %1:2 = call @xla.sdy.manual_computation_body_11(%0#0, %0#1) : (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<1xi64>) + %2:2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1#0, %1#1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, []>, <@mesh_0, [{"b"}]>]>} : (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<2xi64>) + return %2#0, %2#1 : !stablehlo.token, tensor<2xi64> +} + +// Tests stray/dangling boundary custom calls where results are unused: safely stripped. +// CHECK-LABEL: func @stray_unused_manual_computation_custom_calls +func.func @stray_unused_manual_computation_custom_calls(%arg0: tensor<0x16xf32>) -> tensor<0x16xf32> { + // CHECK-NEXT: %[[CONST_0_8:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32> + // CHECK-NEXT: %[[CONST_0_16:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<0x16xf32> + // CHECK-NOT: call @xla.sdy.manual_computation_body + // CHECK-NOT: call @xla.sdy.GlobalToLocalShape + // CHECK-NOT: call @xla.sdy.LocalToGlobalShape + // CHECK-NEXT: return %[[CONST_0_16]] + %c1 = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32> + %c2 = stablehlo.constant dense<0.000000e+00> : tensor<0x16xf32> + %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%arg0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<0x16xf32>) -> tensor<0x8xf32> + %1 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%c1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}, {}]>]>} : (tensor<0x8xf32>) -> tensor<0x16xf32> + return %c2 : tensor<0x16xf32> +} + +// CHECK-NOT: func private @xla.sdy.manual_computation_body( +func.func private @xla.sdy.manual_computation_body(%arg0: tensor<2x8xf32>, %arg1: tensor<8x32xf32>) -> tensor<2x32xf32> { + %0 = stablehlo.add %arg0, %arg0 : tensor<2x8xf32> + %1 = stablehlo.dot %0, %arg1 : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32> + %2 = "stablehlo.all_reduce"(%1) <{replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>}> ({ + ^bb0(%arg2: tensor, %arg3: tensor): + %3 = stablehlo.add %arg2, %arg3 : tensor + stablehlo.return %3 : tensor + }) : (tensor<2x32xf32>) -> tensor<2x32xf32> + return %2 : tensor<2x32xf32> +} + +// CHECK-NOT: func private @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234( +func.func private @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> { + return %arg0 : tensor<2x8xf32> +} + +// CHECK-NOT: func private @xla.sdy.manual_computation_body_0( +func.func private @xla.sdy.manual_computation_body_0(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> { + return %arg0 : tensor<2x8xf32> +} + +// CHECK-NOT: func private @xla.sdy.manual_computation_body_1( +func.func private @xla.sdy.manual_computation_body_1(%arg0: tensor<8x4xf32>) -> tensor<8x4xf32> { + return %arg0 : tensor<8x4xf32> +} + +// CHECK-NOT: func private @xla.sdy.manual_computation_body_2( +func.func private @xla.sdy.manual_computation_body_2(%arg0: tensor<2x4xf32>) -> tensor<2x4xf32> { + %0 = stablehlo.multiply %arg0, %arg0 : tensor<2x4xf32> + return %0 : tensor<2x4xf32> +} + +// CHECK-NOT: func private @xla.sdy.manual_computation_body_3( +func.func private @xla.sdy.manual_computation_body_3(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> { + %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x8xf32>) -> tensor<2x8xf32> + %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<2x8xf32>) -> tensor<2x4xf32> + %1 = call @xla.sdy.manual_computation_body_2(%0) : (tensor<2x4xf32>) -> tensor<2x4xf32> + %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x4xf32>) -> tensor<2x8xf32> + return %2 : tensor<2x8xf32> +} + +// CHECK-NOT: func private @xla.sdy.manual_computation_body_4( +func.func private @xla.sdy.manual_computation_body_4(%arg0: tensor<2x4xf32>) -> tensor<2x4xf32> { + %0 = stablehlo.multiply %arg0, %arg0 : tensor<2x4xf32> + return %0 : tensor<2x4xf32> +} + +// CHECK-NOT: func private @xla.sdy.manual_computation_body_5( +func.func private @xla.sdy.manual_computation_body_5(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> { + %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x8xf32>) -> tensor<2x8xf32> + %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy} : (tensor<2x8xf32>) -> tensor<2x4xf32> + %1 = call @xla.sdy.manual_computation_body_4(%0) : (tensor<2x4xf32>) -> tensor<2x4xf32> + %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x4xf32>) -> tensor<2x8xf32> + %3 = stablehlo.add %2, %2 : tensor<2x8xf32> + return %3 : tensor<2x8xf32> +} + +// CHECK-NOT: func private @xla.sdy.manual_computation_body_6( +func.func private @xla.sdy.manual_computation_body_6() -> tensor<2xi64> { + %0 = stablehlo.constant dense<[2, 3]> : tensor<2xi64> + return %0 : tensor<2xi64> +} + +// CHECK-NOT: func private @xla.sdy.manual_computation_body_7( +func.func private @xla.sdy.manual_computation_body_7(%arg0: tensor<2xi64>) { + stablehlo.custom_call @sdy_testonly(%arg0) : (tensor<2xi64>) -> () + return +} + +// CHECK-NOT: func private @xla.sdy.manual_computation_body_8( +func.func private @xla.sdy.manual_computation_body_8() { + return +} + +// CHECK-NOT: func private @xla.sdy.manual_computation_body_9( +func.func private @xla.sdy.manual_computation_body_9(%arg0: tensor<0x8xf32>, %arg1: tensor<8x32xf32>) -> (tensor<0x32xf32>, tensor<8x32xf32>) { + %0 = stablehlo.dot %arg0, %arg1 : (tensor<0x8xf32>, tensor<8x32xf32>) -> tensor<0x32xf32> + return %0, %arg1 : tensor<0x32xf32>, tensor<8x32xf32> +} + +// CHECK-NOT: func private @xla.sdy.manual_computation_body_11( +func.func private @xla.sdy.manual_computation_body_11(%arg0: !stablehlo.token, %arg1: tensor<1xi64>) -> (!stablehlo.token, tensor<1xi64>) { + return %arg0, %arg1 : !stablehlo.token, tensor<1xi64> +} diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc index 4974a9b69ae604..30b8dfe86497e6 100644 --- a/third_party/xla/xla/service/spmd/spmd_partitioner.cc +++ b/third_party/xla/xla/service/spmd/spmd_partitioner.cc @@ -7496,41 +7496,20 @@ absl::Status SpmdPartitioner::ConvertUnreducedSharding( return res; }; auto convert_unreduced_subgroup_sharding = - [](HloInstruction* hlo, - const HloSharding& sharding) -> absl::StatusOr { - // TODO(b/438306205): Remove this check once the unreduced - // subgroup sharding is compatible with manual. - TF_RET_CHECK(!sharding.IsManualSubgroup()) - << "Incompatible unreduced sharding at " << hlo->ToString(); + [](HloInstruction* hlo, const HloSharding& sharding) -> HloSharding { hlo->add_frontend_attribute(sdy::kHasUnreducedAxes, "true"); - TileAssignment tile_assignment = sharding.tile_assignment(); - if (sharding.HasPartialReplication()) { - // When we have both replicated and unreduced, merge them into one - // in the tile assignment. - int64_t unreduced_dim = sharding.SubgroupUnreducedDim(); - DimensionVector new_dims(tile_assignment.dimensions().begin(), - tile_assignment.dimensions().end()); - new_dims[sharding.SubgroupReplicationDim()] *= - new_dims[unreduced_dim]; - new_dims.erase(new_dims.begin() + unreduced_dim); - tile_assignment = tile_assignment.Reshape(new_dims); - } - HloSharding res = - HloSharding::PartialTile(tile_assignment, sharding.metadata()); + std::vector subgroup_types( + sharding.subgroup_types().begin(), sharding.subgroup_types().end()); + absl::c_replace(subgroup_types, OpSharding::UNREDUCED, + OpSharding::REPLICATED); + HloSharding res = HloSharding::Subgroup( + sharding.tile_assignment(), subgroup_types, sharding.metadata()); res.set_reduction_op(sharding.reduction_op()); return res; }; auto convert_unreduced_named_sharding = - [](HloInstruction* hlo, - const HloSharding& sharding) -> absl::StatusOr { + [](HloInstruction* hlo, const HloSharding& sharding) -> HloSharding { const NamedSharding& named_sharding = sharding.named_sharding(); - // TODO(b/438306205): Remove this check once the unreduced named - // sharding is compatible with manual. - if (!named_sharding.manual_axes().empty()) { - return absl::UnimplementedError( - "NamedSharding with both unreduced and manual axes is not " - "supported."); - } hlo->add_frontend_attribute(sdy::kHasUnreducedAxes, "true"); std::vector new_replicated_axes( named_sharding.replicated_axes().begin(), @@ -7552,11 +7531,10 @@ absl::Status SpmdPartitioner::ConvertUnreducedSharding( for (HloSharding& subsharding : subshardings) { if (subsharding.IsUnreducedSubgroup()) { if (subsharding.UseNamedShardingLeaf()) { - ABSL_ASSIGN_OR_RETURN(subsharding, convert_unreduced_named_sharding( - hlo, subsharding)); + subsharding = convert_unreduced_named_sharding(hlo, subsharding); } else { - ABSL_ASSIGN_OR_RETURN(subsharding, convert_unreduced_subgroup_sharding( - hlo, subsharding)); + subsharding = + convert_unreduced_subgroup_sharding(hlo, subsharding); } should_convert = true; } else if (subsharding.IsUnreduced()) { @@ -7570,14 +7548,10 @@ absl::Status SpmdPartitioner::ConvertUnreducedSharding( } else { if (sharding.IsUnreducedSubgroup()) { if (sharding.UseNamedShardingLeaf()) { - ABSL_ASSIGN_OR_RETURN(HloSharding new_sharding, - convert_unreduced_named_sharding(hlo, sharding)); - hlo->set_sharding(new_sharding); + hlo->set_sharding(convert_unreduced_named_sharding(hlo, sharding)); } else { - ABSL_ASSIGN_OR_RETURN( - HloSharding new_sharding, + hlo->set_sharding( convert_unreduced_subgroup_sharding(hlo, sharding)); - hlo->set_sharding(new_sharding); } } else if (sharding.IsUnreduced()) { hlo->set_sharding(convert_unreduced_sharding(hlo)); diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc index c9deeba34249bc..bbb367d23f79a2 100644 --- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc +++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc @@ -18542,6 +18542,26 @@ ENTRY entry { } } +TEST_P(SpmdPartitioningTest, SubgroupUnreducedAndManual) { + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + a = f32[8,1024]{1,0} parameter(0), sharding={devices=[1,2,2]<=[2,2]T(1,0) last_tile_dims={manual}} + b = f32[1024,256]{1,0} parameter(1), sharding={devices=[2,1,2]<=[2,2]T(1,0) last_tile_dims={manual}} + ROOT dot = f32[8,256]{1,0} dot(a, b), lhs_contracting_dims={1}, rhs_contracting_dims={0}, sharding={devices=[1,1,2,2]<=[4] last_tile_dims={manual,unreduced}} +})"; + SpmdPartitionerOptions options; + for (bool need_resolve_conflicts : {true, false}) { + options.need_resolve_conflicts = need_resolve_conflicts; + ASSERT_OK_AND_ASSIGN( + auto module, + PartitionComputation(hlo_string, /*num_devices=*/4, options)); + EXPECT_THAT(module->entry_computation()->root_instruction(), op::Dot()); + EXPECT_EQ(FindInstruction(module.get(), HloOpcode::kAllReduce), nullptr); + } +} + TEST_P(SpmdPartitioningTest, UnreducedReduce) { absl::string_view hlo_string = R"( HloModule module @@ -19245,6 +19265,26 @@ ENTRY entry { "{mesh['x'=2] [], unreduced={'x'}}}"))); } +TEST_F(SpmdPartitioningV3Test, SubgroupUnreducedAndManualV3) { + absl::string_view hlo_string = R"( +HloModule module + +ENTRY entry { + a = f32[8,1024]{1,0} parameter(0), sharding={mesh['x'=2,'y'=2] [{},{'y'}], manual={'x'}} + b = f32[1024,256]{1,0} parameter(1), sharding={mesh['x'=2,'y'=2] [{'y'},{}], manual={'x'}} + ROOT dot = f32[8,256]{1,0} dot(a, b), lhs_contracting_dims={1}, rhs_contracting_dims={0}, sharding={mesh['x'=2,'y'=2] [{},{}], unreduced={'y'}, manual={'x'}} +})"; + SpmdPartitionerOptions options; + for (bool need_resolve_conflicts : {true, false}) { + options.need_resolve_conflicts = need_resolve_conflicts; + ASSERT_OK_AND_ASSIGN( + auto module, + PartitionComputation(hlo_string, /*num_devices=*/4, options)); + EXPECT_THAT(module->entry_computation()->root_instruction(), op::Dot()); + EXPECT_EQ(FindInstruction(module.get(), HloOpcode::kAllReduce), nullptr); + } +} + TEST_F(SpmdPartitioningV3Test, PatternMatchMergeNamedSharding) { absl::string_view hlo_string = R"( HloModule module diff --git a/third_party/xla/xla/stream_executor/cuda/BUILD b/third_party/xla/xla/stream_executor/cuda/BUILD index 363a2b440b0e2b..d23960e276d69b 100644 --- a/third_party/xla/xla/stream_executor/cuda/BUILD +++ b/third_party/xla/xla/stream_executor/cuda/BUILD @@ -942,6 +942,7 @@ xla_cc_test( tags = [ "cuda-only", "gpu", + "nomsan", "requires-gpu-sm90", ], deps = [ @@ -2407,6 +2408,7 @@ xla_cc_test( tags = [ "cuda-only", "gpu", + "nomsan", ], deps = [ ":compilation_provider_test_lib", @@ -2420,6 +2422,7 @@ xla_cc_test( tags = [ "cuda-only", "gpu", + "nomsan", "notsan", "requires-gpu-sm90", ], @@ -3129,6 +3132,7 @@ cuda_library( "//xla/stream_executor/gpu:gpu_kernel_registry", "//xla/stream_executor/gpu:topk_kernel", "//xla/tsl/lib/math:math_util", + "@com_google_absl//absl/base", ], alwayslink = 1, ) diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc index 68f73ca0aabad7..b45a2271673e4b 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc +++ b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc @@ -5871,7 +5871,7 @@ bool CudnnSupport::GetConvolveBackwardDataAlgorithms( if (CudnnEnvVar::IsEnabled()) { algo_types.push_back(CUDNN_CONVOLUTION_BWD_DATA_ALGO_WINOGRAD_NONFUSED); } - if (engine_options.require_determinism) { + if (!engine_options.require_determinism) { algo_types.push_back(CUDNN_CONVOLUTION_BWD_DATA_ALGO_0); } diff --git a/third_party/xla/xla/stream_executor/cuda/topk_kernel_cuda_common.cu.h b/third_party/xla/xla/stream_executor/cuda/topk_kernel_cuda_common.cu.h index f4a38df8f1cd09..046e65812860d8 100644 --- a/third_party/xla/xla/stream_executor/cuda/topk_kernel_cuda_common.cu.h +++ b/third_party/xla/xla/stream_executor/cuda/topk_kernel_cuda_common.cu.h @@ -23,7 +23,9 @@ limitations under the License. #include #include +#include +#include "absl/base/casts.h" #include "xla/stream_executor/cuda/cuda_platform_id.h" #include "xla/stream_executor/gpu/gpu_kernel_registry.h" #include "xla/stream_executor/gpu/topk_kernel.h" @@ -61,6 +63,26 @@ __device__ __forceinline__ NT GpuShuffle(NT val, uint32_t idx, return res.v; } +// Converts IEEE 754 floating-point keys (32-bit float and 16-bit bfloat16) +// to order-preserving unsigned integers. This establishes a well-defined total +// ordering, properly handling special values such as NaNs and signed zeroes +// during integer sorting. +namespace details { +template +__device__ __forceinline__ auto ToOrdered(T x) { + if constexpr (sizeof(T) == 4 && !std::is_integral_v) { + uint32_t val = absl::bit_cast(x); + return (val & 0x80000000u) ? ~val : (val | 0x80000000u); + } else if constexpr (sizeof(T) == 2 && !std::is_integral_v) { + uint16_t val = absl::bit_cast(x); + return (val & 0x8000u) ? static_cast(~val) + : static_cast(val | 0x8000u); + } else { + return x; + } +} +} // namespace details + // Default implementation for KV holder. Useful for testing while adding support // for a new type, but generally bitpacking those values is more efficient. See // implementations below. @@ -72,7 +94,9 @@ struct Descending { }; __device__ __forceinline__ static bool cmp(const KVT& lhs, const KVT& rhs) { - return lhs.key == rhs.key ? lhs.idx < rhs.idx : lhs.key > rhs.key; + auto l = details::ToOrdered(lhs.key); + auto r = details::ToOrdered(rhs.key); + return l == r ? lhs.idx < rhs.idx : l > r; } }; diff --git a/third_party/xla/xla/stream_executor/rocm/BUILD b/third_party/xla/xla/stream_executor/rocm/BUILD index 9a87e4ba0e043b..f95ca29577e1c5 100644 --- a/third_party/xla/xla/stream_executor/rocm/BUILD +++ b/third_party/xla/xla/stream_executor/rocm/BUILD @@ -1481,6 +1481,7 @@ rocm_library( "//xla/stream_executor/gpu:gpu_kernel_registry", "//xla/stream_executor/gpu:topk_kernel", "//xla/tsl/lib/math:math_util", + "@com_google_absl//absl/base", ], alwayslink = 1, ) diff --git a/third_party/xla/xla/stream_executor/rocm/topk_kernel_rocm_common.cu.h b/third_party/xla/xla/stream_executor/rocm/topk_kernel_rocm_common.cu.h index 9e2fc316a2fce4..df2b8db853fa68 100644 --- a/third_party/xla/xla/stream_executor/rocm/topk_kernel_rocm_common.cu.h +++ b/third_party/xla/xla/stream_executor/rocm/topk_kernel_rocm_common.cu.h @@ -22,7 +22,9 @@ limitations under the License. #include #include +#include +#include "absl/base/casts.h" #include "xla/stream_executor/gpu/gpu_kernel_registry.h" #include "xla/stream_executor/gpu/topk_kernel.h" #include "xla/stream_executor/kernel_symbol_registry.h" @@ -65,6 +67,26 @@ __device__ __forceinline__ NT GpuShuffle(NT val, uint32_t idx, return res.v; } +// Converts IEEE 754 floating-point keys (e.g., 32-bit float and 16-bit +// bfloat16) to order-preserving unsigned integers. This establishes a +// well-defined total ordering, properly handling special values such as NaNs +// and signed zeroes during integer sorting. +namespace details { +template +__device__ __forceinline__ auto ToOrdered(T x) { + if constexpr (sizeof(T) == 4 && !std::is_integral_v) { + uint32_t val = absl::bit_cast(x); + return (val & 0x80000000u) ? ~val : (val | 0x80000000u); + } else if constexpr (sizeof(T) == 2 && !std::is_integral_v) { + uint16_t val = absl::bit_cast(x); + return (val & 0x8000u) ? static_cast(~val) + : static_cast(val | 0x8000u); + } else { + return x; + } +} +} // namespace details + // Default implementation for KV holder. Useful for testing while adding support // for a new type, but generally bitpacking those values is more efficient. See // implementations below. @@ -76,7 +98,9 @@ struct Descending { }; __device__ __forceinline__ static bool cmp(const KVT& lhs, const KVT& rhs) { - return lhs.key == rhs.key ? lhs.idx < rhs.idx : lhs.key > rhs.key; + auto l = details::ToOrdered(lhs.key); + auto r = details::ToOrdered(rhs.key); + return l == r ? lhs.idx < rhs.idx : l > r; } }; diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD index eeebfd10a03b90..184fac20a87fcc 100644 --- a/third_party/xla/xla/tests/BUILD +++ b/third_party/xla/xla/tests/BUILD @@ -2181,6 +2181,37 @@ xla_test( ], ) +xla_test( + name = "topk_special_floats_test", + srcs = ["topk_special_floats_test.cc"], + backend_tags = { + "gpu": [ + # Skip P100/V100 GPUs due to 0.0 and -0.0 interleaving bug in sort. + "requires-gpu-sm80", + ], + }, + backends = [ + "cpu", + "gpu", + ], + deps = [ + ":hlo_pjrt_interpreter_reference_mixin", + ":hlo_test_base", + ":literal_test_util", + ":xla_internal_test_main", + "//xla:array2d", + "//xla:error_spec", + "//xla:literal", + "//xla:literal_util", + "//xla:types", + "//xla/tsl/platform:test", + "@com_google_absl//absl/base", + "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/strings:string_view", + "@com_google_googletest//:gtest", + ], +) + xla_test( name = "runtime_topk_test", srcs = ["runtime_topk_test.cc"], diff --git a/third_party/xla/xla/tests/topk_special_floats_test.cc b/third_party/xla/xla/tests/topk_special_floats_test.cc new file mode 100644 index 00000000000000..e0fda470d77848 --- /dev/null +++ b/third_party/xla/xla/tests/topk_special_floats_test.cc @@ -0,0 +1,255 @@ +/* Copyright 2023 The OpenXLA Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include + +#include +#include "absl/base/casts.h" +#include "absl/strings/str_format.h" +#include "absl/strings/string_view.h" +#include "xla/array2d.h" +#include "xla/error_spec.h" +#include "xla/literal.h" +#include "xla/literal_util.h" +#include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h" +#include "xla/tests/hlo_test_base.h" +#include "xla/tests/literal_test_util.h" +#include "xla/tsl/platform/test.h" +#include "xla/types.h" + +namespace xla { +namespace { + +using TopkTest = HloPjRtInterpreterReferenceMixin; + +TEST_F(TopkTest, TopKWithSpecialFloats) { + // Regression test for TopK TotalOrder with special float values + // (NaN, +-Inf, +-0). + float neg_qnan0 = absl::bit_cast(0xFFC00000u); + float neg_qnan1 = absl::bit_cast(0xFFE00000u); + float neg_qnan2 = absl::bit_cast(0xFFE00001u); + float neg_smallest_qnan = absl::bit_cast(0xFFFFFFFFu); + float pos_qnan0 = absl::bit_cast(0x7FC00000u); + float pos_qnan1 = absl::bit_cast(0x7FE00000u); + float pos_qnan2 = absl::bit_cast(0x7FE00001u); + + float pos_inf = absl::bit_cast(0x7F800000u); + float neg_inf = absl::bit_cast(0xFF800000u); + float pos_zero = absl::bit_cast(0x00000000u); + float neg_zero = absl::bit_cast(0x80000000u); + + float pos_dnrm_1e38 = absl::bit_cast(0x00700000u); // 1.028e-38 + float pos_dnrm_1e40 = absl::bit_cast(0x00020000u); // 1.836e-40 + float neg_dnrm_1e39 = absl::bit_cast(0x80100000u); // -1.469e-39 + float neg_dnrm_9e41 = absl::bit_cast(0x80010000u); // -9.183e-41 + + float pos_64_5 = absl::bit_cast(0x42810000u); + float pos_0_0125 = absl::bit_cast(0x3C4d0000u); + float neg_64_5 = absl::bit_cast(0xC2810000u); + + std::vector input_row_start = { + neg_inf, pos_inf, neg_dnrm_1e39, pos_zero, pos_qnan0, + pos_qnan2, neg_zero, neg_64_5, pos_inf, neg_qnan0, + neg_qnan2, neg_64_5, neg_zero, pos_dnrm_1e40, neg_dnrm_9e41, + neg_qnan1, pos_zero, neg_inf, pos_64_5, pos_qnan1, + pos_0_0125, pos_dnrm_1e38, neg_dnrm_1e39, pos_dnrm_1e40, neg_dnrm_9e41}; + + std::vector expected_values_64 = { + pos_qnan2, pos_qnan1, pos_qnan0, + pos_inf, pos_inf, pos_64_5, + pos_0_0125, pos_dnrm_1e38, pos_dnrm_1e40, + pos_dnrm_1e40, pos_zero, pos_zero, + neg_zero, neg_zero, neg_dnrm_9e41, + neg_dnrm_9e41, neg_dnrm_1e39, neg_dnrm_1e39, + neg_64_5, neg_64_5, neg_inf, + neg_inf, neg_qnan0, neg_qnan1, + neg_qnan2, neg_smallest_qnan, neg_smallest_qnan, + neg_smallest_qnan, neg_smallest_qnan, neg_smallest_qnan, + neg_smallest_qnan, neg_smallest_qnan, neg_smallest_qnan, + neg_smallest_qnan, neg_smallest_qnan, neg_smallest_qnan, + neg_smallest_qnan, neg_smallest_qnan, neg_smallest_qnan, + neg_smallest_qnan, neg_smallest_qnan, neg_smallest_qnan, + neg_smallest_qnan, neg_smallest_qnan, neg_smallest_qnan, + neg_smallest_qnan, neg_smallest_qnan, neg_smallest_qnan, + neg_smallest_qnan, neg_smallest_qnan, neg_smallest_qnan, + neg_smallest_qnan, neg_smallest_qnan, neg_smallest_qnan, + neg_smallest_qnan, neg_smallest_qnan, neg_smallest_qnan, + neg_smallest_qnan, neg_smallest_qnan, neg_smallest_qnan, + neg_smallest_qnan, neg_smallest_qnan, neg_smallest_qnan, + neg_smallest_qnan}; + + std::vector expected_indices_64 = { + 5, 19, 4, 1, 8, 18, 20, 21, 13, 23, 3, 16, 6, 12, 14, 24, + 2, 22, 7, 11, 0, 17, 9, 15, 10, 25, 26, 27, 28, 29, 30, 31, + 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, + 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63}; + + for (int n : {512, 1024, 32768, 65536 + 4096}) { + Array2D input(8, n, neg_smallest_qnan); + for (int i = 0; i < input_row_start.size(); ++i) { + for (int j = 0; j < 8; ++j) { + input(j, i) = input_row_start[i]; + } + } + + auto literal = LiteralUtil::CreateR2FromArray2D(input); + + for (int k : {2, 4, 8, 16, 32, 64}) { + const auto hlo_text_module = absl::StrFormat(R"( + HloModule topk + + ENTRY TopK { + x = f32[8,%d]{1,0} parameter(0) + ROOT topk = (f32[8,%d]{1,0}, s32[8,%d]{1,0}) topk(x), k=%d, largest=true + } + )", + n, k, k, k); + + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_text_module)); + + Array2D expected_vals_arr(8, k); + Array2D expected_idxs_arr(8, k); + for (int i = 0; i < 8; ++i) { + for (int j = 0; j < k; ++j) { + expected_vals_arr(i, j) = expected_values_64[j]; + expected_idxs_arr(i, j) = expected_indices_64[j]; + } + } + + auto expected_vals_literal = + LiteralUtil::CreateR2FromArray2D(expected_vals_arr); + auto expected_idxs_literal = + LiteralUtil::CreateR2FromArray2D(expected_idxs_arr); + auto expected_tuple = LiteralUtil::MakeTuple( + {&expected_vals_literal, &expected_idxs_literal}); + + ASSERT_OK_AND_ASSIGN(Literal actual_tuple, + Execute(std::move(module), {&literal})); + + EXPECT_TRUE(LiteralTestUtil::Near(expected_tuple, actual_tuple, + xla::ErrorSpec{0, 0})); + } + } +} + +bfloat16 ToBfloat16(uint16_t x) { return absl::bit_cast(x); } + +TEST_F(TopkTest, TopKWithSpecialBfloat16) { + // Regression test for TopK TotalOrder with special float values + // (NaN, +-Inf, +-0). + bfloat16 neg_qnan0 = ToBfloat16(0xFFC0u); + bfloat16 neg_qnan1 = ToBfloat16(0xFFE0u); + bfloat16 neg_qnan2 = ToBfloat16(0xFFE1u); + bfloat16 neg_smallest_qnan = ToBfloat16(0xFFFFu); + bfloat16 pos_qnan0 = ToBfloat16(0x7FC0u); + bfloat16 pos_qnan1 = ToBfloat16(0x7FE0u); + bfloat16 pos_qnan2 = ToBfloat16(0x7FE1u); + + bfloat16 pos_inf = ToBfloat16(0x7F80u); + bfloat16 neg_inf = ToBfloat16(0xFF80u); + bfloat16 pos_zero = ToBfloat16(0x0000u); + bfloat16 neg_zero = ToBfloat16(0x8000u); + + bfloat16 pos_dnrm_1e38 = ToBfloat16(0x0070u); // 1.028e-38 + bfloat16 pos_dnrm_1e40 = ToBfloat16(0x0002u); // 1.836e-40 + bfloat16 neg_dnrm_1e39 = ToBfloat16(0x8010u); // -1.469e-39 + bfloat16 neg_dnrm_9e41 = ToBfloat16(0x8001u); // -9.183e-41 + + bfloat16 pos_64_5 = ToBfloat16(0x4281u); + bfloat16 pos_0_0125 = ToBfloat16(0x3C4du); + bfloat16 neg_64_5 = ToBfloat16(0xC281u); + + std::vector input_row_start = { + neg_inf, pos_inf, neg_dnrm_1e39, pos_zero, pos_qnan0, + pos_qnan2, neg_zero, neg_64_5, pos_inf, neg_qnan0, + neg_qnan2, neg_64_5, neg_zero, pos_dnrm_1e40, neg_dnrm_9e41, + neg_qnan1, pos_zero, neg_inf, pos_64_5, pos_qnan1, + pos_0_0125, pos_dnrm_1e38, neg_dnrm_1e39, pos_dnrm_1e40, neg_dnrm_9e41}; + + std::vector expected_values_64 = { + pos_qnan0, pos_qnan0, pos_qnan0, pos_inf, pos_inf, + pos_64_5, pos_0_0125, pos_dnrm_1e38, pos_dnrm_1e40, pos_dnrm_1e40, + pos_zero, pos_zero, neg_zero, neg_zero, neg_dnrm_9e41, + neg_dnrm_9e41, neg_dnrm_1e39, neg_dnrm_1e39, neg_64_5, neg_64_5, + neg_inf, neg_inf, neg_qnan0, neg_qnan0, neg_qnan0, + neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0, + neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0, + neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0, + neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0, + neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0, + neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0, + neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0, + neg_qnan0, neg_qnan0, neg_qnan0, neg_qnan0}; + + std::vector expected_indices_64 = { + 5, 19, 4, 1, 8, 18, 20, 21, 13, 23, 3, 16, 6, 12, 14, 24, + 2, 22, 7, 11, 0, 17, 9, 15, 10, 25, 26, 27, 28, 29, 30, 31, + 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, + 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63}; + + for (int n : {512, 1024, 32768, 65536 + 4096}) { + Array2D input(8, n, neg_smallest_qnan); + for (int i = 0; i < input_row_start.size(); ++i) { + for (int j = 0; j < 8; ++j) { + input(j, i) = input_row_start[i]; + } + } + + auto literal = LiteralUtil::CreateR2FromArray2D(input); + + for (int k : {2, 4, 8, 16, 32, 64}) { + const auto hlo_text_module = absl::StrFormat(R"( + HloModule topk + + ENTRY TopK { + x = bf16[8,%d]{1,0} parameter(0) + ROOT topk = (bf16[8,%d]{1,0}, s32[8,%d]{1,0}) topk(x), k=%d, largest=true + } + )", + n, k, k, k); + + ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_text_module)); + + Array2D expected_vals_arr(8, k); + Array2D expected_idxs_arr(8, k); + for (int i = 0; i < 8; ++i) { + for (int j = 0; j < k; ++j) { + expected_vals_arr(i, j) = expected_values_64[j]; + expected_idxs_arr(i, j) = expected_indices_64[j]; + } + } + + auto expected_vals_literal = + LiteralUtil::CreateR2FromArray2D(expected_vals_arr); + auto expected_idxs_literal = + LiteralUtil::CreateR2FromArray2D(expected_idxs_arr); + auto expected_tuple = LiteralUtil::MakeTuple( + {&expected_vals_literal, &expected_idxs_literal}); + + ASSERT_OK_AND_ASSIGN(Literal actual_tuple, + Execute(std::move(module), {&literal})); + + EXPECT_TRUE(LiteralTestUtil::Near(expected_tuple, actual_tuple, + xla::ErrorSpec{0, 0})); + } + } +} + +} // namespace +} // namespace xla diff --git a/third_party/xla/xla/tsl/framework/allocator.h b/third_party/xla/xla/tsl/framework/allocator.h index 934cd58beea206..997232dd18dece 100644 --- a/third_party/xla/xla/tsl/framework/allocator.h +++ b/third_party/xla/xla/tsl/framework/allocator.h @@ -323,28 +323,8 @@ class AllocatorWrapper : public Allocator { return wrapped_->AllocateRaw(alignment, num_bytes, allocation_attr); } - void* AllocateRawAlignedNew(size_t alignment, size_t num_bytes) override { - return wrapped_->AllocateRawAlignedNew(alignment, num_bytes); - } - - void* AllocateRawAlignedNew( - size_t alignment, size_t num_bytes, - const AllocationAttributes& allocation_attr) override { - return wrapped_->AllocateRawAlignedNew(alignment, num_bytes, - allocation_attr); - } - void DeallocateRaw(void* ptr) override { wrapped_->DeallocateRaw(ptr); } - void DeallocateRaw(void* ptr, size_t alignment, size_t num_bytes) override { - wrapped_->DeallocateRaw(ptr, alignment, num_bytes); - } - - void DeallocateRawAlignedDelete(void* ptr, size_t alignment, - size_t num_bytes) override { - wrapped_->DeallocateRawAlignedDelete(ptr, alignment, num_bytes); - } - bool TracksAllocationSizes() const override { return wrapped_->TracksAllocationSizes(); } @@ -369,20 +349,6 @@ class AllocatorWrapper : public Allocator { return wrapped_->AllocatedSizeSlow(ptr); } - std::optional GetStats() override { - return wrapped_->GetStats(); - } - - bool ClearStats() override { return wrapped_->ClearStats(); } - - void SetSafeFrontier(uint64_t count) override { - wrapped_->SetSafeFrontier(count); - } - - void SetStreamAndPreallocateMemory(void* stream) override { - wrapped_->SetStreamAndPreallocateMemory(stream); - } - AllocatorMemoryType GetMemoryType() const override { return wrapped_->GetMemoryType(); } diff --git a/third_party/xla/xla/util/split_proto/split_gpu_executable_writer.cc b/third_party/xla/xla/util/split_proto/split_gpu_executable_writer.cc index 1309ac3b8db75e..20a080751ba808 100644 --- a/third_party/xla/xla/util/split_proto/split_gpu_executable_writer.cc +++ b/third_party/xla/xla/util/split_proto/split_gpu_executable_writer.cc @@ -122,6 +122,23 @@ absl::Status NormalizeBackendConfig(gpu::GpuExecutableProto& executable) { } else if (!instruction.backend_config().empty()) { instruction.set_backend_config(std::move(backend_config_str)); } + + // The payload table is shared with metadata payloads, so they must be + // re-mapped into the new table as well, otherwise their IDs would dangle. + if (instruction.has_metadata() && + instruction.metadata().has_metadata_payload()) { + Payload* payload = + instruction.mutable_metadata()->mutable_metadata_payload(); + if (payload->has_id()) { + const int64_t id = payload->id(); + if (id < 0 || id >= module->payloads_size()) { + return absl::InvalidArgumentError( + absl::StrCat("Invalid metadata payload id ", id, + " with payloads size ", module->payloads_size())); + } + payload->set_id(get_new_payload_id(module->payloads(id))); + } + } } } module->mutable_payloads()->Assign( diff --git a/third_party/xla/xla/util/split_proto/split_gpu_executable_writer_test.cc b/third_party/xla/xla/util/split_proto/split_gpu_executable_writer_test.cc index c5c2ffdc26b03a..b427b2ddd44516 100644 --- a/third_party/xla/xla/util/split_proto/split_gpu_executable_writer_test.cc +++ b/third_party/xla/xla/util/split_proto/split_gpu_executable_writer_test.cc @@ -220,5 +220,44 @@ TEST(SplitGpuExecutableWriterTest, read_comp->instructions(1).backend_config_payload().id()); } +TEST(SplitGpuExecutableWriterTest, MetadataPayloadsArePreserved) { + GpuExecutableProto proto; + HloModuleProto* module = + proto.mutable_hlo_module_with_config()->mutable_hlo_module(); + + module->add_payloads("backend_config"); // ID 0: backend config + module->add_payloads("metadata_0"); // ID 1: metadata + module->add_payloads("metadata_1"); // ID 2: metadata + + HloComputationProto* comp = module->add_computations(); + HloInstructionProto* instr0 = comp->add_instructions(); + instr0->mutable_backend_config_payload()->set_id(0); + instr0->mutable_metadata()->mutable_metadata_payload()->set_id(2); + HloInstructionProto* instr1 = comp->add_instructions(); + instr1->mutable_metadata()->mutable_metadata_payload()->set_id(1); + + std::string serialized; + ASSERT_OK(WriteSplitGpuExecutable( + proto, riegeli::Maker(&serialized))); + + GpuExecutableProto deserialized; + ASSERT_OK(ReadSplitProto(riegeli::Maker(serialized), + deserialized)); + + const HloModuleProto& read_module = + deserialized.hlo_module_with_config().hlo_module(); + ASSERT_EQ(read_module.payloads_size(), 3); + const HloComputationProto& read_comp = read_module.computations(0); + EXPECT_EQ(read_module.payloads( + read_comp.instructions(0).backend_config_payload().id()), + "backend_config"); + EXPECT_EQ(read_module.payloads( + read_comp.instructions(0).metadata().metadata_payload().id()), + "metadata_1"); + EXPECT_EQ(read_module.payloads( + read_comp.instructions(1).metadata().metadata_payload().id()), + "metadata_0"); +} + } // namespace } // namespace xla