lib_node_expander.h

Go to the documentation of this file.

#pragma once
#include "sdfg/builder/structured_sdfg_builder.h"
#include "sdfg/options.h"
#include "sdfg/structured_control_flow/block.h"
 
namespace sdfg::passes {
 
class LibNodeExpander {
public:
    enum class InputUse { Skip = 0, Scalar, IndirectRead, IndirectWrite, IndirectReadWrite };
 
    struct ExpandOutcome {
        bool applied;
        explicit ExpandOutcome(bool applied) : applied(applied) {
        }
    };
    virtual ~LibNodeExpander() = default;

    virtual const LibNodeExpander* for_lib_node(const data_flow::LibraryNode& node) const {
        return this;
    }

    virtual const data_flow::LibraryNodeCode* matched_code() const {
        return nullptr;
    }
 
    class AccessNodeExpand {
    public:
        virtual ~AccessNodeExpand() = default;
 
        virtual builder::StructuredSDFGBuilder& builder() = 0;

        virtual structured_control_flow::Sequence& replace_with_sequence() = 0;
 
        enum class LoopType { For, Map };
 
        virtual structured_control_flow::StructuredLoop& replace_with_structured_loop(
            LoopType type,
            const symbolic::Symbol indvar,
            const symbolic::Condition condition,
            const symbolic::Expression init,
            const symbolic::Expression update,
            const ScheduleType& schedule_type
        ) = 0;

        virtual data_flow::AccessNode& add_scalar_input_access(structured_control_flow::Block& block, size_t input_idx) = 0;
        virtual data_flow::AccessNode&
        add_indirect_read_access(structured_control_flow::Block& block, size_t input_idx) = 0;
        virtual data_flow::AccessNode&
        add_indirect_write_access(structured_control_flow::Block& block, size_t input_idx) = 0;
        virtual data_flow::AccessNode& add_output_access(structured_control_flow::Block& block, size_t output_idx) = 0;

        virtual ExpandOutcome successfully_expanded() = 0;
    };
 
    class ExpandContext {
    public:
        virtual ~ExpandContext() = default;

        virtual std::unique_ptr<AccessNodeExpand>
        replacement_requires_access_nodes(const std::vector<InputUse>& access_dirs) = 0;
 
        // TODO virtual std::unique_ptr<DataflowExpand> replace_dataflow() = 0;
        // this would be a quicker way, when the expansion only consists of pure dataflow, where base expansion
        // infrastructure needs to migrate edges from the old node to the new nodes. Describing the new targets of edges
        // can be taken from `ElementWiseDataflowTensorNode.expand_operation_dataflow()`

        virtual ExpandOutcome successfully_modified_node_only() = 0;
 
        virtual ExpandOutcome unable() = 0;
 
        virtual ExpandOutcome unapplicable() = 0;

        virtual const Options& options() const = 0;
    };
 
    virtual ExpandOutcome handle_expand(
        ExpandContext& context, structured_control_flow::Block& block, data_flow::LibraryNode& node
    ) const = 0;
};
 
template<typename T>
class TypedLibNodeExpander : public LibNodeExpander {
public:
    using LibNodeExpander::handle_expand;
 
    const LibNodeExpander* for_lib_node(const data_flow::LibraryNode& node) const override {
        if (auto* cast = dynamic_cast<const T*>(&node)) {
            return this;
        }
        return nullptr;
    }
 
    virtual LibNodeExpander::ExpandOutcome
    handle_expand(LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block, T& node) const = 0;
 
    LibNodeExpander::ExpandOutcome handle_expand(
        LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block, data_flow::LibraryNode& node
    ) const override {
        auto& typed_node = static_cast<T&>(node);
        return handle_expand(context, block, typed_node);
    }
};
 
template<typename T>
class CodeLibNodeExpander : public TypedLibNodeExpander<T> {
    const data_flow::LibraryNodeCode code_;
 
public:
    CodeLibNodeExpander(const data_flow::LibraryNodeCode& code) : code_(code) {
    }

    const data_flow::LibraryNodeCode& code() const {
        return code_;
    }
 
    const data_flow::LibraryNodeCode* matched_code() const override {
        return &code_;
    }
 
    const LibNodeExpander* for_lib_node(const data_flow::LibraryNode& node) const override {
        if (node.code() == code_) {
            return this;
        }
        return nullptr;
    }
};
 
} // namespace sdfg::passes