Search for multiple "slice" instructions in an instruction's outputs * which are contiguous slices of the same tensor. */
| 1117 | * which are contiguous slices of the same tensor. |
| 1118 | */ |
| 1119 | static std::vector<instruction_ref> get_splits(instruction_ref ins) |
| 1120 | { |
| 1121 | std::vector<instruction_ref> result; |
| 1122 | std::copy_if(ins->outputs().begin(), |
| 1123 | ins->outputs().end(), |
| 1124 | std::back_inserter(result), |
| 1125 | [&](auto i) { return i->name() == "slice"; }); |
| 1126 | if(result.size() < 2) |
| 1127 | return {}; |
| 1128 | auto get_slice = [](auto& i) -> auto& { return any_cast<op::slice>(i->get_operator()); }; |
| 1129 | auto&& axes = get_slice(result.front()).axes; |
| 1130 | |
| 1131 | // "slice" instructions must all have the same axes |
| 1132 | if(std::any_of(result.begin(), result.end(), [&](auto i) { return get_slice(i).axes != axes; })) |
| 1133 | return {}; |
| 1134 | auto get_start = [&](auto& i) -> auto& { return get_slice(i).starts; }; |
| 1135 | auto get_end = [&](auto& i) -> auto& { return get_slice(i).ends; }; |
| 1136 | |
| 1137 | // Sort the "slice" instructions in order of starts |
| 1138 | std::sort( |
| 1139 | result.begin(), result.end(), [&](auto x, auto y) { return get_start(x) < get_start(y); }); |
| 1140 | if(std::any_of(get_start(result.front()).begin(), get_start(result.front()).end(), [&](auto i) { |
| 1141 | return i != 0; |
| 1142 | })) |
| 1143 | return {}; |
| 1144 | |
| 1145 | // one slice must "start" where the last slice "end" |
| 1146 | auto it = std::adjacent_find( |
| 1147 | result.begin(), result.end(), [&](auto x, auto y) { return get_end(x) != get_start(y); }); |
| 1148 | if(it != result.end()) |
| 1149 | return {}; |
| 1150 | for(std::size_t i = 0; i < axes.size(); i++) |
| 1151 | { |
| 1152 | auto axis = axes[i]; |
| 1153 | if(ins->get_shape().lens()[axis] != get_slice(result.back()).ends[i]) |
| 1154 | return {}; |
| 1155 | } |
| 1156 | return result; |
| 1157 | } |
| 1158 | |
| 1159 | struct find_splits |
| 1160 | { |