1 /* Copyright 2021 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 #ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_PASSES_H_ 17 #define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_PASSES_H_ 18 19 #include <memory> 20 21 #include "mlir/Pass/Pass.h" // from @llvm-project 22 #include "mlir/Transforms/DialectConversion.h" // from @llvm-project 23 #include "tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h" 24 #include "tensorflow/compiler/mlir/tfrt/transforms/tpu_passes.h" 25 26 namespace mlir { 27 class PassManager; 28 } 29 30 namespace tensorflow { 31 32 namespace tfrt_compiler { 33 34 // Create a pass to set shape_invariant attribute for all tf.While ops. 35 std::unique_ptr<mlir::OperationPass<mlir::FuncOp>> 36 CreateSetShapeInvariantInWhileOps(); 37 38 // Create a pass to insert kernels that copy fallback tensors when they are 39 // passed to multiple threads, to avoid atomic contention on their refcounts. 40 std::unique_ptr<mlir::OperationPass<mlir::FuncOp>> 41 CreateInsertFallbackTensorCopyPass(); 42 43 // Create a pass to reorder tf.Assert ops or tf.If ops that contains only 44 // tf.Assert ops to the end of the function, to avoid unnecessary control 45 // dependencies to other ops. 46 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> 47 CreateReorderTfAssertPass(); 48 49 // Create a pass to optimize the side-effect of control flow ops. eg. if both 50 // branches of a tf.If op contains only non-side-effecting ops, its 51 // `is_stateless` attribute will be set to true. 52 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> 53 CreateOptimizeTfControlFlowSideEffectPass(); 54 55 // Create a pass to remove tf.If ops' operands that are produced by tf.Const 56 // ops. 57 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> 58 CreateRemoveTfIfConstArgsPass(); 59 60 // Create a pass to merge non-side-effecting tf.If ops that have the same 61 // operands. 62 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateMergeTfIfOpsPass(); 63 64 // Create a pass to deduplicate the function invoked by tf.BatchFunction with 65 // the same shared_name. 66 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> 67 CreateDeduplicateFunctionsInovkedByBatchFunctionPass(); 68 69 } // namespace tfrt_compiler 70 71 class CoreRTConverter; 72 73 // Create a pass that rewrites tf_saved_model dialect's ops according to TFRT's 74 // requirements. 75 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> 76 CreateLowerTFSavedModelPass(bool hoist_invariant_ops); 77 78 // Create a pass that converts ref variables to resource variables in a limited 79 // number of cases. 80 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> 81 CreateConvertReferenceVariableToResourceVariablePass(); 82 83 // Run *ToCoreRTConversionPassRun as free functions. Useful for 84 // reusing the pass logic in a custom pass with additional conversions. 85 mlir::LogicalResult TFSavedModelToCoreRTConversionPassRun( 86 mlir::MLIRContext* context, mlir::FuncOp func, 87 mlir::ConversionTarget* target, mlir::OwningRewritePatternList* patterns, 88 CoreRTConverter* corert_converter); 89 90 // Create an operation pass that converts each tfrt_dist.remote_execute_func op 91 // into a combination of tfrt_dist.register_tfrt_function op and 92 // tfrt_dist.remote_execute op. 93 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> 94 CreateDistRemoteRunEncapsulatePass(); 95 96 // Create an operation pass that removes the device attribute from every 97 // corert.executeop. 98 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> 99 CreateRemoveDeviceAttributePass(); 100 101 // Create an operation pass that inserts corert.transfer op to make sure any 102 // argument of any op is on the same device of the op itself. 103 std::unique_ptr<mlir::FunctionPass> CreateCrossDeviceTransferPass(); 104 105 struct TfrtPipelineOptions 106 : public mlir::PassPipelineOptions<TfrtPipelineOptions> { 107 Option<std::string> default_device{ 108 *this, "default-device", llvm::cl::desc("default device assignment"), 109 llvm::cl::init("/job:localhost/replica:0/task:0/device:CPU:0")}; 110 Option<bool> enable_optimizer{ 111 *this, "enable-optimizer", 112 llvm::cl::desc("run optimization passes on corert dialect"), 113 llvm::cl::init(false)}; 114 Option<bool> decompose_resource_ops{ 115 *this, "decompose-resource-ops", 116 llvm::cl::desc("decompose composite resource ops into ReadVariableOp and " 117 "non-resource ops. This is currently used in TFRT " 118 "savedmodel pipeline."), 119 llvm::cl::init(false)}; 120 Option<std::string> force_data_format{ 121 *this, "force-data-format", 122 llvm::cl::desc("force data format for all layout sensitive operations")}; 123 // TODO(tfrt-devs): consider making compiler to figure out whether to fold 124 // transpose or not instead of exposing the specific option. 125 Option<bool> skip_fold_transpose_in_ops{ 126 *this, "skip-fold-transpose-in-ops", 127 llvm::cl::desc("Skip folding transpose operands in Ops which can support " 128 "different layouts.")}; 129 Option<bool> target_tpu{*this, "target-tpu", 130 llvm::cl::desc("target TPU programs if true"), 131 llvm::cl::init(false)}; 132 Option<bool> tpu_use_core_selector{ 133 *this, "tpu-use-core-selector", 134 llvm::cl::desc("If true, use ServingCoreSelector to pick TPU core. " 135 "Otherwise, use the assigned core. Currently we use " 136 "core selector for Servo serving use cases."), 137 llvm::cl::init(true)}; 138 Option<bool> tpu_use_bundled_transfer{ 139 *this, "tpu-use-bundled-transfer", 140 llvm::cl::desc("If true, use BundledTransferToTpuOp to transfer " 141 "variables and input tensors to TPU."), 142 llvm::cl::init(true)}; 143 Option<bool> tpu_lower_to_fallback{ 144 *this, "tpu-lower-to-fallback", 145 llvm::cl::desc("If true, lower an TF op that's placed on TPU device " 146 "to be executed by tfrt_fallback.execute."), 147 llvm::cl::init(true)}; 148 // TODO(b/194081364): remove this option once we unify servo TPU serving 149 // result transfer behavior. 150 Option<bool> tpu_transfer_result_to_host{ 151 *this, "tpu-transfer-result-to-host", 152 llvm::cl::desc("If true, transfer the result of tpurt.execute from TPU " 153 "to host."), 154 llvm::cl::init(true)}; 155 Option<bool> enable_native_ops{ 156 *this, "enable-native-ops", 157 llvm::cl::desc( 158 "If true, native ops will be used on an opt-in basis instead of " 159 "fallback ops. If false, no native ops are used."), 160 llvm::cl::init(true)}; 161 Option<bool> func_use_fallback_tensor{ 162 *this, "func-use-fallback-tensor", 163 llvm::cl::desc( 164 "If true, use TF tensor as input/output types in func (and other " 165 "control flow) ops."), 166 llvm::cl::init(false)}; 167 168 Option<bool> hoist_invariant_ops{ 169 *this, "hoist-invariant-ops", 170 llvm::cl::desc("If true, invariant ops in savedmodels will be hoisted " 171 "out to run during loading."), 172 llvm::cl::init(false)}; 173 174 Option<uint64_t> cost_threshold{ 175 *this, "tfrt-cost-threshold", 176 llvm::cl::desc( 177 "The cost threshold to decide whether a sequence of operations is " 178 "cheap, and then whether it can be executed inline."), 179 llvm::cl::init(1)}; 180 181 Option<int64_t> upper_cost_threshold{ 182 *this, "tfrt-upper-cost-threshold", 183 llvm::cl::desc( 184 "The threshold to limit the merging of dependent sequence."), 185 llvm::cl::init(-1)}; 186 187 Option<bool> merge_inter_dependent_streams{ 188 *this, "tfrt-merge-inter-dependent-streams", 189 llvm::cl::desc("If true, streams with inter data depenedencies will be " 190 "preferred to be merged for inline execution."), 191 llvm::cl::init(false)}; 192 193 // A set of flags to control auto-fusion: automatic clustering of Tensorflow 194 // operations and compiling outlined regions using MLIR based compilation 195 // stack. 196 // 197 // WARNING: These flags are experimental and are intended for manual testing 198 // of different auto-fusion strategies. They will be removed in the future. 199 200 ListOption<std::string> auto_fusion_oplist{ 201 *this, "auto-fusion-oplist", 202 llvm::cl::desc("A list of Tensorflow operations to cluster together for " 203 "JIT compilation. Alternatively use 'tier1', ..., 'all' " 204 "to allow clustering for all operations included in the " 205 "given clustering tier."), 206 llvm::cl::MiscFlags::CommaSeparated}; 207 208 Option<int> auto_fusion_min_cluster_size{ 209 *this, "auto-fusion-min-cluster-size", 210 llvm::cl::desc("Minimum size of the cluster that should be outlined for " 211 "compilation"), 212 llvm::cl::init(2)}; 213 }; 214 215 // Create a pass that converts MLIR TF dialect to MLIR TFRT dialect. 216 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> 217 CreateTfToTfrtConversionPass(const TfrtPipelineOptions& options); 218 219 // Creates a pipeline of passes that lowers MLIR TF Executor dialect to TF 220 // dialect for CoreRT purposes. 221 void CreateTFExecutorToTFPipeline(mlir::OpPassManager& pm, 222 const TfrtPipelineOptions& options); 223 224 // Creates a pipeline of passes that lowers MLIR TF dialect from tf.function to 225 // TFRT dialect. SavedModel related conversions are not included. 226 void CreateTfExecutorToTfrtPipeline(mlir::PassManager& pm, 227 const TfrtPipelineOptions& options); 228 229 } // namespace tensorflow 230 231 #endif // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_PASSES_H_ 232