28014415443

Committed 23 Jun 2026 08:54AM UTC coverage: 61.853% (+0.07%) from 61.779%

Build # 28014415443

Build Type

Pull #806

github

Committed by

web-flow

Commit Message

Merge 644dca1bc into 6ee760cce

Pull Request Pull Request #806: Map Collapse for Multiple targets in a neste sequence

Coverage Stats

129 of 137 new or added lines in 1 file covered. (94.16%)

10 existing lines in 2 files now uncovered.

37270 of 60256 relevant lines covered (61.85%)

1013.21 hits per line

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

93.97

/opt/src/transformations/map_collapse.cpp

#include "sdfg/transformations/map_collapse.h"

#include "sdfg/builder/structured_sdfg_builder.h"
#include "sdfg/structured_control_flow/if_else.h"
#include "sdfg/structured_control_flow/map.h"
#include "sdfg/structured_control_flow/sequence.h"
#include "sdfg/structured_control_flow/structured_loop.h"
#include "sdfg/symbolic/symbolic.h"

namespace sdfg {
namespace transformations {

MapCollapse::MapCollapse(structured_control_flow::Map& loop, size_t count) : loop_(loop), count_(count) {}

std::string MapCollapse::name() const { return "MapCollapse"; }

bool MapCollapse::can_be_applied(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {
    // Criterion: count must be at least 2
    if (count_ < 2) {
        return false;
    }

    // Fast path: a chain of perfectly-nested maps.
    if (this->check_perfect_nest()) {
        return true;
    }

    // Imperfect (CUDA-style) collapse is performed one level at a time.
    if (count_ == 2 && this->check_imperfect()) {
        return true;
    }

    return false;
}

bool MapCollapse::check_perfect_nest() {
    // Collect the chain of count_ perfectly-nested maps
    std::vector<structured_control_flow::Map*> maps;
    maps.push_back(&loop_);

    auto* current = &loop_;
    for (size_t i = 1; i < count_; ++i) {
        auto& body = current->root();

        // Criterion: All target maps must be perfectly nested (exactly one child that is a Map)
        if (body.size() != 1) {
            return false;
        }

        auto* next = dynamic_cast<structured_control_flow::Map*>(&body.at(0).first);
        if (!next) {
            return false;
        }

        // Criterion: The Sequence holding each map must have empty transitions
        if (!body.at(0).second.empty()) {
            return false;
        }

        maps.push_back(next);
        current = next;
    }

    // Criterion: All maps must be contiguous (stride 1)
    for (auto* map : maps) {
        if (!map->is_contiguous()) {
            return false;
        }
    }

    // Collect indvars of all maps being collapsed
    symbolic::SymbolSet indvars;
    for (auto* map : maps) {
        indvars.insert(map->indvar());
    }

    // Criterion: Map inits may not depend on any of the loop induction variables
    for (auto* map : maps) {
        auto init = map->init();
        for (auto& iv : indvars) {
            if (symbolic::uses(init, iv)) {
                return false;
            }
        }
    }

    // Criterion: Map bounds may not depend on any of the loop induction variables
    // of the maps being collapsed
    for (auto* map : maps) {
        auto bound = map->canonical_bound();
        if (bound.is_null()) {
            // If we can't even compute a closed-form bound, be conservative and disallow collapsing
            return false;
        }
        for (auto& iv : indvars) {
            if (symbolic::uses(bound, iv)) {
                return false;
            }
        }
    }

    return true;
}

bool MapCollapse::is_collapsible_inner_map(structured_control_flow::Map& map, const symbolic::Symbol& outer_indvar) {
    // Must increment by exactly 1 to participate in the flattened iteration space
    if (!map.is_contiguous()) {
        return false;
    }

    // Init may not depend on the outer induction variable
    if (symbolic::uses(map.init(), outer_indvar)) {
        return false;
    }

    // A closed-form bound is required and may not depend on the outer indvar
    auto bound = map.canonical_bound();
    if (bound.is_null()) {
        return false;
    }
    if (symbolic::uses(bound, outer_indvar)) {
        return false;
    }

    return true;
}

bool MapCollapse::check_imperfect() {
    // The outer map must itself be collapsible (contiguous, closed-form bound
    // that does not depend on its own induction variable).
    auto outer_indvar = loop_.indvar();
    if (!loop_.is_contiguous()) {
        return false;
    }
    if (symbolic::uses(loop_.init(), outer_indvar)) {
        return false;
    }
    auto outer_bound = loop_.canonical_bound();
    if (outer_bound.is_null()) {
        return false;
    }
    if (symbolic::uses(outer_bound, outer_indvar)) {
        return false;
    }

    auto& body = loop_.root();

    size_t num_collapsible = 0;
    for (size_t idx = 0; idx < body.size(); ++idx) {
        // Criterion (v1): transitions between body elements must be empty. Any
        // inter-element assignments would need to be re-guarded, which is not
        // handled yet.
        if (!body.at(idx).second.empty()) {
            return false;
        }

        auto* map = dynamic_cast<structured_control_flow::Map*>(&body.at(idx).first);
        if (map != nullptr && this->is_collapsible_inner_map(*map, outer_indvar)) {
            ++num_collapsible;
        }
        // Everything else (blocks, if-else, loops, non-collapsible maps) is a
        // "skipped" element that will be guarded to run once per outer iteration.
    }

    // Need at least one collapsible inner map to flatten against.
    return num_collapsible >= 1;
}

void MapCollapse::apply(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {
    if (this->check_perfect_nest()) {
        this->apply_perfect(builder, analysis_manager);
    } else {
        this->apply_imperfect(builder, analysis_manager);
    }
}

void MapCollapse::apply_perfect(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {
    auto& sdfg = builder.subject();

    // Step 1: Gather the maps to collapse and their bounds
    std::vector<structured_control_flow::Map*> maps;
    maps.push_back(&loop_);
    auto* current = &loop_;
    for (size_t i = 1; i < count_; ++i) {
        auto* next = dynamic_cast<structured_control_flow::Map*>(&current->root().at(0).first);
        maps.push_back(next);
        current = next;
    }

    std::vector<symbolic::Symbol> indvars;
    std::vector<symbolic::Expression> bounds;
    for (auto* map : maps) {
        indvars.push_back(map->indvar());
        bounds.push_back(map->canonical_bound());
    }

    // Step 2: Compute total iteration count = product of all bounds
    symbolic::Expression total_bound = bounds[0];
    for (size_t i = 1; i < bounds.size(); ++i) {
        total_bound = symbolic::mul(total_bound, bounds[i]);
    }

    // Step 3: Create the collapsed induction variable
    auto civ_name = builder.find_new_name(indvars[0]->get_name() + "_collapsed");
    builder.add_container(civ_name, sdfg.type(indvars[0]->get_name()));
    auto civ = symbolic::symbol(civ_name);

    // Step 4: Find the parent sequence of the outermost map
    auto parent = static_cast<structured_control_flow::Sequence*>(loop_.get_parent());
    size_t index = parent->index(loop_);
    auto& transition = parent->at(index).second;

    // Step 5: Create the new collapsed map before the original
    auto& collapsed_map = builder.add_map_before(
        *parent,
        loop_,
        civ,
        symbolic::Lt(civ, total_bound),
        symbolic::integer(0),
        symbolic::add(civ, symbolic::integer(1)),
        loop_.schedule_type(),
        transition.assignments(),
        loop_.debug_info()
    );

    // Step 6: Add an empty block for indvar recovery before the original contents
    builder.add_block(collapsed_map.root());

    // Step 7: Move the body of the innermost map into the collapsed map
    auto* innermost = maps.back();
    builder.move_children(innermost->root(), collapsed_map.root());

    // Step 8: Add indvar recovery assignments to the transition of the empty
    // block so that all induction variables are defined before the original
    // loop contents.
    //
    // For maps [0..n-1] with bounds [B0, B1, ..., B_{n-1}]:
    //   indvar_0     = civ / (B1 * B2 * ... * B_{n-1})
    //   indvar_k     = (civ / (B_{k+1} * ... * B_{n-1})) % B_k
    //   indvar_{n-1} = civ % B_{n-1}
    auto& first_transition = collapsed_map.root().at(0).second;
    size_t n = indvars.size();

    for (size_t k = 0; k < n; ++k) {
        // Compute suffix product = B_{k+1} * ... * B_{n-1}
        symbolic::Expression suffix = symbolic::integer(1);
        for (size_t j = k + 1; j < n; ++j) {
            suffix = symbolic::mul(suffix, bounds[j]);
        }

        symbolic::Expression value;
        if (k == 0 && n == 1) {
            // Single-loop degenerate case (shouldn't happen with count>=2, but safe)
            value = civ;
        } else if (k == 0) {
            // Outermost: indvar_0 = civ / suffix
            value = symbolic::div(civ, suffix);
        } else if (k == n - 1) {
            // Innermost: indvar_{n-1} = civ % B_{n-1}
            value = symbolic::mod(civ, bounds[k]);
        } else {
            // Middle: indvar_k = (civ / suffix) % B_k
            value = symbolic::mod(symbolic::div(civ, suffix), bounds[k]);
        }

        first_transition.assignments()[indvars[k]] = value;
    }

    // Step 9: Remove the original nest
    // The index shifted by 1 because we inserted a map before
    transition.assignments().clear();
    builder.remove_child(*parent, parent->index(loop_));

    applied_ = true;
    collapsed_loop_ = &collapsed_map;
}

void MapCollapse::apply_imperfect(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {
    auto& sdfg = builder.subject();
    auto& outer = loop_;
    auto outer_indvar = outer.indvar();
    auto& body = outer.root();

    // Step 1: Classify the outer body children (in order) and gather the bounds
    // of the collapsible inner maps.
    struct BodyItem {
        structured_control_flow::ControlFlowNode* node;
        structured_control_flow::Map* map; // non-null only if collapsible
        symbolic::Expression bound; // valid only if collapsible
    };

    std::vector<BodyItem> items;
    std::vector<symbolic::Expression> collapsible_bounds;
    for (size_t idx = 0; idx < body.size(); ++idx) {
        auto& child = body.at(idx).first;
        auto* map = dynamic_cast<structured_control_flow::Map*>(&child);
        if (map != nullptr && this->is_collapsible_inner_map(*map, outer_indvar)) {
            auto bound = map->canonical_bound();
            items.push_back({&child, map, bound});
            collapsible_bounds.push_back(bound);
        } else {
            items.push_back({&child, nullptr, SymEngine::null});
        }
    }

    // Step 2: Virtual inner extent = max of all collapsible bounds.
    symbolic::Expression inner_extent = collapsible_bounds[0];
    for (size_t i = 1; i < collapsible_bounds.size(); ++i) {
        inner_extent = symbolic::max(inner_extent, collapsible_bounds[i]);
    }

    auto outer_bound = outer.canonical_bound();
    auto total_bound = symbolic::mul(outer_bound, inner_extent);

    // Step 3: Create the collapsed induction variable and the virtual inner index.
    auto civ_name = builder.find_new_name(outer_indvar->get_name() + "_collapsed");
    builder.add_container(civ_name, sdfg.type(outer_indvar->get_name()));
    auto civ = symbolic::symbol(civ_name);

    auto inner_name = builder.find_new_name(outer_indvar->get_name() + "_inner");
    builder.add_container(inner_name, sdfg.type(outer_indvar->get_name()));
    auto inner_index = symbolic::symbol(inner_name);

    // Step 4: Find the parent sequence of the outer map.
    auto parent = static_cast<structured_control_flow::Sequence*>(outer.get_parent());
    size_t index = parent->index(outer);
    auto& transition = parent->at(index).second;

    // Step 5: Create the collapsed map before the original outer map.
    auto& collapsed_map = builder.add_map_before(
        *parent,
        outer,
        civ,
        symbolic::Lt(civ, total_bound),
        symbolic::integer(0),
        symbolic::add(civ, symbolic::integer(1)),
        outer.schedule_type(),
        transition.assignments(),
        outer.debug_info()
    );

    // Step 6: Recovery block defining the outer index and the virtual inner index:
    //   outer_indvar = civ / inner_extent
    //   inner_index  = civ % inner_extent
    builder.add_block(collapsed_map.root());
    auto& recovery_transition = collapsed_map.root().at(0).second;
    recovery_transition.assignments()[outer_indvar] = symbolic::div(civ, inner_extent);
    recovery_transition.assignments()[inner_index] = symbolic::mod(civ, inner_extent);

    // Step 7: Move the outer body into the collapsed map (after the recovery block),
    // preserving order.
    builder.move_children(outer.root(), collapsed_map.root());

    // Step 8: Guard each original body element. Collapsible maps run for the valid
    // portion of the inner extent (`inner_index < bound`); every other element runs
    // exactly once per outer iteration (`inner_index == 0`).
    for (auto& item : items) {
        auto* child = item.node;
        if (item.map != nullptr) {
            auto inner_iv = item.map->indvar();

            auto& if_else = builder.add_if_else_before(collapsed_map.root(), *child, {}, outer.debug_info());
            auto& branch = builder.add_case(if_else, symbolic::Lt(inner_index, item.bound), outer.debug_info());

            // Recover the inner induction variable from the virtual index.
            builder.add_block(branch);
            branch.at(0).second.assignments()[inner_iv] = inner_index;

            // Move the map body into the guarded branch and drop the empty map shell.
            builder.move_children(item.map->root(), branch);
            builder.remove_child(collapsed_map.root(), collapsed_map.root().index(*child));
        } else {
            auto& if_else = builder.add_if_else_before(collapsed_map.root(), *child, {}, outer.debug_info());
            auto& branch =
                builder.add_case(if_else, symbolic::Eq(inner_index, symbolic::integer(0)), outer.debug_info());

            builder.move_child(collapsed_map.root(), collapsed_map.root().index(*child), branch);
        }
    }

    // Step 9: Remove the original outer map.
    transition.assignments().clear();
    builder.remove_child(*parent, parent->index(outer));

    applied_ = true;
    collapsed_loop_ = &collapsed_map;

    analysis_manager.invalidate_all();
}

void MapCollapse::to_json(nlohmann::json& j) const {
    std::string loop_type;
    if (dynamic_cast<const structured_control_flow::Map*>(&loop_)) {
        loop_type = "map";
    } else {
        throw InvalidSDFGException("Unsupported loop type for serialization of map: " + loop_.indvar()->get_name());
    }

    j["transformation_type"] = this->name();
    j["subgraph"] = {{"0", {{"element_id", this->loop_.element_id()}, {"type", loop_type}}}};
    j["parameters"] = {{"count", count_}};
}

MapCollapse MapCollapse::from_json(builder::StructuredSDFGBuilder& builder, const nlohmann::json& desc) {
    auto loop_id = desc["subgraph"]["0"]["element_id"].get<size_t>();
    size_t count = desc["parameters"]["count"].get<size_t>();
    auto element = builder.find_element_by_id(loop_id);
    if (!element) {
        throw InvalidTransformationDescriptionException("Element with ID " + std::to_string(loop_id) + " not found.");
    }
    auto loop = dynamic_cast<structured_control_flow::Map*>(element);

    return MapCollapse(*loop, count);
}

structured_control_flow::Map* MapCollapse::collapsed_loop() {
    if (!applied_) {
        return &loop_;
    }

    return collapsed_loop_;
}

} // namespace transformations
} // namespace sdfg

1	#include "sdfg/transformations/map_collapse.h"
2
3	#include "sdfg/builder/structured_sdfg_builder.h"
4	#include "sdfg/structured_control_flow/if_else.h"
5	#include "sdfg/structured_control_flow/map.h"
6	#include "sdfg/structured_control_flow/sequence.h"
7	#include "sdfg/structured_control_flow/structured_loop.h"
8	#include "sdfg/symbolic/symbolic.h"
9
10	namespace sdfg {
11	namespace transformations {
12
13	MapCollapse::MapCollapse(structured_control_flow::Map& loop, size_t count) : loop_(loop), count_(count) {}	64✔
14
15	std::string MapCollapse::name() const { return "MapCollapse"; }	2✔
16
17	bool MapCollapse::can_be_applied(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {	52✔
18	// Criterion: count must be at least 2
19	if (count_ < 2) {	52✔
20	return false;	1✔
21	}	1✔
22
23	// Fast path: a chain of perfectly-nested maps.
24	if (this->check_perfect_nest()) {	51✔
25	return true;	32✔
26	}	32✔
27
28	// Imperfect (CUDA-style) collapse is performed one level at a time.
29	if (count_ == 2 && this->check_imperfect()) {	19✔
30	return true;	12✔
31	}	12✔
32
33	return false;	7✔
34	}	19✔
35
36	bool MapCollapse::check_perfect_nest() {	83✔
37	// Collect the chain of count_ perfectly-nested maps
38	std::vector<structured_control_flow::Map*> maps;	83✔
39	maps.push_back(&loop_);	83✔
40
41	auto* current = &loop_;	83✔
42	for (size_t i = 1; i < count_; ++i) {	159✔
43	auto& body = current->root();	101✔
44
45	// Criterion: All target maps must be perfectly nested (exactly one child that is a Map)
46	if (body.size() != 1) {	101✔
47	return false;	23✔
48	}	23✔
49
50	auto* next = dynamic_cast<structured_control_flow::Map*>(&body.at(0).first);	78✔
51	if (!next) {	78✔
52	return false;	1✔
53	}	1✔
54
55	// Criterion: The Sequence holding each map must have empty transitions
56	if (!body.at(0).second.empty()) {	77✔
57	return false;	1✔
58	}	1✔
59
60	maps.push_back(next);	76✔
61	current = next;	76✔
62	}	76✔
63
64	// Criterion: All maps must be contiguous (stride 1)
65	for (auto* map : maps) {	133✔
66	if (!map->is_contiguous()) {	133✔
UNCOV 67	return false;	×
UNCOV 68	}	×
69	}	133✔
70
71	// Collect indvars of all maps being collapsed
72	symbolic::SymbolSet indvars;	58✔
73	for (auto* map : maps) {	133✔
74	indvars.insert(map->indvar());	133✔
75	}	133✔
76
77	// Criterion: Map inits may not depend on any of the loop induction variables
78	for (auto* map : maps) {	133✔
79	auto init = map->init();	133✔
80	for (auto& iv : indvars) {	322✔
81	if (symbolic::uses(init, iv)) {	322✔
82	return false;	1✔
83	}	1✔
84	}	322✔
85	}	133✔
86
87	// Criterion: Map bounds may not depend on any of the loop induction variables
88	// of the maps being collapsed
89	for (auto* map : maps) {	131✔
90	auto bound = map->canonical_bound();	131✔
91	if (bound.is_null()) {	131✔
92	// If we can't even compute a closed-form bound, be conservative and disallow collapsing
UNCOV 93	return false;	×
UNCOV 94	}	×
95	for (auto& iv : indvars) {	318✔
96	if (symbolic::uses(bound, iv)) {	318✔
97	return false;	1✔
98	}	1✔
99	}	318✔
100	}	131✔
101
102	return true;	56✔
103	}	57✔
104
105	bool MapCollapse::is_collapsible_inner_map(structured_control_flow::Map& map, const symbolic::Symbol& outer_indvar) {	38✔
106	// Must increment by exactly 1 to participate in the flattened iteration space
107	if (!map.is_contiguous()) {	38✔
108	return false;	1✔
109	}	1✔
110
111	// Init may not depend on the outer induction variable
112	if (symbolic::uses(map.init(), outer_indvar)) {	37✔
113	return false;	1✔
114	}	1✔
115
116	// A closed-form bound is required and may not depend on the outer indvar
117	auto bound = map.canonical_bound();	36✔
118	if (bound.is_null()) {	36✔
NEW 119	return false;	×
NEW 120	}	×
121	if (symbolic::uses(bound, outer_indvar)) {	36✔
122	return false;	1✔
123	}	1✔
124
125	return true;	35✔
126	}	36✔
127
128	bool MapCollapse::check_imperfect() {	17✔
129	// The outer map must itself be collapsible (contiguous, closed-form bound
130	// that does not depend on its own induction variable).
131	auto outer_indvar = loop_.indvar();	17✔
132	if (!loop_.is_contiguous()) {	17✔
133	return false;	1✔
134	}	1✔
135	if (symbolic::uses(loop_.init(), outer_indvar)) {	16✔
NEW 136	return false;	×
NEW 137	}	×
138	auto outer_bound = loop_.canonical_bound();	16✔
139	if (outer_bound.is_null()) {	16✔
NEW 140	return false;	×
NEW 141	}	×
142	if (symbolic::uses(outer_bound, outer_indvar)) {	16✔
NEW 143	return false;	×
NEW 144	}	×
145
146	auto& body = loop_.root();	16✔
147
148	size_t num_collapsible = 0;	16✔
149	for (size_t idx = 0; idx < body.size(); ++idx) {	45✔
150	// Criterion (v1): transitions between body elements must be empty. Any
151	// inter-element assignments would need to be re-guarded, which is not
152	// handled yet.
153	if (!body.at(idx).second.empty()) {	30✔
154	return false;	1✔
155	}	1✔
156
157	auto* map = dynamic_cast<structured_control_flow::Map*>(&body.at(idx).first);	29✔
158	if (map != nullptr && this->is_collapsible_inner_map(*map, outer_indvar)) {	29✔
159	++num_collapsible;	20✔
160	}	20✔
161	// Everything else (blocks, if-else, loops, non-collapsible maps) is a
162	// "skipped" element that will be guarded to run once per outer iteration.
163	}	29✔
164
165	// Need at least one collapsible inner map to flatten against.
166	return num_collapsible >= 1;	15✔
167	}	16✔
168
169	void MapCollapse::apply(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {	32✔
170	if (this->check_perfect_nest()) {	32✔
171	this->apply_perfect(builder, analysis_manager);	24✔
172	} else {	24✔
173	this->apply_imperfect(builder, analysis_manager);	8✔
174	}	8✔
175	}	32✔
176
177	void MapCollapse::apply_perfect(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {	24✔
178	auto& sdfg = builder.subject();	24✔
179
180	// Step 1: Gather the maps to collapse and their bounds
181	std::vector<structured_control_flow::Map*> maps;	24✔
182	maps.push_back(&loop_);	24✔
183	auto* current = &loop_;	24✔
184	for (size_t i = 1; i < count_; ++i) {	54✔
185	auto* next = dynamic_cast<structured_control_flow::Map*>(&current->root().at(0).first);	30✔
186	maps.push_back(next);	30✔
187	current = next;	30✔
188	}	30✔
189
190	std::vector<symbolic::Symbol> indvars;	24✔
191	std::vector<symbolic::Expression> bounds;	24✔
192	for (auto* map : maps) {	54✔
193	indvars.push_back(map->indvar());	54✔
194	bounds.push_back(map->canonical_bound());	54✔
195	}	54✔
196
197	// Step 2: Compute total iteration count = product of all bounds
198	symbolic::Expression total_bound = bounds[0];	24✔
199	for (size_t i = 1; i < bounds.size(); ++i) {	54✔
200	total_bound = symbolic::mul(total_bound, bounds[i]);	30✔
201	}	30✔
202
203	// Step 3: Create the collapsed induction variable
204	auto civ_name = builder.find_new_name(indvars[0]->get_name() + "_collapsed");	24✔
205	builder.add_container(civ_name, sdfg.type(indvars[0]->get_name()));	24✔
206	auto civ = symbolic::symbol(civ_name);	24✔
207
208	// Step 4: Find the parent sequence of the outermost map
209	auto parent = static_cast<structured_control_flow::Sequence*>(loop_.get_parent());	24✔
210	size_t index = parent->index(loop_);	24✔
211	auto& transition = parent->at(index).second;	24✔
212
213	// Step 5: Create the new collapsed map before the original
214	auto& collapsed_map = builder.add_map_before(	24✔
215	*parent,	24✔
216	loop_,	24✔
217	civ,	24✔
218	symbolic::Lt(civ, total_bound),	24✔
219	symbolic::integer(0),	24✔
220	symbolic::add(civ, symbolic::integer(1)),	24✔
221	loop_.schedule_type(),	24✔
222	transition.assignments(),	24✔
223	loop_.debug_info()	24✔
224	);	24✔
225
226	// Step 6: Add an empty block for indvar recovery before the original contents
227	builder.add_block(collapsed_map.root());	24✔
228
229	// Step 7: Move the body of the innermost map into the collapsed map
230	auto* innermost = maps.back();	24✔
231	builder.move_children(innermost->root(), collapsed_map.root());	24✔
232
233	// Step 8: Add indvar recovery assignments to the transition of the empty
234	// block so that all induction variables are defined before the original
235	// loop contents.
236	//
237	// For maps [0..n-1] with bounds [B0, B1, ..., B_{n-1}]:
238	// indvar_0 = civ / (B1 * B2 * ... * B_{n-1})
239	// indvar_k = (civ / (B_{k+1} * ... * B_{n-1})) % B_k
240	// indvar_{n-1} = civ % B_{n-1}
241	auto& first_transition = collapsed_map.root().at(0).second;	24✔
242	size_t n = indvars.size();	24✔
243
244	for (size_t k = 0; k < n; ++k) {	78✔
245	// Compute suffix product = B_{k+1} * ... * B_{n-1}
246	symbolic::Expression suffix = symbolic::integer(1);	54✔
247	for (size_t j = k + 1; j < n; ++j) {	91✔
248	suffix = symbolic::mul(suffix, bounds[j]);	37✔
249	}	37✔
250
251	symbolic::Expression value;	54✔
252	if (k == 0 && n == 1) {	54✔
253	// Single-loop degenerate case (shouldn't happen with count>=2, but safe)
UNCOV 254	value = civ;	×
255	} else if (k == 0) {	54✔
256	// Outermost: indvar_0 = civ / suffix
257	value = symbolic::div(civ, suffix);	24✔
258	} else if (k == n - 1) {	30✔
259	// Innermost: indvar_{n-1} = civ % B_{n-1}
260	value = symbolic::mod(civ, bounds[k]);	24✔
261	} else {	24✔
262	// Middle: indvar_k = (civ / suffix) % B_k
263	value = symbolic::mod(symbolic::div(civ, suffix), bounds[k]);	6✔
264	}	6✔
265
266	first_transition.assignments()[indvars[k]] = value;	54✔
267	}	54✔
268
269	// Step 9: Remove the original nest
270	// The index shifted by 1 because we inserted a map before
271	transition.assignments().clear();	24✔
272	builder.remove_child(*parent, parent->index(loop_));	24✔
273
274	applied_ = true;	24✔
275	collapsed_loop_ = &collapsed_map;	24✔
276	}	24✔
277
278	void MapCollapse::apply_imperfect(builder::StructuredSDFGBuilder& builder, analysis::AnalysisManager& analysis_manager) {	8✔
279	auto& sdfg = builder.subject();	8✔
280	auto& outer = loop_;	8✔
281	auto outer_indvar = outer.indvar();	8✔
282	auto& body = outer.root();	8✔
283
284	// Step 1: Classify the outer body children (in order) and gather the bounds
285	// of the collapsible inner maps.
286	struct BodyItem {	8✔
287	structured_control_flow::ControlFlowNode* node;	8✔
288	structured_control_flow::Map* map; // non-null only if collapsible	8✔
289	symbolic::Expression bound; // valid only if collapsible	8✔
290	};	8✔
291
292	std::vector<BodyItem> items;	8✔
293	std::vector<symbolic::Expression> collapsible_bounds;	8✔
294	for (size_t idx = 0; idx < body.size(); ++idx) {	25✔
295	auto& child = body.at(idx).first;	17✔
296	auto* map = dynamic_cast<structured_control_flow::Map*>(&child);	17✔
297	if (map != nullptr && this->is_collapsible_inner_map(*map, outer_indvar)) {	17✔
298	auto bound = map->canonical_bound();	15✔
299	items.push_back({&child, map, bound});	15✔
300	collapsible_bounds.push_back(bound);	15✔
301	} else {	15✔
302	items.push_back({&child, nullptr, SymEngine::null});	2✔
303	}	2✔
304	}	17✔
305
306	// Step 2: Virtual inner extent = max of all collapsible bounds.
307	symbolic::Expression inner_extent = collapsible_bounds[0];	8✔
308	for (size_t i = 1; i < collapsible_bounds.size(); ++i) {	15✔
309	inner_extent = symbolic::max(inner_extent, collapsible_bounds[i]);	7✔
310	}	7✔
311
312	auto outer_bound = outer.canonical_bound();	8✔
313	auto total_bound = symbolic::mul(outer_bound, inner_extent);	8✔
314
315	// Step 3: Create the collapsed induction variable and the virtual inner index.
316	auto civ_name = builder.find_new_name(outer_indvar->get_name() + "_collapsed");	8✔
317	builder.add_container(civ_name, sdfg.type(outer_indvar->get_name()));	8✔
318	auto civ = symbolic::symbol(civ_name);	8✔
319
320	auto inner_name = builder.find_new_name(outer_indvar->get_name() + "_inner");	8✔
321	builder.add_container(inner_name, sdfg.type(outer_indvar->get_name()));	8✔
322	auto inner_index = symbolic::symbol(inner_name);	8✔
323
324	// Step 4: Find the parent sequence of the outer map.
325	auto parent = static_cast<structured_control_flow::Sequence*>(outer.get_parent());	8✔
326	size_t index = parent->index(outer);	8✔
327	auto& transition = parent->at(index).second;	8✔
328
329	// Step 5: Create the collapsed map before the original outer map.
330	auto& collapsed_map = builder.add_map_before(	8✔
331	*parent,	8✔
332	outer,	8✔
333	civ,	8✔
334	symbolic::Lt(civ, total_bound),	8✔
335	symbolic::integer(0),	8✔
336	symbolic::add(civ, symbolic::integer(1)),	8✔
337	outer.schedule_type(),	8✔
338	transition.assignments(),	8✔
339	outer.debug_info()	8✔
340	);	8✔
341
342	// Step 6: Recovery block defining the outer index and the virtual inner index:
343	// outer_indvar = civ / inner_extent
344	// inner_index = civ % inner_extent
345	builder.add_block(collapsed_map.root());	8✔
346	auto& recovery_transition = collapsed_map.root().at(0).second;	8✔
347	recovery_transition.assignments()[outer_indvar] = symbolic::div(civ, inner_extent);	8✔
348	recovery_transition.assignments()[inner_index] = symbolic::mod(civ, inner_extent);	8✔
349
350	// Step 7: Move the outer body into the collapsed map (after the recovery block),
351	// preserving order.
352	builder.move_children(outer.root(), collapsed_map.root());	8✔
353
354	// Step 8: Guard each original body element. Collapsible maps run for the valid
355	// portion of the inner extent (`inner_index < bound`); every other element runs
356	// exactly once per outer iteration (`inner_index == 0`).
357	for (auto& item : items) {	17✔
358	auto* child = item.node;	17✔
359	if (item.map != nullptr) {	17✔
360	auto inner_iv = item.map->indvar();	15✔
361
362	auto& if_else = builder.add_if_else_before(collapsed_map.root(), *child, {}, outer.debug_info());	15✔
363	auto& branch = builder.add_case(if_else, symbolic::Lt(inner_index, item.bound), outer.debug_info());	15✔
364
365	// Recover the inner induction variable from the virtual index.
366	builder.add_block(branch);	15✔
367	branch.at(0).second.assignments()[inner_iv] = inner_index;	15✔
368
369	// Move the map body into the guarded branch and drop the empty map shell.
370	builder.move_children(item.map->root(), branch);	15✔
371	builder.remove_child(collapsed_map.root(), collapsed_map.root().index(*child));	15✔
372	} else {	15✔
373	auto& if_else = builder.add_if_else_before(collapsed_map.root(), *child, {}, outer.debug_info());	2✔
374	auto& branch =	2✔
375	builder.add_case(if_else, symbolic::Eq(inner_index, symbolic::integer(0)), outer.debug_info());	2✔
376
377	builder.move_child(collapsed_map.root(), collapsed_map.root().index(*child), branch);	2✔
378	}	2✔
379	}	17✔
380
381	// Step 9: Remove the original outer map.
382	transition.assignments().clear();	8✔
383	builder.remove_child(*parent, parent->index(outer));	8✔
384
385	applied_ = true;	8✔
386	collapsed_loop_ = &collapsed_map;	8✔
387
388	analysis_manager.invalidate_all();	8✔
389	}	8✔
390
391	void MapCollapse::to_json(nlohmann::json& j) const {	1✔
392	std::string loop_type;	1✔
393	if (dynamic_cast<const structured_control_flow::Map*>(&loop_)) {	1✔
394	loop_type = "map";	1✔
395	} else {	1✔
396	throw InvalidSDFGException("Unsupported loop type for serialization of map: " + loop_.indvar()->get_name());	×
397	}	×
398
399	j["transformation_type"] = this->name();	1✔
400	j["subgraph"] = {{"0", {{"element_id", this->loop_.element_id()}, {"type", loop_type}}}};	1✔
401	j["parameters"] = {{"count", count_}};	1✔
402	}	1✔
403
404	MapCollapse MapCollapse::from_json(builder::StructuredSDFGBuilder& builder, const nlohmann::json& desc) {	1✔
405	auto loop_id = desc["subgraph"]["0"]["element_id"].get<size_t>();	1✔
406	size_t count = desc["parameters"]["count"].get<size_t>();	1✔
407	auto element = builder.find_element_by_id(loop_id);	1✔
408	if (!element) {	1✔
409	throw InvalidTransformationDescriptionException("Element with ID " + std::to_string(loop_id) + " not found.");	×
410	}	×
411	auto loop = dynamic_cast<structured_control_flow::Map*>(element);	1✔
412
413	return MapCollapse(*loop, count);	1✔
414	}	1✔
415
416	structured_control_flow::Map* MapCollapse::collapsed_loop() {	33✔
417	if (!applied_) {	33✔
418	return &loop_;	1✔
419	}	1✔
420
421	return collapsed_loop_;	32✔
422	}	33✔
423
424	} // namespace transformations
425	} // namespace sdfg

daisytuner / docc / 28014415443

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