30005670487

Committed 23 Jul 2026 12:06PM UTC coverage: 64.077% (-0.003%) from 64.08%

Build # 30005670487

Build Type

Pull #872

github

Committed by

web-flow

Commit Message

Merge bdd6eb01e into 7b41ae7c7

Pull Request Pull Request #872: Added mapping for expand ops to PyTorch frontend

Coverage Stats

25 of 38 new or added lines in 3 files covered. (65.79%)

2 existing lines in 1 file now uncovered.

42674 of 66598 relevant lines covered (64.08%)

739.24 hits per line

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

0.0

/sdfg/src/data_flow/library_nodes/math/tensor/broadcast_node.cpp

#include "sdfg/data_flow/library_nodes/math/tensor/broadcast_node.h"
#include "sdfg/builder/structured_sdfg_builder.h"
#include "sdfg/structured_control_flow/for.h"

namespace sdfg {
namespace math {
namespace tensor {

BroadcastNode::BroadcastNode(
    size_t element_id,
    const DebugInfo& debug_info,
    const graph::Vertex vertex,
    data_flow::DataFlowGraph& parent,
    const std::vector<symbolic::Expression>& input_shape,
    const std::vector<symbolic::Expression>& output_shape
)
    : TensorNode(
          element_id,
          debug_info,
          vertex,
          parent,
          LibraryNodeType_Broadcast,
          {},
          {"Y", "X"},
          data_flow::ImplementationType_NONE
      ),
      input_shape_(input_shape), output_shape_(output_shape) {}

void BroadcastNode::validate(const Function& function) const {
    TensorNode::validate(function);

    auto& graph = this->get_parent();

    auto* iedge = graph.in_edge_for_connector(*this, inputs_.at(X_INPUT_IDX));
    auto& shape = static_cast<const types::Tensor&>(iedge->base_type());
    if (!shape.is_scalar()) {
        if (shape.shape().size() != this->input_shape_.size()) {
            throw InvalidSDFGException(
                "Library Node: Tensor shape must match node shape. Tensor shape: " +
                std::to_string(shape.shape().size()) + " Node shape: " + std::to_string(this->input_shape_.size())
            );
        }
        for (size_t i = 0; i < this->input_shape_.size(); ++i) {
            if (!symbolic::eq(shape.shape().at(i), this->input_shape_.at(i))) {
                throw InvalidSDFGException(
                    "Library Node: Tensor shape does not match expected shape. Tensor shape: " +
                    shape.shape().at(i)->__str__() + " Expected shape: " + this->input_shape_.at(i)->__str__()
                );
            }
        }
    }

    auto* oedge = graph.in_edge_for_connector(*this, inputs_.at(RESULT_PTR_IDX));
    auto& output_shape = static_cast<const types::Tensor&>(oedge->base_type());
    if (output_shape.shape().size() != this->output_shape_.size()) {
        throw InvalidSDFGException(
            "Library Node: Output tensor shape must match node shape. Output tensor shape: " +
            std::to_string(output_shape.shape().size()) + " Node shape: " + std::to_string(this->output_shape_.size())
        );
    }

    for (size_t i = 0; i < this->output_shape_.size(); ++i) {
        if (!symbolic::eq(output_shape.shape().at(i), this->output_shape_.at(i))) {
            throw InvalidSDFGException(
                "Library Node: Output tensor shape does not match expected shape. Output tensor shape: " +
                output_shape.shape().at(i)->__str__() + " Expected shape: " + this->output_shape_.at(i)->__str__()
            );
        }
    }
}

symbolic::SymbolSet BroadcastNode::symbols() const {
    symbolic::SymbolSet syms;
    for (const auto& dim : input_shape_) {
        for (auto& atom : symbolic::atoms(dim)) {
            syms.insert(atom);
        }
    }
    for (const auto& dim : output_shape_) {
        for (auto& atom : symbolic::atoms(dim)) {
            syms.insert(atom);
        }
    }
    return syms;
}

void BroadcastNode::replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {
    for (auto& dim : input_shape_) {
        dim = symbolic::subs(dim, old_expression, new_expression);
    }
    for (auto& dim : output_shape_) {
        dim = symbolic::subs(dim, old_expression, new_expression);
    }
}

void BroadcastNode::replace(const symbolic::ExpressionMapping& replacements) {
    for (auto& dim : input_shape_) {
        dim = symbolic::subs(dim, replacements);
    }
    for (auto& dim : output_shape_) {
        dim = symbolic::subs(dim, replacements);
    }
}

passes::LibNodeExpander::ExpandOutcome BroadcastNode::
    expand(passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block) {
    auto& dataflow = this->get_parent();

    if (dataflow.in_degree(*this) != 2 || dataflow.out_degree(*this) != 0) {
        return context.unable();
    }

    auto edges = dataflow.in_edges_by_connector(*this);
    auto& in_edge = *edges.at(X_INPUT_IDX);
    auto& result_ptr_edge = *edges.at(RESULT_PTR_IDX);


    using Use = passes::LibNodeExpander::InputUse;
    auto standalone = context.replacement_requires_access_nodes({Use::IndirectWrite, Use::IndirectRead});

    if (!standalone) {
        return context.unable();
    }

    symbolic::MultiExpression loop_vars;
    auto& builder = standalone->builder();
    structured_control_flow::Sequence* inner_scope = nullptr;

    for (size_t i = 0; i < output_shape_.size(); ++i) {
        std::string var_name = builder.find_new_name("_i" + std::to_string(i));
        builder.add_container(var_name, types::Scalar(types::PrimitiveType::Int64));

        auto sym_var = symbolic::symbol(var_name);
        auto condition = symbolic::Lt(sym_var, output_shape_[i]);
        auto init = symbolic::zero();
        auto update = symbolic::add(sym_var, symbolic::one());

        if (i == 0) {
            auto& loop = standalone->replace_with_structured_loop(
                passes::LibNodeExpander::AccessNodeExpand::LoopType::Map,
                sym_var,
                condition,
                init,
                update,
                structured_control_flow::ScheduleType_Sequential::create()
            );
            inner_scope = &loop.root();
        } else {
            auto& loop = builder.add_map(
                *inner_scope,
                sym_var,
                condition,
                init,
                update,
                structured_control_flow::ScheduleType_Sequential::create(),
                {},
                this->debug_info()
            );
            inner_scope = &loop.root();
        }
        loop_vars.push_back(sym_var);
    }

    auto& tasklet_block = builder.add_block(*inner_scope, {}, this->debug_info());

    auto& in_acc = standalone->add_indirect_read_access(tasklet_block, X_INPUT_IDX);
    auto& out_acc = standalone->add_indirect_write_access(tasklet_block, RESULT_PTR_IDX);

    symbolic::MultiExpression input_subset = {};
    size_t j = 0;
    for (size_t i = 0; i < output_shape_.size(); ++i) {
        if (j >= input_shape_.size()) {
            break;
        } else if (symbolic::eq(output_shape_[i], input_shape_[j])) {
            input_subset.push_back(loop_vars[i]);
            j++;
        } else if (symbolic::eq(input_shape_[j], symbolic::one())) {
            input_subset.push_back(symbolic::zero());
            j++;
        }
    }
    if (j < input_shape_.size()) {
        throw InvalidSDFGException("BroadcastNode: Could not resolve indvars for inputs");
    }
    auto& iedge_tensor = static_cast<const types::Tensor&>(in_edge.base_type());
    if (iedge_tensor.is_scalar()) {
        input_subset = {};
    }

    auto& tasklet =
        builder.add_tasklet(tasklet_block, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info());

    builder.add_computational_memlet(
        tasklet_block, in_acc, tasklet, "_in", input_subset, in_edge.base_type(), this->debug_info()
    );
    builder.add_computational_memlet(
        tasklet_block, tasklet, "_out", out_acc, loop_vars, result_ptr_edge.base_type(), this->debug_info()
    );

    return standalone->successfully_expanded();
}

std::unique_ptr<data_flow::DataFlowNode> BroadcastNode::
    clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {
    return std::unique_ptr<data_flow::DataFlowNode>(
        new BroadcastNode(element_id, this->debug_info(), vertex, parent, input_shape_, output_shape_)
    );
}

data_flow::PointerAccessType BroadcastNode::pointer_access_type(int input_idx) const {
    if (input_idx == RESULT_PTR_IDX) {
        return data_flow::PointerAccessMeta::create_full_write_only(symbolic::__nullptr__(), true);
    } else if (input_idx == X_INPUT_IDX) {
        return data_flow::PointerAccessMeta::create_read_only(symbolic::__nullptr__(), true);
    } else {
        return TensorNode::pointer_access_type(input_idx);
    }
}

} // namespace tensor
} // namespace math
} // namespace sdfg

