30463152603

Committed 29 Jul 2026 02:53PM UTC coverage: 64.817%. First build

Build # 30463152603

Build Type

Pull #910

github

Committed by

web-flow

Commit Message

Merge 318d9aa81 into 3eb01880e

Pull Request Pull Request #910: [PyTorch] Add Support for Embedding Layers

Coverage Stats

517 of 591 new or added lines in 6 files covered. (87.48%)

44583 of 68783 relevant lines covered (64.82%)

715.02 hits per line

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

83.71

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

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

namespace sdfg {
namespace math {
namespace tensor {

EmbeddingNode::EmbeddingNode(
    size_t element_id,
    const DebugInfo& debug_info,
    const graph::Vertex vertex,
    data_flow::DataFlowGraph& parent,
    const std::vector<symbolic::Expression>& weight_shape,
    const std::vector<symbolic::Expression>& index_shape,
    const data_flow::ImplementationType& impl_type
)
    : TensorNode(element_id, debug_info, vertex, parent, LibraryNodeType_Embedding, {}, {"Y", "W", "I"}, impl_type),
      weight_shape_(weight_shape), index_shape_(index_shape) {}

const std::vector<symbolic::Expression>& EmbeddingNode::weight_shape() const { return weight_shape_; }

const std::vector<symbolic::Expression>& EmbeddingNode::index_shape() const { return index_shape_; }

bool EmbeddingNode::supports_integer_types() const { return true; }

void EmbeddingNode::validate(const Function& function) const {
    // NOTE: The base TensorNode::validate enforces that all connected memlets share one
    // primitive type. That does not hold here: the index tensor is integer-typed while
    // the weight/output tensors may be floating-point. We therefore validate at the
    // MathNode level and add our own structural checks.
    MathNode::validate(function);

    if (weight_shape_.size() != 2) {
        throw InvalidSDFGException(
            "EmbeddingNode: weight must be 2-dimensional but got rank " + std::to_string(weight_shape_.size())
        );
    }
    if (index_shape_.empty()) {
        throw InvalidSDFGException("EmbeddingNode: index_shape must not be empty");
    }
}

symbolic::SymbolSet EmbeddingNode::symbols() const {
    symbolic::SymbolSet syms;
    for (const auto& dim : weight_shape_) {
        for (auto& atom : symbolic::atoms(dim)) {
            syms.insert(atom);
        }
    }
    for (const auto& dim : index_shape_) {
        for (auto& atom : symbolic::atoms(dim)) {
            syms.insert(atom);
        }
    }
    return syms;
}

void EmbeddingNode::replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {
    for (auto& dim : weight_shape_) {
        dim = symbolic::subs(dim, old_expression, new_expression);
    }
    for (auto& dim : index_shape_) {
        dim = symbolic::subs(dim, old_expression, new_expression);
    }
}

void EmbeddingNode::replace(const symbolic::ExpressionMapping& replacements) {
    for (auto& dim : weight_shape_) {
        dim = symbolic::subs(dim, replacements);
    }
    for (auto& dim : index_shape_) {
        dim = symbolic::subs(dim, replacements);
    }
}

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

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

    auto edges = dataflow.in_edges_by_connector(*this);
    auto& result_ptr_edge = *edges.at(RESULT_PTR_IDX);
    auto& weight_edge = *edges.at(W_INPUT_IDX);
    auto& index_edge = *edges.at(INDEX_IDX);

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

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

    // Output shape = index (broadcast) shape ++ [embedding_dim].
    size_t nB = index_shape_.size();
    std::vector<symbolic::Expression> output_shape;
    output_shape.reserve(nB + 1);
    for (const auto& s : index_shape_) {
        output_shape.push_back(s);
    }
    output_shape.push_back(weight_shape_[1]);

    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);
    }

    // The leading loop variables index into the index tensor.
    symbolic::MultiExpression index_vars(loop_vars.begin(), loop_vars.begin() + nB);

