23797200186

Committed 31 Mar 2026 12:26PM UTC coverage: 64.476% (-0.009%) from 64.485%

Build # 23797200186

Build Type

Pull #605

github

Committed by

web-flow

Commit Message

Merge 95c5d0639 into faa465ea4

Pull Request Pull Request #605: Move Tenstorrent support

Coverage Stats

1 of 8 new or added lines in 4 files covered. (12.5%)

28612 of 44376 relevant lines covered (64.48%)

388.5 hits per line

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

0.0

/opt/src/transformations/einsum2dot.cpp

#include "sdfg/transformations/einsum2dot.h"

#include <cstddef>
#include <nlohmann/json_fwd.hpp>
#include <string>

#include "sdfg/analysis/analysis.h"
#include "sdfg/builder/structured_sdfg_builder.h"
#include "sdfg/data_flow/access_node.h"
#include "sdfg/data_flow/library_node.h"
#include "sdfg/data_flow/library_nodes/math/blas/dot_node.h"
#include "sdfg/data_flow/library_nodes/math/math.h"
#include "sdfg/einsum/einsum.h"
#include "sdfg/optimization_report/pass_report_consumer.h"
#include "sdfg/structured_control_flow/block.h"
#include "sdfg/symbolic/symbolic.h"
#include "sdfg/targets/cuda/cuda.h"
#include "sdfg/transformations/transformation.h"
#include "sdfg/types/type.h"
#include "sdfg/types/utils.h"

