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_HLO_TO_IR_BINDINGS_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_HLO_TO_IR_BINDINGS_H_ 18 19 #include <unordered_map> 20 21 #include "absl/container/flat_hash_map.h" 22 #include "absl/types/span.h" 23 #include "llvm/IR/IRBuilder.h" 24 #include "llvm/IR/Value.h" 25 #include "tensorflow/compiler/xla/map_util.h" 26 #include "tensorflow/compiler/xla/service/buffer_assignment.h" 27 #include "tensorflow/compiler/xla/service/hlo_instruction.h" 28 #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" 29 30 namespace xla { 31 namespace gpu { 32 33 // This class encapsulates the bindings between HloInstructions and LLVM IR 34 // values that represent their addresses. 35 class HloToIrBindings { 36 public: HloToIrBindings(const HloModule & module,const BufferAssignment * buffer_assignment,llvm::IRBuilder<> * b,llvm::Module * llvm_module,bool is_nested)37 HloToIrBindings(const HloModule& module, 38 const BufferAssignment* buffer_assignment, 39 llvm::IRBuilder<>* b, llvm::Module* llvm_module, 40 bool is_nested) 41 : buffer_assignment_(buffer_assignment), 42 is_nested_(is_nested), 43 b_(b), 44 module_(llvm_module) {} 45 46 void EmitBasePointersForHlos( 47 absl::Span<const HloInstruction* const> io_hlos, 48 absl::Span<const HloInstruction* const> non_io_hlos); 49 50 // Rebinds the given HLO to the LLVM IR value that represent its address. 51 void BindHloToIrValue(const HloInstruction& hlo, llvm::Value* ir_value, 52 ShapeIndexView shape_index = {}); 53 54 // Unbinds all IR values that's defined in an LLVM function, e.g., function 55 // arguments and stack variables. Global variables will be kept in bindings_. 56 // 57 // This method is called after emitting code for each top-level HLO. The local 58 // IR values are out of scope at that point and should not be used. 59 void UnbindAllLocalIrValues(); 60 61 // Returns whether `hlo` is bound to an LLVM IR value. BoundToIrValue(const HloInstruction & hlo)62 bool BoundToIrValue(const HloInstruction& hlo) const { 63 return base_ptrs_.contains(&hlo); 64 } 65 GetTempBufferBase()66 llvm::Value* GetTempBufferBase() const { return temp_buffer_base_; } SetTempBufferBase(llvm::Value * v)67 void SetTempBufferBase(llvm::Value* v) { temp_buffer_base_ = v; } 68 69 // A helper method that returns the base pointer of the IrArray containing the 70 // output of "inst".at the given ShapeIndex. 71 llvm::Value* GetBasePointer(const HloInstruction& hlo, 72 ShapeIndexView shape_index = {}) const { 73 auto it = base_ptrs_.find(&hlo); 74 CHECK(it != base_ptrs_.end()) << hlo.ToString(); 75 return it->second.element(shape_index); 76 } 77 78 // Returns the IrArray which contains the output of hlo. 79 // 80 // consumer is the HLO in which this IrArray is used -- we use this to (try 81 // to) add metadata indicating that the array is invariant within consumer. 82 // 83 // To get the buffer into which hlo should write its own output, call 84 // GetIrArray(hlo, hlo). 85 llvm_ir::IrArray GetIrArray(const HloInstruction& hlo, 86 const HloInstruction& consumer, 87 const ShapeIndex& shape_index = {}); 88 89 string ToString() const; 90 91 private: 92 // Emits IR to resolve (possibly) recursive GetTupleElement instructions. 93 llvm::Value* EmitGetTupleElement(const HloInstruction* gte, 94 llvm::Value* base_ptr); 95 96 // Returns an llvm typed ir representation of 'ir_value' based on 'hlo' shape. 97 llvm::Value* GetTypedIrValue(const HloInstruction& hlo, 98 ShapeIndexView shape_index, 99 llvm::Value* ir_value); 100 101 const BufferAssignment* buffer_assignment_; 102 103 const bool is_nested_; 104 105 llvm::IRBuilder<>* b_; 106 llvm::Module* module_; 107 108 // Stores the underlying llvm::IrArray for each HloInstruction. 109 // For an instruction that generates multiple outputs, the root will be a 110 // tuple shape. The IrArray for each element output is stored in the subnode 111 // in the ShapeTree. 112 absl::flat_hash_map<const HloInstruction*, ShapeTree<llvm::Value*>> 113 base_ptrs_; 114 115 // The address of the memory block that contains all temporary buffers. 116 llvm::Value* temp_buffer_base_ = nullptr; 117 }; 118 119 // Converts `ir_value` with type i8* to a typed LLVM Value* based on `shape`. 120 llvm::Value* CastToTypedValue(const Shape& shape, llvm::Value* ir_value, 121 llvm::IRBuilder<>* b); 122 123 } // namespace gpu 124 } // namespace xla 125 126 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_HLO_TO_IR_BINDINGS_H_ 127