[pull] master from tensorflow:master

tensorflow/compiler/mlir/tensorflow/transforms/BUILD

@@ -717,6 +717,7 @@ cc_library(
         "//tensorflow/core/protobuf:for_core_protos_cc",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:vlog_is_on",
         "@com_google_absl//absl/status",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:FuncDialect",

tensorflow/compiler/mlir/tensorflow/transforms/bridge.cc

@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/tensorflow/transforms/bridge.h"
 
 #include "absl/log/log.h"
+#include "absl/log/vlog_is_on.h"
 #include "llvm/ADT/StringRef.h"
 #include "mlir/IR/BuiltinOps.h"  // from @llvm-project
 #include "mlir/Pass/PassManager.h"  // from @llvm-project

tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD

@@ -241,10 +241,9 @@ tf_cc_test(
         "@xla//xla/pjrt:pjrt_compiler",
         "@xla//xla/pjrt/plugin/xla_cpu:cpu_topology_description",
         "@xla//xla/python/ifrt",
-        "@xla//xla/python/ifrt:mock",
         "@xla//xla/python/ifrt:test_util",
         "@xla//xla/python/pjrt_ifrt",
-        "@xla//xla/python/pjrt_ifrt:tfrt_cpu_client_test_lib",
+        "@xla//xla/python/pjrt_ifrt:pjrt_cpu_client_test_lib",
         "@xla//xla/service:device_assignment",
         "@xla//xla/tsl/platform:statusor",
     ],

tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo_test.cc

@@ -43,7 +43,6 @@ limitations under the License.
 #include "xla/pjrt/pjrt_compiler.h"
 #include "xla/pjrt/plugin/xla_cpu/cpu_topology_description.h"
 #include "xla/python/ifrt/client.h"
-#include "xla/python/ifrt/mock.h"
 #include "xla/python/ifrt/test_util.h"
 #include "xla/python/pjrt_ifrt/pjrt_topology.h"
 #include "xla/service/device_assignment.h"
@@ -514,18 +513,15 @@ TEST_F(Tf2HloTest, GpuCompile) {
   ASSERT_TRUE(mlir_module);
   ASSERT_TRUE(mlir_module.get() != nullptr);
 
-  xla::ifrt::MockClient mock_client;
-  ON_CALL(mock_client, GetDefaultDeviceAssignment)
-      .WillByDefault([]() -> absl::StatusOr<xla::DeviceAssignment> {
-        return xla::DeviceAssignment(1, 1);
-      });
+  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
+                          xla::ifrt::test_util::GetClient());
 
   std::vector<DtypeAndShape> dtype_and_shapes;
   dtype_and_shapes.push_back(DtypeAndShape{DT_FLOAT, {}});
 
   TF_ASSERT_OK_AND_ASSIGN(
       tensorflow::tpu::TPUCompileMetadataProto compile_metadata,
-      GetCompileMetadata(mlir_module.get(), mock_client));
+      GetCompileMetadata(mlir_module.get(), *client));
   TF_ASSERT_OK(UpdateCompileMetadata(compile_metadata, dtype_and_shapes));
 
   std::vector<int> variable_arg_indices;

tensorflow/core/kernels/rnn/lstm_ops.cc