namespace sdfg {
namespace transformations {

Einsum2Dot::Einsum2Dot(einsum::EinsumNode& einsum_node, const std::string& target_tune)
    : einsum_node_(einsum_node), target_tune_(target_tune) {}

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

std::optional<sdfg::data_flow::ImplementationType> Einsum2Dot::get_impl_type(types::PrimitiveType data_type) {
    std::optional<sdfg::data_flow::ImplementationType> impl_type; // TODO make generic for any target
    if (target_tune_ == "sequential") {
        impl_type = sdfg::data_flow::ImplementationType_NONE;
    } else if (target_tune_ == "openmp") {
        impl_type = sdfg::math::blas::ImplementationType_BLAS;
    } else if (target_tune_ == "cuda") {
        impl_type = sdfg::cuda::blas::ImplementationType_CUBLASWithTransfers;
    } else if (target_tune_ == "tenstorrent") {
        if (data_type == types::PrimitiveType::Float) {
            impl_type = data_flow::ImplementationType{"TENSTORRENT_WithTransfers"};
        }
    }

    if (impl_type) {
        return impl_type;
    } else {
        return std::nullopt;
    }
}

bool Einsum2Dot::can_be_applied(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {
    // Check dims
    if (this->einsum_node_.dims().size() != 1 || !symbolic::eq(this->einsum_node_.init(0), symbolic::zero())) {
        return false;
    }
    symbolic::Symbol indvar = this->einsum_node_.indvar(0);

    // Check out indices
    if (this->einsum_node_.out_indices().size() != 0) {
        return false;
    }

    // Check inputs
    if (this->einsum_node_.inputs().size() != 3 || this->einsum_node_.input(2) != this->einsum_node_.output(0)) {
        return false;
    }

    // Check in indices
    if (this->einsum_node_.in_indices(0).size() != 1 || !symbolic::eq(this->einsum_node_.in_index(0, 0), indvar)) {
        return false;
    }
    if (this->einsum_node_.in_indices(1).size() != 1 || !symbolic::eq(this->einsum_node_.in_index(1, 0), indvar)) {
        return false;
    }

    // Get the data flow graph
    auto& dfg = this->einsum_node_.get_parent();

    // Determine and check the base type of output
    auto& oedge = *dfg.out_edges(this->einsum_node_).begin();
    auto data_type = oedge.base_type().primitive_type();
    if (data_type != types::PrimitiveType::Float && data_type != types::PrimitiveType::Double) {
        return false;
    }

    // Check if all inputs have the same primitive type
    for (auto& iedge : dfg.in_edges(this->einsum_node_)) {
        if (iedge.base_type().primitive_type() != data_type) {
            return false;
        }
    }

    if (!get_impl_type(data_type)) { // no implementation for the given tune exists
        return false;
    }

    return true;
}

void Einsum2Dot::apply(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {
    // Get the data flow graph
    auto& dfg = this->einsum_node_.get_parent();

    // Get the block in which the einsum node lives
    auto* block = dynamic_cast<structured_control_flow::Block*>(dfg.get_parent());
    assert(block);

    // Get the number of iterations (n)
    symbolic::Expression n = this->einsum_node_.bound(0);

    // Determine the BLAS precision
    math::blas::BLAS_Precision precision;
    auto& datatype_oedge = *dfg.out_edges(this->einsum_node_).begin();
    types::PrimitiveType data_type = datatype_oedge.base_type().primitive_type();
    if (data_type == types::PrimitiveType::Float) {
        precision = math::blas::BLAS_Precision::s;
    } else {
        precision = math::blas::BLAS_Precision::d;
    }

    // Add the dot node
    auto& libnode = builder.add_library_node<
        math::blas::DotNode,
        const data_flow::ImplementationType&,
        const math::blas::BLAS_Precision&,
        symbolic::Expression>(
        *block, this->einsum_node_.debug_info(), this->get_impl_type(data_type).value(), precision, n
    );

    // Copy the memlets
    data_flow::AccessNode* leftover_access_node = nullptr;
    for (auto& iedge : dfg.in_edges(this->einsum_node_)) {
        if (iedge.dst_conn() == this->einsum_node_.input(0)) {
            builder.add_memlet(
                *block,
                iedge.src(),
                iedge.src_conn(),
                libnode,
                "__x",
                iedge.subset(),
                iedge.base_type(),
                iedge.debug_info()
            );
        } else if (iedge.dst_conn() == this->einsum_node_.input(1)) {
            builder.add_memlet(
                *block,
                iedge.src(),
                iedge.src_conn(),
                libnode,
                "__y",
                iedge.subset(),
                iedge.base_type(),
                iedge.debug_info()
            );
        } else if (iedge.dst_conn() == this->einsum_node_.input(2)) {
            leftover_access_node = dynamic_cast<data_flow::AccessNode*>(&iedge.src());
        }
    }
    for (auto& oedge : dfg.out_edges(this->einsum_node_)) {
        if (oedge.src_conn() == this->einsum_node_.output(0)) {
            builder.add_memlet(
                *block,
                libnode,
                "__out",
                oedge.dst(),
                oedge.dst_conn(),
                oedge.subset(),
                oedge.base_type(),
                oedge.debug_info()
            );
        }
    }

    // Remove the old memlets
    while (dfg.in_edges(this->einsum_node_).begin() != dfg.in_edges(this->einsum_node_).end()) {
        auto& iedge = *dfg.in_edges(this->einsum_node_).begin();
        builder.remove_memlet(*block, iedge);
    }
    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.remove_memlet(*block, oedge);
    }

    // Remove leftover access node
    builder.remove_node(*block, *leftover_access_node);

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

    analysis_manager.invalidate_all();
}

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

Einsum2Dot Einsum2Dot::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());
    assert(j.contains("impl_type"));

    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<einsum::EinsumNode*>(einsum_node_element);
    if (!einsum_node) {
        throw InvalidTransformationDescriptionException(
            "Element with ID " + std::to_string(einsum_node_id) + " is not an EinsumNode"
        );
    }

    std::string target_tune;
    if (j.contains("target_tune")) {
        target_tune = j.at("target_tune").get<std::string>();
    } else {
        target_tune = "none";
    }

    return Einsum2Dot(*einsum_node, target_tune);
}

} // namespace transformations
} // namespace sdfg

