elementwise_node.h

Go to the documentation of this file.

 
#pragma once
 
#include "sdfg/data_flow/library_nodes/math/tensor/tensor_node.h"
 
#include "sdfg/codegen/dispatchers/block_dispatcher.h"
#include "sdfg/serializer/json_serializer.h"
#include "sdfg/structured_control_flow/block.h"
 
namespace sdfg {
 
namespace math {
namespace tensor {

class ElementWiseDataflowTensorNode : public TensorNode {
protected:
    QuantizationType fixed_quantization_;
    std::vector<symbolic::Expression> shape_; 
 
public:
    struct ElementOutput {
        CodeNode* producer = nullptr;
        int output_conn_index = -1;
        types::PrimitiveType type = types::Void;
    };
    struct ElementInput {
        CodeNode* consumer = nullptr;
        int input_conn_index = -1;
        types::PrimitiveType required_type = types::Void;
    };
 
    ElementWiseDataflowTensorNode(
        size_t element_id,
        const DebugInfo& debug_info,
        const graph::Vertex vertex,
        data_flow::DataFlowGraph& parent,
        const data_flow::LibraryNodeCode& code,
        const std::vector<symbolic::Expression>& shape,
        const std::string& modified_tensor_conn,
        const std::vector<std::string>& tensor_inputs,
        QuantizationType quantization = QUANTIZATION_MATCH_INPUTS,
        const data_flow::ImplementationType& impl_type = data_flow::ImplementationType_NONE
    );
 
    static std::vector<std::string>
    build_input_conns(const std::string& modified_tensor_conn, const std::vector<std::string>& inputs);

    virtual int tensor_input_count() const {
        return inputs_.size();
    }

    virtual int mandatory_input_count() const {
        return inputs_.size();
    }
 
    QuantizationType quantization() const {
        return quantization(get_parent());
    }

    QuantizationType quantization(const data_flow::DataFlowGraph& dataflow) const;

    std::optional<QuantizationType> uniform_quantization(const data_flow::DataFlowGraph& dataflow) const;

    QuantizationType fixed_quantization() const;
 
    void set_fixed_quantization(const QuantizationType quant);
 
    void validate_target_tensor(const data_flow::DataFlowGraph& graph) const;
 
    void validate_all_input_tensors(const data_flow::DataFlowGraph& graph) const;

    virtual void validate_non_tensor_inputs(const data_flow::DataFlowGraph& graph) const;
 
    void validate(const Function& function) const override;

    const std::vector<symbolic::Expression>& shape() const {
        return shape_;
    }
 
    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;
 
    static std::pair<structured_control_flow::Sequence*, std::vector<symbolic::Expression>> add_eltwise_scope(
        builder::StructuredSDFGBuilder& builder,
        const DebugInfo& scope_deb_info,
        Sequence& parent,
        const std::vector<symbolic::Expression>& shape
    );
 
    static std::unique_ptr<types::IType> access_type(const std::pair<types::PrimitiveType, const TensorLayout*>& pair);
 
    static bool create_input(
        builder::StructuredSDFGBuilder& builder,
        structured_control_flow::Block& block,
        const data_flow::AccessNode& org_src,
        const std::pair<types::PrimitiveType, const TensorLayout*>& src_type,
        const ElementInput& needed_input,
        const std::vector<symbolic::Expression>& eltwise_subset,
        std::unordered_map<const data_flow::AccessNode*, data_flow::AccessNode*>& new_node_mapping
    );
 
    static void create_output(
        builder::StructuredSDFGBuilder& builder,
        structured_control_flow::Block& block,
        const data_flow::AccessNode& org_dst,
        const types::Tensor& dst_type,
        const ElementOutput& provided_output,
        const std::vector<symbolic::Expression>& eltwise_subset
    );

    passes::LibNodeExpander::ExpandOutcome
    expand(passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block) override;
 
    data_flow::PointerAccessType pointer_access_type(int input_idx) const override;
 
protected:
    virtual ElementOutput expand_operation_dataflow(
        builder::StructuredSDFGBuilder& builder,
        structured_control_flow::Block& block,
        std::vector<ElementInput>& needed_inputs,
        types::PrimitiveType expected_type
    ) = 0;
 
    data_flow::AccessNode& create_tmp_access_node(
        builder::StructuredSDFGBuilder& builder,
        structured_control_flow::Block& block,
        const std::string& prefix,
        const types::IType& type
    ) const;
};
 
class BaseElementWiseDataflowTensorNodeSerializer : public serializer::LibraryNodeSerializer {
public:
    nlohmann::json serialize(const data_flow::LibraryNode& library_node) override;
 
    struct BaseDeser {
        std::vector<symbolic::Expression> shape;
        QuantizationType quantization;
        DebugInfo debug_info;
    };
 
    BaseDeser deserialize_base_values(const nlohmann::json& j);
};
 
template<typename T>
class SimpleElementWiseDataflowTensorNodeSerializer : public BaseElementWiseDataflowTensorNodeSerializer {
public:
    data_flow::LibraryNode& deserialize(
        const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
    ) override {
        // Assertions for required fields
        auto base = deserialize_base_values(j);
 
        return static_cast<ElementWiseDataflowTensorNode&>(builder.add_library_node<
                                                           T>(parent, base.debug_info, base.shape, base.quantization));
    }
};
 
} // namespace tensor
} // namespace math
} // namespace sdfg