path: root/tensorflow/core/framework/model.cc

/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/core/framework/model.h"

#include <memory>

#include "tensorflow/core/lib/gtl/map_util.h"

namespace tensorflow {
namespace data {
namespace model {

// TODO(jsimsa): Use `Node` subclassing instead of types and node statements.
void Node::CollectTunables(
    std::vector<std::shared_ptr<Node::Tunable>>* tunables) {
  mutex_lock l(mu_);
  for (auto input : inputs_) {
    input->CollectTunables(tunables);
  }
  switch (type_) {
    case Type::MAP_AND_BATCH:
    case Type::PARALLEL_INTERLEAVE_V2:
    case Type::PARALLEL_MAP: {
      if (auto* tunable_param =
              gtl::FindOrNull(tunable_params_, "parallelism")) {
        tunables->push_back(*tunable_param);
      }
      return;
    }
    default:
      return;
  }
}

int64 Node::GetParameterValue(const string& name) {
  if (auto* tunable_param = gtl::FindOrNull(tunable_params_, name)) {
    return (*tunable_param)->value;
  }
  return constant_params_[name];
}

int64 Node::ProcessingTimeLocked() {
  switch (type_) {
    case Type::BATCH:
    case Type::MAP_AND_BATCH:
    case Type::PADDED_BATCH: {
      int64 batch_size = GetParameterValue("batch_size");
      return NanosPerElementLocked() + batch_size * ProcessingTimeForInputs();
    }
    case Type::FILTER: {
      std::shared_ptr<Node> input = inputs_.front();
      double ratio = static_cast<double>(input->num_elements()) /
                     static_cast<double>(num_elements_);
      return NanosPerElementLocked() +
             static_cast<int64>(ratio *
                                static_cast<double>(ProcessingTimeForInputs()));
    }
    case Type::FLAT_MAP:
    case Type::INTERLEAVE:
    case Type::PARALLEL_INTERLEAVE:
    case Type::PARALLEL_INTERLEAVE_V2: {
      // TODO(jsimsa): model the first input
      // TODO(jsimsa): use processing time history as a prior for future inputs
      if (inputs_.size() <= 1) {
        return NanosPerElementLocked();
      }
      int64 processing_time =
          ProcessingTimeForInputs() - inputs_.front()->ProcessingTime();
      return NanosPerElementLocked() +
             static_cast<double>(processing_time) /
                 static_cast<double>(inputs_.size() - 1);
    }
    case Type::CACHE:
    case Type::CONCATENATE:
    case Type::MAP:
    case Type::PARALLEL_MAP:
    case Type::PREFETCH:
      // TODO(jsimsa): use processing time history as a prior for future inputs
    case Type::REPEAT:
    case Type::SHUFFLE:
    case Type::SKIP:
    case Type::TAKE:
    case Type::ZIP: {
      return NanosPerElementLocked() + ProcessingTimeForInputs();
    }
    default:
      return NanosPerElementLocked();
  }
}

int64 Node::OutputTimeLocked(std::vector<int64>* input_times) {
  switch (type_) {
    case Type::BATCH:
    case Type::PADDED_BATCH: {
      double batch_size = GetParameterValue("batch_size");
      int64 old_value = (*input_times)[input_times->size() - 1];
      (*input_times)[input_times->size() - 1] = static_cast<int64>(
          static_cast<double>(old_value + NanosPerElementLocked()) /
          batch_size);
      auto cleanup = gtl::MakeCleanup([input_times, old_value]() {
        (*input_times)[input_times->size() - 1] = old_value;
      });
      return NanosPerElementLocked() +
             batch_size * OutputTimeForInputs(input_times);
    }
    case Type::FILTER: {
      std::shared_ptr<Node> input = inputs_.front();
      int64 old_value = (*input_times)[input_times->size() - 1];
      double ratio = static_cast<double>(input->num_elements()) /
                     static_cast<double>(num_elements_);
      (*input_times)[input_times->size() - 1] = static_cast<int64>(
          static_cast<double>(old_value + NanosPerElementLocked()) / ratio);
      auto cleanup = gtl::MakeCleanup([input_times, old_value]() {
        (*input_times)[input_times->size() - 1] = old_value;
      });
      return NanosPerElementLocked() +
             static_cast<int64>(
                 static_cast<double>(OutputTimeForInputs(input_times)) * ratio);
    }
    case Type::FLAT_MAP:
    case Type::INTERLEAVE: {
      // TODO(jsimsa): model the first input
      // TODO(jsimsa): use cycle length metadata instead of `inputs_.size() - 1`
      if (inputs_.size() <= 1) {
        return NanosPerElementLocked();
      }
      int64 delta =
          static_cast<int64>(static_cast<double>(NanosPerElementLocked()) *
                             static_cast<double>(inputs_.size() - 1));
      (*input_times)[input_times->size() - 1] += delta;
      auto cleanup = gtl::MakeCleanup([input_times, delta]() {
        (*input_times)[input_times->size() - 1] -= delta;
      });
      int64 output_time = OutputTimeForInputs(input_times) -
                          inputs_.front()->OutputTime(input_times);
      return NanosPerElementLocked() +
             static_cast<double>(output_time) /
                 static_cast<double>(inputs_.size() - 1);
    }
    case Type::MAP_AND_BATCH: {
      double batch_size = GetParameterValue("batch_size");
      double parallelism = GetParameterValue("parallelism");
      int64 delta =
          static_cast<int64>(static_cast<double>(NanosPerElementLocked()) /
                             (batch_size * parallelism));
      input_times->push_back(delta);
      auto cleanup =
          gtl::MakeCleanup([input_times]() { input_times->pop_back(); });
      int64 output_time = static_cast<int64>(
          static_cast<double>(NanosPerElementLocked()) / parallelism +
          batch_size * OutputTimeForInputs(input_times));
      return std::max(0LL,
                      output_time - input_times->at(input_times->size() - 2));
    }
    case Type::PARALLEL_INTERLEAVE: {
      // TODO(jsimsa): model the first input
      if (inputs_.size() <= 1) {
        return NanosPerElementLocked();
      }
      int64 delta = static_cast<double>(NanosPerElementLocked()) *
                    static_cast<double>(inputs_.size() - 1);
      input_times->push_back(delta);
      auto cleanup =
          gtl::MakeCleanup([input_times]() { input_times->pop_back(); });
      int64 inputs_output_time = OutputTimeForInputs(input_times) -
                                 inputs_.front()->OutputTime(input_times);
      double parallelism = GetParameterValue("parallelism");
      int64 output_time =
          NanosPerElementLocked() + ((static_cast<double>(inputs_output_time) /
                                      static_cast<double>(inputs_.size() - 1)) /
                                     parallelism);
      return std::max(0LL,
                      output_time - input_times->at(input_times->size() - 2));
    }
    case Type::PARALLEL_INTERLEAVE_V2: {
      // TODO(jsimsa): model the first input
      if (inputs_.size() <= 1) {
        return NanosPerElementLocked();
      }
      int64 delta = static_cast<double>(NanosPerElementLocked()) *
                    static_cast<double>(inputs_.size() - 1);
      input_times->push_back(delta);
      auto cleanup =
          gtl::MakeCleanup([input_times]() { input_times->pop_back(); });
      int64 inputs_output_time = OutputTimeForInputs(input_times) -
                                 inputs_.front()->OutputTime(input_times);
      double parallelism =
          std::min(static_cast<int>(GetParameterValue("cycle_length")),
                   static_cast<int>(GetParameterValue("parallelism")));
      int64 output_time =
          NanosPerElementLocked() + ((static_cast<double>(inputs_output_time) /
                                      static_cast<double>(inputs_.size() - 1)) /
                                     parallelism);
      return std::max(0LL,
                      output_time - input_times->at(input_times->size() - 2));
    }
    case Type::PARALLEL_MAP: {
      double parallelism =
          std::min(port::NumSchedulableCPUs(),
                   static_cast<int>(GetParameterValue("parallelism")));
      int64 delta = static_cast<int64>(
          static_cast<double>(NanosPerElementLocked()) / parallelism);
      input_times->push_back(delta);
      auto cleanup =
          gtl::MakeCleanup([input_times]() { input_times->pop_back(); });
      int64 output_time =
          static_cast<double>(NanosPerElementLocked()) / parallelism +
          OutputTimeForInputs(input_times);
      return std::max(0LL,
                      output_time - input_times->at(input_times->size() - 2));
    }
    case Type::PREFETCH: {
      int64 delta = NanosPerElementLocked();
      input_times->push_back(delta);
      auto cleanup =
          gtl::MakeCleanup([input_times]() { input_times->pop_back(); });
      return std::max(0LL, NanosPerElementLocked() +
                               OutputTimeForInputs(input_times) -
                               input_times->at(input_times->size() - 2));
    }
    case Type::CACHE:
    case Type::CONCATENATE:
    case Type::MAP:
    case Type::REPEAT:
    case Type::SHUFFLE:
    case Type::SKIP:
    case Type::TAKE:
    case Type::ZIP: {
      int64 delta = NanosPerElementLocked();
      (*input_times)[input_times->size() - 1] += delta;
      auto cleanup = gtl::MakeCleanup([input_times, delta]() {
        (*input_times)[input_times->size() - 1] -= delta;
      });
      return NanosPerElementLocked() + OutputTimeForInputs(input_times);
    }
    default:
      return NanosPerElementLocked();
  }
}

std::shared_ptr<Node> Model::AddNode(const string& name,
                                     const string& output_name) {
  mutex_lock l(mu_);
  std::shared_ptr<Node> output;
  auto it = lookup_table_.find(output_name);
  if (it != lookup_table_.end()) {
    output = it->second;
  }
  std::shared_ptr<Node> node(new Node(id_counter_++, output));
  if (!output_) {
    output_ = node;
  }
  if (output) {
    output->add_input(node);
  }
  lookup_table_.insert(std::make_pair(name, node));
  return node;
}

std::shared_ptr<Node> Model::LookupNode(const string& name) {
  tf_shared_lock l(mu_);
  std::shared_ptr<Node> result;
  auto it = lookup_table_.find(name);
  if (it != lookup_table_.end()) {
    result = it->second;
  }
  return result;
}

// The optimization algorithm starts by setting all tunable parallelism
// parameters to 1. It then repeatedly identifies the parameter that whose
// increase in parallelism decreases the output time the most. This process is
// repeated until all parameters reach their maximum values or the
// projected output time is less than or equal to the processing time needed to
// produce an element divided by CPU budget.
void Model::Optimize(int64 cpu_budget) {
  mutex_lock l(optimization_mu_);
  std::vector<std::shared_ptr<Node::Tunable>> tunables;
  {
    mutex_lock l2(mu_);
    const int64 processing_time = ProcessingTime();
    tunables = CollectTunables();
    for (auto tunable : tunables) {
      tunable->value = 1;
    }
    while (true) {
      const int64 output_time = OutputTime();
      bool all_tunables = true;
      for (auto& tunable : tunables) {
        if (tunable->value < tunable->max) {
          all_tunables = false;
          break;
        }
      }
      if (output_time < processing_time / cpu_budget || all_tunables) {
        break;
      }
      int64 best_delta = -1;
      Node::Tunable* best_tunable = nullptr;
      for (auto& tunable : tunables) {
        if (tunable->value == tunable->max) {
          continue;
        }
        tunable->value++;
        int64 delta = output_time - OutputTime();
        if (delta > best_delta) {
          best_delta = delta;
          best_tunable = tunable.get();
        }
        tunable->value--;
      }
      if (best_tunable) {
        // NOTE: This can happen because we are performing the optimization
        // while the model data is changing. If this becomes an issue, we should
        // look into performing the optimization using a model snapshot.
        break;
      }
      best_tunable->value++;
    }
  }
  // The `set_fn` functions should be invoked without holding a lock to avoid a
  // potential deadlock.
  for (auto& tunable : tunables) {
    tunable->set_fn(tunable->value);
  }
}

void Model::RemoveNode(const string& prefix) {
  // Nodes are not allowed to be removed when optimization is in progress to
  // prevent the optimization from trying to access an iterator that was
  // concurrently deleted.
  mutex_lock l(optimization_mu_);
  mutex_lock l2(mu_);
  lookup_table_.erase(prefix);
}

std::vector<std::shared_ptr<Node::Tunable>> Model::CollectTunables() {
  std::vector<std::shared_ptr<Node::Tunable>> tunables;
  output_->CollectTunables(&tunables);
  return tunables;
}

int64 Model::OutputTime() {
  std::vector<int64> input_times(1, 0);
  return output_->OutputTime(&input_times);
}

int64 Model::ProcessingTime() { return output_->ProcessingTime(); }

}  // namespace model
}  // namespace data
}  // namespace tensorflow