26306537039

Committed 22 May 2026 07:00PM UTC coverage: 60.884% (+0.05%) from 60.837%

Build # 26306537039

Build Type

Pull #721

github

Committed by

web-flow

Commit Message

Merge 90ef1efb5 into b03adb735

Pull Request Pull Request #721: Make second level tiling possible

Coverage Stats

74 of 92 new or added lines in 2 files covered. (80.43%)

213 existing lines in 9 files now uncovered.

35101 of 57652 relevant lines covered (60.88%)

11087.7 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

84.08

/opt/src/transformations/einsum_extend.cpp

#include "sdfg/transformations/einsum_extend.h"

#include <cstddef>
#include <nlohmann/json_fwd.hpp>
#include <string>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include "sdfg/analysis/analysis.h"
#include "sdfg/builder/structured_sdfg_builder.h"
#include "sdfg/data_flow/access_node.h"
#include "sdfg/data_flow/data_flow_node.h"
#include "sdfg/data_flow/library_nodes/math/tensor/einsum_node.h"
#include "sdfg/data_flow/memlet.h"
#include "sdfg/data_flow/tasklet.h"
#include "sdfg/element.h"
#include "sdfg/structured_control_flow/block.h"
#include "sdfg/transformations/transformation.h"
#include "sdfg/types/type.h"

