einsum_replacer.h

Go to the documentation of this file.

#pragma once
 
#include <memory>
#include <string>
#include <vector>
 
#include <sdfg/builder/structured_sdfg_builder.h>
#include <sdfg/data_flow/access_node.h>
#include <sdfg/options.h>
#include <sdfg/passes/expansion/lib_node_expander.h>
#include <sdfg/structured_control_flow/control_flow_node.h>
#include <sdfg/structured_control_flow/sequence.h>
 
#include "sdfg/einsum/einsum_detection.h"
 
namespace sdfg::einsum {
 
typedef passes::LibNodeExpander::ExpandOutcome ReplaceOutcome;

class EinsumReplacementContext {
public:
    virtual ~EinsumReplacementContext() = default;
 
    virtual std::unique_ptr<passes::LibNodeExpander::AccessNodeExpand> replacement_requires_access_nodes(
        const std::vector<passes::LibNodeExpander::InputUse>& access_dirs, bool leave_unconsumed_indices
    ) = 0;
 
    virtual ReplaceOutcome successfully_modified_node_only() = 0;
    virtual ReplaceOutcome unable() = 0;
    virtual ReplaceOutcome unapplicable() = 0;
};
 

class EinsumReplacer {
public:
    virtual ~EinsumReplacer() = default;
 
    virtual std::string name() const = 0;
 
    virtual ReplaceOutcome replace(EinsumReplacementContext& context, const EinsumCluster& cluster) const = 0;
};

ReplaceOutcome replace_einsum_cluster(
    builder::StructuredSDFGBuilder& builder, const EinsumCluster& cluster, const EinsumReplacer& replacer
);
 
} // namespace sdfg::einsum