conv_node.h

Go to the documentation of this file.

 
#pragma once
 
#include "sdfg/data_flow/library_nodes/math/blas/blas_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/spatial_tensor_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/tensor_node.h"
 
namespace sdfg {
namespace math {
namespace tensor {
 
inline data_flow::LibraryNodeCode LibraryNodeType_Conv("ml::Conv");

class ConvNode : public SpatialTensorNode {
protected:
    symbolic::Expression output_channels_; 
    symbolic::Expression group_; 
    bool with_bias_;
 
public:
    ConvNode(
        size_t element_id,
        const DebugInfo& debug_info,
        const graph::Vertex vertex,
        data_flow::DataFlowGraph& parent,
        const std::vector<symbolic::Expression>& shape,
        const std::vector<symbolic::Expression>& kernel_shape,
        const std::vector<symbolic::Expression>& strides,
        const std::vector<symbolic::Expression>& pads,
        const std::vector<symbolic::Expression>& dilations,
        symbolic::Expression output_channels,
        symbolic::Expression group,
        bool with_bias = false,
        QuantizationType quantization = QUANTIZATION_MATCH_INPUTS,
        const data_flow::ImplementationType& impl_type = data_flow::ImplementationType_NONE
    );
 
    static int constexpr Y_INPUT_IDX = 0;
    static int constexpr X_INPUT_IDX = 1;
    static int constexpr W_INPUT_IDX = 2;
    static int constexpr B_INPUT_IDX = 3;

    symbolic::Expression output_channels() const {
        return output_channels_;
    }

    symbolic::Expression group() const {
        return group_;
    }
 
    void validate(const Function& function) const override;
 
    static blas::BLAS_Precision get_blas_precision(types::Scalar base_type);
 
    symbolic::MultiExpression get_out_shape();
 
    bool has_bias() const;

    passes::LibNodeExpander::ExpandOutcome
    expand(passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block) override;
 
    symbolic::SymbolSet symbols() const override;
 
    void replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) override;
 
    void replace(const symbolic::ExpressionMapping& replacements) override;
 
    bool supports_integer_types() const override {
        return false;
    }
 
    std::unique_ptr<data_flow::DataFlowNode>
    clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const override;
 
    std::string toStr() const override;

    symbolic::Expression num_output_elements() const;

    symbolic::Expression kernel_iteration_count() const;
 
    symbolic::Expression flop() const override;
 
    data_flow::PointerAccessType pointer_access_type(int input_idx) const override;
 
    struct ConvExpandPrerequisits {
        const data_flow::Memlet* iedge_X;
        const data_flow::Memlet* iedge_W;
        const data_flow::Memlet* iedge_B;
        const data_flow::Memlet* iedge_Y;
        const data_flow::AccessNode* access_X;
        const data_flow::AccessNode* access_W;
        const data_flow::AccessNode* access_B;
        const data_flow::AccessNode* access_Y;
        bool has_bias;
        structured_control_flow::Block* block;
        structured_control_flow::Sequence* block_parent;
        size_t block_index;
    };
 
    bool check_expandable(data_flow::DataFlowGraph& dfg, ConvExpandPrerequisits& boundary) const;
};

class ConvNodeSerializer : public SpatialTensorNodeBaseSerializer {
public:
    nlohmann::json serialize(const data_flow::LibraryNode& library_node) override;
 
    data_flow::LibraryNode& deserialize(
        const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
    ) override;
};
 
} // namespace tensor
} // namespace math
} // namespace sdfg