IntenseWebs/openvino
3
0 Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Cleanup in Split mask propagation transformation (#14119)

Browse Source 
* Cleanup in Split mask propagation transformation

* review comments
This commit is contained in:
Mateusz Tabaka authored and GitHub committed 2022-11-22 10:25:58 +04:00
1 parent 81c579fd2a
commit 47acda3476
2 files changed
+46 -73
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/common/offline_transformations/include/mask_attribute.hpp
+1 -1
View File
@@ -203,7 +203,7 @@ public:
while (it != src_set.end() && *it < split_end)
dst_set.insert(*it++ - split_start);
at(i) = dst_set;
} else if (!mask->at(i).empty()) {
} else {
at(i) = mask->at(i);
}
}
src/common/offline_transformations/src/pruning/propagate_masks.cpp
+45 -72
View File
@@ -1432,6 +1432,42 @@ public:
}
};
static ngraph::Mask::Ptr create_connect_split_output_mask(ngraph::Mask::Ptr input_mask,
const int64_t axis,
const uint64_t split_start,
const uint64_t split_end) {
auto output_mask = std::make_shared<ngraph::Mask>();
auto input_mask_raw = input_mask.get();
output_mask->add_callback(
[input_mask_raw, axis, split_start, split_end](ngraph::Mask::Ptr cur_mask) -> bool {
cur_mask->copy_and_slice_mask_from(input_mask_raw, axis, split_start, split_end);
return true;
},
input_mask);
auto output_mask_raw = output_mask.get();
input_mask->add_callback(
[output_mask_raw, axis, split_start, split_end](ngraph::Mask::Ptr cur_mask) -> bool {
auto& dim_mask = cur_mask->at(axis);
auto it = dim_mask.lower_bound(split_start);
while (it != dim_mask.end() && *it < split_end) {
it = dim_mask.erase(it);
}
for (size_t j = 0; j < output_mask_raw->size(); j++) {
const auto& dim_mask = output_mask_raw->at(j);
if (j == axis) {
for (auto d : dim_mask)
cur_mask->at(j).insert(d + split_start);
} else {
cur_mask->at(j) = dim_mask;
}
}
return true;
},
output_mask);
return output_mask;
}
class ov::pass::mask_propagation::VariadicSplit : public MatcherPass {
public:
VariadicSplit() {
@@ -1484,45 +1520,14 @@ public:
uint64_t split_start = 0;
uint64_t split_end = 0;
std::vector<ngraph::Mask::Ptr> output_masks;
std::vector<ngraph::Mask*> output_masks_raw;
auto input_mask_raw = input_mask.get();
for (size_t i = 0; i < split->get_output_size(); i++) {
split_end += split_lengths[i];
auto output_mask = std::make_shared<ngraph::Mask>();
output_mask->copy_and_slice_mask_from(input_mask_raw, axis, split_start, split_end);
output_masks.push_back(output_mask);
output_masks_raw.push_back(output_mask.get());
output_mask->add_callback(
[input_mask_raw, axis, split_start, split_end](ngraph::Mask::Ptr cur_mask) -> bool {
cur_mask->copy_and_slice_mask_from(input_mask_raw, axis, split_start, split_end);
return true;
},
input_mask);
split_start = split_end;
output_masks.push_back(create_connect_split_output_mask(input_mask, axis, split_start, split_end));
ngraph::setMask(split->output(i), output_masks[i]);
split_start = split_end;
}
auto input_mask_callback = [output_masks_raw, axis, split_lengths](ngraph::Mask::Ptr cur_mask) -> bool {
cur_mask->clean_dim_values();
uint64_t split_start = 0;
for (size_t i = 0; i < output_masks_raw.size(); i++) {
auto mask = output_masks_raw[i];
for (size_t j = 0; j < mask->size(); j++) {
const auto& dim_mask = mask->at(j);
if (j == axis) {
for (auto d : dim_mask)
cur_mask->at(j).insert(d + split_start);
} else if (!dim_mask.empty()) {
cur_mask->at(j) = dim_mask;
}
}
split_start += split_lengths[i];
}
return true;
};
for (size_t i = 0; i < split->get_output_size(); i++) {
input_mask->add_callback(input_mask_callback, output_masks[i]);
for (const auto& output_mask : output_masks) {
output_mask->apply_callback(input_mask);
}
return true;
};
@@ -1565,46 +1570,14 @@ public:
auto split_step = split_dim / num_splits;
uint64_t split_end = split_step;
std::vector<ngraph::Mask::Ptr> output_masks;
std::vector<ngraph::Mask*> output_masks_raw;
auto input_mask_raw = input_mask.get();
for (size_t i = 0; i < num_splits; i++) {
auto output_mask = std::make_shared<ngraph::Mask>();
output_mask->copy_and_slice_mask_from(input_mask_raw, axis, split_start, split_end);
output_masks.push_back(output_mask);
output_masks_raw.push_back(output_mask.get());
output_mask->add_callback(
[input_mask_raw, axis, split_start, split_end](ngraph::Mask::Ptr cur_mask) -> bool {
cur_mask->copy_and_slice_mask_from(input_mask_raw, axis, split_start, split_end);
return true;
},
input_mask);
for (size_t i = 0; i < split->get_output_size(); i++) {
output_masks.push_back(create_connect_split_output_mask(input_mask, axis, split_start, split_end));
ngraph::setMask(split->output(i), output_masks[i]);
split_start = split_end;
split_end += split_step;
ngraph::setMask(split->output(i), output_masks[i]);
}
auto input_mask_callback =
[output_masks_raw, axis, split_dim, split_step, num_splits](ngraph::Mask::Ptr cur_mask) -> bool {
cur_mask->clean_dim_values();
uint64_t split_start = 0;
for (auto mask : output_masks_raw) {
for (size_t j = 0; j < mask->size(); j++) {
const auto& dim_mask = mask->at(j);
if (j == axis) {
for (auto d : dim_mask)
cur_mask->at(j).insert(d + split_start);
} else if (!dim_mask.empty()) {
cur_mask->at(j) = dim_mask;
}
}
split_start += split_step;
}
cur_mask->initialize_dependencies();
return true;
};
for (size_t i = 0; i < split->get_output_size(); i++) {
input_mask->add_callback(input_mask_callback, output_masks[i]);
for (const auto& output_mask : output_masks) {
output_mask->apply_callback(input_mask);
}
return true;
};