Separate OpNode to PrepackNode/ExecuteNode (#2102)

Summary:
Pull Request resolved: #2102

Derived classes of `OpNode` are currently used only for prepack or execute, never both. This means they need not have both API.

Inspired by [Stephen's comment](https://www.internalfb.com/diff/D53982441?dst_version_fbid=370105355800543&transaction_fbid=940226924354801), we will build on `ExecuteNode` to be initialized with member functions for `create_params_block()`, `get_shader()`, etc. `PrepackNode` doesn't need these members. Hence, it makes sense to separate the classes.

Feel free to suggest better names. I don't really like mine.
ghstack-source-id: 216481872
exported-using-ghexport

Reviewed By: SS-JIA

Differential Revision: D54042646

fbshipit-source-id: a8965d69c94cdb3fe9837e9b83b6db8a877949f0

Author: jorgep31415

backends/vulkan/runtime/graph/Graph.cpp

@@ -193,8 +193,8 @@ void ComputeGraph::copy_from_staging(
 }
 
 void ComputeGraph::encode_prepack() {
-  for (std::unique_ptr<OpNode>& node : prepack_nodes_) {
-    node->encode_prepack(this);
+  for (std::unique_ptr<PrepackNode>& node : prepack_nodes_) {
+    node->encode(this);
   }
 }
 
@@ -216,8 +216,8 @@ void ComputeGraph::encode_execute() {
     shared_object.bind_users(this);
   }
 
-  for (std::unique_ptr<OpNode>& node : execute_nodes_) {
-    node->encode_execute(this);
+  for (std::unique_ptr<ExecuteNode>& node : execute_nodes_) {
+    node->encode(this);
   }
 }
 

backends/vulkan/runtime/graph/Graph.h

@@ -33,32 +33,52 @@ struct IOValueRef {
 class ComputeGraph;
 
 /*
- * Represents a single op in a ML model. In graph mode, ops will be implemented
- * introducing a derived class that implements encode_execute, which will
- * implement encoding of the shader corresponding to the op into the command
- * buffer of a ComputeGraph, as well as encode_prepack, which will implement
+ * Represents a single prepacking op in a ML model. In graph mode, ops will be
+ * implemented in a derived class that implements encode, which will implement
  * encoding of shaders transferring necessary data (such as weights and biases)
- * to the GPU, wherever prepacking is necessary.
+ * to the GPU.
  */
-class OpNode {
+class PrepackNode {
   friend class ComputeGraph;
 
  public:
-  OpNode(ValueRef input, ValueRef output) : inputs_{input}, outputs_{output} {}
-  OpNode(
+  PrepackNode(ValueRef tref, ValueRef packed) : tref_{tref}, packed_{packed} {}
+
+  virtual ~PrepackNode() = default;
+
+ protected:
+  ValueRef tref_;
+  ValueRef packed_;
+
+ public:
+  virtual void encode(ComputeGraph* graph) const = 0;
+};
+
+/*
+ * Represents a single execution op in a ML model. In graph mode, ops will be
+ * implemented in a derived class that implements encode, which will implement
+ * encoding of the shader corresponding to the op into the command buffer of a
+ * ComputeGraph.
+ */
+class ExecuteNode {
+  friend class ComputeGraph;
+
+ public:
+  ExecuteNode(ValueRef input, ValueRef output)
+      : inputs_{input}, outputs_{output} {}
+  ExecuteNode(
       const std::vector<ValueRef>& inputs,
       const std::vector<ValueRef>& outputs)
       : inputs_(inputs), outputs_(outputs) {}
 
-  virtual ~OpNode() = default;
+  virtual ~ExecuteNode() = default;
 
  protected:
   std::vector<ValueRef> inputs_;
   std::vector<ValueRef> outputs_;
 
  public:
-  virtual void encode_prepack(ComputeGraph* graph) const {}
-  virtual void encode_execute(ComputeGraph* graph) const {}
+  virtual void encode(ComputeGraph* graph) const = 0;
 };
 
 struct SharedObject {
@@ -99,8 +119,8 @@ class ComputeGraph final {
   std::vector<SharedObject> shared_objects_;
   std::vector<Value> values_;
 
-  std::vector<std::unique_ptr<OpNode>> prepack_nodes_;
-  std::vector<std::unique_ptr<OpNode>> execute_nodes_;
+  std::vector<std::unique_ptr<PrepackNode>> prepack_nodes_;
+  std::vector<std::unique_ptr<ExecuteNode>> execute_nodes_;
 
   std::vector<ValueRef> inputs_;
   std::vector<ValueRef> outputs_;
@@ -149,11 +169,11 @@ class ComputeGraph final {
     VK_THROW("Could not get dtype of value with type ", val.type());
   }
 
-  inline std::vector<std::unique_ptr<OpNode>>& prepack_nodes() {
+  inline std::vector<std::unique_ptr<PrepackNode>>& prepack_nodes() {
     return prepack_nodes_;
   }
 
-  inline std::vector<std::unique_ptr<OpNode>>& execute_nodes() {
+  inline std::vector<std::unique_ptr<ExecuteNode>>& execute_nodes() {
     return execute_nodes_;
   }
 

backends/vulkan/runtime/graph/ops/Arithmetic.cpp

@@ -61,11 +61,11 @@ ValueRef add_arithmetic_node(
 }
 
 ArithmeticPrepack::ArithmeticPrepack(const ValueRef tref, const ValueRef packed)
-    : OpNode(tref, packed) {}
+    : PrepackNode(tref, packed) {}
 
-void ArithmeticPrepack::encode_prepack(ComputeGraph* graph) const {
-  TensorRef tref = graph->get_val(inputs_[0]).toTensorRef();
-  vTensor packed = graph->get_val(outputs_[0]).toTensor();
+void ArithmeticPrepack::encode(ComputeGraph* graph) const {
+  TensorRef tref = graph->get_val(tref_).toTensorRef();
+  vTensor packed = graph->get_val(packed_).toTensor();
 
   api::StorageBuffer staging(
       graph->context(), packed.dtype(), packed.gpu_nbytes());
@@ -83,9 +83,9 @@ ArithmeticNode::ArithmeticNode(
     const ValueRef out,
     const float alpha,
     const arithmetic::OpType optype)
-    : OpNode({t1, t2}, {out}), alpha_(alpha), optype_(optype) {}
+    : ExecuteNode({t1, t2}, {out}), alpha_(alpha), optype_(optype) {}
 
-void ArithmeticNode::encode_execute(ComputeGraph* graph) const {
+void ArithmeticNode::encode(ComputeGraph* graph) const {
   vTensor& in1 = graph->get_val(inputs_[0]).toTensor();
   vTensor& in2 = graph->get_val(inputs_[1]).toTensor();
   vTensor& out = graph->get_val(outputs_[0]).toTensor();

backends/vulkan/runtime/graph/ops/Arithmetic.h

@@ -34,14 +34,14 @@ ValueRef add_arithmetic_node(
     const arithmetic::OpType optype,
     const int64_t shared_object_idx = -1);
 
-class ArithmeticPrepack : public virtual OpNode {
+class ArithmeticPrepack : public virtual PrepackNode {
  public:
   explicit ArithmeticPrepack(const ValueRef tref, const ValueRef packed);
 
-  void encode_prepack(ComputeGraph* graph) const override;
+  void encode(ComputeGraph* graph) const override;
 };
 
-class ArithmeticNode : public virtual OpNode {
+class ArithmeticNode : public virtual ExecuteNode {
  public:
   explicit ArithmeticNode(
       const ValueRef t1,
@@ -50,7 +50,7 @@ class ArithmeticNode : public virtual OpNode {
       const float alpha,
       const arithmetic::OpType optype);
 
-  void encode_execute(ComputeGraph* graph) const override;
+  void encode(ComputeGraph* graph) const override;
 
  private:
   float alpha_;

backends/vulkan/runtime/graph/ops/Copy.cpp

@@ -27,9 +27,10 @@ ValueRef add_copy_node(ComputeGraph& graph, const ValueRef from) {
   return to;
 }
 
-CopyNode::CopyNode(const ValueRef from, const ValueRef to) : OpNode(from, to) {}
+CopyNode::CopyNode(const ValueRef from, const ValueRef to)
+    : ExecuteNode(from, to) {}
 
-void CopyNode::encode_execute(ComputeGraph* graph) const {
+void CopyNode::encode(ComputeGraph* graph) const {
   api::PipelineBarrier pipeline_barrier{};
 
   vTensor& from_tensor = graph->get_val(inputs_[0]).toTensor();

backends/vulkan/runtime/graph/ops/Copy.h

@@ -19,11 +19,11 @@ namespace vulkan {
 void add_copy_node(ComputeGraph& graph, const ValueRef from, const ValueRef to);
 ValueRef add_copy_node(ComputeGraph& graph, const ValueRef from);
 
-class CopyNode : public virtual OpNode {
+class CopyNode : public virtual ExecuteNode {
  public:
   explicit CopyNode(const ValueRef from, const ValueRef to);
 
-  void encode_execute(ComputeGraph* graph) const override;
+  void encode(ComputeGraph* graph) const override;
 };
 
 } // namespace vulkan

backends/vulkan/runtime/graph/ops/Staging.cpp

@@ -98,9 +98,9 @@ void encode_copy_from_vtensor(
       VK_NULL_HANDLE);
 }
 
-StagingNode::StagingNode(ValueRef from, ValueRef to) : OpNode(from, to) {}
+StagingNode::StagingNode(ValueRef from, ValueRef to) : ExecuteNode(from, to) {}
 
-void StagingNode::encode_execute(ComputeGraph* graph) const {
+void StagingNode::encode(ComputeGraph* graph) const {
   Value& in_val = graph->get_val(inputs_[0]);
   Value& out_val = graph->get_val(outputs_[0]);
 

backends/vulkan/runtime/graph/ops/Staging.h

@@ -84,11 +84,11 @@ void encode_copy_from_vtensor(
 /*
  * OpNode that allows copying data into and out of a staging buffer.
  */
-class StagingNode : public virtual OpNode {
+class StagingNode : public virtual ExecuteNode {
  public:
   explicit StagingNode(ValueRef from, ValueRef to);
 
-  void encode_execute(ComputeGraph* graph) const override;
+  void encode(ComputeGraph* graph) const override;
 };
 
 } // namespace vulkan