Skip to content

[ET-VK] Separate OpNode to PrepackNode/ExecuteNode #2102

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions backends/vulkan/runtime/graph/Graph.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}

Expand All @@ -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);
}
}

Expand Down
50 changes: 35 additions & 15 deletions backends/vulkan/runtime/graph/Graph.h
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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_;
Expand Down Expand Up @@ -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_;
}

Expand Down
12 changes: 6 additions & 6 deletions backends/vulkan/runtime/graph/ops/Arithmetic.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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());
Expand All @@ -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();
Expand Down
8 changes: 4 additions & 4 deletions backends/vulkan/runtime/graph/ops/Arithmetic.h
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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_;
Expand Down
5 changes: 3 additions & 2 deletions backends/vulkan/runtime/graph/ops/Copy.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
4 changes: 2 additions & 2 deletions backends/vulkan/runtime/graph/ops/Copy.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions backends/vulkan/runtime/graph/ops/Staging.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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]);

Expand Down
4 changes: 2 additions & 2 deletions backends/vulkan/runtime/graph/ops/Staging.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down