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 "llvm/IR/Function.h" 20 #include "llvm/IR/IRBuilder.h" 21 #include "llvm/IR/Module.h" 22 #include "llvm/IR/Value.h" 23 #include "tensorflow/compiler/xla/service/cpu/ir_emission_utils.h" 24 #include "tensorflow/compiler/xla/shape_util.h" 25 #include "tensorflow/compiler/xla/statusor.h" 26 #include "tensorflow/compiler/xla/types.h" 27 #include "tensorflow/core/lib/gtl/array_slice.h" 28 29 namespace xla { 30 namespace cpu { 31 32 // IrFunction creates and encapsulates an llvm::Function, exposing methods to 33 // emitters for function and function argument access. 34 // The llvm::Function is created with the standard function signature 35 // used in the XLA CPU backend (see ir_function.cc for argument details). 36 // In addtion IrFunction saves the callers IR insert point during contruction, 37 // and restores it after desctruction. 38 // 39 // Example usage: 40 // 41 // // Create and initialize new IrFunction. 42 // std::unique_ptr<IrFunction> compute_function(new IrFunction(...)); 43 // // Emit IR for function body using IrFunction helper methods. 44 // ... 45 // // Store reference to llvm::Function for future invocation. 46 // ir_functions.push_back(compute_function.function()); 47 // // Delete IrFunction (finalizes IR function and restores caller insertion 48 // // point). 49 // compute_function.reset(); 50 // 51 52 class IrFunction { 53 public: 54 IrFunction(const string& function_name, llvm::Function::LinkageTypes linkage, 55 const bool optimize_for_size_requested, 56 const bool enable_fast_math, llvm::Module* llvm_module, 57 llvm::IRBuilder<>* ir_builder, 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. 71 llvm::Function* function() { return function_; } 72 73 // Get the llvm::Value* that represents this functions "retval" argument. 74 llvm::Argument* result_arg() { return result_arg_; } 75 76 // Get the xla::ExecutableRunOptions that represents this functions 77 // "run_options" argument. 78 llvm::Value* exec_run_options_arg() { return exec_run_options_arg_; } 79 80 // Get the llvm::Value* that represents this functions parameters argument. 81 llvm::Value* parameters_arg() { return parameters_arg_; } 82 83 // Get the llvm::Value* that represents this functions "temps" argument. 84 llvm::Value* temp_buffers_arg() { return temp_buffers_arg_; } 85 86 // Get the llvm::Value* that represents this functions "prof_counters" 87 // argument. 88 llvm::Value* profile_counters_arg() { return profile_counters_arg_; } 89 90 private: 91 // Initialize an llvm::Function with standard signature based on arguments. 92 void Initialize(const string& function_name, 93 llvm::Function::LinkageTypes linkage, 94 bool optimize_for_size_requested, bool enable_fast_math); 95 96 // Emit ir to read and return the ir value for the dynamic loop bound at 97 // 'offset' from the "dynamic_loop_bounds" argument of this function. 98 llvm::Value* GetDynamicLoopBound(int64 offset); 99 100 llvm::IRBuilder<>* ir_builder_; 101 llvm::Module* llvm_module_; 102 llvm::IRBuilder<>::InsertPointGuard caller_insert_point_guard_; 103 104 int64 num_dynamic_loop_bounds_ = 0; 105 // Encapsulated llvm::Function. 106 llvm::Function* function_; 107 // Function argument IR values. 108 llvm::Argument* result_arg_; 109 llvm::Value* exec_run_options_arg_; 110 llvm::Value* parameters_arg_; 111 llvm::Value* temp_buffers_arg_; 112 llvm::Value* dynamic_loop_bounds_arg_ = nullptr; 113 llvm::Value* profile_counters_arg_; 114 }; 115 116 // Returns an array of compute function call argument ir values. 117 std::vector<llvm::Value*> GetArrayFunctionCallArguments( 118 tensorflow::gtl::ArraySlice<llvm::Value*> parameter_addresses, 119 llvm::IRBuilder<>* ir_builder, tensorflow::StringPiece name, 120 llvm::Value* return_value_buffer, llvm::Value* exec_run_options_arg, 121 llvm::Value* temp_buffers_arg, llvm::Value* profile_counters_arg); 122 123 // Emits a call to a runtime fork/join function which dispatches parallel 124 // calls to 'parallel_function' (and joins threads before returning). 125 Status EmitCallToParallelForkJoin( 126 const std::vector<llvm::Value*>& arguments, const Shape& shape, 127 const std::vector<int64>& dimension_partition_counts, 128 llvm::IRBuilder<>* ir_builder, llvm::Function* parallel_function, 129 const string& name); 130 131 } // namespace cpu 132 } // namespace xla 133 134 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_IR_FUNCTION_H_ 135
