Home
last modified time | relevance | path

Searched refs:src_dst_pairs (Results 1 – 4 of 4) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dfft_handler.cc284 std::vector<std::pair<int64, int64>> src_dst_pairs; in GetFinalFftUsingCollectivePermute() local
290 src_dst_pairs.emplace_back(src_device, dst_device); in GetFinalFftUsingCollectivePermute()
295 &body_b, source_partition_id, src_dst_pairs, (*next_channel_id)++); in GetFinalFftUsingCollectivePermute()
299 &body_b, source_transform, src_dst_pairs, (*next_channel_id)++); in GetFinalFftUsingCollectivePermute()
Dspmd_partitioner_util.cc1630 std::vector<std::pair<int64, int64>>& src_dst_pairs, in GetPerGroupCollectiveOpsCreator()
1633 src_dst_pairs.size() * device_groups.size()); in GetPerGroupCollectiveOpsCreator()
1635 for (int64 i = 0; i < src_dst_pairs.size(); ++i) { in GetPerGroupCollectiveOpsCreator()
1636 expanded_pairs[g * src_dst_pairs.size() + i] = in GetPerGroupCollectiveOpsCreator()
1638 device_groups[g][src_dst_pairs[i].first], in GetPerGroupCollectiveOpsCreator()
1639 device_groups[g][src_dst_pairs[i].second]}; in GetPerGroupCollectiveOpsCreator()
Dspmd_partitioner.h127 std::vector<std::pair<int64, int64>>& src_dst_pairs,
Dspmd_partitioner.cc1338 std::vector<std::pair<int64, int64>> src_dst_pairs; in ReshardWithCollectivePermute() local
1342 src_dst_pairs.emplace_back(src_device, dst_device); in ReshardWithCollectivePermute()
1346 state_.b, hlo(), src_dst_pairs, (*state_.next_channel_id)++); in ReshardWithCollectivePermute()
3449 std::vector<std::pair<int64, int64>>& src_dst_pairs, in GetDefaultCollectiveOpsCreator()
3452 operand->shape(), operand, src_dst_pairs, channel_id)); in GetDefaultCollectiveOpsCreator()