Skip to content

Commit e5ff051

Browse files
authored
[CINN] Fix remove reduce one axis in anchor fusion (#70747)
1 parent 294686b commit e5ff051

File tree

1 file changed

+7
-6
lines changed

1 file changed

+7
-6
lines changed

paddle/cinn/operator_fusion/policy/iters_fusion_policy.cc

+7-6
Original file line numberDiff line numberDiff line change
@@ -257,12 +257,13 @@ std::optional<ItersTransformRoute> ItersFusionPolicy::SearchItersTransformRoute(
257257
auto squeezed_source = source;
258258
if (squeeze_source) {
259259
// Remove iters equal to one in source
260-
auto source_ones = MapVectorIfTrue<std::pair<std::string, int>, int>(
261-
Enumerate(source.loop_iters),
262-
[this](std::pair<std::string, int> p) { return p.second; },
263-
[this](std::pair<std::string, int> p) {
264-
return this->iters_manager_->IterSymbolEqualOne(p.first);
265-
});
260+
std::vector<int> source_ones;
261+
for (int i = 0; i < source.loop_iters.size() - source.reduce_iter_nums;
262+
++i) {
263+
if (iters_manager_->IterSymbolEqualOne(source.loop_iters[i])) {
264+
source_ones.push_back(i);
265+
}
266+
}
266267
if (!source_ones.empty() &&
267268
source_ones.size() != source.loop_iters.size()) {
268269
iters_transforms.emplace_back(RemoveOnesTransform(source_ones));

0 commit comments

Comments
 (0)