    // First load the index value into a scalar symbol so it can be used as the
    // data-dependent row coordinate in the gathering read below.
    std::string idx_name = builder.find_new_name("_idx");
    builder.add_container(idx_name, types::Scalar(types::PrimitiveType::Int64));
    auto idx_sym = symbolic::symbol(idx_name);

    auto& load_block = builder.add_block(*inner_scope, {}, this->debug_info());
    auto& idx_in_acc = standalone->add_indirect_read_access(load_block, INDEX_IDX);
    auto& idx_out_acc = builder.add_access(load_block, idx_name, this->debug_info());
    auto& load_tasklet =
        builder.add_tasklet(load_block, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info());
    builder.add_computational_memlet(
        load_block, idx_in_acc, load_tasklet, "_in", index_vars, index_edge.base_type(), this->debug_info()
    );
    builder.add_computational_memlet(
        load_block, load_tasklet, "_out", idx_out_acc, {}, types::Scalar(types::PrimitiveType::Int64), this->debug_info()
    );

    // Gather: out[b..., j] = W[idx, j]
    auto& tasklet_block = builder.add_block(*inner_scope, {}, this->debug_info());
    auto& in_acc = standalone->add_indirect_read_access(tasklet_block, W_INPUT_IDX);
    auto& out_acc = standalone->add_indirect_write_access(tasklet_block, RESULT_PTR_IDX);

    symbolic::MultiExpression weight_subset = {idx_sym, loop_vars[nB]};

    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", weight_subset, weight_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> EmbeddingNode::
    clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {
    return std::unique_ptr<data_flow::DataFlowNode>(
        new EmbeddingNode(element_id, this->debug_info(), vertex, parent, weight_shape_, index_shape_)
    );
}

data_flow::PointerAccessType EmbeddingNode::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 == W_INPUT_IDX || input_idx == INDEX_IDX) {
        return data_flow::PointerAccessMeta::create_read_only(symbolic::__nullptr__(), true);
    } else {
        return TensorNode::pointer_access_type(input_idx);
    }
}

nlohmann::json EmbeddingNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {
    const EmbeddingNode& embedding_node = static_cast<const EmbeddingNode&>(library_node);
    nlohmann::json j;

    j["code"] = embedding_node.code().value();

    serializer::JSONSerializer serializer;
    j["weight_shape"] = nlohmann::json::array();
    for (auto& dim : embedding_node.weight_shape()) {
        j["weight_shape"].push_back(serializer.expression(dim));
    }
    j["index_shape"] = nlohmann::json::array();
    for (auto& dim : embedding_node.index_shape()) {
        j["index_shape"].push_back(serializer.expression(dim));
    }

    return j;
}

data_flow::LibraryNode& EmbeddingNodeSerializer::deserialize(
    const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
) {
    assert(j.contains("element_id"));
    assert(j.contains("code"));
    assert(j.contains("debug_info"));
    assert(j.contains("weight_shape"));
    assert(j.contains("index_shape"));

    std::vector<symbolic::Expression> weight_shape;
    for (const auto& dim : j["weight_shape"]) {
        weight_shape.push_back(symbolic::parse(dim.get<std::string>()));
    }
    std::vector<symbolic::Expression> index_shape;
    for (const auto& dim : j["index_shape"]) {
        index_shape.push_back(symbolic::parse(dim.get<std::string>()));
    }

    sdfg::serializer::JSONSerializer serializer;
    DebugInfo debug_info = serializer.json_to_debug_info(j["debug_info"]);

    return builder.add_library_node<EmbeddingNode>(parent, debug_info, weight_shape, index_shape);
}

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

