30527237467

Committed 30 Jul 2026 08:34AM UTC coverage: 64.718% (+0.1%) from 64.62%

Build # 30527237467

Build Type

push

github

Committed by

web-flow

Commit Message

Add aten.index parser and implmentation (#905)

Coverage Stats

245 of 275 new or added lines in 3 files covered. (89.09%)

44312 of 68469 relevant lines covered (64.72%)

717.1 hits per line

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

90.09

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

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

namespace sdfg {
namespace math {
namespace tensor {

static std::vector<std::string> make_index_inputs(long long num_indices) {
    std::vector<std::string> inputs = {"Y", "X"};
    for (long long j = 0; j < num_indices; ++j) {
        inputs.push_back("I" + std::to_string(j));
    }
    return inputs;
}

IndexNode::IndexNode(
    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>& index_shape,
    long long dim_offset,
    long long num_indices,
    const data_flow::ImplementationType& impl_type
)
    : TensorNode(
          element_id, debug_info, vertex, parent, LibraryNodeType_Index, {}, make_index_inputs(num_indices), impl_type
      ),
      input_shape_(input_shape), index_shape_(index_shape), dim_offset_(dim_offset), num_indices_(num_indices) {}

const std::vector<symbolic::Expression>& IndexNode::input_shape() const { return input_shape_; }

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

long long IndexNode::dim_offset() const { return dim_offset_; }

long long IndexNode::num_indices() const { return num_indices_; }

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

void IndexNode::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 tensors are integer-typed while
    // the data tensors may be floating-point. We therefore validate at the MathNode level
    // and add our own structural checks.
    MathNode::validate(function);

    if (num_indices_ < 1) {
        throw InvalidSDFGException("IndexNode: expected at least one index tensor but got " + std::to_string(num_indices_));
    }
    if (index_shape_.empty()) {
        throw InvalidSDFGException("IndexNode: index_shape must not be empty");
    }
    if (dim_offset_ < 0 || dim_offset_ + num_indices_ > static_cast<long long>(input_shape_.size())) {
        throw InvalidSDFGException(
            "IndexNode: indexed dimension block out of range. dim_offset: " + std::to_string(dim_offset_) +
            " num_indices: " + std::to_string(num_indices_) + " rank: " + std::to_string(input_shape_.size())
        );
    }
}

symbolic::SymbolSet IndexNode::symbols() const {
    symbolic::SymbolSet syms;
    for (const auto& dim : input_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 IndexNode::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 : index_shape_) {
        dim = symbolic::subs(dim, old_expression, new_expression);
    }
}

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

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

    size_t expected_inputs = 2 + static_cast<size_t>(num_indices_);
    if (dataflow.in_degree(*this) != expected_inputs || 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& in_edge = *edges.at(X_INPUT_IDX);
    std::vector<data_flow::Memlet*> index_edges;
    index_edges.reserve(num_indices_);
    for (long long j = 0; j < num_indices_; ++j) {
        index_edges.push_back(edges.at(FIRST_INDEX_IDX + j));
    }

    using Use = passes::LibNodeExpander::InputUse;
    std::vector<Use> uses = {Use::IndirectWrite, Use::IndirectRead};
    for (long long j = 0; j < num_indices_; ++j) {
        uses.push_back(Use::IndirectRead);
    }
    auto standalone = context.replacement_requires_access_nodes(uses);

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

    // Output shape = leading dims ++ index (broadcast) shape ++ trailing dims.
    size_t nB = index_shape_.size();
    std::vector<symbolic::Expression> output_shape;
    output_shape.reserve(dim_offset_ + nB + (input_shape_.size() - dim_offset_ - num_indices_));
    for (long long d = 0; d < dim_offset_; ++d) {
        output_shape.push_back(input_shape_[d]);
    }
    for (const auto& s : index_shape_) {
        output_shape.push_back(s);
    }
    for (long long d = dim_offset_ + num_indices_; d < static_cast<long long>(input_shape_.size()); ++d) {
        output_shape.push_back(input_shape_[d]);
    }

    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 broadcast loop variables index into the (1:1 shaped) index tensors.
    symbolic::MultiExpression broadcast_vars(loop_vars.begin() + dim_offset_, loop_vars.begin() + dim_offset_ + nB);

    // First load each index value into a scalar symbol so it can be used as a
    // data-dependent coordinate in the gathering read below.
    std::vector<symbolic::Symbol> idx_syms;
    idx_syms.reserve(num_indices_);
    auto& load_block = builder.add_block(*inner_scope, {}, this->debug_info());
    for (long long j = 0; j < num_indices_; ++j) {
        std::string idx_name = builder.find_new_name("_idx" + std::to_string(j));
        builder.add_container(idx_name, types::Scalar(types::PrimitiveType::Int64));
        idx_syms.push_back(symbolic::symbol(idx_name));

        auto& idx_in_acc = standalone->add_indirect_read_access(load_block, FIRST_INDEX_IDX + j);
        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", broadcast_vars, index_edges[j]->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[l..., b..., t...] = X[l..., idx_0, ..., idx_{k-1}, t...]
    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;
    input_subset.reserve(input_shape_.size());
    for (long long d = 0; d < dim_offset_; ++d) {
        input_subset.push_back(loop_vars[d]);
    }
    for (long long j = 0; j < num_indices_; ++j) {
        input_subset.push_back(idx_syms[j]);
    }
    for (size_t i = dim_offset_ + nB; i < output_shape.size(); ++i) {
        input_subset.push_back(loop_vars[i]);
    }

    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> IndexNode::
    clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {
    return std::unique_ptr<data_flow::DataFlowNode>(new IndexNode(
        element_id, this->debug_info(), vertex, parent, input_shape_, index_shape_, dim_offset_, num_indices_
    ));
}

data_flow::PointerAccessType IndexNode::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 && input_idx < FIRST_INDEX_IDX + num_indices_) {
        return data_flow::PointerAccessMeta::create_read_only(symbolic::__nullptr__(), true);
    } else {
        return TensorNode::pointer_access_type(input_idx);
    }
}

nlohmann::json IndexNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {
    const IndexNode& index_node = static_cast<const IndexNode&>(library_node);
    nlohmann::json j;

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

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

    j["dim_offset"] = index_node.dim_offset();
    j["num_indices"] = index_node.num_indices();

    return j;
}

data_flow::LibraryNode& IndexNodeSerializer::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("input_shape"));
    assert(j.contains("index_shape"));
    assert(j.contains("dim_offset"));
    assert(j.contains("num_indices"));

    std::vector<symbolic::Expression> input_shape;
    for (const auto& dim : j["input_shape"]) {
        input_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>()));
    }

    long long dim_offset = j["dim_offset"].get<long long>();
    long long num_indices = j["num_indices"].get<long long>();

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

    return builder.add_library_node<IndexNode>(parent, debug_info, input_shape, index_shape, dim_offset, num_indices);
}

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

