From adfa0ffffb79fd54c0185f6e38c035f51aeb608f Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 10:16:37 +0000 Subject: [PATCH 01/23] Bump ubuntu from `b7f4819` to `2260313` in /tensorflow/tools/gcs_test Bumps ubuntu from `b7f4819` to `2260313`. --- updated-dependencies: - dependency-name: ubuntu dependency-version: '24.04' dependency-type: direct:production ... Signed-off-by: dependabot[bot] --- tensorflow/tools/gcs_test/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 " From bb6b3b8937baca2a29446c27f84046d168ffac4d Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Tue, 1 Sep 2026 23:46:21 -0700 Subject: [PATCH 02/23] Update LLVM patches for XLA Update build.patch and toolchains.patch to resolve patch failures introduced by the last LLVM integrate. PiperOrigin-RevId: 974933166 --- third_party/xla/third_party/llvm/build.patch | 45 ++++++++++++++++--- .../xla/third_party/llvm/toolchains.patch | 20 ++++----- 2 files changed, 49 insertions(+), 16 deletions(-) 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"), From af1345665dbf9796445e3572a9f6cfbd0e8c551d Mon Sep 17 00:00:00 2001 From: Alexandru Paiu Date: Wed, 2 Sep 2026 00:58:52 -0700 Subject: [PATCH 03/23] Re-adds TFLITE_LOG_PROD when dlopen fails in GpuModulePlugin PiperOrigin-RevId: 974962936 --- .../mini_benchmark/gpu_module_plugin.cc | 4 ++++ .../mini_benchmark/gpu_module_plugin_test.cc | 22 +++++++++++++++++-- 2 files changed, 24 insertions(+), 2 deletions(-) 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 From d6d8c11f36aa947281961e3b0a2f16ea462bd9f1 Mon Sep 17 00:00:00 2001 From: Ilya Tikhonovskiy Date: Wed, 2 Sep 2026 01:06:49 -0700 Subject: [PATCH 04/23] [XLA] Add compare_literals tool for single-pair literal verification and error analysis. XLA GPU numerics monitoring requires comparing dumped LiteralProto files against golden references to detect subtle mathematical regressions. Existing tools either operate only on in-memory tensors or lack diagnostic distribution metrics. This change introduces the compare_literals tool (Phase 1): - Evaluates elementwise tolerances (abs, rel) returning pass/fail verdicts. - Computes single-pass online statistics (mean, variance, min, max) without buffer allocations. - Emits a 1D signed relative error distribution histogram with median, mean, and zero markers. - Constructs a 2D (abs, rel) error heatmap identifying Pareto frontiers and highlight target tolerances. - Automatically calculates suggested ErrorSpecs (balanced on the Pareto frontier knee point, pure absolute, pure relative) with 2x safety margins. - Supports exporting structured GitHub Flavored Markdown comparison reports via --output_markdown. PiperOrigin-RevId: 974966726 --- .../xla/xla/tools/compare_literals/BUILD | 72 ++ .../compare_literals/compare_literal_math.cc | 279 ++++++ .../compare_literal_report.cc | 601 ++++++++++++ .../compare_literals/compare_literals.cc | 136 +++ .../tools/compare_literals/compare_literals.h | 180 ++++ .../compare_literals/compare_literals_main.cc | 119 +++ .../compare_literals/compare_literals_test.cc | 865 ++++++++++++++++++ .../compare_literals/element_comparator.h | 347 +++++++ 8 files changed, 2599 insertions(+) create mode 100644 third_party/xla/xla/tools/compare_literals/BUILD create mode 100644 third_party/xla/xla/tools/compare_literals/compare_literal_math.cc create mode 100644 third_party/xla/xla/tools/compare_literals/compare_literal_report.cc create mode 100644 third_party/xla/xla/tools/compare_literals/compare_literals.cc create mode 100644 third_party/xla/xla/tools/compare_literals/compare_literals.h create mode 100644 third_party/xla/xla/tools/compare_literals/compare_literals_main.cc create mode 100644 third_party/xla/xla/tools/compare_literals/compare_literals_test.cc create mode 100644 third_party/xla/xla/tools/compare_literals/element_comparator.h 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_ From 145f5f43a4652aa14864ec505293b35bedfc7c74 Mon Sep 17 00:00:00 2001 From: dlcompilers-infra-bot Date: Wed, 2 Sep 2026 01:18:37 -0700 Subject: [PATCH 05/23] PR #48203: [XLA:GPU] Query NVLink count from NVML MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/48203 📝 Summary of Changes Query the device-specific NVLink count through NVML_FI_DEV_NVLINK_LINK_COUNT instead of relying on hardcoded constant. 🎯 Justification Two paths checked only 18/32 nvlink indices. Rubin exposes 36 NVLink ports, so both paths omitted valid links. 🚀 Kind of Contribution 🐛 Bug Fix, 🧪 Tests 🧪 Unit Tests: CudaExecutorTest.CreateDeviceDescription Copybara import of the project: -- 3a2c005e2e25b16717dc189e9d3f0422ae3e6113 by Andrei Ivanov : [XLA:GPU] Query NVLink count from NVML Merging this change closes #48203 PiperOrigin-RevId: 974972190 --- .../xla/stream_executor/cuda/cuda_executor.cc | 22 +++++++++++++++---- .../cuda/cuda_executor_test.cc | 9 +++++++- 2 files changed, 26 insertions(+), 5 deletions(-) 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())); From fd2073bf4f6e53a94618979425801598b5d35e6e Mon Sep 17 00:00:00 2001 From: Nikita Putikhin Date: Wed, 2 Sep 2026 01:31:01 -0700 Subject: [PATCH 06/23] Extract triton configs and config generation in dot search space This is refactoring to make methods reusable in the fission estimation code coming later PiperOrigin-RevId: 974977062 --- .../xla/xla/backends/gpu/autotuner/triton.cc | 79 +++----- .../xla/backends/gpu/autotuner/triton/BUILD | 8 +- .../gpu/autotuner/triton/dot_search_space.cc | 21 ++- .../gpu/autotuner/triton/dot_search_space.h | 14 +- .../autotuner/triton/dot_search_space_test.cc | 169 ++++++++++++------ .../gpu/autotuner/triton/triton_configs.cc | 41 ++++- .../gpu/autotuner/triton/triton_configs.h | 6 + .../autotuner/triton/triton_configs_test.cc | 36 ++++ 8 files changed, 255 insertions(+), 119 deletions(-) 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 From bcad5676f4ea4da5f796858495ab4185242cfe26 Mon Sep 17 00:00:00 2001 From: lao fan <928784910@qq.com> Date: Wed, 2 Sep 2026 01:45:14 -0700 Subject: [PATCH 07/23] PR #48186: Gate unsafe exponential rewrites on fast math MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/48186 These rewrites can change NaN and infinity behavior. Preserve strict floating-point semantics by applying them only when fast math is enabled, and test both modes. 📝 Summary of Changes Gate four exponential algebraic rewrites behind `enable_fast_math`: - `exp(A) * exp(B) -> exp(A + B)` - `exp(A) / exp(B) -> exp(A - B)` - `A / exp(B) -> A * exp(-B)` - `pow(exp(A), B) -> exp(A * B)` Strict floating-point mode now preserves the original HLO, while fast-math mode retains the existing simplifications. Inspired by tensorflow/tensorflow/pull/123991. 🎯 Justification These rewrites can change NaN and infinity behavior. This was observed in practice in tensorflow/tensorflow#123169, where an exponential rewrite changed a NaN result to `1`. 🚀 Kind of Contribution: 🐛 Bug Fix, 🧪 Tests Updated the four existing algebraic simplifier tests to verify both strict and fast-math behavior. `//xla/hlo/transforms/simplifiers:algebraic_simplifier_test`: 1998 tests passed. No new execution test. The change is covered at the HLO pass level, and the practical failure is documented in tensorflow/ tensorflow#123169. Copybara import of the project: -- 29764815682e1e825b7b5a89a650669698fb634f by wuren <928784910@qq.com>: Gate unsafe exponential rewrites on fast math These rewrites can change NaN and infinity behavior. Preserve strict floating-point semantics by applying them only when fast math is enabled, and test both modes. Merging this change closes #48186 PiperOrigin-RevId: 974982874 --- .../simplifiers/algebraic_simplifier.cc | 12 +++-- .../simplifiers/algebraic_simplifier_test.cc | 54 ++++++++++++++++--- 2 files changed, 54 insertions(+), 12 deletions(-) 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(), From 234ece16a11785540ac0fa6df8f215a030ae19ee Mon Sep 17 00:00:00 2001 From: Maxim Ermilov Date: Wed, 2 Sep 2026 01:51:54 -0700 Subject: [PATCH 08/23] Introduce CommandBuffer (Create|Update)Memcpy(H2D|D2H) PiperOrigin-RevId: 974985859 --- .../xla/xla/stream_executor/command_buffer.h | 20 ++++ .../cuda/cuda_command_buffer.cc | 101 ++++++++++++++++++ .../cuda/cuda_command_buffer.h | 17 +++ .../cuda/cuda_command_buffer_test.cc | 54 ++++++++++ .../stream_executor/gpu/gpu_command_buffer.cc | 41 +++++++ .../stream_executor/gpu/gpu_command_buffer.h | 33 ++++++ .../xla/stream_executor/mock_command_buffer.h | 16 +++ .../rocm/rocm_command_buffer.cc | 65 +++++++++++ .../rocm/rocm_command_buffer.h | 17 +++ 9 files changed, 364 insertions(+) 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/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 { From 5c54c66bcb035418b5d7154691031b8132c6b970 Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Wed, 2 Sep 2026 02:06:12 -0700 Subject: [PATCH 09/23] Do not bit-pack PRED elements in Literal serialization Serialize PRED elements with 8-bit storage instead of packing them at 1-bit per element with LSB masking. This allows PRED literals preserving non-zero bytes (e.g., from bitcast-convert) to round-trip through serialization and deserialization without losing upper bits. PiperOrigin-RevId: 974992146 --- third_party/xla/xla/BUILD | 1 - third_party/xla/xla/literal.h | 14 ++++++------- third_party/xla/xla/shape_util.cc | 6 +----- third_party/xla/xla/tests/BUILD | 3 --- .../xla/xla/tests/bitcast_convert_test.cc | 21 +------------------ 5 files changed, 9 insertions(+), 36 deletions(-) diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD index fc24f6ecb62a80..5d78fa465ee4a9 100644 --- a/third_party/xla/xla/BUILD +++ b/third_party/xla/xla/BUILD @@ -521,7 +521,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/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/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/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 From 85d7b06d6bcddc2902f8fe8c644fab2a250a25ca Mon Sep 17 00:00:00 2001 From: destro4evr Date: Wed, 2 Sep 2026 02:44:07 -0700 Subject: [PATCH 10/23] PR #46604: Fix integer overflow in IFRT proxy byte-strides size check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/46604 ## Problem `ArrayMemRegion::FromZerothElementPointer` in `xla/python/ifrt_proxy/common/array_util.cc` accumulates the byte span of a strided array using an unchecked `int64_t` multiplication: ```cpp last_element_byte_offset += (stride * (shape.dims()[i] - 1)); ``` When `stride * (dims[i] - 1)` overflows `int64_t` and wraps to zero (e.g. `stride = 2^62`, `dims[i] = 5`: `2^62 × 4 = 2^64 ≡ 0`), the computed memory region size equals `byte_size` regardless of the true array span. `FromMinimalMemRegion` then compares this against the caller-supplied data size, which an IFRT proxy client also controls - so the size check passes and validation is bypassed. This affects two gRPC handlers in `IfrtBackend`: - **`MakeArrayFromHostBuffer`**: the PjRt backend is handed a small host buffer declared as a multi-element strided array and reads elements at offsets up to `(N-1) × stride` bytes beyond it (OOB read). - **`CopyToHostBuffer`**: the server allocates a `byte_size`-byte `std::string` for the output, then the PjRt backend writes `N` strided elements into it (heap buffer overflow). ## Fix Use `OverflowSafeMultiply` (already available in `xla/overflow_util.h`) before accumulating `last_element_byte_offset`, and return `InvalidArgumentError` on overflow. Add `//xla:overflow_util` to the `array_util` BUILD target. ## Changes - `xla/python/ifrt_proxy/common/array_util.cc`: replace unchecked multiplication with `OverflowSafeMultiply`; add `#include "xla/overflow_util.h"` - `xla/python/ifrt_proxy/common/BUILD`: add `//xla:overflow_util` dep Copybara import of the project: -- 767ad4ff23c76f5c98d24d60be61823e6ea4c1b4 by destro4evr-rgb : Fix integer overflow in IFRT proxy byte-strides size check stride * (shape.dims()[i] - 1) in FromZerothElementPointer was an unchecked int64_t multiplication. With stride=2^62 and dims[i]=5 the product wraps to zero, making the computed memory region size equal to byte_size regardless of the true array span. FromMinimalMemRegion then compares this against the caller-supplied data size which an IFRT proxy client also controls, so the size check passes and validation is bypassed. This allows an IFRT proxy client to trigger an OOB read via MakeArrayFromHostBuffer (PjRt backend reads elements at offsets up to (N-1)*stride bytes beyond a byte_size-byte buffer) and a heap buffer overflow via CopyToHostBuffer (server allocates byte_size bytes then PjRt writes N strided elements into it). Fix: use OverflowSafeMultiply before accumulating last_element_byte_offset and return InvalidArgumentError on overflow. Add //xla:overflow_util to the array_util BUILD target. -- 17c3961eaca239c9ece44d7f73eb0dff32422909 by destro4evr-rgb : Add unit tests for byte-stride overflow rejection in ArrayMemRegion Cover the case where stride * (dim - 1) wraps to zero on overflow, which would otherwise bypass the size check in FromMinimalMemRegion. -- 1d9d9c48ff35155f50eb670c920f5e7a3ea185dc by destro4evr-rgb : Address review nit: use dim[i] in overflow error message for clarity -- 4fcd4bf21bd59f6a415f28f0fe64e8c4438a1c9e by destro4evr-rgb : Fix error message: dim -> dim[i] only in formula part per review suggestion -- 1501c18038b6359e2f32f09980f6f32d09ade373 by destro4evr-rgb : address review: use __builtin_mul_overflow instead of OverflowSafeMultiply -- 9fb74f2c39dd0a1a05c3456ca411f0379fd305f7 by destro4evr-rgb : address review: remove overflow_util BUILD dependency Merging this change closes #46604 PiperOrigin-RevId: 975008318 --- .../xla/xla/python/ifrt_proxy/common/array_util.cc | 8 +++++++- .../xla/python/ifrt_proxy/common/array_util_test.cc | 13 ++++++++++++- 2 files changed, 19 insertions(+), 2 deletions(-) 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(); From 4681ff91fb017672253642d9662cb52bc164c89d Mon Sep 17 00:00:00 2001 From: kgotlinux <60880393+kgotlinux@users.noreply.github.com> Date: Wed, 2 Sep 2026 02:45:03 -0700 Subject: [PATCH 11/23] PR #48240: Fix powerpc64le build MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/48240 📝 Summary of Changes Two fixes necessary to fix PowerPC builds. 🎯 Justification Without those, building on PowerPC doesn't work. 🚀 Kind of Contribution 🐛 Bug Fix Copybara import of the project: -- 141e584924fbf2d4dc57fbb2664dc10c30e76cc5 by Katarzyna Kubaj : Include in builtin_fp16.h The last branch of the XlaF16ABIType selection falls back to uint16_t, but the header includes nothing, so uint16_t is undeclared: xla/backends/cpu/codegen/builtin_fp16.h:32:23: error: unknown type name 'uint16_t' followed by a cascade of "unknown type name 'XlaF16ABIType'" wherever the alias is used. The fallback is unreachable on x86_64 and on AArch64, which is why this has gone unnoticed; it is taken on any target whose compiler does not provide _Float16, such as clang on powerpc64. -- 5f4811d70aca4733c6fc9d8b9045ebf9b0a11e75 by Katarzyna Kubaj : Require _Float16 support for the vectorised CPU intrinsics vector_ops.h and eigen_unary.cc are guarded on ext_vector_type and __builtin_vectorelements, but the body of vector_ops.h then declares typedef _Float16 Vec8h __attribute__((vector_size(16))); clang provides both of those attributes on powerpc64 while rejecting _Float16 there: xla/codegen/intrinsic/cpp/vector_ops.h:31:9: error: _Float16 is not supported on this target so the guard admits targets that cannot compile the code. Test __FLT16_MANT_DIG__ as well, which is the same predicate builtin_fp16.h already uses to decide whether _Float16 is available. Both files then compile to empty translation units on such targets, which is the behaviour the existing guard was written to provide. x86_64 and AArch64 are unaffected: clang defines __FLT16_MANT_DIG__ there (on x86 whenever SSE2 is available, i.e. always on x86_64). Merging this change closes #48240 PiperOrigin-RevId: 975008760 --- third_party/xla/xla/backends/cpu/codegen/builtin_fp16.h | 2 ++ third_party/xla/xla/codegen/intrinsic/cpp/eigen_unary.cc | 5 +++-- third_party/xla/xla/codegen/intrinsic/cpp/vector_ops.h | 5 +++-- 3 files changed, 8 insertions(+), 4 deletions(-) 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/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 From dc18aa14aaf1a9bc0373296d74783d15520c32dc Mon Sep 17 00:00:00 2001 From: Bill Varcho Date: Wed, 2 Sep 2026 02:47:19 -0700 Subject: [PATCH 12/23] [SDY][Bug-Fix] use preorder walk to match comment in shard_map export. PiperOrigin-RevId: 975009593 --- .../spmd/shardy/stablehlo_round_trip/shard_map_export.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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); }); From c9f7155e0350a4396e730aa9b67563e983b446a3 Mon Sep 17 00:00:00 2001 From: Penporn Koanantakool Date: Wed, 2 Sep 2026 02:58:01 -0700 Subject: [PATCH 13/23] [XLA:CPU:oneDNN] Add missing `copts` for `onednn_threadpool` `OneDnnThreadpool` checks for `ENABLE_ONEDNN_ASYNC` macro. This is defined in `tsl_copts` and is missing from the current `copts`. PiperOrigin-RevId: 975014094 --- third_party/xla/xla/backends/cpu/runtime/onednn/BUILD | 1 + 1 file changed, 1 insertion(+) 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", From a30ea91b11ca72297b9f4f00b2a710e84e55cbef Mon Sep 17 00:00:00 2001 From: Bill Varcho Date: Wed, 2 Sep 2026 03:04:49 -0700 Subject: [PATCH 14/23] [SDY][1/2] Support subaxes in ReplicaGroupMeshAxesAttr PiperOrigin-RevId: 975017508 --- .../xla/third_party/stablehlo/temporary.patch | 571 ++++++++++++++++++ 1 file changed, 571 insertions(+) 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 From 5b3f7e035a0d20dd154384198125fc3e1b9c7cef Mon Sep 17 00:00:00 2001 From: abhinav srivastava Date: Wed, 2 Sep 2026 03:08:50 -0700 Subject: [PATCH 15/23] PR #48202: [XLA:GPU][oneAPI][Bugfix] Prevent compile-time hangs in VectorizeLoad trunc-user walk Imported from GitHub PR https://github.com/openxla/xla/pull/48202 This PR fixes exponential re-exploration in the recursive walk used during the transitive search over arith.trunci users that narrow to a sub-byte type. Because the walk did not track visited operations, reconverging use graphs could cause compile-time hangs, even though the generated result was unchanged. testcase fixed: random_lax_test.py::DistributionsTest.testDirichlet0 random_lax_test.py::DistributionsTest.testDirichlet1 Copybara import of the project: -- 3efdc9351c894156759a2a36fcd7a22c34972305 by abhinav srivastava : Prevent compile-time hangs in VectorizeLoad trunc-user walk -- 2ce582918457d7cbabebdfafabc7a5974e4926c9 by abhinav srivastava : Added braces to single-statement ifs in the lambda Merging this change closes #48202 PiperOrigin-RevId: 975019469 --- .../codegen/emitters/transforms/vectorize_loads_stores.cc | 5 +++++ 1 file changed, 5 insertions(+) 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())) From 5c0c477a413c7232209b265b8ffd701251f988eb Mon Sep 17 00:00:00 2001 From: Penporn Koanantakool Date: Wed, 2 Sep 2026 03:12:10 -0700 Subject: [PATCH 16/23] [XLA:CPU:oneDNN] Delete a stray, unused `InitializeXnnPack` forward declaration. PiperOrigin-RevId: 975020955 --- .../xla/xla/backends/cpu/runtime/onednn/onednn_interop.h | 3 --- 1 file changed, 3 deletions(-) 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)) { From bba1dfc2c9a46a76abb38a4ae0e0c15de359bf53 Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Wed, 2 Sep 2026 03:23:52 -0700 Subject: [PATCH 17/23] Remove deprecated wrapper Removes the deprecated createExpandHloTuplesPass wrapper taking entryFunctionName from mhlo/transforms/passes.h. PiperOrigin-RevId: 975025831 --- third_party/xla/xla/mlir_hlo/mhlo/transforms/passes.h | 8 -------- 1 file changed, 8 deletions(-) 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" From eb0f8c70aa1eb46ba1be2bbb9ec5415115bec578 Mon Sep 17 00:00:00 2001 From: Kanvi Khanna Date: Wed, 2 Sep 2026 03:35:13 -0700 Subject: [PATCH 18/23] PR #47656: [XLA:GPU][oneAPI] Add F64 support for matmul MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Imported from GitHub PR https://github.com/openxla/xla/pull/47656 📝 Summary of Changes This PR adds F64 data-type support for matmul for oneAPI 🚀 Kind of Contribution ✨ New Feature Copybara import of the project: -- e788342caf3961a175c7cd213f6ee6e64f659624 by Kanvi Khanna : Add F64 support for matmul -- 5a196e8306f6843a51c61c538b6d10bd2cf61146 by Kanvi Khanna : update comment -- f8be87dbe282c22d9608ec20c9f677cf521a50f3 by Kanvi Khanna : Address review comments and format -- 1afeb93d6cbc0627755a675f9bc5411ec33c1251 by Kanvi Khanna : use absl_assign_or_return Merging this change closes #47656 PiperOrigin-RevId: 975030465 --- .../backends/gpu/transforms/gemm_rewriter.cc | 4 +- .../xla/xla/stream_executor/sycl/BUILD | 3 + .../stream_executor/sycl/gemm_sycl_test.cc | 3 +- .../xla/stream_executor/sycl/onednn_util.cc | 25 +++++ .../xla/stream_executor/sycl/onednn_util.h | 6 + .../stream_executor/sycl/sycl_matmul_utils.cc | 103 +++--------------- 6 files changed, 51 insertions(+), 93 deletions(-) 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/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 From fccb487ea3390c6dd1389225fb5dfaeeee2c610a Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Wed, 2 Sep 2026 04:11:21 -0700 Subject: [PATCH 19/23] Upgrade rules_cc to 0.2.20 and bazel_skylib to 1.9.0 in TensorFlow - Explicitly define rules_cc 0.2.20 in workspace3.bzl. - Call compatibility_proxy_repo() in WORKSPACE and workspace1.bzl to define @cc_compatibility_proxy in WORKSPACE mode. - Bump rules_cc to 0.2.20 in MODULE.bazel. - Bump bazel_skylib to 1.9.0 in workspace3.bzl and MODULE.bazel to support the 'scope' attribute on bool_flag used by rules_cc 0.2.20. - In third_party/googleapis/build_rules.bzl.oss, load cc_proto_library from protobuf instead of rules_cc defs.bzl, avoiding errors from the removal of cc_proto_library in rules_cc >= 0.2.18. PiperOrigin-RevId: 975045752 --- MODULE.bazel | 4 ++-- WORKSPACE | 4 ++++ tensorflow/workspace1.bzl | 3 +++ tensorflow/workspace3.bzl | 13 +++++++++++-- third_party/googleapis/build_rules.bzl | 2 +- 5 files changed, 21 insertions(+), 5 deletions(-) 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/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: From d53df5ed3bedad45c6ad465b888360b08915887a Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 2 Sep 2026 04:11:31 -0700 Subject: [PATCH 20/23] Automated Code Change PiperOrigin-RevId: 975045823 --- tensorflow/core/kernels/data/BUILD | 5 +++++ tensorflow/core/kernels/data/cache_dataset_ops.cc | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD index 3c5b008c1a74fb..fafd1230216aa6 100644 --- a/tensorflow/core/kernels/data/BUILD +++ b/tensorflow/core/kernels/data/BUILD @@ -80,11 +80,16 @@ tf_kernel_library( "//tensorflow/core/data:name_utils", "//tensorflow/core/data:serialization_utils", "//tensorflow/core/framework:dataset_options_proto_cc", + "//tensorflow/core/framework:types_proto_cc", "//tensorflow/core/util/tensor_bundle", "//tensorflow/core/util/tensor_bundle:naming", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/strings:string_view", + "@com_google_absl//absl/types:span", ], ) diff --git a/tensorflow/core/kernels/data/cache_dataset_ops.cc b/tensorflow/core/kernels/data/cache_dataset_ops.cc index 9bfcbe60a7ffb1..7866214dac8ef6 100644 --- a/tensorflow/core/kernels/data/cache_dataset_ops.cc +++ b/tensorflow/core/kernels/data/cache_dataset_ops.cc @@ -22,9 +22,13 @@ limitations under the License. #include #include +#include "absl/log/check.h" +#include "absl/log/log.h" #include "absl/status/status.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 "tensorflow/core/data/global_shuffle_utils.h" #include "tensorflow/core/data/name_utils.h" #include "tensorflow/core/data/serialization_utils.h" @@ -33,6 +37,7 @@ limitations under the License. #include "tensorflow/core/framework/partial_tensor_shape.h" #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/kernels/data/cache_ops.h" #include "tensorflow/core/kernels/data/iterator_ops.h" #include "tensorflow/core/lib/core/errors.h" From 6e8b7d3b4483bf9534c502a4de9205933aa592bf Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Wed, 2 Sep 2026 04:29:59 -0700 Subject: [PATCH 21/23] Add ComparisonOrder to HloProto, IR, and HLO parser `Comparison::Order` was introduced to distinguish data type classification from mathematical ordering, but the migration was never completed across HloProto and HLO text format. This change adds `comparison_order` to `HloInstructionProto`, supports parsing `order=` on `compare` HLO instructions, and fixes `Comparison` constructor initialization so `type_` and `order_` stay synchronized during the migration. PiperOrigin-RevId: 975052031 --- third_party/xla/xla/BUILD | 1 + third_party/xla/xla/comparison_util.cc | 95 ++++++++++++------- third_party/xla/xla/comparison_util.h | 12 ++- third_party/xla/xla/comparison_util_test.cc | 24 +++++ third_party/xla/xla/hlo/ir/BUILD | 2 + third_party/xla/xla/hlo/ir/hlo_instruction.cc | 34 ++++--- third_party/xla/xla/hlo/ir/hlo_instruction.h | 2 +- .../xla/xla/hlo/ir/hlo_instruction_test.cc | 58 ++++++++--- .../xla/xla/hlo/ir/hlo_instructions.cc | 13 ++- third_party/xla/xla/hlo/ir/hlo_instructions.h | 13 ++- third_party/xla/xla/hlo/parser/BUILD | 1 + third_party/xla/xla/hlo/parser/hlo_parser.cc | 44 ++++++++- .../xla/xla/hlo/parser/hlo_parser_test.cc | 49 +++++++++- third_party/xla/xla/service/hlo.proto | 5 +- third_party/xla/xla/service/hlo_verifier.cc | 26 ++--- .../xla/xla/service/hlo_verifier_test.cc | 47 ++++----- 16 files changed, 301 insertions(+), 125 deletions(-) diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD index 5d78fa465ee4a9..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", ], ) 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/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) { From 738e30ccb05fcce6447f62704a0c7ebe300b5f3f Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Wed, 2 Sep 2026 05:05:25 -0700 Subject: [PATCH 22/23] Migrate createTileLoopsPass caller in kernel_creator Migrates LowerHloToLoops in kernel_creator.cc to construct TileLoopsPassOptions instead of calling the legacy pass wrapper. PiperOrigin-RevId: 975064328 --- .../compiler/mlir/tools/kernel_gen/kernel_creator.cc | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) 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()); From 21c7f59253e88831b8889a15a8c09ad9a7ce345e Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 2 Sep 2026 05:13:59 -0700 Subject: [PATCH 23/23] Reverts d53df5ed3bedad45c6ad465b888360b08915887a PiperOrigin-RevId: 975067190 --- tensorflow/core/kernels/data/BUILD | 5 ----- tensorflow/core/kernels/data/cache_dataset_ops.cc | 5 ----- 2 files changed, 10 deletions(-) diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD index fafd1230216aa6..3c5b008c1a74fb 100644 --- a/tensorflow/core/kernels/data/BUILD +++ b/tensorflow/core/kernels/data/BUILD @@ -80,16 +80,11 @@ tf_kernel_library( "//tensorflow/core/data:name_utils", "//tensorflow/core/data:serialization_utils", "//tensorflow/core/framework:dataset_options_proto_cc", - "//tensorflow/core/framework:types_proto_cc", "//tensorflow/core/util/tensor_bundle", "//tensorflow/core/util/tensor_bundle:naming", - "@com_google_absl//absl/log", - "@com_google_absl//absl/log:check", "@com_google_absl//absl/status", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", - "@com_google_absl//absl/strings:string_view", - "@com_google_absl//absl/types:span", ], ) diff --git a/tensorflow/core/kernels/data/cache_dataset_ops.cc b/tensorflow/core/kernels/data/cache_dataset_ops.cc index 7866214dac8ef6..9bfcbe60a7ffb1 100644 --- a/tensorflow/core/kernels/data/cache_dataset_ops.cc +++ b/tensorflow/core/kernels/data/cache_dataset_ops.cc @@ -22,13 +22,9 @@ limitations under the License. #include #include -#include "absl/log/check.h" -#include "absl/log/log.h" #include "absl/status/status.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 "tensorflow/core/data/global_shuffle_utils.h" #include "tensorflow/core/data/name_utils.h" #include "tensorflow/core/data/serialization_utils.h" @@ -37,7 +33,6 @@ limitations under the License. #include "tensorflow/core/framework/partial_tensor_shape.h" #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/framework/types.pb.h" #include "tensorflow/core/kernels/data/cache_ops.h" #include "tensorflow/core/kernels/data/iterator_ops.h" #include "tensorflow/core/lib/core/errors.h"