Add distributed annotations

Summary: Annotations for DAI

Reviewed By: duc0

Differential Revision: D9805867

fbshipit-source-id: 9ce2d9f3984817510ec8362a281f39878aad55e7

Author: bwasti

Diff for: caffe2/opt/annotations.cc

@@ -0,0 +1,86 @@
+#include "caffe2/opt/annotations.h"
+
+namespace caffe2 {
+
+using namespace nom::repr;
+
+void Caffe2Annotation::setOperatorDef(const caffe2::OperatorDef& opDef) {
+  OpDef = opDef;
+  OpDefExists = true;
+}
+
+bool Caffe2Annotation::hasOperatorDef() const {
+  return OpDefExists;
+}
+
+const caffe2::OperatorDef& Caffe2Annotation::getOperatorDef() const {
+  CAFFE_ENFORCE(
+      OpDefExists,
+      "OperatorDef was never set.  Use Caffe2Annotation::setOperatorDef.");
+  return OpDef;
+}
+caffe2::OperatorDef* Caffe2Annotation::getMutableOperatorDef() {
+  CAFFE_ENFORCE(
+      OpDefExists,
+      "OperatorDef was never set.  Use Caffe2Annotation::setOperatorDef.");
+  return &OpDef;
+}
+
+// Distributed annotations
+void Caffe2Annotation::setDevice(std::string device) {
+  Device = device;
+}
+const std::string Caffe2Annotation::getDevice() const {
+  return Device;
+}
+
+void Caffe2Annotation::setDeviceType(int device) {
+  DeviceType = device;
+}
+int Caffe2Annotation::getDeviceType() const {
+  return DeviceType;
+}
+
+void Caffe2Annotation::setParallelization(
+    Caffe2Annotation::ParallelizationScheme s,
+    int num) {
+  parallelization_scheme_ = s;
+  parallelization_ = num;
+}
+
+Caffe2Annotation::ParallelizationScheme
+Caffe2Annotation::getParallelizationScheme() const {
+  return parallelization_scheme_;
+}
+
+int Caffe2Annotation::getParallelization() const {
+  return parallelization_;
+}
+
+void Caffe2Annotation::setKeyNode(NNGraph::NodeRef n) {
+  key_node_ = n;
+}
+const NNGraph::NodeRef& Caffe2Annotation::getKeyNode() const {
+  CAFFE_ENFORCE(key_node_, "No key node has been annotated");
+  return key_node_;
+}
+void Caffe2Annotation::setLengthNode(NNGraph::NodeRef n) {
+  length_node_ = n;
+}
+const NNGraph::NodeRef& Caffe2Annotation::getLengthNode() const {
+  CAFFE_ENFORCE(length_node_, "No length node has been annotated");
+  return length_node_;
+}
+
+void Caffe2Annotation::setComponentLevels(std::vector<std::string> components) {
+  component_levels_ = components;
+}
+std::vector<std::string> Caffe2Annotation::getComponentLevels() const {
+  return component_levels_;
+}
+
+bool Caffe2Annotation::classof(const Annotation* A) {
+  return A->getKind() == AnnotationKind::Caffe2;
+}
+
+} // namespace caffe2

Diff for: caffe2/opt/annotations.h

@@ -0,0 +1,63 @@
+#pragma once
+
+#include "caffe2/core/common.h"
+#include "caffe2/core/logging.h"
+#include "caffe2/proto/caffe2_pb.h"
+#include "nomnigraph/Representations/NeuralNet.h"
+
+namespace caffe2 {
+
+class Caffe2Annotation : public nom::repr::Annotation {
+ public:
+  Caffe2Annotation() : Annotation(AnnotationKind::Caffe2) {}
+  Caffe2Annotation(std::string device)
+      : Annotation(AnnotationKind::Caffe2), Device(device) {}
+  virtual ~Caffe2Annotation() {}
+
+  void setOperatorDef(const caffe2::OperatorDef& opDef);
+  bool hasOperatorDef() const;
+  const caffe2::OperatorDef& getOperatorDef() const;
+  caffe2::OperatorDef* getMutableOperatorDef();
+
+  // Distributed annotations
+  void setDevice(std::string device);
+  const std::string getDevice() const;
+  void setDeviceType(int device);
+  int getDeviceType() const;
+
+  enum class ParallelizationScheme {
+    none,
+    split_by_batch,
+    split_by_length,
+    shard,
+    shard_by_number
+  };
+  void setParallelization(ParallelizationScheme, int num = -1);
+  ParallelizationScheme getParallelizationScheme() const;
+  int getParallelization() const;
+
+  void setKeyNode(nom::repr::NNGraph::NodeRef);
+  const nom::repr::NNGraph::NodeRef& getKeyNode() const;
+  void setLengthNode(nom::repr::NNGraph::NodeRef);
+  const nom::repr::NNGraph::NodeRef& getLengthNode() const;
+
+  void setComponentLevels(std::vector<std::string> components);
+  std::vector<std::string> getComponentLevels() const;
+
+  static bool classof(const Annotation* A);
+
+ private:
+  std::string Device = "";
+  caffe2::OperatorDef OpDef;
+  bool OpDefExists = false;
+
+  // Distributed annotations
+  int DeviceType = caffe2::DeviceTypeProto::PROTO_CPU;
+  ParallelizationScheme parallelization_scheme_ = ParallelizationScheme::none;
+  int parallelization_ = -1;
+  nom::repr::NNGraph::NodeRef key_node_ = nullptr;
+  nom::repr::NNGraph::NodeRef length_node_ = nullptr;
+  std::vector<std::string> component_levels_;
+};
+
+} // namespace caffe2

Diff for: caffe2/opt/converter.h

