Update visitor to support available operations (#320)

Update visitor to support available operations

Updates visitor base class to support currently implemented operations

Signed-off-by: Vladimir Plazun <v.plazun@partner.samsung.com>

diff --git a/contrib/nnc/libs/core/include/nnc/core/IR/model/visitor/visitor.h b/contrib/nnc/libs/core/include/nnc/core/IR/model/visitor/visitor.h
index 1bb8185..0d16e83 100644 (file)
--- a/contrib/nnc/libs/core/include/nnc/core/IR/model/visitor/visitor.h
+++ b/contrib/nnc/libs/core/include/nnc/core/IR/model/visitor/visitor.h
@@ -12,12 +12,34 @@ namespace ADT {
   class INode;
 }
 
+//Forward declare operations as we don't need anything but references
+namespace ops
+{
+  class ConcatOp;
+  class Conv2DOp;
+  class DepthwiseConv2DOp;
+  class SoftmaxOp;
+  class PoolOp;
+  class FullyConnectedOp;
+  class CappedReluOp;
+  class BiasAddOp;
+  class VariableOp;
+  class ReluOp;
+}
+
 class Visitor {
  public:
-  // Operation-specific methods like:
-  // virtual void visit(INode* node, OpDescriptionSubclass& op) = 0;
+  virtual void visit(ADT::INode *node, ops::ConcatOp &op) = 0;
+  virtual void visit(ADT::INode *node, ops::Conv2DOp &op) = 0;
+  virtual void visit(ADT::INode *node, ops::DepthwiseConv2DOp &op) = 0;
+  virtual void visit(ADT::INode *node, ops::SoftmaxOp &op) = 0;
+  virtual void visit(ADT::INode *node, ops::PoolOp &op) = 0;
+  virtual void visit(ADT::INode *node, ops::FullyConnectedOp &op) = 0;
+  virtual void visit(ADT::INode *node, ops::CappedReluOp &op) = 0;
+  virtual void visit(ADT::INode *node, ops::BiasAddOp &op) = 0;
+  virtual void visit(ADT::INode *node, ops::VariableOp &op) = 0;
+  virtual void visit(ADT::INode *node, ops::ReluOp &op) = 0;
 
-  virtual void visit(ADT::INode *node) = 0;
   virtual ~Visitor() = default;
 };