computation_graph.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <memory>
#include <optional>
#include <string>
#include <vector>
 
#include "pulsatrix/node.hpp"
#include "pulsatrix/op_type.hpp"
#include "pulsatrix/shape.hpp"
 
namespace pulsatrix {
 
class ComputationGraph {
public:
    NodeId add_node(OpType op_type, Shape shape, std::optional<std::string> label = std::nullopt,
                     std::vector<NodeId> parent_ids = {});
 
    [[nodiscard]] const Node& node(NodeId id) const;
 
    [[nodiscard]] size_t node_count() const { return nodes_.size(); }
 
    [[nodiscard]] std::vector<NodeId> nodes_by_op_type(OpType op_type) const;
 
    [[nodiscard]] std::vector<NodeId> topological_order() const;
 
private:
    std::vector<std::unique_ptr<Node>> nodes_;
};
 
}  // namespace pulsatrix