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 string& function_name, llvm::Function::LinkageTypes linkage, 56 const HloModuleConfig& module_config, llvm::Module* llvm_module, 57 llvm::IRBuilder<>* b, int64 num_dynamic_loop_bounds); 58 ~IrFunction(); 59 60 // Emit ir to read and return the set of ir values representing the dynamic 61 // loop bounds argument of this function. 62 // Each element in returned vector is a pair of ir values representing 63 // the loop bounds for a specific dimension, where the first element of the 64 // pair is the dimension start index, and the second element of the pair 65 // is the dimension limit. 66 // EX: [dimension_i_index_start_ir_value, dimension_i_index_limit_ir_value] 67 // 68 DynamicLoopBounds GetDynamicLoopBounds(); 69 70 // Returns the encapculated llvm::Function. function()71 llvm::Function* function() { return function_; } 72 73 // Get the llvm::Value* that represents this functions "retval" argument. result_arg()74 llvm::Argument* result_arg() { return result_arg_; } 75 76 // Get the xla::ExecutableRunOptions that represents this functions 77 // "run_options" argument. exec_run_options_arg()78 llvm::Value* exec_run_options_arg() { return exec_run_options_arg_; } 79 80 // Get the llvm::Value* that represents this functions parameters argument. parameters_arg()81 llvm::Value* parameters_arg() { return parameters_arg_; } 82 83 // Get the llvm::Value* that represents this functions "buffer_table" 84 // argument. buffer_table_arg()85 llvm::Value* buffer_table_arg() { return buffer_table_arg_; } 86 87 // Get the llvm::Value* that represents this functions "prof_counters" 88 // argument. profile_counters_arg()89 llvm::Value* profile_counters_arg() { return profile_counters_arg_; } 90 91 private: 92 // Initialize an llvm::Function with standard signature based on arguments. 93 void Initialize(const string& function_name, 94 llvm::Function::LinkageTypes linkage, 95 const HloModuleConfig& module_config); 96 97 // Emit ir to read and return the ir value for the dynamic loop bound at 98 // 'offset' from the "dynamic_loop_bounds" argument of this function. 99 llvm::Value* GetDynamicLoopBound(int64 offset); 100 101 llvm::IRBuilder<>* b_; 102 llvm::Module* llvm_module_; 103 llvm::IRBuilder<>::InsertPointGuard caller_insert_point_guard_; 104 105 int64 num_dynamic_loop_bounds_ = 0; 106 // Encapsulated llvm::Function. 107 llvm::Function* function_; 108 // Function argument IR values. 109 llvm::Argument* result_arg_; 110 llvm::Value* exec_run_options_arg_; 111 llvm::Value* parameters_arg_; 112 llvm::Value* buffer_table_arg_; 113 llvm::Value* dynamic_loop_bounds_arg_ = nullptr; 114 llvm::Value* profile_counters_arg_; 115 }; 116 117 // Returns arguments in `arguments` encoded as a single buffer, suitable for a 118 // function call. 119 llvm::Value* EncodeArrayFunctionArguments( 120 absl::Span<llvm::Value* const> arguments, absl::string_view name, 121 llvm::IRBuilder<>* b); 122 123 // Returns an array of compute function call argument ir values. 124 std::vector<llvm::Value*> GetArrayFunctionCallArguments( 125 absl::Span<llvm::Value* const> parameter_addresses, llvm::IRBuilder<>* b, 126 absl::string_view name, llvm::Value* return_value_buffer, 127 llvm::Value* exec_run_options_arg, llvm::Value* buffer_table_arg, 128 llvm::Value* profile_counters_arg); 129 130 // Emits a call to a runtime fork/join function which dispatches parallel 131 // calls to 'parallel_function' (and joins threads before returning). 132 Status EmitCallToParallelForkJoin( 133 const std::vector<llvm::Value*>& arguments, const Shape& shape, 134 const std::vector<int64>& dimension_partition_counts, llvm::IRBuilder<>* b, 135 llvm::Function* parallel_function, const string& name); 136 137 } // namespace cpu 138 } // namespace xla 139 140 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_IR_FUNCTION_H_ 141