@@ -593,6 +593,35 @@ class LSTMBlockCellGradOp : public OpKernel {
     const Tensor* h_grad_tensor = nullptr;
     OP_REQUIRES_OK(ctx, ctx->input("h_grad", &h_grad_tensor));
 
+    // The dimensions read below, and the sanity checks that follow them, index
+    // dimension 1 of these tensors. That is a fatal check rather than an error
+    // when a tensor has a lower rank, so validate the ranks first, the way
+    // LSTMBlockCellOp does before using the same inputs.
+    const auto check_rank = [](const Tensor* t, const char* name,
+                               int expected) -> absl::Status {
+      if (t->dims() != expected) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            name, " must be rank ", expected, " but is rank ", t->dims()));
+      }
+      return absl::OkStatus();
+    };
+    OP_REQUIRES_OK(ctx, check_rank(x_tensor, "x", 2));
+    OP_REQUIRES_OK(ctx, check_rank(cs_prev_tensor, "cs_prev", 2));
+    OP_REQUIRES_OK(ctx, check_rank(h_prev_tensor, "h_prev", 2));
+    OP_REQUIRES_OK(ctx, check_rank(w_tensor, "w", 2));
+    OP_REQUIRES_OK(ctx, check_rank(wci_tensor, "wci", 1));
+    OP_REQUIRES_OK(ctx, check_rank(wcf_tensor, "wcf", 1));
+    OP_REQUIRES_OK(ctx, check_rank(wco_tensor, "wco", 1));
+    OP_REQUIRES_OK(ctx, check_rank(b_tensor, "b", 1));
+    OP_REQUIRES_OK(ctx, check_rank(i_tensor, "i", 2));
+    OP_REQUIRES_OK(ctx, check_rank(cs_tensor, "cs", 2));
+    OP_REQUIRES_OK(ctx, check_rank(f_tensor, "f", 2));
+    OP_REQUIRES_OK(ctx, check_rank(o_tensor, "o", 2));
+    OP_REQUIRES_OK(ctx, check_rank(ci_tensor, "ci", 2));
+    OP_REQUIRES_OK(ctx, check_rank(co_tensor, "co", 2));
+    OP_REQUIRES_OK(ctx, check_rank(cs_grad_tensor, "cs_grad", 2));
+    OP_REQUIRES_OK(ctx, check_rank(h_grad_tensor, "h_grad", 2));
+
     const int64_t batch_size = x_tensor->dim_size(0);
     const int64_t input_size = x_tensor->dim_size(1);
     const int64_t cell_size = cs_prev_tensor->dim_size(1);
@@ -1292,6 +1321,50 @@ class BlockLSTMGradOp : public OpKernel {
     const Tensor* h_grad = nullptr;
     OP_REQUIRES_OK(ctx, ctx->input("h_grad", &h_grad));
 
+    // `timelen`, `batch_size` and `cell_size` are taken from x and w, and every
+    // other input is sliced with them below. A tensor whose dimensions
+    // disagree is read out of bounds rather than failing the op: a zero batch
+    // on one side and a non-zero batch on the other crashes, and a shorter
+    // time dimension trips a fatal check while slicing. LSTMBlockCellGradOp
+    // already validates its own inputs this way. The rank is checked here as
+    // well so that the dimension lookups below are valid for any caller.
+    const auto check_dims = [&](const Tensor* t, const char* name,
+                                bool has_time) -> absl::Status {
+      const int expected_dims = has_time ? 3 : 2;
+      if (t->dims() != expected_dims) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            name, " must be rank ", expected_dims, " but is rank ", t->dims()));
+      }
+      const int batch_dim = has_time ? 1 : 0;
+      if (has_time && t->dim_size(0) != timelen) {
+        return absl::InvalidArgumentError(
+            absl::StrCat(name, ".dim_size(0) != timelen: ", t->dim_size(0),
+                         " vs. ", timelen));
+      }
+      if (t->dim_size(batch_dim) != batch_size) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            name, ".dim_size(", batch_dim,
+            ") != batch_size: ", t->dim_size(batch_dim), " vs. ", batch_size));
+      }
+      if (t->dim_size(batch_dim + 1) != cell_size) {
+        return absl::InvalidArgumentError(
+            absl::StrCat(name, ".dim_size(", batch_dim + 1, ") != cell_size: ",
+                         t->dim_size(batch_dim + 1), " vs. ", cell_size));
+      }
+      return absl::OkStatus();
+    };
+    OP_REQUIRES_OK(ctx, check_dims(cs_prev_tensor, "cs_prev", false));
+    OP_REQUIRES_OK(ctx, check_dims(h_prev_tensor, "h_prev", false));
+    OP_REQUIRES_OK(ctx, check_dims(i_out, "i", true));
+    OP_REQUIRES_OK(ctx, check_dims(cs_out, "cs", true));
+    OP_REQUIRES_OK(ctx, check_dims(f_out, "f", true));
+    OP_REQUIRES_OK(ctx, check_dims(o_out, "o", true));
+    OP_REQUIRES_OK(ctx, check_dims(ci_out, "ci", true));
+    OP_REQUIRES_OK(ctx, check_dims(co_out, "co", true));
+    OP_REQUIRES_OK(ctx, check_dims(h_out, "h", true));
+    OP_REQUIRES_OK(ctx, check_dims(cs_grad, "cs_grad", true));
+    OP_REQUIRES_OK(ctx, check_dims(h_grad, "h_grad", true));
+
     TensorShape batch_input_shape({timelen, batch_size, input_size});
     Tensor* x_grad;
     OP_REQUIRES_OK(ctx,

tensorflow/python/kernel_tests/nn_ops/rnn_cell_test.py

@@ -1468,6 +1468,102 @@ def testLSTMBlockInvalidArgument(self):
           )
       )
 
