[pull] master from tensorflow:master

third_party/xla/MODULE.bazel

@@ -108,6 +108,7 @@ use_repo(
     "KleidiAI",
     "XNNPACK",
     "cpuinfo",
+    "cuda_tile",
     "cudnn_frontend_archive",
     "cutlass_archive",
     "cutlass_cutedsl_runtime",
@@ -147,6 +148,7 @@ use_repo(
     "slinky",
     "spdlog",
     "stablehlo",
+    "tensor_ir",
     "transformer_engine",
     "triton",
     "uv",

third_party/xla/third_party/extensions/third_party.bzl

@@ -19,6 +19,7 @@ load("//third_party/benchmark:workspace.bzl", benchmark = "repo")
 load("//third_party/brotli:workspace.bzl", brotli = "repo")
 load("//third_party/compute_library:workspace.bzl", compute_library = "repo")
 load("//third_party/cpuinfo:workspace.bzl", cpuinfo = "repo")
+load("//third_party/cuda_tile:workspace.bzl", cuda_tile = "repo")
 load("//third_party/cudnn_frontend:workspace.bzl", cudnn_frontend = "repo")
 load("//third_party/cutlass:workspace.bzl", cutlass = "repo")
 load("//third_party/cutlass_cutedsl_runtime:workspace.bzl", cutlass_cutedsl_runtime = "repo")
@@ -59,6 +60,7 @@ load("//third_party/slinky:workspace.bzl", slinky = "repo")
 load("//third_party/spdlog:workspace.bzl", spdlog = "repo")
 load("//third_party/sqlite:workspace.bzl", sqlite = "repo")
 load("//third_party/stablehlo:workspace.bzl", stablehlo = "repo")
+load("//third_party/tensor_ir:workspace.bzl", tensor_ir = "repo")
 load("//third_party/tensorrt:workspace.bzl", tensorrt = "repo")
 load("//third_party/transformer_engine:workspace.bzl", transformer_engine = "repo")
 load("//third_party/triton:workspace.bzl", triton = "repo")
@@ -72,6 +74,7 @@ def _third_party_ext_impl(mctx):  # @unused
     brotli()  # Needed for Riegeli
     compute_library()
     cpuinfo()
+    cuda_tile()
     cudnn_frontend()
     cutlass()
     cutlass_cutedsl_runtime()
@@ -112,6 +115,7 @@ def _third_party_ext_impl(mctx):  # @unused
     spdlog()
     sqlite()
     stablehlo()
+    tensor_ir()
     tensorrt()
     transformer_engine()
     triton()

third_party/xla/xla/backends/cpu/codegen/BUILD

@@ -228,7 +228,6 @@ cc_library(
         "//xla/tsl/framework/mlir:status_scoped_diagnostic_handler",
         "@com_google_absl//absl/base:config",
         "@com_google_absl//absl/functional:any_invocable",
-        "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",

third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc

@@ -22,7 +22,6 @@ limitations under the License.
 #include <utility>
 
 #include "absl/base/config.h"  // IWYU pragma: keep
-#include "absl/functional/function_ref.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
@@ -133,6 +132,14 @@ limitations under the License.
 namespace xla::cpu {
 namespace {
 
+emitters::SimplifyArithPassOptions GetSimplifyArithPassOptions(
+    bool fast_min_max) {
+  emitters::SimplifyArithPassOptions options;
+  options.fast_min_max_ = fast_min_max;
+  options.explicit_nan_propagation_ = false;
+  return options;
+}
+
 absl::Status RunPassPipeline(mlir::ModuleOp module, mlir::PassManager& pm,
                              mlir::interpreter::MlirCompilationTrace* trace,
                              int32_t verification_level) {
@@ -231,11 +238,8 @@ void AddScalarOptimizationPasses(mlir::OpPassManager& pm,
 // These passes are primarily responsible for lowering individual ops to
 // their LLVM equivalent.
 void AddGenericLoweringPasses(mlir::OpPassManager& pm, bool fast_min_max) {
-  emitters::SimplifyArithPassOptions simplify_arith_options;
-  simplify_arith_options.fast_min_max_ = fast_min_max;
-  simplify_arith_options.explicit_nan_propagation_ = false;
-  pm.addNestedPass<mlir::func::FuncOp>(
-      emitters::createSimplifyArithPass(simplify_arith_options));
+  pm.addNestedPass<mlir::func::FuncOp>(emitters::createSimplifyArithPass(
+      GetSimplifyArithPassOptions(fast_min_max)));
   pm.addPass(emitters::createExpandIntegerPowerPass());
   pm.addPass(emitters::createSimplifyAffinePass());
   pm.addPass(mlir::createCanonicalizerPass());
@@ -405,9 +409,13 @@ void AddNewXtileToVectorPasses(mlir::OpPassManager& pm) {
   pm.addPass(xtile::createLegalizeUnsignedIntegersAsSignlessPass());
   pm.addPass(mlir::createCanonicalizerPass());
   pm.addPass(cpu::createVectorizeXTilePass());
-
   pm.addPass(cpu::createLowerXTileEntryPass());
 
+  pm.addNestedPass<mlir::func::FuncOp>(emitters::createSimplifyArithPass(
+      GetSimplifyArithPassOptions(/*fast_min_max=*/false)));
+  pm.addPass(mlir::createCanonicalizerPass());
+  pm.addPass(mlir::createCSEPass());
+
   pm.addNestedPass<mlir::func::FuncOp>(
       mlir::stablehlo::createStablehloTargetIndependentOptimizationPass());
 

third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/minimum_mask.hlo

@@ -67,23 +67,17 @@ fusion {
 // NEW-VECTOR:           %[[CMP0:.*]] = arith.cmpi sge, %[[BOUNDS]]#0, %{{.*}} : index
 // NEW-VECTOR:           %[[CMP1:.*]] = arith.cmpi sge, %[[BOUNDS]]#1, %{{.*}} : index
 // NEW-VECTOR:           %[[COND:.*]] = arith.andi %[[CMP0]], %[[CMP1]] : i1
-// NEW-VECTOR:           %[[R0:.*]] = scf.if %[[COND]] -> (vector<16x32xf32>) {
+// NEW-VECTOR:           %[[R:.*]]:2 = scf.if %[[COND]] -> (vector<16x32xf32>, vector<16x32xf32>) {
 // NEW-VECTOR:             %[[R0_IN:.*]] = vector.transfer_read %[[ARG0]][%{{.*}}, %{{.*}}], %[[PAD]] {in_bounds = [true, true]} : memref<100x100xf32>, vector<16x32xf32>
-// NEW-VECTOR:             scf.yield %[[R0_IN]]
+// NEW-VECTOR:             %[[R1_IN:.*]] = vector.transfer_read %[[ARG1]][%{{.*}}, %{{.*}}], %[[PAD]] {in_bounds = [true, true]} : memref<100x100xf32>, vector<16x32xf32>
+// NEW-VECTOR:             scf.yield %[[R0_IN]], %[[R1_IN]]
 // NEW-VECTOR:           } else {
 // NEW-VECTOR:             %[[MASK0:.*]] = vector.create_mask %{{.*}}, %{{.*}} : vector<16x32xi1>
 // NEW-VECTOR:             %[[R0_MASK:.*]] = vector.transfer_read %[[ARG0]][%{{.*}}, %{{.*}}], %[[PAD]], %[[MASK0]] : memref<100x100xf32>, vector<16x32xf32>
-// NEW-VECTOR:             scf.yield %[[R0_MASK]]
-// NEW-VECTOR:           }
-// NEW-VECTOR:           %[[R1:.*]] = scf.if %[[COND]] -> (vector<16x32xf32>) {
-// NEW-VECTOR:             %[[R1_IN:.*]] = vector.transfer_read %[[ARG1]][%{{.*}}, %{{.*}}], %[[PAD]] {in_bounds = [true, true]} : memref<100x100xf32>, vector<16x32xf32>
-// NEW-VECTOR:             scf.yield %[[R1_IN]]
-// NEW-VECTOR:           } else {
-// NEW-VECTOR:             %[[MASK1:.*]] = vector.create_mask %{{.*}}, %{{.*}} : vector<16x32xi1>
-// NEW-VECTOR:             %[[R1_MASK:.*]] = vector.transfer_read %[[ARG1]][%{{.*}}, %{{.*}}], %[[PAD]], %[[MASK1]] : memref<100x100xf32>, vector<16x32xf32>
-// NEW-VECTOR:             scf.yield %[[R1_MASK]]
+// NEW-VECTOR:             %[[R1_MASK:.*]] = vector.transfer_read %[[ARG1]][%{{.*}}, %{{.*}}], %[[PAD]], %[[MASK0]] : memref<100x100xf32>, vector<16x32xf32>
+// NEW-VECTOR:             scf.yield %[[R0_MASK]], %[[R1_MASK]]
 // NEW-VECTOR:           }
-// NEW-VECTOR:           %[[MIN:.*]] = arith.minimumf %[[R0]], %[[R1]]
+// NEW-VECTOR:           %[[MIN:.*]] = arith.minimumf %[[R]]#0, %[[R]]#1
 // NEW-VECTOR:           scf.if %[[COND]] {
 // NEW-VECTOR:             vector.transfer_write %[[MIN]], %[[ARG2]][%{{.*}}, %{{.*}}] {in_bounds = [true, true]} : vector<16x32xf32>, memref<100x100xf32>
 // NEW-VECTOR:           } else {

third_party/xla/xla/backends/cpu/codegen/tiled/transforms/BUILD

@@ -89,6 +89,7 @@ cc_library(
         "//xla/codegen/xtile/ir:xtile",
         "//xla/codegen/xtile/ir/transforms:passes",
         "//xla/hlo/analysis:indexing_analysis",
+        "//xla/hlo/analysis:interval",
         "//xla/hlo/analysis:symbolic_map",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/log:check",

third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/vectorize_xtile.mlir

@@ -626,7 +626,7 @@ func.func @test_extract_aligned(%arg0: memref<128xf32>, %arg1: index) -> tensor<
   %0 = xtile.extract %arg0[%c0] [8] [1] : memref<128xf32> -> tensor<8xf32>
   return %0 : tensor<8xf32>
 }
-// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8){{.*}}">
+// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8), domain: d0 in [0, 0]">
 // CHECK-LABEL: @test_extract_aligned
 // CHECK-DAG: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32
 // CHECK-DAG: %[[INDEXING:[^:]+]] = xla.apply_indexing #indexing_map(%{{.*}})
@@ -642,11 +642,12 @@ func.func @test_extract_aligned(%arg0: memref<128xf32>, %arg1: index) -> tensor<
 
 // -----
 
-func.func @test_extract_unaligned(%arg0: memref<128xf32>, %arg1: index) -> tensor<8xf32> {
+func.func @test_extract_unaligned(%arg0: memref<128xf32>,
+    %arg1: index {xla.range = [0 : index, 12 : index]}) -> tensor<8xf32> {
   %0 = xtile.extract %arg0[%arg1] [8] [1] : memref<128xf32> -> tensor<8xf32>
   return %0 : tensor<8xf32>
 }
-// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8){{.*}}">
+// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8), domain: d0 in [0, 12]">
 // CHECK-LABEL: @test_extract_unaligned
 // CHECK-DAG: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32
 // CHECK-DAG: %[[INDEXING:[^:]+]] = xla.apply_indexing #indexing_map(%{{.*}})
@@ -667,7 +668,7 @@ func.func @test_insert_aligned(%arg0: tensor<8xf32>, %arg1: memref<128xf32>) {
   xtile.insert %arg0 into %arg1[%c0] [8] [1] : tensor<8xf32> -> memref<128xf32>
   return
 }
-// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8){{.*}}">
+// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8), domain: d0 in [0, 0]">
 // CHECK-LABEL: @test_insert_aligned
 // CHECK: %[[INDEXING:[^:]+]] = xla.apply_indexing #indexing_map(%{{.*}})
 // CHECK: %[[COND:.*]] = arith.cmpi sge, %[[INDEXING]], %{{.*}} : index
@@ -680,11 +681,12 @@ func.func @test_insert_aligned(%arg0: tensor<8xf32>, %arg1: memref<128xf32>) {
 
 // -----
 
-func.func @test_insert_unaligned(%arg0: tensor<8xf32>, %arg1: memref<128xf32>, %arg2: index) {
+func.func @test_insert_unaligned(%arg0: tensor<8xf32>, %arg1: memref<128xf32>,
+    %arg2: index {xla.range = [0 : index, 42 : index]}) {
   xtile.insert %arg0 into %arg1[%arg2] [8] [1] : tensor<8xf32> -> memref<128xf32>
   return
 }
-// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8){{.*}}">
+// CHECK: #indexing_map = #xla.indexing_map<"(d0) -> (-d0 + 128 - 8), domain: d0 in [0, 42]">
 // CHECK-LABEL: @test_insert_unaligned
 // CHECK: %[[INDEXING:[^:]+]] = xla.apply_indexing #indexing_map(%{{.*}})
 // CHECK: %[[COND:.*]] = arith.cmpi sge, %[[INDEXING]], %{{.*}} : index

third_party/xla/xla/backends/cpu/codegen/tiled/transforms/vectorize_xtile.cc

@@ -54,6 +54,7 @@ limitations under the License.
 #include "xla/codegen/xtile/ir/xtile_dialect.h"
 #include "xla/codegen/xtile/ir/xtile_ops.h"
 #include "xla/hlo/analysis/indexing_map.h"
+#include "xla/hlo/analysis/interval.h"
 #include "xla/hlo/analysis/symbolic_expr.h"
 #include "xla/hlo/analysis/symbolic_map.h"
 #include "xla/util.h"
@@ -74,6 +75,17 @@ namespace shlo = ::mlir::stablehlo;
 using ::mlir::Value;
 using ::mlir::ValueRange;
 
+std::vector<xla::IndexingMap::Variable> GetVars(ValueRange values) {
+  std::vector<xla::IndexingMap::Variable> vars;
+  vars.reserve(values.size());
+  for (Value offset : values) {
+    vars.push_back(xla::IndexingMap::Variable{GetRange(offset).value_or(
+        Interval{std::numeric_limits<int64_t>::min(),
+                 std::numeric_limits<int64_t>::max()})});
+  }
+  return vars;
+}
+
 xla::SymbolicMap GetBoundsCheckSymbolicMap(mlir::MLIRContext* ctx,
                                            llvm::ArrayRef<int64_t> memref_shape,
                                            llvm::ArrayRef<int64_t> tile_shape) {
@@ -92,17 +104,12 @@ Value GetIsInBoundsCondition(mlir::OpBuilder& builder, mlir::Location loc,
                              ValueRange offsets, Value memref,
                              llvm::ArrayRef<int64_t> tile_shape) {
   auto memref_shape = mlir::cast<mlir::MemRefType>(memref.getType()).getShape();
-  int rank = memref_shape.size();
 
   xla::SymbolicMap symbolic_map =
       GetBoundsCheckSymbolicMap(builder.getContext(), memref_shape, tile_shape);
 
-  std::vector<xla::IndexingMap::Variable> vars(
-      rank, xla::IndexingMap::Variable{std::numeric_limits<int64_t>::min(),
-                                       std::numeric_limits<int64_t>::max()});
-
   xla::IndexingMap indexing_map(symbolic_map,
-                                /*dimensions=*/std::move(vars),
+                                /*dimensions=*/GetVars(offsets),
                                 /*range_vars=*/{}, /*rt_vars=*/{});
 
   auto apply_indexing =

third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD

@@ -0,0 +1,57 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":friends"],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    includes = [
+        "//xla:friends",
+    ],
+)
+
+cc_library(
+    name = "conversion",
+    srcs = ["conversion.cc"],
+    hdrs = ["conversion.h"],
+    deps = [
+        "//xla:comparison_util",
+        "//xla:literal",
+        "//xla:literal_util",
+        "//xla:shape_util",
+        "//xla:util",
+        "//xla:xla_data_proto_cc",
+        "//xla/hlo/ir:hlo",
+        "//xla/service/llvm_ir:llvm_util",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/cleanup",
+        "@com_google_absl//absl/log",
+        "@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/types:span",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+        "@tensor_ir//:NVTensorIRDialect",
+    ],
+)
+
+cc_library(
+    name = "support",
+    srcs = ["support.cc"],
+    hdrs = ["support.h"],
+    deps = [
+        "//xla:comparison_util",
+        "//xla:xla_data_proto_cc",
+        "//xla/hlo/ir:hlo",
+        "//xla/service:decision",
+        "@com_google_absl//absl/strings",
+    ],
+)