all_reduce_folder.cc
imports tsl/platform/errors.hStructured object names a sequence field
SHA-25642bf807c138f7a35b5e70f0ade9a3443e08720ecb2918b4685f0bca7e3b8bcde
MaleculeH(Db)Md(In)
Evidence
1/* Copyright 2021 The OpenXLA Authors.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is …
11:42… IND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#include "xla/hlo/transforms/simplifiers/all_reduce_folder.h"
17
18#include <algorithm>
19#include <cstdint>
20#include <optional>
21#include <vector>
22
23#include "absl/algorithm/container.h"
24#include "absl/container/flat_hash_map.h"
25#include "absl/container/flat_hash_set.h"
26#include "absl/log/log.h"
27#include "absl/status/statusor.h"
28#include "absl/strings/string_view.h"
29#include "absl/types/span.h"
30#include "xla/hlo/ir/hlo_casting_utils.h"
31#include "xla/hlo/ir/hlo_instruction.h"
32#include "xla/hlo/ir/hlo_instructions.h"
33#include "xla/hlo/ir/hlo_module.h"
34#include "xla/hlo/ir/hlo_opcode.h"
35#include "xla/hlo/ir/replica_group.h"
36#include "xla/hlo/utils/hlo_query.h"
37#include "xla/service/all_reduce_key.h"
38#include "xla/xla_data.pb.h"
39#include "tsl/platform/errors.h"
40
41namespace xla {
⋯4 lines
58:60… replicas which
59 // contribute to the output of that replica by essentially tracing through
60 // the 2 sets of replica groups.
61
62 // For each replic …
109:44… t_id;
110 }
111
112 // All replica id in the group have the same set of contributors.
113 for (int64_t id : rg.replica_ids()) {
114 contributing_repl …
190:51… educe for that participant.
191 // For example, for the below sequence:
192 // ar0 = all-reduce(x) replica_groups={{0,1},{2,3},{4,5},{6,7}}
193 // ar1 = all-reduce(ar0) replica_groups={{0,2 …