+  def _blockLSTMGradInputs(
+      self, timelen, batch_size, input_size, cell_size, x_batch_size=None
+  ):
+    """Returns a consistent set of BlockLSTMGrad inputs, x batch overridable."""
+    if x_batch_size is None:
+      x_batch_size = batch_size
+
+    def rand(*shape):
+      return random_ops.random_uniform(list(shape), dtype=dtypes.float32)
+
+    def per_step():
+      return rand(timelen, batch_size, cell_size)
+
+    return dict(
+        seq_len_max=constant_op.constant(timelen, dtype=dtypes.int64),
+        x=rand(timelen, x_batch_size, input_size),
+        cs_prev=rand(batch_size, cell_size),
+        h_prev=rand(batch_size, cell_size),
+        w=rand(input_size + cell_size, 4 * cell_size),
+        wci=rand(cell_size),
+        wcf=rand(cell_size),
+        wco=rand(cell_size),
+        b=rand(4 * cell_size),
+        i=per_step(),
+        cs=per_step(),
+        f=per_step(),
+        o=per_step(),
+        ci=per_step(),
+        co=per_step(),
+        h=per_step(),
+        cs_grad=per_step(),
+        h_grad=per_step(),
+        use_peephole=True,
+    )
+
+  def testBlockLSTMGradBatchSizeMismatch(self):
+    # Test case for GitHub issue 112746. x carried a different batch size from
+    # the other inputs, which were then sliced with the batch size taken from
+    # x and read out of bounds, crashing the process.
+    kwargs = self._blockLSTMGradInputs(3, 2, 4, 5, x_batch_size=0)
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, "batch_size"):
+      self.evaluate(gen_rnn_ops.BlockLSTMGrad(**kwargs))
+
+  def testBlockLSTMGradTimeLenMismatch(self):
+    # A shorter time dimension on one of the per step inputs was sliced past
+    # its end.
+    kwargs = self._blockLSTMGradInputs(3, 2, 4, 5)
+    kwargs["i"] = random_ops.random_uniform([1, 2, 5], dtype=dtypes.float32)
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, "timelen"):
+      self.evaluate(gen_rnn_ops.BlockLSTMGrad(**kwargs))
+
+  def testBlockLSTMGradConsistentShapesStillWork(self):
+    # Consistent shapes must keep working, including a zero batch size.
+    for batch_size in (0, 2):
+      kwargs = self._blockLSTMGradInputs(3, batch_size, 4, 5)
+      self.evaluate(gen_rnn_ops.BlockLSTMGrad(**kwargs))
+
+  def testLSTMBlockCellGradInvalidRank(self):
+    # Test case for GitHub issue 112746. A rank 1 x was indexed at dimension 1
+    # before any shape check ran, which aborted the process. Eager execution
+    # reaches the kernel check ("x must be rank 2"), while graph construction
+    # is stopped earlier by the op's shape function ("Shape must be rank 2"),
+    # which surfaces as a ValueError, so both are accepted.
+    batch_size, input_size, cell_size = 2, 4, 5
+
+    def rand(*shape):
+      return random_ops.random_uniform(list(shape), dtype=dtypes.float32)
+
+    def per_step():
+      return rand(batch_size, cell_size)
+
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError), "must be rank 2"
+    ):
+      self.evaluate(
+          gen_rnn_ops.LSTMBlockCellGrad(
+              x=constant_op.constant([], dtype=dtypes.float32),
+              cs_prev=per_step(),
+              h_prev=per_step(),
+              w=rand(input_size + cell_size, 4 * cell_size),
+              wci=rand(cell_size),
+              wcf=rand(cell_size),
+              wco=rand(cell_size),
+              b=rand(4 * cell_size),
+              i=per_step(),
+              cs=per_step(),
+              f=per_step(),
+              o=per_step(),
+              ci=per_step(),
+              co=per_step(),
+              cs_grad=per_step(),
+              h_grad=per_step(),
+              use_peephole=True,
+          )
+      )
+
   @test_util.run_in_graph_and_eager_modes
   def testBlockLSTMSeqLenMaxBounds(self):
     timelen, batch_size, input_size, cell_size = 4, 2, 3, 5

third_party/xla/build_tools/lint/tags.py

@@ -104,6 +104,7 @@
     "xla_interpreter": "Uses interpreter backend.",
     "xla_cpu": "Uses CPU backend.",
     "xla_amdgpu_any": "Uses ROCm backend.",
