| 140 | } |
| 141 | |
| 142 | Maybe<std::vector<int64_t>> GetSelectedParallelIds(const Shape& hierarchy_shape, |
| 143 | const std::vector<int>& axis2is_selected, |
| 144 | int64_t parallel_id) { |
| 145 | CHECK_EQ_OR_RETURN(hierarchy_shape.NumAxes(), axis2is_selected.size()); |
| 146 | Stride hierarchy_strides(hierarchy_shape); |
| 147 | IndexVector indexes{}; |
| 148 | JUST(GetIndexesFromOffset(hierarchy_strides, parallel_id, &indexes)); |
| 149 | std::function<void(const DimVector&, DimVector*)> SelectedIndex2OriginIndex; |
| 150 | JUST(GetSelectedIndex2OriginIndex(indexes, axis2is_selected, &SelectedIndex2OriginIndex)); |
| 151 | const auto& broadcast_shape = JUST(GetSelectedShape(hierarchy_shape, axis2is_selected)); |
| 152 | Stride broadcast_strides(*broadcast_shape); |
| 153 | const auto& origin_offsets = std::make_shared<std::vector<int64_t>>(broadcast_shape->elem_cnt()); |
| 154 | for (int64_t i = 0; i < broadcast_shape->elem_cnt(); ++i) { |
| 155 | IndexVector broadcast_indexes{}; |
| 156 | JUST(GetIndexesFromOffset(broadcast_strides, i, &broadcast_indexes)); |
| 157 | IndexVector origin_indexes{}; |
| 158 | SelectedIndex2OriginIndex(broadcast_indexes, &origin_indexes); |
| 159 | int64_t origin_offset = -1; |
| 160 | JUST(GetOffsetFromIndexes(hierarchy_strides, origin_indexes, &origin_offset)); |
| 161 | origin_offsets->at(i) = origin_offset; |
| 162 | } |
| 163 | return origin_offsets; |
| 164 | } |
| 165 | |
| 166 | Maybe<Symbol<ParallelDesc>> GetBroadcastSubParallelDesc(Symbol<ParallelDesc> parallel_desc, |
| 167 | Symbol<NdSbp> nd_sbp) { |
no test coverage detected