einsum_detection.h

Go to the documentation of this file.

#pragma once
 
#include <cstdint>
#include <list>
#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
 
#include <sdfg/data_flow/access_node.h>
#include <sdfg/data_flow/data_flow_node.h>
#include <sdfg/data_flow/memlet.h>
#include <sdfg/data_flow/tasklet.h>
#include <sdfg/einsum/einsum_node.h>
#include <sdfg/structured_control_flow/block.h>
#include <sdfg/structured_control_flow/structured_loop.h>
#include <sdfg/symbolic/conjunctive_normal_form.h>
#include <sdfg/symbolic/symbolic.h>
#include <sdfg/visitor/structured_sdfg_visitor.h>
 
namespace sdfg::einsum {

struct EinsumIndexContribution {
    symbolic::Symbol indvar;
    size_t index = 0;
    int64_t factor = 1;
 
    symbolic::Expression factor_expr() const;
};

struct EinsumIndexing {
    std::vector<EinsumIndexContribution> contributions;
    size_t covered = 0;

    std::vector<symbolic::Expression> get_strides(
        const std::vector<EinsumIndexContribution>& ordered_contributions,
        const std::vector<symbolic::Expression>& shape
    ) const;
};

struct OperandContribution {
    bool contributes = false;
    size_t index = 0;
    int64_t factor = 1;
};

struct PromotionCheck {
    bool ok = false;
    OperandContribution output;
    std::vector<OperandContribution> inputs;
};
 
class EinsumEdge {
public:
    virtual ~EinsumEdge() = default;
 
    virtual const data_flow::Subset& subset() const = 0;
    virtual const types::IType& base_type() const = 0;
    virtual data_flow::Memlet* original_edge() const = 0;
};
 
class EinsumEdgeOriginal : public EinsumEdge {
    data_flow::Memlet* edge;
 
public:
    EinsumEdgeOriginal(data_flow::Memlet* edge) : edge(edge) {
    }
 
    const data_flow::Subset& subset() const override {
        return edge->subset();
    }
    const types::IType& base_type() const override {
        return edge->base_type();
    }
    data_flow::Memlet* original_edge() const override {
        return edge;
    }
};
 
class EinsumEdgeVirtual : public EinsumEdge {
    data_flow::Subset subset_;
    const data_flow::Memlet* edge;
 
public:
    EinsumEdgeVirtual(const data_flow::Memlet* edge, const data_flow::Subset& subset) : subset_(subset), edge(edge) {
    }
 
    const data_flow::Subset& subset() const override {
        return subset_;
    }
    const types::IType& base_type() const override {
        return edge->base_type();
    }
    data_flow::Memlet* original_edge() const override {
        return nullptr;
    }
};

class EinsumCluster {
public:
    structured_control_flow::Block* block = nullptr;
 
    // Core einsum properties (mirroring EinsumNode, but excluding the implicit
    // "__einsum_out" reduction input, which is represented by output_node/out_indices).
    std::vector<EinsumDimension> dims;
    std::vector<std::unique_ptr<EinsumEdge>> in_edges;
    std::vector<std::string> inputs;

    std::vector<EinsumIndexing> einsum_in_indices;
    EinsumIndexing einsum_out_indices;

    std::unique_ptr<EinsumEdge> output_edge = nullptr;
    data_flow::AccessNode* output_node = nullptr;
    std::unordered_map<std::string, data_flow::DataFlowNode*> input_nodes;
    bool subtraction = false;

    std::list<data_flow::DataFlowNode*> consumed_nodes;
    std::list<structured_control_flow::StructuredLoop*> consumed_loops;
 
    std::pair<data_flow::AccessNode*, const data_flow::Subset&> get_input_for(size_t input_idx);
 
    const std::vector<EinsumIndexing>& get_input_indexings() const;
 
    std::string toStr() const;
 
    const data_flow::Subset& get_input_subset(size_t input_idx) const;
 
    const EinsumIndexing& get_output_indexing() const;
 
    symbolic::Expression get_linearized_outer_offset(int input_idx) const;
 
    const data_flow::Subset& out_indices() const;
};

class EinsumDetection : public visitor::ActualStructuredSDFGVisitor {
protected:
    struct BlockState {
        std::unordered_set<data_flow::DataFlowNode*> consumed_nodes;
        std::list<EinsumCluster*> clusters;
    };
    struct LoopScopeState {
        bool non_reducible = false;
        EinsumCluster* representing_last_block_ = nullptr;
    };
 

    std::vector<std::unique_ptr<EinsumCluster>> all_einsums_;
    std::vector<LoopScopeState> loop_stack_;
 
    static bool subsets_eq(const data_flow::Subset& subset1, const data_flow::Subset& subset2);
 
    EinsumCluster& new_cluster(structured_control_flow::Block& block);

    EinsumCluster* try_lift(structured_control_flow::Block& block, data_flow::Tasklet& tasklet);

    EinsumCluster* from_einsum_node(structured_control_flow::Block& block, EinsumNode& einsum_node);

    static bool try_extend(BlockState& state, EinsumCluster& cluster);
 
    static symbolic::Expression cnf_to_upper_bound(const symbolic::CNF& cnf, const symbolic::Symbol indvar);
    static bool subset_contains_symbol(const data_flow::Subset& subset, const symbolic::Symbol& symbol);

    PromotionCheck can_promote(EinsumCluster& cluster, structured_control_flow::StructuredLoop& loop);
    void apply_promote(EinsumCluster& cluster, structured_control_flow::StructuredLoop& loop, const PromotionCheck& check);

    void recompute_einsum_indices(EinsumCluster& cluster);
 
    bool block_fully_consumed(BlockState& state, structured_control_flow::Block& block);
    void register_cluster_in_scope(EinsumCluster* cluster);
 
public:
    EinsumDetection() = default;

    void find_einsum_core_ops(BlockState& state, structured_control_flow::Block& block);

    void consume_input_operations(BlockState& state);

    bool try_to_consume_loop(EinsumCluster& einsum, structured_control_flow::StructuredLoop& loop);
 
    void filter_for_coverage();

    size_t run(structured_control_flow::ControlFlowNode& start);

    const std::vector<std::unique_ptr<EinsumCluster>>& einsums() const {
        return all_einsums_;
    }
 
    using visitor::ActualStructuredSDFGVisitor::visit;
 
    bool visit(structured_control_flow::Block& node) override;
    bool visit(structured_control_flow::AssignmentBlock& node) override;
    bool visit(structured_control_flow::IfElse& node) override;
    bool visit(structured_control_flow::While& node) override;
    bool visit(structured_control_flow::Return& node) override;
    bool handleStructuredLoop(structured_control_flow::StructuredLoop& loop) override;
};
 
} // namespace sdfg::einsum