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
LoopEmitter(const BodyEmitter & body_emitter,const Shape & shape,llvm::IRBuilder<> * ir_builder)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
LoopEmitter(const ElementGenerator & target_element_generator,const IrArray & target_array,llvm::IRBuilder<> * ir_builder)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
MakeBodyEmitterForMultiOutputFusion(const ElementGenerator & target_element_generator,const std::vector<IrArray> & target_arrays,llvm::IRBuilder<> * ir_builder)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
LoopEmitter(const ElementGenerator & target_element_generator,tensorflow::gtl::ArraySlice<IrArray> target_arrays,llvm::IRBuilder<> * ir_builder)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
EmitIndexAndSetExitBasicBlock(tensorflow::StringPiece loop_name)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
EmitLoop(tensorflow::StringPiece loop_name)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