einsum2matmul.h

Go to the documentation of this file.

#pragma once
 
#include <optional>
#include <string>
#include <vector>
 
#include <sdfg/passes/expansion/lib_node_expander.h>
 
#include "sdfg/einsum/einsum_detection.h"
#include "sdfg/einsum/einsum_replacer.h"
#include "sdfg/types/type.h"
 
namespace sdfg::einsum {

class Einsum2MatMul : public EinsumReplacer {
public:
    std::string name() const override {
        return "Einsum2MatMul";
    }
 
    math::tensor::TensorLayout build_tensor_layout(
        const EinsumIndexing& indexing,
        bool swapped,
        const symbolic::Expression& outer_dim,
        const symbolic::Expression& inner_dim,
        const symbolic::Expression& outer_offset,
        const symbolic::Expression& inner_offset,
        const symbolic::Expression& external_offset
    ) const;
 
    ReplaceOutcome replace(EinsumReplacementContext& context, const EinsumCluster& cluster) const override;

    virtual bool can_be_applied(const EinsumCluster& cluster) const;

    struct MatMulAnalysis {
        symbolic::Symbol indvar_outer_1 = SymEngine::null;
        symbolic::Symbol indvar_outer_2 = SymEngine::null;
        symbolic::Symbol indvar_inner = SymEngine::null;
        symbolic::Expression m = SymEngine::null;
        symbolic::Expression n = SymEngine::null;
        symbolic::Expression k = SymEngine::null;
        symbolic::Expression m_init = SymEngine::null;
        symbolic::Expression n_init = SymEngine::null;
        symbolic::Expression k_init = SymEngine::null;
        int a_idx = -1;
        int b_idx = -1;
        std::optional<math::tensor::TensorLayout> layout_a;
        std::optional<math::tensor::TensorLayout> layout_b;
        std::optional<math::tensor::TensorLayout> layout_y;
        types::PrimitiveType input_type = types::PrimitiveType::Void;
        types::PrimitiveType output_type = types::PrimitiveType::Void;
    };

    bool analyze(const EinsumCluster& cluster, MatMulAnalysis& result) const;
};
 
} // namespace sdfg::einsum