1	#include "sdfg/data_flow/library_nodes/math/tensor/broadcast_node.h"
2	#include "sdfg/builder/structured_sdfg_builder.h"
3	#include "sdfg/structured_control_flow/for.h"
4
5	namespace sdfg {
6	namespace math {
7	namespace tensor {
8
9	BroadcastNode::BroadcastNode(
10	size_t element_id,
11	const DebugInfo& debug_info,
12	const graph::Vertex vertex,
13	data_flow::DataFlowGraph& parent,
14	const std::vector<symbolic::Expression>& input_shape,
15	const std::vector<symbolic::Expression>& output_shape
16	)
17	: TensorNode(	×
18	element_id,	×
19	debug_info,	×
20	vertex,	×
21	parent,	×
22	LibraryNodeType_Broadcast,	×
23	{},	×
24	{"Y", "X"},	×
25	data_flow::ImplementationType_NONE	×
26	),	×
27	input_shape_(input_shape), output_shape_(output_shape) {}	×
28
29	void BroadcastNode::validate(const Function& function) const {	×
30	TensorNode::validate(function);	×
31
32	auto& graph = this->get_parent();	×
33
34	auto* iedge = graph.in_edge_for_connector(*this, inputs_.at(X_INPUT_IDX));	×
35	auto& shape = static_cast<const types::Tensor&>(iedge->base_type());	×
36	if (!shape.is_scalar()) {	×
37	if (shape.shape().size() != this->input_shape_.size()) {	×
38	throw InvalidSDFGException(	×
39	"Library Node: Tensor shape must match node shape. Tensor shape: " +	×
40	std::to_string(shape.shape().size()) + " Node shape: " + std::to_string(this->input_shape_.size())	×
41	);	×
42	}	×
43	for (size_t i = 0; i < this->input_shape_.size(); ++i) {	×
44	if (!symbolic::eq(shape.shape().at(i), this->input_shape_.at(i))) {	×
45	throw InvalidSDFGException(	×
46	"Library Node: Tensor shape does not match expected shape. Tensor shape: " +	×
47	shape.shape().at(i)->__str__() + " Expected shape: " + this->input_shape_.at(i)->__str__()	×
48	);	×
49	}	×
50	}	×
51	}	×
52
53	auto* oedge = graph.in_edge_for_connector(*this, inputs_.at(RESULT_PTR_IDX));	×
54	auto& output_shape = static_cast<const types::Tensor&>(oedge->base_type());	×
55	if (output_shape.shape().size() != this->output_shape_.size()) {	×
56	throw InvalidSDFGException(	×
57	"Library Node: Output tensor shape must match node shape. Output tensor shape: " +	×
58	std::to_string(output_shape.shape().size()) + " Node shape: " + std::to_string(this->output_shape_.size())	×
59	);	×
60	}	×
61
62	for (size_t i = 0; i < this->output_shape_.size(); ++i) {	×
63	if (!symbolic::eq(output_shape.shape().at(i), this->output_shape_.at(i))) {	×
64	throw InvalidSDFGException(	×
65	"Library Node: Output tensor shape does not match expected shape. Output tensor shape: " +	×
66	output_shape.shape().at(i)->__str__() + " Expected shape: " + this->output_shape_.at(i)->__str__()	×
67	);	×
68	}	×
69	}	×
70	}	×
71
72	symbolic::SymbolSet BroadcastNode::symbols() const {	×
73	symbolic::SymbolSet syms;	×
74	for (const auto& dim : input_shape_) {	×
75	for (auto& atom : symbolic::atoms(dim)) {	×
76	syms.insert(atom);	×
77	}	×
78	}	×
79	for (const auto& dim : output_shape_) {	×
80	for (auto& atom : symbolic::atoms(dim)) {	×
81	syms.insert(atom);	×
82	}	×
83	}	×
84	return syms;	×
85	}	×
86
87	void BroadcastNode::replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {	×
88	for (auto& dim : input_shape_) {	×
89	dim = symbolic::subs(dim, old_expression, new_expression);	×
90	}	×
91	for (auto& dim : output_shape_) {	×
92	dim = symbolic::subs(dim, old_expression, new_expression);	×
93	}	×
94	}	×
95
96	void BroadcastNode::replace(const symbolic::ExpressionMapping& replacements) {	×
97	for (auto& dim : input_shape_) {	×
98	dim = symbolic::subs(dim, replacements);	×
99	}	×
100	for (auto& dim : output_shape_) {	×
101	dim = symbolic::subs(dim, replacements);	×
102	}	×
103	}	×
104
105	passes::LibNodeExpander::ExpandOutcome BroadcastNode::
106	expand(passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block) {	×
107	auto& dataflow = this->get_parent();	×
108
109	if (dataflow.in_degree(this) != 2 \|\| dataflow.out_degree(this) != 0) {	×
110	return context.unable();	×
111	}	×
112
113	auto edges = dataflow.in_edges_by_connector(*this);	×
114	auto& in_edge = *edges.at(X_INPUT_IDX);	×
115	auto& result_ptr_edge = *edges.at(RESULT_PTR_IDX);	×
116
117
118	using Use = passes::LibNodeExpander::InputUse;	×
119	auto standalone = context.replacement_requires_access_nodes({Use::IndirectWrite, Use::IndirectRead});	×
120
121	if (!standalone) {	×
122	return context.unable();	×
123	}	×
124
125	symbolic::MultiExpression loop_vars;	×
126	auto& builder = standalone->builder();	×
127	structured_control_flow::Sequence* inner_scope = nullptr;	×
128
129	for (size_t i = 0; i < output_shape_.size(); ++i) {	×
130	std::string var_name = builder.find_new_name("_i" + std::to_string(i));	×
131	builder.add_container(var_name, types::Scalar(types::PrimitiveType::Int64));	×
132
133	auto sym_var = symbolic::symbol(var_name);	×
134	auto condition = symbolic::Lt(sym_var, output_shape_[i]);	×
135	auto init = symbolic::zero();	×
136	auto update = symbolic::add(sym_var, symbolic::one());	×
137
138	if (i == 0) {	×
139	auto& loop = standalone->replace_with_structured_loop(	×
140	passes::LibNodeExpander::AccessNodeExpand::LoopType::Map,	×
141	sym_var,	×
142	condition,	×
143	init,	×
144	update,	×
145	structured_control_flow::ScheduleType_Sequential::create()	×
146	);	×
147	inner_scope = &loop.root();	×
148	} else {	×
149	auto& loop = builder.add_map(	×
150	*inner_scope,	×
151	sym_var,	×
152	condition,	×
153	init,	×
154	update,	×
155	structured_control_flow::ScheduleType_Sequential::create(),	×
156	{},	×
157	this->debug_info()	×
158	);	×
159	inner_scope = &loop.root();	×
160	}	×
161	loop_vars.push_back(sym_var);	×
162	}	×
163
164	auto& tasklet_block = builder.add_block(*inner_scope, {}, this->debug_info());	×
165
166	auto& in_acc = standalone->add_indirect_read_access(tasklet_block, X_INPUT_IDX);	×
167	auto& out_acc = standalone->add_indirect_write_access(tasklet_block, RESULT_PTR_IDX);	×
168
169	symbolic::MultiExpression input_subset = {};	×
170	size_t j = 0;	×
171	for (size_t i = 0; i < output_shape_.size(); ++i) {	×
172	if (j >= input_shape_.size()) {	×
173	break;	×
174	} else if (symbolic::eq(output_shape_[i], input_shape_[j])) {	×
175	input_subset.push_back(loop_vars[i]);	×
176	j++;	×
NEW 177	} else if (symbolic::eq(input_shape_[j], symbolic::one())) {	×
NEW 178	input_subset.push_back(symbolic::zero());	×
NEW 179	j++;	×
180	}	×
181	}	×
182	if (j < input_shape_.size()) {	×
183	throw InvalidSDFGException("BroadcastNode: Could not resolve indvars for inputs");	×
184	}	×
185	auto& iedge_tensor = static_cast<const types::Tensor&>(in_edge.base_type());	×
186	if (iedge_tensor.is_scalar()) {	×
187	input_subset = {};	×
188	}	×
189
190	auto& tasklet =	×
191	builder.add_tasklet(tasklet_block, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info());	×
192
193	builder.add_computational_memlet(	×
194	tasklet_block, in_acc, tasklet, "_in", input_subset, in_edge.base_type(), this->debug_info()	×
195	);	×
196	builder.add_computational_memlet(	×
197	tasklet_block, tasklet, "_out", out_acc, loop_vars, result_ptr_edge.base_type(), this->debug_info()	×
198	);	×
199
200	return standalone->successfully_expanded();	×
201	}	×
202
203	std::unique_ptr<data_flow::DataFlowNode> BroadcastNode::
204	clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {	×
205	return std::unique_ptr<data_flow::DataFlowNode>(	×
206	new BroadcastNode(element_id, this->debug_info(), vertex, parent, input_shape_, output_shape_)	×
207	);	×
208	}	×
209
210	data_flow::PointerAccessType BroadcastNode::pointer_access_type(int input_idx) const {	×
211	if (input_idx == RESULT_PTR_IDX) {	×
212	return data_flow::PointerAccessMeta::create_full_write_only(symbolic::__nullptr__(), true);	×
213	} else if (input_idx == X_INPUT_IDX) {	×
214	return data_flow::PointerAccessMeta::create_read_only(symbolic::__nullptr__(), true);	×
215	} else {	×
216	return TensorNode::pointer_access_type(input_idx);	×
217	}	×
218	}	×
219
220	} // namespace tensor
221	} // namespace math
222	} // namespace sdfg

daisytuner / docc / 30005670487

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