+    "xla_mi350": "Runs on MI350 backend.",
     "xla_nvgpu_any": "Uses NVIDIA GPU backend.",
     "xla_intelgpu_any": "Uses Intel GPU backend.",
     # Below tags are emitted alongside `requires-gpu-x` tags, which is what the

third_party/xla/third_party/cuda_tile/cuda_tile.BUILD

@@ -156,6 +156,35 @@ gentbl_cc_library(
     deps = [":CudaTileTdFiles"],
 )
 
+gentbl_cc_library(
+    name = "CudaTileAttrCodeGenIncGen",
+    tbl_outs = [
+        (
+            ["-gen-enum-attr-list"],
+            "include/cuda_tile/Dialect/CudaTile/IR/CudaTileEnumAttrs.inc",
+        ),
+        (
+            ["-gen-gpu-arch-def"],
+            "include/cuda_tile/Dialect/CudaTile/IR/GpuArchitectures.inc",
+        ),
+        (
+            ["-gen-gpu-arch-constants"],
+            "include/cuda_tile/Dialect/CudaTile/IR/GpuArchConstants.inc",
+        ),
+        (
+            ["-gen-hint-key-accessors"],
+            "include/cuda_tile/Dialect/CudaTile/IR/HintKeyAccessors.inc",
+        ),
+        (
+            ["-gen-hint-key-impl"],
+            "include/cuda_tile/Dialect/CudaTile/IR/HintKeyImpl.inc",
+        ),
+    ],
+    tblgen = ":cuda-tile-tblgen",
+    td_file = "include/cuda_tile/Dialect/CudaTile/IR/AttrDefs.td",
+    deps = [":CudaTileTdFiles"],
+)
+
 gentbl_cc_library(
     name = "CudaTileRemarksIncGen",
     tbl_outs = [
@@ -219,12 +248,15 @@ cc_library(
     ],
     visibility = ["//visibility:public"],
     deps = [
+        ":CudaTileAttrCodeGenIncGen",
         ":CudaTileAttrDefsIncGen",
         ":CudaTileDialectIncGen",
         ":CudaTileInterfacesIncGen",
         ":CudaTileOpsCanonicalizationIncGen",
         ":CudaTileOpsIncGen",
+        ":CudaTileRemarksIncGen",
         ":CudaTileTypesIncGen",
+        "@llvm-project//llvm:Core",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:BytecodeOpInterface",
@@ -276,6 +308,7 @@ cc_library(
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:SideEffectInterfaces",
         "@llvm-project//mlir:TransformUtils",
     ],
 )
@@ -363,7 +396,9 @@ cc_library(
         ":CudaTileBytecodeTypeIncGen",
         ":CudaTileDialect",
         "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FunctionInterfaces",
         "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TranslateLib",
     ],
 )

third_party/xla/third_party/cuda_tile/patches/constructor.patch

@@ -3,7 +3,7 @@ Remove redundant defaulted copy/move constructors to preserve aggregate status.
 diff --git a/tools/cuda-tile-tblgen/CudaTileOp.h b/tools/cuda-tile-tblgen/CudaTileOp.h
 --- a/tools/cuda-tile-tblgen/CudaTileOp.h
 +++ b/tools/cuda-tile-tblgen/CudaTileOp.h
-@@ -55,11 +55,6 @@
+@@ -56,11 +56,6 @@
    std::optional<std::vector<std::string>> selectedVariants;
    ParameterType type;
 

third_party/xla/third_party/cuda_tile/patches/pow_op_alias.patch

@@ -0,0 +1,13 @@
+diff --git a/include/cuda_tile/Dialect/CudaTile/IR/Ops.h b/include/cuda_tile/Dialect/CudaTile/IR/Ops.h
+--- a/include/cuda_tile/Dialect/CudaTile/IR/Ops.h
++++ b/include/cuda_tile/Dialect/CudaTile/IR/Ops.h
+@@ -59,6 +59,9 @@
+
+ namespace mlir::cuda_tile {
+
++/// Backwards compatibility alias for PowOp which was renamed to FPowFOp in v13.4.0.
++using PowOp = FPowFOp;
++
+ // Helper function to extract cuda_tile::ModuleOp
+ cuda_tile::ModuleOp extractCudaTileModuleOp(Operation *op);
+

third_party/xla/third_party/cuda_tile/patches/series

@@ -1,2 +1,3 @@
 constructor.patch
-symbol_op_interface.patch
+unused_variable.patch
+pow_op_alias.patch