1 /* Copyright 2018 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_NESTED_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_NESTED_H_ 18 19 #include "llvm/IR/Function.h" 20 #include "tensorflow/compiler/xla/service/gpu/ir_emitter.h" 21 22 namespace xla { 23 namespace gpu { 24 25 // Emits LLVM IR for a "nested computation" into a non-kernel device function. 26 // 27 // This is used to emit code for HloComputations that don't require a separate 28 // kernel call. For example, IrEmitterNested is used to emit code for a kReduce 29 // HLO's elementwise reduction computation. Notably, IrEmitterNested is *not* 30 // used to emit code for fusion nodes -- fusion nodes use FusedIrEmitter, which 31 // is a different beast altogether. 32 // 33 // IrEmitterNested generates a non-kernel function with the following 34 // parameters: 35 // 36 // - N pointers to the buffers of each of the N parameters to the computation, 37 // - a pointer to the output buffer of the computation, and 38 // - a pointer to the top-level temp buffer. 39 // 40 class IrEmitterNested : public IrEmitter { 41 public: 42 // Constructs IrEmitterNested and emits code for global constants, and also 43 // populates related information to 'ir_emitter_context_' for large-constant 44 // initializations. Large constants don't get initializers in the generated 45 // code and so must be initialized by XLA. The value of these constants will 46 // be stored in 'content'. Constants with initializers in the generated code 47 // will have empty 'content'. 48 // 49 // The allocation index for these constants will always be -1 (i.e. doesn't 50 // correspond to any allocation). 51 static StatusOr<std::unique_ptr<IrEmitterNested>> Create( 52 const HloModuleConfig& hlo_module_config, 53 const HloComputation& nested_computation, 54 IrEmitterContext* ir_emitter_context); 55 56 IrEmitterNested(const IrEmitterNested&) = delete; 57 IrEmitterNested& operator=(const IrEmitterNested&) = delete; 58 59 // Overrides the default empty implementation. Binds the given instruction 60 // "parameter" with the parameter of the IR function. 61 Status HandleParameter(HloInstruction* parameter) override; 62 GetEmittedFunction()63 llvm::Function* GetEmittedFunction() const { return emitted_function_; } 64 65 // Generate the code for the computation passed in the constructor. 66 Status CodegenNestedComputation(); 67 68 protected: 69 Status EmitTargetElementLoop( 70 const HloInstruction& hlo, 71 const llvm_ir::ElementGenerator& body_emitter) override; 72 73 private: 74 // Constructs an LLVM IR emitter for a nested HLO computation. `function` is 75 // the containing IR function this emitter produces IR to. See 76 // IrEmitter::IrEmitter for the meanings of other arguments. 77 IrEmitterNested(const HloModuleConfig& hlo_module_config, 78 const HloComputation& nested_computation, 79 IrEmitterContext* ir_emitter_context); 80 81 // Emits constants to generated LLVM IR, and also populates related 82 // information to 'ir_emitter_context_' for large-constant initializations. 83 Status EmitConstants(const HloComputation& computation); 84 85 const HloComputation& nested_computation_; 86 llvm::Function* emitted_function_; 87 }; 88 89 } // namespace gpu 90 } // namespace xla 91 92 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_NESTED_H_ 93