structured_loop.h

Go to the documentation of this file.

#pragma once
 
#include <optional>
#include <string>
#include <vector>
 
#include "sdfg/structured_control_flow/control_flow_node.h"
#include "sdfg/structured_control_flow/sequence.h"
#include "sdfg/symbolic/symbolic.h"
 
namespace sdfg {
 
namespace builder {
class StructuredSDFGBuilder;
}
namespace structured_control_flow {

enum class ScheduleTypeCategory { Offloader, Parallelizer, Vectorizer, None };

class ScheduleType {
private:
    std::unordered_map<std::string, std::string> properties_;
    std::string value_;
    ScheduleTypeCategory category_;
 
public:
    ScheduleType(std::string value, ScheduleTypeCategory category) : value_(value), category_(category) {
    }

    const std::string& value() const {
        return value_;
    }

    const std::unordered_map<std::string, std::string>& properties() const {
        return properties_;
    }

    ScheduleTypeCategory category() const {
        return category_;
    }

    void set_property(const std::string& key, const std::string& value) {
        if (properties_.find(key) == properties_.end()) {
            properties_.insert({key, value});
            return;
        }
        properties_.at(key) = value;
    }

    void remove_property(const std::string& key) {
        properties_.erase(key);
    }
 
    void operator=(const ScheduleType& rhs) {
        value_ = rhs.value_;
        properties_.clear();
        for (const auto& entry : rhs.properties_) {
            properties_.insert(entry);
        }
        category_ = rhs.category_;
    }
};

class ScheduleType_Sequential {
public:
    static const std::string value() {
        return "SEQUENTIAL";
    }
    static ScheduleType create() {
        return ScheduleType(value(), ScheduleTypeCategory::None);
    }
};

class ScheduleType_Unroll {
public:
    static void set(ScheduleType& schedule, bool full = true) {
        schedule.set_property("unroll", full ? "full" : "none");
    }

    static bool is_set(const ScheduleType& schedule) {
        auto it = schedule.properties().find("unroll");
        return it != schedule.properties().end() && it->second == "full";
    }
};

class ScheduleType_NoFusion {
public:
    static void set(ScheduleType& schedule, bool no_fusion = true) {
        schedule.set_property("fusion", no_fusion ? "none" : "default");
    }
 
    static bool is_set(const ScheduleType& schedule) {
        auto it = schedule.properties().find("fusion");
        return it != schedule.properties().end() && it->second == "none";
    }
};

class ScheduleType_ZeroInit {
public:
    static constexpr const char* KEY = "init";
 
    static void set(ScheduleType& schedule, bool zero_init = true) {
        schedule.set_property(KEY, zero_init ? "zero" : "none");
    }
 
    static bool is_set(const ScheduleType& schedule) {
        auto it = schedule.properties().find(KEY);
        return it != schedule.properties().end() && it->second == "zero";
    }
};

class ScheduleType_NestAnnotations {
public:
    static const std::vector<std::string>& keys() {
        static const std::vector<std::string> keys{"fusion", ScheduleType_ZeroInit::KEY};
        return keys;
    }

    static void copy(const ScheduleType& from, ScheduleType& to) {
        for (const auto& key : keys()) {
            auto it = from.properties().find(key);
            if (it != from.properties().end()) {
                to.set_property(key, it->second);
            }
        }
    }

    static void inherit(const ScheduleType& from, ScheduleType& to) {
        for (const auto& key : keys()) {
            auto it = from.properties().find(key);
            if (it != from.properties().end() && !to.properties().contains(key)) {
                to.set_property(key, it->second);
            }
        }
    }

    static void strip(ScheduleType& schedule) {
        for (const auto& key : keys()) {
            schedule.remove_property(key);
        }
    }
};
 

class StructuredLoop : public ControlFlowNode {
    friend class sdfg::builder::StructuredSDFGBuilder;
 
protected:
    symbolic::Symbol indvar_;
    symbolic::Expression init_;
    symbolic::Expression update_;
    symbolic::Condition condition_;
 
    std::unique_ptr<Sequence> root_;
 
    ScheduleType schedule_type_;
 
    // Memoized header analyses; every writer of indvar/init/update/condition must reset them.
    mutable std::optional<symbolic::Integer> stride_cache_;
    mutable std::optional<symbolic::Expression> canonical_bound_cache_;
    mutable std::optional<symbolic::Expression> num_iterations_cache_;
 
    void reset_cache();
 
    symbolic::Integer compute_stride() const;
 
    symbolic::Expression compute_canonical_bound() const;
 
    symbolic::Expression compute_num_iterations() const;
 
    StructuredLoop(
        size_t element_id,
        const DebugInfo& debug_info,
        ControlFlowNode* parent,
        symbolic::Symbol indvar,
        symbolic::Expression init,
        symbolic::Expression update,
        symbolic::Condition condition,
        const ScheduleType& schedule_type = ScheduleType_Sequential::create()
    );
 
    static constexpr size_t REQUIRED_ELEMENT_IDS = 2;
 
public:
    static constexpr ElementType TypeGroup = ElementType::For | ElementType::Map | ElementType::Reduce;

    static bool classof(const Element& element) {
        return is_a(element.type_id(), TypeGroup);
    }
 
    virtual ~StructuredLoop() = default;
 
    StructuredLoop(const StructuredLoop& node) = delete;
    StructuredLoop& operator=(const StructuredLoop&) = delete;
 
    void validate(const Function& function) const override;

    const symbolic::Symbol indvar() const;

    const symbolic::Expression init() const;

    const symbolic::Expression update() const;

    const symbolic::Condition condition() const;

    Sequence& root() const;

    const ScheduleType& schedule_type() const;

    void replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) override;
 
    void replace(const symbolic::ExpressionMapping& replacements) override;

    symbolic::Integer stride() const;

    bool is_contiguous() const;

    bool is_monotonic() const;

    symbolic::Expression canonical_bound() const;

    symbolic::Expression canonical_bound_upper() const;

    symbolic::Expression canonical_bound_lower() const;

    symbolic::Expression num_iterations() const;

    symbolic::Expression num_iterations_approx() const;

    bool is_loop_normal_form() const;
};
 
} // namespace structured_control_flow
} // namespace sdfg