[Graph class] add State to validate order of steps

Author: artemrakhov-glow

include/glow/Graph/Graph.h

@@ -21,6 +21,18 @@ using UnsignedArrayRef = llvm::ArrayRef<size_t>;
 
 /// Represents the compute graph.
 class Graph final : public Named {
+public:
+  enum class State {
+    Created,
+    Differentiated,
+    Lowered,
+    Optimized,
+    IRGenerated,
+  };
+
+private:
+  /// A current state of the graph.
+  State state_{State::Created};
   /// A uniqued list of types in the module. Types in this list can be equated
   /// by comparing their addresses.
   TypesList types_{};
@@ -49,13 +61,17 @@ class Graph final : public Named {
 
   /// Inserts the node \p N to the list of nodes, and returns the inserted node.
   template <class NodeTy> NodeTy *addNode(NodeTy *N) {
+    assert(state_ < State::IRGenerated &&
+           "Trying to add Node when IR is already generated.");
     uniqueNames(N);
     nodes_.push_back(N);
     return N;
   }
 
   /// Inserts the variable \p V to the list of variables.
   Variable *addVar(Variable *V) {
+    assert(state_ < State::IRGenerated &&
+           "Trying to add Variable when IR is already generated.");
     uniqueNames(V);
     vars_.push_back(V);
     return V;
@@ -199,6 +215,13 @@ class Graph final : public Named {
   /// Returns nullptr if there is no gradient variable
   /// related to this variable.
   Variable *getGradientVariable(Variable *V);
+
+  /// Resets current state to Created.
+  void resetState();
+
+  /// Verifies that current state of the graph is not later then \p s
+  /// and assigns current state to be \p s.
+  void advanceState(State s);
 };
 
 struct TrainingConfig;

src/glow/ExecutionEngine/ExecutionEngine.cpp

@@ -88,6 +88,7 @@ void ExecutionEngine::compile(CompilationMode mode) {
   // Wipe out the module and start a new compilation process.
   M_->clear();
   IP_->clear();
+  G_->resetState();
 
   if (mode != CompilationMode::Infer) {
     generateGradientNodes(*G_, getConfig(), mode);

src/glow/Graph/Graph.cpp

@@ -91,6 +91,7 @@ std::string Graph::uniqueName(llvm::StringRef name) {
 void Graph::uniqueNames(Node *N) { N->setName(uniqueName(N->getName())); }
 
 void Graph::addGradientVariable(Variable *V, Variable *GradV) {
+  advanceState(State::Differentiated);
   grads_.push_back({V, GradV});
 }
 
@@ -650,3 +651,10 @@ void Graph::verify() const {
     llvm_unreachable("Multiple nodes with the same name");
   }
 }
+
+void Graph::resetState() { state_ = State::Created; }
+
+void Graph::advanceState(State s) {
+  assert(state_ <= s && "Wrong order of actions with a graph.");
+  state_ = s;
+}

src/glow/IR/IRGen.cpp

@@ -530,6 +530,7 @@ struct IRGenVisitor : NodeWalker {
 } // namespace
 
 void Module::generateIR(CompilationMode mode) {
+  G_->advanceState(Graph::State::IRGenerated);
   G_->verify();
   IRGenVisitor irgen(this);
 

src/glow/Optimizer/GraphOptimizer.cpp

@@ -546,6 +546,8 @@ static void OptimizeSliceOfSplat(Graph &G) {
 }
 
 void glow::optimize(Graph &G, CompilationMode mode) {
+  G.advanceState(Graph::State::Optimized);
+
   // Sink transpose operations in an attempt to cancel them out.
   sinkCode(G);
 

src/glow/Optimizer/Lower.cpp

@@ -223,6 +223,7 @@ void lowerReluNode(Graph &graph, ReluNode &R) {
 }
 
 void glow::lower(Graph &G, CompilationMode mode) {
+  G.advanceState(Graph::State::Lowered);
   auto &nodes = G.getNodes();
 
   for (auto const &node : nodes) {