[XLA] Fix dynamic expression edge cases and batch bucketing

tensorflow/compiler/jit/BUILD

@@ -780,6 +780,8 @@ cc_library(
         ":flags_headers",
         ":tf_graph_to_hlo_compiler",
         ":xla_compile_util",
+        ":xla_batch_matcher",
+        "//tensorflow/compiler/jit:flags",
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/core:framework",
         "//tensorflow/core:framework_lite",
@@ -1006,6 +1008,7 @@ cc_library(
     hdrs = ["shape_inference.h"],
     visibility = [":friends"],
     deps = [
+        ":flags",
         ":shape_inference_helpers",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
@@ -2005,3 +2008,23 @@ tf_cuda_cc_test(
         "@local_xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
     ],
 )
+
+cc_library(
+    name = "xla_batch_matcher",
+    srcs = ["xla_batch_matcher.cc"],
+    hdrs = ["xla_batch_matcher.h"],
+    deps = [
+        "//tensorflow/core/platform:logging",
+        "@local_xla//xla:debug_options_flags",
+    ],
+)
+
+tf_cc_test(
+    name = "xla_batch_matcher_test",
+    srcs = ["xla_batch_matcher_test.cc"],
+    deps = [
+        ":xla_batch_matcher",
+        "//tensorflow/core:test",
+        "@com_google_googletest//:gtest_main",
+    ],
+)

tensorflow/compiler/jit/device_compilation_cluster_signature.cc

@@ -15,21 +15,31 @@ limitations under the License.
 
 #include "tensorflow/compiler/jit/device_compilation_cluster_signature.h"
 
+#include "absl/strings/str_cat.h"
 #include <string>
 #include <utility>
 #include <variant>
 
 namespace tensorflow {
 namespace {
 using Signature = DeviceCompilationClusterSignature;
+using ConstantTensor = Signature::ConstantTensor;
 using TensorTypeAndShape = Signature::TensorTypeAndShape;
 
 // Functor that converts a Signature's arg to a human readable string.
 struct SignatureHumanStringAppender {
   explicit SignatureHumanStringAppender(std::string* dest) : dest(dest) {}
   std::string* dest;
-  void operator()(const Tensor& arg) {
-    absl::StrAppend(dest, "; ", arg.DebugString());
+  void operator()(const ConstantTensor& arg) {
+    absl::StrAppend(dest, "; ", arg.value.DebugString());
+    if (!arg.contents.empty()) {
+      absl::StrAppend(dest, " contents=[");
+      for (int i = 0; i < arg.contents.size(); ++i) {
+        if (i > 0) absl::StrAppend(dest, ",");
+        absl::StrAppend(dest, arg.contents[i].DebugString());
+      }
+      absl::StrAppend(dest, "]");
+    }
   }
   void operator()(const TensorTypeAndShape& arg) {
     absl::StrAppend(dest, ",", DataTypeString(arg.first));
@@ -40,18 +50,29 @@ struct SignatureHumanStringAppender {
 // Functor that compares the arg values of two different signatures. Returns
 // true when the args are not equal.
 struct SignatureNotEqual {
-  bool operator()(const Tensor& arg, const Tensor& other) {
-    return arg.dtype() != other.dtype() || arg.shape() != other.shape() ||
-           arg.tensor_data() != other.tensor_data();
+  bool operator()(const ConstantTensor& arg, const ConstantTensor& other) {
+    if (arg.value.dtype() != other.value.dtype() ||
+        arg.value.shape() != other.value.shape() ||
+        arg.value.tensor_data() != other.value.tensor_data() ||
+        arg.contents.size() != other.contents.size()) {
+      return true;
+    }
+    for (int i = 0; i < arg.contents.size(); ++i) {
+      if (arg.contents[i].SerializeAsString() !=
+          other.contents[i].SerializeAsString()) {
+        return true;
+      }
+    }
+    return false;
   }
   bool operator()(const TensorTypeAndShape& arg,
                   const TensorTypeAndShape& other) {
     return arg.first != other.first || arg.second != other.second;
   }
-  bool operator()(const Tensor& arg, const TensorTypeAndShape& other) {
+  bool operator()(const ConstantTensor& arg, const TensorTypeAndShape& other) {
     return true;
   }
-  bool operator()(const TensorTypeAndShape& arg, const Tensor& other) {
+  bool operator()(const TensorTypeAndShape& arg, const ConstantTensor& other) {
     return true;
   }
 };
@@ -61,12 +82,16 @@ struct SignatureNotEqual {
 struct SignatureHashCombiner {
   explicit SignatureHashCombiner(const uint64 h) : h(h) {}
   uint64 h;
-  uint64 operator()(const Tensor& arg) {
-    h = Hash64Combine(h, std::hash<int>()(static_cast<int>(arg.dtype())));
+  uint64 operator()(const ConstantTensor& arg) {
+    h = Hash64Combine(h, std::hash<int>()(static_cast<int>(arg.value.dtype())));
     h = Hash64Combine(
-        h, Hash64(arg.tensor_data().data(), arg.tensor_data().size()));
-    for (int dim = 0; dim < arg.dims(); ++dim) {
-      h = Hash64Combine(h, std::hash<int>()(arg.dim_size(dim)));
+        h, Hash64(arg.value.tensor_data().data(), arg.value.tensor_data().size()));
+    for (int dim = 0; dim < arg.value.dims(); ++dim) {
+      h = Hash64Combine(h, std::hash<int>()(arg.value.dim_size(dim)));
+    }
+    for (const xla::ExpressionProto& expr : arg.contents) {
+      std::string serialized = expr.SerializeAsString();
+      h = Hash64Combine(h, Hash64(serialized.data(), serialized.size()));
     }
     return h;
   }
@@ -120,7 +145,8 @@ absl::StatusOr<Signature> Signature::Build(
     switch (arg.kind) {
       case XlaCompiler::Argument::kConstant:
       case XlaCompiler::Argument::kConstantResource:
-        signature.args.push_back(arg.constant_value);
+        signature.args.push_back(
+            ConstantTensor{arg.constant_value, arg.constant_value_expressions});
         break;
       case XlaCompiler::Argument::kParameter:
       case XlaCompiler::Argument::kResource:

tensorflow/compiler/jit/device_compilation_cluster_signature.h

@@ -18,8 +18,10 @@ limitations under the License.
 
 #include <utility>
 #include <variant>
+#include <vector>
 
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
 
 namespace tensorflow {
 
@@ -34,7 +36,11 @@ struct DeviceCompilationClusterSignature {
   // argument number. Tensors must be in host memory.
   using TensorTypeAndShape =
       std::pair<DataType, absl::InlinedVector<int64_t, 4>>;
-  absl::InlinedVector<std::variant<Tensor, TensorTypeAndShape>, 8> args;
+  struct ConstantTensor {
+    Tensor value;
+    std::vector<xla::ExpressionProto> contents;
+  };
+  absl::InlinedVector<std::variant<ConstantTensor, TensorTypeAndShape>, 8> args;
 
   bool operator==(const DeviceCompilationClusterSignature& other) const;
 

tensorflow/compiler/jit/device_compilation_profiler.cc

@@ -21,6 +21,7 @@ limitations under the License.
 #include <utility>
 
 #include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/jit/xla_activity.pb.h"
 #include "tensorflow/compiler/jit/xla_activity_listener.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
@@ -32,9 +33,30 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 bool ShouldBeMegamorphic(int64_t compile_count, int64_t execution_count) {
-  const int64_t kCompileThreshold = 10;
+  int64_t kCompileThreshold = 10;
   const int64_t kMinExecutionsPerCompile = 50;
 
+  int64_t tf_xla_threshold_for_megamorphic =
+      GetMarkForCompilationPassFlags()->tf_xla_threshold_for_megamorphic;
+
+  // Negative values other that -1 cannot be used
+  if (tf_xla_threshold_for_megamorphic < -1) {
+    LOG(FATAL) << "The value for the tf_xla_threshold_for_megamorphic flag "
+               << "is out of range.\n"
+               << "Allowed ranges are (-1) to "
+               << std::numeric_limits<int64_t>::max()
+               << " got " << tf_xla_threshold_for_megamorphic << ".";
+  }
+
+  // -1: setting clusters as Megamorphic is disabled
+  //  0 Default behaviour in Tensorflow
+  //  Any other number sets the compilation threshold
+  if (tf_xla_threshold_for_megamorphic == -1) {
+    return false;
+  } else if (tf_xla_threshold_for_megamorphic > 0) {
+    kCompileThreshold = tf_xla_threshold_for_megamorphic;
+  }
+
   // This heuristic is trying to capture the following property: have we sunk a
   // certain minimum amount of compile time into the cluster that didn't quite
   // "pay off"?

tensorflow/compiler/jit/device_compiler.h

@@ -36,6 +36,7 @@ limitations under the License.
 #include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/jit/tf_graph_to_hlo_compiler.h"
 #include "tensorflow/compiler/jit/xla_compile_util.h"
+#include "tensorflow/compiler/jit/xla_batch_matcher.h"
 #include "tensorflow/compiler/tf2xla/xla_compiler.h"
 #include "tensorflow/core/framework/metrics.h"
 #include "tensorflow/core/framework/op_kernel.h"
@@ -125,6 +126,7 @@ class DeviceCompiler : public ResourceBase {
   DeviceCompilerClient<ExecutableType, ClientType>* compiler_client() {
     return compiler_client_.get();
   }
+  XlaBatchMatcher* xla_batch_matcher() { return xla_batch_matcher_.get(); }
 
   string DebugString() const override;
 
@@ -177,6 +179,9 @@ class DeviceCompiler : public ResourceBase {
   // Pool of threads for asynchronous compilations.
   std::unique_ptr<thread::ThreadPool> async_compiler_threads_;
 
+  // Specified dynamic batch padding values.
+  std::unique_ptr<XlaBatchMatcher> xla_batch_matcher_;
+
   mutex cluster_mutexes_mu_;
   absl::flat_hash_map<DeviceCompilationClusterSignature, std::unique_ptr<mutex>,
                       DeviceCompilationClusterSignature::Hash>
@@ -225,6 +230,11 @@ DeviceCompiler<ExecutableType, ClientType>::DeviceCompiler(
   async_compiler_threads_ = std::make_unique<tensorflow::thread::ThreadPool>(
       tensorflow::Env::Default(), "async_compiler_threads",
       kNumAsyncDeviceCompilerThreads);
+
+  MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags();
+  if (flags->tf_xla_enable_dynamic_sizes) {
+    xla_batch_matcher_ = std::make_unique<XlaBatchMatcher>();
+  }
 }
 
 template <typename ExecutableType, typename ClientType>

tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc

@@ -26,6 +26,7 @@ limitations under the License.
 #include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "absl/types/optional.h"
+#include "tensorflow/compiler/jit/encapsulate_util.h"
 #include "tensorflow/compiler/jit/flags.h"
 #include "tensorflow/compiler/jit/mark_for_compilation_pass.h"
 #include "tensorflow/compiler/jit/shape_inference_helpers.h"
@@ -54,9 +55,14 @@ limitations under the License.
 #include "tensorflow/core/public/version.h"
 #include "tensorflow/core/util/device_name_utils.h"
 #include "tensorflow/core/util/dump_graph.h"
-
+#include "tensorflow/core/framework/tensor_shape.pb.h"
 namespace tensorflow {
 
+static const absl::flat_hash_set<absl::string_view> kFailingOps = {
+    "Where",
+    // add more here
+};
+
 const char* const kXlaCompiledKernelAttr = "_XlaCompiledKernel";
 const char* const kXlaNumConstantArgsAttr = "_XlaNumConstantArgs";
 const char* const kXlaNumResourceArgsAttr = "_XlaNumResourceArgs";
@@ -114,6 +120,41 @@ void MarkGuaranteedConstants(
   }
 }
 
+// Helper to convert ExpressionProto to a readable string.
+std::string ExprProtoToString(const ExpressionProto& e) {
+  switch (e.node_type_case()) {
+    case ExpressionProto::kConstantValue:
+      return std::to_string(e.constant_value());
+    case ExpressionProto::kVariableId:
+      return absl::StrCat("Var(", e.variable_id(), ")");
+    case ExpressionProto::kAddNode:
+      return absl::StrCat("(", ExprProtoToString(e.add_node().lhs()), " + ",
+                         ExprProtoToString(e.add_node().rhs()), ")");
+    case ExpressionProto::kSubNode:
+      return absl::StrCat("(", ExprProtoToString(e.sub_node().lhs()), " - ",
+                         ExprProtoToString(e.sub_node().rhs()), ")");
+    case ExpressionProto::kMulNode:
+      return absl::StrCat("(", ExprProtoToString(e.mul_node().lhs()), " * ",
+                         ExprProtoToString(e.mul_node().rhs()), ")");
+    case ExpressionProto::kDivNode:
+      return absl::StrCat("(", ExprProtoToString(e.div_node().lhs()), " / ",
+                         ExprProtoToString(e.div_node().rhs()), ")");
+    case ExpressionProto::kMaxNode:
+      return absl::StrCat("max(", ExprProtoToString(e.max_node().lhs()), ", ",
+                          ExprProtoToString(e.max_node().rhs()), ")");
+    case ExpressionProto::kGtNode:
+      return absl::StrCat("(", ExprProtoToString(e.gt_node().lhs()), " > ",
+                          ExprProtoToString(e.gt_node().rhs()), ")");
+    case ExpressionProto::kSelectNode:
+      return absl::StrCat("select(", ExprProtoToString(e.select_node().pred()),
+                          ", ", ExprProtoToString(e.select_node().on_true()),
+                          ", ", ExprProtoToString(e.select_node().on_false()),
+                          ")");
+    default:
+      return "<none>";
+  }
+}
+
 struct OutputInputTensorPairHasher {
   uint64 operator()(std::pair<OutputTensor, InputTensor> const& s) const {
     return Hash64Combine(OutputTensor::Hash()(s.first),
@@ -369,6 +410,19 @@ class Encapsulator {
 
 namespace {
 
+bool BuildOutputShapeProto(const Node& node, int output_slot,
+                           TensorShapeProto* proto) {
+  AttrSlice attrs = node.attrs();
+  auto shape_attr =
+      attrs.FindByString(kXlaInferredOutputTensorShapesAttrName);
+  if (shape_attr == nullptr || !shape_attr->has_list() ||
+      shape_attr->list().shape_size() <= output_slot) {
+    return false;
+  }
+  *proto = shape_attr->list().shape(output_slot);
+  return true;
+}
+
 // Return in 'sorted' a topological sort of clusters according to the
 // dependencies encoded in ancestors. clusters is the list of all clusters
 // including clusters that are not present in the ancestors map. has_successors
@@ -470,6 +524,26 @@ absl::Status Encapsulator::Subgraph::RecordArg(
     DataType dtype = edge->dst()->input_type(edge->dst_input());
     builder.Attr("T", dtype);
     builder.Attr("index", arg_index);
+    AttrSlice attrs = src_node->attrs();
+    TensorShapeProto output_shape_proto;
+    if (BuildOutputShapeProto(*src_node, src_slot, &output_shape_proto)) {
+      VLOG(1) << "Adding following output shapes for node " << src_node->name()
+              << " : " << output_shape_proto.DebugString();
+      builder.Attr("_output_shapes", {output_shape_proto});
+      builder.Attr(kXlaInferredOutputShapesAttrName, {output_shape_proto});
+    } else {
+      // if cluster argument is the real argument.
+      auto build_attr = attrs.FindByString("_dynamic_dim");
+      if (build_attr) {
+        VLOG(1) << "Found Dynamic dimension in " << src_node->name() << ":"
+                << src_slot;
+        builder.Attr("_dynamic_dim", *build_attr);
+      }
+    }
+    auto shape_derived_attr = attrs.FindByString(kXlaShapeDerivedAttrName);
+    if (shape_derived_attr) {
+      builder.Attr(kXlaShapeDerivedAttrName, *shape_derived_attr);
+    }
     absl::Status s = builder.Finalize(&arg_def);
     if (!s.ok()) return s;
 
@@ -1143,6 +1217,14 @@ static absl::Status RenumberArguments(Graph* graph,
   return absl::OkStatus();
 }
 
+static bool SubgraphHasFailingOps(const Graph& g) {
+  for (Node* n : g.op_nodes()) {
+    if (n->IsRetval()) continue;
+    if (kFailingOps.contains(n->def().op())) return true;
+  }
+  return false;
+}
+
 absl::Status EncapsulateSubgraphsPass::Run(
     const GraphOptimizationPassOptions& options) {
   VLOG(1) << "EncapsulateSubgraphsPass::Run";
@@ -1289,8 +1371,8 @@ absl::Status EncapsulateSubgraphsPass::Run(
 
         // TODO(phawkins): add a forward is-constant analysis, similarly split
         // outputs into host-memory constants and device-memory non-constants.
-
-        AddNodeAttr(kXlaCompiledKernelAttr, true, node);
+        bool compile_enabled = !SubgraphHasFailingOps(**subgraph);
+        AddNodeAttr(kXlaCompiledKernelAttr, compile_enabled, node);
         AddNodeAttr(kXlaNumConstantArgsAttr, num_consts, node);
         AddNodeAttr(kXlaNumResourceArgsAttr, num_resources, node);
         return absl::OkStatus();

tensorflow/compiler/jit/encapsulate_util.cc

@@ -303,6 +303,11 @@ absl::Status PostprocessControlEdgesBetweenOutsideCompilations(
 }  // namespace
 
 const char kXlaInferredShapesAttrName[] = "_xla_inferred_shapes";
+const char kXlaInferredOutputTensorShapesAttrName[] =
+    "_xla_inferred_output_tensor_shapes";
+const char kXlaInferredOutputShapesAttrName[] =
+    "_xla_inferred_output_shapes";
+const char kXlaShapeDerivedAttrName[] = "_xla_shape_derived";
 
 const char kXlaConnectedToXlaComputationAttrName[] =
     "_xla_connected_to_xla_computation";