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/tuple_ops.h"
17
18 #include <stddef.h>
19
20 #include <string>
21 #include <vector>
22
23 #include "llvm/IR/Instructions.h"
24 #include "tensorflow/compiler/xla/service/llvm_ir/llvm_type_conversion_util.h"
25 #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h"
26 #include "tensorflow/compiler/xla/shape_util.h"
27 #include "tensorflow/compiler/xla/types.h"
28 #include "tensorflow/compiler/xla/xla_data.pb.h"
29 #include "tensorflow/core/platform/logging.h"
30
31 namespace xla {
32 namespace llvm_ir {
33
getModuleFromBuilder(llvm::IRBuilder<> * b)34 static llvm::Module* getModuleFromBuilder(llvm::IRBuilder<>* b) {
35 return b->GetInsertBlock()->getModule();
36 }
37
EmitTuple(const IrArray & tuple,absl::Span<llvm::Value * const> operands,llvm::IRBuilder<> * b)38 void EmitTuple(const IrArray& tuple, absl::Span<llvm::Value* const> operands,
39 llvm::IRBuilder<>* b) {
40 llvm::Module* module = getModuleFromBuilder(b);
41 for (size_t i = 0; i < operands.size(); ++i) {
42 auto* cast =
43 b->CreatePointerCast(operands[i], PrimitiveTypeToIrType(TUPLE, module));
44 auto* store = b->CreateStore(
45 cast,
46 b->CreateInBoundsGEP(tuple.GetBasePointeeType(), tuple.GetBasePointer(),
47 {b->getInt64(0), b->getInt64(i)}));
48 tuple.AnnotateLoadStoreInstructionWithMetadata(store);
49 }
50 }
51
EmitTuple(const IrArray & tuple,absl::Span<const IrArray> buffers,llvm::IRBuilder<> * b)52 void EmitTuple(const IrArray& tuple, absl::Span<const IrArray> buffers,
53 llvm::IRBuilder<>* b) {
54 std::vector<llvm::Value*> buffer_ptrs;
55 buffer_ptrs.reserve(buffers.size());
56 absl::c_transform(
57 buffers, std::back_inserter(buffer_ptrs),
58 [](const llvm_ir::IrArray& buffer) { return buffer.GetBasePointer(); });
59 llvm_ir::EmitTuple(tuple, buffer_ptrs, b);
60 }
61
EmitTupleAllocasAtFunctionEntry(const Shape & tuple_shape,llvm::IRBuilder<> * b)62 std::vector<llvm::Value*> EmitTupleAllocasAtFunctionEntry(
63 const Shape& tuple_shape, llvm::IRBuilder<>* b) {
64 llvm::Module* module = b->GetInsertBlock()->getModule();
65
66 llvm::IRBuilder<>::InsertPointGuard guard(*b);
67 llvm::Function* function = b->GetInsertBlock()->getParent();
68 b->SetInsertPoint(&function->getEntryBlock(),
69 function->getEntryBlock().getFirstInsertionPt());
70 CHECK(tuple_shape.IsTuple());
71 int tuple_size = tuple_shape.tuple_shapes_size();
72
73 std::vector<llvm::Value*> generated_allocas;
74 for (int i = 0; i < tuple_size; i++) {
75 const Shape& element_shape = tuple_shape.tuple_shapes(i);
76 CHECK(ShapeUtil::IsScalar(element_shape));
77 llvm::Type* type =
78 llvm_ir::PrimitiveTypeToIrType(element_shape.element_type(), module);
79 llvm::AllocaInst* alloca = b->CreateAlloca(
80 type,
81 /*ArraySize=*/nullptr, AsStringRef(absl::StrCat("tuple_element_", i)));
82 generated_allocas.push_back(alloca);
83 }
84
85 return generated_allocas;
86 }
87
EmitGetTupleElement(const Shape & target_shape,int64_t index,int alignment,llvm::Value * operand,llvm::Type * operand_pointee_type,llvm::IRBuilder<> * b)88 llvm::Value* EmitGetTupleElement(const Shape& target_shape, int64_t index,
89 int alignment, llvm::Value* operand,
90 llvm::Type* operand_pointee_type,
91 llvm::IRBuilder<>* b) {
92 CHECK(llvm::cast<llvm::PointerType>(operand->getType())
93 ->isOpaqueOrPointeeTypeMatches(operand_pointee_type));
94 llvm::Module* module = getModuleFromBuilder(b);
95 const std::vector<llvm::Value*> gep_index = {b->getInt64(0),
96 b->getInt64(index)};
97 llvm::Value* element_ptr =
98 b->CreateInBoundsGEP(operand_pointee_type, operand, gep_index);
99 llvm::Type* element_pointee_type =
100 llvm::GetElementPtrInst::getIndexedType(operand_pointee_type, gep_index);
101 llvm::LoadInst* src_buffer = b->CreateLoad(element_pointee_type, element_ptr);
102
103 // Mark the loaded pointer as dereferenceable if we know its shape.
104 if (!target_shape.IsOpaque()) {
105 SetDereferenceableMetadataForLoad(
106 src_buffer,
107 ByteSizeOf(target_shape, src_buffer->getModule()->getDataLayout()));
108 }
109 SetAlignmentMetadataForLoad(src_buffer, alignment);
110
111 llvm::Type* element_type = ShapeToIrType(target_shape, module);
112 llvm::Value* ret_val =
113 b->CreateBitCast(src_buffer, element_type->getPointerTo());
114 return ret_val;
115 }
116
117 } // namespace llvm_ir
118 } // namespace xla
119