30434616029

Committed 29 Jul 2026 07:58AM UTC coverage: 64.62% (+0.3%) from 64.324%

Build # 30434616029

Build Type

Pull #900

github

Committed by

web-flow

Commit Message

Merge 9ffae1a3f into f1a0f70d9

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

Coverage Stats

191 of 222 new or added lines in 3 files covered. (86.04%)

176 existing lines in 5 files now uncovered.

44062 of 68186 relevant lines covered (64.62%)

719.5 hits per line

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

89.35

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

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

namespace sdfg {
namespace math {
namespace tensor {

SliceNode::SliceNode(
    size_t element_id,
    const DebugInfo& debug_info,
    const graph::Vertex vertex,
    data_flow::DataFlowGraph& parent,
    const std::vector<symbolic::Expression>& input_shape,
    long long dim,
    long long start,
    long long end,
    long long step,
    const data_flow::ImplementationType& impl_type
)
    : TensorNode(element_id, debug_info, vertex, parent, LibraryNodeType_Slice, {}, {"Y", "X"}, impl_type),
      input_shape_(input_shape), dim_(dim), start_(start), end_(end), step_(step) {}

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

long long SliceNode::dim() const { return dim_; }

long long SliceNode::start() const { return start_; }

long long SliceNode::end() const { return end_; }

long long SliceNode::step() const { return step_; }

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

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

    if (dim_ < 0 || static_cast<size_t>(dim_) >= input_shape_.size()) {
        throw InvalidSDFGException(
            "SliceNode: dim out of range. dim: " + std::to_string(dim_) +
            " rank: " + std::to_string(input_shape_.size())
        );
    }
    if (step_ <= 0) {
        throw InvalidSDFGException("SliceNode: step must be positive but got " + std::to_string(step_));
    }
    if (start_ < 0 || end_ < start_) {
        throw InvalidSDFGException(
            "SliceNode: expected 0 <= start <= end but got start: " + std::to_string(start_) +
            " end: " + std::to_string(end_)
        );
    }
}

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

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

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

passes::LibNodeExpander::ExpandOutcome SliceNode::
    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();
    }

    // Output shape equals the input shape except along the sliced dimension.
    std::vector<symbolic::Expression> output_shape = input_shape_;
    long long sliced_dim = (end_ - start_ + step_ - 1) / step_;
    output_shape[dim_] = symbolic::integer(sliced_dim);

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

    // Source subset: `start + _i * step` along the sliced dimension, identity otherwise.
    symbolic::MultiExpression input_subset;
    input_subset.reserve(output_shape.size());
    for (size_t i = 0; i < output_shape.size(); ++i) {
        if (static_cast<long long>(i) == dim_) {
            input_subset
                .push_back(symbolic::add(symbolic::integer(start_), symbolic::mul(loop_vars[i], symbolic::integer(step_)))
                );
        } else {
            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> SliceNode::
    clone(size_t element_id, const graph::Vertex vertex, data_flow::DataFlowGraph& parent) const {
    return std::unique_ptr<data_flow::DataFlowNode>(
        new SliceNode(element_id, this->debug_info(), vertex, parent, input_shape_, dim_, start_, end_, step_)
    );
}

data_flow::PointerAccessType SliceNode::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);
    }
}

nlohmann::json SliceNodeSerializer::serialize(const data_flow::LibraryNode& library_node) {
    const SliceNode& slice_node = static_cast<const SliceNode&>(library_node);
    nlohmann::json j;

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

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

    j["dim"] = slice_node.dim();
    j["start"] = slice_node.start();
    j["end"] = slice_node.end();
    j["step"] = slice_node.step();

    return j;
}

data_flow::LibraryNode& SliceNodeSerializer::deserialize(
    const nlohmann::json& j, builder::StructuredSDFGBuilder& builder, structured_control_flow::Block& parent
) {
    // Assertions for required fields
    assert(j.contains("element_id"));
    assert(j.contains("code"));
    assert(j.contains("debug_info"));
    assert(j.contains("input_shape"));
    assert(j.contains("dim"));
    assert(j.contains("start"));
    assert(j.contains("end"));
    assert(j.contains("step"));

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

    long long dim = j["dim"].get<long long>();
    long long start = j["start"].get<long long>();
    long long end = j["end"].get<long long>();
    long long step = j["step"].get<long long>();

    // Extract debug info using JSONSerializer
    sdfg::serializer::JSONSerializer serializer;
    DebugInfo debug_info = serializer.json_to_debug_info(j["debug_info"]);

    return builder.add_library_node<SliceNode>(parent, debug_info, input_shape, dim, start, end, step);
}

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

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

daisytuner / docc / 30434616029

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