22020750556

Committed 14 Feb 2026 04:38PM UTC coverage: 64.828% (-1.5%) from 66.315%

Build # 22020750556

Build Type

Pull #524

github

Committed by

web-flow

Commit Message

Merge 2784aa264 into 9d01cacd5

Pull Request Pull Request #524: Native Tensor Support - Step 2: Use tensor types on memlets of tensor nodes

Run Details

245 of 570 new or added lines in 24 files covered. (42.98%)

458 existing lines in 18 files now uncovered.

23080 of 35602 relevant lines covered (64.83%)

371.57 hits per line

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

79.77

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

#include "sdfg/data_flow/library_nodes/math/tensor/reduce_node.h"

#include "sdfg/analysis/analysis.h"
#include "sdfg/analysis/scope_analysis.h"
#include "sdfg/builder/structured_sdfg_builder.h"

#include <algorithm>

namespace sdfg {
namespace math {
namespace tensor {

ReduceNode::ReduceNode(
    size_t element_id,
    const DebugInfo& debug_info,
    const graph::Vertex vertex,
    data_flow::DataFlowGraph& parent,
    const data_flow::LibraryNodeCode& code,
    const std::vector<symbolic::Expression>& shape,
    const std::vector<int64_t>& axes,
    bool keepdims
)
    : TensorNode(element_id, debug_info, vertex, parent, code, {"Y"}, {"X"}, data_flow::ImplementationType_NONE),
      shape_(shape), axes_(axes), keepdims_(keepdims) {}

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

    auto& graph = this->get_parent();

    auto& iedge = *graph.in_edges(*this).begin();
    auto& tensor_input = static_cast<const types::Tensor&>(iedge.base_type());
    if (tensor_input.shape().size() != this->shape_.size()) {
        throw InvalidSDFGException(
            "Library Node: Input tensor shape must match node shape. Input shape: " +
            std::to_string(tensor_input.shape().size()) + " Node shape: " + std::to_string(this->shape_.size())
        );
    }
    for (size_t i = 0; i < shape_.size(); ++i) {
        if (!symbolic::eq(tensor_input.shape().at(i), shape_.at(i))) {
            throw InvalidSDFGException(
                "Library Node: Input tensor shape must match node shape. Input shape at dim " + std::to_string(i) +
                ": " + tensor_input.shape().at(i)->__str__() + " Node shape at dim " + std::to_string(i) + ": " +
                shape_.at(i)->__str__()
            );
        }
    }

    // Calculate expected output shape based on axes and keepdims
    std::vector<int64_t> sorted_axes = axes_;
    // Normalize negative axes
    for (auto& axis : sorted_axes) {
        if (axis < 0) {
            axis = static_cast<int64_t>(shape_.size()) + axis;
        }
        // Validate axis is in bounds
        if (axis < 0 || axis >= static_cast<int64_t>(shape_.size())) {
            throw InvalidSDFGException(
                "Library Node: Axis value out of bounds. Axis: " + std::to_string(axis) +
                " Shape size: " + std::to_string(shape_.size())
            );
        }
    }
    std::sort(sorted_axes.begin(), sorted_axes.end());

    std::vector<symbolic::Expression> expected_output_shape;
    for (size_t i = 0; i < shape_.size(); ++i) {
        bool is_axis = false;
        for (auto axis : sorted_axes) {
            if (axis == (int64_t) i) {
                is_axis = true;
                break;
            }
        }

        if (is_axis) {
            if (keepdims_) {
                expected_output_shape.push_back(symbolic::one());
            }
        } else {
            expected_output_shape.push_back(shape_[i]);
        }
    }

