diff --git a/MODULE.bazel b/MODULE.bazel index f3177c24804133..539d8107d2fa2f 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -20,7 +20,7 @@ bazel_dep(name = "nlohmann_json", version = "3.12.0.bcr.1", repo_name = "nlohman bazel_dep(name = "abseil-py", version = "2.1.0", repo_name = "absl_py") bazel_dep(name = "rules_python", version = "1.6.1") bazel_dep(name = "rules_shell", version = "0.6.1") -bazel_dep(name = "bazel_skylib", version = "1.8.1") +bazel_dep(name = "bazel_skylib", version = "1.9.0") bazel_dep(name = "platforms", version = "1.0.0") bazel_dep(name = "rules_license", version = "1.0.0") bazel_dep(name = "rules_pkg", version = "1.0.1") @@ -35,7 +35,7 @@ single_version_override( bazel_dep(name = "google_cloud_cpp", version = "3.0.0-rc1", repo_name = "com_github_googlecloudplatform_google_cloud_cpp") 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.11") +bazel_dep(name = "rules_cc", version = "0.2.20") bazel_dep(name = "curl", version = "8.11.0.bcr.5") 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") diff --git a/WORKSPACE b/WORKSPACE index a4147874be9920..f36369929ba12e 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -44,6 +44,10 @@ load("@bazel_features//:deps.bzl", "bazel_features_deps") bazel_features_deps() +load("@rules_cc//cc:extensions.bzl", "compatibility_proxy_repo") + +compatibility_proxy_repo() + load("@rules_shell//shell:repositories.bzl", "rules_shell_dependencies", "rules_shell_toolchains") rules_shell_dependencies() diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.cc b/tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.cc index 0f5f6dbd4dfa96..72808bf6f96d97 100644 --- a/tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.cc +++ b/tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.cc @@ -249,8 +249,12 @@ absl::Status LowerHlotoLoops(mlir::ModuleOp module, pm.addNestedPass(::mlir::createCSEPass()); // Collapse and tile parallel loops for GPU only. pm.addNestedPass(mlir::createCollapseParallelLoopsTo1DPass()); - pm.addNestedPass( - mlir::createTileLoopsPass(tile_sizes, unroll_factors)); + mlir::TileLoopsPassOptions tile_loops_options; + tile_loops_options.tile_sizes_ = + llvm::SmallVector(tile_sizes.begin(), tile_sizes.end()); + tile_loops_options.unroll_factors_ = + llvm::SmallVector(unroll_factors.begin(), unroll_factors.end()); + pm.addNestedPass(mlir::createTileLoopsPass(tile_loops_options)); pm.addNestedPass(::mlir::createCanonicalizerPass()); pm.addNestedPass(::mlir::createCSEPass()); diff --git a/tensorflow/lite/experimental/acceleration/mini_benchmark/gpu_module_plugin.cc b/tensorflow/lite/experimental/acceleration/mini_benchmark/gpu_module_plugin.cc index 57a964c3e100d4..4be81c5332a672 100644 --- a/tensorflow/lite/experimental/acceleration/mini_benchmark/gpu_module_plugin.cc +++ b/tensorflow/lite/experimental/acceleration/mini_benchmark/gpu_module_plugin.cc @@ -75,6 +75,10 @@ GpuModulePlugin::GpuModulePlugin(const TFLiteSettings& tflite_settings) { ->c_str(), dlopen_flags); if (!module_) { + TFLITE_LOG_PROD(TFLITE_LOG_WARNING, "Failed to load Gpu Module from %s", + tflite_settings_->stable_delegate_loader_settings() + ->delegate_path() + ->c_str()); error_code_ = kMinibenchmarkCannotLoadGpuModule; return; } diff --git a/tensorflow/lite/experimental/acceleration/mini_benchmark/gpu_module_plugin_test.cc b/tensorflow/lite/experimental/acceleration/mini_benchmark/gpu_module_plugin_test.cc index a9787cc284d264..b2d1e7895e13ef 100644 --- a/tensorflow/lite/experimental/acceleration/mini_benchmark/gpu_module_plugin_test.cc +++ b/tensorflow/lite/experimental/acceleration/mini_benchmark/gpu_module_plugin_test.cc @@ -65,7 +65,8 @@ void VerifyPluginCanLoadAndCreateDelegate(const TFLiteSettings& settings) { } // Verifies that the plugin initialization handles dlopen failures gracefully -// when provided with an invalid path (i.e. it doesn't crash). +// when provided with an invalid path (i.e. it doesn't crash) and logs a +// warning. TEST_F(GpuModulePluginTest, DlopenFlags) { const ComputeSettings* settings = nullptr; auto fbb = @@ -73,8 +74,25 @@ TEST_F(GpuModulePluginTest, DlopenFlags) { ASSERT_NE(settings, nullptr); ASSERT_NE(settings->tflite_settings(), nullptr); + // Stderr capture using GTest is only reliable on host platforms. While + // TFLITE_LOG_PROD also writes to stderr on Android, GTest's stream capture + // relies on creating a temporary file. Android does not have /tmp and + // /data/local/tmp may not be writable, so CaptureStderr silently fails + // unless TMPDIR is explicitly set to a writable location. +#if !defined(__ANDROID__) + testing::internal::CaptureStderr(); +#endif auto plugin = GpuModulePlugin::New(*settings->tflite_settings()); - ASSERT_NE(plugin.get(), nullptr); +#if !defined(__ANDROID__) + std::string captured_stderr = testing::internal::GetCapturedStderr(); +#endif + + ASSERT_NE(plugin, nullptr); +#if !defined(__ANDROID__) + EXPECT_NE(captured_stderr.find("Failed to load Gpu Module from " + "invalid_path_to_force_dlopen_fail.so"), + std::string::npos); +#endif } // Verifies that the plugin can be successfully loaded from a shared library diff --git a/tensorflow/tools/gcs_test/Dockerfile b/tensorflow/tools/gcs_test/Dockerfile index 6261ab8d6bb48f..dd2f74ecf8997d 100644 --- a/tensorflow/tools/gcs_test/Dockerfile +++ b/tensorflow/tools/gcs_test/Dockerfile @@ -13,7 +13,7 @@ # limitations under the License. # ============================================================================== -FROM ubuntu:26.04@sha256:b7f48194d4d8b763a478a621cdc81c27be222ba2206ca3ca6bc42b49685f3d9e +FROM ubuntu:26.04@sha256:2260313b31c8c011cd2eebe728008efac1b3982be73eb71348ea2648d2c0e09b LABEL maintainer="Shanqing Cai " diff --git a/tensorflow/workspace1.bzl b/tensorflow/workspace1.bzl index a058028e255b0e..1fa2077a8975c0 100644 --- a/tensorflow/workspace1.bzl +++ b/tensorflow/workspace1.bzl @@ -18,6 +18,7 @@ load("@com_google_benchmark//:bazel/benchmark_deps.bzl", "benchmark_deps") load("@grpc//bazel:grpc_deps.bzl", "grpc_deps") load("@io_bazel_rules_closure//closure:defs.bzl", "closure_repositories") +load("@rules_cc//cc:extensions.bzl", "compatibility_proxy_repo") load("@rules_pkg//:deps.bzl", "rules_pkg_dependencies") load("@xla//third_party/llvm:setup.bzl", "llvm_setup") load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls") @@ -33,6 +34,8 @@ def workspace(with_rules_cc = True): llvm_setup(name = "llvm-project") native.register_toolchains("@local_config_python//:py_toolchain") rules_pkg_dependencies() + if "cc_compatibility_proxy" not in native.existing_rules(): + compatibility_proxy_repo() closure_repositories() diff --git a/tensorflow/workspace3.bzl b/tensorflow/workspace3.bzl index 68dd7e9ad16d63..06809b6207741a 100644 --- a/tensorflow/workspace3.bzl +++ b/tensorflow/workspace3.bzl @@ -36,9 +36,9 @@ def workspace(): # https://github.com/bazelbuild/bazel-skylib/releases tf_http_archive( name = "bazel_skylib", - sha256 = "bc283cdfcd526a52c3201279cda4bc298652efa898b10b4db0837dc51652756f", + sha256 = "3b5b49006181f5f8ff626ef8ddceaa95e9bb8ad294f7b5d7b11ea9f7ddaf8c59", urls = tf_mirror_urls( - "https://github.com/bazelbuild/bazel-skylib/releases/download/1.7.1/bazel-skylib-1.7.1.tar.gz", + "https://github.com/bazelbuild/bazel-skylib/releases/download/1.9.0/bazel-skylib-1.9.0.tar.gz", ), ) @@ -67,6 +67,15 @@ def workspace(): ), ) + tf_http_archive( + name = "rules_cc", + sha256 = "69e05df29f0010ba248ef8dafc1f084c8fd2f5c553da634422d8167f5c4b277b", + strip_prefix = "rules_cc-0.2.20", + urls = tf_mirror_urls( + "https://github.com/bazelbuild/rules_cc/releases/download/0.2.20/rules_cc-0.2.20.tar.gz", + ), + ) + # Toolchains for ML projects hermetic builds. # Details: https://github.com/google-ml-infra/rules_ml_toolchain tf_http_archive( diff --git a/third_party/googleapis/build_rules.bzl b/third_party/googleapis/build_rules.bzl index 969ab719e093ab..a3067fb7d61ea3 100644 --- a/third_party/googleapis/build_rules.bzl +++ b/third_party/googleapis/build_rules.bzl @@ -16,8 +16,8 @@ Utilities for building grpc and proto libraries from googleapis. """ +load("@com_google_protobuf//bazel:cc_proto_library.bzl", native_cc_proto_library = "cc_proto_library") load("@grpc//bazel:generate_cc.bzl", "generate_cc") -load("@rules_cc//cc:defs.bzl", native_cc_proto_library = "cc_proto_library") def _tf_cc_headers(ctx): if len(ctx.attr.deps) != 1: diff --git a/third_party/xla/third_party/llvm/build.patch b/third_party/xla/third_party/llvm/build.patch index 868226d56a92b7..641bdb95bc262c 100644 --- a/third_party/xla/third_party/llvm/build.patch +++ b/third_party/xla/third_party/llvm/build.patch @@ -13,10 +13,9 @@ # 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 a7e652c..5b8ac5e 100644 --- a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel -@@ -234,19 +234,19 @@ write_file( +@@ -317,19 +317,19 @@ config_setting( name = "is_windows_clang_mingw", constraint_values = ["@platforms//os:windows"], @@ -39,7 +38,34 @@ index a7e652c..5b8ac5e 100644 ) config_setting( -@@ -255,7 +255,7 @@ config_setting( +@@ -338,7 +338,7 @@ + "@platforms//cpu:aarch64", + "@platforms//os:windows", + ], +- flag_values = {"@rules_cc//cc/compiler:compiler": "clang"}, ++ flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang"}, + ) + + config_setting( +@@ -347,7 +347,7 @@ + "@platforms//cpu:aarch64", + "@platforms//os:windows", + ], +- flag_values = {"@rules_cc//cc/compiler:compiler": "clang-cl"}, ++ flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang-cl"}, + ) + + config_setting( +@@ -356,7 +356,7 @@ + "@platforms//cpu:aarch64", + "@platforms//os:windows", + ], +- flag_values = {"@rules_cc//cc/compiler:compiler": "msvc-cl"}, ++ flag_values = {"@rules_cc//cc/private/toolchain:compiler": "msvc-cl"}, + ) + + config_setting( +@@ -365,7 +365,7 @@ "@platforms//cpu:x86_64", "@platforms//os:windows", ], @@ -47,8 +73,17 @@ index a7e652c..5b8ac5e 100644 + flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang"}, ) + config_setting( +@@ -374,7 +374,7 @@ + "@platforms//cpu:x86_64", + "@platforms//os:windows", + ], +- flag_values = {"@rules_cc//cc/compiler:compiler": "clang-cl"}, ++ flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang-cl"}, + ) + BLAKE3_x86_64_ASM_SOURCE_PATTERNS = [ -@@ -308,7 +308,8 @@ cc_library( +@@ -430,7 +430,8 @@ "@platforms//cpu:aarch64": [ "lib/Support/BLAKE3/blake3_neon.c", ], @@ -58,7 +93,7 @@ index a7e652c..5b8ac5e 100644 "//conditions:default": [ ], }), -@@ -337,8 +338,9 @@ cc_library( +@@ -459,8 +460,9 @@ ], "//conditions:default": ["BLAKE3_USE_NEON=0"], }) + select({ diff --git a/third_party/xla/third_party/llvm/toolchains.patch b/third_party/xla/third_party/llvm/toolchains.patch index 759b47a69df0f8..1170c7610ab573 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 9affa75801b7..2f681c82c298 100644 --- a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel +++ b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel -@@ -42,6 +42,36 @@ exports_files([ +@@ -57,6 +57,36 @@ + # This one is needed for building and vendoring out lldb from off tree. "utils/lldbDataFormatters.py", ]) - ++ +config_setting( + name = "macos_arm64", + constraint_values = [ @@ -49,23 +49,21 @@ index 9affa75801b7..2f681c82c298 100644 + "@platforms//cpu:x86_64", + ], +) -+ + config_setting( name = "darwin_arm64", - constraint_values = [ diff --git a/utils/bazel/llvm-project-overlay/llvm/config.bzl b/utils/bazel/llvm-project-overlay/llvm/config.bzl -index d592f9e5bfcc..7b87dad91e1b 100644 --- a/utils/bazel/llvm-project-overlay/llvm/config.bzl +++ b/utils/bazel/llvm-project-overlay/llvm/config.bzl -@@ -100,8 +100,9 @@ builtin_thread_pointer = select({ - - # TODO: We should split out host vs. target here. - llvm_config_defines = os_defines + builtin_thread_pointer + select({ +@@ -147,8 +147,9 @@ + Label("//llvm:is_aarch64_windows_clang_cl"): native_arch_defines("AArch64", "aarch64-pc-windows-msvc"), + Label("//llvm:is_aarch64_windows_msvc"): native_arch_defines("AArch64", "aarch64-pc-windows-msvc"), + Label("//llvm:is_x86_64_windows_clang_mingw"): native_arch_defines("X86", "x86_64-w64-windows-gnu"), - Label("//llvm:darwin_arm64"): native_arch_defines("AArch64", "arm64-apple-darwin"), - Label("//llvm:darwin_x86_64"): native_arch_defines("X86", "x86_64-unknown-darwin"), + Label("//llvm:macos_arm64"): native_arch_defines("AArch64", "arm64-apple-darwin"), + Label("//llvm:macos_x86_64"): native_arch_defines("X86", "x86_64-unknown-darwin"), + Label("//llvm:macos_x86_64_default"): native_arch_defines("X86", "x86_64-unknown-darwin"), Label("//llvm:linux_aarch64"): native_arch_defines("AArch64", "aarch64-unknown-linux-gnu"), + Label("//llvm:linux_armv7"): native_arch_defines("ARM", "armv7-linux-gnueabihf"), Label("//llvm:linux_ppc64le"): native_arch_defines("PowerPC", "powerpc64le-unknown-linux-gnu"), - Label("//llvm:linux_riscv64"): native_arch_defines("RISCV", "riscv64-unknown-linux-gnu"), diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch index ebb9dda1528d7c..e1b5fa813d250c 100644 --- a/third_party/xla/third_party/stablehlo/temporary.patch +++ b/third_party/xla/third_party/stablehlo/temporary.patch @@ -53,6 +53,295 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md * `is_type_name(x: Value | Placeholder | Type) -> Value`. Available for all types. For example, `is_float(x)` returns `true` if `x` is a `FloatType`. If `x` is a value or placeholder, this function is a shortcut for +diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp +--- stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp ++++ stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp +@@ -20,6 +20,7 @@ + #include + + #include "llvm/ADT/DenseSet.h" ++#include "llvm/ADT/STLExtras.h" + #include "llvm/ADT/SmallVector.h" + #include "llvm/ADT/StringRef.h" + #include "mlir/IR/Attributes.h" +@@ -31,57 +32,53 @@ + namespace mlir { + namespace stablehlo { + +-static SmallVector> ++namespace { ++ ++struct ReindexedAxes { ++ SmallVector splitAxisSizes; ++ SmallVector groupedAxisIndices; ++}; ++ ++// Generates replica groups from the reshaped mesh axis sizes and the indices of ++// the communication axes using a Reshape-Transpose permutation. ++SmallVector> + flattenedReplicaGroupsFromTransposePermutation( +- const SmallVector& meshAxisNames, +- const SmallVector& commAxisNames, +- const llvm::DenseSet& commAxisSet, +- const SmallVector& axisSizes, +- const SmallVector& deviceIds, int64_t totalDevices) { +- // Reshape and Transpose equivalence bridging XLA TileAssignment behavior. ++ ArrayRef axisSizes, ArrayRef groupedAxisIndices, ++ ArrayRef deviceIds, int64_t totalDevices) { ++ llvm::DenseSet groupedAxisSet(groupedAxisIndices.begin(), ++ groupedAxisIndices.end()); + SmallVector transposeAxes; + // Non-grouped axes first +- for (size_t i = 0; i < meshAxisNames.size(); ++i) { +- if (!commAxisSet.count(meshAxisNames[i])) { ++ for (size_t i = 0; i < axisSizes.size(); ++i) { ++ if (!groupedAxisSet.count(i)) { + transposeAxes.push_back(i); + } + } +- // Grouped axes +- for (const auto& name : commAxisNames) { +- for (size_t i = 0; i < meshAxisNames.size(); ++i) { +- if (meshAxisNames[i] == name) { +- transposeAxes.push_back(i); +- break; +- } +- } +- } +- +- SmallVector transposedSizes(meshAxisNames.size()); +- for (size_t i = 0; i < meshAxisNames.size(); ++i) { ++ // Grouped axes in the specified order ++ for (int64_t idx : groupedAxisIndices) { ++ transposeAxes.push_back(idx); ++ } ++ ++ SmallVector transposedSizes(axisSizes.size()); ++ for (size_t i = 0; i < axisSizes.size(); ++i) { + transposedSizes[i] = axisSizes[transposeAxes[i]]; + } + +- // Compute strides for original shape +- SmallVector originalStrides(meshAxisNames.size(), 1); +- for (int i = static_cast(meshAxisNames.size()) - 2; i >= 0; --i) { ++ // Compute strides for reshaped shape ++ SmallVector originalStrides(axisSizes.size(), 1); ++ for (int i = static_cast(axisSizes.size()) - 2; i >= 0; --i) { + originalStrides[i] = originalStrides[i + 1] * axisSizes[i + 1]; + } + + // Compute strides for transposed shape +- SmallVector transposedStrides(meshAxisNames.size(), 1); +- for (int i = static_cast(meshAxisNames.size()) - 2; i >= 0; --i) { ++ SmallVector transposedStrides(axisSizes.size(), 1); ++ for (int i = static_cast(axisSizes.size()) - 2; i >= 0; --i) { + transposedStrides[i] = transposedStrides[i + 1] * transposedSizes[i + 1]; + } + +- // Generate chunks + int64_t numDevicesPerGroup = 1; +- for (auto name : commAxisNames) { +- for (size_t i = 0; i < meshAxisNames.size(); ++i) { +- if (meshAxisNames[i] == name) { +- numDevicesPerGroup *= axisSizes[i]; +- break; +- } +- } ++ for (int64_t idx : groupedAxisIndices) { ++ numDevicesPerGroup *= axisSizes[idx]; + } + int64_t numGroups = totalDevices / numDevicesPerGroup; + +@@ -93,7 +90,7 @@ + for (int64_t j = 0; j < numDevicesPerGroup; ++j) { + int64_t linearTransposeIdx = i * numDevicesPerGroup + j; + int64_t originalIndex = 0; +- for (size_t k = 0; k < meshAxisNames.size(); ++k) { ++ for (size_t k = 0; k < axisSizes.size(); ++k) { + int64_t coord = + (linearTransposeIdx / transposedStrides[k]) % transposedSizes[k]; + originalIndex += coord * originalStrides[transposeAxes[k]]; +@@ -102,9 +99,128 @@ + } + groups.push_back(std::move(group)); + } +- + return groups; + } ++ ++// Splits mesh axes based on sub-axis references and computes the corresponding ++// indices for the communication axes. ++FailureOr computeReindexedAxes(ArrayRef axesInMesh, ++ ArrayAttr commAxes, ++ Location loc) { ++ ReindexedAxes result; ++ ++ // Validate commAxes and verify that all mesh axes exist and have valid sizes. ++ for (auto attr : commAxes) { ++ auto shloAxisRef = llvm::dyn_cast(attr); ++ if (!shloAxisRef) { ++ return emitError(loc) << "expected AxisRefAttr in comm_axes"; ++ } ++ StringRef axisName = shloAxisRef.getName(); ++ bool found = false; ++ for (auto meshAxis : axesInMesh) { ++ if (meshAxis.getName() == axisName) { ++ found = true; ++ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { ++ int64_t preSize = subAxisInfo.getPreSize(); ++ int64_t size = subAxisInfo.getSize(); ++ if (preSize < 1 || size < 1) { ++ return emitError(loc) ++ << "sub-axis pre_size and size must be at least 1"; ++ } ++ int64_t nextPreSize = preSize * size; ++ if (nextPreSize > meshAxis.getSize() || ++ meshAxis.getSize() % nextPreSize != 0) { ++ return emitError(loc) ++ << "sub-axis (pre_size * size) must divide mesh axis size"; ++ } ++ } ++ break; ++ } ++ } ++ if (!found) { ++ return emitError(loc) ++ << "axis '" << axisName << "' not found in mesh definition"; ++ } ++ } ++ ++ // Split each mesh axis according to the referenced subaxes. ++ struct SplitDim { ++ StringRef axisName; ++ int64_t preSize; ++ int64_t size; ++ int64_t dimIndex; ++ }; ++ SmallVector splitDims; ++ ++ for (auto meshAxis : axesInMesh) { ++ StringRef axisName = meshAxis.getName(); ++ int64_t axisSize = meshAxis.getSize(); ++ ++ SmallVector preSizes = {1, axisSize}; ++ for (auto attr : commAxes) { ++ auto shloAxisRef = llvm::cast(attr); ++ if (shloAxisRef.getName() == axisName) { ++ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { ++ preSizes.push_back(subAxisInfo.getPreSize()); ++ preSizes.push_back(subAxisInfo.getPreSize() * subAxisInfo.getSize()); ++ } ++ } ++ } ++ ++ llvm::sort(preSizes); ++ preSizes.erase(llvm::unique(preSizes), preSizes.end()); ++ ++ for (size_t j = 0; j < preSizes.size() - 1; ++j) { ++ int64_t segPreSize = preSizes[j]; ++ int64_t segSize = preSizes[j + 1] / segPreSize; ++ int64_t dimIdx = result.splitAxisSizes.size(); ++ result.splitAxisSizes.push_back(segSize); ++ splitDims.push_back({axisName, segPreSize, segSize, dimIdx}); ++ } ++ } ++ ++ // Map each communication axis to its corresponding split dimension. ++ llvm::DenseSet groupedSet; ++ for (auto attr : commAxes) { ++ auto shloAxisRef = llvm::cast(attr); ++ StringRef axisName = shloAxisRef.getName(); ++ int64_t reqPreSize = 1; ++ int64_t reqSize = 0; ++ if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) { ++ reqPreSize = subAxisInfo.getPreSize(); ++ reqSize = subAxisInfo.getSize(); ++ } else { ++ for (auto meshAxis : axesInMesh) { ++ if (meshAxis.getName() == axisName) { ++ reqSize = meshAxis.getSize(); ++ break; ++ } ++ } ++ } ++ ++ bool matched = false; ++ for (const auto& splitDim : splitDims) { ++ if (splitDim.axisName == axisName && splitDim.preSize == reqPreSize && ++ splitDim.size == reqSize) { ++ if (!groupedSet.insert(splitDim.dimIndex).second) { ++ return emitError(loc) ++ << "Duplicate or overlapping communication axis: " << axisName; ++ } ++ result.groupedAxisIndices.push_back(splitDim.dimIndex); ++ matched = true; ++ break; ++ } ++ } ++ if (!matched) { ++ return emitError(loc) << "Invalid or overlapping communication axis on '" ++ << axisName << "'"; ++ } ++ } ++ ++ return result; ++} ++ ++} // namespace + + FailureOr>> flattenReplicaGroupMeshAxes( + Attribute meshAttr, ArrayAttr commAxes, Location loc) { +@@ -120,34 +236,13 @@ + if (!mesh) + return emitOptionalError(loc, "expected stablehlo.mesh for mesh attribute"); + +- auto axesInMesh = mesh.getAxes(); +- +- // Identify which axes are communication axes. +- llvm::SmallVector commAxisNames; +- llvm::DenseSet commAxisSet; +- for (auto attr : commAxes) { +- auto shloAxisRef = llvm::dyn_cast(attr); +- if (!shloAxisRef) { +- return emitError(loc) << "expected AxisRefAttr in comm_axes"; +- } +- if (shloAxisRef.getSubAxisInfo()) { +- return emitError(loc) << "Subaxes are not supported in " +- "flattenReplicaGroupMeshAxes"; +- } +- commAxisNames.push_back(shloAxisRef.getName()); +- commAxisSet.insert(shloAxisRef.getName()); +- } +- +- // Calculate total devices and axis sizes ++ FailureOr reindexedAxes = ++ computeReindexedAxes(mesh.getAxes(), commAxes, loc); ++ if (failed(reindexedAxes)) return failure(); + + int64_t totalDevices = 1; +- SmallVector axisSizes; +- SmallVector meshAxisNames; +- for (auto meshAxis : axesInMesh) { +- auto typedMeshAxis = llvm::cast(meshAxis); +- axisSizes.push_back(typedMeshAxis.getSize()); +- meshAxisNames.push_back(typedMeshAxis.getName()); +- totalDevices *= typedMeshAxis.getSize(); ++ for (auto meshAxis : mesh.getAxes()) { ++ totalDevices *= llvm::cast(meshAxis).getSize(); + } + + SmallVector deviceIds; +@@ -160,8 +255,8 @@ + } + + return flattenedReplicaGroupsFromTransposePermutation( +- meshAxisNames, commAxisNames, commAxisSet, axisSizes, deviceIds, +- totalDevices); ++ reindexedAxes->splitAxisSizes, reindexedAxes->groupedAxisIndices, ++ deviceIds, totalDevices); + } + + } // namespace stablehlo diff --ruN a/stablehlo/stablehlo/dialect/Serialization.cpp b/stablehlo/stablehlo/dialect/Serialization.cpp --- stablehlo/stablehlo/dialect/Serialization.cpp +++ stablehlo/stablehlo/dialect/Serialization.cpp @@ -487,6 +776,196 @@ diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/st // CHECK-LABEL: func.func @ragged_dot_mode_3( // CHECK-SAME: %[[ARG0:.*]]: tensor<2x3x5xf32>, // CHECK-SAME: %[[ARG1:.*]]: tensor<2x5x7xf32>, +diff --ruN a/stablehlo/stablehlo/tests/interpret/all_gather.mlir b/stablehlo/stablehlo/tests/interpret/all_gather.mlir +--- stablehlo/stablehlo/tests/interpret/all_gather.mlir ++++ stablehlo/stablehlo/tests/interpret/all_gather.mlir +@@ -133,3 +133,41 @@ + func.return + } + } ++ ++// ----- ++ ++module @mesh_axes_subaxis { ++ func.func @all_gather(%operand : tensor<1xi64>) -> tensor<2xi64> { ++ %result = "stablehlo.all_gather"(%operand) { ++ all_gather_dim = 0 : i64, ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> ++ } : (tensor<1xi64>) -> tensor<2xi64> ++ return %result : tensor<2xi64> ++ } ++ func.func @main() { ++ %p0 = stablehlo.constant dense<[0]> : tensor<1xi64> ++ %p1 = stablehlo.constant dense<[1]> : tensor<1xi64> ++ %p2 = stablehlo.constant dense<[2]> : tensor<1xi64> ++ %p3 = stablehlo.constant dense<[3]> : tensor<1xi64> ++ %p4 = stablehlo.constant dense<[4]> : tensor<1xi64> ++ %p5 = stablehlo.constant dense<[5]> : tensor<1xi64> ++ %p6 = stablehlo.constant dense<[6]> : tensor<1xi64> ++ %p7 = stablehlo.constant dense<[7]> : tensor<1xi64> ++ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { ++ programs=[[@all_gather], [@all_gather], [@all_gather], [@all_gather], ++ [@all_gather], [@all_gather], [@all_gather], [@all_gather]] ++ } : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, ++ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> ++ (tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, ++ tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) ++ check.expect_eq_const %results#0, dense<[0, 2]> : tensor<2xi64> ++ check.expect_eq_const %results#1, dense<[1, 3]> : tensor<2xi64> ++ check.expect_eq_const %results#2, dense<[0, 2]> : tensor<2xi64> ++ check.expect_eq_const %results#3, dense<[1, 3]> : tensor<2xi64> ++ check.expect_eq_const %results#4, dense<[4, 6]> : tensor<2xi64> ++ check.expect_eq_const %results#5, dense<[5, 7]> : tensor<2xi64> ++ check.expect_eq_const %results#6, dense<[4, 6]> : tensor<2xi64> ++ check.expect_eq_const %results#7, dense<[5, 7]> : tensor<2xi64> ++ func.return ++ } ++} +diff --ruN a/stablehlo/stablehlo/tests/interpret/all_reduce.mlir b/stablehlo/stablehlo/tests/interpret/all_reduce.mlir +--- stablehlo/stablehlo/tests/interpret/all_reduce.mlir ++++ stablehlo/stablehlo/tests/interpret/all_reduce.mlir +@@ -135,3 +135,45 @@ + func.return + } + } ++ ++// ----- ++ ++module @mesh_axes_subaxis { ++ func.func @all_reduce(%operand : tensor<1xi64>) -> tensor<1xi64> { ++ %result = "stablehlo.all_reduce"(%operand) ({ ++ ^bb0(%arg0: tensor, %arg1: tensor): ++ %0 = stablehlo.add %arg0, %arg1 : tensor ++ stablehlo.return %0 : tensor ++ }) { ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]>, ++ channel_handle = #stablehlo.channel_handle ++ } : (tensor<1xi64>) -> tensor<1xi64> ++ return %result : tensor<1xi64> ++ } ++ func.func @main() { ++ %p0 = stablehlo.constant dense<[10]> : tensor<1xi64> ++ %p1 = stablehlo.constant dense<[20]> : tensor<1xi64> ++ %p2 = stablehlo.constant dense<[30]> : tensor<1xi64> ++ %p3 = stablehlo.constant dense<[40]> : tensor<1xi64> ++ %p4 = stablehlo.constant dense<[50]> : tensor<1xi64> ++ %p5 = stablehlo.constant dense<[60]> : tensor<1xi64> ++ %p6 = stablehlo.constant dense<[70]> : tensor<1xi64> ++ %p7 = stablehlo.constant dense<[80]> : tensor<1xi64> ++ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { ++ programs=[[@all_reduce], [@all_reduce], [@all_reduce], [@all_reduce], ++ [@all_reduce], [@all_reduce], [@all_reduce], [@all_reduce]] ++ } : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, ++ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> ++ (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, ++ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) ++ check.expect_eq_const %results#0, dense<[40]> : tensor<1xi64> ++ check.expect_eq_const %results#1, dense<[60]> : tensor<1xi64> ++ check.expect_eq_const %results#2, dense<[40]> : tensor<1xi64> ++ check.expect_eq_const %results#3, dense<[60]> : tensor<1xi64> ++ check.expect_eq_const %results#4, dense<[120]> : tensor<1xi64> ++ check.expect_eq_const %results#5, dense<[140]> : tensor<1xi64> ++ check.expect_eq_const %results#6, dense<[120]> : tensor<1xi64> ++ check.expect_eq_const %results#7, dense<[140]> : tensor<1xi64> ++ func.return ++ } ++} +diff --ruN a/stablehlo/stablehlo/tests/interpret/all_to_all.mlir b/stablehlo/stablehlo/tests/interpret/all_to_all.mlir +--- stablehlo/stablehlo/tests/interpret/all_to_all.mlir ++++ stablehlo/stablehlo/tests/interpret/all_to_all.mlir +@@ -172,3 +172,43 @@ + func.return %results#0, %results#1, %results#2, %results#3 : tensor<4x2xi64>, tensor<6x2xi32>, tensor<4x2xi64>, tensor<6x2xi32> + } + } ++ ++// ----- ++ ++module @mesh_axes_subaxis { ++ func.func @all_to_all(%operand : tensor<2x1xi64>) -> tensor<1x2xi64> { ++ %result = "stablehlo.all_to_all"(%operand) { ++ split_dimension = 0 : i64, ++ concat_dimension = 1 : i64, ++ split_count = 2 : i64, ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> ++ } : (tensor<2x1xi64>) -> tensor<1x2xi64> ++ return %result : tensor<1x2xi64> ++ } ++ func.func @main() { ++ %p0 = stablehlo.constant dense<[[1], [2]]> : tensor<2x1xi64> ++ %p1 = stablehlo.constant dense<[[3], [4]]> : tensor<2x1xi64> ++ %p2 = stablehlo.constant dense<[[10], [20]]> : tensor<2x1xi64> ++ %p3 = stablehlo.constant dense<[[30], [40]]> : tensor<2x1xi64> ++ %p4 = stablehlo.constant dense<[[5], [6]]> : tensor<2x1xi64> ++ %p5 = stablehlo.constant dense<[[7], [8]]> : tensor<2x1xi64> ++ %p6 = stablehlo.constant dense<[[50], [60]]> : tensor<2x1xi64> ++ %p7 = stablehlo.constant dense<[[70], [80]]> : tensor<2x1xi64> ++ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { ++ programs=[[@all_to_all], [@all_to_all], [@all_to_all], [@all_to_all], ++ [@all_to_all], [@all_to_all], [@all_to_all], [@all_to_all]] ++ } : (tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, ++ tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>) -> ++ (tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, ++ tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>) ++ check.expect_eq_const %results#0, dense<[[1, 10]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#1, dense<[[3, 30]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#2, dense<[[2, 20]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#3, dense<[[4, 40]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#4, dense<[[5, 50]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#5, dense<[[7, 70]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#6, dense<[[6, 60]]> : tensor<1x2xi64> ++ check.expect_eq_const %results#7, dense<[[8, 80]]> : tensor<1x2xi64> ++ func.return ++ } ++} +diff --ruN a/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir b/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir +--- stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir ++++ stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir +@@ -90,3 +90,45 @@ + func.return + } + } ++ ++// ----- ++ ++module @mesh_axes_subaxis { ++ func.func @reduce_scatter(%operand : tensor<2xi64>) -> tensor<1xi64> { ++ %result = "stablehlo.reduce_scatter"(%operand) ({ ++ ^bb0(%arg0: tensor, %arg1: tensor): ++ %0 = stablehlo.add %arg0, %arg1 : tensor ++ stablehlo.return %0 : tensor ++ }) { ++ scatter_dimension = 0 : i64, ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> ++ } : (tensor<2xi64>) -> tensor<1xi64> ++ return %result : tensor<1xi64> ++ } ++ func.func @main() { ++ %p0 = stablehlo.constant dense<[1, 2]> : tensor<2xi64> ++ %p1 = stablehlo.constant dense<[3, 4]> : tensor<2xi64> ++ %p2 = stablehlo.constant dense<[10, 20]> : tensor<2xi64> ++ %p3 = stablehlo.constant dense<[30, 40]> : tensor<2xi64> ++ %p4 = stablehlo.constant dense<[5, 6]> : tensor<2xi64> ++ %p5 = stablehlo.constant dense<[7, 8]> : tensor<2xi64> ++ %p6 = stablehlo.constant dense<[50, 60]> : tensor<2xi64> ++ %p7 = stablehlo.constant dense<[70, 80]> : tensor<2xi64> ++ %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) { ++ programs=[[@reduce_scatter], [@reduce_scatter], [@reduce_scatter], [@reduce_scatter], ++ [@reduce_scatter], [@reduce_scatter], [@reduce_scatter], [@reduce_scatter]] ++ } : (tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, ++ tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> ++ (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, ++ tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) ++ check.expect_eq_const %results#0, dense<[11]> : tensor<1xi64> ++ check.expect_eq_const %results#1, dense<[33]> : tensor<1xi64> ++ check.expect_eq_const %results#2, dense<[22]> : tensor<1xi64> ++ check.expect_eq_const %results#3, dense<[44]> : tensor<1xi64> ++ check.expect_eq_const %results#4, dense<[55]> : tensor<1xi64> ++ check.expect_eq_const %results#5, dense<[77]> : tensor<1xi64> ++ check.expect_eq_const %results#6, dense<[66]> : tensor<1xi64> ++ check.expect_eq_const %results#7, dense<[88]> : tensor<1xi64> ++ func.return ++ } ++} diff --ruN a/stablehlo/stablehlo/tests/ops_broadcasting.mlir b/stablehlo/stablehlo/tests/ops_broadcasting.mlir --- stablehlo/stablehlo/tests/ops_broadcasting.mlir +++ stablehlo/stablehlo/tests/ops_broadcasting.mlir @@ -505,6 +984,98 @@ diff --ruN a/stablehlo/stablehlo/tests/ops_broadcasting.mlir b/stablehlo/stableh + return %0 : tensor<3x4x5xf64> +} + +diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir b/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir +--- stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir ++++ stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir +@@ -5,7 +5,7 @@ + + // CHECK-LABEL: @all_reduce_rgv3 + func.func @all_reduce_rgv3(%arg0: tensor<4xf32>) -> tensor<4xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_reduce"(%arg0) ({ + ^bb0(%arg1: tensor, %arg2: tensor): + %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor +@@ -19,7 +19,7 @@ + + // CHECK-LABEL: @all_gather_rgv3 + func.func @all_gather_rgv3(%arg0: tensor<4xf32>) -> tensor<8xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 1], [2, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 1], [2, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_gather"(%arg0) { + all_gather_dim = 0 : i64, + replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> +@@ -30,7 +30,7 @@ + + // CHECK-LABEL: @all_to_all_rgv3 + func.func @all_to_all_rgv3(%arg0: tensor<4xf32>) -> tensor<4xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_to_all"(%arg0) { + concat_dimension = 0 : i64, + split_dimension = 0 : i64, +@@ -44,7 +44,7 @@ + + // CHECK-LABEL: @all_reduce_sdy_mesh + func.func @all_reduce_sdy_mesh(%arg0: tensor<4xf32>) -> tensor<4xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_reduce"(%arg0) ({ + ^bb0(%arg1: tensor, %arg2: tensor): + %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor +@@ -59,7 +59,7 @@ + + // CHECK-LABEL: @all_reduce_sdy_mesh_dev + func.func @all_reduce_sdy_mesh_dev(%arg0: tensor<4xf32>) -> tensor<4xf32> { +- // CHECK: replica_groups = dense<{{\[\[}}0, 1], [2, 3]]> : tensor<2x2xi64> ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 1], [2, 3]]> : tensor<2x2xi64> + %0 = "stablehlo.all_reduce"(%arg0) ({ + ^bb0(%arg1: tensor, %arg2: tensor): + %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor +@@ -84,4 +84,43 @@ + } : (tensor<4xf32>) -> tensor<4xf32> + return %0 : tensor<4xf32> + } ++ ++ // CHECK-LABEL: @all_reduce_subaxis ++ func.func @all_reduce_subaxis(%arg0: tensor<4xf32>) -> tensor<4xf32> { ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3], [4, 6], [5, 7]]> : tensor<4x2xi64> ++ %0 = "stablehlo.all_reduce"(%arg0) ({ ++ ^bb0(%arg1: tensor, %arg2: tensor): ++ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor ++ "stablehlo.return"(%1) : (tensor) -> () ++ }) { ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref]> ++ } : (tensor<4xf32>) -> tensor<4xf32> ++ return %0 : tensor<4xf32> ++ } ++ ++ // CHECK-LABEL: @all_reduce_subaxis_order_1 ++ func.func @all_reduce_subaxis_order_1(%arg0: tensor<4xf32>) -> tensor<4xf32> { ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], [20, 21, 22, 23, 24, 25, 26, 27, 28, 29]]> : tensor<3x10xi64> ++ %0 = "stablehlo.all_reduce"(%arg0) ({ ++ ^bb0(%arg1: tensor, %arg2: tensor): ++ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor ++ "stablehlo.return"(%1) : (tensor) -> () ++ }) { ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> ++ } : (tensor<4xf32>) -> tensor<4xf32> ++ return %0 : tensor<4xf32> ++ } ++ ++ // CHECK-LABEL: @all_reduce_subaxis_order_2 ++ func.func @all_reduce_subaxis_order_2(%arg0: tensor<4xf32>) -> tensor<4xf32> { ++ // CHECK{LITERAL}: replica_groups = dense<[[0, 5, 1, 6, 2, 7, 3, 8, 4, 9], [10, 15, 11, 16, 12, 17, 13, 18, 14, 19], [20, 25, 21, 26, 22, 27, 23, 28, 24, 29]]> : tensor<3x10xi64> ++ %0 = "stablehlo.all_reduce"(%arg0) ({ ++ ^bb0(%arg1: tensor, %arg2: tensor): ++ %1 = "stablehlo.add"(%arg1, %arg2) : (tensor, tensor) -> tensor ++ "stablehlo.return"(%1) : (tensor) -> () ++ }) { ++ replica_groups = #stablehlo.replica_group_mesh_axes, #stablehlo.mesh_axis]>, axes = [#stablehlo.axis_ref, #stablehlo.axis_ref]> ++ } : (tensor<4xf32>) -> tensor<4xf32> ++ return %0 : tensor<4xf32> ++ } + } diff --ruN a/stablehlo/stablehlo/tests/verify_convolution.mlir b/stablehlo/stablehlo/tests/verify_convolution.mlir --- stablehlo/stablehlo/tests/verify_convolution.mlir +++ stablehlo/stablehlo/tests/verify_convolution.mlir diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD index fc24f6ecb62a80..3f55250582219f 100644 --- a/third_party/xla/xla/BUILD +++ b/third_party/xla/xla/BUILD @@ -146,6 +146,7 @@ xla_cc_test( ":xla_data_proto_cc", "//xla/hlo/testlib:test", "//xla/tsl/platform:test_main", + "@com_google_absl//absl/status:status_matchers", "@com_google_googletest//:gtest", ], ) @@ -521,7 +522,6 @@ cc_library( "@com_google_absl//absl/functional:function_ref", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", - "@com_google_absl//absl/numeric:bits", "@com_google_absl//absl/status", "@com_google_absl//absl/status:status_macros", "@com_google_absl//absl/status:statusor", diff --git a/third_party/xla/xla/backends/cpu/codegen/builtin_fp16.h b/third_party/xla/xla/backends/cpu/codegen/builtin_fp16.h index cbe75ec2e18684..b24cee4bcaab79 100644 --- a/third_party/xla/xla/backends/cpu/codegen/builtin_fp16.h +++ b/third_party/xla/xla/backends/cpu/codegen/builtin_fp16.h @@ -16,6 +16,8 @@ limitations under the License. #ifndef XLA_BACKENDS_CPU_CODEGEN_BUILTIN_FP16_H_ #define XLA_BACKENDS_CPU_CODEGEN_BUILTIN_FP16_H_ +#include + // _Float16 always gets us the correct ABI type, so use that if available. // AArch64 GCC defines __FLT16_MANT_DIG__ even when _Float16 is not available. #if defined(__FLT16_MANT_DIG__) && \ diff --git a/third_party/xla/xla/backends/cpu/runtime/onednn/BUILD b/third_party/xla/xla/backends/cpu/runtime/onednn/BUILD index 6c63fab99c4af5..f9fd6fd0b94494 100644 --- a/third_party/xla/xla/backends/cpu/runtime/onednn/BUILD +++ b/third_party/xla/xla/backends/cpu/runtime/onednn/BUILD @@ -99,6 +99,7 @@ cc_library( name = "onednn_threadpool", hdrs = ["onednn_threadpool.h"], # copybara:uncomment compatible_with = ["//buildenv/target:non_prod"], + copts = tsl_copts(), deps = [ "//xla/backends/cpu/runtime:work_queue", "//xla/tsl/concurrency:async_value", diff --git a/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_interop.h b/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_interop.h index 50b321be6b6acb..367c27becfb125 100644 --- a/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_interop.h +++ b/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_interop.h @@ -40,9 +40,6 @@ namespace xla::cpu { } \ } while (0) -// Statically initializes XNNPACK for the current process. -absl::Status InitializeXnnPack(); - // Converts oneDNN status to absl::Status. inline absl::Status OneDnnStatusToStatus(dnnl::graph::status status) { if (ABSL_PREDICT_TRUE(status == dnnl::graph::status::success)) { diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton.cc b/third_party/xla/xla/backends/gpu/autotuner/triton.cc index e52708f236831e..3f4755c4ab40e7 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/triton.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/triton.cc @@ -73,40 +73,6 @@ namespace xla { namespace gpu { namespace { -std::vector GetDefaultTritonConfigs( - se::GpuComputeCapability compute_capability) { - if (compute_capability.IsRocm()) { - const auto* rocm_cc = compute_capability.rocm_compute_capability(); - if (rocm_cc->gfx9_mi300()) { - return GetTritonConfigsForPlatform(TritonConfigsPlatform::kMI300); - } - if (rocm_cc->gfx9_mi350()) { - return GetTritonConfigsForPlatform(TritonConfigsPlatform::kMI350); - } - return GetTritonConfigsForPlatform(TritonConfigsPlatform::kDefaultRocm); - } - - CHECK(compute_capability.IsCuda()); - auto* cuda_compute_capability = compute_capability.cuda_compute_capability(); - std::vector configs; - - if (cuda_compute_capability->IsBlackwell()) { - // SM 10.0 (datacenter: B200, B100) - configs = GetTritonConfigsForPlatform(TritonConfigsPlatform::kBlackwell); - } else if (cuda_compute_capability->IsAtLeastBlackwell()) { - // SM 11.0+ / 12.0+ (consumer: RTX 5090, etc.) - configs = - GetTritonConfigsForPlatform(TritonConfigsPlatform::kBlackwellConsumer); - } else if (cuda_compute_capability->IsHopper()) { - configs = GetTritonConfigsForPlatform(TritonConfigsPlatform::kHopper); - } else if (cuda_compute_capability->IsAmpere()) { - configs = GetTritonConfigsForPlatform(TritonConfigsPlatform::kAmpere); - } else { - configs = GetTritonConfigsForPlatform(TritonConfigsPlatform::kDefaultCuda); - } - - return configs; -} bool IsWarpSpecializationAvailable( se::GpuComputeCapability compute_capability) { @@ -204,28 +170,31 @@ TritonBackend::GetSupportedConfigsForDot(const HloInstruction* instr) { VLOG(1) << "Generating configs from search space: " << search_space.ToString(); - // We don't need to consider small_dot here. The new search space will - // already generate a unique config for small problems. - std::vector gemm_configs = search_space.GenerateConfigs( - /*autotune_warp_specialization=*/autotune_warp_specialization); - - if (!debug_options().xla_gpu_exhaustive_tiling_search()) { - VLOG(1) << "Restricting configs to the default set."; - std::vector all_configs = gemm_configs; - gemm_configs = search_space.OptimizeConfigSet( - gemm_configs, /*hints=*/GetDefaultTritonConfigs( - target_config().device_description.gpu_compute_capability())); - - if (!debug_options() - .xla_gpu_experimental_cost_model_gemm_tiling_options() - .empty()) { - ABSL_ASSIGN_OR_RETURN(gemm_configs, OptimizeConfigsWithCostModel( - dot, all_configs, gemm_configs, - target_config().device_description, - debug_options(), mlir_context_)); - } + + if (debug_options().xla_gpu_exhaustive_tiling_search()) { + return search_space.GenerateConfigs(autotune_warp_specialization); } - return gemm_configs; + + const std::vector& default_configs = + GetDefaultTritonConfigs( + target_config().device_description.gpu_compute_capability()); + + if (!debug_options() + .xla_gpu_experimental_cost_model_gemm_tiling_options() + .empty()) { + VLOG(1) << "Optimizing configs with the cost model."; + std::vector all_configs = + search_space.GenerateConfigs(autotune_warp_specialization); + std::vector candidate_configs = + search_space.OptimizeConfigSet(all_configs, default_configs); + return OptimizeConfigsWithCostModel(dot, all_configs, candidate_configs, + target_config().device_description, + debug_options(), mlir_context_); + } + + VLOG(1) << "Restricting configs to the default set."; + return search_space.GenerateAndOptimizeConfigs(default_configs, + autotune_warp_specialization); } absl::StatusOr> diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton/BUILD b/third_party/xla/xla/backends/gpu/autotuner/triton/BUILD index 8abaeebf46edec..0d25e1b2a1aee5 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/triton/BUILD +++ b/third_party/xla/xla/backends/gpu/autotuner/triton/BUILD @@ -54,6 +54,7 @@ cc_library( "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/types:span", "@com_google_protobuf//:protobuf_lite", "@llvm-project//llvm:Support", "@tsl//tsl/platform:protobuf", @@ -74,7 +75,6 @@ xla_cc_test( "//xla/stream_executor:device_description_proto_cc", "//xla/stream_executor/cuda:cuda_compute_capability", "//xla/stream_executor/rocm:rocm_compute_capability", - "//xla/tsl/platform:statusor", "@com_google_absl//absl/log:check", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings:str_format", @@ -112,6 +112,9 @@ cc_library( ":embed_default_configs", "//xla:autotuning_proto_cc", "//xla/service/gpu:matmul_utils", + "//xla/stream_executor:device_description", + "//xla/stream_executor/cuda:cuda_compute_capability", + "//xla/stream_executor/rocm:rocm_compute_capability", "@com_google_absl//absl/base:no_destructor", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log", @@ -126,6 +129,9 @@ xla_cc_test( srcs = ["triton_configs_test.cc"], deps = [ ":triton_configs", + "//xla/stream_executor:device_description", + "//xla/stream_executor/cuda:cuda_compute_capability", + "//xla/stream_executor/rocm:rocm_compute_capability", "@com_google_googletest//:gtest_main", ], ) 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 cd4a29a57927dd..67777810bcda26 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 @@ -26,6 +26,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/log/log.h" #include "absl/strings/str_format.h" +#include "absl/types/span.h" #include "llvm/ADT/STLExtras.h" #include "google/protobuf/repeated_field.h" #include "xla/backends/gpu/codegen/triton/tma_utils.h" @@ -155,11 +156,23 @@ std::vector TritonDotFusionSearchSpace::GenerateConfigs( return result; } +std::vector +TritonDotFusionSearchSpace::GenerateAndOptimizeConfigs( + absl::Span hints, + bool autotune_warp_specialization) const { + std::vector gemm_configs = + GenerateConfigs(autotune_warp_specialization); + if (hints.empty()) { + return gemm_configs; + } + return OptimizeConfigSet(gemm_configs, hints); +} + std::vector TritonDotFusionSearchSpace::OptimizeConfigSet( - const std::vector& configs, - const std::vector& hints) const { + absl::Span configs, + absl::Span hints) const { if (hints.empty() || configs.empty()) { - return configs; + return std::vector(configs.begin(), configs.end()); } absl::flat_hash_set filter; @@ -195,7 +208,7 @@ std::vector TritonDotFusionSearchSpace::OptimizeConfigSet( "sufficiently match the hints. Maybe the hints set does " "not contain a good representative set of valid configs? " "Working around this by using the full hints set instead."; - return hints; + return std::vector(hints.begin(), hints.end()); } return result_configs; } 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 252a9a291cb0e6..28b9530a753bcf 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 @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include "absl/types/span.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_instructions.h" #include "xla/service/gpu/matmul_utils.h" @@ -49,6 +50,15 @@ class TritonDotFusionSearchSpace { std::vector GenerateConfigs( bool autotune_warp_specialization = false) const; + // Generates the list of promising configs in the search space and optimizes + // them against the provided hints. + // + // If true, `autotune_warp_specialization` extends the search space with warp + // specialization support. + std::vector GenerateAndOptimizeConfigs( + absl::Span hints, + bool autotune_warp_specialization = false) const; + // Restrict the set of configs to the ones compatible with the hints list. // Generally, this will mean that configs are restricted to the ones that // appear in hints. The implementation is allowed to deviate though, and @@ -57,8 +67,8 @@ class TritonDotFusionSearchSpace { // hints list is larger than the problem's RHS side, it might restrict that // config to the problem's RHS size). std::vector OptimizeConfigSet( - const std::vector& configs, - const std::vector& hints) const; + absl::Span configs, + absl::Span hints) const; // Serializes the search space to a human-readable string. std::string ToString() const; 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 7b83b6ca845d6c..62340d81b85372 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 @@ -34,7 +34,6 @@ limitations under the License. #include "xla/stream_executor/device_description.h" #include "xla/stream_executor/device_description.pb.h" #include "xla/stream_executor/rocm/rocm_compute_capability.h" -#include "xla/tsl/platform/statusor.h" namespace xla::gpu { @@ -137,8 +136,8 @@ ENTRY e { }; TEST_F(DefaultDeviceDotSearchSpaceTest, ReturnsValidConfigList) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule()); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); EXPECT_THAT(search_space.GenerateConfigs(), Not(IsEmpty())); @@ -173,8 +172,8 @@ TEST_F(DotSearchSpaceTest, ExhaustiveSearchSpaceIsLargerThanDefault) { lhs_contracting_dims={1}, rhs_contracting_dims={1} })"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(kModuleText)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kModuleText)); auto default_search_space = MakeSearchSpace(module.get()); std::vector default_configs = default_search_space.GenerateConfigs(); @@ -190,7 +189,7 @@ TEST_F(DotSearchSpaceTest, ExhaustiveSearchSpaceIsLargerThanDefault) { } TEST_F(DotSearchSpaceTest, SerializesSearchSpace) { - TF_ASSERT_OK_AND_ASSIGN( + ASSERT_OK_AND_ASSIGN( std::unique_ptr module, GetDefaultDotModule(/*lhs_parallel_dim=*/1024, /*rhs_parallel_dim=*/1024, /*contracting_dim=*/1024)); @@ -204,8 +203,8 @@ TEST_F(DotSearchSpaceTest, SerializesSearchSpace) { } TEST_F(DotSearchSpaceTest, ReturnsValidConfigList) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule()); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); EXPECT_THAT(search_space.GenerateConfigs(), @@ -213,9 +212,9 @@ TEST_F(DotSearchSpaceTest, ReturnsValidConfigList) { } TEST_F(DotSearchSpaceTest, FindsGoodDataReuseOutputTiles) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule(/*lhs_parallel_dim=*/1024, - /*rhs_parallel_dim=*/1024)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule(/*lhs_parallel_dim=*/1024, + /*rhs_parallel_dim=*/1024)); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); EXPECT_THAT(search_space.GenerateConfigs(), @@ -223,9 +222,9 @@ TEST_F(DotSearchSpaceTest, FindsGoodDataReuseOutputTiles) { } TEST_F(DotSearchSpaceTest, RestrictsOutputToSquareishTiles) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule(/*lhs_parallel_dim=*/1024, - /*rhs_parallel_dim=*/1024)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule(/*lhs_parallel_dim=*/1024, + /*rhs_parallel_dim=*/1024)); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); EXPECT_THAT( @@ -242,8 +241,8 @@ ENTRY e { ROOT r = f16[4096,4096] dot(e0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0} })"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(kModuleText)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kModuleText)); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); EXPECT_THAT(search_space.GenerateConfigs(), @@ -259,8 +258,8 @@ ENTRY e { ROOT r = f16[4096,4096] dot(p0, e1), lhs_contracting_dims={1}, rhs_contracting_dims={0} })"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(kModuleText)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kModuleText)); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); EXPECT_THAT(search_space.GenerateConfigs(), @@ -268,20 +267,20 @@ ENTRY e { } TEST_F(DotSearchSpaceTest, PadsTilesForSmallParallelDimension) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule(/*lhs_parallel_dim=*/1024, - /*rhs_parallel_dim=*/15, - /*contracting_dim=*/1024)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule(/*lhs_parallel_dim=*/1024, + /*rhs_parallel_dim=*/15, + /*contracting_dim=*/1024)); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); EXPECT_THAT(search_space.GenerateConfigs(), Contains(BlockNIs(Eq(16)))); } TEST_F(DotSearchSpaceTest, HonorsMinimumOutputTileSizeForTinyProblem) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule(/*lhs_parallel_dim=*/12, - /*rhs_parallel_dim=*/8, - /*contracting_dim=*/16)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule(/*lhs_parallel_dim=*/12, + /*rhs_parallel_dim=*/8, + /*contracting_dim=*/16)); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); EXPECT_THAT( @@ -290,9 +289,9 @@ TEST_F(DotSearchSpaceTest, HonorsMinimumOutputTileSizeForTinyProblem) { } TEST_F(DotSearchSpaceTest, DoesNotBreakCtaSizeLimits) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule(/*lhs_parallel_dim=*/1024 * 16, - /*rhs_parallel_dim=*/1024 * 16)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule(/*lhs_parallel_dim=*/1024 * 16, + /*rhs_parallel_dim=*/1024 * 16)); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); EXPECT_THAT(search_space.GenerateConfigs(), @@ -300,9 +299,9 @@ TEST_F(DotSearchSpaceTest, DoesNotBreakCtaSizeLimits) { } TEST_F(DotSearchSpaceTest, ConsidersAppropriateCtaSizeForTileSize) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule(/*lhs_parallel_dim=*/4096, - /*rhs_parallel_dim=*/4096)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule(/*lhs_parallel_dim=*/4096, + /*rhs_parallel_dim=*/4096)); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); EXPECT_THAT(search_space.GenerateConfigs(), @@ -314,7 +313,7 @@ TEST_F(DotSearchSpaceTest, ConsidersAppropriateCtaSizeForTileSize) { // TODO: b/422419331 - Remove this once Triton properly handles 32-bit dots. TEST_F(DotSearchSpaceTest, ConsidersSmallCtasFor32BitDot) { - TF_ASSERT_OK_AND_ASSIGN( + ASSERT_OK_AND_ASSIGN( std::unique_ptr module, GetDefaultDotModule(/*lhs_parallel_dim=*/8 * 1024, /*rhs_parallel_dim=*/8 * 1024, @@ -327,7 +326,7 @@ TEST_F(DotSearchSpaceTest, ConsidersSmallCtasFor32BitDot) { } TEST_F(DotSearchSpaceTest, FindsFullCacheLineContractingTileSize) { - TF_ASSERT_OK_AND_ASSIGN( + ASSERT_OK_AND_ASSIGN( std::unique_ptr module, GetDefaultDotModule(/*lhs_parallel_dim=*/1024, /*rhs_parallel_dim=*/1024, /*contracting_dim=*/1024)); @@ -337,7 +336,7 @@ TEST_F(DotSearchSpaceTest, FindsFullCacheLineContractingTileSize) { } TEST_F(DotSearchSpaceTest, HonorsSharedMemoryLimit) { - TF_ASSERT_OK_AND_ASSIGN( + ASSERT_OK_AND_ASSIGN( std::unique_ptr module, GetDefaultDotModule(/*lhs_parallel_dim=*/4096, /*rhs_parallel_dim=*/4096, /*contracting_dim=*/4096)); @@ -355,7 +354,7 @@ TEST_F(DotSearchSpaceTest, HonorsSharedMemoryLimit) { } TEST_F(DotSearchSpaceTest, EnsuresContractingTileSizeFitsInstructonShape) { - TF_ASSERT_OK_AND_ASSIGN( + ASSERT_OK_AND_ASSIGN( std::unique_ptr module, GetDefaultDotModule(/*lhs_parallel_dim=*/1024, /*rhs_parallel_dim=*/1024, /*contracting_dim=*/4)); @@ -366,8 +365,8 @@ TEST_F(DotSearchSpaceTest, EnsuresContractingTileSizeFitsInstructonShape) { } TEST_F(DotSearchSpaceTest, FindReasonablePipeliningStageCount) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule()); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); EXPECT_THAT(search_space.GenerateConfigs(), @@ -376,7 +375,7 @@ TEST_F(DotSearchSpaceTest, FindReasonablePipeliningStageCount) { } TEST_F(DotSearchSpaceTest, ConsidersFewWarpsPerCtaAndMmaForSmallProblem) { - TF_ASSERT_OK_AND_ASSIGN( + ASSERT_OK_AND_ASSIGN( std::unique_ptr module, GetDefaultDotModule(/*lhs_parallel_dim=*/128, /*rhs_parallel_dim=*/128, /*contracting_dim=*/128)); @@ -388,10 +387,10 @@ TEST_F(DotSearchSpaceTest, ConsidersFewWarpsPerCtaAndMmaForSmallProblem) { } TEST_F(DotSearchSpaceTest, EnsuresWgmmaShapeForLargeProblem) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule(/*lhs_parallel_dim=*/16 * 1024, - /*rhs_parallel_dim=*/16 * 1024, - /*contracting_dim=*/4096)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule(/*lhs_parallel_dim=*/16 * 1024, + /*rhs_parallel_dim=*/16 * 1024, + /*contracting_dim=*/4096)); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); EXPECT_THAT( @@ -401,8 +400,8 @@ TEST_F(DotSearchSpaceTest, EnsuresWgmmaShapeForLargeProblem) { } TEST_F(DotSearchSpaceTest, ReturnsAllConfigsIfNoHints) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule()); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); std::vector configs = search_space.GenerateConfigs(); @@ -411,8 +410,8 @@ TEST_F(DotSearchSpaceTest, ReturnsAllConfigsIfNoHints) { } TEST_F(DotSearchSpaceTest, OptimizesEmptyConfigSet) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule()); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); TritonGemmConfig hint = {/*block_m=*/32, /*block_n=*/32, /*block_k=*/32, @@ -423,8 +422,8 @@ TEST_F(DotSearchSpaceTest, OptimizesEmptyConfigSet) { } TEST_F(DotSearchSpaceTest, RestrictsConfigsToHints) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule()); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); TritonGemmConfig matching_hint = { /*block_m=*/32, /*block_n=*/32, /*block_k=*/32, @@ -446,9 +445,9 @@ TEST_F(DotSearchSpaceTest, RestrictsConfigsToHints) { } TEST_F(DotSearchSpaceTest, ReturnsNonEmptySetForUnusualHints) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule(/*lhs_parallel_dim=*/4096, - /*rhs_parallel_dim=*/4096)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule(/*lhs_parallel_dim=*/4096, + /*rhs_parallel_dim=*/4096)); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); TritonGemmConfig hint = {/*block_m=*/1024, /*block_n=*/1024, @@ -461,9 +460,67 @@ TEST_F(DotSearchSpaceTest, ReturnsNonEmptySetForUnusualHints) { Not(IsEmpty())); } +TEST_F(DotSearchSpaceTest, GenerateAndOptimizeConfigsFiltersConfigsByHints) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule()); + TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); + std::vector all_configs = search_space.GenerateConfigs(); + ASSERT_FALSE(all_configs.empty()); + + TritonGemmConfig hint = all_configs.front(); + std::vector candidate_configs = + search_space.GenerateAndOptimizeConfigs({hint}); + EXPECT_THAT(candidate_configs, ElementsAre(hint)); +} + +TEST_F(DotSearchSpaceTest, + GenerateAndOptimizeConfigsWithEmptyDefaultConfigsReturnsAllConfigs) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule()); + TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); + std::vector all_configs = search_space.GenerateConfigs(); + ASSERT_FALSE(all_configs.empty()); + + std::vector candidate_configs = + search_space.GenerateAndOptimizeConfigs({}); + EXPECT_THAT(candidate_configs, ElementsAreArray(all_configs)); +} + +TEST_F(DotSearchSpaceTest, + GenerateAndOptimizeConfigsWithNonMatchingHintsFallsBackToHints) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule()); + TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); + + TritonGemmConfig non_matching_hint = { + /*block_m=*/9999, /*block_n=*/9999, + /*block_k=*/9999, /*num_stages=*/99, + /*num_warps=*/99, /*num_ctas=*/1}; + std::vector candidate_configs = + search_space.GenerateAndOptimizeConfigs({non_matching_hint}); + EXPECT_THAT(candidate_configs, ElementsAre(non_matching_hint)); +} + +TEST_F(DotSearchSpaceTest, + GenerateAndOptimizeConfigsWithWarpSpecializationPassesFlag) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule()); + TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); + std::vector no_ws_configs = + search_space.GenerateAndOptimizeConfigs( + {}, /*autotune_warp_specialization=*/false); + std::vector ws_configs = + search_space.GenerateAndOptimizeConfigs( + {}, /*autotune_warp_specialization=*/true); + EXPECT_GT(ws_configs.size(), no_ws_configs.size()); + EXPECT_THAT( + ws_configs, + Contains(Field(&TritonGemmConfig::is_warp_specialization_allowed, true))); +} + TEST_F(DotSearchSpaceTest, CudaDoesNotGenerateWavesPerEuConfigs) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule()); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); EXPECT_THAT(search_space.GenerateConfigs(), @@ -485,8 +542,8 @@ class RocmDotSearchSpaceTest : public DefaultDeviceDotSearchSpaceTest { }; TEST_F(RocmDotSearchSpaceTest, GeneratesWavesPerEuConfigs) { - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - GetDefaultDotModule()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + GetDefaultDotModule()); TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get()); std::vector configs = search_space.GenerateConfigs(); diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton/triton_configs.cc b/third_party/xla/xla/backends/gpu/autotuner/triton/triton_configs.cc index 36712718cacd96..8f0a28227adbcc 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/triton/triton_configs.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/triton/triton_configs.cc @@ -28,6 +28,9 @@ limitations under the License. #include "xla/autotuning.pb.h" #include "xla/backends/gpu/autotuner/triton/embed_default_configs.h" #include "xla/service/gpu/matmul_utils.h" +#include "xla/stream_executor/cuda/cuda_compute_capability.h" +#include "xla/stream_executor/device_description.h" +#include "xla/stream_executor/rocm/rocm_compute_capability.h" namespace xla::gpu { namespace { @@ -44,7 +47,7 @@ std::vector ParseConfig(absl::string_view config_str) { configs.push_back(*config); } return configs; -}; +} absl::string_view GetDefaultConfigStr(absl::string_view filename) { const struct FileToc* toc = configs::embed_default_configs_create(); @@ -81,4 +84,40 @@ const std::vector& GetTritonConfigsForPlatform( return kConfigs->at(platform); } +const std::vector& GetDefaultTritonConfigs( + const stream_executor::GpuComputeCapability& compute_capability) { + if (compute_capability.IsRocm()) { + const stream_executor::RocmComputeCapability* rocm_cc = + compute_capability.rocm_compute_capability(); + if (rocm_cc->gfx9_mi300()) { + return GetTritonConfigsForPlatform(TritonConfigsPlatform::kMI300); + } + if (rocm_cc->gfx9_mi350()) { + return GetTritonConfigsForPlatform(TritonConfigsPlatform::kMI350); + } + return GetTritonConfigsForPlatform(TritonConfigsPlatform::kDefaultRocm); + } + + CHECK(compute_capability.IsCuda()); + const stream_executor::CudaComputeCapability* cuda_compute_capability = + compute_capability.cuda_compute_capability(); + + if (cuda_compute_capability->IsBlackwell()) { + // SM 10.0 (datacenter: B200, B100) + return GetTritonConfigsForPlatform(TritonConfigsPlatform::kBlackwell); + } + if (cuda_compute_capability->IsAtLeastBlackwell()) { + // SM 11.0+ / 12.0+ (consumer: RTX 5090, etc.) + return GetTritonConfigsForPlatform( + TritonConfigsPlatform::kBlackwellConsumer); + } + if (cuda_compute_capability->IsHopper()) { + return GetTritonConfigsForPlatform(TritonConfigsPlatform::kHopper); + } + if (cuda_compute_capability->IsAmpere()) { + return GetTritonConfigsForPlatform(TritonConfigsPlatform::kAmpere); + } + return GetTritonConfigsForPlatform(TritonConfigsPlatform::kDefaultCuda); +} + } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton/triton_configs.h b/third_party/xla/xla/backends/gpu/autotuner/triton/triton_configs.h index f8a3658876c0fc..b39e218993aa82 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/triton/triton_configs.h +++ b/third_party/xla/xla/backends/gpu/autotuner/triton/triton_configs.h @@ -19,6 +19,7 @@ limitations under the License. #include #include "xla/service/gpu/matmul_utils.h" +#include "xla/stream_executor/device_description.h" namespace xla::gpu { @@ -36,6 +37,11 @@ enum class TritonConfigsPlatform { const std::vector& GetTritonConfigsForPlatform( TritonConfigsPlatform); +// Returns the default set of Triton GEMM configurations for the given GPU +// compute capability. +const std::vector& GetDefaultTritonConfigs( + const stream_executor::GpuComputeCapability& compute_capability); + } // namespace xla::gpu #endif // XLA_BACKENDS_GPU_AUTOTUNER_TRITON_TRITON_CONFIGS_H_ diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton/triton_configs_test.cc b/third_party/xla/xla/backends/gpu/autotuner/triton/triton_configs_test.cc index eb60b261978d13..c0b7f4643de7ce 100644 --- a/third_party/xla/xla/backends/gpu/autotuner/triton/triton_configs_test.cc +++ b/third_party/xla/xla/backends/gpu/autotuner/triton/triton_configs_test.cc @@ -17,6 +17,9 @@ limitations under the License. #include #include +#include "xla/stream_executor/cuda/cuda_compute_capability.h" +#include "xla/stream_executor/device_description.h" +#include "xla/stream_executor/rocm/rocm_compute_capability.h" namespace xla::gpu { namespace { @@ -40,5 +43,38 @@ TEST(TritonConfigsTest, PlatformsReturnNonEmptyConfig) { SizeIs(58)); } +TEST(TritonConfigsTest, GetDefaultTritonConfigsCuda) { + se::CudaComputeCapability hopper_cc{se::CudaComputeCapability::kHopper, 0}; + EXPECT_EQ(GetDefaultTritonConfigs(se::GpuComputeCapability{hopper_cc}), + GetTritonConfigsForPlatform(TritonConfigsPlatform::kHopper)); + + se::CudaComputeCapability ampere_cc{se::CudaComputeCapability::kAmpere, 0}; + EXPECT_EQ(GetDefaultTritonConfigs(se::GpuComputeCapability{ampere_cc}), + GetTritonConfigsForPlatform(TritonConfigsPlatform::kAmpere)); + + se::CudaComputeCapability blackwell_cc{se::CudaComputeCapability::kBlackwell, + 0}; + EXPECT_EQ(GetDefaultTritonConfigs(se::GpuComputeCapability{blackwell_cc}), + GetTritonConfigsForPlatform(TritonConfigsPlatform::kBlackwell)); + + se::CudaComputeCapability volta_cc{se::CudaComputeCapability::kVolta, 0}; + EXPECT_EQ(GetDefaultTritonConfigs(se::GpuComputeCapability{volta_cc}), + GetTritonConfigsForPlatform(TritonConfigsPlatform::kDefaultCuda)); +} + +TEST(TritonConfigsTest, GetDefaultTritonConfigsRocm) { + se::RocmComputeCapability mi300_cc("gfx942"); + EXPECT_EQ(GetDefaultTritonConfigs(se::GpuComputeCapability{mi300_cc}), + GetTritonConfigsForPlatform(TritonConfigsPlatform::kMI300)); + + se::RocmComputeCapability mi350_cc("gfx950"); + EXPECT_EQ(GetDefaultTritonConfigs(se::GpuComputeCapability{mi350_cc}), + GetTritonConfigsForPlatform(TritonConfigsPlatform::kMI350)); + + se::RocmComputeCapability default_rocm_cc("gfx908"); + EXPECT_EQ(GetDefaultTritonConfigs(se::GpuComputeCapability{default_rocm_cc}), + GetTritonConfigsForPlatform(TritonConfigsPlatform::kDefaultRocm)); +} + } // namespace } // namespace xla::gpu diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter.cc index 651e66366443e2..14c5bd07fe259b 100644 --- a/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter.cc +++ b/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter.cc @@ -1802,8 +1802,8 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor { // compatible with Epilogue Fusion. DEPRECATED: This standalone function has // been moved to GemmRewriterVisitor as a member function to allow bool SupportsEpilogueFusion(PrimitiveType type) { - // ROCm doesn't support F64 epilogue fusion - if (gpu_version_.IsRocm() && type == F64) { + // ROCm/oneAPI doesn't support F64 epilogue fusion + if ((gpu_version_.IsRocm() || gpu_version_.IsOneAPI()) && type == F64) { return false; } diff --git a/third_party/xla/xla/codegen/emitters/transforms/vectorize_loads_stores.cc b/third_party/xla/xla/codegen/emitters/transforms/vectorize_loads_stores.cc index 04cb4743fb1d17..366b13f3ed3209 100644 --- a/third_party/xla/xla/codegen/emitters/transforms/vectorize_loads_stores.cc +++ b/third_party/xla/xla/codegen/emitters/transforms/vectorize_loads_stores.cc @@ -23,6 +23,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/numeric/bits.h" #include "llvm/ADT/APInt.h" +#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/SmallVector.h" #include "llvm/Support/ErrorHandling.h" @@ -296,8 +297,12 @@ struct VectorizeLoad : mlir::OpRewritePattern { // trunc (extractelement <4 x i8> %X, i64 0) to i2 -> // extractelement <16 x i2> (bitcast <4 x i8> %X to <16 x i2>), i64 0. The // sub-byte vector types are not supported in the LLVM SPIR-V backend. + llvm::DenseSet visited; std::function has_sub_byte_trunc_user = [&](mlir::Operation* op) { + if (!visited.insert(op).second) { + return false; + } return absl::c_any_of(op->getUsers(), [&](mlir::Operation* user) { auto trunc = mlir::dyn_cast(user); if (trunc && IsSubByteIntOrFloatType(trunc.getResult().getType())) diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/eigen_unary.cc b/third_party/xla/xla/codegen/intrinsic/cpp/eigen_unary.cc index ddb5cecd40908c..1c3ded98ce8a8b 100644 --- a/third_party/xla/xla/codegen/intrinsic/cpp/eigen_unary.cc +++ b/third_party/xla/xla/codegen/intrinsic/cpp/eigen_unary.cc @@ -13,8 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#if defined(__has_attribute) && __has_attribute(ext_vector_type) && \ - defined(__has_builtin) && __has_builtin(__builtin_vectorelements) +#if defined(__FLT16_MANT_DIG__) && defined(__has_attribute) && \ + __has_attribute(ext_vector_type) && defined(__has_builtin) && \ + __has_builtin(__builtin_vectorelements) #include "xla/codegen/intrinsic/cpp/eigen_unary.h" diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops.h b/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops.h index af671549fed392..4218fb219b11b4 100644 --- a/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops.h +++ b/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops.h @@ -16,8 +16,9 @@ limitations under the License. #ifndef XLA_CODEGEN_INTRINSIC_CPP_VECTOR_OPS_H_ #define XLA_CODEGEN_INTRINSIC_CPP_VECTOR_OPS_H_ -#if defined(__has_attribute) && __has_attribute(ext_vector_type) && \ - defined(__has_builtin) && __has_builtin(__builtin_vectorelements) +#if defined(__FLT16_MANT_DIG__) && defined(__has_attribute) && \ + __has_attribute(ext_vector_type) && defined(__has_builtin) && \ + __has_builtin(__builtin_vectorelements) #include #include diff --git a/third_party/xla/xla/comparison_util.cc b/third_party/xla/xla/comparison_util.cc index 32b255823f9dfe..ddf6179c928a9e 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 { -// Verifies that this is a valid Comparison: (1) not a partial ordering on -// integers, and (2) a valid PrimitiveType. -bool IsValidComparison(xla::PrimitiveType type, Comparison::Order order) { - if (primitive_util::IsFloatingPointType(type) || - primitive_util::IsComplexType(type)) { - return true; - } - if (primitive_util::IsIntegralType(type) || type == PRED) { - return order == Comparison::Order::kTotal; - } - LOG(FATAL) << "Unsupported type: " << PrimitiveType_Name(type); -} - // Returns the X32 primitive type for each Type. PrimitiveType DefaultPrimitiveType(Comparison::Type type) { switch (type) { @@ -58,20 +45,8 @@ PrimitiveType DefaultPrimitiveType(Comparison::Type type) { } } -// Returns the default ordering for each Comparison::Type. -Comparison::Order 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 DefaultOrdering(PrimitiveType type) { +Comparison::Order DefaultPrimitiveOrdering(PrimitiveType type) { if (primitive_util::IsFloatingPointType(type) || primitive_util::IsComplexType(type)) { return Comparison::Order::kPartial; @@ -165,6 +140,15 @@ absl::string_view ComparisonOrderToString(Comparison::Order order) { } } +absl::string_view ComparisonOrderToShortString(Comparison::Order order) { + switch (order) { + case Comparison::Order::kPartial: + return "PARTIAL"; + case Comparison::Order::kTotal: + return "TOTAL"; + } +} + absl::StatusOr StringToComparisonDirection( absl::string_view direction) { static auto* const map = @@ -183,6 +167,20 @@ absl::StatusOr StringToComparisonDirection( return it->second; } +absl::StatusOr ShortStringToComparisonOrder( + absl::string_view order) { + static auto* const map = + new absl::flat_hash_map({ + {"TOTAL", Comparison::Order::kTotal}, + {"PARTIAL", Comparison::Order::kPartial}, + }); + auto it = map->find(order); + if (it == map->end()) { + return InvalidArgument("Unknown comparison order: %s", order); + } + return it->second; +} + absl::StatusOr StringToComparisonType( absl::string_view comparison) { static auto* const map = @@ -213,29 +211,54 @@ Comparison::Type Comparison::DefaultComparisonType(PrimitiveType type) { 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; + } +} + +namespace { +Comparison::Type ComparisonTypeFromPrimitiveTypeAndOrder( + PrimitiveType type, Comparison::Order order) { + if (primitive_util::IsFloatingPointType(type) || + primitive_util::IsComplexType(type)) { + return order == Comparison::Order::kTotal + ? Comparison::Type::kFloatTotalOrder + : Comparison::Type::kFloat; + } + if (primitive_util::IsSignedIntegralType(type)) { + return Comparison::Type::kSigned; + } + if (primitive_util::IsUnsignedIntegralType(type) || type == PRED) { + return Comparison::Type::kUnsigned; + } + LOG(FATAL) << "Unexpected: " << PrimitiveType_Name(type); +} +} // namespace + Comparison::Comparison(Direction dir, PrimitiveType type, Order order) : dir_(dir), primitive_type_(type), order_(order), - type_(DefaultComparisonType(type)) { - CHECK(IsValidComparison(primitive_type_, order_)); -} + type_(ComparisonTypeFromPrimitiveTypeAndOrder(type, order)) {} Comparison::Comparison(Direction dir, PrimitiveType type) : dir_(dir), primitive_type_(type), - order_(DefaultOrdering(type)), - type_(DefaultComparisonType(type)) { - CHECK(IsValidComparison(primitive_type_, order_)); -} + order_(DefaultPrimitiveOrdering(type)), + type_(DefaultComparisonType(type)) {} Comparison::Comparison(Direction dir, Type type) : dir_(dir), primitive_type_(DefaultPrimitiveType(type)), order_(DefaultOrdering(type)), - type_(type) { - CHECK(IsValidComparison(primitive_type_, order_)); -} + type_(type) {} Comparison Comparison::Converse() const { return Comparison(xla::Converse(dir_), primitive_type_, order_); diff --git a/third_party/xla/xla/comparison_util.h b/third_party/xla/xla/comparison_util.h index 64a49126151e8f..86df8e2a3c6fd3 100644 --- a/third_party/xla/xla/comparison_util.h +++ b/third_party/xla/xla/comparison_util.h @@ -66,6 +66,8 @@ class Comparison { }; friend absl::string_view ComparisonOrderToString(Comparison::Order order); + friend absl::string_view ComparisonOrderToShortString( + Comparison::Order order); template friend void AbslStringify(Sink& sink, const Order& p) { @@ -165,7 +167,7 @@ class Comparison { // Returns optional value because not all inversions may be supported. std::optional Inverse() const; - // Returns a string version of this comparison, e.g., ".GT.F32.TOTALORDER" + // Returns a string version of this comparison, e.g., ".GT.F32.TOTAL" std::string ToString(std::string prefix1 = ".", std::string prefix2 = ".", std::string prefix3 = ".") const; @@ -210,6 +212,10 @@ class Comparison { return GetComparator()(a, b); } + // Returns the Comparison::Order corresponding to the deprecated + // Comparison::Type. + static Comparison::Order DefaultOrdering(Type 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. @@ -238,6 +244,8 @@ 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); template void AbslStringify(Sink& sink, const ComparisonDirection& direction) { @@ -246,6 +254,8 @@ void AbslStringify(Sink& sink, const ComparisonDirection& direction) { absl::StatusOr StringToComparisonDirection( absl::string_view direction); +absl::StatusOr ShortStringToComparisonOrder( + absl::string_view order); absl::StatusOr StringToComparisonType( absl::string_view comparison); diff --git a/third_party/xla/xla/comparison_util_test.cc b/third_party/xla/xla/comparison_util_test.cc index e1f3f4c73e1e86..a0f10c173e148c 100644 --- a/third_party/xla/xla/comparison_util_test.cc +++ b/third_party/xla/xla/comparison_util_test.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "absl/status/status_matchers.h" #include "xla/hlo/testlib/test.h" #include "xla/types.h" #include "xla/xla_data.pb.h" @@ -26,6 +27,7 @@ limitations under the License. namespace xla { namespace { +using ::absl_testing::IsOkAndHolds; using ::testing::Eq; TEST(Comparison, FloatsDefaultToPartialOrder) { @@ -176,6 +178,28 @@ TEST(Comparison, ToString) { "_1_GE_2_C128_3_PARTIALORDER"); } +TEST(Comparison, ComparisonOrderToString) { + EXPECT_EQ(ComparisonOrderToString(Comparison::Order::kTotal), "TOTALORDER"); + EXPECT_EQ(ComparisonOrderToString(Comparison::Order::kPartial), + "PARTIALORDER"); +} + +TEST(Comparison, ComparisonOrderToShortString) { + EXPECT_EQ(ComparisonOrderToShortString(Comparison::Order::kTotal), "TOTAL"); + EXPECT_EQ(ComparisonOrderToShortString(Comparison::Order::kPartial), + "PARTIAL"); +} + +TEST(Comparison, ShortStringToComparisonOrder) { + EXPECT_THAT(ShortStringToComparisonOrder("TOTAL"), + IsOkAndHolds(Comparison::Order::kTotal)); + EXPECT_THAT(ShortStringToComparisonOrder("PARTIAL"), + IsOkAndHolds(Comparison::Order::kPartial)); + EXPECT_FALSE(ShortStringToComparisonOrder("TOTALORDER").ok()); + EXPECT_FALSE(ShortStringToComparisonOrder("PARTIALORDER").ok()); + EXPECT_FALSE(ShortStringToComparisonOrder("INVALID").ok()); +} + TEST(Comparison, TotalOrderFloatComparison) { EXPECT_TRUE(Comparison(Comparison::Direction::kEq, PrimitiveType::F32, Comparison::Order::kTotal) diff --git a/third_party/xla/xla/hlo/ir/BUILD b/third_party/xla/xla/hlo/ir/BUILD index 09a64187bd450f..5e27d9c78ef4af 100644 --- a/third_party/xla/xla/hlo/ir/BUILD +++ b/third_party/xla/xla/hlo/ir/BUILD @@ -186,6 +186,7 @@ xla_cc_test( srcs = ["hlo_instruction_test.cc"], deps = [ ":hlo", + "//xla:comparison_util", "//xla:frontend_attributes", "//xla:literal_util", "//xla:printer", @@ -200,6 +201,7 @@ xla_cc_test( "//xla/tsl/lib/core:status_test_util", "//xla/tsl/platform:statusor", "//xla/tsl/platform:test_main", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/status", "@com_google_absl//absl/strings:string_view", "@com_google_googletest//:gtest", diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc index 73012d47597a92..680243c17b4afa 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.cc @@ -518,18 +518,28 @@ absl::StatusOr> HloInstruction::CreateFromProto( comparison_direction, StringToComparisonDirection(proto.comparison_direction())); } - 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)); + 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_type); + *comparison_direction, comparison_order); } else { - // Allow the specify 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); + 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); + } } break; } @@ -1770,9 +1780,9 @@ HloInstruction::CreateRngBitGenerator(const Shape& shape, HloInstruction* state, /* static */ std::unique_ptr HloInstruction::CreateCompare( const Shape& shape, HloInstruction* lhs, HloInstruction* rhs, - ComparisonDirection direction, std::optional type) { + ComparisonDirection direction, std::optional order) { return std::make_unique(shape, lhs, rhs, direction, - type); + order); } /* static */ std::unique_ptr diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.h b/third_party/xla/xla/hlo/ir/hlo_instruction.h index b0a90e9fe7fd3f..33513636280a9c 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction.h +++ b/third_party/xla/xla/hlo/ir/hlo_instruction.h @@ -477,7 +477,7 @@ class HloInstruction { static std::unique_ptr CreateCompare( const Shape& shape, HloInstruction* lhs, HloInstruction* rhs, Comparison::Direction direction, - std::optional type = std::nullopt); + std::optional order = std::nullopt); static std::unique_ptr CreateTriangularSolve( const Shape& shape, HloInstruction* a, HloInstruction* b, diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc b/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc index fb5a4904ebadd7..44691380b98efb 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc @@ -24,11 +24,14 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_map.h" #include "absl/status/status.h" #include "absl/strings/string_view.h" +#include "xla/comparison_util.h" #include "xla/frontend_attributes.h" #include "xla/hlo/ir/dfs_hlo_visitor_with_default.h" #include "xla/hlo/ir/hlo_casting_utils.h" +#include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instructions.h" #include "xla/hlo/ir/hlo_opcode.h" #include "xla/hlo/ir/hlo_print_options.h" @@ -523,8 +526,8 @@ ENTRY main { ROOT call-done.0 = s32[] call-done(call-start.0) })"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(kHlo)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHlo)); ASSERT_TRUE(module->has_schedule()); TF_ASSERT_OK(module->schedule().Verify()); @@ -604,8 +607,8 @@ ENTRY main { ROOT collective-permute.0 = (f32[32,32]{1,0}, f32[32,32]{1,0}) collective-permute(arg.0, arg.0), channel_id=388, source_target_pairs={{0,0},{4,1}} })"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(kHlo)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHlo)); HloInstruction* cp = module->entry_computation()->root_instruction(); ASSERT_EQ(cp->opcode(), HloOpcode::kCollectivePermute); @@ -622,8 +625,7 @@ TEST_F(HloInstructionTest, PrintCompareOpWorksIfDead) { ROOT result = pred[] compare(p0, p1), direction=GT, type=TOTALORDER } )"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnVerifiedModule(kModuleStr)); + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kModuleStr)); HloInstruction* root = module->entry_computation()->root_instruction(); EXPECT_EQ( root->ToString(), @@ -642,8 +644,8 @@ TEST_F(HloInstructionTest, PrintCompareOpWorksIfDead) { } TEST_F(HloInstructionTest, CanonicalPrintingSupportsInt64) { - TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( - R"( + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + R"( HloModule m ENTRY main { p0 = f32[] parameter(0) @@ -685,8 +687,8 @@ TEST_F(HloInstructionTest, CanonicalPrintingSupportsInt64) { } TEST_F(HloInstructionTest, CanonicalPrintingSupportsCustomCall) { - TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( - R"( + ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule( + R"( HloModule custom_call_with_comp max_F32 { @@ -901,8 +903,8 @@ ENTRY main { } )"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnUnverifiedModule(kHlo)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(kHlo)); HloInstruction* start = FindInstruction(module.get(), "start"); HloInstruction* update = FindInstruction(module.get(), "update"); HloInstruction* done = FindInstruction(module.get(), "done"); @@ -973,8 +975,8 @@ ENTRY main { } )"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnUnverifiedModule(kHlo)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(kHlo)); HloInstruction* start = FindInstruction(module.get(), "start"); HloInstruction* update = FindInstruction(module.get(), "update"); HloInstruction* done = FindInstruction(module.get(), "done"); @@ -1021,5 +1023,33 @@ ENTRY main { } } +TEST_F(HloInstructionTest, CompareProtoRoundTripWithOrder) { + auto module = CreateNewVerifiedModule(); + HloComputation::Builder builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + HloInstruction* p0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")); + HloInstruction* p1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")); + HloInstruction* cmp = builder.AddInstruction(HloInstruction::CreateCompare( + ShapeUtil::MakeShape(PRED, {4}), p0, p1, ComparisonDirection::kLt, + ComparisonOrder::kTotal)); + module->AddEntryComputation(builder.Build()); + HloInstructionProto proto = cmp->ToProto(); + EXPECT_EQ(proto.comparison_direction(), "LT"); + EXPECT_EQ(proto.comparison_order(), "TOTAL"); + absl::flat_hash_map instruction_map; + instruction_map[p0->unique_id()] = p0; + instruction_map[p1->unique_id()] = p1; + + ASSERT_OK_AND_ASSIGN(auto clone, + HloInstruction::CreateFromProto(proto, instruction_map)); + EXPECT_EQ(clone->opcode(), HloOpcode::kCompare); + const auto* compare_clone = + static_cast(clone.get()); + EXPECT_EQ(compare_clone->direction(), ComparisonDirection::kLt); + EXPECT_EQ(compare_clone->order(), ComparisonOrder::kTotal); +} + } // namespace } // namespace xla diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.cc b/third_party/xla/xla/hlo/ir/hlo_instructions.cc index 0ed3f092b02857..4a04449f0504f5 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.cc +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.cc @@ -946,10 +946,10 @@ HloCopyStartInstruction::CloneWithNewOperandsImpl( HloCompareInstruction::HloCompareInstruction( const Shape& shape, HloInstruction* lhs, HloInstruction* rhs, - ComparisonDirection direction, std::optional type) + ComparisonDirection direction, std::optional order) : HloInstruction(HloOpcode::kCompare, shape), - compare_(type.has_value() - ? Comparison(direction, *type) + compare_(order.has_value() + ? Comparison(direction, lhs->shape().element_type(), *order) : Comparison(direction, lhs->shape().element_type())) { AppendOperand(lhs); AppendOperand(rhs); @@ -959,6 +959,8 @@ void HloCompareInstruction::ToProto(HloInstructionProto* proto) const { HloInstruction::ToProto(proto); proto->set_comparison_direction( ComparisonDirectionToString(compare_.GetDirection())); + proto->set_comparison_order( + ComparisonOrderToShortString(compare_.GetOrder())); proto->set_comparison_type(ComparisonTypeToString(compare_.GetType())); } @@ -983,7 +985,8 @@ bool HloCompareInstruction::IdenticalSlowPath( absl::FunctionRef eq_computations) const { const auto& casted_other = static_cast(other); - return direction() == casted_other.direction(); + return direction() == casted_other.direction() && + order() == casted_other.order(); } std::unique_ptr HloCompareInstruction::CloneWithNewOperandsImpl( @@ -991,7 +994,7 @@ std::unique_ptr HloCompareInstruction::CloneWithNewOperandsImpl( HloCloneContext* context) const { CHECK_EQ(new_operands.size(), 2); return std::make_unique( - shape, new_operands[0], new_operands[1], direction(), type()); + shape, new_operands[0], new_operands[1], direction(), order()); } namespace { diff --git a/third_party/xla/xla/hlo/ir/hlo_instructions.h b/third_party/xla/xla/hlo/ir/hlo_instructions.h index 7d1495baa9244f..10ae4cba26b71d 100644 --- a/third_party/xla/xla/hlo/ir/hlo_instructions.h +++ b/third_party/xla/xla/hlo/ir/hlo_instructions.h @@ -448,13 +448,16 @@ class HloCopyStartInstruction : public HloInstruction { class HloCompareInstruction : public HloInstruction { public: - explicit HloCompareInstruction(const Shape& shape, HloInstruction* lhs, - HloInstruction* rhs, - ComparisonDirection direction, - std::optional type); + explicit HloCompareInstruction( + const Shape& shape, HloInstruction* lhs, HloInstruction* rhs, + ComparisonDirection direction, + std::optional order = std::nullopt); ComparisonDirection direction() const { return compare_.GetDirection(); } ComparisonOrder order() const { return compare_.GetOrder(); } - Comparison::Type type() const { return compare_.GetType(); } + [[deprecated("Use order()")]] Comparison::Type type() const { + return compare_.GetType(); + } + const Comparison& comparison() const { return compare_; } void ToProto(HloInstructionProto* proto) const override; static bool ClassOf(const HloInstruction* hlo) { diff --git a/third_party/xla/xla/hlo/parser/BUILD b/third_party/xla/xla/hlo/parser/BUILD index 4b468fa1c389d3..6b90eb2a45b9d3 100644 --- a/third_party/xla/xla/hlo/parser/BUILD +++ b/third_party/xla/xla/hlo/parser/BUILD @@ -78,6 +78,7 @@ xla_cc_test( ":hlo_lexer", ":hlo_parser", "//xla:array", + "//xla:comparison_util", "//xla:shape_util", "//xla:window_util", "//xla:xla_data_proto_cc", diff --git a/third_party/xla/xla/hlo/parser/hlo_parser.cc b/third_party/xla/xla/hlo/parser/hlo_parser.cc index a5ec279715662d..21abddde1b2211 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser.cc @@ -323,6 +323,7 @@ class HloParserImpl : public HloParser { kFftType, kPaddingType, kComparisonDirection, + kComparisonOrder, kComparisonType, kWindow, kConvolutionDimensionNumbers, @@ -629,6 +630,7 @@ class HloParserImpl : public HloParser { bool ParsePaddingType(PaddingType* result); bool ParsePrimitiveType(PrimitiveType* result); bool ParseComparisonDirection(ComparisonDirection* result); + bool ParseComparisonOrder(Comparison::Order* result); bool ParseComparisonType(Comparison::Type* result); bool ParseFusionKind(HloInstruction::FusionKind* result); bool ParseRandomDistribution(RandomDistribution* result); @@ -2851,9 +2853,11 @@ HloInstruction* HloParserImpl::CreateInstruction( // NOLINT } case HloOpcode::kCompare: { optional direction; + optional order; optional type; attrs["direction"] = {/*required=*/true, AttrTy::kComparisonDirection, &direction}; + attrs["order"] = {/*required=*/false, AttrTy::kComparisonOrder, &order}; attrs["type"] = {/*required=*/false, AttrTy::kComparisonType, &type}; if ((!preset_operands && !ParseOperands(&operands, builder, /*expected_size=*/2)) || @@ -2866,8 +2870,22 @@ HloInstruction* HloParserImpl::CreateInstruction( // NOLINT })) { return nullptr; } + if (order.has_value() && type.has_value()) { + TokenError( + "Cannot specify both 'type' and 'order' attributes on compare"); + return nullptr; + } + if (order.has_value()) { + return builder->AddInstruction(HloInstruction::CreateCompare( + *shape, operands[0], operands[1], *direction, *order)); + } + if (type.has_value()) { + return builder->AddInstruction(HloInstruction::CreateCompare( + *shape, operands[0], operands[1], *direction, + Comparison::DefaultOrdering(*type))); + } return builder->AddInstruction(HloInstruction::CreateCompare( - *shape, operands[0], operands[1], *direction, type)); + *shape, operands[0], operands[1], *direction)); } case HloOpcode::kCholesky: { CholeskyOptions options; @@ -6117,6 +6135,15 @@ bool HloParserImpl::ParseAttributeHelper( ->emplace(result); return true; } + case AttrTy::kComparisonOrder: { + Comparison::Order result; + if (!ParseComparisonOrder(&result)) { + return false; + } + static_cast*>(attr_out_ptr) + ->emplace(result); + return true; + } case AttrTy::kComparisonType: { Comparison::Type result; if (!ParseComparisonType(&result)) { @@ -8371,6 +8398,21 @@ bool HloParserImpl::ParseComparisonDirection(ComparisonDirection* result) { return true; } +bool HloParserImpl::ParseComparisonOrder(Comparison::Order* result) { + VLOG(kDebugLevel) << "ParseComparisonOrder"; + if (lexer_.GetKind() != TokKind::kIdent) { + return TokenError("expects comparison order"); + } + std::string val = lexer_.GetStrVal(); + auto status_or_result = ShortStringToComparisonOrder(val); + if (!status_or_result.ok()) { + return TokenError(StrFormat("expects comparison order but sees: %s", val)); + } + *result = status_or_result.value(); + lexer_.Lex(); + return true; +} + bool HloParserImpl::ParseComparisonType(Comparison::Type* result) { VLOG(kDebugLevel) << "ParseComparisonType"; if (lexer_.GetKind() != TokKind::kIdent) { 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 ae3857c563526f..0de052fc0a1420 100644 --- a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc +++ b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc @@ -34,6 +34,7 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/types/span.h" #include "xla/array.h" +#include "xla/comparison_util.h" #include "xla/hlo/builder/xla_builder.h" #include "xla/hlo/ir/hlo_casting_utils.h" #include "xla/hlo/ir/hlo_instruction.h" @@ -3728,6 +3729,50 @@ ENTRY %configuration_test() -> s32[] { ->raw_backend_config_string()); } +TEST_F(HloParserTest, CompareWithOrder) { + const std::string original = R"(HloModule CompareWithOrder +ENTRY %entry(p0: f32[], p1: f32[]) -> pred[] { + %p0 = f32[] parameter(0) + %p1 = f32[] parameter(1) + ROOT %cmp = pred[] compare(f32[] %p0, f32[] %p1), direction=GT, order=TOTAL +})"; + 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::kGt); + EXPECT_EQ(compare->order(), ComparisonOrder::kTotal); +} + +TEST_F(HloParserTest, CompareBothTypeAndOrderFails) { + const std::string original = R"(HloModule CompareBothTypeAndOrderFails +ENTRY %entry(p0: f32[], p1: f32[]) -> pred[] { + %p0 = f32[] parameter(0) + %p1 = f32[] parameter(1) + ROOT %cmp = pred[] compare(f32[] %p0, f32[] %p1), direction=GT, type=FLOAT, order=TOTAL +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr( + result.status().message(), + "Cannot specify both 'type' and 'order' attributes on compare"); +} + +TEST_F(HloParserTest, CompareTotalOrderRejected) { + const std::string original = R"(HloModule CompareTotalOrderRejected +ENTRY %entry(p0: f32[], p1: f32[]) -> pred[] { + %p0 = f32[] parameter(0) + %p1 = f32[] parameter(1) + ROOT %cmp = pred[] compare(f32[] %p0, f32[] %p1), direction=GT, order=TOTALORDER +})"; + auto result = ParseAndReturnUnverifiedModule(original); + EXPECT_NE(absl::OkStatus(), result.status()); + ExpectHasSubstr(result.status().message(), + "expects comparison order but sees: TOTALORDER"); +} + TEST_F(HloParserTest, LiteralDimensionsError) { const std::string original = R"(HloModule some_2x3_module @@ -6709,8 +6754,8 @@ ENTRY AsyncDoneWithTransparentIntermediaries { ROOT async-done = f32[3,2] async-done(copy) } )"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnUnverifiedModule(hlo_string)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnUnverifiedModule(hlo_string)); HloInstruction* done = module->entry_computation()->root_instruction(); EXPECT_EQ(done->opcode(), HloOpcode::kAsyncDone); const HloInstruction* producer = 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 00480dc82bf636..a72a150261dba4 100644 --- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc +++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc @@ -2479,7 +2479,8 @@ absl::Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { Shape* shape; // exp(A)/exp(B) => exp(A-B) - if (Match(divide, m::Divide(m::Exp(m::Op(&a)), m::Exp(m::Op(&b))) + if (options_.enable_fast_math() && + Match(divide, m::Divide(m::Exp(m::Op(&a)), m::Exp(m::Op(&b))) .WithShape(m::Shape(&shape)))) { VLOG(10) << "transform [exp(A)/exp(B) => exp(A-B)]: " << divide->ToString(); HloInstruction* subtract = divide->AddInstruction( @@ -2489,7 +2490,8 @@ absl::Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide) { } // A/exp(B) => A*exp(-B) - if (Match(divide, m::Divide(m::Op(&a), m::Exp(m::Op(&b))))) { + if (options_.enable_fast_math() && + Match(divide, m::Divide(m::Op(&a), m::Exp(m::Op(&b))))) { VLOG(10) << "transform [A/exp(B) => A*exp(-B)]: " << divide->ToString(); HloInstruction* negate = divide->AddInstruction( HloInstruction::CreateUnary(divide->shape(), HloOpcode::kNegate, b)); @@ -5164,7 +5166,8 @@ absl::Status AlgebraicSimplifierVisitor::HandleMultiply( VLOG(10) << "trying to transform exp(LHS) * exp(RHS) => exp(LHS+RHS) " << multiply->ToString(); - if (Match(multiply, m::Multiply(m::Exp(m::Op(&lhs)), m::Exp(m::Op(&rhs))))) { + if (options_.enable_fast_math() && + Match(multiply, m::Multiply(m::Exp(m::Op(&lhs)), m::Exp(m::Op(&rhs))))) { auto add = multiply->AddInstruction(HloInstruction::CreateBinary( multiply->shape(), HloOpcode::kAdd, lhs, rhs)); return ReplaceWithNewInstruction( @@ -6194,7 +6197,8 @@ absl::Status AlgebraicSimplifierVisitor::HandlePower(HloInstruction* power) { // pow(exp(A),B) => exp(A*B) HloInstruction *a, *b; - if (Match(power, m::Power(m::Exp(m::Op(&a)), m::Op(&b)))) { + if (options_.enable_fast_math() && + Match(power, m::Power(m::Exp(m::Op(&a)), m::Op(&b)))) { auto a_times_b = power->AddInstruction(HloInstruction::CreateBinary( power->shape(), HloOpcode::kMultiply, a, b)); return ReplaceWithNewInstruction( 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 294e6653bb57a9..5c827aec5f682e 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 @@ -2524,7 +2524,7 @@ TEST_F(AlgebraicSimplifierTest, DivOfDivAndDiv) { m::Multiply(m::Parameter(1), m::Parameter(2))))); } -// Test that A/exp(B) is simplified to A*exp(-B). +// Test that A/exp(B) is simplified only with fast-math. TEST_F(AlgebraicSimplifierTest, DivOfExp) { auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); @@ -2544,7 +2544,16 @@ TEST_F(AlgebraicSimplifierTest, DivOfExp) { GmockMatch(m::Divide(m::Parameter(0), m::Exp(m::Parameter(1))))); AlgebraicSimplifier simplifier(default_options_); - ASSERT_TRUE(simplifier.Run(m.get()).value()); + TF_ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(m.get())); + EXPECT_FALSE(changed); + + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Divide(m::Parameter(0), m::Exp(m::Parameter(1))))); + + AlgebraicSimplifierOptions options = default_options_; + options.set_enable_fast_math(true); + AlgebraicSimplifier fast_math_simplifier(options); + ASSERT_TRUE(fast_math_simplifier.Run(m.get()).value()); EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Multiply(m::Parameter(0), @@ -2841,7 +2850,7 @@ TEST_F(AlgebraicSimplifierTest, SelectMakeTuple) { EXPECT_THAT(root, GmockMatch(m::Add(m::Parameter(1), m::Parameter(2)))); } -// Test that exp(A)/exp(B) is simplified to exp(A-B) +// Test that exp(A)/exp(B) is simplified only with fast-math. TEST_F(AlgebraicSimplifierTest, ExpDiv) { auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); @@ -2864,14 +2873,24 @@ TEST_F(AlgebraicSimplifierTest, ExpDiv) { GmockMatch(m::Divide(m::Exp(m::Parameter(0)), m::Exp(m::Parameter(1))))); AlgebraicSimplifier simplifier(default_options_); - ASSERT_TRUE(simplifier.Run(m.get()).value()); + TF_ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(m.get())); + EXPECT_FALSE(changed); + + EXPECT_THAT( + computation->root_instruction(), + GmockMatch(m::Divide(m::Exp(m::Parameter(0)), m::Exp(m::Parameter(1))))); + + AlgebraicSimplifierOptions options = default_options_; + options.set_enable_fast_math(true); + AlgebraicSimplifier fast_math_simplifier(options); + ASSERT_TRUE(fast_math_simplifier.Run(m.get()).value()); EXPECT_THAT( computation->root_instruction(), GmockMatch(m::Exp(m::Subtract(m::Parameter(0), m::Parameter(1))))); } -// Test that exp(A)*exp(B) is simplified to exp(A+B) +// Test that exp(A)*exp(B) is simplified only with fast-math. TEST_F(AlgebraicSimplifierTest, ExpMul) { auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); @@ -2894,13 +2913,23 @@ TEST_F(AlgebraicSimplifierTest, ExpMul) { m::Exp(m::Parameter(1))))); AlgebraicSimplifier simplifier(default_options_); - ASSERT_TRUE(simplifier.Run(m.get()).value()); + TF_ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(m.get())); + EXPECT_FALSE(changed); + + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Multiply(m::Exp(m::Parameter(0)), + m::Exp(m::Parameter(1))))); + + AlgebraicSimplifierOptions options = default_options_; + options.set_enable_fast_math(true); + AlgebraicSimplifier fast_math_simplifier(options); + ASSERT_TRUE(fast_math_simplifier.Run(m.get()).value()); EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Exp(m::Add(m::Parameter(0), m::Parameter(1))))); } -// Test that pow(exp(A), B) is simplified to exp(A*B) +// Test that pow(exp(A), B) is simplified only with fast-math. TEST_F(AlgebraicSimplifierTest, PowExp) { auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); @@ -2920,7 +2949,16 @@ TEST_F(AlgebraicSimplifierTest, PowExp) { GmockMatch(m::Power(m::Exp(m::Parameter(0)), m::Parameter(1)))); AlgebraicSimplifier simplifier(default_options_); - ASSERT_TRUE(simplifier.Run(m.get()).value()); + TF_ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(m.get())); + EXPECT_FALSE(changed); + + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Power(m::Exp(m::Parameter(0)), m::Parameter(1)))); + + AlgebraicSimplifierOptions options = default_options_; + options.set_enable_fast_math(true); + AlgebraicSimplifier fast_math_simplifier(options); + ASSERT_TRUE(fast_math_simplifier.Run(m.get()).value()); EXPECT_THAT( computation->root_instruction(), diff --git a/third_party/xla/xla/literal.h b/third_party/xla/xla/literal.h index 9324772cce53d4..5a965de7394736 100644 --- a/third_party/xla/xla/literal.h +++ b/third_party/xla/xla/literal.h @@ -597,14 +597,14 @@ class LiteralBase { void WriteElement(NativeT element) { constexpr PrimitiveType primitive_type = primitive_util::NativeToPrimitiveType(); - static_assert(primitive_util::BitWidth(primitive_type) % 8 == 0); + static_assert(primitive_util::StorageBitWidth(primitive_type) % 8 == 0); if constexpr (primitive_util::IsComplexType(primitive_type)) { WriteElement(element.real()); WriteElement(element.imag()); } else { constexpr PrimitiveType unsigned_type = primitive_util::UnsignedIntegralTypeForBitWidth( - primitive_util::BitWidth(primitive_type)); + primitive_util::StorageBitWidth(primitive_type)); using UnsignedT = primitive_util::NativeTypeOf; UnsignedT unsigned_element = absl::bit_cast(element); if constexpr (sizeof(UnsignedT) == 1) { @@ -625,7 +625,7 @@ class LiteralBase { constexpr PrimitiveType primitive_type = primitive_util::NativeToPrimitiveType(); constexpr int bits_per_element = primitive_util::BitWidth(primitive_type); - if constexpr (bits_per_element < 8) { + if constexpr (primitive_util::IsSubByteNonPredType(primitive_type)) { static_assert(!primitive_util::IsComplexType(primitive_type)); static_assert(8 % bits_per_element == 0); @@ -688,7 +688,7 @@ class LiteralBase { ABSL_MUST_USE_RESULT bool ReadElement(NativeT& element) { constexpr PrimitiveType primitive_type = primitive_util::NativeToPrimitiveType(); - static_assert(primitive_util::BitWidth(primitive_type) % 8 == 0); + static_assert(primitive_util::StorageBitWidth(primitive_type) % 8 == 0); if constexpr (primitive_util::IsComplexType(primitive_type)) { using ComponentT = primitive_util::NativeTypeOf; if constexpr (sizeof(UnsignedT) == 1) { if (at_end()) { @@ -736,7 +736,7 @@ class LiteralBase { constexpr PrimitiveType primitive_type = primitive_util::NativeToPrimitiveType(); constexpr int bits_per_element = primitive_util::BitWidth(primitive_type); - if constexpr (bits_per_element < 8) { + if constexpr (primitive_util::IsSubByteNonPredType(primitive_type)) { static_assert(!primitive_util::IsComplexType(primitive_type)); static_assert(8 % bits_per_element == 0); @@ -1758,7 +1758,7 @@ bool LiteralBase::Piece::DeserializeData( // - If a piece is dynamic, we first write the sizes of the dynamic dimensions. // // - The elements of the piece are then written. Elements smaller than a single -// byte (PRED, S4, U4) are packed into bytes. Otherwise, they are written in +// byte (e.g. S4, U4) are packed into bytes. Otherwise, they are written in // little-endian byte order. template absl::Status LiteralBase::SerializeWithShapeProto(const ShapeProto& shape_proto, diff --git a/third_party/xla/xla/mlir_hlo/mhlo/transforms/passes.h b/third_party/xla/xla/mlir_hlo/mhlo/transforms/passes.h index b45c9d6d7a47fb..a1142a7759bea0 100644 --- a/third_party/xla/xla/mlir_hlo/mhlo/transforms/passes.h +++ b/third_party/xla/xla/mlir_hlo/mhlo/transforms/passes.h @@ -44,14 +44,6 @@ ChloLegalizeToHighLevelMhloPassOptions getDefaultChloToHighLevelMhloOptions(); /// Returns options for the ChloLegalizeToHighLevelMhloPass for the GPU backend. ChloLegalizeToHighLevelMhloPassOptions getGpuChloToHighLevelMhloOptions(); -// TODO(b/397167511): Remove legacy wrapper once callers are migrated. -inline std::unique_ptr createExpandHloTuplesPass( - const std::string& entryFunctionName) { - ExpandHloTuplesPassOptions options; - options.entry_function_name_ = entryFunctionName; - return createExpandHloTuplesPass(options); -} - #define GEN_PASS_REGISTRATION #include "mhlo/transforms/mhlo_passes.h.inc" diff --git a/third_party/xla/xla/python/ifrt_proxy/common/array_util.cc b/third_party/xla/xla/python/ifrt_proxy/common/array_util.cc index 0de6ba869f654d..1bae9d8549fe88 100644 --- a/third_party/xla/xla/python/ifrt_proxy/common/array_util.cc +++ b/third_party/xla/xla/python/ifrt_proxy/common/array_util.cc @@ -125,7 +125,13 @@ absl::StatusOr ArrayMemRegion::FromZerothElementPointer( // `shape.dims()[i]` cannot be negative (we explicitly check for this // above) or zero (we return early for `shape.num_elements() == 0`). DCHECK_GT(shape.dims()[i], 0); - last_element_byte_offset += (stride * (shape.dims()[i] - 1)); + int64_t product; + if (__builtin_mul_overflow(stride, shape.dims()[i] - 1, &product)) { + return absl::InvalidArgumentError(absl::StrCat( + "byte_stride[", i, "] * (dim[", i, + "] - 1) overflows: stride=", stride, " dim=", shape.dims()[i])); + } + last_element_byte_offset += static_cast(product); } } return ArrayMemRegion(mem_region_start, last_element_byte_offset + byte_size); diff --git a/third_party/xla/xla/python/ifrt_proxy/common/array_util_test.cc b/third_party/xla/xla/python/ifrt_proxy/common/array_util_test.cc index 490992379f2c63..011991663a4e90 100644 --- a/third_party/xla/xla/python/ifrt_proxy/common/array_util_test.cc +++ b/third_party/xla/xla/python/ifrt_proxy/common/array_util_test.cc @@ -137,7 +137,18 @@ INSTANTIATE_TEST_SUITE_P( TC{"SmallerByteStrideThanDataType", kS32, {5, 5}, Strides({1, 1})}, TC{"ByteStrideIndivisibleByDataType", kS32, {5, 5}, Strides({7, 7})}, // Bad arguments - TC{"NegativeShapeDimension", kS32, {-5, -5}, Strides({20, 4})}), + TC{"NegativeShapeDimension", kS32, {-5, -5}, Strides({20, 4})}, + // Stride * (dim - 1) overflows int64_t — must be rejected, not silently + // wrap to zero (which would bypass the size check in + // FromMinimalMemRegion). + TC{"ByteStrideOverflowSingleDim", + kS32, + {5}, + Strides({INT64_C(4611686018427387904)})}, // 2^62 * 4 = 2^64 wraps + TC{"ByteStrideOverflowMultiDim", + kS32, + {5, 3}, + Strides({INT64_C(4611686018427387904), 4})}), testing::PrintToStringParamName()); TEST_P(ArrayMemRegionFailure, TestCase) { const TC tc = GetParam(); diff --git a/third_party/xla/xla/service/hlo.proto b/third_party/xla/xla/service/hlo.proto index b2df4f33b6ec37..167b2f779d64fc 100644 --- a/third_party/xla/xla/service/hlo.proto +++ b/third_party/xla/xla/service/hlo.proto @@ -128,7 +128,7 @@ enum CustomCallApiVersion { } // Serialization of HloInstruction. -// Next ID: 102 +// Next ID: 103 message HloInstructionProto { reserved 10; reserved "parameter_name"; @@ -439,6 +439,9 @@ message HloInstructionProto { // Convolution block scaling config. BlockScalingConfig block_scaling_config = 101; + + // Comparison order (TOTAL, PARTIAL). + string comparison_order = 102; } // Serialization of HloComputation. diff --git a/third_party/xla/xla/service/hlo_verifier.cc b/third_party/xla/xla/service/hlo_verifier.cc index d79888ecabdbfc..a384adb659d3ba 100644 --- a/third_party/xla/xla/service/hlo_verifier.cc +++ b/third_party/xla/xla/service/hlo_verifier.cc @@ -3945,28 +3945,18 @@ absl::Status CheckElementwiseInstruction(HloInstruction* instruction) { } if (auto* comparison = DynCast(instruction)) { - const Shape& operand_shape = comparison->operand(1)->shape(); + const Shape& operand_shape = comparison->operand(0)->shape(); PrimitiveType operand_element_type = operand_shape.element_type(); - Comparison::Type default_comparison_type = - Comparison::DefaultComparisonType(operand_element_type); - if (primitive_util::IsFloatingPointType(operand_element_type)) { - if (comparison->type() != Comparison::Type::kFloat && - comparison->type() != Comparison::Type::kFloatTotalOrder) { + if (primitive_util::IsIntegralType(operand_element_type) || + operand_element_type == PRED) { + if (comparison->order() != ComparisonOrder::kTotal) { return FailedPrecondition( - "Expected comparison type %s or %s.\n" - "actual: %s\noperand: %s\n", - ComparisonTypeToString(Comparison::Type::kFloat), - ComparisonTypeToString(Comparison::Type::kFloatTotalOrder), - ComparisonTypeToString(comparison->type()), + "Expected comparison order %s for integral/pred operand, but got " + "%s.\noperand: %s\n", + ComparisonOrderToShortString(ComparisonOrder::kTotal), + ComparisonOrderToShortString(comparison->order()), ShapeUtil::HumanString(operand_shape)); } - } else if (comparison->type() != default_comparison_type) { - return FailedPrecondition( - "Expected comparison type %s.\n" - "actual: %s\noperand: %s\n", - ComparisonTypeToString(default_comparison_type), - ComparisonTypeToString(comparison->type()), - ShapeUtil::HumanString(operand_shape)); } } return absl::OkStatus(); diff --git a/third_party/xla/xla/service/hlo_verifier_test.cc b/third_party/xla/xla/service/hlo_verifier_test.cc index 4981119fcfea9e..60edc3cae66153 100644 --- a/third_party/xla/xla/service/hlo_verifier_test.cc +++ b/third_party/xla/xla/service/hlo_verifier_test.cc @@ -3782,69 +3782,58 @@ TEST_F(HloVerifierTest, CollectivePermuteDoneNoCollectivePermuteStart) { "needs to be collective-permute-start, found tuple")); } -TEST_F(HloVerifierTest, ComparisonTypeFloat) { +TEST_F(HloVerifierTest, ComparisonOrderSigned) { const char* const hlo_string = R"( HloModule Module - ENTRY RngOperandElementTypesNotMatch { - p0 = f32[] parameter(0) - ROOT cmp = pred[] compare(f32[] p0, f32[] p0), direction=LT, type=UNSIGNED - } - )"; - ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_string)); - - auto status = verifier().Run(module.get()).status(); - ASSERT_FALSE(status.ok()); - EXPECT_THAT(status.message(), - HasSubstr("Expected comparison type FLOAT or TOTALORDER")); -} - -TEST_F(HloVerifierTest, ComparisonTypeSigned) { - const char* const hlo_string = R"( - HloModule Module - - ENTRY RngOperandElementTypesNotMatch { + ENTRY CompareSignedPartial { p0 = s32[] parameter(0) - ROOT cmp = pred[] compare(s32[] p0, s32[] p0), direction=LT, type=UNSIGNED + ROOT cmp = pred[] compare(s32[] p0, s32[] p0), direction=LT, order=PARTIAL } )"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_string)); auto status = verifier().Run(module.get()).status(); ASSERT_FALSE(status.ok()); - EXPECT_THAT(status.message(), HasSubstr("Expected comparison type SIGNED")); + EXPECT_THAT(status.message(), + HasSubstr("Expected comparison order TOTAL for integral/pred " + "operand, but got PARTIAL")); } -TEST_F(HloVerifierTest, ComparisonTypeUnsigned) { +TEST_F(HloVerifierTest, ComparisonOrderUnsigned) { const char* const hlo_string = R"( HloModule Module - ENTRY RngOperandElementTypesNotMatch { + ENTRY CompareUnsignedPartial { p0 = u32[] parameter(0) - ROOT cmp = pred[] compare(u32[] p0, u32[] p0), direction=LT, type=SIGNED + ROOT cmp = pred[] compare(u32[] p0, u32[] p0), direction=LT, order=PARTIAL } )"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_string)); auto status = verifier().Run(module.get()).status(); ASSERT_FALSE(status.ok()); - EXPECT_THAT(status.message(), HasSubstr("Expected comparison type UNSIGNED")); + EXPECT_THAT(status.message(), + HasSubstr("Expected comparison order TOTAL for integral/pred " + "operand, but got PARTIAL")); } -TEST_F(HloVerifierTest, ComparisonTypePred) { +TEST_F(HloVerifierTest, ComparisonOrderPred) { const char* const hlo_string = R"( HloModule Module - ENTRY RngOperandElementTypesNotMatch { + ENTRY ComparePredPartial { p0 = pred[] parameter(0) - ROOT cmp = pred[] compare(pred[] p0, pred[] p0), direction=LT, type=SIGNED + ROOT cmp = pred[] compare(pred[] p0, pred[] p0), direction=LT, order=PARTIAL } )"; ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_string)); auto status = verifier().Run(module.get()).status(); ASSERT_FALSE(status.ok()); - EXPECT_THAT(status.message(), HasSubstr("Expected comparison type UNSIGNED")); + EXPECT_THAT(status.message(), + HasSubstr("Expected comparison order TOTAL for integral/pred " + "operand, but got PARTIAL")); } TEST_F(HloVerifierTest, UseGlobalDeviceIdsEmptyReplicaGroup) { diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc index 3d2053b44eadc9..0366d458356da7 100644 --- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc +++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc @@ -260,7 +260,7 @@ void setFuncManualAxesRecursively( funcOp.setAllResultAttrs(newResultAttrs); // Walk in preorder of blocks in order to stop walks on manual computations. - funcOp->walk([&](Operation* op) { + funcOp->walk([&](Operation* op) { return setManualAxes(op, manualAxes, meshOrRef, symbolTable, parentManualCompAxes); }); diff --git a/third_party/xla/xla/shape_util.cc b/third_party/xla/xla/shape_util.cc index 4db4a138d7b5b9..668ff5abe239b8 100644 --- a/third_party/xla/xla/shape_util.cc +++ b/third_party/xla/xla/shape_util.cc @@ -1086,11 +1086,7 @@ Shape ShapeUtil::PrependMajorDimension(int64_t bound, Shape shape) { if (subshape.is_dynamic()) { size += sizeof(DynamicSizeType) * subshape.dimensions().size(); } - if (subshape.element_type() == PRED) { - // PRED is packed 8 elements per byte. - size += CeilOfRatio(ElementsIn(subshape), 8); - } else if (primitive_util::IsSubByteNonPredType( - subshape.element_type())) { + if (primitive_util::IsSubByteNonPredType(subshape.element_type())) { // 4-bit types are packed 2 elements per byte. size += CeilOfRatio( ElementsIn(subshape), diff --git a/third_party/xla/xla/stream_executor/command_buffer.h b/third_party/xla/xla/stream_executor/command_buffer.h index e76b041cfb0b42..5e3b34410226ad 100644 --- a/third_party/xla/xla/stream_executor/command_buffer.h +++ b/third_party/xla/xla/stream_executor/command_buffer.h @@ -232,6 +232,26 @@ class CommandBuffer { const DeviceAddressBase& src, uint64_t size) = 0; + // Creates a device-to-host memory copy. + virtual absl::StatusOr CreateMemcpyD2H( + void* dst, const DeviceAddressBase& src, uint64_t size, + absl::Span dependencies) = 0; + + // Updates a device-to-host memory copy. + virtual absl::Status UpdateMemcpyD2H(const Command* command, void* dst, + const DeviceAddressBase& src, + uint64_t size) = 0; + + // Creates a host-to-device memory copy. + virtual absl::StatusOr CreateMemcpyH2D( + DeviceAddressBase* dst, const void* src, uint64_t size, + absl::Span dependencies) = 0; + + // Updates a host-to-device memory copy. + virtual absl::Status UpdateMemcpyH2D(const Command* command, + DeviceAddressBase* dst, const void* src, + uint64_t size) = 0; + // Creates a memset command. virtual absl::StatusOr CreateMemset( DeviceAddressBase* dst, BitPattern bit_pattern, size_t num_elements, diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc index 8e0c7173015ce1..29e7dc349269f9 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc +++ b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc @@ -446,6 +446,107 @@ absl::Status CudaCommandBuffer::UpdateMemcpyD2DNode( "Failed to set memcpy d2d node params"); } +absl::StatusOr CudaCommandBuffer::CreateMemcpyD2HNode( + absl::Span dependencies, void* destination, + DeviceAddressBase source, uint64_t size) { + VLOG(2) << "Add memcpy d2h node to a graph " << graph_ + << "; dst: " << destination << "; src: " << source.opaque() + << "; size: " << size << "; context: " << cuda_context_->context() + << "; deps(" << dependencies.size() + << "): " << FormatGraphNodeHandles(dependencies); + + CUDA_MEMCPY3D params{}; + params.srcMemoryType = CU_MEMORYTYPE_DEVICE; + params.srcDevice = AsDevicePtr(source); + params.dstMemoryType = CU_MEMORYTYPE_HOST; + params.dstHost = destination; + params.WidthInBytes = size; + params.Height = 1; + params.Depth = 1; + + std::vector deps = ToCudaGraphHandles(dependencies); + + CUgraphNode node_handle = nullptr; + ABSL_RETURN_IF_ERROR(cuda::ToStatus( + cuGraphAddMemcpyNode(&node_handle, graph_, deps.data(), deps.size(), + ¶ms, cuda_context_->context()), + "Failed to add memcpy d2h node to a CUDA graph")); + return FromCudaGraphHandle(node_handle); +} + +absl::Status CudaCommandBuffer::UpdateMemcpyD2HNode(GraphNodeHandle node_handle, + void* destination, + DeviceAddressBase source, + uint64_t size) { + VLOG(2) << "Set memcpy d2h node params " << node_handle + << " in graph executable " << graph_exec() << "; dst: " << destination + << "; src: " << source.opaque() << "; size: " << size + << "; context: " << cuda_context_->context(); + + CUDA_MEMCPY3D params{}; + params.srcMemoryType = CU_MEMORYTYPE_DEVICE; + params.srcDevice = AsDevicePtr(source); + params.dstMemoryType = CU_MEMORYTYPE_HOST; + params.dstHost = destination; + params.WidthInBytes = size; + params.Height = 1; + params.Depth = 1; + return cuda::ToStatus(cuGraphExecMemcpyNodeSetParams( + graph_exec(), ToCudaGraphHandle(node_handle), + ¶ms, cuda_context_->context()), + "Failed to set memcpy d2h node params"); +} + +absl::StatusOr CudaCommandBuffer::CreateMemcpyH2DNode( + absl::Span dependencies, + DeviceAddressBase destination, const void* source, uint64_t size) { + VLOG(2) << "Add memcpy h2d node to a graph " << graph_ + << "; dst: " << destination.opaque() << "; src: " << source + << "; size: " << size << "; context: " << cuda_context_->context() + << "; deps(" << dependencies.size() + << "): " << FormatGraphNodeHandles(dependencies); + + CUDA_MEMCPY3D params{}; + params.srcMemoryType = CU_MEMORYTYPE_HOST; + params.srcHost = source; + params.dstMemoryType = CU_MEMORYTYPE_DEVICE; + params.dstDevice = AsDevicePtr(destination); + params.WidthInBytes = size; + params.Height = 1; + params.Depth = 1; + + std::vector deps = ToCudaGraphHandles(dependencies); + + CUgraphNode node_handle = nullptr; + ABSL_RETURN_IF_ERROR(cuda::ToStatus( + cuGraphAddMemcpyNode(&node_handle, graph_, deps.data(), deps.size(), + ¶ms, cuda_context_->context()), + "Failed to add memcpy h2d node to a CUDA graph")); + return FromCudaGraphHandle(node_handle); +} + +absl::Status CudaCommandBuffer::UpdateMemcpyH2DNode( + GraphNodeHandle node_handle, DeviceAddressBase destination, + const void* source, uint64_t size) { + VLOG(2) << "Set memcpy h2d node params " << node_handle + << " in graph executable " << graph_exec() + << "; dst: " << destination.opaque() << "; src: " << source + << "; size: " << size << "; context: " << cuda_context_->context(); + + CUDA_MEMCPY3D params{}; + params.srcMemoryType = CU_MEMORYTYPE_HOST; + params.srcHost = source; + params.dstMemoryType = CU_MEMORYTYPE_DEVICE; + params.dstDevice = AsDevicePtr(destination); + params.WidthInBytes = size; + params.Height = 1; + params.Depth = 1; + return cuda::ToStatus(cuGraphExecMemcpyNodeSetParams( + graph_exec(), ToCudaGraphHandle(node_handle), + ¶ms, cuda_context_->context()), + "Failed to set memcpy h2d node params"); +} + absl::Status CudaCommandBuffer::PopulateDnnGraphNode( dnn::DnnGraph& dnn_graph, Stream& stream, absl::Span operands) { diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h index f728afcd29c2a5..1e73aff2c0557d 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h +++ b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h @@ -125,6 +125,23 @@ class CudaCommandBuffer final : public GpuCommandBuffer { DeviceAddressBase source, uint64_t size) override; + absl::StatusOr CreateMemcpyD2HNode( + absl::Span dependencies, void* destination, + DeviceAddressBase source, uint64_t size) override; + + absl::Status UpdateMemcpyD2HNode(GraphNodeHandle node_handle, + void* destination, DeviceAddressBase source, + uint64_t size) override; + + absl::StatusOr CreateMemcpyH2DNode( + absl::Span dependencies, + DeviceAddressBase destination, const void* source, + uint64_t size) override; + + absl::Status UpdateMemcpyH2DNode(GraphNodeHandle node_handle, + DeviceAddressBase destination, + const void* source, uint64_t size) override; + absl::Status PopulateDnnGraphNode( dnn::DnnGraph&, Stream&, absl::Span operands) override; diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_test.cc b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_test.cc index 87068826a6739d..d577f73d484c1b 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_test.cc +++ b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_test.cc @@ -304,5 +304,59 @@ TEST(CudaCommandBufferTest, LaunchClusterKernelWithClusterDimsSucceeds) { ASSERT_OK(stream->BlockHostUntilDone()); } +TEST(CudaCommandBufferTest, MemcpyH2D2H) { + Platform* platform = CudaPlatform(); + ASSERT_OK_AND_ASSIGN(StreamExecutor * executor, + platform->ExecutorForDevice(0)); + ASSERT_OK_AND_ASSIGN(std::unique_ptr stream, + executor->CreateStream()); + ASSERT_OK_AND_ASSIGN(std::unique_ptr cmd_buffer, + executor->CreateCommandBuffer(primary)); + DeviceAddress device_buf = executor->AllocateArray(1); + + int32_t src = 123; + ASSERT_OK_AND_ASSIGN( + const CommandBuffer::Command* h2d_cmd, + cmd_buffer->CreateMemcpyH2D(&device_buf, &src, sizeof(int32_t), {})); + ASSERT_NE(h2d_cmd, nullptr); + + int32_t dst = 0; + ASSERT_OK_AND_ASSIGN(const CommandBuffer::Command* d2h_cmd, + cmd_buffer->CreateMemcpyD2H(&dst, device_buf, + sizeof(int32_t), {h2d_cmd})); + ASSERT_NE(d2h_cmd, nullptr); + + ASSERT_OK(cmd_buffer->Finalize()); + ASSERT_OK(cmd_buffer->Submit(stream.get())); + ASSERT_OK(stream->BlockHostUntilDone()); + EXPECT_EQ(dst, 123); + + int32_t src2 = 456; + int32_t dst2 = 0; + ASSERT_OK(cmd_buffer->Update()); + ASSERT_OK(cmd_buffer->UpdateMemcpyH2D(h2d_cmd, &device_buf, &src2, + sizeof(int32_t))); + ASSERT_OK( + cmd_buffer->UpdateMemcpyD2H(d2h_cmd, &dst2, device_buf, sizeof(int32_t))); + ASSERT_OK(cmd_buffer->Finalize()); + ASSERT_OK(cmd_buffer->Submit(stream.get())); + ASSERT_OK(stream->BlockHostUntilDone()); + EXPECT_EQ(dst, 123); + EXPECT_EQ(dst2, 456); + + DeviceAddress device_buf2 = executor->AllocateArray(1); + int32_t src3 = 789; + int32_t dst3 = 0; + ASSERT_OK(cmd_buffer->Update()); + ASSERT_OK(cmd_buffer->UpdateMemcpyH2D(h2d_cmd, &device_buf2, &src3, + sizeof(int32_t))); + ASSERT_OK(cmd_buffer->UpdateMemcpyD2H(d2h_cmd, &dst3, device_buf2, + sizeof(int32_t))); + ASSERT_OK(cmd_buffer->Finalize()); + ASSERT_OK(cmd_buffer->Submit(stream.get())); + ASSERT_OK(stream->BlockHostUntilDone()); + EXPECT_EQ(dst2, 456); + EXPECT_EQ(dst3, 789); +} } // namespace } // namespace stream_executor::cuda diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc index bfe34fcad8ef10..ca3ace26d0aee6 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc +++ b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc @@ -638,11 +638,25 @@ absl::StatusOr GetDevicePcieBandwidth(nvmlDevice_t nvml_device) { return lane_speed * link_width; } +absl::StatusOr GetNvLinkCount(nvmlDevice_t nvml_device) { + nvmlFieldValue_t field_value = {}; + field_value.fieldId = NVML_FI_DEV_NVLINK_LINK_COUNT; + ABSL_RETURN_IF_ERROR( + ToStatus(nvmlDeviceGetFieldValues(nvml_device, 1, &field_value))); + ABSL_RETURN_IF_ERROR(ToStatus(field_value.nvmlReturn)); + if (field_value.valueType != NVML_VALUE_TYPE_UNSIGNED_INT) { + return absl::InternalError( + absl::StrFormat("Unexpected NVLink count value type: %d", + static_cast(field_value.valueType))); + } + return field_value.value.uiVal; +} + absl::StatusOr GetNumberOfActiveP2PNvlinks(nvmlDevice_t nvml_device) { int p2p_links = 0; - constexpr int kBlackwellNvLinkCount = 18; - for (unsigned int i = 0; i < kBlackwellNvLinkCount; i++) { + ABSL_ASSIGN_OR_RETURN(unsigned int nvlink_count, GetNvLinkCount(nvml_device)); + for (unsigned int i = 0; i < nvlink_count; i++) { nvmlEnableState_t is_active = NVML_FEATURE_DISABLED; nvmlReturn_t result = nvmlDeviceGetNvLinkState(nvml_device, i, &is_active); if (result == NVML_ERROR_NOT_SUPPORTED) { @@ -1926,8 +1940,8 @@ absl::StatusOr CudaExecutor::GetInterconnectStatus() const { // 2. NVLink Status (as a proxy for IMEX/NVLink health) absl::StrAppend(&status_msg, "NVLinks: "); bool first = true; - constexpr int kMaxNvLinks = 32; - for (unsigned int i = 0; i < kMaxNvLinks; ++i) { + ABSL_ASSIGN_OR_RETURN(unsigned int nvlink_count, GetNvLinkCount(nvml_device)); + for (unsigned int i = 0; i < nvlink_count; ++i) { nvmlEnableState_t isActive; nvmlReturn_t r = nvmlDeviceGetNvLinkState(nvml_device, i, &isActive); if (r == NVML_ERROR_INVALID_ARGUMENT) { diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_executor_test.cc b/third_party/xla/xla/stream_executor/cuda/cuda_executor_test.cc index 823863373a073b..4847321b0bedd4 100644 --- a/third_party/xla/xla/stream_executor/cuda/cuda_executor_test.cc +++ b/third_party/xla/xla/stream_executor/cuda/cuda_executor_test.cc @@ -79,7 +79,14 @@ TEST(CudaExecutorTest, CreateDeviceDescription) { DeviceInterconnectInfo info = result->device_interconnect_info(); if (result->cuda_compute_capability().IsAtLeastHopper() && info.active_links) { - EXPECT_GE(info.active_links, 18); + const auto cc = result->cuda_compute_capability(); + if (cc.major == 10 && cc.minor == 7) { + EXPECT_EQ(info.active_links, 36); + } else if (cc.major == 10 && (cc.minor == 0 || cc.minor == 3)) { + EXPECT_EQ(info.active_links, 18); + } else { + EXPECT_GE(info.active_links, 18); + } // nvmlDeviceGetGpuFabricInfoV is only available in driver r545+ if (result->kernel_mode_driver_version().major_version() >= 545) { EXPECT_THAT(info.clique_id, Not(IsEmpty())); diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc index 983c80db9a7c25..8719039f1d67f7 100644 --- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc +++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.cc @@ -292,6 +292,47 @@ absl::Status GpuCommandBuffer::UpdateMemcpyD2D(const Command* command, return UpdateMemcpyD2DNode(gpu_command->handle, *dst, src, size); } +absl::StatusOr GpuCommandBuffer::CreateMemcpyD2H( + void* dst, const DeviceAddressBase& src, uint64_t size, + absl::Span dependencies) { + ABSL_RETURN_IF_ERROR(CheckInState(State::kCreate)); + + ABSL_ASSIGN_OR_RETURN(GraphNodeHandle handle, + CreateMemcpyD2HNode(ToGraphNodeDependencies(dependencies), + dst, src, size)); + + return AppendCommand(GpuCommand{handle}); +} + +absl::Status GpuCommandBuffer::UpdateMemcpyD2H(const Command* command, + void* dst, + const DeviceAddressBase& src, + uint64_t size) { + ABSL_RETURN_IF_ERROR(CheckInState(State::kUpdate)); + auto* gpu_command = absl::down_cast(command); + return UpdateMemcpyD2HNode(gpu_command->handle, dst, src, size); +} + +absl::StatusOr GpuCommandBuffer::CreateMemcpyH2D( + DeviceAddressBase* dst, const void* src, uint64_t size, + absl::Span dependencies) { + ABSL_RETURN_IF_ERROR(CheckInState(State::kCreate)); + + ABSL_ASSIGN_OR_RETURN(GraphNodeHandle handle, + CreateMemcpyH2DNode(ToGraphNodeDependencies(dependencies), + *dst, src, size)); + + return AppendCommand(GpuCommand{handle}); +} + +absl::Status GpuCommandBuffer::UpdateMemcpyH2D(const Command* command, + DeviceAddressBase* dst, + const void* src, uint64_t size) { + ABSL_RETURN_IF_ERROR(CheckInState(State::kUpdate)); + auto* gpu_command = absl::down_cast(command); + return UpdateMemcpyH2DNode(gpu_command->handle, *dst, src, size); +} + absl::StatusOr GpuCommandBuffer::CreateMemset( DeviceAddressBase* dst, BitPattern bit_pattern, size_t num_elements, absl::Span dependencies) { diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h index c6c4be5bb3f248..3bfaaa70999d58 100644 --- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h +++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer.h @@ -164,6 +164,21 @@ class GpuCommandBuffer : public CommandBuffer { const DeviceAddressBase& src, uint64_t size) override; + absl::StatusOr CreateMemcpyD2H( + void* dst, const DeviceAddressBase& src, uint64_t size, + absl::Span dependencies) override; + + absl::Status UpdateMemcpyD2H(const Command* command, void* dst, + const DeviceAddressBase& src, + uint64_t size) override; + + absl::StatusOr CreateMemcpyH2D( + DeviceAddressBase* dst, const void* src, uint64_t size, + absl::Span dependencies) override; + + absl::Status UpdateMemcpyH2D(const Command* command, DeviceAddressBase* dst, + const void* src, uint64_t size) override; + absl::StatusOr CreateMemset( DeviceAddressBase* dst, BitPattern bit_pattern, size_t num_elements, absl::Span dependencies) override; @@ -348,6 +363,24 @@ class GpuCommandBuffer : public CommandBuffer { DeviceAddressBase source, uint64_t size) = 0; + virtual absl::StatusOr CreateMemcpyD2HNode( + absl::Span dependencies, void* destination, + DeviceAddressBase source, uint64_t size) = 0; + + virtual absl::Status UpdateMemcpyD2HNode(GraphNodeHandle node_handle, + void* destination, + DeviceAddressBase source, + uint64_t size) = 0; + + virtual absl::StatusOr CreateMemcpyH2DNode( + absl::Span dependencies, + DeviceAddressBase destination, const void* source, uint64_t size) = 0; + + virtual absl::Status UpdateMemcpyH2DNode(GraphNodeHandle node_handle, + DeviceAddressBase destination, + const void* source, + uint64_t size) = 0; + virtual absl::Status PopulateDnnGraphNode( dnn::DnnGraph&, Stream&, absl::Span operands) = 0; diff --git a/third_party/xla/xla/stream_executor/mock_command_buffer.h b/third_party/xla/xla/stream_executor/mock_command_buffer.h index 984256d34c7178..3dfe4fe7d35d53 100644 --- a/third_party/xla/xla/stream_executor/mock_command_buffer.h +++ b/third_party/xla/xla/stream_executor/mock_command_buffer.h @@ -84,6 +84,22 @@ class MockCommandBuffer : public CommandBuffer { (const Command* command, DeviceAddressBase* dst, const DeviceAddressBase& src, uint64_t size), (override)); + MOCK_METHOD(absl::StatusOr, CreateMemcpyD2H, + (void* dst, const DeviceAddressBase& src, uint64_t size, + absl::Span dependencies), + (override)); + MOCK_METHOD(absl::Status, UpdateMemcpyD2H, + (const Command* command, void* dst, const DeviceAddressBase& src, + uint64_t size), + (override)); + MOCK_METHOD(absl::StatusOr, CreateMemcpyH2D, + (DeviceAddressBase * dst, const void* src, uint64_t size, + absl::Span dependencies), + (override)); + MOCK_METHOD(absl::Status, UpdateMemcpyH2D, + (const Command* command, DeviceAddressBase* dst, const void* src, + uint64_t size), + (override)); MOCK_METHOD(absl::StatusOr, CreateMemset, (DeviceAddressBase * dst, BitPattern bit_pattern, size_t num_elements, diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc index cc5ec78e7742d2..b81487137e927d 100644 --- a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc +++ b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc @@ -217,6 +217,71 @@ absl::Status RocmCommandBuffer::UpdateMemcpyD2DNode( "Failed to set memcpy d2d node params"); } +absl::StatusOr RocmCommandBuffer::CreateMemcpyD2HNode( + absl::Span dependencies, void* destination, + DeviceAddressBase source, uint64_t size) { + VLOG(2) << "Add memcpy d2h node to a graph " << graph_ + << "; dst: " << destination << "; src: " << source.opaque() + << "; size: " << size << "; deps: " << dependencies.size(); + + std::vector deps = ToHipGraphHandles(dependencies); + + hipGraphNode_t node_handle = nullptr; + ABSL_RETURN_IF_ERROR(ToStatus( + hipGraphAddMemcpyNode1D(&node_handle, graph_, deps.data(), deps.size(), + destination, AsDevicePtr(source), size, + hipMemcpyDeviceToHost), + "Failed to add memcpy d2h node to a HIP graph")); + return FromHipGraphHandle(node_handle); +} + +absl::Status RocmCommandBuffer::UpdateMemcpyD2HNode(GraphNodeHandle node_handle, + void* destination, + DeviceAddressBase source, + uint64_t size) { + VLOG(2) << "Set memcpy d2h node params " << node_handle + << " in graph executable " << exec_ << "; dst: " << destination + << "; src: " << source.opaque() << "; size: " << size; + + return ToStatus(hipGraphExecMemcpyNodeSetParams1D( + exec_, ToHipGraphHandle(node_handle), destination, + AsDevicePtr(source), size, hipMemcpyDeviceToHost), + "Failed to set memcpy d2h node params"); +} + +absl::StatusOr RocmCommandBuffer::CreateMemcpyH2DNode( + absl::Span dependencies, + DeviceAddressBase destination, const void* source, uint64_t size) { + VLOG(2) << "Add memcpy h2d node to a graph " << graph_ + << "; dst: " << destination.opaque() << "; src: " << source + << "; size: " << size << "; deps: " << dependencies.size(); + + std::vector deps = ToHipGraphHandles(dependencies); + + hipGraphNode_t node_handle = nullptr; + ABSL_RETURN_IF_ERROR( + ToStatus(hipGraphAddMemcpyNode1D(&node_handle, graph_, deps.data(), + deps.size(), AsDevicePtr(destination), + source, size, hipMemcpyHostToDevice), + "Failed to add memcpy h2d node to a HIP graph")); + return FromHipGraphHandle(node_handle); +} + +absl::Status RocmCommandBuffer::UpdateMemcpyH2DNode( + GraphNodeHandle node_handle, DeviceAddressBase destination, + const void* source, uint64_t size) { + VLOG(2) << "Set memcpy h2d node params " << node_handle + << " in graph executable " << exec_ + << "; dst: " << destination.opaque() << "; src: " << source + << "; size: " << size; + + return ToStatus( + hipGraphExecMemcpyNodeSetParams1D(exec_, ToHipGraphHandle(node_handle), + AsDevicePtr(destination), source, size, + hipMemcpyHostToDevice), + "Failed to set memcpy h2d node params"); +} + absl::StatusOr RocmCommandBuffer::CreateClonedChildNode( absl::Span dependencies, const CommandBuffer& nested) { diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h index ba6e5d2c6f43cc..876be379f02bfd 100644 --- a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h +++ b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.h @@ -107,6 +107,23 @@ class RocmCommandBuffer : public GpuCommandBuffer { DeviceAddressBase source, uint64_t size) override; + absl::StatusOr CreateMemcpyD2HNode( + absl::Span dependencies, void* destination, + DeviceAddressBase source, uint64_t size) override; + + absl::Status UpdateMemcpyD2HNode(GraphNodeHandle node_handle, + void* destination, DeviceAddressBase source, + uint64_t size) override; + + absl::StatusOr CreateMemcpyH2DNode( + absl::Span dependencies, + DeviceAddressBase destination, const void* source, + uint64_t size) override; + + absl::Status UpdateMemcpyH2DNode(GraphNodeHandle node_handle, + DeviceAddressBase destination, + const void* source, uint64_t size) override; + absl::Status PopulateDnnGraphNode( dnn::DnnGraph&, Stream&, absl::Span operands) override { diff --git a/third_party/xla/xla/stream_executor/sycl/BUILD b/third_party/xla/xla/stream_executor/sycl/BUILD index 0511a4e3894db6..a36c9b44045de2 100644 --- a/third_party/xla/xla/stream_executor/sycl/BUILD +++ b/third_party/xla/xla/stream_executor/sycl/BUILD @@ -592,11 +592,14 @@ cc_library( "oneapi-only", ], deps = [ + "//xla:shape_util", + "//xla:xla_data_proto_cc", "//xla/tsl/mkl:onednn", "//xla/tsl/platform:errors", "//xla/tsl/platform:logging", "//xla/tsl/util:env_var", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/synchronization", "@tsl//tsl/platform:logging", "@tsl//tsl/platform:str_util", diff --git a/third_party/xla/xla/stream_executor/sycl/gemm_sycl_test.cc b/third_party/xla/xla/stream_executor/sycl/gemm_sycl_test.cc index 578321fb883af8..7ff564ee9c340c 100644 --- a/third_party/xla/xla/stream_executor/sycl/gemm_sycl_test.cc +++ b/third_party/xla/xla/stream_executor/sycl/gemm_sycl_test.cc @@ -33,7 +33,8 @@ class GemmSyclTest : public HloInterpreterReferenceMixin { protected: void TestGemmWithTypeVariations(absl::string_view hlo_template) { std::vector> - type_combinations = {{"f32", "f32"}, {"f16", "f16"}, {"bf16", "bf16"}}; + type_combinations = { + {"f32", "f32"}, {"f16", "f16"}, {"bf16", "bf16"}, {"f64", "f64"}}; for (const auto& type_combination : type_combinations) { VLOG(3) << "Testing type combination: " << std::get<0>(type_combination) diff --git a/third_party/xla/xla/stream_executor/sycl/onednn_util.cc b/third_party/xla/xla/stream_executor/sycl/onednn_util.cc index 340ba64132376e..f08d1e71ed3747 100644 --- a/third_party/xla/xla/stream_executor/sycl/onednn_util.cc +++ b/third_party/xla/xla/stream_executor/sycl/onednn_util.cc @@ -21,6 +21,7 @@ limitations under the License. #include "absl/synchronization/mutex.h" #include "dnnl.hpp" #include "dnnl_sycl.hpp" +#include "xla/primitive_util.h" #include "xla/tsl/platform/errors.h" #include "xla/tsl/platform/logging.h" #include "xla/tsl/util/env_var.h" @@ -79,5 +80,29 @@ dnnl::memory CreateDnnlMemory(const dnnl::memory::desc& md, } } +absl::StatusOr ToOneDnnDataType( + xla::PrimitiveType xla_type) { + switch (xla_type) { + case xla::PrimitiveType::F16: + return dnnl::memory::data_type::f16; + case xla::PrimitiveType::BF16: + return dnnl::memory::data_type::bf16; + case xla::PrimitiveType::F32: + return dnnl::memory::data_type::f32; + case xla::PrimitiveType::F64: + return dnnl::memory::data_type::f64; + case xla::PrimitiveType::S8: + return dnnl::memory::data_type::s8; + case xla::PrimitiveType::U8: + return dnnl::memory::data_type::u8; + case xla::PrimitiveType::S32: + return dnnl::memory::data_type::s32; + default: + return absl::InvalidArgumentError(absl::StrCat( + "Unsupported element type: ", + xla::primitive_util::LowercasePrimitiveTypeName(xla_type))); + } +} + } // namespace sycl } // namespace stream_executor diff --git a/third_party/xla/xla/stream_executor/sycl/onednn_util.h b/third_party/xla/xla/stream_executor/sycl/onednn_util.h index eec74487f5d987..198ae7b18c9095 100644 --- a/third_party/xla/xla/stream_executor/sycl/onednn_util.h +++ b/third_party/xla/xla/stream_executor/sycl/onednn_util.h @@ -17,10 +17,12 @@ limitations under the License. #define XLA_STREAM_EXECUTOR_SYCL_ONEDNN_UTIL_H_ #include "absl/container/flat_hash_map.h" +#include "absl/status/statusor.h" #include "absl/synchronization/mutex.h" #include "dnnl.hpp" #include "dnnl_sycl.hpp" #include "xla/tsl/util/env_var.h" +#include "xla/xla_data.pb.h" #include "tsl/platform/str_util.h" namespace stream_executor { @@ -43,6 +45,10 @@ dnnl::fpmath_mode GetFP32MathMode(); dnnl::memory CreateDnnlMemory(const dnnl::memory::desc& md, const dnnl::engine& engine, void* data_handle = nullptr); + +// Converts XLA primitive type to oneDNN data type. +absl::StatusOr ToOneDnnDataType( + xla::PrimitiveType xla_type); } // namespace sycl } // namespace stream_executor #endif // XLA_STREAM_EXECUTOR_SYCL_ONEDNN_UTIL_H_ diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_matmul_utils.cc b/third_party/xla/xla/stream_executor/sycl/sycl_matmul_utils.cc index fa3cb917d31375..8a14f7753d1daa 100644 --- a/third_party/xla/xla/stream_executor/sycl/sycl_matmul_utils.cc +++ b/third_party/xla/xla/stream_executor/sycl/sycl_matmul_utils.cc @@ -161,6 +161,10 @@ struct PrimitiveTypeToNative { using type = float; }; template <> +struct PrimitiveTypeToNative { + using type = double; +}; +template <> struct PrimitiveTypeToNative { using type = ::sycl::half; }; @@ -320,32 +324,8 @@ absl::StatusOr ShapeToMemDesc(const xla::Shape& shape) { if (dims.empty()) { return dnnl::memory::desc{}; } - dnnl::memory::data_type dtype; - switch (shape.element_type()) { - case xla::PrimitiveType::F16: - dtype = dnnl::memory::data_type::f16; - break; - case xla::PrimitiveType::BF16: - dtype = dnnl::memory::data_type::bf16; - break; - case xla::PrimitiveType::F32: - dtype = dnnl::memory::data_type::f32; - break; - case xla::PrimitiveType::S8: - dtype = dnnl::memory::data_type::s8; - break; - case xla::PrimitiveType::U8: - dtype = dnnl::memory::data_type::u8; - break; - case xla::PrimitiveType::S32: - dtype = dnnl::memory::data_type::s32; - break; - default: - return absl::InvalidArgumentError( - absl::StrFormat("Unsupported element type: %s", - xla::primitive_util::LowercasePrimitiveTypeName( - shape.element_type()))); - } + ABSL_ASSIGN_OR_RETURN(dnnl::memory::data_type dtype, + sycl::ToOneDnnDataType(shape.element_type())); return dnnl::memory::desc(dims, dtype, strides); } @@ -419,70 +399,12 @@ CreateMatMulPrimDescFromGemmConfig( } // Get OneDNN data type from layout - dnnl::memory::data_type lhs_dtype, rhs_dtype, output_dtype; - switch (lhs_layout.dtype) { - case xla::PrimitiveType::F16: - lhs_dtype = dnnl::memory::data_type::f16; - break; - case xla::PrimitiveType::BF16: - lhs_dtype = dnnl::memory::data_type::bf16; - break; - case xla::PrimitiveType::F32: - lhs_dtype = dnnl::memory::data_type::f32; - break; - case xla::PrimitiveType::S8: - lhs_dtype = dnnl::memory::data_type::s8; - break; - case xla::PrimitiveType::S32: - lhs_dtype = dnnl::memory::data_type::s32; - break; - default: - return absl::InvalidArgumentError(absl::StrFormat( - "Unsupported LHS element type: %s", - xla::primitive_util::LowercasePrimitiveTypeName(lhs_layout.dtype))); - } - - switch (rhs_layout.dtype) { - case xla::PrimitiveType::F16: - rhs_dtype = dnnl::memory::data_type::f16; - break; - case xla::PrimitiveType::BF16: - rhs_dtype = dnnl::memory::data_type::bf16; - break; - case xla::PrimitiveType::F32: - rhs_dtype = dnnl::memory::data_type::f32; - break; - case xla::PrimitiveType::S8: - rhs_dtype = dnnl::memory::data_type::s8; - break; - case xla::PrimitiveType::S32: - rhs_dtype = dnnl::memory::data_type::s32; - break; - default: - return absl::InvalidArgumentError(absl::StrFormat( - "Unsupported RHS element type: %s", - xla::primitive_util::LowercasePrimitiveTypeName(rhs_layout.dtype))); - } - - switch (output_layout.dtype) { - case xla::PrimitiveType::F16: - output_dtype = dnnl::memory::data_type::f16; - break; - case xla::PrimitiveType::BF16: - output_dtype = dnnl::memory::data_type::bf16; - break; - case xla::PrimitiveType::F32: - output_dtype = dnnl::memory::data_type::f32; - break; - case xla::PrimitiveType::S32: - output_dtype = dnnl::memory::data_type::s32; - break; - default: - return absl::InvalidArgumentError( - absl::StrFormat("Unsupported output element type: %s", - xla::primitive_util::LowercasePrimitiveTypeName( - output_layout.dtype))); - } + ABSL_ASSIGN_OR_RETURN(dnnl::memory::data_type lhs_dtype, + sycl::ToOneDnnDataType(lhs_layout.dtype)); + ABSL_ASSIGN_OR_RETURN(dnnl::memory::data_type rhs_dtype, + sycl::ToOneDnnDataType(rhs_layout.dtype)); + ABSL_ASSIGN_OR_RETURN(dnnl::memory::data_type output_dtype, + sycl::ToOneDnnDataType(output_layout.dtype)); auto lhs_md = dnnl::memory::desc(lhs_dims, lhs_dtype, lhs_strides); auto rhs_md = dnnl::memory::desc(rhs_dims, rhs_dtype, rhs_strides); @@ -714,6 +636,7 @@ absl::Status RunGemm(const gpu::GemmConfig& config, TYPED_GEMM(xla::BF16, xla::BF16, xla::F32) TYPED_GEMM(xla::F16, xla::F16, xla::F16) TYPED_GEMM(xla::F16, xla::F16, xla::F32) + TYPED_GEMM(xla::F64, xla::F64, xla::F64) TYPED_GEMM(xla::S8, xla::S8, xla::S32) // TODO (intel-tf): Add support for more combinations of input/output types diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD index 27a2b4fdc7f22a..55382cafa7ae96 100644 --- a/third_party/xla/xla/tests/BUILD +++ b/third_party/xla/xla/tests/BUILD @@ -2756,15 +2756,12 @@ xla_test( ":client_library_test_runner_mixin", ":hlo_pjrt_interpreter_reference_mixin", ":hlo_pjrt_test_base", - ":hlo_runner_agnostic_reference_mixin", ":xla_internal_test_main", "//xla:error_spec", "//xla:shape_util", "//xla:xla_data_proto_cc", "//xla/hlo/builder:xla_builder", - "//xla/hlo/evaluator:hlo_evaluator", "//xla/pjrt/interpreter:interpreter_client", - "//xla/service:hlo_runner_pjrt", "//xla/tsl/platform:test", "@com_google_absl//absl/strings:string_view", "@tsl//tsl/platform:ml_dtypes", diff --git a/third_party/xla/xla/tests/bitcast_convert_test.cc b/third_party/xla/xla/tests/bitcast_convert_test.cc index 49ca31d1fc1328..1899c086204d48 100644 --- a/third_party/xla/xla/tests/bitcast_convert_test.cc +++ b/third_party/xla/xla/tests/bitcast_convert_test.cc @@ -23,14 +23,11 @@ limitations under the License. #include "absl/strings/string_view.h" #include "xla/error_spec.h" #include "xla/hlo/builder/xla_builder.h" -#include "xla/hlo/evaluator/hlo_evaluator.h" #include "xla/pjrt/interpreter/interpreter_client.h" -#include "xla/service/hlo_runner_pjrt.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" -#include "xla/tests/hlo_runner_agnostic_reference_mixin.h" #include "xla/tsl/platform/test.h" #include "xla/xla_data.pb.h" #include "tsl/platform/ml_dtypes.h" @@ -216,23 +213,7 @@ ENTRY main { EXPECT_TRUE(RunAndCompare(hlo_string, ErrorSpec{1e-5, 1e-5})); } -template -class HloPjRtInterpreterReferenceMixinNoAot - : public HloRunnerAgnosticReferenceMixin { - protected: - template - explicit HloPjRtInterpreterReferenceMixinNoAot(BaseArgs&&... base_args) - : HloRunnerAgnosticReferenceMixin( - std::make_unique(std::make_unique( - []() { return std::make_unique(); })), - std::forward(base_args)...) {} - ~HloPjRtInterpreterReferenceMixinNoAot() override = default; -}; - -class BitcastConvertNoAotTest - : public HloPjRtInterpreterReferenceMixinNoAot {}; - -TEST_F(BitcastConvertNoAotTest, S8ToPred) { +TEST_F(BitcastConvertTest, S8ToPred) { absl::string_view hlo_string = R"( HloModule bitcast_to_smaller diff --git a/third_party/xla/xla/tools/compare_literals/BUILD b/third_party/xla/xla/tools/compare_literals/BUILD new file mode 100644 index 00000000000000..3dcd7bf72f934a --- /dev/null +++ b/third_party/xla/xla/tools/compare_literals/BUILD @@ -0,0 +1,72 @@ +# Tools for comparing XLA output literals. + +load( + "//xla:xla.default.bzl", + "xla_cc_binary", + "xla_cc_test", +) +load("//xla/tsl/platform:rules_cc.bzl", "cc_library") + +package( + # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"], + default_visibility = ["//xla:internal"], + licenses = ["notice"], +) + +cc_library( + name = "compare_literals_lib", + srcs = [ + "compare_literal_math.cc", + "compare_literal_report.cc", + "compare_literals.cc", + ], + hdrs = [ + "compare_literals.h", + "element_comparator.h", + ], + deps = [ + "//xla:literal", + "//xla:shape_util", + "//xla:xla_data_proto_cc", + "//xla/tsl/platform:env", + "@com_google_absl//absl/algorithm:container", + "@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/strings:str_format", + "@com_google_absl//absl/types:span", + ], +) + +xla_cc_binary( + name = "compare_literals", + srcs = ["compare_literals_main.cc"], + deps = [ + ":compare_literals_lib", + "//xla/tsl/platform:env", + "@com_google_absl//absl/flags:flag", + "@com_google_absl//absl/flags:parse", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:statusor", + "@com_google_absl//absl/strings:string_view", + "@tsl//tsl/platform:platform_port", + ], +) + +xla_cc_test( + name = "compare_literals_test", + srcs = ["compare_literals_test.cc"], + deps = [ + ":compare_literals_lib", + "//xla:literal", + "//xla:literal_util", + "//xla:shape_util", + "//xla:types", + "//xla/tsl/platform:env", + "@com_google_absl//absl/status", + "@com_google_absl//absl/status:status_matchers", + "@com_google_googletest//:gtest_main", + "@tsl//tsl/platform:path", + ], +) diff --git a/third_party/xla/xla/tools/compare_literals/compare_literal_math.cc b/third_party/xla/xla/tools/compare_literals/compare_literal_math.cc new file mode 100644 index 00000000000000..17255b1a9e8534 --- /dev/null +++ b/third_party/xla/xla/tools/compare_literals/compare_literal_math.cc @@ -0,0 +1,279 @@ +/* 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 +#include +#include +#include +#include +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/types/span.h" +#include "xla/tools/compare_literals/compare_literals.h" + +namespace xla::compare_literals { + +std::vector BuildThresholds(double target) { + std::vector thresholds; + thresholds.reserve(35); + + for (int e = -7; e <= 2; ++e) { + double base = std::pow(10.0, e); + for (double m : kDefaultMultipliers) { + thresholds.push_back(m * base); + } + } + thresholds.push_back(1e3); + + if (target > 0.0) { + thresholds.push_back(target); + } + absl::c_sort(thresholds); + thresholds.erase(std::unique(thresholds.begin(), thresholds.end()), + thresholds.end()); + return thresholds; +} + +std::vector CreateDefaultRelBins() { + constexpr double kInfinity = std::numeric_limits::infinity(); + + // Positive boundaries: 0.0, 2e-6, 5e-6, 1e-5, ..., 10.0, +inf + std::vector pos_bounds = {0.0}; + for (int e = -6; e <= 0; ++e) { + double base = std::pow(10.0, e); + for (double m : kDefaultMultipliers) { + if (e == -6 && m == 1.0) { + continue; + } + pos_bounds.push_back(m * base); + } + } + pos_bounds.push_back(10.0); + pos_bounds.push_back(kInfinity); + + std::vector bins; + bins.reserve(2 * pos_bounds.size() - 1); + + // Negative bins: ordered ascending from -inf to 0 + for (int i = static_cast(pos_bounds.size()) - 1; i >= 1; --i) { + bins.push_back({-pos_bounds[i], -pos_bounds[i - 1], 0, false}); + } + + // Exact zero bin + bins.push_back({0.0, 0.0, 0, true}); + + // Positive bins: ordered ascending from 0 to +inf + for (size_t i = 1; i < pos_bounds.size(); ++i) { + bins.push_back({pos_bounds[i - 1], pos_bounds[i], 0, false}); + } + + CHECK(std::isinf(bins.begin()->lower)); + CHECK(std::isinf(bins.rbegin()->upper)); + return bins; +} + +int FindRelBin(double rel_err, absl::Span bins) { + if (rel_err == 0.0) { + for (size_t i = 0; i < bins.size(); ++i) { + if (bins[i].is_exact_zero) { + return static_cast(i); + } + } + } + for (size_t i = 0; i < bins.size(); ++i) { + if (bins[i].is_exact_zero) { + continue; + } + if (rel_err >= bins[i].lower && rel_err < bins[i].upper) { + return static_cast(i); + } + } + if (rel_err < bins.front().upper) { + return 0; + } + return static_cast(bins.size() - 1); +} + +std::vector> ComputeHeatmapMismatchCounts( + const std::vector>& hist_2d, int num_rel_thresh, + int num_abs_thresh) { + std::vector> mismatch_counts( + num_rel_thresh, std::vector(num_abs_thresh, 0)); + + std::vector> suffix( + num_rel_thresh + 2, std::vector(num_abs_thresh + 2, 0)); + + for (int r = num_rel_thresh; r >= 0; --r) { + for (int a = num_abs_thresh; a >= 0; --a) { + suffix[r][a] = hist_2d[r][a] + suffix[r + 1][a] + suffix[r][a + 1] - + suffix[r + 1][a + 1]; + } + } + + for (int r = 0; r < num_rel_thresh; ++r) { + for (int a = 0; a < num_abs_thresh; ++a) { + mismatch_counts[r][a] = suffix[r + 1][a + 1]; + } + } + + return mismatch_counts; +} + +std::optional ComputeSuggestedErrorSpec( + const ErrorHeatmap& heatmap, double max_abs_error, double max_rel_error) { + if (!std::isfinite(max_abs_error) || !std::isfinite(max_rel_error)) { + return std::nullopt; + } + if (heatmap.abs_thresholds.empty() || heatmap.rel_thresholds.empty() || + heatmap.mismatch_counts.empty()) { + return std::nullopt; + } + if (max_abs_error == 0.0 && max_rel_error == 0.0) { + return SuggestedErrorSpec{0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}; + } + + SuggestedErrorSpec spec; + + // Pure absolute bound: smallest threshold >= max_abs_error + auto abs_pure_it = absl::c_lower_bound(heatmap.abs_thresholds, max_abs_error); + spec.pure_abs_bound = (abs_pure_it != heatmap.abs_thresholds.end()) + ? *abs_pure_it + : max_abs_error; + auto m_abs_pure_it = + absl::c_lower_bound(heatmap.abs_thresholds, 2.0 * spec.pure_abs_bound); + spec.margin_pure_abs_bound = (m_abs_pure_it != heatmap.abs_thresholds.end()) + ? *m_abs_pure_it + : 2.0 * spec.pure_abs_bound; + + // When mismatches exist solely against reference values of 0.0, max_rel_error + // remains 0.0. A relative tolerance cannot satisfy differences against + // reference zeros, so a pure relative bound is impossible (+inf) and no + // (a, r) trade-off exists. + if (max_abs_error > 0.0 && max_rel_error == 0.0) { + constexpr double kInf = std::numeric_limits::infinity(); + spec.abs_bound = spec.pure_abs_bound; + spec.rel_bound = 0.0; + spec.margin_abs_bound = spec.margin_pure_abs_bound; + spec.margin_rel_bound = 0.0; + spec.pure_rel_bound = kInf; + spec.margin_pure_rel_bound = kInf; + return spec; + } + + // Pure relative bound: smallest threshold >= max_rel_error + auto rel_pure_it = absl::c_lower_bound(heatmap.rel_thresholds, max_rel_error); + spec.pure_rel_bound = (rel_pure_it != heatmap.rel_thresholds.end()) + ? *rel_pure_it + : max_rel_error; + auto m_rel_pure_it = + absl::c_lower_bound(heatmap.rel_thresholds, 2.0 * spec.pure_rel_bound); + spec.margin_pure_rel_bound = (m_rel_pure_it != heatmap.rel_thresholds.end()) + ? *m_rel_pure_it + : 2.0 * spec.pure_rel_bound; + + // Pareto frontier for balanced (a, r): + // For each abs threshold a, find smallest rel threshold r where + // mismatch_counts[r][a] == 0. + const int num_abs = heatmap.abs_thresholds.size(); + const int num_rel = heatmap.rel_thresholds.size(); + + struct Candidate { + double a_val; + double r_val; + }; + std::vector frontier; + int prev_r = num_rel; + for (int a = 0; a < num_abs; ++a) { + double a_val = heatmap.abs_thresholds[a]; + if (a_val >= spec.pure_abs_bound) { + break; + } + + int found_r = -1; + for (int r = 0; r < num_rel; ++r) { + if (heatmap.mismatch_counts[r][a] == 0) { + found_r = r; + break; + } + } + if (found_r != -1 && found_r < prev_r) { + double r_val = heatmap.rel_thresholds[found_r]; + if (r_val < spec.pure_rel_bound) { + frontier.push_back({a_val, r_val}); + prev_r = found_r; + } + } + } + + if (frontier.empty()) { + spec.abs_bound = spec.pure_abs_bound; + spec.rel_bound = spec.pure_rel_bound; + spec.margin_abs_bound = spec.margin_pure_abs_bound; + spec.margin_rel_bound = spec.margin_pure_rel_bound; + return spec; + } + + if (frontier.size() == 1) { + spec.abs_bound = frontier[0].a_val; + spec.rel_bound = frontier[0].r_val; + } else { + // Find knee in normalized log-space. + // In frontier, a_val increases from front to back, and r_val decreases from + // front to back. + double min_u = std::log10(std::max(1e-15, frontier.front().a_val)); + double max_u = std::log10(std::max(1e-15, frontier.back().a_val)); + double min_v = std::log10(std::max(1e-15, frontier.back().r_val)); + double max_v = std::log10(std::max(1e-15, frontier.front().r_val)); + + double span_u = max_u - min_u; + double span_v = max_v - min_v; + + double best_dist = std::numeric_limits::infinity(); + int best_idx = 0; + + for (size_t i = 0; i < frontier.size(); ++i) { + double u = std::log10(std::max(1e-15, frontier[i].a_val)); + double v = std::log10(std::max(1e-15, frontier[i].r_val)); + double norm_u = span_u > 1e-9 ? (u - min_u) / span_u : 0.0; + double norm_v = span_v > 1e-9 ? (v - min_v) / span_v : 0.0; + double dist = norm_u * norm_u + norm_v * norm_v; + if (dist < best_dist) { + best_dist = dist; + best_idx = i; + } + } + + spec.abs_bound = frontier[best_idx].a_val; + spec.rel_bound = frontier[best_idx].r_val; + } + + auto m_abs_it = + absl::c_lower_bound(heatmap.abs_thresholds, 2.0 * spec.abs_bound); + spec.margin_abs_bound = (m_abs_it != heatmap.abs_thresholds.end()) + ? *m_abs_it + : 2.0 * spec.abs_bound; + + auto m_rel_it = + absl::c_lower_bound(heatmap.rel_thresholds, 2.0 * spec.rel_bound); + spec.margin_rel_bound = (m_rel_it != heatmap.rel_thresholds.end()) + ? *m_rel_it + : 2.0 * spec.rel_bound; + + return spec; +} + +} // namespace xla::compare_literals diff --git a/third_party/xla/xla/tools/compare_literals/compare_literal_report.cc b/third_party/xla/xla/tools/compare_literals/compare_literal_report.cc new file mode 100644 index 00000000000000..b00c88b82dbaa9 --- /dev/null +++ b/third_party/xla/xla/tools/compare_literals/compare_literal_report.cc @@ -0,0 +1,601 @@ +/* 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 +#include +#include +#include +#include +#include +#include + +#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 "xla/tools/compare_literals/compare_literals.h" + +namespace xla::compare_literals { +namespace { + +std::string FormatCompactSci(double v) { + if (std::isnan(v)) { + return "nan"; + } + if (std::isinf(v)) { + return v < 0 ? "-inf" : "inf"; + } + if (v == 0.0) { + return "0"; + } + bool negative = v < 0.0; + if (negative) { + v = -v; + } + + int exp = static_cast(std::floor(std::log10(v) + 1e-9)); + exp = std::clamp(exp, std::numeric_limits::min_exponent10, + std::numeric_limits::max_exponent10); + double divisor = std::pow(10.0, exp); + double mantissa = + (divisor > 0.0) ? std::round((v / divisor) * 1e6) / 1e6 : 0.0; + if (mantissa == 0.0) { + return "0"; + } + if (mantissa >= 10.0) { + mantissa /= 10.0; + exp += 1; + } + + return absl::StrFormat("%s%ge%d", negative ? "-" : "", mantissa, exp); +} + +std::string FormatBar(int64_t count, int64_t max_count, int max_bar_width) { + int bar_len = (max_count > 0) + ? static_cast(std::round(static_cast(count) * + max_bar_width / max_count)) + : 0; + if (count > 0 && bar_len == 0) { + bar_len = 1; + } + return std::string(bar_len, '#'); +} + +struct HeatmapWindow { + int col_start; + int col_end; + int row_start; + int row_end; +}; + +HeatmapWindow ComputeHeatmapWindow( + const std::vector>& mismatch_counts, + int num_cols_total, int num_rows_total, int target_abs_idx, + int target_rel_idx, int max_cols, int max_rows) { + CHECK_GT(num_cols_total, 0); + CHECK_GT(num_rows_total, 0); + const int target_abs = std::clamp(target_abs_idx >= 0 ? target_abs_idx : 0, 0, + num_cols_total - 1); + const int target_rel = std::clamp(target_rel_idx >= 0 ? target_rel_idx : 0, 0, + num_rows_total - 1); + + // Find first column where all rows are zero. + // Because row 0 is the strictest relative threshold, if + // mismatch_counts[0][a] == 0, then mismatch_counts[r][a] == 0 for all r >= 0. + int col_zero = num_cols_total - 1; + for (int a = 0; a < num_cols_total; ++a) { + if (mismatch_counts[0][a] == 0) { + col_zero = a; + break; + } + } + + // Find first row where all columns are zero. + // Because col 0 is the strictest absolute threshold, if + // mismatch_counts[r][0] == 0, then mismatch_counts[r][a] == 0 for all a >= 0. + int row_zero = num_rows_total - 1; + for (int r = 0; r < num_rows_total; ++r) { + if (mismatch_counts[r][0] == 0) { + row_zero = r; + break; + } + } + + int col_end = std::min(num_cols_total - 1, std::max(target_abs, col_zero)); + if (col_end == 0) { + col_end = std::min(num_cols_total - 1, max_cols - 1); + } + int col_start = std::max(0, col_end - max_cols + 1); + if (col_start > target_abs) { + col_start = target_abs; + col_end = std::min(num_cols_total - 1, col_start + max_cols - 1); + } + + int row_end = std::min(num_rows_total - 1, std::max(target_rel, row_zero)); + if (row_end == 0) { + row_end = std::min(num_rows_total - 1, max_rows - 1); + } + int row_start = std::max(0, row_end - max_rows + 1); + if (row_start > target_rel) { + row_start = target_rel; + row_end = std::min(num_rows_total - 1, row_start + max_rows - 1); + } + + return {col_start, col_end, row_start, row_end}; +} + +} // namespace + +std::string RelErrorHistogram::ToString(int max_bar_width) const { + std::string out = + "1D Signed Relative Error Distribution ((actual - expected) / " + "|expected|):\n"; + int64_t max_bin_count = 0; + for (const auto& b : bins) { + max_bin_count = std::max(max_bin_count, b.count); + } + + for (size_t i = 0; i < bins.size(); ++i) { + const auto& b = bins[i]; + if (b.count == 0 && static_cast(i) != median_bin_index && + !b.is_exact_zero) { + continue; + } + + std::string bar = FormatBar(b.count, max_bin_count, max_bar_width); + + std::string markers; + if (static_cast(i) == median_bin_index) { + markers += " <--- median"; + } + if (b.is_exact_zero) { + markers += " <--- exact match (zero)"; + } else if (mean_rel_error >= b.lower && mean_rel_error < b.upper) { + markers += " <--- mean"; + } + + std::string range_str; + if (b.is_exact_zero) { + range_str = " [ 0 ] "; + } else { + char left_bracket = std::isinf(b.lower) ? '(' : '['; + char right_bracket = ')'; + std::string lower_str = + std::isinf(b.lower) ? "-inf" : FormatCompactSci(b.lower); + std::string upper_str = + std::isinf(b.upper) ? "+inf" : FormatCompactSci(b.upper); + range_str = absl::StrFormat("%c%6s, %6s%c", left_bracket, lower_str, + upper_str, right_bracket); + } + + double pct = total_samples > 0 + ? (100.0 * static_cast(b.count) / total_samples) + : 0.0; + std::string pct_str; + if (b.count > 0 && pct < 0.05) { + pct_str = " <0.1%"; + } else { + pct_str = absl::StrFormat("%5.1f%%", pct); + } + + absl::StrAppendFormat(&out, " %2d: %s %8lld (%s) %s%s\n", i, range_str, + b.count, pct_str, bar, markers); + } + + absl::StrAppendFormat( + &out, + " Summary: min = %1.3e | max = %1.3e | mean = %1.3e | std_dev = %1.3e\n", + min_rel_error, max_rel_error, mean_rel_error, std_dev_rel_error); + return out; +} + +std::string RelErrorHistogram::ToMarkdown(int max_bar_width) const { + if (bins.empty() || total_samples == 0) { + return ""; + } + + int64_t max_bin_count = 0; + for (const auto& b : bins) { + max_bin_count = std::max(max_bin_count, b.count); + } + + std::string out; + absl::StrAppend(&out, + "| Markers | Range | Count | Percent | Distribution |\n"); + absl::StrAppend(&out, "| :--- | :--- | :---: | :---: | :--- |\n"); + + for (size_t i = 0; i < bins.size(); ++i) { + const auto& b = bins[i]; + if (b.count == 0 && static_cast(i) != median_bin_index && + !b.is_exact_zero) { + continue; + } + + double pct = 100.0 * static_cast(b.count) / total_samples; + std::string pct_str = + (b.count > 0 && pct < 0.05) ? "<0.1%" : absl::StrFormat("%.1f%%", pct); + + std::vector marker_labels; + if (static_cast(i) == median_bin_index) { + marker_labels.push_back("**Median**"); + } + if (b.is_exact_zero) { + marker_labels.push_back("**Zero**"); + } + if (!b.is_exact_zero && mean_rel_error >= b.lower && + mean_rel_error < b.upper) { + marker_labels.push_back("**Mean**"); + } + std::string marker_str = absl::StrJoin(marker_labels, ", "); + + std::string range_str; + if (b.is_exact_zero) { + range_str = "`[0]`"; + } else if (std::isinf(b.lower)) { + range_str = absl::StrFormat("`(-inf, %s)`", FormatCompactSci(b.upper)); + } else if (std::isinf(b.upper)) { + range_str = absl::StrFormat("`(%s, +inf)`", FormatCompactSci(b.lower)); + } else { + range_str = absl::StrFormat("`[%s, %s)`", FormatCompactSci(b.lower), + FormatCompactSci(b.upper)); + } + + std::string bar = FormatBar(b.count, max_bin_count, max_bar_width); + absl::StrAppendFormat(&out, "| %s | %s | %lld | %s | `%s` |\n", marker_str, + range_str, b.count, pct_str, bar); + } + + return out; +} + +std::string ErrorHeatmap::ToString(bool use_color) const { + if (abs_thresholds.empty() || rel_thresholds.empty()) { + return ""; + } + + const int num_cols_total = static_cast(abs_thresholds.size()); + const int num_rows_total = static_cast(rel_thresholds.size()); + constexpr int kMaxCols = 13; + constexpr int kMaxRows = 25; + + HeatmapWindow win = + ComputeHeatmapWindow(mismatch_counts, num_cols_total, num_rows_total, + target_abs_idx, target_rel_idx, kMaxCols, kMaxRows); + const int col_start = win.col_start; + const int col_end = win.col_end; + const int row_start = win.row_start; + const int row_end = win.row_end; + + std::string out = + "2D Error Heatmap (Percentage of elements failing: abs_diff > X AND " + "rel_diff > Y):\n"; + + // Table header (from largest abs_threshold down to smallest) + absl::StrAppend(&out, " Rel \\ Abs |"); + for (int a = col_end; a >= col_start; --a) { + std::string col_hdr = FormatCompactSci(abs_thresholds[a]); + if (a == target_abs_idx) { + col_hdr = absl::StrCat("*", col_hdr); + } + if (col_hdr.size() > 8) { + col_hdr = col_hdr.substr(0, 8); + } + absl::StrAppendFormat(&out, " %8s |", col_hdr); + } + absl::StrAppend(&out, "\n -----------+"); + for (int a = col_end; a >= col_start; --a) { + absl::StrAppend(&out, "----------+"); + } + absl::StrAppend(&out, "\n"); + + // Rows from smallest rel_threshold up to largest + for (int r = row_start; r <= row_end; ++r) { + std::string row_hdr = FormatCompactSci(rel_thresholds[r]); + if (r == target_rel_idx) { + row_hdr = absl::StrCat("*", row_hdr); + } + if (row_hdr.size() > 10) { + row_hdr = row_hdr.substr(0, 10); + } + absl::StrAppendFormat(&out, " %10s |", row_hdr); + + for (int a = col_end; a >= col_start; --a) { + int64_t count = mismatch_counts[r][a]; + double pct = total_elements > 0 + ? (100.0 * static_cast(count) / total_elements) + : 0.0; + bool is_target = (r == target_rel_idx && a == target_abs_idx); + + std::string val_str; + if (count == 0) { + val_str = "0.0%"; + } else if (pct < 0.05) { + val_str = "<0.1%"; + } else { + val_str = absl::StrFormat("%.1f%%", pct); + } + + std::string cell_str; + if (is_target) { + cell_str = absl::StrFormat("[%s]", val_str); + if (cell_str.size() < 8) { + int left = (8 - static_cast(cell_str.size())) / 2; + int right = 8 - static_cast(cell_str.size()) - left; + cell_str = absl::StrCat(std::string(left, ' '), cell_str, + std::string(right, ' ')); + } + } else { + int pad = 8 - static_cast(val_str.size()); + int left = pad / 2; + int right = pad - left; + cell_str = absl::StrCat(std::string(left, ' '), val_str, + std::string(right, ' ')); + } + + if (use_color) { + if (count == 0) { + cell_str = absl::StrCat("\033[32m", cell_str, "\033[0m"); + } else if (pct <= yellow_threshold_pct) { + cell_str = absl::StrCat("\033[33m", cell_str, "\033[0m"); + } else { + cell_str = absl::StrCat("\033[31m", cell_str, "\033[0m"); + } + if (is_target) { + cell_str = absl::StrCat("\033[1m", cell_str); + } + } + absl::StrAppendFormat(&out, " %s |", cell_str); + } + absl::StrAppend(&out, "\n"); + } + + absl::StrAppend(&out, " -----------+"); + for (int a = col_end; a >= col_start; --a) { + absl::StrAppend(&out, "----------+"); + } + absl::StrAppend(&out, "\n"); + absl::StrAppendFormat( + &out, + " Legend: [*] Target Tolerance (abs, rel) | Green = 0.0%% (100%% " + "pass), Yellow <= %.1f%%, Red > %.1f%% failures\n", + yellow_threshold_pct, yellow_threshold_pct); + + return out; +} + +std::string ErrorHeatmap::ToMarkdown() const { + if (abs_thresholds.empty() || rel_thresholds.empty()) { + return ""; + } + + const int num_cols_total = static_cast(abs_thresholds.size()); + const int num_rows_total = static_cast(rel_thresholds.size()); + constexpr int kMaxCols = 13; + constexpr int kMaxRows = 25; + + HeatmapWindow win = + ComputeHeatmapWindow(mismatch_counts, num_cols_total, num_rows_total, + target_abs_idx, target_rel_idx, kMaxCols, kMaxRows); + const int col_start = win.col_start; + const int col_end = win.col_end; + const int row_start = win.row_start; + const int row_end = win.row_end; + + std::string out; + absl::StrAppend(&out, "| Rel \\ Abs |"); + for (int a = col_end; a >= col_start; --a) { + std::string col_hdr = FormatCompactSci(abs_thresholds[a]); + if (a == target_abs_idx) { + absl::StrAppendFormat(&out, " **%s** *(Target)* |", col_hdr); + } else { + absl::StrAppendFormat(&out, " %s |", col_hdr); + } + } + absl::StrAppend(&out, "\n| :--- |"); + for (int a = col_end; a >= col_start; --a) { + absl::StrAppend(&out, " :---: |"); + } + absl::StrAppend(&out, "\n"); + + for (int r = row_start; r <= row_end; ++r) { + std::string row_hdr = FormatCompactSci(rel_thresholds[r]); + if (r == target_rel_idx) { + absl::StrAppendFormat(&out, "| **%s** *(Target)* |", row_hdr); + } else { + absl::StrAppendFormat(&out, "| %s |", row_hdr); + } + + for (int a = col_end; a >= col_start; --a) { + int64_t count = mismatch_counts[r][a]; + double pct = total_elements > 0 + ? (100.0 * static_cast(count) / total_elements) + : 0.0; + bool is_target = (r == target_rel_idx && a == target_abs_idx); + + std::string val_str; + const char* tile = "🟩"; + if (count == 0) { + tile = "🟩"; + val_str = "0.0%"; + } else { + val_str = (pct < 0.05) ? "<0.1%" : absl::StrFormat("%.1f%%", pct); + if (pct <= yellow_threshold_pct) { + tile = "🟨"; + } else { + tile = "🟥"; + } + } + + if (is_target) { + absl::StrAppendFormat(&out, " %s **[%s]** |", tile, val_str); + } else { + absl::StrAppendFormat(&out, " %s %s |", tile, val_str); + } + } + absl::StrAppend(&out, "\n"); + } + + absl::StrAppendFormat( + &out, + "\n*Legend: 🟩 0.0%% failures (100%% pass), 🟨 <= %.1f%% failures, " + "🟥 > %.1f%% failures. Cells indicate percentage of elements failing: " + "`abs_diff > Abs` AND `rel_diff > Rel`. Target tolerance `(abs, rel)` " + "is highlighted with `[ ]`.*\n", + yellow_threshold_pct, yellow_threshold_pct); + + return out; +} + +std::string ComparisonResult::SummaryToString(bool use_color) const { + std::string out; + std::string pass_str = + use_color ? "\033[32mPASS (MATCH)\033[0m" : "PASS (MATCH)"; + std::string fail_str = + use_color ? "\033[31mFAIL (MISMATCH)\033[0m" : "FAIL (MISMATCH)"; + absl::StrAppendFormat(&out, "Verdict: %s\n", passed ? pass_str : fail_str); + absl::StrAppendFormat(&out, " Element Type: %s\n", element_type); + absl::StrAppendFormat(&out, " Shape: %s\n", shape_str); + absl::StrAppendFormat(&out, " Total Elements: %lld\n", total_elements); + double match_pct = + total_elements > 0 ? (100.0 * exact_matches / total_elements) : 0.0; + absl::StrAppendFormat(&out, " Exact Matches: %lld (%1.2f%%)\n", + exact_matches, match_pct); + double mismatch_pct = + total_elements > 0 ? (100.0 * mismatches / total_elements) : 0.0; + absl::StrAppendFormat(&out, + " Mismatches (exceeding tolerance): %lld (%1.4f%%)\n", + mismatches, mismatch_pct); + absl::StrAppendFormat(&out, " NaN Mismatches: %lld\n", nan_mismatches); + absl::StrAppendFormat(&out, " Inf Mismatches: %lld\n", inf_mismatches); + absl::StrAppendFormat(&out, " Max Absolute Error: %1.4e\n", max_abs_error); + absl::StrAppendFormat(&out, " Max Relative Error: %1.4e\n", max_rel_error); + + if (!top_mismatches.empty()) { + absl::StrAppend(&out, "\nFirst Mismatches (up to 10):\n"); + for (const auto& m : top_mismatches) { + absl::StrAppendFormat( + &out, + " Index %lld: clean = %s, dirty = %s (abs = %.4e, rel = %.4e)\n", + m.linear_index, m.clean_str, m.dirty_str, m.abs_diff, m.rel_diff); + } + } + + if (suggested_error_spec.has_value()) { + absl::StrAppend(&out, "\n", suggested_error_spec->ToString(), "\n"); + } + + return out; +} + +std::string ComparisonResult::SummaryToMarkdown() const { + std::string out; + absl::StrAppend(&out, "# Comparison Report\n\n"); + absl::StrAppendFormat(&out, "**Verdict**: %s\n\n", + passed ? "✅ **PASS**" : "❌ **FAIL**"); + + absl::StrAppend(&out, "## Summary Statistics\n\n"); + absl::StrAppend(&out, "| Metric | Value |\n"); + absl::StrAppend(&out, "| :--- | :--- |\n"); + absl::StrAppendFormat(&out, "| Element Type | `%s` |\n", element_type); + absl::StrAppendFormat(&out, "| Shape | `%s` |\n", shape_str); + absl::StrAppendFormat(&out, "| Total Elements | %lld |\n", total_elements); + double match_pct = + total_elements > 0 ? (100.0 * exact_matches / total_elements) : 0.0; + absl::StrAppendFormat(&out, "| Exact Matches | %lld (%.2f%%) |\n", + exact_matches, match_pct); + double mismatch_pct = + total_elements > 0 ? (100.0 * mismatches / total_elements) : 0.0; + absl::StrAppendFormat( + &out, "| Mismatches (Exceeding Tolerance) | %lld (%.4f%%) |\n", + mismatches, mismatch_pct); + absl::StrAppendFormat(&out, "| NaN Mismatches | %lld |\n", nan_mismatches); + absl::StrAppendFormat(&out, "| Inf Mismatches | %lld |\n", inf_mismatches); + absl::StrAppendFormat(&out, "| Max Absolute Error | `%.4e` |\n", + max_abs_error); + absl::StrAppendFormat(&out, "| Max Relative Error | `%.4e` |\n", + max_rel_error); + + if (histogram.total_samples > 0) { + absl::StrAppend(&out, "\n## 1D Signed Relative Error Distribution\n\n"); + absl::StrAppend(&out, histogram.ToMarkdown()); + absl::StrAppendFormat( + &out, + "\n*Summary: min = `%.3e`, max = `%.3e`, mean = `%.3e`, std_dev = " + "`%.3e`*\n", + histogram.min_rel_error, histogram.max_rel_error, + histogram.mean_rel_error, histogram.std_dev_rel_error); + } + + if (!heatmap.abs_thresholds.empty()) { + absl::StrAppend(&out, "\n## 2D Error Heatmap\n\n"); + absl::StrAppend(&out, heatmap.ToMarkdown()); + } + + if (!top_mismatches.empty()) { + absl::StrAppend(&out, "\n## First Mismatches\n\n"); + absl::StrAppend(&out, "| Index | Clean | Dirty | Abs Diff | Rel Diff |\n"); + absl::StrAppend(&out, "| :---: | :--- | :--- | :---: | :---: |\n"); + for (const auto& m : top_mismatches) { + absl::StrAppendFormat(&out, "| %lld | `%s` | `%s` | `%.4e` | `%.4e` |\n", + m.linear_index, m.clean_str, m.dirty_str, + m.abs_diff, m.rel_diff); + } + } + + if (suggested_error_spec.has_value()) { + const auto& spec = *suggested_error_spec; + absl::StrAppend(&out, "\n## Suggested ErrorSpec\n\n"); + absl::StrAppendFormat( + &out, + "- **Balanced**: abs = `%s`, rel = `%s` (with 2x margin: abs = `%s`, " + "rel = `%s`)\n", + FormatCompactSci(spec.abs_bound), FormatCompactSci(spec.rel_bound), + FormatCompactSci(spec.margin_abs_bound), + FormatCompactSci(spec.margin_rel_bound)); + absl::StrAppendFormat( + &out, "- **Pure Absolute**: abs = `%s` (with 2x margin: abs = `%s`)\n", + FormatCompactSci(spec.pure_abs_bound), + FormatCompactSci(spec.margin_pure_abs_bound)); + absl::StrAppendFormat( + &out, "- **Pure Relative**: rel = `%s` (with 2x margin: rel = `%s`)\n", + FormatCompactSci(spec.pure_rel_bound), + FormatCompactSci(spec.margin_pure_rel_bound)); + } + + return out; +} + +std::string SuggestedErrorSpec::ToString() const { + std::string s; + absl::StrAppend(&s, "Suggested ErrorSpec (to pass all elements):\n"); + absl::StrAppendFormat( + &s, + " Balanced: abs = %s, rel = %s (with 2x margin: abs = %s, rel = " + "%s)\n", + FormatCompactSci(abs_bound), FormatCompactSci(rel_bound), + FormatCompactSci(margin_abs_bound), FormatCompactSci(margin_rel_bound)); + absl::StrAppendFormat( + &s, + " Pure Absolute: abs = %s (with 2x margin: abs = %s)\n", + FormatCompactSci(pure_abs_bound), + FormatCompactSci(margin_pure_abs_bound)); + absl::StrAppendFormat( + &s, " Pure Relative: rel = %s (with 2x margin: rel = %s)", + FormatCompactSci(pure_rel_bound), + FormatCompactSci(margin_pure_rel_bound)); + return s; +} + +} // namespace xla::compare_literals diff --git a/third_party/xla/xla/tools/compare_literals/compare_literals.cc b/third_party/xla/xla/tools/compare_literals/compare_literals.cc new file mode 100644 index 00000000000000..436a75f5f54998 --- /dev/null +++ b/third_party/xla/xla/tools/compare_literals/compare_literals.cc @@ -0,0 +1,136 @@ +/* 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/tools/compare_literals/compare_literals.h" + +#include +#include +#include + +#include "absl/status/status.h" +#include "absl/status/status_macros.h" +#include "absl/status/statusor.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_format.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "xla/layout_util.h" +#include "xla/literal.h" +#include "xla/primitive_util.h" +#include "xla/shape_util.h" +#include "xla/tools/compare_literals/element_comparator.h" +#include "xla/tsl/platform/env.h" +#include "xla/xla_data.pb.h" + +namespace xla::compare_literals { +namespace { + +template +ComparisonResult CompareArrayValues(const LiteralSlice& clean, + const LiteralSlice& dirty, + const ComparisonOptions& options) { + ElementComparator comparator(options, + ShapeUtil::ElementsIn(clean.shape())); + + if (LayoutUtil::Equal(dirty.shape().layout(), clean.shape().layout()) && + clean.shape().layout().element_size_in_bits() == 0 && + clean.shape().is_static() && dirty.shape().is_static()) { + absl::Span clean_span = clean.data(); + absl::Span dirty_span = dirty.data(); + const int64_t num_elements = clean_span.size(); + for (int64_t i = 0; i < num_elements; ++i) { + comparator.RecordElement(i, clean_span[i], dirty_span[i]); + } + } else { + std::vector multi_index(clean.shape().dimensions_size(), 0); + const int64_t num_elements = ShapeUtil::ElementsIn(clean.shape()); + for (int64_t i = 0; i < num_elements; ++i) { + comparator.RecordElement(i, clean.Get(multi_index), + dirty.Get(multi_index)); + for (int d = static_cast(multi_index.size()) - 1; d >= 0; --d) { + if (++multi_index[d] < clean.shape().dimensions(d)) { + break; + } + multi_index[d] = 0; + } + } + } + + return comparator.Finalize(); +} + +} // namespace + +absl::StatusOr CompareLiterals( + const LiteralSlice& clean, const LiteralSlice& dirty, + const ComparisonOptions& options) { + if (!ShapeUtil::Compatible(clean.shape(), dirty.shape())) { + return absl::InvalidArgumentError( + absl::StrFormat("Shapes must be equal; clean: %s, dirty: %s", + ShapeUtil::HumanString(clean.shape()), + ShapeUtil::HumanString(dirty.shape()))); + } + + if (!clean.shape().IsArray()) { + return absl::InvalidArgumentError( + absl::StrCat("Only array literals are supported; got: ", + ShapeUtil::HumanString(clean.shape()))); + } + + if (!primitive_util::IsArrayType(clean.shape().element_type())) { + return absl::InvalidArgumentError( + absl::StrCat("Unsupported element type for literal comparison: ", + primitive_util::LowercasePrimitiveTypeName( + clean.shape().element_type()))); + } + + ComparisonResult result = primitive_util::ArrayTypeSwitch( + [&](auto type_constant) -> ComparisonResult { + using NativeT = primitive_util::NativeTypeOf; + return CompareArrayValues(clean, dirty, options); + }, + clean.shape().element_type()); + + result.element_type = + primitive_util::LowercasePrimitiveTypeName(clean.shape().element_type()); + result.shape_str = ShapeUtil::HumanString(clean.shape()); + return result; +} + +absl::StatusOr CompareLiteralProtos( + const LiteralProto& clean_proto, const LiteralProto& dirty_proto, + const ComparisonOptions& options) { + ABSL_ASSIGN_OR_RETURN(Literal clean, Literal::CreateFromProto(clean_proto)); + ABSL_ASSIGN_OR_RETURN(Literal dirty, Literal::CreateFromProto(dirty_proto)); + return CompareLiterals(clean, dirty, options); +} + +absl::StatusOr CompareLiteralFiles( + absl::string_view clean_file, absl::string_view dirty_file, + const ComparisonOptions& options) { + LiteralProto clean_proto; + ABSL_RETURN_IF_ERROR( + tsl::ReadBinaryProto(tsl::Env::Default(), clean_file, &clean_proto)) + << absl::StrCat("Failed to read clean literal file '", clean_file, "'"); + + LiteralProto dirty_proto; + ABSL_RETURN_IF_ERROR( + tsl::ReadBinaryProto(tsl::Env::Default(), dirty_file, &dirty_proto)) + << absl::StrCat("Failed to read dirty literal file '", dirty_file, "'"); + + return CompareLiteralProtos(clean_proto, dirty_proto, options); +} + +} // namespace xla::compare_literals diff --git a/third_party/xla/xla/tools/compare_literals/compare_literals.h b/third_party/xla/xla/tools/compare_literals/compare_literals.h new file mode 100644 index 00000000000000..191c2307c23241 --- /dev/null +++ b/third_party/xla/xla/tools/compare_literals/compare_literals.h @@ -0,0 +1,180 @@ +/* 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_TOOLS_COMPARE_LITERALS_COMPARE_LITERALS_H_ +#define XLA_TOOLS_COMPARE_LITERALS_COMPARE_LITERALS_H_ + +#include +#include +#include +#include + +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "xla/literal.h" +#include "xla/xla_data.pb.h" + +namespace xla::compare_literals { + +struct ComparisonOptions { + double abs_error_bound = 1e-3; + double rel_error_bound = 1e-3; + int max_mismatches_to_record = 10; + double heatmap_yellow_pct = 0.5; +}; + +// Represents a bin in the 1D relative error distribution. +struct RelErrorBin { + double lower = 0.0; + double upper = 0.0; + int64_t count = 0; + bool is_exact_zero = false; +}; + +// 1D Relative Error Histogram with ASCII formatting. +struct RelErrorHistogram { + std::vector bins; + int64_t total_samples = 0; + double min_rel_error = 0.0; + double max_rel_error = 0.0; + double mean_rel_error = 0.0; + double std_dev_rel_error = 0.0; + int median_bin_index = -1; + + // Formats as an ASCII bar chart similar to dot_algorithms_test.cc. + std::string ToString(int max_bar_width = 40) const; + // Formats natively as a Markdown table. + std::string ToMarkdown(int max_bar_width = 30) const; +}; + +// 2D Heatmap of element mismatches for pairs of (abs_threshold, rel_threshold). +struct ErrorHeatmap { + // Sorted threshold boundaries. + std::vector abs_thresholds; + std::vector rel_thresholds; + + // 2D grid: mismatch_counts[rel_idx][abs_idx] is the number of elements + // having abs_diff > abs_thresholds[abs_idx] AND + // rel_diff > rel_thresholds[rel_idx]. + std::vector> mismatch_counts; + + // The user's target tolerance parameters. + double target_abs = 0.0; + double target_rel = 0.0; + int target_abs_idx = -1; + int target_rel_idx = -1; + int64_t total_elements = 0; + double yellow_threshold_pct = 0.5; + + // Formats the 2D matrix into a terminal-friendly table with ANSI colors. + std::string ToString(bool use_color = true) const; + // Formats natively as a Markdown table. + std::string ToMarkdown() const; +}; + +// Detailed info for an individual element mismatch. +struct MismatchDetail { + int64_t linear_index = 0; + std::string clean_str; + std::string dirty_str; + double abs_diff = 0.0; + double rel_diff = 0.0; +}; + +// Suggested error specification (abs and rel bounds) to make comparison pass. +struct SuggestedErrorSpec { + // Balanced point on Pareto frontier (knee in log-log space). + double abs_bound = 0.0; + double rel_bound = 0.0; + double margin_abs_bound = 0.0; + double margin_rel_bound = 0.0; + + // Pure absolute bound (ignoring relative error). + double pure_abs_bound = 0.0; + double margin_pure_abs_bound = 0.0; + + // Pure relative bound (ignoring absolute error). + double pure_rel_bound = 0.0; + double margin_pure_rel_bound = 0.0; + + std::string ToString() const; +}; + +// Result of comparing two literals. +struct ComparisonResult { + bool passed = false; + std::string element_type; + std::string shape_str; + int64_t total_elements = 0; + int64_t exact_matches = 0; + int64_t mismatches = 0; + int64_t nan_mismatches = 0; + int64_t inf_mismatches = 0; + double max_abs_error = 0.0; + double max_rel_error = 0.0; + + std::vector top_mismatches; + RelErrorHistogram histogram; + ErrorHeatmap heatmap; + std::optional suggested_error_spec; + + std::string SummaryToString(bool use_color = true) const; + std::string SummaryToMarkdown() const; +}; + +// Standard decade multipliers (1-2-5 sequence) used for threshold grids and +// histogram binning. +inline constexpr double kDefaultMultipliers[] = {1.0, 2.0, 5.0}; + +// Builds a wide grid of thresholds covering 10^-7 to 10^2 with {1, 2, 5} +// steps per decade, ensuring the target tolerance is always included. +std::vector BuildThresholds(double target); + +// Subdivided boundaries with {1, 2, 5} steps per decade. +std::vector CreateDefaultRelBins(); + +// Finds the 1D histogram bin index for a given relative error. +int FindRelBin(double rel_err, absl::Span bins); + +// Computes the 2D suffix sum matrix from the raw 2D histogram. +std::vector> ComputeHeatmapMismatchCounts( + const std::vector>& hist_2d, int num_rel_thresh, + int num_abs_thresh); + +// Computes suggested ErrorSpec (balanced on Pareto frontier, pure abs, pure +// rel) based on the 2D heatmap suffix sums and max recorded errors. +std::optional ComputeSuggestedErrorSpec( + const ErrorHeatmap& heatmap, double max_abs_error, double max_rel_error); + +// Compares two LiteralSlice objects and computes statistics, 1D histogram, and +// 2D heatmap in a single pass. +absl::StatusOr CompareLiterals( + const LiteralSlice& clean, const LiteralSlice& dirty, + const ComparisonOptions& options); + +// Compares two LiteralProto objects. +absl::StatusOr CompareLiteralProtos( + const LiteralProto& clean_proto, const LiteralProto& dirty_proto, + const ComparisonOptions& options); + +// Reads two LiteralProto binary files from disk and compares them. +absl::StatusOr CompareLiteralFiles( + absl::string_view clean_file, absl::string_view dirty_file, + const ComparisonOptions& options); + +} // namespace xla::compare_literals + +#endif // XLA_TOOLS_COMPARE_LITERALS_COMPARE_LITERALS_H_ diff --git a/third_party/xla/xla/tools/compare_literals/compare_literals_main.cc b/third_party/xla/xla/tools/compare_literals/compare_literals_main.cc new file mode 100644 index 00000000000000..9e6408879e790f --- /dev/null +++ b/third_party/xla/xla/tools/compare_literals/compare_literals_main.cc @@ -0,0 +1,119 @@ +/* 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 +#include +#include + +#include "absl/flags/flag.h" +#include "absl/flags/parse.h" +#include "absl/status/status.h" +#include "absl/status/statusor.h" +#include "absl/strings/string_view.h" +#include "xla/tools/compare_literals/compare_literals.h" +#include "xla/tsl/platform/env.h" +#include "tsl/platform/init_main.h" + +ABSL_FLAG(double, abs_error_bound, 1e-3, "Absolute error tolerance bound."); +ABSL_FLAG(double, rel_error_bound, 1e-3, "Relative error tolerance bound."); +ABSL_FLAG(bool, show_histogram, false, "Display 1D relative error histogram."); +ABSL_FLAG(bool, show_heatmap, false, "Display 2D error heatmap."); +ABSL_FLAG(bool, suggest_error_spec, true, + "Display suggested ErrorSpec in console output."); +ABSL_FLAG(double, heatmap_yellow_pct, 0.5, + "Failure percentage threshold (0-100) below which heatmap cells are " + "colored yellow."); +ABSL_FLAG(std::string, output_markdown, "", + "Path to write Markdown report file."); +ABSL_FLAG(int, max_bar_width, 40, "Maximum character width of histogram bars."); +ABSL_FLAG(bool, color, true, "Use ANSI colors in terminal output."); + +namespace { + +using ::xla::compare_literals::CompareLiteralFiles; +using ::xla::compare_literals::ComparisonOptions; +using ::xla::compare_literals::ComparisonResult; + +constexpr int kExitPass = 0; +constexpr int kExitMismatch = 1; +constexpr int kExitError = 2; + +} // namespace + +int main(int argc, char** argv) { + constexpr absl::string_view kUsage = + "Usage: compare_literals [flags]"; + tsl::port::InitMain(kUsage.data(), &argc, &argv); + std::vector positional_args = absl::ParseCommandLine(argc, argv); + + if (positional_args.size() != 3) { + std::cerr << "Error: Exactly two positional file paths must be provided.\n"; + std::cerr << "Usage:\n " << argv[0] + << " [flags]\n"; + return kExitError; + } + + std::string clean = positional_args[1]; + std::string dirty = positional_args[2]; + + ComparisonOptions options; + options.abs_error_bound = absl::GetFlag(FLAGS_abs_error_bound); + options.rel_error_bound = absl::GetFlag(FLAGS_rel_error_bound); + options.heatmap_yellow_pct = absl::GetFlag(FLAGS_heatmap_yellow_pct); + + absl::StatusOr result_or = + CompareLiteralFiles(clean, dirty, options); + + if (!result_or.ok()) { + std::cerr << "Comparison error: " << result_or.status() << "\n"; + return kExitError; + } + + const ComparisonResult& result = *result_or; + + std::cout << "Comparing:\n"; + std::cout << " Clean: " << clean << "\n"; + std::cout << " Dirty: " << dirty << "\n"; + std::cout << " Bounds: abs = " << options.abs_error_bound + << ", rel = " << options.rel_error_bound << "\n\n"; + + bool use_color = absl::GetFlag(FLAGS_color); + + std::cout << result.SummaryToString(use_color) << "\n"; + + if (result.total_elements > 0) { + if (absl::GetFlag(FLAGS_show_histogram)) { + std::cout << result.histogram.ToString(absl::GetFlag(FLAGS_max_bar_width)) + << "\n"; + } + + if (absl::GetFlag(FLAGS_show_heatmap)) { + std::cout << result.heatmap.ToString(absl::GetFlag(FLAGS_color)) << "\n"; + } + } + + std::string md_path = absl::GetFlag(FLAGS_output_markdown); + if (!md_path.empty()) { + std::string md = result.SummaryToMarkdown(); + absl::Status s = tsl::WriteStringToFile(tsl::Env::Default(), md_path, md); + if (!s.ok()) { + std::cerr << "Failed to write markdown report to '" << md_path + << "': " << s << "\n"; + return kExitError; + } + } + + return result.passed ? kExitPass : kExitMismatch; +} diff --git a/third_party/xla/xla/tools/compare_literals/compare_literals_test.cc b/third_party/xla/xla/tools/compare_literals/compare_literals_test.cc new file mode 100644 index 00000000000000..e6dc0008d90920 --- /dev/null +++ b/third_party/xla/xla/tools/compare_literals/compare_literals_test.cc @@ -0,0 +1,865 @@ +/* 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/tools/compare_literals/compare_literals.h" + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include "absl/status/status.h" +#include "absl/status/status_matchers.h" +#include "xla/layout_util.h" +#include "xla/literal.h" +#include "xla/literal_util.h" +#include "xla/tools/compare_literals/element_comparator.h" +#include "xla/tsl/platform/env.h" +#include "xla/types.h" +#include "tsl/platform/path.h" + +namespace xla::compare_literals { +namespace { + +using ::absl_testing::StatusIs; +using ::testing::HasSubstr; +using ::testing::Not; + +TEST(CompareLiteralsTest, ExactMatch) { + Literal lit1 = LiteralUtil::CreateR1({1.0f, 2.0f, 3.0f, 4.0f}); + Literal lit2 = LiteralUtil::CreateR1({1.0f, 2.0f, 3.0f, 4.0f}); + + ComparisonOptions options; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + + EXPECT_TRUE(result.passed); + EXPECT_EQ(result.element_type, "f32"); + EXPECT_EQ(result.shape_str, "f32[4]"); + EXPECT_EQ(result.total_elements, 4); + EXPECT_EQ(result.exact_matches, 4); + EXPECT_EQ(result.mismatches, 0); + EXPECT_DOUBLE_EQ(result.max_abs_error, 0.0); + EXPECT_DOUBLE_EQ(result.max_rel_error, 0.0); + EXPECT_THAT(result.SummaryToString(), HasSubstr("Element Type: f32")); + EXPECT_THAT(result.SummaryToString(), HasSubstr("Shape: f32[4]")); +} + +TEST(CompareLiteralsTest, WithinTolerance) { + Literal lit1 = LiteralUtil::CreateR1({1.0f, 10.0f, 100.0f}); + Literal lit2 = LiteralUtil::CreateR1({1.0001f, 10.001f, 100.01f}); + + ComparisonOptions options; + options.abs_error_bound = 1e-3; + options.rel_error_bound = 1e-3; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + + EXPECT_TRUE(result.passed); + EXPECT_EQ(result.total_elements, 3); + EXPECT_EQ(result.mismatches, 0); + EXPECT_NEAR(result.max_abs_error, 0.01, 1e-4); + EXPECT_NEAR(result.max_rel_error, 1e-4, 1e-5); +} + +TEST(CompareLiteralsTest, ExceedsTolerance) { + Literal lit1 = LiteralUtil::CreateR1({1.0f, 2.0f, 3.0f}); + Literal lit2 = LiteralUtil::CreateR1({1.0f, 2.5f, 3.0f}); + + ComparisonOptions options; + options.abs_error_bound = 1e-3; + options.rel_error_bound = 1e-3; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + + EXPECT_FALSE(result.passed); + EXPECT_EQ(result.total_elements, 3); + EXPECT_EQ(result.exact_matches, 2); + EXPECT_EQ(result.mismatches, 1); + EXPECT_NEAR(result.max_abs_error, 0.5, 1e-5); + EXPECT_NEAR(result.max_rel_error, 0.25, 1e-5); + ASSERT_EQ(result.top_mismatches.size(), 1); + EXPECT_EQ(result.top_mismatches[0].linear_index, 1); + EXPECT_EQ(result.top_mismatches[0].clean_str, "2"); + EXPECT_EQ(result.top_mismatches[0].dirty_str, "2.5"); +} + +TEST(CompareLiteralsTest, NaNHandling) { + constexpr float kNaN = std::numeric_limits::quiet_NaN(); + Literal lit1 = LiteralUtil::CreateR1({1.0f, kNaN}); + Literal lit2 = LiteralUtil::CreateR1({1.0f, kNaN}); + + ComparisonOptions options; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + + EXPECT_TRUE(result.passed); + EXPECT_EQ(result.exact_matches, 2); +} + +TEST(CompareLiteralsTest, NaNMismatch) { + constexpr float kNaN = std::numeric_limits::quiet_NaN(); + Literal lit1 = LiteralUtil::CreateR1({1.0f, 2.0f}); + Literal lit2 = LiteralUtil::CreateR1({1.0f, kNaN}); + + ComparisonOptions options; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + + EXPECT_FALSE(result.passed); + EXPECT_EQ(result.nan_mismatches, 1); + EXPECT_EQ(result.mismatches, 1); + EXPECT_TRUE(std::isinf(result.max_abs_error)); + EXPECT_TRUE(std::isinf(result.max_rel_error)); +} + +TEST(CompareLiteralsTest, ShapeMismatch) { + Literal lit1 = LiteralUtil::CreateR1({1.0f, 2.0f}); + Literal lit2 = LiteralUtil::CreateR1({1.0f, 2.0f, 3.0f}); + + ComparisonOptions options; + EXPECT_THAT(CompareLiterals(lit1, lit2, options), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Shapes must be equal"))); +} + +TEST(CompareLiteralsTest, HistogramAndHeatmapOutput) { + Literal lit1 = LiteralUtil::CreateR1({1.0f, 2.0f, 4.0f}); + Literal lit2 = LiteralUtil::CreateR1({1.01f, 2.02f, 4.04f}); + + ComparisonOptions options; + options.abs_error_bound = 1e-3; + options.rel_error_bound = 1e-3; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + + std::string hist_str = result.histogram.ToString(); + EXPECT_THAT(hist_str, HasSubstr("Summary: min =")); + + std::string heatmap_str = result.heatmap.ToString(/*use_color=*/false); + EXPECT_THAT(heatmap_str, HasSubstr("2D Error Heatmap")); + EXPECT_THAT(heatmap_str, HasSubstr("Legend:")); +} + +TEST(ElementComparatorTest, RecordIndividualElements) { + ComparisonOptions options; + options.abs_error_bound = 1e-3; + options.rel_error_bound = 1e-3; + + ElementComparator comparator(options, /*total_elements=*/3); + comparator.RecordElement(0, 1.0f, 1.0f); // exact match + comparator.RecordElement(1, 10.0f, 10.005f); // within tolerance + comparator.RecordElement(2, 2.0f, 2.5f); // mismatch + + ComparisonResult result = comparator.Finalize(); + EXPECT_FALSE(result.passed); + EXPECT_EQ(result.total_elements, 3); + EXPECT_EQ(result.exact_matches, 1); + EXPECT_EQ(result.mismatches, 1); + EXPECT_NEAR(result.max_abs_error, 0.5, 1e-5); + EXPECT_NEAR(result.max_rel_error, 0.25, 1e-5); + ASSERT_EQ(result.top_mismatches.size(), 1); + EXPECT_EQ(result.top_mismatches[0].clean_str, "2"); + EXPECT_EQ(result.top_mismatches[0].dirty_str, "2.5"); +} + +TEST(ElementComparatorTest, ComplexNumbers) { + ComparisonOptions options; + options.abs_error_bound = 1e-2; + options.rel_error_bound = 1e-2; + + ElementComparator> comparator(options, + /*total_elements=*/2); + comparator.RecordElement(0, {1.0f, 2.0f}, {1.0f, 2.0f}); + comparator.RecordElement(1, {1.0f, 0.0f}, {2.0f, 0.0f}); + + ComparisonResult result = comparator.Finalize(); + EXPECT_FALSE(result.passed); + EXPECT_EQ(result.exact_matches, 1); + EXPECT_EQ(result.mismatches, 1); + EXPECT_DOUBLE_EQ(result.max_abs_error, 1.0); +} + +TEST(ElementComparatorTest, Int8AndUint8FormattingAndDiff) { + ComparisonOptions options; + options.abs_error_bound = 0.0; + options.rel_error_bound = 0.0; + + ElementComparator int8_comp(options, /*total_elements=*/1); + int8_comp.RecordElement(0, static_cast(65), static_cast(66)); + ComparisonResult int8_result = int8_comp.Finalize(); + EXPECT_FALSE(int8_result.passed); + ASSERT_EQ(int8_result.top_mismatches.size(), 1); + // Must format as number "65", not ASCII 'A'. + EXPECT_EQ(int8_result.top_mismatches[0].clean_str, "65"); + EXPECT_EQ(int8_result.top_mismatches[0].dirty_str, "66"); + + ElementComparator uint8_comp(options, /*total_elements=*/1); + uint8_comp.RecordElement(0, static_cast(48), + static_cast(49)); + ComparisonResult uint8_result = uint8_comp.Finalize(); + EXPECT_FALSE(uint8_result.passed); + ASSERT_EQ(uint8_result.top_mismatches.size(), 1); + // Must format as number "48", not ASCII '0'. + EXPECT_EQ(uint8_result.top_mismatches[0].clean_str, "48"); + EXPECT_EQ(uint8_result.top_mismatches[0].dirty_str, "49"); +} + +TEST(CompareLiteralsTest, NegativeValues) { + Literal lit1 = LiteralUtil::CreateR1({-10.0f, -100.0f}); + Literal lit2 = LiteralUtil::CreateR1({-10.005f, -99.95f}); + + ComparisonOptions options; + options.abs_error_bound = 1e-1; + options.rel_error_bound = 1e-3; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + + EXPECT_TRUE(result.passed); + EXPECT_NEAR(result.max_abs_error, 0.05, 1e-5); + EXPECT_NEAR(result.max_rel_error, 5e-4, 1e-6); + + // Exceeding tolerance on negative numbers + Literal lit3 = LiteralUtil::CreateR1({-2.0f}); + Literal lit4 = LiteralUtil::CreateR1({-2.5f}); + ASSERT_OK_AND_ASSIGN(ComparisonResult result2, + CompareLiterals(lit3, lit4, options)); + EXPECT_FALSE(result2.passed); + EXPECT_NEAR(result2.max_abs_error, 0.5, 1e-5); + EXPECT_NEAR(result2.max_rel_error, 0.25, 1e-5); +} + +TEST(CompareLiteralsTest, ExactMatchesIncludedInWelford) { + // 3 exact matches, 1 element with 0.04 relative error. + Literal lit1 = LiteralUtil::CreateR1({10.0f, 10.0f, 10.0f, 10.0f}); + Literal lit2 = LiteralUtil::CreateR1({10.0f, 10.0f, 10.0f, 10.4f}); + + ComparisonOptions options; + options.abs_error_bound = 1.0; + options.rel_error_bound = 0.1; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + + EXPECT_TRUE(result.passed); + // Mean relative error across 4 elements: (0 + 0 + 0 + 0.04) / 4 = 0.01. + EXPECT_NEAR(result.histogram.mean_rel_error, 0.01, 1e-5); +} + +TEST(CompareLiteralsTest, LargeInt64Comparison) { + // Values above 2^53 that differ by 1. + constexpr int64_t kBase = 9007199254740992LL; // 2^53 + Literal lit1 = LiteralUtil::CreateR1({kBase, kBase + 1}); + Literal lit2 = LiteralUtil::CreateR1({kBase, kBase}); + + ComparisonOptions options; + options.abs_error_bound = 0.0; + options.rel_error_bound = 0.0; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + + EXPECT_FALSE(result.passed); + EXPECT_EQ(result.exact_matches, 1); + EXPECT_EQ(result.mismatches, 1); + EXPECT_DOUBLE_EQ(result.max_abs_error, 1.0); +} + +TEST(CompareLiteralsTest, SingleElementMedian) { + Literal lit1 = LiteralUtil::CreateR1({100.0f}); + Literal lit2 = LiteralUtil::CreateR1({102.0f}); // +2% rel error + + ComparisonOptions options; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + + ASSERT_GE(result.histogram.median_bin_index, 0); + EXPECT_EQ(result.histogram.bins[result.histogram.median_bin_index].count, 1); +} + +TEST(CompareLiteralsTest, SuggestedErrorSpecFailingComparison) { + Literal lit1 = LiteralUtil::CreateR1({1.0f, 2.0f, 3.0f}); + Literal lit2 = LiteralUtil::CreateR1({1.0f, 2.5f, 3.0f}); + + ComparisonOptions options; + options.abs_error_bound = 1e-3; + options.rel_error_bound = 1e-3; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + + EXPECT_FALSE(result.passed); + ASSERT_TRUE(result.suggested_error_spec.has_value()); + const auto& spec = *result.suggested_error_spec; + + EXPECT_GE(spec.pure_abs_bound, 0.5); + EXPECT_GE(spec.pure_rel_bound, 0.25); + EXPECT_GE(spec.margin_abs_bound, spec.abs_bound); + EXPECT_GE(spec.margin_rel_bound, spec.rel_bound); + + // Crucial verification: running CompareLiterals with suggested balanced + // bounds MUST pass! + ComparisonOptions passing_options; + passing_options.abs_error_bound = spec.abs_bound; + passing_options.rel_error_bound = spec.rel_bound; + ASSERT_OK_AND_ASSIGN(ComparisonResult passing_result, + CompareLiterals(lit1, lit2, passing_options)); + EXPECT_TRUE(passing_result.passed); + EXPECT_EQ(passing_result.mismatches, 0); + + // Pure absolute bound verification (with rel = 0) + ComparisonOptions pure_abs_options; + pure_abs_options.abs_error_bound = spec.pure_abs_bound; + pure_abs_options.rel_error_bound = 0.0; + ASSERT_OK_AND_ASSIGN(ComparisonResult pure_abs_result, + CompareLiterals(lit1, lit2, pure_abs_options)); + EXPECT_TRUE(pure_abs_result.passed); + + // Pure relative bound verification (with abs = 0) + ComparisonOptions pure_rel_options; + pure_rel_options.abs_error_bound = 0.0; + pure_rel_options.rel_error_bound = spec.pure_rel_bound; + ASSERT_OK_AND_ASSIGN(ComparisonResult pure_rel_result, + CompareLiterals(lit1, lit2, pure_rel_options)); + EXPECT_TRUE(pure_rel_result.passed); + + // Output formatting verification + EXPECT_THAT(result.SummaryToString(), HasSubstr("Suggested ErrorSpec")); + EXPECT_THAT(result.SummaryToString(), HasSubstr("Balanced:")); + EXPECT_THAT(result.SummaryToString(), HasSubstr("Pure Absolute:")); + EXPECT_THAT(result.SummaryToString(), HasSubstr("Pure Relative:")); +} + +TEST(CompareLiteralsTest, SuggestedErrorSpecIncludedOnPassingComparison) { + Literal lit1 = LiteralUtil::CreateR1({1.0f, 2.0f, 3.0f}); + Literal lit2 = LiteralUtil::CreateR1({1.0f, 2.01f, 3.0f}); + + ComparisonOptions options; + options.abs_error_bound = 0.1; + options.rel_error_bound = 0.1; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + + EXPECT_TRUE(result.passed); + ASSERT_TRUE(result.suggested_error_spec.has_value()); + + // Verify SummaryToString defaults to including Suggested ErrorSpec on PASS + EXPECT_THAT(result.SummaryToString(), HasSubstr("Suggested ErrorSpec")); + EXPECT_THAT(result.SummaryToString(), HasSubstr("Balanced:")); + + // Verify SummaryToString with use_color = false emits plain text without ANSI + EXPECT_THAT(result.SummaryToString(/*use_color=*/false), + HasSubstr("PASS (MATCH)")); +} + +TEST(CompareLiteralsTest, SuggestedErrorSpecNulloptOnNanMismatches) { + constexpr float kNaN = std::numeric_limits::quiet_NaN(); + Literal lit1 = LiteralUtil::CreateR1({1.0f, 2.0f}); + Literal lit2 = LiteralUtil::CreateR1({1.0f, kNaN}); + + ComparisonOptions options; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + + EXPECT_FALSE(result.passed); + EXPECT_EQ(result.nan_mismatches, 1); + EXPECT_FALSE(result.suggested_error_spec.has_value()); +} + +TEST(CompareLiteralsTest, InfinityMatchesAndMismatches) { + constexpr float kInf = std::numeric_limits::infinity(); + Literal clean_inf = LiteralUtil::CreateR1({kInf, -kInf}); + Literal dirty_inf = LiteralUtil::CreateR1({kInf, -kInf}); + + ComparisonOptions options; + ASSERT_OK_AND_ASSIGN(ComparisonResult match_result, + CompareLiterals(clean_inf, dirty_inf, options)); + EXPECT_TRUE(match_result.passed); + EXPECT_EQ(match_result.exact_matches, 2); + EXPECT_EQ(match_result.inf_mismatches, 0); + + Literal clean_mix = LiteralUtil::CreateR1({1.0f, -kInf}); + Literal dirty_mix = LiteralUtil::CreateR1({kInf, -kInf}); + ASSERT_OK_AND_ASSIGN(ComparisonResult mismatch_result, + CompareLiterals(clean_mix, dirty_mix, options)); + EXPECT_FALSE(mismatch_result.passed); + EXPECT_EQ(mismatch_result.inf_mismatches, 1); + EXPECT_EQ(mismatch_result.mismatches, 1); + EXPECT_EQ(mismatch_result.max_abs_error, + std::numeric_limits::infinity()); + EXPECT_FALSE(mismatch_result.suggested_error_spec.has_value()); + ASSERT_FALSE(mismatch_result.top_mismatches.empty()); + EXPECT_EQ(mismatch_result.top_mismatches[0].linear_index, 0); +} + +TEST(CompareLiteralsTest, MismatchedLayoutFallback) { + Literal lit_row = LiteralUtil::CreateR2({{1.0f, 2.0f}, {3.0f, 4.0f}}); + Literal lit_col = lit_row.Relayout(LayoutUtil::MakeLayout({0, 1})); + ASSERT_FALSE( + LayoutUtil::Equal(lit_row.shape().layout(), lit_col.shape().layout())); + + ComparisonOptions options; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit_row, lit_col, options)); + EXPECT_TRUE(result.passed); + EXPECT_EQ(result.exact_matches, 4); + EXPECT_EQ(result.mismatches, 0); +} + +TEST(CompareLiteralsTest, SummaryToMarkdownReport) { + Literal lit1 = LiteralUtil::CreateR1({1.0f, 2.0f}); + Literal lit2 = LiteralUtil::CreateR1({1.0f, 2.5f}); + + ComparisonOptions options; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + + std::string md = result.SummaryToMarkdown(); + EXPECT_THAT(md, HasSubstr("# Comparison Report")); + EXPECT_THAT(md, HasSubstr("## Summary Statistics")); + EXPECT_THAT(md, HasSubstr("## 1D Signed Relative Error Distribution")); + EXPECT_THAT( + md, HasSubstr("| Markers | Range | Count | Percent | Distribution |")); + EXPECT_THAT(md, HasSubstr("## 2D Error Heatmap")); + EXPECT_THAT(md, HasSubstr("| Rel \\ Abs |")); + EXPECT_THAT(md, HasSubstr("🟩 0.0%")); + EXPECT_THAT(md, HasSubstr("Legend: 🟩 0.0% failures")); + EXPECT_THAT(md, HasSubstr("## First Mismatches")); + EXPECT_THAT(md, HasSubstr("## Suggested ErrorSpec")); + // Heatmap and histogram must be native Markdown tables without code fences. + EXPECT_THAT(md, Not(HasSubstr("```"))); +} + +TEST(CompareLiteralsTest, HeatmapPercentageAndYellowThresholdConfigurable) { + // 100 elements: 98 matching, 2 differing by 0.5 (abs_diff = 0.5, rel_diff = + // 0.25). Total elements = 100, failures at tight tolerances = 2 (2.0%). + std::vector clean_vals(100, 2.0f); + std::vector dirty_vals = clean_vals; + dirty_vals[0] = 2.5f; + dirty_vals[1] = 2.5f; + + Literal lit1 = LiteralUtil::CreateR1(clean_vals); + Literal lit2 = LiteralUtil::CreateR1(dirty_vals); + + // Run with default yellow threshold (0.5%): 2.0% > 0.5% so cell should be RED + // (🟥). + ComparisonOptions options_default; + options_default.abs_error_bound = 1e-3; + options_default.rel_error_bound = 1e-3; + ASSERT_OK_AND_ASSIGN(ComparisonResult result_default, + CompareLiterals(lit1, lit2, options_default)); + + std::string md_default = result_default.heatmap.ToMarkdown(); + EXPECT_THAT(md_default, HasSubstr("2.0%")); + // At target tolerance (1e-3, 1e-3), failure rate is 2.0% > 0.5%, so it is + // red: + EXPECT_THAT(md_default, HasSubstr("🟥 **[2.0%]**")); + + // Run with custom yellow threshold (5.0%): 2.0% <= 5.0% so cell should be + // YELLOW (🟨). + ComparisonOptions options_custom; + options_custom.abs_error_bound = 1e-3; + options_custom.rel_error_bound = 1e-3; + options_custom.heatmap_yellow_pct = 5.0; + ASSERT_OK_AND_ASSIGN(ComparisonResult result_custom, + CompareLiterals(lit1, lit2, options_custom)); + + std::string md_custom = result_custom.heatmap.ToMarkdown(); + EXPECT_THAT(md_custom, HasSubstr("🟨 **[2.0%]**")); + EXPECT_THAT(md_custom, + HasSubstr("Legend: 🟩 0.0% failures (100% pass), 🟨 <= 5.0% " + "failures, 🟥 > 5.0% failures")); + + // Console output should also display percentages and yellow threshold: + std::string console = result_custom.heatmap.ToString(/*use_color=*/false); + EXPECT_THAT(console, HasSubstr("2.0%")); + EXPECT_THAT(console, HasSubstr("Yellow <= 5.0%")); +} + +TEST(CompareLiteralsTest, CompareLiteralProtosAndFiles) { + Literal lit1 = LiteralUtil::CreateR1({1.0f, 2.0f}); + Literal lit2 = LiteralUtil::CreateR1({1.0f, 2.0f}); + + LiteralProto proto1 = lit1.ToProto(); + LiteralProto proto2 = lit2.ToProto(); + + ComparisonOptions options; + ASSERT_OK_AND_ASSIGN(ComparisonResult proto_result, + CompareLiteralProtos(proto1, proto2, options)); + EXPECT_TRUE(proto_result.passed); + + std::string clean_path = tsl::io::JoinPath(testing::TempDir(), "clean.pb"); + std::string dirty_path = tsl::io::JoinPath(testing::TempDir(), "dirty.pb"); + + ASSERT_OK(tsl::WriteBinaryProto(tsl::Env::Default(), clean_path, proto1)); + ASSERT_OK(tsl::WriteBinaryProto(tsl::Env::Default(), dirty_path, proto2)); + + ASSERT_OK_AND_ASSIGN(ComparisonResult file_result, + CompareLiteralFiles(clean_path, dirty_path, options)); + EXPECT_TRUE(file_result.passed); + + EXPECT_THAT( + CompareLiteralFiles("/nonexistent/path/clean.pb", dirty_path, options), + StatusIs(absl::StatusCode::kNotFound, + HasSubstr("Failed to read clean literal file"))); + EXPECT_THAT( + CompareLiteralFiles(clean_path, "/nonexistent/path/dirty.pb", options), + StatusIs(absl::StatusCode::kNotFound, + HasSubstr("Failed to read dirty literal file"))); +} + +TEST(CompareLiteralsTest, AllZerosCleanAndDirty) { + Literal clean = LiteralUtil::CreateR1({0.0f, 0.0f, 0.0f, 0.0f}); + Literal dirty = LiteralUtil::CreateR1({0.0f, 0.0f, 0.0f, 0.0f}); + + ComparisonOptions options; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(clean, dirty, options)); + + EXPECT_TRUE(result.passed); + EXPECT_EQ(result.total_elements, 4); + EXPECT_EQ(result.exact_matches, 4); + EXPECT_EQ(result.mismatches, 0); + EXPECT_DOUBLE_EQ(result.max_abs_error, 0.0); + EXPECT_DOUBLE_EQ(result.max_rel_error, 0.0); + EXPECT_TRUE( + result.histogram.bins[result.histogram.median_bin_index].is_exact_zero); + EXPECT_EQ(result.histogram.bins[result.histogram.median_bin_index].count, 4); +} + +TEST(CompareLiteralsTest, SignedZerosEquivalent) { + Literal clean = LiteralUtil::CreateR1({+0.0f, -0.0f}); + Literal dirty = LiteralUtil::CreateR1({-0.0f, +0.0f}); + + ComparisonOptions options; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(clean, dirty, options)); + + EXPECT_TRUE(result.passed); + EXPECT_EQ(result.exact_matches, 2); + EXPECT_EQ(result.mismatches, 0); +} + +TEST(CompareLiteralsTest, ZeroReferenceWithNonZeroDirty) { + Literal clean = LiteralUtil::CreateR1({0.0f, 0.0f}); + Literal dirty = LiteralUtil::CreateR1({1e-4f, 0.5f}); + + ComparisonOptions options; + options.abs_error_bound = 1e-3; + options.rel_error_bound = 1e-3; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(clean, dirty, options)); + + EXPECT_FALSE(result.passed); + EXPECT_EQ(result.exact_matches, 0); + EXPECT_EQ(result.mismatches, 1); + EXPECT_NEAR(result.max_abs_error, 0.5, 1e-5); + EXPECT_DOUBLE_EQ(result.max_rel_error, 0.0); + ASSERT_TRUE(result.suggested_error_spec.has_value()); + EXPECT_GE(result.suggested_error_spec->pure_abs_bound, 0.5); + EXPECT_TRUE(std::isinf(result.suggested_error_spec->pure_rel_bound)); + EXPECT_TRUE(std::isinf(result.suggested_error_spec->margin_pure_rel_bound)); + EXPECT_GE(result.suggested_error_spec->abs_bound, 0.5); + EXPECT_DOUBLE_EQ(result.suggested_error_spec->rel_bound, 0.0); +} + +TEST(CompareLiteralsTest, Rank0ScalarLiterals) { + Literal clean_scalar = LiteralUtil::CreateR0(42.0f); + Literal dirty_scalar = LiteralUtil::CreateR0(42.0f); + + ComparisonOptions options; + ASSERT_OK_AND_ASSIGN(ComparisonResult match_result, + CompareLiterals(clean_scalar, dirty_scalar, options)); + EXPECT_TRUE(match_result.passed); + EXPECT_EQ(match_result.shape_str, "f32[]"); + EXPECT_EQ(match_result.total_elements, 1); + EXPECT_EQ(match_result.exact_matches, 1); + + Literal mismatch_scalar = LiteralUtil::CreateR0(43.0f); + ASSERT_OK_AND_ASSIGN(ComparisonResult mismatch_result, + CompareLiterals(clean_scalar, mismatch_scalar, options)); + EXPECT_FALSE(mismatch_result.passed); + EXPECT_EQ(mismatch_result.mismatches, 1); + EXPECT_DOUBLE_EQ(mismatch_result.max_abs_error, 1.0); +} + +TEST(CompareLiteralsTest, Rank3TensorLiterals) { + Literal clean = LiteralUtil::CreateR3( + {{{1.0f, 2.0f}, {3.0f, 4.0f}}, {{5.0f, 6.0f}, {7.0f, 8.0f}}}); + Literal dirty = LiteralUtil::CreateR3( + {{{1.0f, 2.0f}, {3.0f, 4.0f}}, {{5.0f, 6.0f}, {7.0f, 8.5f}}}); + + ComparisonOptions options; + options.abs_error_bound = 1e-3; + options.rel_error_bound = 1e-3; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(clean, dirty, options)); + + EXPECT_FALSE(result.passed); + EXPECT_EQ(result.shape_str, "f32[2,2,2]"); + EXPECT_EQ(result.total_elements, 8); + EXPECT_EQ(result.exact_matches, 7); + EXPECT_EQ(result.mismatches, 1); + EXPECT_NEAR(result.max_abs_error, 0.5, 1e-5); + ASSERT_EQ(result.top_mismatches.size(), 1); + EXPECT_EQ(result.top_mismatches[0].clean_str, "8"); + EXPECT_EQ(result.top_mismatches[0].dirty_str, "8.5"); +} + +TEST(CompareLiteralsTest, EmptyLiteralZeroElements) { + Literal clean = LiteralUtil::CreateR1({}); + Literal dirty = LiteralUtil::CreateR1({}); + + ComparisonOptions options; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(clean, dirty, options)); + + EXPECT_TRUE(result.passed); + EXPECT_EQ(result.total_elements, 0); + EXPECT_EQ(result.exact_matches, 0); + EXPECT_EQ(result.mismatches, 0); + EXPECT_TRUE(result.suggested_error_spec.has_value()); + EXPECT_THAT(result.SummaryToString(), HasSubstr("Total Elements: 0")); +} + +TEST(CompareLiteralsTest, HeatmapIncludesAllZeroColumnAndRow) { + Literal clean = LiteralUtil::CreateR1({1.0f, 2.0f, 3.0f}); + Literal dirty = LiteralUtil::CreateR1({1.0f, 2.05f, 3.0f}); + + ComparisonOptions options; + options.abs_error_bound = 1e-4; + options.rel_error_bound = 1e-4; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(clean, dirty, options)); + + EXPECT_FALSE(result.passed); + std::string md = result.heatmap.ToMarkdown(); + std::string console = result.heatmap.ToString(/*use_color=*/false); + + // The true max absolute error is 0.05. The next grid threshold is 5e-2. + // The table must include the 5e-2 column (or higher), which has all 0.0%. + EXPECT_THAT(md, HasSubstr("5e-2")); + EXPECT_THAT(console, HasSubstr("5e-2")); + // Target tolerance *1e-4 must also be visible: + EXPECT_THAT(md, HasSubstr("1e-4")); + EXPECT_THAT(console, HasSubstr("1e-4")); +} + +TEST(CompareLiteralsTest, ExactToleranceBoundaryPasses) { + Literal lit1 = LiteralUtil::CreateR1({1.0f}); + Literal lit2 = LiteralUtil::CreateR1({1.125f}); + ComparisonOptions options; + options.abs_error_bound = 0.125; + options.rel_error_bound = 0.125; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(lit1, lit2, options)); + EXPECT_TRUE(result.passed); + EXPECT_EQ(result.mismatches, 0); +} + +TEST(CompareLiteralsTest, StandaloneMathHelpers) { + std::vector thresholds = BuildThresholds(0.035); + EXPECT_THAT(thresholds, ::testing::Contains(0.035)); + EXPECT_TRUE(std::is_sorted(thresholds.begin(), thresholds.end())); + + // Verify non-positive targets are not added (kills Mutant 1) + std::vector zero_thresholds = BuildThresholds(0.0); + EXPECT_THAT(zero_thresholds, ::testing::Not(::testing::Contains(0.0))); + std::vector neg_thresholds = BuildThresholds(-1.0); + EXPECT_THAT(neg_thresholds, ::testing::Not(::testing::Contains(-1.0))); + + std::vector bins = CreateDefaultRelBins(); + EXPECT_GE(bins.size(), 30); + + // Verify 1.0 multiplier boundaries for e > -6 exist (kills Mutant 2) + bool has_one_milli = false; + bool has_one = false; + for (const auto& b : bins) { + if (std::abs(b.lower - 1e-3) < 1e-9 || std::abs(b.upper - 1e-3) < 1e-9) { + has_one_milli = true; + } + if (std::abs(b.lower - 1.0) < 1e-9 || std::abs(b.upper - 1.0) < 1e-9) { + has_one = true; + } + } + EXPECT_TRUE(has_one_milli); + EXPECT_TRUE(has_one); + + int zero_idx = FindRelBin(0.0, bins); + EXPECT_TRUE(bins[zero_idx].is_exact_zero); + int pos_idx = FindRelBin(0.05, bins); + EXPECT_GE(pos_idx, zero_idx); + int neg_idx = FindRelBin(-0.05, bins); + EXPECT_LE(neg_idx, zero_idx); + + constexpr double kInf = std::numeric_limits::infinity(); + EXPECT_EQ(FindRelBin(-kInf, bins), 0); + EXPECT_EQ(FindRelBin(kInf, bins), static_cast(bins.size() - 1)); + + // Verify asymmetric bins where exact_zero is NOT at mid (kills Mutant 3) + std::vector zero_at_start = { + {0.0, 0.0, 0, true}, + {0.0, 1.0, 0, false}, + {1.0, 2.0, 0, false}, + }; + EXPECT_EQ(FindRelBin(0.0, zero_at_start), 0); + + // Verify bins where rel_err >= lower is strictly required (kills Mutant 25) + std::vector reverse_bins = { + {0.5, 1.0, 0, false}, + {0.0, 0.5, 0, false}, + }; + EXPECT_EQ(FindRelBin(0.2, reverse_bins), 1); + + std::vector> hist_2d = { + {1, 2, 0}, + {3, 4, 0}, + {0, 0, 0}, + }; + auto mismatch_counts = ComputeHeatmapMismatchCounts(hist_2d, 2, 2); + EXPECT_EQ(mismatch_counts[0][0], 4); + + // Verify SuggestedErrorSpec when max_abs_error > 0 but max_rel_error == 0 + // (kills Mutant 26) + ErrorHeatmap dummy_heatmap; + dummy_heatmap.abs_thresholds = {1e-3, 1e-2, 1e-1}; + dummy_heatmap.rel_thresholds = {1e-3, 1e-2, 1e-1}; + dummy_heatmap.mismatch_counts = { + {0, 0, 0}, + {0, 0, 0}, + {0, 0, 0}, + }; + auto spec_abs_only = ComputeSuggestedErrorSpec(dummy_heatmap, 0.05, 0.0); + ASSERT_TRUE(spec_abs_only.has_value()); + EXPECT_GE(spec_abs_only->pure_abs_bound, 0.05); + EXPECT_TRUE(std::isinf(spec_abs_only->pure_rel_bound)); + EXPECT_TRUE(std::isinf(spec_abs_only->margin_pure_rel_bound)); + EXPECT_GE(spec_abs_only->abs_bound, 0.05); + EXPECT_DOUBLE_EQ(spec_abs_only->rel_bound, 0.0); +} + +TEST(CompareLiteralsTest, MultiPointParetoKneeSelection) { + Literal clean = LiteralUtil::CreateR1({1000.0f, 0.01f}); + Literal dirty = LiteralUtil::CreateR1({1005.0f, 0.012f}); + + ComparisonOptions options; + options.abs_error_bound = 1e-4; + options.rel_error_bound = 1e-4; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(clean, dirty, options)); + + EXPECT_FALSE(result.passed); + ASSERT_TRUE(result.suggested_error_spec.has_value()); + const auto& spec = *result.suggested_error_spec; + + ComparisonOptions balanced_options; + balanced_options.abs_error_bound = spec.abs_bound; + balanced_options.rel_error_bound = spec.rel_bound; + ASSERT_OK_AND_ASSIGN(ComparisonResult balanced_result, + CompareLiterals(clean, dirty, balanced_options)); + EXPECT_TRUE(balanced_result.passed); +} + +TEST(CompareLiteralsTest, Bfloat16AndHalfSupport) { + Literal clean_bf16 = LiteralUtil::CreateR1( + {bfloat16(1.0f), bfloat16(2.0f), bfloat16(3.0f)}); + Literal dirty_bf16 = LiteralUtil::CreateR1( + {bfloat16(1.0f), bfloat16(2.5f), bfloat16(3.0f)}); + + ComparisonOptions options; + options.abs_error_bound = 0.05; + options.rel_error_bound = 0.05; + ASSERT_OK_AND_ASSIGN(ComparisonResult bf16_result, + CompareLiterals(clean_bf16, dirty_bf16, options)); + EXPECT_FALSE(bf16_result.passed); + EXPECT_EQ(bf16_result.mismatches, 1); + + Literal clean_f16 = + LiteralUtil::CreateR1({half(1.0f), half(2.0f), half(3.0f)}); + Literal dirty_f16 = + LiteralUtil::CreateR1({half(1.0f), half(2.0f), half(3.0f)}); + ASSERT_OK_AND_ASSIGN(ComparisonResult f16_result, + CompareLiterals(clean_f16, dirty_f16, options)); + EXPECT_TRUE(f16_result.passed); +} + +TEST(CompareLiteralsTest, DoubleAndInt32AndBoolSupport) { + Literal clean_f64 = LiteralUtil::CreateR1({1.0, 2.0}); + Literal dirty_f64 = LiteralUtil::CreateR1({1.0, 2.0}); + ComparisonOptions options; + ASSERT_OK_AND_ASSIGN(ComparisonResult f64_result, + CompareLiterals(clean_f64, dirty_f64, options)); + EXPECT_TRUE(f64_result.passed); + + Literal clean_s32 = LiteralUtil::CreateR1({10, 20}); + Literal dirty_s32 = LiteralUtil::CreateR1({10, 25}); + ASSERT_OK_AND_ASSIGN(ComparisonResult s32_result, + CompareLiterals(clean_s32, dirty_s32, options)); + EXPECT_FALSE(s32_result.passed); + EXPECT_EQ(s32_result.mismatches, 1); + + Literal clean_bool = LiteralUtil::CreateR1({true, false}); + Literal dirty_bool = LiteralUtil::CreateR1({true, false}); + ASSERT_OK_AND_ASSIGN(ComparisonResult bool_result, + CompareLiterals(clean_bool, dirty_bool, options)); + EXPECT_TRUE(bool_result.passed); +} + +TEST(CompareLiteralsTest, ComplexLiteralsEndToEnd) { + Literal clean = LiteralUtil::CreateR1( + {complex64(1.0f, 2.0f), complex64(3.0f, 4.0f)}); + Literal dirty = LiteralUtil::CreateR1( + {complex64(1.0f, 2.0f), complex64(4.0f, 4.0f)}); + + ComparisonOptions options; + options.abs_error_bound = 0.1; + options.rel_error_bound = 0.1; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(clean, dirty, options)); + EXPECT_FALSE(result.passed); + EXPECT_EQ(result.exact_matches, 1); + EXPECT_EQ(result.mismatches, 1); +} + +TEST(CompareLiteralsTest, NonArrayTupleLiteralRejected) { + Literal lit1 = LiteralUtil::MakeTupleFromSlices({}); + Literal lit2 = LiteralUtil::MakeTupleFromSlices({}); + + ComparisonOptions options; + auto result = CompareLiterals(lit1, lit2, options); + EXPECT_THAT(result.status(), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("Only array literals are supported"))); +} + +TEST(CompareLiteralsTest, MaxMismatchesCapRespected) { + std::vector clean_vals(20, 1.0f); + std::vector dirty_vals(20, 2.0f); + Literal clean = LiteralUtil::CreateR1(clean_vals); + Literal dirty = LiteralUtil::CreateR1(dirty_vals); + + ComparisonOptions options; + options.max_mismatches_to_record = 5; + ASSERT_OK_AND_ASSIGN(ComparisonResult result, + CompareLiterals(clean, dirty, options)); + EXPECT_FALSE(result.passed); + EXPECT_EQ(result.mismatches, 20); + EXPECT_EQ(result.top_mismatches.size(), 5); +} + +} // namespace +} // namespace xla::compare_literals diff --git a/third_party/xla/xla/tools/compare_literals/element_comparator.h b/third_party/xla/xla/tools/compare_literals/element_comparator.h new file mode 100644 index 00000000000000..c96ba4df78c323 --- /dev/null +++ b/third_party/xla/xla/tools/compare_literals/element_comparator.h @@ -0,0 +1,347 @@ +/* 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_TOOLS_COMPARE_LITERALS_ELEMENT_COMPARATOR_H_ +#define XLA_TOOLS_COMPARE_LITERALS_ELEMENT_COMPARATOR_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_format.h" +#include "xla/tools/compare_literals/compare_literals.h" + +namespace xla::compare_literals { + +namespace internal { +// Type trait helpers for complex values. +template +struct IsComplex : std::false_type {}; +template +struct IsComplex> : std::true_type {}; +template +inline constexpr bool is_complex_v = IsComplex::value; +} // namespace internal + +template +struct ValueTraits { + static bool IsNan(T val) { + if constexpr (std::is_floating_point_v || !std::is_integral_v) { + return std::isnan(static_cast(val)); + } + return false; + } + static bool IsInf(T val) { + if constexpr (std::is_floating_point_v || !std::is_integral_v) { + return std::isinf(static_cast(val)); + } + return false; + } + static double AbsDiff(T a, T b) { + if constexpr (std::is_integral_v) { + if (a == b) { + return 0.0; + } + uint64_t diff = (a > b) + ? static_cast(a) - static_cast(b) + : static_cast(b) - static_cast(a); + return std::max(1.0, static_cast(diff)); + } else { + return std::abs(static_cast(a) - static_cast(b)); + } + } + static double Magnitude(T val) { return std::abs(static_cast(val)); } + static std::string Format(T val) { + if constexpr (std::is_same_v) { + return val ? "true" : "false"; + } else if constexpr (std::is_integral_v) { + return absl::StrCat(+val); + } else { + return absl::StrCat(static_cast(val)); + } + } +}; + +template +struct ValueTraits> { + static bool IsNan(const std::complex& val) { + return std::isnan(val.real()) || std::isnan(val.imag()); + } + static bool IsInf(const std::complex& val) { + return std::isinf(val.real()) || std::isinf(val.imag()); + } + static double AbsDiff(const std::complex& a, const std::complex& b) { + return std::abs(std::complex(a.real(), a.imag()) - + std::complex(b.real(), b.imag())); + } + static double Magnitude(const std::complex& val) { + return std::abs(std::complex(val.real(), val.imag())); + } + static std::string Format(const std::complex& val) { + return absl::StrFormat("(%s, %s)", absl::StrCat(val.real()), + absl::StrCat(val.imag())); + } +}; + +// Element-level comparator and metric accumulator. +template +class ElementComparator { + public: + ElementComparator(const ComparisonOptions& options, int64_t total_elements) + : options_(options) { + result_.total_elements = total_elements; + + ErrorHeatmap& heatmap = result_.heatmap; + heatmap.total_elements = total_elements; + heatmap.target_abs = options.abs_error_bound; + heatmap.target_rel = options.rel_error_bound; + heatmap.abs_thresholds = BuildThresholds(options.abs_error_bound); + heatmap.rel_thresholds = BuildThresholds(options.rel_error_bound); + heatmap.yellow_threshold_pct = options.heatmap_yellow_pct; + + auto abs_it = absl::c_find(heatmap.abs_thresholds, options.abs_error_bound); + heatmap.target_abs_idx = + abs_it != heatmap.abs_thresholds.end() + ? std::distance(heatmap.abs_thresholds.begin(), abs_it) + : -1; + + auto rel_it = absl::c_find(heatmap.rel_thresholds, options.rel_error_bound); + heatmap.target_rel_idx = + rel_it != heatmap.rel_thresholds.end() + ? std::distance(heatmap.rel_thresholds.begin(), rel_it) + : -1; + + const int num_abs_thresh = heatmap.abs_thresholds.size(); + const int num_rel_thresh = heatmap.rel_thresholds.size(); + hist_2d_.assign(num_rel_thresh + 1, + std::vector(num_abs_thresh + 1, 0)); + + result_.histogram.bins = CreateDefaultRelBins(); + } + + // Compares an individual element pair at linear index `idx`. + void RecordElement(int64_t idx, NativeT clean_val, NativeT dirty_val) { + constexpr double kInfinity = std::numeric_limits::infinity(); + double abs_diff = 0.0; + double rel_diff = 0.0; + double signed_rel = 0.0; + bool is_nan = false; + bool is_inf = false; + bool has_rel_error = false; + + if (ValueTraits::IsNan(clean_val) || + ValueTraits::IsNan(dirty_val)) { + if (ValueTraits::IsNan(clean_val) && + ValueTraits::IsNan(dirty_val)) { + result_.exact_matches++; + abs_diff = 0.0; + rel_diff = 0.0; + signed_rel = 0.0; + } else { + is_nan = true; + result_.nan_mismatches++; + result_.mismatches++; + abs_diff = kInfinity; + rel_diff = kInfinity; + signed_rel = kInfinity; + result_.max_abs_error = std::max(result_.max_abs_error, abs_diff); + result_.max_rel_error = std::max(result_.max_rel_error, rel_diff); + if (result_.top_mismatches.size() < options_.max_mismatches_to_record) { + result_.top_mismatches.push_back( + {idx, ValueTraits::Format(clean_val), + ValueTraits::Format(dirty_val), abs_diff, rel_diff}); + } + } + } else if (ValueTraits::IsInf(clean_val) || + ValueTraits::IsInf(dirty_val)) { + if (clean_val == dirty_val) { + result_.exact_matches++; + abs_diff = 0.0; + rel_diff = 0.0; + signed_rel = 0.0; + } else { + is_inf = true; + result_.inf_mismatches++; + result_.mismatches++; + abs_diff = kInfinity; + rel_diff = kInfinity; + signed_rel = kInfinity; + result_.max_abs_error = std::max(result_.max_abs_error, abs_diff); + result_.max_rel_error = std::max(result_.max_rel_error, rel_diff); + if (result_.top_mismatches.size() < options_.max_mismatches_to_record) { + result_.top_mismatches.push_back( + {idx, ValueTraits::Format(clean_val), + ValueTraits::Format(dirty_val), abs_diff, rel_diff}); + } + } + } else { + if (clean_val == dirty_val) { + result_.exact_matches++; + abs_diff = 0.0; + rel_diff = 0.0; + signed_rel = 0.0; + + double clean_mag = ValueTraits::Magnitude(clean_val); + if (clean_mag != 0.0) { + has_rel_error = true; + // Online Welford update for exact matches (signed_rel = 0.0) + finite_rel_samples_++; + double delta = 0.0 - mean_signed_rel_; + mean_signed_rel_ += delta / finite_rel_samples_; + double delta2 = 0.0 - mean_signed_rel_; + m2_signed_rel_ += delta * delta2; + + min_signed_rel_ = std::min(min_signed_rel_, 0.0); + max_signed_rel_ = std::max(max_signed_rel_, 0.0); + } + } else { + abs_diff = ValueTraits::AbsDiff(dirty_val, clean_val); + result_.max_abs_error = std::max(result_.max_abs_error, abs_diff); + + double clean_mag = ValueTraits::Magnitude(clean_val); + if (clean_mag != 0.0) { + has_rel_error = true; + if constexpr (internal::is_complex_v) { + rel_diff = abs_diff / clean_mag; + signed_rel = rel_diff; + } else { + rel_diff = abs_diff / clean_mag; + signed_rel = (dirty_val >= clean_val ? rel_diff : -rel_diff); + } + result_.max_rel_error = std::max(result_.max_rel_error, rel_diff); + + // Online Welford update guarded against subnormal overflow + if (std::isfinite(signed_rel)) { + finite_rel_samples_++; + double delta = signed_rel - mean_signed_rel_; + mean_signed_rel_ += delta / finite_rel_samples_; + double delta2 = signed_rel - mean_signed_rel_; + m2_signed_rel_ += delta * delta2; + + min_signed_rel_ = std::min(min_signed_rel_, signed_rel); + max_signed_rel_ = std::max(max_signed_rel_, signed_rel); + } + } + + // Mismatch check against target bounds: + // When clean is non-zero, both bounds must be exceeded. + // When clean is zero, relative error is undefined so only abs bound + // applies. + bool is_mismatch = + (abs_diff > options_.abs_error_bound) && + (!has_rel_error || rel_diff > options_.rel_error_bound); + + if (is_mismatch) { + result_.mismatches++; + if (result_.top_mismatches.size() < + options_.max_mismatches_to_record) { + result_.top_mismatches.push_back( + {idx, ValueTraits::Format(clean_val), + ValueTraits::Format(dirty_val), abs_diff, + has_rel_error ? rel_diff : 0.0}); + } + } + } + } + + // 1D histogram binning (only for finite relative errors or exact zeros) + if (!is_nan && !is_inf && std::isfinite(signed_rel) && + (clean_val == dirty_val || has_rel_error)) { + int bin_idx = FindRelBin(signed_rel, result_.histogram.bins); + result_.histogram.bins[bin_idx].count++; + result_.histogram.total_samples++; + } + + // 2D heatmap binning + int a_bin = absl::c_lower_bound(result_.heatmap.abs_thresholds, abs_diff) - + result_.heatmap.abs_thresholds.begin(); + int r_bin = + has_rel_error + ? (absl::c_lower_bound(result_.heatmap.rel_thresholds, rel_diff) - + result_.heatmap.rel_thresholds.begin()) + : (is_nan || is_inf || abs_diff > 0.0 + ? static_cast(result_.heatmap.rel_thresholds.size()) + : 0); + hist_2d_[r_bin][a_bin]++; + } + + // Finalizes summary statistics (mean, stddev, median, 2D suffix sums) + // and returns the final ComparisonResult. + ComparisonResult Finalize() { + RelErrorHistogram& histogram = result_.histogram; + if (finite_rel_samples_ > 0) { + histogram.min_rel_error = min_signed_rel_; + histogram.max_rel_error = max_signed_rel_; + histogram.mean_rel_error = mean_signed_rel_; + histogram.std_dev_rel_error = + finite_rel_samples_ > 1 ? std::sqrt(std::max(0.0, m2_signed_rel_) / + (finite_rel_samples_ - 1)) + : 0.0; + } + + if (histogram.total_samples > 0) { + const int64_t target = (histogram.total_samples + 1) / 2; + int64_t cumulative = 0; + for (size_t i = 0; i < histogram.bins.size(); ++i) { + cumulative += histogram.bins[i].count; + if (cumulative >= target) { + histogram.median_bin_index = static_cast(i); + break; + } + } + } + + // 2D Suffix sum to compute mismatch counts + ErrorHeatmap& heatmap = result_.heatmap; + const int num_abs_thresh = heatmap.abs_thresholds.size(); + const int num_rel_thresh = heatmap.rel_thresholds.size(); + heatmap.mismatch_counts = + ComputeHeatmapMismatchCounts(hist_2d_, num_rel_thresh, num_abs_thresh); + + result_.passed = (result_.mismatches == 0 && result_.nan_mismatches == 0 && + result_.inf_mismatches == 0); + + result_.suggested_error_spec = ComputeSuggestedErrorSpec( + result_.heatmap, result_.max_abs_error, result_.max_rel_error); + + return result_; + } + + const ComparisonResult& result() const { return result_; } + + private: + ComparisonOptions options_; + ComparisonResult result_; + std::vector> hist_2d_; + + double min_signed_rel_ = std::numeric_limits::infinity(); + double max_signed_rel_ = -std::numeric_limits::infinity(); + double mean_signed_rel_ = 0.0; + double m2_signed_rel_ = 0.0; + int64_t finite_rel_samples_ = 0; +}; + +} // namespace xla::compare_literals + +#endif // XLA_TOOLS_COMPARE_LITERALS_ELEMENT_COMPARATOR_H_