1 /* Copyright 2017 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_XLA_SERVICE_CPU_IR_FUNCTION_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_IR_FUNCTION_H_ 18 19 #include "absl/types/span.h" 20 #include "llvm/IR/Function.h" 21 #include "llvm/IR/IRBuilder.h" 22 #include "llvm/IR/Module.h" 23 #include "llvm/IR/Value.h" 24 #include "tensorflow/compiler/xla/service/cpu/ir_emission_utils.h" 25 #include "tensorflow/compiler/xla/service/hlo_module_config.h" 26 #include "tensorflow/compiler/xla/shape_util.h" 27 #include "tensorflow/compiler/xla/statusor.h" 28 #include "tensorflow/compiler/xla/types.h" 29 30 namespace xla { 31 namespace cpu { 32 33 // IrFunction creates and encapsulates an llvm::Function, exposing methods to 34 // emitters for function and function argument access. 35 // The llvm::Function is created with the standard function signature 36 // used in the XLA CPU backend (see ir_function.cc for argument details). 37 // In addition IrFunction saves the callers IR insert point during construction, 38 // and restores it after destruction. 39 // 40 // Example usage: 41 // 42 // // Create and initialize new IrFunction. 43 // std::unique_ptr<IrFunction> compute_function(new IrFunction(...)); 44 // // Emit IR for function body using IrFunction helper methods. 45 // ... 46 // // Store reference to llvm::Function for future invocation. 47 // ir_functions.push_back(compute_function.function()); 48 // // Delete IrFunction (finalizes IR function and restores caller insertion 49 // // point). 50 // compute_function.reset(); 51 // 52 53 class IrFunction { 54 public: 55 IrFunction(const std::string& function_name, 56 llvm::Function::LinkageTypes linkage, 57 const HloModuleConfig& module_config, llvm::Module* llvm_module, 58 llvm::IRBuilder<>* b, int64_t num_dynamic_loop_bounds); 59 ~IrFunction(); 60 61 // Emit IR to read and return the set of IR values representing the dynamic 62 // loop bounds argument of this function. These bounds delimit the subset 63 // of the output that will be written by the computation's root instruction at 64 // runtime. This is used for parallel computations, where a single computation 65 // is partitioned into N calls to a function with parallel loop bounds, and 66 // then called N times in parallel with loop bounds limiting each call to 67 // producing 1/N of the output. 68 // 69 // Each element in returned vector is a pair of ir values representing the 70 // loop bounds for a specific dimension, where the first element of the pair 71 // is the dimension start index, and the second element of the pair is the 72 // dimension limit. 73 // 74 // EX: [dimension_i_index_start_ir_value, // dimension_i_index_limit_ir_value] 75 DynamicLoopBounds GetDynamicLoopBounds(); 76 77 // Returns the encapculated llvm::Function. function()78 llvm::Function* function() { return function_; } 79 80 // Get the llvm::Value* that represents this functions "retval" argument. result_arg()81 llvm::Argument* result_arg() { return result_arg_; } 82 83 // Get the xla::ExecutableRunOptions that represents this functions 84 // "run_options" argument. exec_run_options_arg()85 llvm::Value* exec_run_options_arg() { return exec_run_options_arg_; } 86 87 // Get the llvm::Value* that represents this functions parameters argument. parameters_arg()88 llvm::Value* parameters_arg() { return parameters_arg_; } 89 90 // Get the llvm::Value* that represents this functions "buffer_table" 91 // argument. buffer_table_arg()92 llvm::Value* buffer_table_arg() { return buffer_table_arg_; } 93 94 // Get the llvm::Value* that represents this functions "prof_counters" 95 // argument. profile_counters_arg()96 llvm::Value* profile_counters_arg() { return profile_counters_arg_; } 97 98 // Get the llvm::BasicBlock* that contains this function's "ret" instruction. return_block()99 llvm::BasicBlock* return_block() { return return_block_; } 100 101 // Get the llvm::Value* that represents this function's "status" argument. status_arg()102 llvm::Value* status_arg() { return status_arg_; } 103 104 private: 105 // Initialize an llvm::Function with standard signature based on arguments. 106 void Initialize(const std::string& function_name, 107 llvm::Function::LinkageTypes linkage, 108 const HloModuleConfig& module_config); 109 110 // Emit ir to read and return the ir value for the dynamic loop bound at 111 // 'offset' from the "dynamic_loop_bounds" argument of this function. 112 llvm::Value* GetDynamicLoopBound(int64_t offset); 113 114 llvm::IRBuilder<>* b_; 115 llvm::Module* llvm_module_; 116 llvm::IRBuilder<>::InsertPointGuard caller_insert_point_guard_; 117 118 int64_t num_dynamic_loop_bounds_ = 0; 119 // Encapsulated llvm::Function. 120 llvm::Function* function_; 121 // Function argument IR values. 122 llvm::Argument* result_arg_; 123 llvm::Value* exec_run_options_arg_; 124 llvm::Value* parameters_arg_; 125 llvm::Value* buffer_table_arg_; 126 llvm::Value* dynamic_loop_bounds_arg_ = nullptr; 127 llvm::Value* profile_counters_arg_; 128 llvm::Value* status_arg_; 129 // Basic block containing return. 130 llvm::BasicBlock* return_block_; 131 }; 132 133 // Returns arguments in `arguments` encoded as a single buffer, suitable for a 134 // function call. 135 llvm::Value* EncodeArrayFunctionArguments( 136 absl::Span<llvm::Value* const> arguments, absl::string_view name, 137 llvm::IRBuilder<>* b); 138 139 // Returns an array of compute function call argument ir values. 140 std::vector<llvm::Value*> GetArrayFunctionCallArguments( 141 absl::Span<llvm::Value* const> parameter_addresses, llvm::IRBuilder<>* b, 142 absl::string_view name, llvm::Value* return_value_buffer, 143 llvm::Value* exec_run_options_arg, llvm::Value* buffer_table_arg, 144 llvm::Value* status_arg, llvm::Value* profile_counters_arg); 145 146 // Emits a call to a runtime fork/join function which dispatches parallel 147 // calls to 'parallel_function' (and joins threads before returning). 148 Status EmitCallToParallelForkJoin( 149 const std::vector<llvm::Value*>& arguments, const Shape& shape, 150 const std::vector<int64_t>& dimension_partition_counts, 151 llvm::IRBuilder<>* b, llvm::Function* parallel_function, 152 const std::string& name); 153 154 } // namespace cpu 155 } // namespace xla 156 157 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_IR_FUNCTION_H_ 158