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/kernel_support_library.h"
17
18 #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h"
19
20 namespace xla {
ForWithStatus(absl::string_view name,llvm::Value * start,llvm::Value * end,llvm::Value * step,const std::function<Status (llvm::Value *,bool)> & for_body_generator)21 Status KernelSupportLibrary::ForWithStatus(
22 absl::string_view name, llvm::Value* start, llvm::Value* end,
23 llvm::Value* step,
24 const std::function<Status(llvm::Value*, bool)>& for_body_generator) {
25 return IfWithStatus(b_->CreateICmpSLT(start, end), [&]() -> Status {
26 TF_RETURN_IF_ERROR(for_body_generator(start, /*is_first_iteration=*/true));
27 return ForWithStatus(
28 name, b_->CreateAdd(start, step), end, step,
29 [&](llvm::Value* iv) { return for_body_generator(iv, false); });
30 });
31 }
32
ForWithStatus(absl::string_view name,llvm::Value * start,llvm::Value * end,llvm::Value * step,bool peel_first_iteration,const std::function<Status (llvm::Value *,llvm::Value *)> & for_body_generator)33 Status KernelSupportLibrary::ForWithStatus(
34 absl::string_view name, llvm::Value* start, llvm::Value* end,
35 llvm::Value* step, bool peel_first_iteration,
36 const std::function<Status(llvm::Value*, llvm::Value*)>&
37 for_body_generator) {
38 if (peel_first_iteration) {
39 return ForWithStatus(
40 name, start, end, step, true,
41 [&](llvm::Value* indvar, bool is_first_iteration) -> Status {
42 return for_body_generator(indvar, b_->getInt1(is_first_iteration));
43 });
44 } else {
45 std::unique_ptr<llvm_ir::ForLoop> loop = llvm_ir::ForLoop::EmitForLoop(
46 name, start, end, step, b_,
47 /*unroll_mode=*/unroll_mode_,
48 /*prevent_vectorization=*/prevent_vectorization_);
49 b_->SetInsertPoint(&loop->GetBodyBasicBlock()->back());
50 TF_RETURN_IF_ERROR(
51 for_body_generator(loop->GetIndVarValue(),
52 /*is_first_iteration=*/b_->CreateICmpEQ(
53 loop->GetIndVarValue(), start)));
54 llvm_ir::SetToLastInsertPoint(loop->GetExitBasicBlock(), b_);
55 return Status::OK();
56 }
57 }
58
IfWithStatus(absl::string_view name,llvm::Value * condition,const std::function<Status ()> & true_block_generator,const std::function<Status ()> & false_block_generator)59 Status KernelSupportLibrary::IfWithStatus(
60 absl::string_view name, llvm::Value* condition,
61 const std::function<Status()>& true_block_generator,
62 const std::function<Status()>& false_block_generator) {
63 llvm_ir::LlvmIfData if_data =
64 llvm_ir::EmitIfThenElse(condition, name, b_,
65 /*emit_else=*/false_block_generator != nullptr);
66 b_->SetInsertPoint(&if_data.true_block->back());
67 TF_RETURN_IF_ERROR(true_block_generator());
68 if (false_block_generator != nullptr) {
69 b_->SetInsertPoint(&if_data.false_block->back());
70 TF_RETURN_IF_ERROR(false_block_generator());
71 }
72 llvm_ir::SetToLastInsertPoint(if_data.after_block, b_);
73 return Status::OK();
74 }
75
EmitAndCallOutlinedKernel(const HloModuleConfig & module_config,llvm::IRBuilder<> * b,absl::string_view kernel_name,KernelSupportLibrary::ArgumentVector arguments,const std::function<void (KernelSupportLibrary::ArgumentVector)> & kernel_body_generator)76 void KernelSupportLibrary::EmitAndCallOutlinedKernel(
77 const HloModuleConfig& module_config, llvm::IRBuilder<>* b,
78 absl::string_view kernel_name,
79 KernelSupportLibrary::ArgumentVector arguments,
80 const std::function<void(KernelSupportLibrary::ArgumentVector)>&
81 kernel_body_generator) {
82 llvm::Module* module = b->GetInsertBlock()->getModule();
83 llvm::Function* function =
84 module->getFunction(llvm_ir::AsStringRef(kernel_name));
85
86 int64 null_arg_idx = -1;
87 std::vector<llvm::Value*> sanitized_args;
88 sanitized_args.reserve(arguments.size());
89 for (int64 i = 0, e = arguments.size(); i < e; i++) {
90 if (arguments[i]) {
91 sanitized_args.push_back(arguments[i]);
92 } else {
93 CHECK_EQ(null_arg_idx, -1);
94 null_arg_idx = i;
95 }
96 }
97
98 if (!function) {
99 VLOG(2) << "Generating kernel for " << kernel_name;
100 std::vector<llvm::Type*> arg_types;
101 std::transform(sanitized_args.begin(), sanitized_args.end(),
102 std::back_inserter(arg_types),
103 [](llvm::Value* arg) { return arg->getType(); });
104
105 auto* function_type =
106 llvm::FunctionType::get(b->getVoidTy(), arg_types, /*isVarArg=*/false);
107
108 function = llvm_ir::CreateCpuFunction(function_type,
109 llvm::GlobalValue::InternalLinkage,
110 module_config, kernel_name, module);
111
112 llvm::IRBuilder<>::InsertPointGuard guard(*b);
113
114 auto* entry_bb =
115 llvm::BasicBlock::Create(b->getContext(), "entry", function);
116 auto* return_inst = llvm::ReturnInst::Create(b->getContext(),
117 /*retVal=*/nullptr, entry_bb);
118 // Set the insert point to before return_inst.
119 b->SetInsertPoint(return_inst);
120
121 std::vector<llvm::Value*> arg_values;
122 /*
123 * clang on OSX doesn't like std::transform or range for loop here.
124 * See https://github.com/tensorflow/tensorflow/issues/15196
125 */
126 for (llvm::Function::arg_iterator arg = function->arg_begin(),
127 arg_e = function->arg_end();
128 arg != arg_e; ++arg) {
129 arg_values.push_back(arg);
130 }
131 if (null_arg_idx != -1) {
132 arg_values.insert(arg_values.begin() + null_arg_idx, nullptr);
133 }
134 kernel_body_generator(arg_values);
135 } else {
136 VLOG(3) << "Re-using kernel for " << kernel_name;
137 }
138
139 b->CreateCall(function, llvm_ir::AsArrayRef(sanitized_args));
140 }
141
142 } // namespace xla
143