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 an LLVM IR emitter for a nested HLO computation. `function` is 43 // the containing IR function this emitter produces IR to. See 44 // IrEmitter::IrEmitter for the meanings of other arguments. 45 IrEmitterNested(const HloModuleConfig& hlo_module_config, 46 const HloComputation& nested_computation, 47 IrEmitterContext* ir_emitter_context); 48 IrEmitterNested(const IrEmitterNested&) = delete; 49 IrEmitterNested& operator=(const IrEmitterNested&) = delete; 50 51 // Overrides the default empty implementation. Binds the given instruction 52 // "parameter" with the parameter of the IR function. 53 Status HandleParameter(HloInstruction* parameter) override; 54 GetEmittedFunction()55 llvm::Function* GetEmittedFunction() const { return emitted_function_; } 56 57 Status EmitTargetElementLoop( 58 const HloInstruction& hlo, 59 const llvm_ir::ElementGenerator& body_emitter) override; 60 61 private: 62 llvm::Function* EmitBasePointersForNestedComputation( 63 const HloComputation& nested_computation, 64 std::vector<const HloInstruction*>* io_hlos); 65 66 llvm::Function* emitted_function_; 67 }; 68 69 } // namespace gpu 70 } // namespace xla 71 72 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_NESTED_H_ 73