1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 // Rewrites computations generated by the tpu.replicate() Python code into 17 // TPUReplicate operators. 18 // 19 // The tpu.replicate() does two main things: 20 // a) marks operators that make up a TPU computation with the attribute 21 // _tpu_replicate=XYZ, where XYZ is a unique key. 22 // b) adds TPUReplicatedInput and TPUReplicatedOutput nodes to represent 23 // replicated inputs. These nodes are not marked with the _tpu_replicate 24 // attribute. 25 26 #ifndef TENSORFLOW_CORE_TPU_GRAPH_REWRITES_ENCAPSULATE_TPU_COMPUTATIONS_PASS_H_ 27 #define TENSORFLOW_CORE_TPU_GRAPH_REWRITES_ENCAPSULATE_TPU_COMPUTATIONS_PASS_H_ 28 29 #include "tensorflow/compiler/jit/encapsulate_util.h" 30 #include "tensorflow/core/common_runtime/optimization_registry.h" 31 #include "tensorflow/core/graph/graph.h" 32 33 namespace tensorflow { 34 35 // Encapsulates nodes marked with the _tpu_replicate attribute into 36 // TPUReplicate operators. 37 class EncapsulateTPUComputationsPass : public GraphOptimizationPass { 38 public: 39 Status Run(const GraphOptimizationPassOptions& options) override; 40 41 // The following methods are public only for unit tests. 42 43 // This pass has two stages: 44 // a) first, we call the EncapsulateSubgraphsPass to encapsulate all nodes 45 // marked with the same _tpu_replicate attribute into functions. These 46 // functions contain the computations to be passed to TPUReplicate. During 47 // encapsulation, we sort the arguments into the order expected by 48 // TPUReplicate. 49 static Status Encapsulate(std::unique_ptr<Graph>* graph, 50 FunctionLibraryDefinition* flib_def); 51 52 // b) we rewrite the function calls generated in phase (a) into TPUReplicate 53 // operators. We also flatten the TPUReplicatedInput and 54 // TPUReplicatedOutput replicated input and output nodes of the function 55 // call into the replicated input and outputs of the TPUReplicate operator. 56 static Status BuildTPUReplicateOps(Graph* graph); 57 }; 58 59 // Graph optimization pass that calls `ExtractOutsideCompilation` for all XLA 60 // computation nodes. 61 class ExtractOutsideCompilationPass : public GraphOptimizationPass { 62 public: 63 Status Run(const GraphOptimizationPassOptions& options) override; 64 65 static Status ProcessHeadTailOutsideCompilation( 66 const string& outside_compilation_attr_name, int* lifted_arg_count, 67 std::unordered_map<string, XlaClusterInfo>* clusters, Graph* g, 68 FunctionLibraryRuntime* flr, FunctionLibraryDefinition* fld); 69 }; 70 71 } // namespace tensorflow 72 73 #endif // TENSORFLOW_CORE_TPU_GRAPH_REWRITES_ENCAPSULATE_TPU_COMPUTATIONS_PASS_H_ 74