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 #include "tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h" 17 18 #include <memory> 19 #include <utility> 20 21 #include "tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h" 22 #include "tensorflow/compiler/xla/shape_util.h" 23 #include "tensorflow/compiler/xla/status_macros.h" 24 #include "tensorflow/compiler/xla/types.h" 25 #include "tensorflow/compiler/xla/xla_data.pb.h" 26 #include "tensorflow/core/lib/core/errors.h" 27 #include "tensorflow/core/lib/strings/stringprintf.h" 28 #include "tensorflow/core/platform/logging.h" 29 #include "tensorflow/core/platform/protobuf.h" 30 #include "tensorflow/core/platform/types.h" 31 32 namespace xla { 33 namespace llvm_ir { 34 35 LoopEmitter::LoopEmitter(const BodyEmitter& body_emitter, const Shape& shape, 36 llvm::IRBuilder<>* ir_builder) 37 : body_emitter_(body_emitter), shape_(shape), ir_builder_(ir_builder) {} 38 39 LoopEmitter::LoopEmitter(const ElementGenerator& target_element_generator, 40 const IrArray& target_array, 41 llvm::IRBuilder<>* ir_builder) 42 : body_emitter_([=](const llvm_ir::IrArray::Index array_index) 43 -> ::tensorflow::Status { 44 // Convert target_element_generator to a BodyEmitter. 45 TF_ASSIGN_OR_RETURN(llvm::Value * target_element, 46 target_element_generator(array_index)); 47 target_array.EmitWriteArrayElement(array_index, target_element, 48 ir_builder); 49 return tensorflow::Status::OK(); 50 }), 51 shape_(target_array.GetShape()), 52 ir_builder_(ir_builder) {} 53 54 static LoopEmitter::BodyEmitter MakeBodyEmitterForMultiOutputFusion( 55 const ElementGenerator& target_element_generator, 56 const std::vector<IrArray>& target_arrays, llvm::IRBuilder<>* ir_builder) { 57 return [=](const llvm_ir::IrArray::Index array_index) { 58 TF_ASSIGN_OR_RETURN(llvm::Value * target_element, 59 target_element_generator(array_index)); 60 CHECK(target_element->getType()->isStructTy()) 61 << "This BodyEmitter is for multi-output fusion, but target element " 62 "generator does not produce values of struct type."; 63 CHECK_EQ(target_element->getType()->getStructNumElements(), 64 target_arrays.size()); 65 66 for (int64 i = 0; i < target_arrays.size(); ++i) { 67 target_arrays[i].EmitWriteArrayElement( 68 array_index, ir_builder->CreateExtractValue(target_element, i), 69 ir_builder); 70 } 71 return Status::OK(); 72 }; 73 } 74 75 LoopEmitter::LoopEmitter(const ElementGenerator& target_element_generator, 76 tensorflow::gtl::ArraySlice<IrArray> target_arrays, 77 llvm::IRBuilder<>* ir_builder) 78 : body_emitter_(MakeBodyEmitterForMultiOutputFusion( 79 target_element_generator, 80 std::vector<IrArray>(target_arrays.begin(), target_arrays.end()), 81 ir_builder)), 82 shape_(target_arrays[0].GetShape()), 83 ir_builder_(ir_builder) { 84 // Sanity check: In multi-output fusion, all shapes produced must have the 85 // same dimensions. 86 for (const IrArray& array : target_arrays) { 87 CHECK(ShapeUtil::SameDimensions(shape_, array.GetShape())); 88 } 89 } 90 91 IrArray::Index LoopEmitter::EmitIndexAndSetExitBasicBlock( 92 tensorflow::StringPiece loop_name) { 93 if (ShapeUtil::IsScalar(shape_)) { 94 // No loop needed, so set exit_bb_ to nullptr. 95 exit_bb_ = nullptr; 96 return IrArray::Index(); 97 } 98 99 // Create loop nest with one for-loop for each dimension of the target shape. 100 // Loops are added from outermost to innermost order with the ForLoopNest 101 // class so emit loops in order from most-major dimension down to most-minor 102 // dimension (of the target shape). 103 ForLoopNest loop_nest(loop_name, ir_builder_); 104 IrArray::Index array_index(shape_.dimensions_size()); 105 for (int i = 0; i < LayoutUtil::MinorToMajor(shape_).size(); ++i) { 106 int64 dimension = LayoutUtil::Major(shape_.layout(), i); 107 std::unique_ptr<ForLoop> loop = loop_nest.AddLoop( 108 /*start_index=*/0, 109 /*end_index=*/shape_.dimensions(dimension), 110 /*suffix=*/tensorflow::strings::Printf("dim.%lld", dimension)); 111 array_index[dimension] = loop->GetIndVarValue(); 112 } 113 114 // Set IR builder insertion point to the loop body basic block of the 115 // innermost loop. 116 llvm::BasicBlock* innermost_body_bb = loop_nest.GetInnerLoopBodyBasicBlock(); 117 ir_builder_->SetInsertPoint(innermost_body_bb, 118 innermost_body_bb->getFirstInsertionPt()); 119 120 // Set exit_bb_ to the exit block of the loop nest. 121 exit_bb_ = loop_nest.GetOuterLoopExitBasicBlock(); 122 CHECK_NOTNULL(exit_bb_); 123 124 return array_index; 125 } 126 127 tensorflow::Status LoopEmitter::EmitLoop(tensorflow::StringPiece loop_name) { 128 IrArray::Index array_index = EmitIndexAndSetExitBasicBlock(loop_name); 129 TF_RETURN_IF_ERROR(body_emitter_(array_index)); 130 131 // Set the insertion point of ir_builder_ to the loop exit, so that 132 // code emitted for later instructions will be correctly placed. 133 if (exit_bb_ != nullptr) { 134 ir_builder_->SetInsertPoint(exit_bb_); 135 } 136 return tensorflow::Status::OK(); 137 } 138 139 } // namespace llvm_ir 140 } // namespace xla 141