1	#include "sdfg/data_flow/library_nodes/math/tensor/embedding_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	EmbeddingNode::EmbeddingNode(
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>& weight_shape,
15	const std::vector<symbolic::Expression>& index_shape,
16	const data_flow::ImplementationType& impl_type
17	)
18	: TensorNode(element_id, debug_info, vertex, parent, LibraryNodeType_Embedding, {}, {"Y", "W", "I"}, impl_type),	6✔
19	weight_shape_(weight_shape), index_shape_(index_shape) {}	6✔
20
21	const std::vector<symbolic::Expression>& EmbeddingNode::weight_shape() const { return weight_shape_; }	2✔
22
23	const std::vector<symbolic::Expression>& EmbeddingNode::index_shape() const { return index_shape_; }	2✔
24
NEW 25	bool EmbeddingNode::supports_integer_types() const { return true; }	×
26
27	void EmbeddingNode::validate(const Function& function) const {	7✔
28	// NOTE: The base TensorNode::validate enforces that all connected memlets share one
29	// primitive type. That does not hold here: the index tensor is integer-typed while
30	// the weight/output tensors may be floating-point. We therefore validate at the
31	// MathNode level and add our own structural checks.
32	MathNode::validate(function);	7✔
33
34	if (weight_shape_.size() != 2) {	7✔
35	throw InvalidSDFGException(	1✔
36	"EmbeddingNode: weight must be 2-dimensional but got rank " + std::to_string(weight_shape_.size())	1✔
37	);	1✔
38	}	1✔
39	if (index_shape_.empty()) {	6✔
NEW 40	throw InvalidSDFGException("EmbeddingNode: index_shape must not be empty");	×
NEW 41	}	×
42	}	6✔
43
44	symbolic::SymbolSet EmbeddingNode::symbols() const {	2✔
45	symbolic::SymbolSet syms;	2✔
46	for (const auto& dim : weight_shape_) {	4✔
47	for (auto& atom : symbolic::atoms(dim)) {	4✔
48	syms.insert(atom);	4✔
49	}	4✔
50	}	4✔
51	for (const auto& dim : index_shape_) {	2✔
52	for (auto& atom : symbolic::atoms(dim)) {	2✔
53	syms.insert(atom);	2✔
54	}	2✔
55	}	2✔
56	return syms;	2✔
57	}	2✔
58
59	void EmbeddingNode::replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {	1✔
60	for (auto& dim : weight_shape_) {	2✔
61	dim = symbolic::subs(dim, old_expression, new_expression);	2✔
62	}	2✔
63	for (auto& dim : index_shape_) {	1✔
64	dim = symbolic::subs(dim, old_expression, new_expression);	1✔
65	}	1✔
66	}	1✔
67
NEW 68	void EmbeddingNode::replace(const symbolic::ExpressionMapping& replacements) {	×
NEW 69	for (auto& dim : weight_shape_) {	×
NEW 70	dim = symbolic::subs(dim, replacements);	×
NEW 71	}	×
NEW 72	for (auto& dim : index_shape_) {	×
NEW 73	dim = symbolic::subs(dim, replacements);	×
NEW 74	}	×
NEW 75	}	×
76
77	passes::LibNodeExpander::ExpandOutcome EmbeddingNode::
78	expand(passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block) {	2✔
79	auto& dataflow = this->get_parent();	2✔
80
81	if (dataflow.in_degree(this) != 3 \|\| dataflow.out_degree(this) != 0) {	2✔
NEW 82	return context.unable();	×
NEW 83	}	×
84
85	auto edges = dataflow.in_edges_by_connector(*this);	2✔
86	auto& result_ptr_edge = *edges.at(RESULT_PTR_IDX);	2✔
87	auto& weight_edge = *edges.at(W_INPUT_IDX);	2✔
88	auto& index_edge = *edges.at(INDEX_IDX);	2✔
89
90	using Use = passes::LibNodeExpander::InputUse;	2✔
91	std::vector<Use> uses = {Use::IndirectWrite, Use::IndirectRead, Use::IndirectRead};	2✔
92	auto standalone = context.replacement_requires_access_nodes(uses);	2✔
93
94	if (!standalone) {	2✔
NEW 95	return context.unable();	×
NEW 96	}	×
97
98	// Output shape = index (broadcast) shape ++ [embedding_dim].
99	size_t nB = index_shape_.size();	2✔
100	std::vector<symbolic::Expression> output_shape;	2✔
101	output_shape.reserve(nB + 1);	2✔
102	for (const auto& s : index_shape_) {	3✔
103	output_shape.push_back(s);	3✔
104	}	3✔
105	output_shape.push_back(weight_shape_[1]);	2✔
106
107	symbolic::MultiExpression loop_vars;	2✔
108	auto& builder = standalone->builder();	2✔
109	structured_control_flow::Sequence* inner_scope = nullptr;	2✔
110
111	for (size_t i = 0; i < output_shape.size(); ++i) {	7✔
112	std::string var_name = builder.find_new_name("_i" + std::to_string(i));	5✔
113	builder.add_container(var_name, types::Scalar(types::PrimitiveType::Int64));	5✔
114
115	auto sym_var = symbolic::symbol(var_name);	5✔
116	auto condition = symbolic::Lt(sym_var, output_shape[i]);	5✔
117	auto init = symbolic::zero();	5✔
118	auto update = symbolic::add(sym_var, symbolic::one());	5✔
119
120	if (i == 0) {	5✔
121	auto& loop = standalone->replace_with_structured_loop(	2✔
122	passes::LibNodeExpander::AccessNodeExpand::LoopType::Map,	2✔
123	sym_var,	2✔
124	condition,	2✔
125	init,	2✔
126	update,	2✔
127	structured_control_flow::ScheduleType_Sequential::create()	2✔
128	);	2✔
129	inner_scope = &loop.root();	2✔
130	} else {	3✔
131	auto& loop = builder.add_map(	3✔
132	*inner_scope,	3✔
133	sym_var,	3✔
134	condition,	3✔
135	init,	3✔
136	update,	3✔
137	structured_control_flow::ScheduleType_Sequential::create(),	3✔
138	this->debug_info()	3✔
139	);	3✔
140	inner_scope = &loop.root();	3✔
141	}	3✔
142	loop_vars.push_back(sym_var);	5✔
143	}	5✔
144
145	// The leading loop variables index into the index tensor.
146	symbolic::MultiExpression index_vars(loop_vars.begin(), loop_vars.begin() + nB);	2✔
147
148	// First load the index value into a scalar symbol so it can be used as the
149	// data-dependent row coordinate in the gathering read below.
150	std::string idx_name = builder.find_new_name("_idx");	2✔
151	builder.add_container(idx_name, types::Scalar(types::PrimitiveType::Int64));	2✔
152	auto idx_sym = symbolic::symbol(idx_name);	2✔
153
154	auto& load_block = builder.add_block(*inner_scope, {}, this->debug_info());	2✔
155	auto& idx_in_acc = standalone->add_indirect_read_access(load_block, INDEX_IDX);	2✔
156	auto& idx_out_acc = builder.add_access(load_block, idx_name, this->debug_info());	2✔
157	auto& load_tasklet =	2✔
158	builder.add_tasklet(load_block, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info());	2✔
159	builder.add_computational_memlet(	2✔
160	load_block, idx_in_acc, load_tasklet, "_in", index_vars, index_edge.base_type(), this->debug_info()	2✔
161	);	2✔
162	builder.add_computational_memlet(	2✔
163	load_block, load_tasklet, "_out", idx_out_acc, {}, types::Scalar(types::PrimitiveType::Int64), this->debug_info()	2✔
164	);	2✔
165
166	// Gather: out[b..., j] = W[idx, j]
167	auto& tasklet_block = builder.add_block(*inner_scope, {}, this->debug_info());	2✔
168	auto& in_acc = standalone->add_indirect_read_access(tasklet_block, W_INPUT_IDX);	2✔
169	auto& out_acc = standalone->add_indirect_write_access(tasklet_block, RESULT_PTR_IDX);	2✔
170
171	symbolic::MultiExpression weight_subset = {idx_sym, loop_vars[nB]};	2✔
172
173	auto& tasklet =	2✔
174	builder.add_tasklet(tasklet_block, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info());	2✔
175	builder.add_computational_memlet(	2✔
176	tasklet_block, in_acc, tasklet, "_in", weight_subset, weight_edge.base_type(), this->debug_info()	2✔
177	);	2✔
178	builder.add_computational_memlet(	2✔
179	tasklet_block, tasklet, "_out", out_acc, loop_vars, result_ptr_edge.base_type(), this->debug_info()	2✔
180	);	2✔
181
182	return standalone->successfully_expanded();	2✔
183	}	2✔
184
185	std::unique_ptr<data_flow::DataFlowNode> EmbeddingNode::
NEW 186	clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {	×
NEW 187	return std::unique_ptr<data_flow::DataFlowNode>(	×
NEW 188	new EmbeddingNode(element_id, this->debug_info(), vertex, parent, weight_shape_, index_shape_)	×
NEW 189	);	×
NEW 190	}	×
191
NEW 192	data_flow::PointerAccessType EmbeddingNode::pointer_access_type(int input_idx) const {	×
NEW 193	if (input_idx == RESULT_PTR_IDX) {	×
NEW 194	return data_flow::PointerAccessMeta::create_full_write_only(symbolic::__nullptr__(), true);	×
NEW 195	} else if (input_idx == W_INPUT_IDX \|\| input_idx == INDEX_IDX) {	×
NEW 196	return data_flow::PointerAccessMeta::create_read_only(symbolic::__nullptr__(), true);	×
NEW 197	} else {	×
NEW 198	return TensorNode::pointer_access_type(input_idx);	×
NEW 199	}	×
NEW 200	}	×
201
202	nlohmann::json EmbeddingNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {	1✔
203	const EmbeddingNode& embedding_node = static_cast<const EmbeddingNode&>(library_node);	1✔
204	nlohmann::json j;	1✔
205
206	j["code"] = embedding_node.code().value();	1✔
207
208	serializer::JSONSerializer serializer;	1✔
209	j["weight_shape"] = nlohmann::json::array();	1✔
210	for (auto& dim : embedding_node.weight_shape()) {	2✔
211	j["weight_shape"].push_back(serializer.expression(dim));	2✔
212	}	2✔
213	j["index_shape"] = nlohmann::json::array();	1✔
214	for (auto& dim : embedding_node.index_shape()) {	1✔
215	j["index_shape"].push_back(serializer.expression(dim));	1✔
216	}	1✔
217
218	return j;	1✔
219	}	1✔
220
221	data_flow::LibraryNode& EmbeddingNodeSerializer::deserialize(
222	const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
223	) {	1✔
224	assert(j.contains("element_id"));	1✔
225	assert(j.contains("code"));	1✔
226	assert(j.contains("debug_info"));	1✔
227	assert(j.contains("weight_shape"));	1✔
228	assert(j.contains("index_shape"));	1✔
229
230	std::vector<symbolic::Expression> weight_shape;	1✔
231	for (const auto& dim : j["weight_shape"]) {	2✔
232	weight_shape.push_back(symbolic::parse(dim.get<std::string>()));	2✔
233	}	2✔
234	std::vector<symbolic::Expression> index_shape;	1✔
235	for (const auto& dim : j["index_shape"]) {	1✔
236	index_shape.push_back(symbolic::parse(dim.get<std::string>()));	1✔
237	}	1✔
238
239	sdfg::serializer::JSONSerializer serializer;	1✔
240	DebugInfo debug_info = serializer.json_to_debug_info(j["debug_info"]);	1✔
241
242	return builder.add_library_node<EmbeddingNode>(parent, debug_info, weight_shape, index_shape);	1✔
243	}	1✔
244
245	} // namespace tensor
246	} // namespace math
247	} // namespace sdfg

daisytuner / docc / 30463152603

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