message: [tf.data] Avoiding excessive buffering of large batches in the fused map+batch kernel when autotuning is enabled.
diff: @@ -56,6 +56,9 @@ namespace experimental {
  MapAndBatchDatasetOp::kPreserveCardinality;
  // Maximum number of batch results to buffer.
  +
  +namespace {
  +
  constexpr int64 kMaxBatchResults = 16;
  constexpr char kParallelism[] = "parallelism";
  constexpr char kCallCounter[] = "call_counter";
  @@ -72,6 +75,11 @@ constexpr char kStatus[] = "status";
  constexpr char kCode[] = "code";
  constexpr char kMessage[] = "msg";
  +// Computes ceil(x / y).
  +inline int64 CeilDiv(int64 x, int64 y) { return (x + y - 1) / y; }
  +
  +} // namespace
  +
  class MapAndBatchDatasetOp::Dataset : public DatasetBase {
  public:
  Dataset(OpKernelContext* ctx, const DatasetBase* input, int64 batch_size,
  @@ -175,14 +183,19 @@ class MapAndBatchDatasetOp::Dataset : public DatasetBase {
  mu_(std::make_shared<mutex>()),
  cond_var_(std::make_shared<condition_variable>()),
  num_parallel_calls_(std::make_shared<model::SharedState>(
  - params.dataset->num_parallel_calls_, mu_, cond_var_)),
  - max_batch_results_(
  - params.dataset->num_parallel_calls_ == model::kAutotune
  - ? kMaxBatchResults
  - : std::min(kMaxBatchResults,
  - (params.dataset->num_parallel_calls_ +
  - params.dataset->batch_size_ - 1) /
  - params.dataset->batch_size_)) {}
  + params.dataset->num_parallel_calls_, mu_, cond_var_)) {
  + // To mitigate the effect of stragglers (i.e. map invocations that take
  + // much longer than others), we allow the kernel to pre-compute batches
  + // ahead of time and store them in an internal buffer. The maximum number
  + // of batches to buffer is a trade-off between performance and memory and
  + // we derive it from the degree of parallelism and the batch size.
  + max_batch_results_ = std::min(
  + kMaxBatchResults,
  + CeilDiv(params.dataset->num_parallel_calls_ == model::kAutotune
  + ? port::NumSchedulableCPUs() // maximum parallelism
  + : params.dataset->num_parallel_calls_,
  + params.dataset->batch_size_));
  + }
  ~Iterator() override {
  CancelThreads(/*wait=*/true);
  @@ -272,13 +285,17 @@ class MapAndBatchDatasetOp::Dataset : public DatasetBase {
  TraceMeMetadata GetTraceMeMetadata() const override {
  int64 parallelism = -1;
  + int64 max_batch_results = -1;
  // NOTE: We only set the parallelism value if the lock can be acquired
  // right away to avoid introducing tracing overhead.
  if (mu_->try_lock()) {
  parallelism = num_parallel_calls_->value;
  + max_batch_results = max_batch_results_;
  mu_->unlock();
  }
  auto result = dataset()->traceme_metadata_;
  + result.push_back(std::make_pair(
  + "max_batch_results", strings::Printf("%lld", max_batch_results)));
  result.push_back(
  std::make_pair("parallelism", strings::Printf("%lld", parallelism)));
  return result;
  