[pull] master from tensorflow:master

tensorflow/core/kernels/image/extract_volume_patches_op.cc

@@ -130,9 +130,9 @@ class ExtractVolumePatchesOp : public UnaryOp<T> {
         MultiplyWithoutOverflow(ksize_rows,
                                 MultiplyWithoutOverflow(ksize_cols, depth)));
     OP_REQUIRES(context, patch_size >= 0,
-                errors::InvalidArgument(
+                absl::InvalidArgumentError(absl::StrCat(
                     "Output size would overflow: ", ksize_planes, " x ",
-                    ksize_rows, " x ", ksize_cols, " x ", depth));
+                    ksize_rows, " x ", ksize_cols, " x ", depth)));
 
     const std::vector<int64_t> out_sizes = {batch, out_planes, out_rows,
                                             out_cols, patch_size};

tensorflow/lite/experimental/acceleration/compatibility/BUILD

@@ -100,9 +100,9 @@ genrule(
     ],
     outs = ["devicedb-sample.bin"],
     cmd = """
-    $(location :json_to_fb) \
-        --fbs=$(location :database.fbs) \
-        --json_input=$(location :devicedb-sample.json) \
+    $(execpath :json_to_fb) \
+        --fbs=$(execpath :database.fbs) \
+        --json_input=$(execpath :devicedb-sample.json) \
         --fb_output=$(@)
     """,
     tools = [":json_to_fb"],