namespace sdfg {
namespace transformations {

EinsumExtend::EinsumExtend(math::tensor::EinsumNode& einsum_node)
    : einsum_node_(einsum_node), new_einsum_node_(nullptr) {}

std::string EinsumExtend::name() const { return "EinsumExtend"; }

bool EinsumExtend::can_be_applied(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {
    // Skip EinsumNodes with dimensions
    if (this->einsum_node_.dims().size() > 0) {
        return false;
    }

    size_t muls = 0;
    auto& dfg = this->einsum_node_.get_parent();
    for (auto& iedge : dfg.in_edges(this->einsum_node_)) {
        // Skip reduction container (connector "__einsum_out")
        if (iedge.dst_conn() == this->einsum_node_.inputs().back()) {
            continue;
        }

        // Skip constant nodes and access nodes without in edges
        auto& access_node = static_cast<data_flow::AccessNode&>(iedge.src());
        if (dynamic_cast<data_flow::ConstantNode*>(&access_node) || dfg.in_degree(access_node) == 0) {
            continue;
        }

        // Count the multiplication tasklets whose output access nodes are input access nodes of the EinsumNode
        for (auto& access_node_iedge : dfg.in_edges(access_node)) {
            auto* tasklet = dynamic_cast<data_flow::Tasklet*>(&access_node_iedge.src());
            if (tasklet && tasklet->code() == data_flow::TaskletCode::fp_mul) {
                muls++;
            }
        }
    }

    return muls > 0;
}

void EinsumExtend::apply(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {
    auto& dfg = this->einsum_node_.get_parent();
    auto* block = dynamic_cast<structured_control_flow::Block*>(dfg.get_parent());
    assert(block);

    // Construct inputs, in indices, and a map for new in edges
    std::vector<std::string> inputs;
    std::vector<data_flow::Subset> in_indices;
    std::unordered_map<std::string, std::tuple<data_flow::AccessNode&, const types::IType&, DebugInfo>> new_iedges_map;
    std::unordered_set<data_flow::Memlet*> memlets_for_removal;
    std::unordered_set<data_flow::DataFlowNode*> nodes_for_removal;
    DebugInfo new_deb_info(this->einsum_node_.debug_info());
    for (size_t i = 0; i < this->einsum_node_.inputs().size() - 1; i++) {
        auto& conn = this->einsum_node_.input(i);

        // Find corresponding in edge
        data_flow::Memlet* iedge = nullptr;
        for (auto& in_edge : dfg.in_edges(this->einsum_node_)) {
            if (in_edge.dst_conn() == conn) {
                iedge = &in_edge;
                break;
            }
        }
        assert(iedge);

        // Check if at the access node there is a multiplication tasklet
        auto& access_node = static_cast<data_flow::AccessNode&>(iedge->src());
        data_flow::Tasklet* tasklet = nullptr;
        data_flow::Memlet* tasklet_oedge = nullptr;
        if (!dynamic_cast<data_flow::ConstantNode*>(&access_node) && dfg.in_degree(access_node) > 0) {
            for (auto& access_node_iedge : dfg.in_edges(access_node)) {
                auto* tmp_taskelt = dynamic_cast<data_flow::Tasklet*>(&access_node_iedge.src());
                if (tmp_taskelt && tmp_taskelt->code() == data_flow::TaskletCode::fp_mul) {
                    tasklet = tmp_taskelt;
                    tasklet_oedge = &access_node_iedge;
                    break;
                }
            }
        }

        // Fill the data ...
        if (tasklet) {
            // ... with new access nodes and connectors
            for (auto& tasklet_iedge : dfg.in_edges(*tasklet)) {
                auto& tasklet_access_node = static_cast<data_flow::AccessNode&>(tasklet_iedge.src());
                std::string new_conn = iedge->dst_conn() + tasklet_iedge.dst_conn();
                inputs.push_back(new_conn);
                in_indices.push_back(tasklet_iedge.subset());
                new_iedges_map.insert(
                    {new_conn,
                     {tasklet_access_node,
                      tasklet_iedge.base_type(),
                      DebugInfo::merge(iedge->debug_info(), tasklet_iedge.debug_info())}}
                );
                memlets_for_removal.insert(&tasklet_iedge);
            }

            // Mark tasklet and its memlets for removal
            memlets_for_removal.insert(tasklet_oedge);
            new_deb_info = DebugInfo::merge(new_deb_info, tasklet->debug_info());
            nodes_for_removal.insert(tasklet);

            // Mark acess node for removal if not used elsewhere
            if (dfg.in_degree(access_node) == 1 && dfg.out_degree(access_node) == 1) {
                nodes_for_removal.insert(&access_node);
            }
        } else {
            // ... with the old stuff
            inputs.push_back(conn);
            in_indices.push_back(this->einsum_node_.in_indices(i));
            new_iedges_map.insert({conn, {access_node, iedge->base_type(), iedge->debug_info()}});
        }

        // Mark in edge for removal
        memlets_for_removal.insert(iedge);
    }

    // Special handling for the reduction input
    {
        auto& conn = this->einsum_node_.inputs().back();

        // Find corresponding in edge
        data_flow::Memlet* iedge = nullptr;
        for (auto& in_edge : dfg.in_edges(this->einsum_node_)) {
            if (in_edge.dst_conn() == conn) {
                iedge = &in_edge;
                break;
            }
        }
        assert(iedge);

        // Mapping and marking for removal
        auto& access_node = static_cast<data_flow::AccessNode&>(iedge->src());
        new_iedges_map.insert({conn, {access_node, iedge->base_type(), iedge->debug_info()}});
        memlets_for_removal.insert(iedge);
    }

    // Create new einsum node
    auto& new_libnode = builder.add_library_node<
        math::tensor::EinsumNode,
        const std::vector<std::string>&,
        const std::vector<math::tensor::EinsumDimension>&,
        const data_flow::Subset&,
        const std::vector<
            data_flow::Subset>&>(*block, new_deb_info, inputs, {}, this->einsum_node_.out_indices(), in_indices);
    this->new_einsum_node_ = static_cast<math::tensor::EinsumNode*>(&new_libnode);

    // Construct in edges
    for (auto& conn : this->new_einsum_node_->inputs()) {
        auto [access_node, type, deb_info] = new_iedges_map.at(conn);
        builder.add_memlet(*block, access_node, "void", new_libnode, conn, {}, type, deb_info);
    }

    // Remove marked memlets & nodes
    for (auto* memlet : memlets_for_removal) {
        builder.remove_memlet(*block, *memlet);
    }
    for (auto* node : nodes_for_removal) {
        builder.remove_node(*block, *node);
    }

    // Redirect out edges
    while (dfg.out_edges(this->einsum_node_).begin() != dfg.out_edges(this->einsum_node_).end()) {
        auto& oedge = *dfg.out_edges(this->einsum_node_).begin();
        builder.add_memlet(
            *block,
            new_libnode,
            oedge.src_conn(),
            oedge.dst(),
            oedge.dst_conn(),
            oedge.subset(),
            oedge.base_type(),
            oedge.debug_info()
        );
        builder.remove_memlet(*block, oedge);
    }

    // Remove old einsum node
    builder.remove_node(*block, this->einsum_node_);

    analysis_manager.invalidate_all();
}

math::tensor::EinsumNode* EinsumExtend::new_einsum_node() { return this->new_einsum_node_; }

void EinsumExtend::to_json(nlohmann::json& j) const {
    j["transformation_type"] = this->name();
    j["einsum_node_element_id"] = this->einsum_node_.element_id();
}

EinsumExtend EinsumExtend::from_json(builder::StructuredSDFGBuilder& builder, const nlohmann::json& j) {
    assert(j.contains("einsum_node_element_id"));
    assert(j["einsum_node_element_id"].is_number_unsigned());
    size_t einsum_node_id = j["einsum_node_element_id"].get<size_t>();
    auto* einsum_node_element = builder.find_element_by_id(einsum_node_id);
    if (!einsum_node_element) {
        throw InvalidTransformationDescriptionException(
            "Element with ID " + std::to_string(einsum_node_id) + " not found"
        );
    }
    auto* einsum_node = dynamic_cast<math::tensor::EinsumNode*>(einsum_node_element);
    if (!einsum_node) {
        throw InvalidTransformationDescriptionException(
            "Element with ID " + std::to_string(einsum_node_id) + " is not an EinsumNode"
        );
    }

    return EinsumExtend(*einsum_node);
}

} // namespace transformations
} // namespace sdfg

1	#include "sdfg/transformations/einsum_extend.h"
2
3	#include <cstddef>
4	#include <nlohmann/json_fwd.hpp>
5	#include <string>
6	#include <tuple>
7	#include <unordered_map>
8	#include <unordered_set>
9	#include <vector>
10
11	#include "sdfg/analysis/analysis.h"
12	#include "sdfg/builder/structured_sdfg_builder.h"
13	#include "sdfg/data_flow/access_node.h"
14	#include "sdfg/data_flow/data_flow_node.h"
15	#include "sdfg/data_flow/library_nodes/math/tensor/einsum_node.h"
16	#include "sdfg/data_flow/memlet.h"
17	#include "sdfg/data_flow/tasklet.h"
18	#include "sdfg/element.h"
19	#include "sdfg/structured_control_flow/block.h"
20	#include "sdfg/transformations/transformation.h"
21	#include "sdfg/types/type.h"
22
23	namespace sdfg {
24	namespace transformations {
25
26	EinsumExtend::EinsumExtend(math::tensor::EinsumNode& einsum_node)
27	: einsum_node_(einsum_node), new_einsum_node_(nullptr) {}	12✔
28
29	std::string EinsumExtend::name() const { return "EinsumExtend"; }	2✔
30
31	bool EinsumExtend::can_be_applied(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {	12✔
32	// Skip EinsumNodes with dimensions
33	if (this->einsum_node_.dims().size() > 0) {	12✔
34	return false;	×
UNCOV 35	}	×
36
37	size_t muls = 0;	12✔
38	auto& dfg = this->einsum_node_.get_parent();	12✔
39	for (auto& iedge : dfg.in_edges(this->einsum_node_)) {	33✔
40	// Skip reduction container (connector "__einsum_out")
41	if (iedge.dst_conn() == this->einsum_node_.inputs().back()) {	33✔
42	continue;	12✔
43	}	12✔
44
45	// Skip constant nodes and access nodes without in edges
46	auto& access_node = static_cast<data_flow::AccessNode&>(iedge.src());	21✔
47	if (dynamic_cast<data_flow::ConstantNode*>(&access_node) \|\| dfg.in_degree(access_node) == 0) {	21✔
48	continue;	14✔
49	}	14✔
50
51	// Count the multiplication tasklets whose output access nodes are input access nodes of the EinsumNode
52	for (auto& access_node_iedge : dfg.in_edges(access_node)) {	7✔
53	auto* tasklet = dynamic_cast<data_flow::Tasklet*>(&access_node_iedge.src());	7✔
54	if (tasklet && tasklet->code() == data_flow::TaskletCode::fp_mul) {	7✔
55	muls++;	7✔
56	}	7✔
57	}	7✔
58	}	7✔
59
60	return muls > 0;	12✔
61	}	12✔
62
63	void EinsumExtend::apply(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {	6✔
64	auto& dfg = this->einsum_node_.get_parent();	6✔
65	auto* block = dynamic_cast<structured_control_flow::Block*>(dfg.get_parent());	6✔
66	assert(block);	6✔
67
68	// Construct inputs, in indices, and a map for new in edges
69	std::vector<std::string> inputs;	6✔
70	std::vector<data_flow::Subset> in_indices;	6✔
71	std::unordered_map<std::string, std::tuple<data_flow::AccessNode&, const types::IType&, DebugInfo>> new_iedges_map;	6✔
72	std::unordered_set<data_flow::Memlet*> memlets_for_removal;	6✔
73	std::unordered_set<data_flow::DataFlowNode*> nodes_for_removal;	6✔
74	DebugInfo new_deb_info(this->einsum_node_.debug_info());	6✔
75	for (size_t i = 0; i < this->einsum_node_.inputs().size() - 1; i++) {	17✔
76	auto& conn = this->einsum_node_.input(i);	11✔
77
78	// Find corresponding in edge
79	data_flow::Memlet* iedge = nullptr;	11✔
80	for (auto& in_edge : dfg.in_edges(this->einsum_node_)) {	17✔
81	if (in_edge.dst_conn() == conn) {	17✔
82	iedge = &in_edge;	11✔
83	break;	11✔
84	}	11✔
85	}	17✔
86	assert(iedge);	11✔
87
88	// Check if at the access node there is a multiplication tasklet
89	auto& access_node = static_cast<data_flow::AccessNode&>(iedge->src());	11✔
90	data_flow::Tasklet* tasklet = nullptr;	11✔
91	data_flow::Memlet* tasklet_oedge = nullptr;	11✔
92	if (!dynamic_cast<data_flow::ConstantNode*>(&access_node) && dfg.in_degree(access_node) > 0) {	11✔
93	for (auto& access_node_iedge : dfg.in_edges(access_node)) {	7✔
94	auto* tmp_taskelt = dynamic_cast<data_flow::Tasklet*>(&access_node_iedge.src());	7✔
95	if (tmp_taskelt && tmp_taskelt->code() == data_flow::TaskletCode::fp_mul) {	7✔
96	tasklet = tmp_taskelt;	7✔
97	tasklet_oedge = &access_node_iedge;	7✔
98	break;	7✔
99	}	7✔
100	}	7✔
101	}	7✔
102
103	// Fill the data ...
104	if (tasklet) {	11✔
105	// ... with new access nodes and connectors
106	for (auto& tasklet_iedge : dfg.in_edges(*tasklet)) {	14✔
107	auto& tasklet_access_node = static_cast<data_flow::AccessNode&>(tasklet_iedge.src());	14✔
108	std::string new_conn = iedge->dst_conn() + tasklet_iedge.dst_conn();	14✔
109	inputs.push_back(new_conn);	14✔
110	in_indices.push_back(tasklet_iedge.subset());	14✔
111	new_iedges_map.insert(	14✔
112	{new_conn,	14✔
113	{tasklet_access_node,	14✔
114	tasklet_iedge.base_type(),	14✔
115	DebugInfo::merge(iedge->debug_info(), tasklet_iedge.debug_info())}}	14✔
116	);	14✔
117	memlets_for_removal.insert(&tasklet_iedge);	14✔
118	}	14✔
119
120	// Mark tasklet and its memlets for removal
121	memlets_for_removal.insert(tasklet_oedge);	7✔
122	new_deb_info = DebugInfo::merge(new_deb_info, tasklet->debug_info());	7✔
123	nodes_for_removal.insert(tasklet);	7✔
124
125	// Mark acess node for removal if not used elsewhere
126	if (dfg.in_degree(access_node) == 1 && dfg.out_degree(access_node) == 1) {	7✔
127	nodes_for_removal.insert(&access_node);	7✔
128	}	7✔
129	} else {	7✔
130	// ... with the old stuff
131	inputs.push_back(conn);	4✔
132	in_indices.push_back(this->einsum_node_.in_indices(i));	4✔
133	new_iedges_map.insert({conn, {access_node, iedge->base_type(), iedge->debug_info()}});	4✔
134	}	4✔
135
136	// Mark in edge for removal
137	memlets_for_removal.insert(iedge);	11✔
138	}	11✔
139
140	// Special handling for the reduction input
141	{	6✔
142	auto& conn = this->einsum_node_.inputs().back();	6✔
143
144	// Find corresponding in edge
145	data_flow::Memlet* iedge = nullptr;	6✔
146	for (auto& in_edge : dfg.in_edges(this->einsum_node_)) {	17✔
147	if (in_edge.dst_conn() == conn) {	17✔
148	iedge = &in_edge;	6✔
149	break;	6✔
150	}	6✔
151	}	17✔
152	assert(iedge);	6✔
153
154	// Mapping and marking for removal
155	auto& access_node = static_cast<data_flow::AccessNode&>(iedge->src());	6✔
156	new_iedges_map.insert({conn, {access_node, iedge->base_type(), iedge->debug_info()}});	6✔
157	memlets_for_removal.insert(iedge);	6✔
158	}	6✔
159
160	// Create new einsum node
UNCOV 161	auto& new_libnode = builder.add_library_node<	×
162	math::tensor::EinsumNode,	6✔
163	const std::vector<std::string>&,	6✔
164	const std::vector<math::tensor::EinsumDimension>&,	6✔
165	const data_flow::Subset&,	6✔
166	const std::vector<	6✔
167	data_flow::Subset>&>(*block, new_deb_info, inputs, {}, this->einsum_node_.out_indices(), in_indices);	6✔
168	this->new_einsum_node_ = static_cast<math::tensor::EinsumNode*>(&new_libnode);	6✔
169
170	// Construct in edges
171	for (auto& conn : this->new_einsum_node_->inputs()) {	24✔
172	auto [access_node, type, deb_info] = new_iedges_map.at(conn);	24✔
173	builder.add_memlet(*block, access_node, "void", new_libnode, conn, {}, type, deb_info);	24✔
174	}	24✔
175
176	// Remove marked memlets & nodes
177	for (auto* memlet : memlets_for_removal) {	38✔
178	builder.remove_memlet(block, memlet);	38✔
179	}	38✔
180	for (auto* node : nodes_for_removal) {	14✔
181	builder.remove_node(block, node);	14✔
182	}	14✔
183
184	// Redirect out edges
185	while (dfg.out_edges(this->einsum_node_).begin() != dfg.out_edges(this->einsum_node_).end()) {	12✔
186	auto& oedge = *dfg.out_edges(this->einsum_node_).begin();	6✔
187	builder.add_memlet(	6✔
188	*block,	6✔
189	new_libnode,	6✔
190	oedge.src_conn(),	6✔
191	oedge.dst(),	6✔
192	oedge.dst_conn(),	6✔
193	oedge.subset(),	6✔
194	oedge.base_type(),	6✔
195	oedge.debug_info()	6✔
196	);	6✔
197	builder.remove_memlet(*block, oedge);	6✔
198	}	6✔
199
200	// Remove old einsum node
201	builder.remove_node(*block, this->einsum_node_);	6✔
202
203	analysis_manager.invalidate_all();	6✔
204	}	6✔
205
206	math::tensor::EinsumNode* EinsumExtend::new_einsum_node() { return this->new_einsum_node_; }	2✔
207
208	void EinsumExtend::to_json(nlohmann::json& j) const {	×
UNCOV 209	j["transformation_type"] = this->name();	×
210	j["einsum_node_element_id"] = this->einsum_node_.element_id();	×
211	}	×
212
213	EinsumExtend EinsumExtend::from_json(builder::StructuredSDFGBuilder& builder, const nlohmann::json& j) {	×
214	assert(j.contains("einsum_node_element_id"));	×
215	assert(j["einsum_node_element_id"].is_number_unsigned());	×
216	size_t einsum_node_id = j["einsum_node_element_id"].get<size_t>();	×
217	auto* einsum_node_element = builder.find_element_by_id(einsum_node_id);	×
218	if (!einsum_node_element) {	×
219	throw InvalidTransformationDescriptionException(	×
220	"Element with ID " + std::to_string(einsum_node_id) + " not found"	×
221	);	×
222	}	×
223	auto* einsum_node = dynamic_cast<math::tensor::EinsumNode*>(einsum_node_element);	×
224	if (!einsum_node) {	×
225	throw InvalidTransformationDescriptionException(	×
UNCOV 226	"Element with ID " + std::to_string(einsum_node_id) + " is not an EinsumNode"	×
227	);	×
228	}	×
229
UNCOV 230	return EinsumExtend(*einsum_node);	×
UNCOV 231	}	×
232
233	} // namespace transformations
234	} // namespace sdfg

daisytuner / docc / 26306537039

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous