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_GPU_IR_EMITTER_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_H_ 18 19 #include <functional> 20 #include <map> 21 #include <memory> 22 #include <utility> 23 #include <vector> 24 25 #include "absl/strings/string_view.h" 26 #include "absl/types/span.h" 27 #include "llvm/IR/Function.h" 28 #include "llvm/IR/IRBuilder.h" 29 #include "llvm/IR/Value.h" 30 #include "tensorflow/compiler/xla/service/buffer_assignment.h" 31 #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" 32 #include "tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h" 33 #include "tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h" 34 #include "tensorflow/compiler/xla/service/gpu/ir_emitter_context.h" 35 #include "tensorflow/compiler/xla/service/gpu/kernel_thunk.h" 36 #include "tensorflow/compiler/xla/service/gpu/thunk.h" 37 #include "tensorflow/compiler/xla/service/hlo_computation.h" 38 #include "tensorflow/compiler/xla/service/hlo_instruction.h" 39 #include "tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.h" 40 #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" 41 #include "tensorflow/compiler/xla/service/llvm_ir/ir_builder_mixin.h" 42 #include "tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h" 43 #include "tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h" 44 #include "tensorflow/compiler/xla/statusor.h" 45 #include "tensorflow/compiler/xla/types.h" 46 #include "tensorflow/compiler/xla/xla_data.pb.h" 47 #include "tensorflow/core/platform/types.h" 48 49 namespace xla { 50 namespace gpu { 51 52 // Abstract base class for translating HLO graphs to LLVM IR for a GPU. 53 // 54 // There are two concrete subclasses of IrEmitter: IrEmitterNested and 55 // IrEmitterUnnested. In the unnested variety, each HLO gets its own kernel 56 // function, whereas in the nested version the whole computation is emitted as 57 // one *non-kernel* function. 58 // 59 // In XLA, kernel functions never call other kernel functions. This means that 60 // if we have a kernel -- e.g. implementing a kReduce HLO -- that wants to use 61 // an HLO computation as a "subroutine" -- e.g. the HLO computation that 62 // specifies how to reduce two elements -- then the subroutine computation must 63 // be emitted using IrEmitterNested. 64 // 65 // Fusion nodes are a special case. A fusion node is emitted using 66 // IrEmitterUnnested, but the code is generated using FusedIrEmitter, which is 67 // not a subclass of gpu::IrEmitter, and in fact is better understood as an IR 68 // generator generator. See comments on that class. 69 class IrEmitter : public DfsHloVisitorWithDefault, 70 public IrBuilderMixin<IrEmitter> { 71 public: 72 using GeneratorForOperandIrArrays = 73 std::function<std::vector<llvm_ir::IrArray>()>; 74 75 IrEmitter(const IrEmitter&) = delete; 76 IrEmitter& operator=(const IrEmitter&) = delete; 77 78 Status DefaultAction(HloInstruction* hlo) override; 79 Status HandleConstant(HloInstruction* constant) override; 80 Status HandleGetTupleElement(HloInstruction* get_tuple_element) override; 81 Status HandleConvolution(HloInstruction* convolution) override; 82 Status HandleFft(HloInstruction* fft) override; 83 Status HandleAllReduce(HloInstruction* crs) override; 84 Status HandleInfeed(HloInstruction* infeed) override; 85 Status HandleOutfeed(HloInstruction* outfeed) override; 86 Status HandleSend(HloInstruction* send) override; 87 Status HandleSendDone(HloInstruction* send_done) override; 88 Status HandleRecv(HloInstruction* recv) override; 89 Status HandleRecvDone(HloInstruction* recv_done) override; 90 Status HandleParameter(HloInstruction* parameter) override; 91 Status HandleTuple(HloInstruction* tuple) override; 92 Status HandleScatter(HloInstruction* scatter) override; 93 Status HandleTupleSelect(HloInstruction* tuple_select) override; 94 Status HandleFusion(HloInstruction* fusion) override; 95 Status HandleCall(HloInstruction* call) override; 96 Status HandleCustomCall(HloInstruction* custom_call) override; 97 Status HandleBatchNormInference(HloInstruction* batch_norm) override; 98 Status HandleBatchNormTraining(HloInstruction* batch_norm) override; 99 Status HandleBatchNormGrad(HloInstruction* batch_norm) override; 100 Status HandleAddDependency(HloInstruction* add_dependency) override; 101 FinishVisit(HloInstruction * root)102 Status FinishVisit(HloInstruction* root) override { return Status::OK(); } 103 builder()104 llvm::IRBuilder<>* builder() { return &b_; } 105 106 // Emits constants to generated LLVM IR, and also populate related 107 // inforamtion to ir_emitter_context for large-constant initializations. 108 Status EmitConstants(const HloComputation& computation); 109 110 protected: 111 // Constructs an IrEmitter with the given IrEmitter context. 112 // ir_emitter_context is owned by the caller and should outlive the IrEmitter 113 // object. 114 explicit IrEmitter(const HloModuleConfig& hlo_module_config, 115 IrEmitterContext* ir_emitter_context, bool is_nested); 116 117 // Helper for calling HloToIrBindings::GetIrArray. 118 // 119 // Gets the IrArray which contains inst. This array has metadata that makes 120 // it valid only within the IR that implements consumer. If you are 121 // implementing an HLO and want to get its own output buffer, call 122 // GetIrArray(hlo, hlo). 123 llvm_ir::IrArray GetIrArray(const HloInstruction& inst, 124 const HloInstruction& consumer, 125 const ShapeIndex& shape_index = {}) { 126 return bindings_.GetIrArray(inst, consumer, shape_index); 127 } 128 // A convenient helper for calling HloToIrBindings::GetBasePointer. 129 llvm::Value* GetBasePointer(const HloInstruction& inst, 130 ShapeIndexView shape_index = {}) const { 131 return bindings_.GetBasePointer(inst, shape_index); 132 } 133 134 // Generates the IrArray for each output of an hlo instruction and returns 135 // a vector containing such IrArrays. 136 std::vector<llvm_ir::IrArray> ConstructIrArrayForOutputs( 137 const HloInstruction& hlo); 138 139 // Emit a singlethreaded or multithreaded loop that computes every element in 140 // the result of the given HLO instruction. This produces a series of nested 141 // loops (e.g. one for each dimension of the `hlo`'s shape). The body of the 142 // inner-most loop is provided by the body_emitter function. 143 virtual Status EmitTargetElementLoop( 144 const HloInstruction& hlo, 145 const llvm_ir::ElementGenerator& body_emitter) = 0; 146 147 // Emits a call in IR to the given nested computation with the given operands 148 // and output. If no IR function has been previously emitted for the 149 // computation, also emits such a function. 150 Status EmitCallToNestedComputation(const HloComputation& nested_computation, 151 absl::Span<llvm::Value* const> operands, 152 llvm::Value* output); 153 154 // Emits an atomic operation that implements `nested_computation` in the 155 // sequentially consistent memory model. `output_address` and `source_address` 156 // are the arguments of the nested computation. For example, 157 // atomicAdd(output_address, *source_address). 158 Status EmitAtomicOperationForNestedComputation( 159 const HloComputation& nested_computation, llvm::Value* output_address, 160 llvm::Value* source_address); 161 GetNestedComputer()162 GpuElementalIrEmitter::NestedComputer GetNestedComputer() { 163 return std::bind(&IrEmitter::ComputeNestedElement, this, 164 std::placeholders::_1, std::placeholders::_2); 165 } 166 167 IrEmitterContext* ir_emitter_context_; 168 llvm::Module* module_; 169 170 // The following fields track the IR emission state. According to LLVM memory 171 // management rules, their memory is owned by the module. 172 llvm::IRBuilder<> b_; 173 174 // Mapping from HLO to its underlying LLVM value. 175 HloToIrBindings bindings_; 176 177 // Hlo configuration data used during code generation. 178 const HloModuleConfig& hlo_module_config_; 179 180 protected: 181 // Bind all argument IrArrays of `fusion` to `fused_emitter`. 182 void BindFusionArguments(const HloInstruction* fusion, 183 FusedIrEmitter* fused_emitter); 184 185 private: 186 // A helper method for EmitAtomicOperationForNestedComputation. Certain 187 // computations, such as floating-point addition and integer maximization, can 188 // be simply implemented using an LLVM atomic instruction. If "computation" is 189 // one of this kind, emits code to do that and returns true; otherwise, 190 // returns false. 191 bool MaybeEmitDirectAtomicOperation(const HloComputation& computation, 192 llvm::Value* output_address, 193 llvm::Value* source_address); 194 195 // A helper method for EmitAtomicOperationForNestedComputation. It implements 196 // binary atomic operations using atomicCAS with special handling to support 197 // small data types. 198 Status EmitAtomicOperationUsingCAS(const HloComputation& computation, 199 llvm::Value* output_address, 200 llvm::Value* source_address); 201 202 // A helper method for HandleSort(). It adds the inner comparison loop where 203 // we compare elements pointed to by 'keys_index' and 'compare_keys_index'. 204 void EmitCompareLoop(int64_t dimension_to_sort, 205 const llvm_ir::IrArray::Index& keys_index, 206 const llvm_ir::IrArray::Index& compare_keys_index, 207 const llvm_ir::IrArray& keys_array); 208 209 StatusOr<std::vector<llvm::Value*>> ComputeNestedElement( 210 const HloComputation& computation, 211 absl::Span<llvm::Value* const> parameter_elements); 212 213 // Emits an atomic operation that implements `nested_computation` in the 214 // sequentially consistent memory model. `output_address` and `source_address` 215 // are the arguments of the nested computation. For example, 216 // atomicAdd(output_address, *source_address). 217 StatusOr<llvm::Function*> EmitAtomicFunctionForNestedComputation( 218 const HloComputation& nested_computation, llvm::Type* element_ir_type); 219 220 // Map nested computations to emitted IR functions. This serves as a cache so 221 // that IrEmitter does not emit multiple functions for the same 222 // HloComputation. 223 std::map<const HloComputation*, llvm::Function*> computation_to_ir_function_; 224 }; 225 226 } // namespace gpu 227 } // namespace xla 228 229 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_H_ 230