1	#include "sdfg/transformations/einsum2dot.h"
2
3	#include <cstddef>
4	#include <nlohmann/json_fwd.hpp>
5	#include <string>
6
7	#include "sdfg/analysis/analysis.h"
8	#include "sdfg/builder/structured_sdfg_builder.h"
9	#include "sdfg/data_flow/access_node.h"
10	#include "sdfg/data_flow/library_node.h"
11	#include "sdfg/data_flow/library_nodes/math/blas/dot_node.h"
12	#include "sdfg/data_flow/library_nodes/math/math.h"
13	#include "sdfg/einsum/einsum.h"
14	#include "sdfg/optimization_report/pass_report_consumer.h"
15	#include "sdfg/structured_control_flow/block.h"
16	#include "sdfg/symbolic/symbolic.h"
17	#include "sdfg/targets/cuda/cuda.h"
18	#include "sdfg/transformations/transformation.h"
19	#include "sdfg/types/type.h"
20	#include "sdfg/types/utils.h"
21
22	namespace sdfg {
23	namespace transformations {
24
25	Einsum2Dot::Einsum2Dot(einsum::EinsumNode& einsum_node, const std::string& target_tune)
26	: einsum_node_(einsum_node), target_tune_(target_tune) {}	×
27
28	std::string Einsum2Dot::name() const { return "Einsum2Dot"; }	×
29
30	std::optional<sdfg::data_flow::ImplementationType> Einsum2Dot::get_impl_type(types::PrimitiveType data_type) {	×
NEW 31	std::optional<sdfg::data_flow::ImplementationType> impl_type; // TODO make generic for any target	×
32	if (target_tune_ == "sequential") {	×
33	impl_type = sdfg::data_flow::ImplementationType_NONE;	×
34	} else if (target_tune_ == "openmp") {	×
35	impl_type = sdfg::math::blas::ImplementationType_BLAS;	×
36	} else if (target_tune_ == "cuda") {	×
37	impl_type = sdfg::cuda::blas::ImplementationType_CUBLASWithTransfers;	×
38	} else if (target_tune_ == "tenstorrent") {	×
39	if (data_type == types::PrimitiveType::Float) {	×
40	impl_type = data_flow::ImplementationType{"TENSTORRENT_WithTransfers"};	×
41	}	×
42	}	×
43
44	if (impl_type) {	×
45	return impl_type;	×
46	} else {	×
47	return std::nullopt;	×
48	}	×
49	}	×
50
51	bool Einsum2Dot::can_be_applied(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {	×
52	// Check dims
53	if (this->einsum_node_.dims().size() != 1 \|\| !symbolic::eq(this->einsum_node_.init(0), symbolic::zero())) {	×
54	return false;	×
55	}	×
56	symbolic::Symbol indvar = this->einsum_node_.indvar(0);	×
57
58	// Check out indices
59	if (this->einsum_node_.out_indices().size() != 0) {	×
60	return false;	×
61	}	×
62
63	// Check inputs
64	if (this->einsum_node_.inputs().size() != 3 \|\| this->einsum_node_.input(2) != this->einsum_node_.output(0)) {	×
65	return false;	×
66	}	×
67
68	// Check in indices
69	if (this->einsum_node_.in_indices(0).size() != 1 \|\| !symbolic::eq(this->einsum_node_.in_index(0, 0), indvar)) {	×
70	return false;	×
71	}	×
72	if (this->einsum_node_.in_indices(1).size() != 1 \|\| !symbolic::eq(this->einsum_node_.in_index(1, 0), indvar)) {	×
73	return false;	×
74	}	×
75
76	// Get the data flow graph
77	auto& dfg = this->einsum_node_.get_parent();	×
78
79	// Determine and check the base type of output
80	auto& oedge = *dfg.out_edges(this->einsum_node_).begin();	×
81	auto data_type = oedge.base_type().primitive_type();	×
82	if (data_type != types::PrimitiveType::Float && data_type != types::PrimitiveType::Double) {	×
83	return false;	×
84	}	×
85
86	// Check if all inputs have the same primitive type
87	for (auto& iedge : dfg.in_edges(this->einsum_node_)) {	×
88	if (iedge.base_type().primitive_type() != data_type) {	×
89	return false;	×
90	}	×
91	}	×
92
93	if (!get_impl_type(data_type)) { // no implementation for the given tune exists	×
94	return false;	×
95	}	×
96
97	return true;	×
98	}	×
99
100	void Einsum2Dot::apply(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {	×
101	// Get the data flow graph
102	auto& dfg = this->einsum_node_.get_parent();	×
103
104	// Get the block in which the einsum node lives
105	auto* block = dynamic_cast<structured_control_flow::Block*>(dfg.get_parent());	×
106	assert(block);	×
107
108	// Get the number of iterations (n)
109	symbolic::Expression n = this->einsum_node_.bound(0);	×
110
111	// Determine the BLAS precision
112	math::blas::BLAS_Precision precision;	×
113	auto& datatype_oedge = *dfg.out_edges(this->einsum_node_).begin();	×
114	types::PrimitiveType data_type = datatype_oedge.base_type().primitive_type();	×
115	if (data_type == types::PrimitiveType::Float) {	×
116	precision = math::blas::BLAS_Precision::s;	×
117	} else {	×
118	precision = math::blas::BLAS_Precision::d;	×
119	}	×
120
121	// Add the dot node
122	auto& libnode = builder.add_library_node<	×
123	math::blas::DotNode,	×
124	const data_flow::ImplementationType&,	×
125	const math::blas::BLAS_Precision&,	×
126	symbolic::Expression>(	×
127	*block, this->einsum_node_.debug_info(), this->get_impl_type(data_type).value(), precision, n	×
128	);	×
129
130	// Copy the memlets
131	data_flow::AccessNode* leftover_access_node = nullptr;	×
132	for (auto& iedge : dfg.in_edges(this->einsum_node_)) {	×
133	if (iedge.dst_conn() == this->einsum_node_.input(0)) {	×
134	builder.add_memlet(	×
135	*block,	×
136	iedge.src(),	×
137	iedge.src_conn(),	×
138	libnode,	×
139	"__x",	×
140	iedge.subset(),	×
141	iedge.base_type(),	×
142	iedge.debug_info()	×
143	);	×
144	} else if (iedge.dst_conn() == this->einsum_node_.input(1)) {	×
145	builder.add_memlet(	×
146	*block,	×
147	iedge.src(),	×
148	iedge.src_conn(),	×
149	libnode,	×
150	"__y",	×
151	iedge.subset(),	×
152	iedge.base_type(),	×
153	iedge.debug_info()	×
154	);	×
155	} else if (iedge.dst_conn() == this->einsum_node_.input(2)) {	×
156	leftover_access_node = dynamic_cast<data_flow::AccessNode*>(&iedge.src());	×
157	}	×
158	}	×
159	for (auto& oedge : dfg.out_edges(this->einsum_node_)) {	×
160	if (oedge.src_conn() == this->einsum_node_.output(0)) {	×
161	builder.add_memlet(	×
162	*block,	×
163	libnode,	×
164	"__out",	×
165	oedge.dst(),	×
166	oedge.dst_conn(),	×
167	oedge.subset(),	×
168	oedge.base_type(),	×
169	oedge.debug_info()	×
170	);	×
171	}	×
172	}	×
173
174	// Remove the old memlets
175	while (dfg.in_edges(this->einsum_node_).begin() != dfg.in_edges(this->einsum_node_).end()) {	×
176	auto& iedge = *dfg.in_edges(this->einsum_node_).begin();	×
177	builder.remove_memlet(*block, iedge);	×
178	}	×
179	while (dfg.out_edges(this->einsum_node_).begin() != dfg.out_edges(this->einsum_node_).end()) {	×
180	auto& oedge = *dfg.out_edges(this->einsum_node_).begin();	×
181	builder.remove_memlet(*block, oedge);	×
182	}	×
183
184	// Remove leftover access node
185	builder.remove_node(block, leftover_access_node);	×
186
187	// Remove the einsum node
188	builder.remove_node(*block, this->einsum_node_);	×
189
190	analysis_manager.invalidate_all();	×
191	}	×
192
193	void Einsum2Dot::to_json(nlohmann::json& j) const {	×
194	j["transformation_type"] = this->name();	×
195	j["einsum_node_element_id"] = this->einsum_node_.element_id();	×
196	j["target_tune"] = this->target_tune_;	×
197	}	×
198
199	Einsum2Dot Einsum2Dot::from_json(builder::StructuredSDFGBuilder& builder, const nlohmann::json& j) {	×
200	assert(j.contains("einsum_node_element_id"));	×
201	assert(j["einsum_node_element_id"].is_number_unsigned());	×
202	assert(j.contains("impl_type"));	×
203
204	size_t einsum_node_id = j["einsum_node_element_id"].get<size_t>();	×
205	auto* einsum_node_element = builder.find_element_by_id(einsum_node_id);	×
206	if (!einsum_node_element) {	×
207	throw InvalidTransformationDescriptionException(	×
208	"Element with ID " + std::to_string(einsum_node_id) + " not found"	×
209	);	×
210	}	×
211	auto* einsum_node = dynamic_cast<einsum::EinsumNode*>(einsum_node_element);	×
212	if (!einsum_node) {	×
213	throw InvalidTransformationDescriptionException(	×
214	"Element with ID " + std::to_string(einsum_node_id) + " is not an EinsumNode"	×
215	);	×
216	}	×
217
218	std::string target_tune;	×
219	if (j.contains("target_tune")) {	×
220	target_tune = j.at("target_tune").get<std::string>();	×
221	} else {	×
222	target_tune = "none";	×
223	}	×
224
225	return Einsum2Dot(*einsum_node, target_tune);	×
226	}	×
227
228	} // namespace transformations
229	} // namespace sdfg

daisytuner / docc / 23797200186

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