@@ -125,10 +125,10 @@ genrule(
     # convert_file_to_c_source for some reason doesn't define the global with
     # 'extern', which is needed for global const variables in C++.
     cmd = """
-    $(location :convert_binary_to_cc_source) \
-        --input_binary_file $(location :devicedb-sample.bin) \
-        --output_header_file $(location :devicedb-sample.h) \
-        --output_source_file $(location :devicedb-sample.cc) \
+    $(execpath :convert_binary_to_cc_source) \
+        --input_binary_file $(execpath :devicedb-sample.bin) \
+        --output_header_file $(execpath :devicedb-sample.h) \
+        --output_source_file $(execpath :devicedb-sample.cc) \
         --array_variable_name g_tflite_acceleration_devicedb_sample_binary
     """,
     tools = [":convert_binary_to_cc_source"],
@@ -162,6 +162,7 @@ cc_test(
 exports_files([
     "database.fbs",
     "gpu_compatibility.bin",
+    "mldrift_compatibility.bin",
 ])
 
 genrule(
@@ -174,15 +175,34 @@ genrule(
     # convert_file_to_c_source for some reason doesn't define the global with
     # 'extern', which is needed for global const variables in C++.
     cmd = """
-    $(location :convert_binary_to_cc_source) \
-        --input_binary_file $(location :gpu_compatibility.bin) \
-        --output_header_file $(location :gpu_compatibility_binary.h) \
-        --output_source_file $(location :gpu_compatibility_binary.cc) \
+    $(execpath :convert_binary_to_cc_source) \
+        --input_binary_file $(execpath :gpu_compatibility.bin) \
+        --output_header_file $(execpath :gpu_compatibility_binary.h) \
+        --output_source_file $(execpath :gpu_compatibility_binary.cc) \
         --array_variable_name g_tflite_acceleration_gpu_compatibility_binary
     """,
     tools = [":convert_binary_to_cc_source"],
 )
 
+genrule(
+    name = "mldrift_compatibility_binary",
+    srcs = ["mldrift_compatibility.bin"],
+    outs = [
+        "mldrift_compatibility_binary.h",
+        "mldrift_compatibility_binary.cc",
+    ],
+    # The array variable name must differ from the gpu_compatibility one above:
+    # both allowlists may end up linked into the same binary.
+    cmd = """
+    $(execpath :convert_binary_to_cc_source) \
+        --input_binary_file $(execpath :mldrift_compatibility.bin) \
+        --output_header_file $(execpath :mldrift_compatibility_binary.h) \
+        --output_source_file $(execpath :mldrift_compatibility_binary.cc) \
+        --array_variable_name g_tflite_acceleration_mldrift_compatibility_binary
+    """,
+    tools = [":convert_binary_to_cc_source"],
+)
+
 cc_library(
     name = "android_info",
     srcs = ["android_info.cc"],
@@ -198,6 +218,16 @@ cc_library(
     hdrs = [":gpu_compatibility_binary.h"],
 )
 
+# Deliberately a target of its own rather than an addition to the library above:
+# a client that only cares about ML Drift should not link the (much larger)
+# TFLite GPU allowlist, and vice versa. Both binaries are consumed through the
+# generic GPUCompatibilityList::Create(flatbuffer, length) overload.
+cc_library(
+    name = "mldrift_compatibility_binary_embed",
+    srcs = [":mldrift_compatibility_binary.cc"],
+    hdrs = [":mldrift_compatibility_binary.h"],
+)
+
 cc_library(
     name = "gpu_compatibility",
     hdrs = ["gpu_compatibility.h"],

third_party/xla/xla/service/BUILD

@@ -4780,8 +4780,7 @@ cc_library(
         "//xla/hlo/pass:hlo_pass",
         "//xla/hlo/transforms/simplifiers:hlo_dce",
         "//xla/hlo/transforms/simplifiers:tuple_simplifier",
-        "//xla/tsl/platform:errors",
-        "//xla/tsl/platform:statusor",
+        "//xla/tsl/platform:logging",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
@@ -4810,6 +4809,7 @@ xla_cc_test(
         "//xla/tests:xla_internal_test_main",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_googletest//:gtest",
     ],
 )

third_party/xla/xla/service/while_loop_invariant_code_motion.cc

@@ -27,6 +27,7 @@ limitations under the License.
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "xla/hlo/analysis/while_loop_analysis.h"
 #include "xla/hlo/ir/hlo_computation.h"
@@ -41,8 +42,7 @@ limitations under the License.
 #include "xla/service/while_util.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
-#include "xla/tsl/platform/errors.h"
-#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/platform/logging.h"
 #include "xla/util.h"
 
 namespace xla {
@@ -100,7 +100,11 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody(
     return false;
   }
 
-  std::string while_instr_name = while_instr->ToString(print_no_metadata);
+  // Only used by VLOG, and the string carries the whole loop state shape.
+  std::string while_instr_name;
+  if (VLOG_IS_ON(1)) {
+    while_instr_name = while_instr->ToString(print_no_metadata);
+  }
   VLOG(2) << "Trying to hoist from " << while_instr_name;
 
   auto maybe_upper_bound = ComputeWhileLoopTripCountUpperBound(while_instr);
@@ -146,7 +150,8 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody(
     return false;
   }
 
-  for (auto* instruction : while_body->MakeInstructionPostOrder()) {
+  // This scan bails on any match, so it does not need the post order.
+  for (const auto* instruction : while_body->instructions()) {
     // LICM in the presence of domain instructions is complex, bail.
     if (instruction->opcode() == HloOpcode::kDomain ||
         instruction->IsCustomCall("SPMDFullToShardShape") ||
@@ -167,6 +172,12 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody(
   std::vector<HloInstruction*> instructions_to_replace;
   std::vector<HloInstruction*> replacement_instructions;
 
+  const auto is_invariant = [&](HloInstruction* op) {
+    return op->opcode() == HloOpcode::kConstant ||
+           hoisted_instructions.contains(op) ||
+           unhoisted_invariant_instructions.contains(op);
+  };
+
   for (auto* instruction : while_body->MakeInstructionPostOrder()) {
     allowance->DeductCost(1);
     if (!allowance->ContinueAnalysis()) {
@@ -185,6 +196,13 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody(
         instruction->opcode() != HloOpcode::kReshape) {
       continue;
     }
+
+    // Cheap check first. Both checks are pure filters, so their order does not
+    // change which instructions are hoisted.
+    if (!absl::c_all_of(instruction->operands(), is_invariant)) {
+      continue;
+    }
+
     // Constants don't inflate, so size inflation check doesn't make sense for
     // constants.
     if (hoist_size_inflation_ratio_ &&
@@ -220,16 +238,6 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody(
       }
     }
 
-    auto is_invariant = [&](HloInstruction* op) {
-      return hoisted_instructions.find(op) != hoisted_instructions.end() ||
-             unhoisted_invariant_instructions.contains(op) ||
-             op->opcode() == HloOpcode::kConstant;
-    };
-
-    if (!absl::c_all_of(instruction->operands(), is_invariant)) {
-      continue;
-    }
-
     if (NotWorthHoistingIndividually(*instruction)) {
       VLOG(2) << "Adding " << instruction->ToString(print_no_metadata)
               << " to unhoisted invariant set.";