    auto& oedge = *graph.out_edges(*this).begin();
    auto& tensor_output = static_cast<const types::Tensor&>(oedge.base_type());
    if (tensor_output.shape().size() != expected_output_shape.size()) {
        throw InvalidSDFGException(
            "Library Node: Output tensor shape must match expected reduced shape. Output shape size: " +
            std::to_string(tensor_output.shape().size()) +
            " Expected shape size: " + std::to_string(expected_output_shape.size())
        );
    }
    for (size_t i = 0; i < expected_output_shape.size(); ++i) {
        if (!symbolic::eq(tensor_output.shape().at(i), expected_output_shape.at(i))) {
            throw InvalidSDFGException(
                "Library Node: Output tensor shape must match expected reduced shape. Output shape at dim " +
                std::to_string(i) + ": " + tensor_output.shape().at(i)->__str__() + " Expected shape at dim " +
                std::to_string(i) + ": " + expected_output_shape.at(i)->__str__()
            );
        }
    }
}


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

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

bool ReduceNode::expand(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {
    auto& dataflow = this->get_parent();
    auto& block = static_cast<structured_control_flow::Block&>(*dataflow.get_parent());

    if (dataflow.in_degree(*this) != 1 || dataflow.out_degree(*this) != 1) {
        return false;
    }

    auto& scope_analysis = analysis_manager.get<analysis::ScopeAnalysis>();
    auto& parent = static_cast<structured_control_flow::Sequence&>(*scope_analysis.parent_scope(&block));
    int index = parent.index(block);
    auto& transition = parent.at(index).second;

    auto& iedge = *dataflow.in_edges(*this).begin();
    auto& oedge = *dataflow.out_edges(*this).begin();

    auto& input_node = static_cast<data_flow::AccessNode&>(iedge.src());
    auto& output_node = static_cast<data_flow::AccessNode&>(oedge.dst());

    if (dataflow.in_degree(input_node) != 0 || dataflow.out_degree(output_node) != 0) {
        return false;
    }

    // Calculate output shape
    std::vector<symbolic::Expression> output_shape;
    std::vector<int64_t> sorted_axes = axes_;
    // Normalize negative axes
    for (auto& axis : sorted_axes) {
        if (axis < 0) {
            axis = static_cast<int64_t>(shape_.size()) + axis;
        }
        // Validate axis is in bounds
        if (axis < 0 || axis >= static_cast<int64_t>(shape_.size())) {
            throw InvalidSDFGException(
                "Library Node: Axis value out of bounds. Axis: " + std::to_string(axis) +
                " Shape size: " + std::to_string(shape_.size())
            );
        }
    }
    std::sort(sorted_axes.begin(), sorted_axes.end());

    for (size_t i = 0; i < shape_.size(); ++i) {
        bool is_axis = false;
        for (auto axis : sorted_axes) {
            if (axis == (int64_t) i) {
                is_axis = true;
                break;
            }
        }

        if (is_axis) {
            if (keepdims_) {
                output_shape.push_back(symbolic::one());
            }
        } else {
            output_shape.push_back(shape_[i]);
        }
    }

    sdfg::types::Scalar element_type(oedge.base_type().primitive_type());
    types::Tensor scalar_tensor(element_type.primitive_type(), {});

    // Add new sequence
    auto& new_sequence = builder.add_sequence_before(parent, block, transition.assignments(), block.debug_info());

    // 1. Initialization Loop
    {
        data_flow::Subset init_subset;
        structured_control_flow::Sequence* last_scope = &new_sequence;
        structured_control_flow::Map* last_map = nullptr;

        for (size_t i = 0; i < output_shape.size(); ++i) {
            std::string indvar_str = builder.find_new_name("_i_init");
            builder.add_container(indvar_str, types::Scalar(types::PrimitiveType::Int64));

            auto indvar = symbolic::symbol(indvar_str);
            auto init = symbolic::zero();
            auto update = symbolic::add(indvar, symbolic::one());
            auto condition = symbolic::Lt(indvar, output_shape[i]);

            last_map = &builder.add_map(
                *last_scope,
                indvar,
                condition,
                init,
                update,
                structured_control_flow::ScheduleType_Sequential::create(),
                {},
                block.debug_info()
            );
            last_scope = &last_map->root();
            init_subset.push_back(indvar);
        }

        // Add initialization tasklet
        auto& init_block = builder.add_block(*last_scope, {}, block.debug_info());
        auto& init_tasklet =
            builder.add_tasklet(init_block, data_flow::TaskletCode::assign, {"_out"}, {"_in"}, block.debug_info());

        auto& const_node = builder.add_constant(init_block, this->identity(), element_type, block.debug_info());
        auto& out_access = builder.add_access(init_block, output_node.data(), block.debug_info());

        builder
            .add_computational_memlet(init_block, const_node, init_tasklet, "_in", {}, scalar_tensor, block.debug_info());
        builder.add_computational_memlet(
            init_block, init_tasklet, "_out", out_access, init_subset, oedge.base_type(), block.debug_info()
        );
    }

    // 2. Reduction Loop
    {
        data_flow::Subset input_subset;
        data_flow::Subset output_subset;

        structured_control_flow::Sequence* last_scope = &new_sequence;
        structured_control_flow::StructuredLoop* last_loop = nullptr;

        std::map<size_t, symbolic::Expression> loop_vars_map;
        std::vector<size_t> outer_dims;
        std::vector<size_t> inner_dims;

        // Partition dimensions
        for (size_t i = 0; i < shape_.size(); ++i) {
            bool is_axis = false;
            for (auto axis : sorted_axes) {
                if (axis == (int64_t) i) {
                    is_axis = true;
                    break;
                }
            }
            if (is_axis) {
                inner_dims.push_back(i);
            } else {
                outer_dims.push_back(i);
            }
        }

        // Generate outer parallel loops (Maps)
        for (size_t dim_idx : outer_dims) {
            std::string indvar_str = builder.find_new_name("_i");
            builder.add_container(indvar_str, types::Scalar(types::PrimitiveType::Int64));

            auto indvar = symbolic::symbol(indvar_str);
            auto init = symbolic::zero();
            auto update = symbolic::add(indvar, symbolic::one());
            auto condition = symbolic::Lt(indvar, shape_[dim_idx]);

            auto& map = builder.add_map(
                *last_scope,
                indvar,
                condition,
                init,
                update,
                structured_control_flow::ScheduleType_Sequential::create(),
                {},
                block.debug_info()
            );
            last_scope = &map.root();
            loop_vars_map[dim_idx] = indvar;
        }

        // Generate inner sequential loops (Fors)
        for (size_t dim_idx : inner_dims) {
            std::string indvar_str = builder.find_new_name("_k");
            builder.add_container(indvar_str, types::Scalar(types::PrimitiveType::Int64));

            auto indvar = symbolic::symbol(indvar_str);
            auto init = symbolic::zero();
            auto update = symbolic::add(indvar, symbolic::one());
            auto condition = symbolic::Lt(indvar, shape_[dim_idx]);

            last_loop = &builder.add_for(*last_scope, indvar, condition, init, update, {}, block.debug_info());
            last_scope = &last_loop->root();
            loop_vars_map[dim_idx] = indvar;
        }

        // Construct output indices
        std::vector<symbolic::Expression> input_indices;
        std::vector<symbolic::Expression> output_indices;
        for (size_t i = 0; i < shape_.size(); ++i) {
            input_indices.push_back(loop_vars_map[i]);
            bool is_axis = false;
            for (auto axis : sorted_axes) {
                if (axis == (int64_t) i) {
                    is_axis = true;
                    break;
                }
            }

            if (is_axis) {
                if (keepdims_) {
                    output_indices.push_back(symbolic::zero());
                }
            } else {
                output_indices.push_back(loop_vars_map[i]);
            }
        }

        this->expand_reduction(
            builder,
            analysis_manager,
            *last_scope,
            input_node.data(),
            output_node.data(),
            static_cast<const types::Tensor&>(iedge.base_type()),
            static_cast<const types::Tensor&>(oedge.base_type()),
            input_indices,
            output_indices
        );
    }

    // Clean up block
    builder.remove_memlet(block, iedge);
    builder.remove_memlet(block, oedge);
    builder.remove_node(block, input_node);
    builder.remove_node(block, output_node);
    builder.remove_node(block, *this);
    builder.remove_child(parent, index + 1);


    return true;
}

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

1	#include "sdfg/data_flow/library_nodes/math/tensor/reduce_node.h"
2
3	#include "sdfg/analysis/analysis.h"
4	#include "sdfg/analysis/scope_analysis.h"
5	#include "sdfg/builder/structured_sdfg_builder.h"
6
7	#include <algorithm>
8
9	namespace sdfg {
10	namespace math {
11	namespace tensor {
12
13	ReduceNode::ReduceNode(
14	size_t element_id,
15	const DebugInfo& debug_info,
16	const graph::Vertex vertex,
17	data_flow::DataFlowGraph& parent,
18	const data_flow::LibraryNodeCode& code,
19	const std::vector<symbolic::Expression>& shape,
20	const std::vector<int64_t>& axes,
21	bool keepdims
22	)
23	: TensorNode(element_id, debug_info, vertex, parent, code, {"Y"}, {"X"}, data_flow::ImplementationType_NONE),	18✔
24	shape_(shape), axes_(axes), keepdims_(keepdims) {}	18✔
25
26	void ReduceNode::validate(const Function& function) const {	12✔
27	TensorNode::validate(function);	12✔
28
29	auto& graph = this->get_parent();	12✔
30
31	auto& iedge = graph.in_edges(this).begin();	12✔
32	auto& tensor_input = static_cast<const types::Tensor&>(iedge.base_type());	12✔
33	if (tensor_input.shape().size() != this->shape_.size()) {	12✔
NEW 34	throw InvalidSDFGException(	×
NEW 35	"Library Node: Input tensor shape must match node shape. Input shape: " +	×
NEW 36	std::to_string(tensor_input.shape().size()) + " Node shape: " + std::to_string(this->shape_.size())	×
NEW 37	);	×
NEW 38	}	×
39	for (size_t i = 0; i < shape_.size(); ++i) {	32✔
40	if (!symbolic::eq(tensor_input.shape().at(i), shape_.at(i))) {	20✔
NEW 41	throw InvalidSDFGException(	×
NEW 42	"Library Node: Input tensor shape must match node shape. Input shape at dim " + std::to_string(i) +	×
NEW 43	": " + tensor_input.shape().at(i)->__str__() + " Node shape at dim " + std::to_string(i) + ": " +	×
NEW 44	shape_.at(i)->__str__()	×
NEW 45	);	×
NEW 46	}	×
47	}	20✔
48
49	// Calculate expected output shape based on axes and keepdims
50	std::vector<int64_t> sorted_axes = axes_;	12✔
51	// Normalize negative axes
52	for (auto& axis : sorted_axes) {	13✔
53	if (axis < 0) {	13✔
54	axis = static_cast<int64_t>(shape_.size()) + axis;	2✔
55	}	2✔
56	// Validate axis is in bounds
57	if (axis < 0 \|\| axis >= static_cast<int64_t>(shape_.size())) {	13✔
NEW 58	throw InvalidSDFGException(	×
NEW 59	"Library Node: Axis value out of bounds. Axis: " + std::to_string(axis) +	×
NEW 60	" Shape size: " + std::to_string(shape_.size())	×
NEW 61	);	×
NEW 62	}	×
63	}	13✔
64	std::sort(sorted_axes.begin(), sorted_axes.end());	12✔
65
66	std::vector<symbolic::Expression> expected_output_shape;	12✔
67	for (size_t i = 0; i < shape_.size(); ++i) {	32✔
68	bool is_axis = false;	20✔
69	for (auto axis : sorted_axes) {	22✔
70	if (axis == (int64_t) i) {	22✔
71	is_axis = true;	13✔
72	break;	13✔
73	}	13✔
74	}	22✔
75
76	if (is_axis) {	20✔
77	if (keepdims_) {	13✔
78	expected_output_shape.push_back(symbolic::one());	1✔
79	}	1✔
80	} else {	13✔
81	expected_output_shape.push_back(shape_[i]);	7✔
82	}	7✔
83	}	20✔
84
85	auto& oedge = graph.out_edges(this).begin();	12✔
86	auto& tensor_output = static_cast<const types::Tensor&>(oedge.base_type());	12✔
87	if (tensor_output.shape().size() != expected_output_shape.size()) {	12✔
NEW 88	throw InvalidSDFGException(	×
NEW 89	"Library Node: Output tensor shape must match expected reduced shape. Output shape size: " +	×
NEW 90	std::to_string(tensor_output.shape().size()) +	×
NEW 91	" Expected shape size: " + std::to_string(expected_output_shape.size())	×
NEW 92	);	×
NEW 93	}	×
94	for (size_t i = 0; i < expected_output_shape.size(); ++i) {	20✔
95	if (!symbolic::eq(tensor_output.shape().at(i), expected_output_shape.at(i))) {	8✔
NEW 96	throw InvalidSDFGException(	×
NEW 97	"Library Node: Output tensor shape must match expected reduced shape. Output shape at dim " +	×
NEW 98	std::to_string(i) + ": " + tensor_output.shape().at(i)->__str__() + " Expected shape at dim " +	×
NEW 99	std::to_string(i) + ": " + expected_output_shape.at(i)->__str__()	×
NEW 100	);	×
NEW 101	}	×
102	}	8✔
103	}	12✔
104
105
106	symbolic::SymbolSet ReduceNode::symbols() const {	×
107	symbolic::SymbolSet syms;	×
108	for (const auto& dim : shape_) {	×
109	for (auto& atom : symbolic::atoms(dim)) {	×
110	syms.insert(atom);	×
111	}	×
112	}	×
113	return syms;	×
114	}	×
115
116	void ReduceNode::replace(const symbolic::Expression old_expression, const symbolic::Expression new_expression) {	×
117	for (auto& dim : shape_) {	×
118	dim = symbolic::subs(dim, old_expression, new_expression);	×
119	}	×
120	}	×
121
122	bool ReduceNode::expand(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {	6✔
123	auto& dataflow = this->get_parent();	6✔
124	auto& block = static_cast<structured_control_flow::Block&>(*dataflow.get_parent());	6✔
125
126	if (dataflow.in_degree(this) != 1 \|\| dataflow.out_degree(this) != 1) {	6✔
127	return false;	×
128	}	×
129
130	auto& scope_analysis = analysis_manager.get<analysis::ScopeAnalysis>();	6✔
131	auto& parent = static_cast<structured_control_flow::Sequence&>(*scope_analysis.parent_scope(&block));	6✔
132	int index = parent.index(block);	6✔
133	auto& transition = parent.at(index).second;	6✔
134
135	auto& iedge = dataflow.in_edges(this).begin();	6✔
136	auto& oedge = dataflow.out_edges(this).begin();	6✔
137
138	auto& input_node = static_cast<data_flow::AccessNode&>(iedge.src());	6✔
139	auto& output_node = static_cast<data_flow::AccessNode&>(oedge.dst());	6✔
140
141	if (dataflow.in_degree(input_node) != 0 \|\| dataflow.out_degree(output_node) != 0) {	6✔
142	return false;	×
143	}	×
144
145	// Calculate output shape
146	std::vector<symbolic::Expression> output_shape;	6✔
147	std::vector<int64_t> sorted_axes = axes_;	6✔
148	// Normalize negative axes
149	for (auto& axis : sorted_axes) {	7✔
150	if (axis < 0) {	7✔
NEW 151	axis = static_cast<int64_t>(shape_.size()) + axis;	×
NEW 152	}	×
153	// Validate axis is in bounds
154	if (axis < 0 \|\| axis >= static_cast<int64_t>(shape_.size())) {	7✔
NEW 155	throw InvalidSDFGException(	×
NEW 156	"Library Node: Axis value out of bounds. Axis: " + std::to_string(axis) +	×
NEW 157	" Shape size: " + std::to_string(shape_.size())	×
NEW 158	);	×
NEW 159	}	×
160	}	7✔
161	std::sort(sorted_axes.begin(), sorted_axes.end());	6✔
162
163	for (size_t i = 0; i < shape_.size(); ++i) {	18✔
164	bool is_axis = false;	12✔
165	for (auto axis : sorted_axes) {	14✔
166	if (axis == (int64_t) i) {	14✔
167	is_axis = true;	7✔
168	break;	7✔
169	}	7✔
170	}	14✔
171
172	if (is_axis) {	12✔
173	if (keepdims_) {	7✔
174	output_shape.push_back(symbolic::one());	1✔
175	}	1✔
176	} else {	7✔
177	output_shape.push_back(shape_[i]);	5✔
178	}	5✔
179	}	12✔
180
181	sdfg::types::Scalar element_type(oedge.base_type().primitive_type());	6✔
182	types::Tensor scalar_tensor(element_type.primitive_type(), {});	6✔
183
184	// Add new sequence
185	auto& new_sequence = builder.add_sequence_before(parent, block, transition.assignments(), block.debug_info());	6✔
186
187	// 1. Initialization Loop
188	{	6✔
189	data_flow::Subset init_subset;	6✔
190	structured_control_flow::Sequence* last_scope = &new_sequence;	6✔
191	structured_control_flow::Map* last_map = nullptr;	6✔
192
193	for (size_t i = 0; i < output_shape.size(); ++i) {	12✔
194	std::string indvar_str = builder.find_new_name("_i_init");	6✔
195	builder.add_container(indvar_str, types::Scalar(types::PrimitiveType::Int64));	6✔
196
197	auto indvar = symbolic::symbol(indvar_str);	6✔
198	auto init = symbolic::zero();	6✔
199	auto update = symbolic::add(indvar, symbolic::one());	6✔
200	auto condition = symbolic::Lt(indvar, output_shape[i]);	6✔
201
202	last_map = &builder.add_map(	6✔
203	*last_scope,	6✔
204	indvar,	6✔
205	condition,	6✔
206	init,	6✔
207	update,	6✔
208	structured_control_flow::ScheduleType_Sequential::create(),	6✔
209	{},	6✔
210	block.debug_info()	6✔
211	);	6✔
212	last_scope = &last_map->root();	6✔
213	init_subset.push_back(indvar);	6✔
214	}	6✔
215
216	// Add initialization tasklet
217	auto& init_block = builder.add_block(*last_scope, {}, block.debug_info());	6✔
218	auto& init_tasklet =	6✔
219	builder.add_tasklet(init_block, data_flow::TaskletCode::assign, {"_out"}, {"_in"}, block.debug_info());	6✔
220
221	auto& const_node = builder.add_constant(init_block, this->identity(), element_type, block.debug_info());	6✔
222	auto& out_access = builder.add_access(init_block, output_node.data(), block.debug_info());	6✔
223
224	builder	6✔
225	.add_computational_memlet(init_block, const_node, init_tasklet, "_in", {}, scalar_tensor, block.debug_info());	6✔
226	builder.add_computational_memlet(	6✔
227	init_block, init_tasklet, "_out", out_access, init_subset, oedge.base_type(), block.debug_info()	6✔
228	);	6✔
229	}	6✔
230
231	// 2. Reduction Loop
232	{	6✔
233	data_flow::Subset input_subset;	6✔
234	data_flow::Subset output_subset;	6✔
235
236	structured_control_flow::Sequence* last_scope = &new_sequence;	6✔
237	structured_control_flow::StructuredLoop* last_loop = nullptr;	6✔
238
239	std::map<size_t, symbolic::Expression> loop_vars_map;	6✔
240	std::vector<size_t> outer_dims;	6✔
241	std::vector<size_t> inner_dims;	6✔
242
243	// Partition dimensions
244	for (size_t i = 0; i < shape_.size(); ++i) {	18✔
245	bool is_axis = false;	12✔
246	for (auto axis : sorted_axes) {	14✔
247	if (axis == (int64_t) i) {	14✔
248	is_axis = true;	7✔
249	break;	7✔
250	}	7✔
251	}	14✔
252	if (is_axis) {	12✔
253	inner_dims.push_back(i);	7✔
254	} else {	7✔
255	outer_dims.push_back(i);	5✔
256	}	5✔
257	}	12✔
258
259	// Generate outer parallel loops (Maps)
260	for (size_t dim_idx : outer_dims) {	6✔
261	std::string indvar_str = builder.find_new_name("_i");	5✔
262	builder.add_container(indvar_str, types::Scalar(types::PrimitiveType::Int64));	5✔
263
264	auto indvar = symbolic::symbol(indvar_str);	5✔
265	auto init = symbolic::zero();	5✔
266	auto update = symbolic::add(indvar, symbolic::one());	5✔
267	auto condition = symbolic::Lt(indvar, shape_[dim_idx]);	5✔
268
269	auto& map = builder.add_map(	5✔
270	*last_scope,	5✔
271	indvar,	5✔
272	condition,	5✔
273	init,	5✔
274	update,	5✔
275	structured_control_flow::ScheduleType_Sequential::create(),	5✔
276	{},	5✔
277	block.debug_info()	5✔
278	);	5✔
279	last_scope = &map.root();	5✔
280	loop_vars_map[dim_idx] = indvar;	5✔
281	}	5✔
282
283	// Generate inner sequential loops (Fors)
284	for (size_t dim_idx : inner_dims) {	7✔
285	std::string indvar_str = builder.find_new_name("_k");	7✔
286	builder.add_container(indvar_str, types::Scalar(types::PrimitiveType::Int64));	7✔
287
288	auto indvar = symbolic::symbol(indvar_str);	7✔
289	auto init = symbolic::zero();	7✔
290	auto update = symbolic::add(indvar, symbolic::one());	7✔
291	auto condition = symbolic::Lt(indvar, shape_[dim_idx]);	7✔
292
293	last_loop = &builder.add_for(*last_scope, indvar, condition, init, update, {}, block.debug_info());	7✔
294	last_scope = &last_loop->root();	7✔
295	loop_vars_map[dim_idx] = indvar;	7✔
296	}	7✔
297
298	// Construct output indices
299	std::vector<symbolic::Expression> input_indices;	6✔
300	std::vector<symbolic::Expression> output_indices;	6✔
301	for (size_t i = 0; i < shape_.size(); ++i) {	18✔
302	input_indices.push_back(loop_vars_map[i]);	12✔
303	bool is_axis = false;	12✔
304	for (auto axis : sorted_axes) {	14✔
305	if (axis == (int64_t) i) {	14✔
306	is_axis = true;	7✔
307	break;	7✔
308	}	7✔
309	}	14✔
310
311	if (is_axis) {	12✔
312	if (keepdims_) {	7✔
313	output_indices.push_back(symbolic::zero());	1✔
314	}	1✔
315	} else {	7✔
316	output_indices.push_back(loop_vars_map[i]);	5✔
317	}	5✔
318	}	12✔
319
320	this->expand_reduction(	6✔
321	builder,	6✔
322	analysis_manager,	6✔
323	*last_scope,	6✔
324	input_node.data(),	6✔
325	output_node.data(),	6✔
326	static_cast<const types::Tensor&>(iedge.base_type()),	6✔
327	static_cast<const types::Tensor&>(oedge.base_type()),	6✔
328	input_indices,	6✔
329	output_indices	6✔
330	);	6✔
331	}	6✔
332
333	// Clean up block
334	builder.remove_memlet(block, iedge);	6✔
335	builder.remove_memlet(block, oedge);	6✔
336	builder.remove_node(block, input_node);	6✔
337	builder.remove_node(block, output_node);	6✔
338	builder.remove_node(block, *this);	6✔
339	builder.remove_child(parent, index + 1);	6✔
340
341
342	return true;	6✔
343	}	6✔
344
345	} // namespace tensor
346	} // namespace math
347	} // namespace sdfg

daisytuner / docc / 22020750556

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