[pull] master from tensorflow:master

third_party/xla/third_party/triton/BUILD

@@ -0,0 +1,17 @@
+# copybara:uncomment package(default_applicable_licenses = ["//third_party/tensorflow:license"])
+
+filegroup(
+    name = "patch_files",
+    srcs = glob([
+        "common/**",
+        "oss_only/**",
+    ]),
+    visibility = ["//visibility:public"],
+)
+
+filegroup(
+    name = "workspace",
+    srcs = ["workspace.bzl"],
+    # copybara:uncomment compatible_with = ["//buildenv/target:non_prod"],
+    visibility = ["//visibility:public"],
+)

third_party/xla/third_party/triton/oss_only/build_files.patch

@@ -5,7 +5,7 @@ diff --git b/BUILD b/BUILD
 new file mode 100644
 --- /dev/null
 +++ b/BUILD
-@@ -0,0 +1,1290 @@
+@@ -0,0 +1,1300 @@
 +# This package imports OpenAI's Triton (https://github.com/triton-lang/triton).
 +
 +# copybara:uncomment load("@rules:copybara.bzl", "copybara_config_test")
@@ -71,9 +71,20 @@ new file mode 100644
 +# generate version.h for building TritonPluginUtils
 +genrule(
 +    name = "triton_version_h_gen",
-+    srcs = ["include/triton/Version.h.in"],
++    srcs = [
++        "include/triton/Version.h.in",
++        "@xla//third_party/triton:workspace",
++    ],
 +    outs = ["include/triton/Version.h"],
-+    cmd = "sed 's/@TRITON_VERSION@//g' $< > $@",
++    cmd = "TRITON_COMMIT=$$(sed -n 's/.*TRITON_COMMIT = \"\\(.*\\)\".*/\\1/p' $(location @xla//third_party/triton:workspace)) && " +
++          "sed \"s/@TRITON_VERSION@/$$TRITON_COMMIT/g\" $(location include/triton/Version.h.in) > $@",
++)
++
++cc_library(
++    name = "triton_version",
++    hdrs = ["include/triton/Version.h"],
++    strip_include_prefix = "include",
++    deps = [":triton_version_h_gen"],
 +)
 +
 +td_library(
@@ -1088,14 +1099,13 @@ new file mode 100644
 +    srcs = ["lib/Tools/PluginUtils.cpp"],
 +    hdrs = [
 +        "include/triton/Tools/PluginUtils.h",
-+        "include/triton/Version.h",
 +    ],
 +    copts = _no_unused_variable,
 +    includes = ["include"],
 +    deps = [
 +        ":Dump",
 +        ":GetEnv",
-+        ":triton_version_h_gen",
++        ":triton_version",
 +        "@llvm-project//llvm:Support",
 +        "@llvm-project//mlir:IR",
 +        "@llvm-project//mlir:Pass",
@@ -1266,7 +1276,7 @@ new file mode 100644
 +#         "@triton//:leakr_badwords.dic",
 +#         "@triton//patches:patch_files",
 +#         "//third_party/xla:copybara_library",
-+#         "//third_party/xla/third_party/triton:patch_files",
++#         "@xla//third_party/triton:patch_files",
 +#     ],
 +# )
 +# copybara:uncomment_end

third_party/xla/xla/backends/autotuner/BUILD

@@ -491,3 +491,19 @@ tf_proto_library(
         "@com_google_protobuf//:duration_proto",
     ],
 )
+
+xla_cc_test(
+    name = "triton_version_test",
+    srcs = ["triton_version_test.cc"],
+    data = [
+        "//third_party/triton:workspace",
+    ],
+    env = {
+        "WORKSPACE_BZL_PATH": "$(rootpath //third_party/triton:workspace)",
+    },
+    deps = [
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@triton//:triton_version",
+    ],
+)

third_party/xla/xla/backends/autotuner/triton_version_test.cc

@@ -0,0 +1,75 @@
+/* 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 <cstdlib>
+#include <fstream>
+#include <iterator>
+#include <optional>
+#include <string>
+
+#include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
+#include "triton/Version.h"
+
+namespace {
+
+constexpr absl::string_view kTritonCommitPattern = "TRITON_COMMIT = \"";
+
+std::optional<std::string> ReadFile(const std::string& path) {
+  std::ifstream file(path);
+  if (!file.is_open()) {
+    return std::nullopt;
+  }
+  return std::string((std::istreambuf_iterator<char>(file)),
+                     std::istreambuf_iterator<char>());
+}
+
+std::optional<absl::string_view> ExtractExpectedVersion(
+    absl::string_view content) {
+  size_t pos = content.find(kTritonCommitPattern);
+  if (pos == absl::string_view::npos) {
+    return std::nullopt;
+  }
+  size_t start = pos + kTritonCommitPattern.size();
+  size_t end = content.find('\"', start);
+  if (end == absl::string_view::npos) {
+    return std::nullopt;
+  }
+  return content.substr(start, end - start);
+}
+
+TEST(VersionTest, MatchExpected) {
+  const char* env_path = std::getenv("WORKSPACE_BZL_PATH");
+  ASSERT_NE(env_path, nullptr)
+      << "WORKSPACE_BZL_PATH environment variable is not set";
+  std::string path(env_path);
+  ASSERT_FALSE(path.empty())
+      << "WORKSPACE_BZL_PATH environment variable is empty";
+
+  std::optional<std::string> content = ReadFile(path);
+  ASSERT_TRUE(content.has_value()) << "Failed to read file: " << path;
+
+  std::optional<absl::string_view> expected_version =
+      ExtractExpectedVersion(*content);
+  ASSERT_TRUE(expected_version.has_value())
+      << "Failed to extract TRITON_COMMIT from " << path;
+  ASSERT_FALSE(expected_version->empty())
+      << "Expected version should not be empty";
+  ASSERT_FALSE(std::string(TRITON_VERSION).empty())
+      << "Triton version should not be empty";
+  EXPECT_EQ(TRITON_VERSION, expected_version.value());
+}
+
+}  // namespace

third_party/xla/xla/backends/cpu/benchmarks/elementwise_benchmark_test.cc

@@ -90,6 +90,58 @@ static void BM_AddBF16(benchmark::State& state, HloBenchmarkOptions options) {
       RunHloBenchmark(state, hlo, args, {{"$d0", absl::StrCat(d0)}}, options));
 }
 
+static void BM_AddU32(benchmark::State& state, HloBenchmarkOptions options) {
+  int64_t d0 = state.range(0);
+
+  absl::string_view hlo = R"(
+    HloModule add_u32_$d0
+
+    ENTRY e {
+      p0 = u32[1,2,1,$d0,256] parameter(0)
+      p1 = u32[1,2,1,$d0,256] parameter(1)
+      ROOT add = u32[1,2,1,$d0,256] add(p0, p1)
+    }
+  )";
+
+  std::minstd_rand0 engine;
+
+  auto shape = ShapeUtil::MakeShape(U32, {1, 2, 1, d0, 256});
+  ASSERT_OK_AND_ASSIGN(Literal p0, LiteralUtil::CreateRandomLiteral<U32>(
+                                       shape, &engine, 100, 1));
+  ASSERT_OK_AND_ASSIGN(Literal p1, LiteralUtil::CreateRandomLiteral<U32>(
+                                       shape, &engine, 100, 1));
+
+  std::vector<const Literal*> args = {&p0, &p1};
+  CHECK_OK(
+      RunHloBenchmark(state, hlo, args, {{"$d0", absl::StrCat(d0)}}, options));
+}
+
+static void BM_AddU64(benchmark::State& state, HloBenchmarkOptions options) {
+  int64_t d0 = state.range(0);
+
+  absl::string_view hlo = R"(
+    HloModule add_u64_$d0
+
+    ENTRY e {
+      p0 = u64[1,2,1,$d0,256] parameter(0)
+      p1 = u64[1,2,1,$d0,256] parameter(1)
+      ROOT add = u64[1,2,1,$d0,256] add(p0, p1)
+    }
+  )";
+
+  std::minstd_rand0 engine;
+
+  auto shape = ShapeUtil::MakeShape(U64, {1, 2, 1, d0, 256});
+  ASSERT_OK_AND_ASSIGN(Literal p0, LiteralUtil::CreateRandomLiteral<U64>(
+                                       shape, &engine, 100, 1));
+  ASSERT_OK_AND_ASSIGN(Literal p1, LiteralUtil::CreateRandomLiteral<U64>(
+                                       shape, &engine, 100, 1));
+
+  std::vector<const Literal*> args = {&p0, &p1};
+  CHECK_OK(
+      RunHloBenchmark(state, hlo, args, {{"$d0", absl::StrCat(d0)}}, options));
+}
+
 static void BM_UnaryOp(benchmark::State& state,
                        const HloBenchmarkOptions& options, PrimitiveType type,
                        HloOpcode op) {
@@ -173,6 +225,8 @@ static void BM_ConvertF32ToBF16(benchmark::State& state,
 
 BENCHMARK_SIZES(BM_AddF32);
 BENCHMARK_SIZES(BM_AddBF16);
+BENCHMARK_SIZES(BM_AddU32);
+BENCHMARK_SIZES(BM_AddU64);
 BENCHMARK_SIZES(BM_ConvertF32ToBF16);
 
 #define BM_UNARY_OP(OP, TYPE)                              \

third_party/xla/xla/backends/cpu/benchmarks/reduction_benchmark_test.cc

@@ -15,6 +15,7 @@ limitations under the License.
 
 #include <cstdint>
 #include <random>
+#include <utility>
 #include <vector>
 
 #include "absl/strings/str_cat.h"
@@ -24,6 +25,7 @@ limitations under the License.
 #include "xla/backends/cpu/benchmarks/multi_benchmark_config.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
+#include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/tsl/platform/logging.h"
 #include "xla/tsl/platform/test_benchmark.h"
@@ -123,6 +125,70 @@ static void BM_ReduceAddF64(benchmark::State& state,
       RunHloBenchmark(state, hlo, args, {{"$d0", absl::StrCat(d0)}}, options));
 }
 
+static void BM_ReduceAddU32(benchmark::State& state,
+                            HloBenchmarkOptions options) {
+  int64_t d0 = state.range(0);
+
+  absl::string_view hlo = R"(
+    HloModule reduce_add_u32_$d0
+
+    add {
+      p0 = u32[] parameter(0)
+      p1 = u32[] parameter(1)
+      ROOT add = u32[] add(p0, p1)
+    }
+
+    ENTRY e {
+      p0 = u32[1,2,1,$d0,256] parameter(0)
+      c0 = u32[] constant(0)
+      ROOT reduce = u32[1,2] reduce(p0, c0), dimensions={2,3,4}, to_apply=add
+    }
+  )";
+
+  std::minstd_rand0 engine;
+
+  Shape shape = ShapeUtil::MakeShape(U32, {1, 2, 1, d0, 256});
+  auto p0_or = LiteralUtil::CreateRandomLiteral<U32>(shape, &engine, 100, 1);
+  CHECK_OK(p0_or.status());
+  Literal p0 = std::move(p0_or).value();
+
+  std::vector<const Literal*> args = {&p0};
+  CHECK_OK(
+      RunHloBenchmark(state, hlo, args, {{"$d0", absl::StrCat(d0)}}, options));
+}
+
+static void BM_ReduceAddU64(benchmark::State& state,
+                            HloBenchmarkOptions options) {
+  int64_t d0 = state.range(0);
+
+  absl::string_view hlo = R"(
+    HloModule reduce_add_u64_$d0
+
+    add {
+      p0 = u64[] parameter(0)
+      p1 = u64[] parameter(1)
+      ROOT add = u64[] add(p0, p1)
+    }
+
+    ENTRY e {
+      p0 = u64[1,2,1,$d0,256] parameter(0)
+      c0 = u64[] constant(0)
+      ROOT reduce = u64[1,2] reduce(p0, c0), dimensions={2,3,4}, to_apply=add
+    }
+  )";
+
+  std::minstd_rand0 engine;
+
+  Shape shape = ShapeUtil::MakeShape(U64, {1, 2, 1, d0, 256});
+  auto p0_or = LiteralUtil::CreateRandomLiteral<U64>(shape, &engine, 100, 1);
+  CHECK_OK(p0_or.status());
+  Literal p0 = std::move(p0_or).value();
+
+  std::vector<const Literal*> args = {&p0};
+  CHECK_OK(
+      RunHloBenchmark(state, hlo, args, {{"$d0", absl::StrCat(d0)}}, options));
+}
+
 static void BM_SumOfSquaresF32(benchmark::State& state,
                                HloBenchmarkOptions options) {
   int64_t d0 = state.range(0);
@@ -348,6 +414,8 @@ static void BM_ReduceWindowAddF32OverlappingWindows(
 BENCHMARK_SIZES(BM_ReduceAddF32);
 BENCHMARK_SIZES(BM_ReduceAddBF16);
 BENCHMARK_SIZES(BM_ReduceAddF64);
+BENCHMARK_SIZES(BM_ReduceAddU32);
+BENCHMARK_SIZES(BM_ReduceAddU64);
 BENCHMARK_SIZES(BM_SumOfSquaresF32);
 
 XLA_CPU_BENCHMARK(BM_ReduceAddF32OverDimension)