[pull] master from tensorflow:master

third_party/xla/MODULE.bazel

@@ -17,7 +17,7 @@ bazel_dep(name = "gutil", version = "20250502.0", repo_name = "com_google_gutil"
 bazel_dep(name = "jsoncpp", version = "1.9.6", repo_name = "jsoncpp_git")
 bazel_dep(name = "or-tools", version = "9.12", repo_name = "com_google_ortools")
 bazel_dep(name = "platforms", version = "1.0.0")
-bazel_dep(name = "protobuf", version = "31.1", repo_name = "com_google_protobuf")
+bazel_dep(name = "protobuf", version = "32.1", repo_name = "com_google_protobuf")
 bazel_dep(name = "pybind11_abseil", version = "202402.0")
 bazel_dep(name = "pybind11_bazel", version = "2.13.6")
 bazel_dep(name = "pybind11_protobuf", version = "0.0.0-20250210-f02a2b7")

third_party/xla/build_tools/ci/build.py

@@ -259,7 +259,7 @@ def commands(self) -> List[List[str]]:
     return cmds
 
 
-_CUDA_COMPUTE_CAPABILITIES = (60, 70, 80, 90, 100, 103)
+_CUDA_COMPUTE_CAPABILITIES = (60, 70, 80, 90, 100, 103, 120)
 
 
 def _tag_filters_only_for_compute_capability(

third_party/xla/xla/backends/cpu/benchmarks/e2e/gemma2/flax_2b/requirements.txt

@@ -18,7 +18,7 @@ opt_einsum==3.4.0
 optax==0.2.4
 orbax-checkpoint==0.10.2
 protobuf==6.33.5
-Pygments==2.18.0
+Pygments==2.20.0
 PyYAML==6.0.2
 rich==13.9.4
 scipy==1.14.1

third_party/xla/xla/backends/gpu/codegen/emitters/BUILD

@@ -330,7 +330,6 @@ cc_library(
         "//xla/codegen/emitters:utils",
         "//xla/codegen/emitters/ir:xla",
         "//xla/hlo/analysis:indexing_analysis",
-        "//xla/hlo/analysis:symbolic_map",
         "//xla/hlo/ir:hlo",
         "//xla/service:scatter_simplifier",
         "//xla/service/gpu:gpu_fusible",

third_party/xla/xla/backends/gpu/codegen/emitters/reduction.cc

@@ -267,7 +267,7 @@ SmallVector<Value> ReductionFusion::EmitterState::WriteToSharedMemory(
     const HloValueMap& values, std::optional<int> padding) {
   SmallVector<int64_t> shape;
   auto map = owner.GetSharedMemoryWriteMap(mlir_context);
-  for (auto result : map.GetSymbolicMap().GetResults()) {
+  for (auto result : map.GetAffineMap().getResults()) {
     shape.push_back(
         map.GetRangeEvaluator().ComputeExpressionRange(result).upper + 1);
   }
@@ -567,9 +567,9 @@ std::optional<IndexingMap> ReductionFusion::ComputeThreadIdToOutputIndexing(
   auto projected_indexing = ComputeReductionOutputIndexing(mlir_context);
   auto output_shape = reduction_dimensions_.GetOutputShape();
   CHECK_EQ(output_shape.size(),
-           projected_indexing.GetSymbolicMap().GetNumResults());
+           projected_indexing.GetAffineMap().getNumResults());
   for (auto [result, dim_size] : llvm::zip(
-           projected_indexing.GetSymbolicMap().GetResults(), output_shape)) {
+           projected_indexing.GetAffineMap().getResults(), output_shape)) {
     projected_indexing.AddConstraint(result, {0, dim_size - 1});
   }
   AddGroupIdConstraint(projected_indexing, root_index, groups_);
@@ -796,7 +796,7 @@ IndexingMap SmallColumnReductionFusion::ComputeReductionInputIndexing(
                     mlir_context);
 
   for (auto [result, dim_size] :
-       llvm::zip(map.GetSymbolicMap().GetResults(), input_shape_)) {
+       llvm::zip(map.GetAffineMap().getResults(), input_shape_)) {
     map.AddConstraint(result, {0, dim_size - 1});
   }
   return map;

third_party/xla/xla/backends/gpu/codegen/emitters/scatter.cc

@@ -55,7 +55,6 @@ limitations under the License.
 #include "xla/codegen/emitters/utils.h"
 #include "xla/hlo/analysis/indexing_analysis.h"
 #include "xla/hlo/analysis/indexing_map.h"
-#include "xla/hlo/analysis/symbolic_map.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
@@ -103,7 +102,6 @@ using mlir::func::FuncOp;
 using mlir::func::ReturnOp;
 using primitive_util::IsUnsignedIntegralType;
 
-constexpr int64_t kGpuGridDims = 6;
 constexpr int64_t kNumWarpsPerBlock = 4;
 constexpr int64_t kMaxVectorizedBits = 128;
 constexpr int64_t kScatterOperandIndex = 0;
@@ -431,9 +429,10 @@ void ScatterWithDistributedUpdates::ComputeIndexing(
   if (indices_map) {
     // Create a map from scatter update to scatter indices.
     *indices_map = IndexingMap{
-        SymbolicMap::Get(mlir_context, kGpuGridDims, /*num_symbols=*/1,
-                         {scatter_update_map.GetSymbolicMap().GetResult(0),
-                          CreateSymbolExpr(0, kGpuGridDims, mlir_context)}),
+        AffineMap::get(6, 1,
+                       {scatter_update_map.GetAffineMap().getResult(0),
+                        getAffineSymbolExpr(0, mlir_context)},
+                       mlir_context),
         DimVarsFromGPUGrid({num_warps_ * warp_size_, 1, 1, num_blocks_, 1, 1}),
         RangeVarsFromTensorSizes({description_.index_vector_length}),
         /*rt_vars=*/{}};
@@ -480,8 +479,8 @@ void EmitNaiveImplementation(ImplicitLocOpBuilder& b,
           .scatter_dims_to_operand_dims();
   MLIRContext* mlir_context = b.getContext();
   auto thread_id_to_update_id_map = IndexingMap(
-      SymbolicMap::Get(mlir_context, kGpuGridDims, /*num_symbols=*/0,
-                       {updates_map.GetSymbolicMap().GetResult(0)}),
+      AffineMap::get(6, 0, {updates_map.GetAffineMap().getResult(0)},
+                     mlir_context),
       updates_map.GetDimVars(),
       /*range_vars = */ {}, /*rt vars = */ {});
   Value thread_id_to_index_id_value =
@@ -676,13 +675,13 @@ absl::Status ScatterWithDistributedIndices::EmitEntryFunctionImpl(
   MLIRContext* mlir_context = b.getContext();
 
   auto thread_id_to_update_id_map = IndexingMap(
-      SymbolicMap::Get(mlir_context, kGpuGridDims, /*num_symbols=*/2,
-                       {indices_map.GetSymbolicMap().GetResult(0)}),
+      AffineMap::get(6, 2, {indices_map.GetAffineMap().getResult(0)},
+                     mlir_context),
       indices_map.GetDimVars(),
       /*range_vars = */
       {indices_map.GetRangeVars().begin(),
        indices_map.GetRangeVars().begin() + 2},
-      /*rt vars = */ {}, indices_map.GetSymbolicConstraints());
+      /*rt vars = */ {}, indices_map.GetConstraints());
 
   // Convert index_id_loop and index_vector_id to dimension variables.
   IndexingMap slice_indexing =

third_party/xla/xla/backends/gpu/codegen/emitters/transpose.cc

@@ -541,7 +541,7 @@ IndexingMap TransposeFusion::GetIndexing(bool input, const xla::Shape& shape,
       input ? ShapeUtil::MakeShape(shape.element_type(), input_shape_)
             : ShapeUtil::MakeShape(shape.element_type(), transpose_.dimensions);
   for (auto [size, dim] : llvm::zip(normalized_shape.dimensions(),
-                                    result.GetSymbolicMap().GetResults())) {
+                                    result.GetAffineMap().getResults())) {
     result.AddConstraint(dim, {0, size - 1});
   }
   result = ComposeIndexingMaps(

third_party/xla/xla/backends/gpu/codegen/triton/BUILD

@@ -296,36 +296,6 @@ cc_library(
     ],
 )
 
-xla_cc_test(
-    name = "fusion_emitter_deviceless_test",
-    srcs = ["fusion_emitter_deviceless_test.cc"],
-    tags = ["no_oss"],  # Doesn't pass in OSS when building with the `fusion_emitter_stub`.
-    deps = [
-        ":xtile_compiler",
-        "//xla:xla_proto_cc",
-        "//xla/hlo/analysis:symbolic_map",
-        "//xla/hlo/ir:hlo",
-        "//xla/hlo/testlib:filecheck",
-        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
-        "//xla/hlo/testlib:verified_hlo_module",
-        "//xla/service/gpu:backend_configs_cc",
-        "//xla/service/gpu:gpu_device_info_for_tests",
-        "//xla/service/gpu:target_constants",
-        "//xla/service/gpu/model:block_level_parameters",
-        "//xla/stream_executor:device_description",
-        "//xla/stream_executor/cuda:cuda_compute_capability",
-        "//xla/tests:xla_internal_test_main",
-        "//xla/tsl/platform:statusor",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:status_matchers",
-        "@com_google_absl//absl/strings:string_view",
-        "@com_google_googletest//:gtest",
-        "@llvm-project//llvm:TargetParser",
-        "@llvm-project//llvm:ir_headers",
-        "@llvm-project//mlir:IR",
-    ],
-)
-
 xla_test(
     name = "triton_gemm_fusion_test",
     srcs = ["triton_gemm_fusion_test.cc"],
@@ -379,35 +349,6 @@ xla_test(
     ],
 )
 
-xla_test(
-    name = "fusion_emitter_int4_device_test",
-    srcs = ["fusion_emitter_int4_device_test.cc"],
-    backends = [
-        "a100",
-        "h100",
-        "b200",
-        "amdgpu_any",
-    ],
-    shard_count = 10,
-    tags = ["no_mac"],
-    deps = [
-        "//xla:autotuning_proto_cc",
-        "//xla:error_spec",
-        "//xla:xla_proto_cc",
-        "//xla/backends/gpu/tests:hlo_pjrt_gpu_test_base",
-        "//xla/hlo/ir:hlo",
-        "//xla/hlo/testlib:filecheck",
-        "//xla/service/gpu:backend_configs_cc",
-        "//xla/stream_executor:device_description",
-        "//xla/tests:hlo_pjrt_interpreter_reference_mixin",
-        "//xla/tests:xla_internal_test_main",  # fixdeps: keep
-        "//xla/tsl/platform:statusor",
-        "@com_google_absl//absl/strings",
-        "@com_google_googletest//:gtest",
-        "@tsl//tsl/platform:path",
-    ],
-)
-
 xla_test(
     name = "dot_algorithms_test",
     srcs = ["dot_algorithms_test.cc"],
@@ -472,67 +413,6 @@ xla_test(
     ],
 )
 
-xla_test(
-    name = "fusion_emitter_device_test",
-    srcs = ["fusion_emitter_device_test.cc"],
-    backends = [
-        "a100",
-        "h100",
-        "b200",
-        "amdgpu_any",
-    ],
-    shard_count = 10,
-    tags = ["no_mac"],
-    deps = [
-        ":support",
-        ":test_utils",
-        ":xtile_compiler",
-        ":xtile_test_base",
-        "//xla:autotuning_proto_cc",
-        "//xla:error_spec",
-        "//xla:literal",
-        "//xla:literal_util",
-        "//xla:shape_util",
-        "//xla:types",
-        "//xla:util",
-        "//xla:xla_data_proto_cc",
-        "//xla:xla_proto_cc",
-        "//xla/backends/gpu/tests:gpu_pjrt_codegen_test",
-        "//xla/hlo/analysis:symbolic_map",
-        "//xla/hlo/ir:hlo",
-        "//xla/hlo/testlib:filecheck",
-        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
-        "//xla/hlo/testlib:verified_hlo_module",
-        "//xla/service:algorithm_util",
-        "//xla/service/gpu:backend_configs_cc",
-        "//xla/service/gpu:gpu_device_info_for_tests",
-        "//xla/service/gpu:target_constants",
-        "//xla/service/gpu/model:block_level_parameters",
-        "//xla/stream_executor:device_description",
-        "//xla/stream_executor/cuda:cuda_compute_capability",
-        "//xla/stream_executor/rocm:rocm_compute_capability",
-        "//xla/tests:hlo_pjrt_interpreter_reference_mixin",
-        "//xla/tests:test_utils",
-        "//xla/tests:xla_internal_test_main",  # fixdeps: keep
-        "//xla/tsl/lib/core:status_test_util",
-        "//xla/tsl/platform:env",
-        "//xla/tsl/platform:errors",
-        "//xla/tsl/platform:statusor",
-        "//xla/tsl/platform:test",
-        "@com_google_absl//absl/log",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:status_matchers",
-        "@com_google_absl//absl/strings",
-        "@com_google_googletest//:gtest",
-        "@eigen_archive//:eigen3",
-        "@llvm-project//llvm:TargetParser",
-        "@llvm-project//llvm:ir_headers",
-        "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:Pass",
-        "@tsl//tsl/platform:path",
-    ],
-)
-
 cc_library(
     name = "test_utils",
     testonly = True,
@@ -632,89 +512,6 @@ cc_library(
     ],
 )
 
-xla_test(
-    name = "fusion_emitter_large_test",
-    srcs = ["fusion_emitter_large_test.cc"],
-    backend_tags = {
-        "h100": [
-            # Needs full H100, otherwise we run OOM.
-            "full",
-        ],
-    },
-    backends = [
-        "a100",
-        "h100",
-        "b200",
-        "amdgpu_any",
-    ],
-    shard_count = 3,
-    tags = [
-        "no_mac",
-        "no_oss",  # requires-mem:16g tag doesn't work in open source
-        "nozapfhahn",  # Times out under coverage
-    ] + if_google([
-        "requires-mem:16g",
-    ]),
-    deps = [
-        "//xla:error_spec",
-        "//xla:xla_proto_cc",
-        "//xla/tests:hlo_pjrt_interpreter_reference_mixin",
-        "//xla/tests:hlo_pjrt_test_base",
-        "//xla/tests:xla_internal_test_main",  # fixdeps: keep
-        "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/strings:string_view",
-        "@com_google_googletest//:gtest",
-    ],
-)
-
-xla_test(
-    name = "fusion_emitter_parametrized_test",
-    srcs = ["fusion_emitter_parametrized_test.cc"],
-    backends = [
-        "a100",
-        "h100",
-        "b200",
-        "amdgpu_any",
-    ],
-    tags = ["no_mac"],
-    use_legacy_runtime = True,
-    deps = [
-        ":support",
-        ":test_utils",
-        "//xla:comparison_util",
-        "//xla:error_spec",
-        "//xla:xla_data_proto_cc",
-        "//xla:xla_proto_cc",
-        "//xla/backends/gpu/tests:gpu_codegen_test",
-        "//xla/hlo/ir:hlo",
-        "//xla/stream_executor:device_description",
-        "//xla/tests:xla_internal_test_main",  # fixdeps: keep
-        "//xla/tsl/lib/core:status_test_util",
-        "//xla/tsl/platform:statusor",
-        "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/strings",
-        "@com_google_googletest//:gtest",
-    ],
-)
-
-xla_cc_test(
-    name = "fusion_emitter_shared_dialect_test",
-    srcs = ["fusion_emitter_shared_dialect_test.cc"],
-    tags = ["no_mac"],
-    deps = [
-        ":xtile_test_base",
-        "//xla/hlo/ir:hlo",
-        "//xla/hlo/parser:hlo_parser",
-        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
-        "//xla/service/gpu/model:block_level_parameters",
-        "//xla/tests:xla_internal_test_main",  # fixdeps: keep
-        "//xla/tsl/lib/core:status_test_util",
-        "//xla/tsl/platform:statusor",
-        "@com_google_absl//absl/strings:string_view",
-        "@com_google_googletest//:gtest",
-    ],
-)
-
 cc_library(
     name = "support",
     srcs = [

third_party/xla/xla/backends/gpu/codegen/triton/support_legacy.cc

@@ -339,13 +339,6 @@ CodegenDecision CanTritonHandleGEMM(
     return decision;
   }
 
-  const DotDimensionNumbers& dim_numbers = dot.dot_dimension_numbers();
-
-  // TODO(b/269580541): support multiple batch dimensions.
-  if (dim_numbers.lhs_batch_dimensions().size() > 1) {
-    return CodegenDecision::Forbid("Multiple batch dimensions.");
-  }
-
   return CodegenDecision::Allow();
 }
 

third_party/xla/xla/backends/gpu/codegen/triton/support_legacy_test.cc

@@ -448,10 +448,8 @@ ENTRY e {
                               kHloTest, /*data_type=*/{}, HloOpcode::kDot));
   const se::DeviceDescription dev_info =
       TestGpuDeviceInfo::RTXA6000DeviceInfo(GetComputeCapability());
-  EXPECT_THAT(legacy_triton::IsTritonSupportedInstruction(
-                  ti.Instruction(), GetComputeCapability())
-                  .Explain(),
-              ::testing::HasSubstr("Multiple batch dimensions"));
+  EXPECT_TRUE(legacy_triton::IsTritonSupportedInstruction(
+      ti.Instruction(), GetComputeCapability()));
   auto block_level_parameters =
       BlockLevelParameters::FromBlockLevelFusionConfig(
           ti.TritonFusion()