1	#include "sdfg/data_flow/library_nodes/math/tensor/index_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	static std::vector<std::string> make_index_inputs(long long num_indices) {	7✔
10	std::vector<std::string> inputs = {"Y", "X"};	7✔
11	for (long long j = 0; j < num_indices; ++j) {	15✔
12	inputs.push_back("I" + std::to_string(j));	8✔
13	}	8✔
14	return inputs;	7✔
15	}	7✔
16
17	IndexNode::IndexNode(
18	size_t element_id,
19	const DebugInfo& debug_info,
20	const graph::Vertex vertex,
21	data_flow::DataFlowGraph& parent,
22	const std::vector<symbolic::Expression>& input_shape,
23	const std::vector<symbolic::Expression>& index_shape,
24	long long dim_offset,
25	long long num_indices,
26	const data_flow::ImplementationType& impl_type
27	)
28	: TensorNode(	7✔
29	element_id, debug_info, vertex, parent, LibraryNodeType_Index, {}, make_index_inputs(num_indices), impl_type	7✔
30	),	7✔
31	input_shape_(input_shape), index_shape_(index_shape), dim_offset_(dim_offset), num_indices_(num_indices) {}	7✔
32
33	const std::vector<symbolic::Expression>& IndexNode::input_shape() const { return input_shape_; }	1✔
34
35	const std::vector<symbolic::Expression>& IndexNode::index_shape() const { return index_shape_; }	1✔
36
37	long long IndexNode::dim_offset() const { return dim_offset_; }	1✔
38
39	long long IndexNode::num_indices() const { return num_indices_; }	1✔
40
NEW 41	bool IndexNode::supports_integer_types() const { return true; }	×
42
43	void IndexNode::validate(const Function& function) const {	8✔
44	// NOTE: The base TensorNode::validate enforces that all connected memlets share one
45	// primitive type. That does not hold here: the index tensors are integer-typed while
46	// the data tensors may be floating-point. We therefore validate at the MathNode level
47	// and add our own structural checks.
48	MathNode::validate(function);	8✔
49
50	if (num_indices_ < 1) {	8✔
NEW 51	throw InvalidSDFGException("IndexNode: expected at least one index tensor but got " + std::to_string(num_indices_));	×
NEW 52	}	×
53	if (index_shape_.empty()) {	8✔
NEW 54	throw InvalidSDFGException("IndexNode: index_shape must not be empty");	×
NEW 55	}	×
56	if (dim_offset_ < 0 \|\| dim_offset_ + num_indices_ > static_cast<long long>(input_shape_.size())) {	8✔
57	throw InvalidSDFGException(	1✔
58	"IndexNode: indexed dimension block out of range. dim_offset: " + std::to_string(dim_offset_) +	1✔
59	" num_indices: " + std::to_string(num_indices_) + " rank: " + std::to_string(input_shape_.size())	1✔
60	);	1✔
61	}	1✔
62	}	8✔
63
64	symbolic::SymbolSet IndexNode::symbols() const {	2✔
65	symbolic::SymbolSet syms;	2✔
66	for (const auto& dim : input_shape_) {	4✔
67	for (auto& atom : symbolic::atoms(dim)) {	4✔
68	syms.insert(atom);	4✔
69	}	4✔
70	}	4✔
71	for (const auto& dim : index_shape_) {	2✔
72	for (auto& atom : symbolic::atoms(dim)) {	2✔
73	syms.insert(atom);	2✔
74	}	2✔
75	}	2✔
76	return syms;	2✔
77	}	2✔
78
79	void IndexNode::replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {	1✔
80	for (auto& dim : input_shape_) {	2✔
81	dim = symbolic::subs(dim, old_expression, new_expression);	2✔
82	}	2✔
83	for (auto& dim : index_shape_) {	1✔
84	dim = symbolic::subs(dim, old_expression, new_expression);	1✔
85	}	1✔
86	}	1✔
87
88	void IndexNode::replace(const symbolic::ExpressionMapping& replacements) {	1✔
89	for (auto& dim : input_shape_) {	2✔
90	dim = symbolic::subs(dim, replacements);	2✔
91	}	2✔
92	for (auto& dim : index_shape_) {	1✔
93	dim = symbolic::subs(dim, replacements);	1✔
94	}	1✔
95	}	1✔
96
97	passes::LibNodeExpander::ExpandOutcome IndexNode::
98	expand(passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block) {	3✔
99	auto& dataflow = this->get_parent();	3✔
100
101	size_t expected_inputs = 2 + static_cast<size_t>(num_indices_);	3✔
102	if (dataflow.in_degree(this) != expected_inputs \|\| dataflow.out_degree(this) != 0) {	3✔
NEW 103	return context.unable();	×
NEW 104	}	×
105
106	auto edges = dataflow.in_edges_by_connector(*this);	3✔
107	auto& result_ptr_edge = *edges.at(RESULT_PTR_IDX);	3✔
108	auto& in_edge = *edges.at(X_INPUT_IDX);	3✔
109	std::vector<data_flow::Memlet*> index_edges;	3✔
110	index_edges.reserve(num_indices_);	3✔
111	for (long long j = 0; j < num_indices_; ++j) {	7✔
112	index_edges.push_back(edges.at(FIRST_INDEX_IDX + j));	4✔
113	}	4✔
114
115	using Use = passes::LibNodeExpander::InputUse;	3✔
116	std::vector<Use> uses = {Use::IndirectWrite, Use::IndirectRead};	3✔
117	for (long long j = 0; j < num_indices_; ++j) {	7✔
118	uses.push_back(Use::IndirectRead);	4✔
119	}	4✔
120	auto standalone = context.replacement_requires_access_nodes(uses);	3✔
121
122	if (!standalone) {	3✔
NEW 123	return context.unable();	×
NEW 124	}	×
125
126	// Output shape = leading dims ++ index (broadcast) shape ++ trailing dims.
127	size_t nB = index_shape_.size();	3✔
128	std::vector<symbolic::Expression> output_shape;	3✔
129	output_shape.reserve(dim_offset_ + nB + (input_shape_.size() - dim_offset_ - num_indices_));	3✔
130	for (long long d = 0; d < dim_offset_; ++d) {	4✔
131	output_shape.push_back(input_shape_[d]);	1✔
132	}	1✔
133	for (const auto& s : index_shape_) {	3✔
134	output_shape.push_back(s);	3✔
135	}	3✔
136	for (long long d = dim_offset_ + num_indices_; d < static_cast<long long>(input_shape_.size()); ++d) {	5✔
137	output_shape.push_back(input_shape_[d]);	2✔
138	}	2✔
139
140	symbolic::MultiExpression loop_vars;	3✔
141	auto& builder = standalone->builder();	3✔
142	structured_control_flow::Sequence* inner_scope = nullptr;	3✔
143
144	for (size_t i = 0; i < output_shape.size(); ++i) {	9✔
145	std::string var_name = builder.find_new_name("_i" + std::to_string(i));	6✔
146	builder.add_container(var_name, types::Scalar(types::PrimitiveType::Int64));	6✔
147
148	auto sym_var = symbolic::symbol(var_name);	6✔
149	auto condition = symbolic::Lt(sym_var, output_shape[i]);	6✔
150	auto init = symbolic::zero();	6✔
151	auto update = symbolic::add(sym_var, symbolic::one());	6✔
152
153	if (i == 0) {	6✔
154	auto& loop = standalone->replace_with_structured_loop(	3✔
155	passes::LibNodeExpander::AccessNodeExpand::LoopType::Map,	3✔
156	sym_var,	3✔
157	condition,	3✔
158	init,	3✔
159	update,	3✔
160	structured_control_flow::ScheduleType_Sequential::create()	3✔
161	);	3✔
162	inner_scope = &loop.root();	3✔
163	} else {	3✔
164	auto& loop = builder.add_map(	3✔
165	*inner_scope,	3✔
166	sym_var,	3✔
167	condition,	3✔
168	init,	3✔
169	update,	3✔
170	structured_control_flow::ScheduleType_Sequential::create(),	3✔
171	this->debug_info()	3✔
172	);	3✔
173	inner_scope = &loop.root();	3✔
174	}	3✔
175	loop_vars.push_back(sym_var);	6✔
176	}	6✔
177
178	// The broadcast loop variables index into the (1:1 shaped) index tensors.
179	symbolic::MultiExpression broadcast_vars(loop_vars.begin() + dim_offset_, loop_vars.begin() + dim_offset_ + nB);	3✔
180
181	// First load each index value into a scalar symbol so it can be used as a
182	// data-dependent coordinate in the gathering read below.
183	std::vector<symbolic::Symbol> idx_syms;	3✔
184	idx_syms.reserve(num_indices_);	3✔
185	auto& load_block = builder.add_block(*inner_scope, {}, this->debug_info());	3✔
186	for (long long j = 0; j < num_indices_; ++j) {	7✔
187	std::string idx_name = builder.find_new_name("_idx" + std::to_string(j));	4✔
188	builder.add_container(idx_name, types::Scalar(types::PrimitiveType::Int64));	4✔
189	idx_syms.push_back(symbolic::symbol(idx_name));	4✔
190
191	auto& idx_in_acc = standalone->add_indirect_read_access(load_block, FIRST_INDEX_IDX + j);	4✔
192	auto& idx_out_acc = builder.add_access(load_block, idx_name, this->debug_info());	4✔
193	auto& load_tasklet =	4✔
194	builder.add_tasklet(load_block, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info());	4✔
195	builder.add_computational_memlet(	4✔
196	load_block, idx_in_acc, load_tasklet, "_in", broadcast_vars, index_edges[j]->base_type(), this->debug_info()	4✔
197	);	4✔
198	builder.add_computational_memlet(	4✔
199	load_block,	4✔
200	load_tasklet,	4✔
201	"_out",	4✔
202	idx_out_acc,	4✔
203	{},	4✔
204	types::Scalar(types::PrimitiveType::Int64),	4✔
205	this->debug_info()	4✔
206	);	4✔
207	}	4✔
208
209	// Gather: out[l..., b..., t...] = X[l..., idx_0, ..., idx_{k-1}, t...]
210	auto& tasklet_block = builder.add_block(*inner_scope, {}, this->debug_info());	3✔
211	auto& in_acc = standalone->add_indirect_read_access(tasklet_block, X_INPUT_IDX);	3✔
212	auto& out_acc = standalone->add_indirect_write_access(tasklet_block, RESULT_PTR_IDX);	3✔
213
214	symbolic::MultiExpression input_subset;	3✔
215	input_subset.reserve(input_shape_.size());	3✔
216	for (long long d = 0; d < dim_offset_; ++d) {	4✔
217	input_subset.push_back(loop_vars[d]);	1✔
218	}	1✔
219	for (long long j = 0; j < num_indices_; ++j) {	7✔
220	input_subset.push_back(idx_syms[j]);	4✔
221	}	4✔
222	for (size_t i = dim_offset_ + nB; i < output_shape.size(); ++i) {	5✔
223	input_subset.push_back(loop_vars[i]);	2✔
224	}	2✔
225
226	auto& tasklet =	3✔
227	builder.add_tasklet(tasklet_block, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info());	3✔
228	builder.add_computational_memlet(	3✔
229	tasklet_block, in_acc, tasklet, "_in", input_subset, in_edge.base_type(), this->debug_info()	3✔
230	);	3✔
231	builder.add_computational_memlet(	3✔
232	tasklet_block, tasklet, "_out", out_acc, loop_vars, result_ptr_edge.base_type(), this->debug_info()	3✔
233	);	3✔
234
235	return standalone->successfully_expanded();	3✔
236	}	3✔
237
238	std::unique_ptr<data_flow::DataFlowNode> IndexNode::
NEW 239	clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {	×
NEW 240	return std::unique_ptr<data_flow::DataFlowNode>(new IndexNode(	×
NEW 241	element_id, this->debug_info(), vertex, parent, input_shape_, index_shape_, dim_offset_, num_indices_	×
NEW 242	));	×
NEW 243	}	×
244
NEW 245	data_flow::PointerAccessType IndexNode::pointer_access_type(int input_idx) const {	×
NEW 246	if (input_idx == RESULT_PTR_IDX) {	×
NEW 247	return data_flow::PointerAccessMeta::create_full_write_only(symbolic::__nullptr__(), true);	×
NEW 248	} else if (input_idx >= X_INPUT_IDX && input_idx < FIRST_INDEX_IDX + num_indices_) {	×
NEW 249	return data_flow::PointerAccessMeta::create_read_only(symbolic::__nullptr__(), true);	×
NEW 250	} else {	×
NEW 251	return TensorNode::pointer_access_type(input_idx);	×
NEW 252	}	×
NEW 253	}	×
254
255	nlohmann::json IndexNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {	1✔
256	const IndexNode& index_node = static_cast<const IndexNode&>(library_node);	1✔
257	nlohmann::json j;	1✔
258
259	j["code"] = index_node.code().value();	1✔
260
261	serializer::JSONSerializer serializer;	1✔
262	j["input_shape"] = nlohmann::json::array();	1✔
263	for (auto& dim : index_node.input_shape()) {	2✔
264	j["input_shape"].push_back(serializer.expression(dim));	2✔
265	}	2✔
266	j["index_shape"] = nlohmann::json::array();	1✔
267	for (auto& dim : index_node.index_shape()) {	1✔
268	j["index_shape"].push_back(serializer.expression(dim));	1✔
269	}	1✔
270
271	j["dim_offset"] = index_node.dim_offset();	1✔
272	j["num_indices"] = index_node.num_indices();	1✔
273
274	return j;	1✔
275	}	1✔
276
277	data_flow::LibraryNode& IndexNodeSerializer::deserialize(
278	const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
279	) {	1✔
280	assert(j.contains("element_id"));	1✔
281	assert(j.contains("code"));	1✔
282	assert(j.contains("debug_info"));	1✔
283	assert(j.contains("input_shape"));	1✔
284	assert(j.contains("index_shape"));	1✔
285	assert(j.contains("dim_offset"));	1✔
286	assert(j.contains("num_indices"));	1✔
287
288	std::vector<symbolic::Expression> input_shape;	1✔
289	for (const auto& dim : j["input_shape"]) {	2✔
290	input_shape.push_back(symbolic::parse(dim.get<std::string>()));	2✔
291	}	2✔
292	std::vector<symbolic::Expression> index_shape;	1✔
293	for (const auto& dim : j["index_shape"]) {	1✔
294	index_shape.push_back(symbolic::parse(dim.get<std::string>()));	1✔
295	}	1✔
296
297	long long dim_offset = j["dim_offset"].get<long long>();	1✔
298	long long num_indices = j["num_indices"].get<long long>();	1✔
299
300	sdfg::serializer::JSONSerializer serializer;	1✔
301	DebugInfo debug_info = serializer.json_to_debug_info(j["debug_info"]);	1✔
302
303	return builder.add_library_node<IndexNode>(parent, debug_info, input_shape, index_shape, dim_offset, num_indices);	1✔
304	}	1✔
305
306	} // namespace tensor
307	} // namespace math
308	} // namespace sdfg

daisytuner / docc / 30527237467

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