tensorflow--tensorflow/tensorflow/python/grappler/item_wrapper.cc

/* Copyright 2020 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 <algorithm>
#include <memory>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "pybind11/pybind11.h"  // from @pybind11
#include "pybind11/stl.h"  // from @pybind11
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/grappler/costs/graph_properties.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/grappler_item_builder.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/grappler/utils/topological_sort.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/python/lib/core/pybind11_status.h"

namespace py = pybind11;

// Manages disjoint sets of colocation groups using a union-find data structure
// with path compression and union by rank.
class ColocationGroups {
 public:
  // Ensures a node is tracked in the disjoint-set structure.
  void RegisterNode(absl::string_view node_name) { Find(node_name); }

  // Unions the colocation sets containing nodes x and y.
  void Group(absl::string_view x, absl::string_view y) {
    Rep* x_root = Find(x);
    Rep* y_root = Find(y);

    // x and y are already in the same set
    if (x_root == y_root) {
      return;
    }
    // x and y are not in same set, so we merge them
    // Use the occasion to strengthen what we know about the handle by merging
    // the information about the 2 subsets.
    if (x_root->rank < y_root->rank) {
      x_root->parent = y_root;
    } else if (x_root->rank > y_root->rank) {
      y_root->parent = x_root;
    } else {
      // Arbitrarily make one root the new parent
      y_root->parent = x_root;
      ++x_root->rank;
    }
  }

  // Extracts all disjoint colocation groups as a list of node name lists.
  std::vector<std::vector<std::string>> ExtractGroups() {
    std::vector<std::vector<std::string>> groups;
    groups.reserve(nodes_.size());
    absl::flat_hash_map<const Rep*, int> group_ids;
    for (const auto& [node_name, rep_ptr] : nodes_) {
      Rep* r = Find(rep_ptr.get());
      auto [it, inserted] = group_ids.try_emplace(r, groups.size());
      int id = it->second;
      if (inserted) {
        // If inserted, this is a new group. The value stored (groups.size())
        // is the index where the new group should be added.
        groups.emplace_back();
      }
      groups[id].push_back(node_name);
    }
    for (auto& g : groups) {
      std::sort(g.begin(), g.end());
    }
    std::sort(groups.begin(), groups.end());
    return groups;
  }

 private:
  struct Rep {
    // Parent in the tree used to encode the set.
    Rep* parent;
    // Rank in the tree, used to figure out how to compress the path to the root
    // of the tree.
    int rank;
  };

  Rep* Find(absl::string_view n) {
    // Try to emplace a new Rep. If the key already exists, try_emplace does
    // nothing and returns an iterator to the existing element. Otherwise,
    // it inserts a new unique_ptr<Rep> and returns an iterator to it.
    auto [it, inserted] = nodes_.try_emplace(n, std::make_unique<Rep>());
    if (inserted) {
      // First time processing this handle, initialize the new entry.
      Rep* raw_node = it->second.get();
      raw_node->parent = raw_node;
      raw_node->rank = 0;
      return raw_node;
    }
    return Find(it->second.get());
  }

  Rep* Find(Rep* node) {
    Rep* root = node->parent;
    while (root != root->parent) {
      root = root->parent;
    }
    while (node->parent != root) {
      Rep* next = node->parent;
      node->parent = root;
      node = next;
    }
    return root;
  }

  absl::flat_hash_map<std::string, std::unique_ptr<Rep>> nodes_;
};

PYBIND11_MAKE_OPAQUE(tensorflow::grappler::GrapplerItem);

PYBIND11_MODULE(_pywrap_tf_item, m) {
  py::class_<tensorflow::grappler::GrapplerItem> grappler_item(m,
                                                               "GrapplerItem");

  m.def("TF_NewItem",
        [](const py::bytes& serialized_metagraph, bool ignore_colocation,
           bool ignore_user_placement)
            -> std::unique_ptr<tensorflow::grappler::GrapplerItem> {
          tensorflow::MetaGraphDef metagraph;
          py::buffer_info info = py::buffer(serialized_metagraph).request();
          if (!metagraph.ParseFromArray(info.ptr, info.size)) {
            throw std::invalid_argument(
                "The MetaGraphDef could not be parsed as a valid protocol "
                "buffer");
          }
          if (metagraph.collection_def().count("train_op") == 0) {
            tsl::MaybeRaiseRegisteredFromStatus(absl::InvalidArgumentError(
                "train_op not specified in the metagraph"));
          }

          tensorflow::grappler::ItemConfig cfg;
          cfg.ignore_user_placement = ignore_user_placement;
          cfg.ignore_colocation = ignore_colocation;
          std::unique_ptr<tensorflow::grappler::GrapplerItem> item =
              tensorflow::grappler::GrapplerItemFromMetaGraphDef(
                  "item", metagraph, cfg);
          if (item == nullptr) {
            tsl::MaybeRaiseRegisteredFromStatus(
                absl::InvalidArgumentError("Invalid metagraph"));
          }
          return item;
        });

  m.def("TF_IdentifyImportantOps",
        [](tensorflow::grappler::GrapplerItem* item,
           bool sort_topologically) -> std::vector<std::string> {
          std::vector<const tensorflow::NodeDef*> main_ops =
              item->MainOpsFanin();
          std::vector<const tensorflow::NodeDef*> enqueue_ops =
              item->EnqueueOpsFanin();
          absl::flat_hash_set<std::string> op_names;
          for (const tensorflow::NodeDef* op : main_ops) {
            op_names.insert(op->name());
          }
          for (const tensorflow::NodeDef* op : enqueue_ops) {
            op_names.insert(op->name());
          }

          std::vector<std::string> ops;
          if (sort_topologically) {
            tensorflow::GraphDef subgraph;
            for (const tensorflow::NodeDef& node : item->graph.node()) {
              if (op_names.find(node.name()) != op_names.end()) {
                *subgraph.add_node() = node;
              }
            }
            tsl::MaybeRaiseRegisteredFromStatus(
                tensorflow::grappler::TopologicalSort(&subgraph));
            for (const tensorflow::NodeDef& node : subgraph.node()) {
              ops.push_back(node.name());
            }
          } else {
            ops.reserve(op_names.size());
            for (const auto& op_name : op_names) {
              ops.push_back(op_name);
            }
            std::sort(ops.begin(), ops.end());
          }
          return ops;
        });

  m.def("TF_GetOpProperties",
        [](tensorflow::grappler::GrapplerItem* item)
            -> std::unordered_map<std::string, std::vector<py::bytes>> {
          tensorflow::grappler::GraphProperties properties(*item);
          tsl::MaybeRaiseRegisteredFromStatus(
              properties.InferStatically(false));

          std::unordered_map<std::string, std::vector<py::bytes>> props;
          for (const tensorflow::NodeDef& node : item->graph.node()) {
            const std::string& node_name = node.name();
            const std::vector<tensorflow::OpInfo::TensorProperties>&
                output_props = properties.GetOutputProperties(node_name);

            std::vector<py::bytes> prop;
            prop.reserve(output_props.size());
            for (const tensorflow::OpInfo::TensorProperties& output_prop :
                 output_props) {
              prop.push_back(output_prop.SerializeAsString());
            }
            props[node_name] = std::move(prop);
          }
          return props;
        });

  m.def("TF_GetColocationGroups",
        [](tensorflow::grappler::GrapplerItem* item)
            -> std::vector<std::vector<std::string>> {
          ColocationGroups groupings;
          tensorflow::OpRegistry* registry = tensorflow::OpRegistry::Global();
          for (const tensorflow::NodeDef& node : item->graph.node()) {
            const tensorflow::OpDef* op_def;
            if (!registry->LookUpOpDef(node.op(), &op_def).ok()) {
              continue;
            }
            tensorflow::NameRangeMap inputs;
            if (!tensorflow::NameRangesForNode(node, *op_def, &inputs, nullptr)
                     .ok()) {
              continue;
            }
            for (const tensorflow::OpDef::ArgDef& arg : op_def->input_arg()) {
              if (!arg.is_ref()) {
                continue;
              }
              const auto& range = inputs[arg.name()];
              for (int i = range.first;
                   i < range.second && i < node.input_size(); ++i) {
                groupings.Group(node.name(),
                                tensorflow::grappler::NodeName(node.input(i)));
              }
            }
          }

          return groupings.ExtractGroups();
        });
}