@@ -3,6 +3,7 @@
 
 #include "caffe2/core/common.h"
 #include "caffe2/core/logging.h"
+#include "caffe2/opt/annotations.h"
 #include "caffe2/proto/caffe2_pb.h"
 #include "nomnigraph/Graph/Graph.h"
 #include "nomnigraph/Representations/ControlFlow.h"
@@ -12,62 +13,7 @@
 
 namespace caffe2 {
 
-class Caffe2Annotation : public nom::repr::Annotation {
-public:
-  Caffe2Annotation() : Annotation(AnnotationKind::Caffe2) {}
-  Caffe2Annotation(std::string device)
-      : Annotation(AnnotationKind::Caffe2), Device(device) {}
-  virtual ~Caffe2Annotation() {}
-
-  void setDevice(std::string device) {
-    Device = device;
-  }
-  const std::string getDevice() const {
-    return Device;
-  }
-
-  void setDeviceType(int device) {
-    DeviceType = device;
-  }
-  int getDeviceType() const {
-    return DeviceType;
-  }
-
-  void setOperatorDef(const caffe2::OperatorDef& opDef) {
-    OpDef = opDef;
-    OpDefExists = true;
-  }
-
-  bool hasOperatorDef() const {
-    return OpDefExists;
-  }
-
-  const caffe2::OperatorDef& getOperatorDef() const {
-    CAFFE_ENFORCE(
-        OpDefExists,
-        "OperatorDef was never set.  Use Caffe2Annotation::setOperatorDef.");
-    return OpDef;
-  }
-  caffe2::OperatorDef* getMutableOperatorDef() {
-    CAFFE_ENFORCE(
-        OpDefExists,
-        "OperatorDef was never set.  Use Caffe2Annotation::setOperatorDef.");
-    return &OpDef;
-  }
-
-  static bool classof(const Annotation *A) {
-    return A->getKind() == AnnotationKind::Caffe2;
-  }
-
-private:
-  std::string Device = "";
-  caffe2::OperatorDef OpDef;
-  bool OpDefExists = false;
-  int DeviceType = caffe2::DeviceTypeProto::PROTO_CPU;
-};
-
 CAFFE2_API nom::repr::NNModule convertToNNModule(caffe2::NetDef &net, bool strict = false);
-
 CAFFE2_API caffe2::NetDef convertToCaffe2Proto(nom::repr::NNModule&);
 
 // Pass in an oldNet to copy all the attributes of that network.

Diff for: caffe2/python/nomnigraph_test.py

@@ -223,3 +223,23 @@ def test_annotation_from_graph(self):
         node.setAnnotation(annot)
         new_annot = node.getAnnotation()
         assert new_annot.getDeviceType() == 7
+
+    def test_distribute_annotations(self):
+        nn = ng.NNModule()
+        key = nn.dataFlow.createNode(ng.NeuralNetData("key"))
+        length = nn.dataFlow.createNode(ng.NeuralNetData("length"))
+        node = nn.dataFlow.createNode(ng.NeuralNetOperator("TestOp"))
+
+        annot = ng.Annotation()
+        annot.setKeyNode(key)
+        annot.setLengthNode(length)
+        annot.setComponentLevels(["", "test", "woot"])
+
+        node.setAnnotation(annot)
+
+        new_annot = node.getAnnotation()
+        #assert new_annot.getLengthNode() == length
+        assert new_annot.getKeyNode() == key
+        assert len(new_annot.getComponentLevels()) == 3
+        assert new_annot.getComponentLevels()[0] == ""
+        assert new_annot.getComponentLevels()[2] == "woot"

Diff for: caffe2/python/pybind_state_nomni.cc

@@ -199,9 +199,19 @@ void addNomnigraphMethods(pybind11::module& m) {
             return nn::get<nom::repr::Tensor>(n);
           },
           py::return_value_policy::reference_internal)
+      .def_property(
+          "annotation",
+          [](NNGraph::NodeRef n) { return getOrAddCaffe2Annotation(n); },
+          [](NNGraph::NodeRef n, Caffe2Annotation annot) {
+            auto* nnOp = nn::get<NeuralNetOperator>(n);
+            nnOp->setAnnotation(
+                nom::util::make_unique<Caffe2Annotation>(annot));
+          },
+          py::return_value_policy::copy)
       .def(
           "getAnnotation",
-          [](NNGraph::NodeRef n) { return getOrAddCaffe2Annotation(n); })
+          [](NNGraph::NodeRef n) { return getOrAddCaffe2Annotation(n); },
+          py::return_value_policy::copy)
       .def(
           "setAnnotation",
           [](NNGraph::NodeRef n, Caffe2Annotation annot) {
@@ -327,7 +337,19 @@ void addNomnigraphMethods(pybind11::module& m) {
       .def("setDevice", &Caffe2Annotation::setDevice)
       .def("getDevice", &Caffe2Annotation::getDevice)
       .def("setDeviceType", &Caffe2Annotation::setDeviceType)
-      .def("getDeviceType", &Caffe2Annotation::getDeviceType);
+      .def("getDeviceType", &Caffe2Annotation::getDeviceType)
+      .def("setKeyNode", &Caffe2Annotation::setKeyNode)
+      .def(
+          "getKeyNode",
+          &Caffe2Annotation::getKeyNode,
+          py::return_value_policy::reference)
+      .def("setLengthNode", &Caffe2Annotation::setLengthNode)
+      .def(
+          "getLengthNode",
+          &Caffe2Annotation::getLengthNode,
+          py::return_value_policy::reference)
+      .def("setComponentLevels", &Caffe2Annotation::setComponentLevels)
+      .def("getComponentLevels", &Caffe2Annotation::getComponentLevels);
 }
 
 REGISTER_PYBIND_ADDITION(addNomnigraphMethods);