30377441042

Committed 28 Jul 2026 04:16PM UTC coverage: 64.171% (-0.2%) from 64.324%

Build # 30377441042

Build Type

Pull #900

github

Committed by

web-flow

Commit Message

Merge c3fd195b7 into 9ecbee2f9

Pull Request Pull Request #900: [PyTorch] Add support for aten.slice

Coverage Stats

39 of 222 new or added lines in 3 files covered. (17.57%)

43457 of 67721 relevant lines covered (64.17%)

723.43 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/slice_node.cpp

1	#include "sdfg/data_flow/library_nodes/math/tensor/slice_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	SliceNode::SliceNode(
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	long long dim,
16	long long start,
17	long long end,
18	long long step,
19	const data_flow::ImplementationType& impl_type
20	)
NEW 21	: TensorNode(element_id, debug_info, vertex, parent, LibraryNodeType_Slice, {}, {"Y", "X"}, impl_type),	×
NEW 22	input_shape_(input_shape), dim_(dim), start_(start), end_(end), step_(step) {}	×
23
NEW 24	const std::vector<symbolic::Expression>& SliceNode::input_shape() const { return input_shape_; }	×
25
NEW 26	long long SliceNode::dim() const { return dim_; }	×
27
NEW 28	long long SliceNode::start() const { return start_; }	×
29
NEW 30	long long SliceNode::end() const { return end_; }	×
31
NEW 32	long long SliceNode::step() const { return step_; }	×
33
NEW 34	bool SliceNode::supports_integer_types() const { return true; }	×
35
NEW 36	void SliceNode::validate(const Function& function) const {	×
NEW 37	TensorNode::validate(function);	×
38
NEW 39	if (dim_ < 0 \|\| static_cast<size_t>(dim_) >= input_shape_.size()) {	×
NEW 40	throw InvalidSDFGException(	×
NEW 41	"SliceNode: dim out of range. dim: " + std::to_string(dim_) +	×
NEW 42	" rank: " + std::to_string(input_shape_.size())	×
NEW 43	);	×
NEW 44	}	×
NEW 45	if (step_ <= 0) {	×
NEW 46	throw InvalidSDFGException("SliceNode: step must be positive but got " + std::to_string(step_));	×
NEW 47	}	×
NEW 48	if (start_ < 0 \|\| end_ < start_) {	×
NEW 49	throw InvalidSDFGException(	×
NEW 50	"SliceNode: expected 0 <= start <= end but got start: " + std::to_string(start_) +	×
NEW 51	" end: " + std::to_string(end_)	×
NEW 52	);	×
NEW 53	}	×
NEW 54	}	×
55
NEW 56	symbolic::SymbolSet SliceNode::symbols() const {	×
NEW 57	symbolic::SymbolSet syms;	×
NEW 58	for (const auto& dim : input_shape_) {	×
NEW 59	for (auto& atom : symbolic::atoms(dim)) {	×
NEW 60	syms.insert(atom);	×
NEW 61	}	×
NEW 62	}	×
NEW 63	return syms;	×
NEW 64	}	×
65
NEW 66	void SliceNode::replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {	×
NEW 67	for (auto& dim : input_shape_) {	×
NEW 68	dim = symbolic::subs(dim, old_expression, new_expression);	×
NEW 69	}	×
NEW 70	}	×
71
NEW 72	void SliceNode::replace(const symbolic::ExpressionMapping& replacements) {	×
NEW 73	for (auto& dim : input_shape_) {	×
NEW 74	dim = symbolic::subs(dim, replacements);	×
NEW 75	}	×
NEW 76	}	×
77
78	passes::LibNodeExpander::ExpandOutcome SliceNode::
NEW 79	expand(passes::LibNodeExpander::ExpandContext& context, structured_control_flow::Block& block) {	×
NEW 80	auto& dataflow = this->get_parent();	×
81
NEW 82	if (dataflow.in_degree(this) != 2 \|\| dataflow.out_degree(this) != 0) {	×
NEW 83	return context.unable();	×
NEW 84	}	×
85
NEW 86	auto edges = dataflow.in_edges_by_connector(*this);	×
NEW 87	auto& in_edge = *edges.at(X_INPUT_IDX);	×
NEW 88	auto& result_ptr_edge = *edges.at(RESULT_PTR_IDX);	×
89
NEW 90	using Use = passes::LibNodeExpander::InputUse;	×
NEW 91	auto standalone = context.replacement_requires_access_nodes({Use::IndirectWrite, Use::IndirectRead});	×
92
NEW 93	if (!standalone) {	×
NEW 94	return context.unable();	×
NEW 95	}	×
96
97	// Output shape equals the input shape except along the sliced dimension.
NEW 98	std::vector<symbolic::Expression> output_shape = input_shape_;	×
NEW 99	long long sliced_dim = (end_ - start_ + step_ - 1) / step_;	×
NEW 100	output_shape[dim_] = symbolic::integer(sliced_dim);	×
101
NEW 102	symbolic::MultiExpression loop_vars;	×
NEW 103	auto& builder = standalone->builder();	×
NEW 104	structured_control_flow::Sequence* inner_scope = nullptr;	×
105
NEW 106	for (size_t i = 0; i < output_shape.size(); ++i) {	×
NEW 107	std::string var_name = builder.find_new_name("_i" + std::to_string(i));	×
NEW 108	builder.add_container(var_name, types::Scalar(types::PrimitiveType::Int64));	×
109
NEW 110	auto sym_var = symbolic::symbol(var_name);	×
NEW 111	auto condition = symbolic::Lt(sym_var, output_shape[i]);	×
NEW 112	auto init = symbolic::zero();	×
NEW 113	auto update = symbolic::add(sym_var, symbolic::one());	×
114
NEW 115	if (i == 0) {	×
NEW 116	auto& loop = standalone->replace_with_structured_loop(	×
NEW 117	passes::LibNodeExpander::AccessNodeExpand::LoopType::Map,	×
NEW 118	sym_var,	×
NEW 119	condition,	×
NEW 120	init,	×
NEW 121	update,	×
NEW 122	structured_control_flow::ScheduleType_Sequential::create()	×
NEW 123	);	×
NEW 124	inner_scope = &loop.root();	×
NEW 125	} else {	×
NEW 126	auto& loop = builder.add_map(	×
NEW 127	*inner_scope,	×
NEW 128	sym_var,	×
NEW 129	condition,	×
NEW 130	init,	×
NEW 131	update,	×
NEW 132	structured_control_flow::ScheduleType_Sequential::create(),	×
NEW 133	this->debug_info()	×
NEW 134	);	×
NEW 135	inner_scope = &loop.root();	×
NEW 136	}	×
NEW 137	loop_vars.push_back(sym_var);	×
NEW 138	}	×
139
NEW 140	auto& tasklet_block = builder.add_block(*inner_scope, {}, this->debug_info());	×
141
NEW 142	auto& in_acc = standalone->add_indirect_read_access(tasklet_block, X_INPUT_IDX);	×
NEW 143	auto& out_acc = standalone->add_indirect_write_access(tasklet_block, RESULT_PTR_IDX);	×
144
145	// Source subset: `start + _i * step` along the sliced dimension, identity otherwise.
NEW 146	symbolic::MultiExpression input_subset;	×
NEW 147	input_subset.reserve(output_shape.size());	×
NEW 148	for (size_t i = 0; i < output_shape.size(); ++i) {	×
NEW 149	if (static_cast<long long>(i) == dim_) {	×
NEW 150	input_subset	×
NEW 151	.push_back(symbolic::add(symbolic::integer(start_), symbolic::mul(loop_vars[i], symbolic::integer(step_)))	×
NEW 152	);	×
NEW 153	} else {	×
NEW 154	input_subset.push_back(loop_vars[i]);	×
NEW 155	}	×
NEW 156	}	×
157
NEW 158	auto& tasklet =	×
NEW 159	builder.add_tasklet(tasklet_block, data_flow::TaskletCode::assign, "_out", {"_in"}, this->debug_info());	×
160
NEW 161	builder.add_computational_memlet(	×
NEW 162	tasklet_block, in_acc, tasklet, "_in", input_subset, in_edge.base_type(), this->debug_info()	×
NEW 163	);	×
NEW 164	builder.add_computational_memlet(	×
NEW 165	tasklet_block, tasklet, "_out", out_acc, loop_vars, result_ptr_edge.base_type(), this->debug_info()	×
NEW 166	);	×
167
NEW 168	return standalone->successfully_expanded();	×
NEW 169	}	×
170
171	std::unique_ptr<data_flow::DataFlowNode> SliceNode::
NEW 172	clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {	×
NEW 173	return std::unique_ptr<data_flow::DataFlowNode>(	×
NEW 174	new SliceNode(element_id, this->debug_info(), vertex, parent, input_shape_, dim_, start_, end_, step_)	×
NEW 175	);	×
NEW 176	}	×
177
NEW 178	data_flow::PointerAccessType SliceNode::pointer_access_type(int input_idx) const {	×
NEW 179	if (input_idx == RESULT_PTR_IDX) {	×
NEW 180	return data_flow::PointerAccessMeta::create_full_write_only(symbolic::__nullptr__(), true);	×
NEW 181	} else if (input_idx == X_INPUT_IDX) {	×
NEW 182	return data_flow::PointerAccessMeta::create_read_only(symbolic::__nullptr__(), true);	×
NEW 183	} else {	×
NEW 184	return TensorNode::pointer_access_type(input_idx);	×
NEW 185	}	×
NEW 186	}	×
187
NEW 188	nlohmann::json SliceNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {	×
NEW 189	const SliceNode& slice_node = static_cast<const SliceNode&>(library_node);	×
NEW 190	nlohmann::json j;	×
191
NEW 192	j["code"] = slice_node.code().value();	×
193
NEW 194	serializer::JSONSerializer serializer;	×
NEW 195	j["input_shape"] = nlohmann::json::array();	×
NEW 196	for (auto& dim : slice_node.input_shape()) {	×
NEW 197	j["input_shape"].push_back(serializer.expression(dim));	×
NEW 198	}	×
199
NEW 200	j["dim"] = slice_node.dim();	×
NEW 201	j["start"] = slice_node.start();	×
NEW 202	j["end"] = slice_node.end();	×
NEW 203	j["step"] = slice_node.step();	×
204
NEW 205	return j;	×
NEW 206	}	×
207
208	data_flow::LibraryNode& SliceNodeSerializer::deserialize(
209	const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
NEW 210	) {	×
211	// Assertions for required fields
NEW 212	assert(j.contains("element_id"));	×
NEW 213	assert(j.contains("code"));	×
NEW 214	assert(j.contains("debug_info"));	×
NEW 215	assert(j.contains("input_shape"));	×
NEW 216	assert(j.contains("dim"));	×
NEW 217	assert(j.contains("start"));	×
NEW 218	assert(j.contains("end"));	×
NEW 219	assert(j.contains("step"));	×
220
NEW 221	std::vector<symbolic::Expression> input_shape;	×
NEW 222	for (const auto& dim : j["input_shape"]) {	×
NEW 223	input_shape.push_back(symbolic::parse(dim.get<std::string>()));	×
NEW 224	}	×
225
NEW 226	long long dim = j["dim"].get<long long>();	×
NEW 227	long long start = j["start"].get<long long>();	×
NEW 228	long long end = j["end"].get<long long>();	×
NEW 229	long long step = j["step"].get<long long>();	×
230
231	// Extract debug info using JSONSerializer
NEW 232	sdfg::serializer::JSONSerializer serializer;	×
NEW 233	DebugInfo debug_info = serializer.json_to_debug_info(j["debug_info"]);	×
234
NEW 235	return builder.add_library_node<SliceNode>(parent, debug_info, input_shape, dim, start, end, step);	×
NEW 236	}	×
237
238	} // namespace tensor
239	} // namespace math
240	} // namespace sdfg

daisytuner / docc / 30377441042

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