distributed_tpu_configuration_rewrite_pass.cc
imports tensorflow/core/platform/status.himports xla/tsl/platform/errors.h
SHA-2567e8e071e1eaa4f64e8b5138333fba806b2fe6bc4696c6c75dc0b66a7bfbdd982
MaleculeH(Db)Md(In₄)
Evidence
1/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
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================================================================ …
14:77… ==*/
15
16// Configuration for distributed TPU jobs
17
18#include "tensorflow/core/tpu/graph_rewrite/distributed_tpu_configuration_rewrite_pass.h"
19
20#include <cstdint>
21#include <string>
22#include <vector>
23
24#include "absl/log/log.h"
25#include "absl/status/status.h"
26#include "xla/status_macros.h"
27#include "xla/tsl/platform/errors.h"
28#include "xla/tsl/platform/logging.h" // IWYU pragma: keep
29#include "xla/tsl/platform/statusor.h"
30#include "tensorflow/core/common_runtime/optimization_registry.h"
31#include "tensorflow/core/framework/device.h"
32#include "tensorflow/core/framework/node_def.pb.h"
33#include "tensorflow/core/framework/node_def_util.h"
34#include "tensorflow/core/framework/types.pb.h"
35#include "tensorflow/core/graph/graph.h"
36#include "tensorflow/core/graph/graph_node_util.h"
37#include "tensorflow/core/platform/status.h"
38#include "tensorflow/core/tpu/graph_rewrite/distributed_tpu_rewrite_helpers.h"
39#include "tensorflow/core/tpu/kernels/tpu_compile_op_options.h"
40#include "tensorflow/core/tpu/tpu_init_mode.h"
41#include "tensorflow/core/util/dump_graph.h"
42
43namespace tensorflow {
44namespace {
45
46constexpr char kIdentityOp[] = "Identity";
47constexpr char kConfigureOp[] = "ConfigureDistributedTPU";
48constexpr char kInternalConfigureOp[] = "_ConfigureDistributedTPU";
49constexpr char kWaitOp[] = "_WaitForDistributedTPU";
50constexpr char kHostConfigureOp[] = "_InitializeHostForDistributedTPU";
51constexpr char kGlobalTPUArrayOp[] = "_SetGlobalTPUArray";
52constexpr char kShutdownOp[] = "ShutdownDistributedTPU";
53constexpr char kInternalShutdownOp[] = "_ShutdownDistributedTPU";
54constexpr char kHostDisconnectOp[] = "_DisconnectHostFromDistributedTPUSystem";
55constexpr char kE …