• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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_LLVM_IR_LLVM_UTIL_H_
17 #define TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_LLVM_UTIL_H_
18 
19 #include <stdint.h>
20 #include <string>
21 #include <vector>
22 
23 #include "absl/strings/string_view.h"
24 #include "absl/types/span.h"
25 #include "llvm/ADT/StringRef.h"
26 #include "llvm/IR/BasicBlock.h"
27 #include "llvm/IR/GlobalVariable.h"
28 #include "llvm/IR/IRBuilder.h"
29 #include "llvm/IR/Instructions.h"
30 #include "llvm/IR/Module.h"
31 #include "llvm/IR/Value.h"
32 #include "llvm/Support/raw_ostream.h"
33 #include "tensorflow/compiler/xla/literal.h"
34 #include "tensorflow/compiler/xla/service/hlo_instruction.h"
35 #include "tensorflow/compiler/xla/service/hlo_module_config.h"
36 #include "tensorflow/compiler/xla/types.h"
37 #include "tensorflow/compiler/xla/xla_data.pb.h"
38 #include "tensorflow/core/platform/types.h"
39 
40 namespace llvm {
41 class FastMathFlags;
42 class TargetOptions;
43 };
44 
45 namespace xla {
46 namespace llvm_ir {
47 
48 // Convert a absl::string_view to a llvm::StringRef. Note: both
49 // absl::string_view and llvm::StringRef are non-owning pointers into a
50 // string in memory. This method is used to feed strings to LLVM
51 // & Clang APIs that expect llvm::StringRef.
AsStringRef(absl::string_view str)52 inline llvm::StringRef AsStringRef(absl::string_view str) {
53   return llvm::StringRef(str.data(), str.size());
54 }
55 
56 template <typename T>
AsArrayRef(const std::vector<T> & vec)57 llvm::ArrayRef<T> AsArrayRef(const std::vector<T>& vec) {
58   return llvm::ArrayRef<T>(vec.data(), vec.size());
59 }
60 
61 template <typename T>
AsArrayRef(const absl::Span<const T> & slice)62 llvm::ArrayRef<T> AsArrayRef(const absl::Span<const T>& slice) {
63   return llvm::ArrayRef<T>(slice.data(), slice.size());
64 }
65 
66 // Dump the given LLVM entity to a string. This works for Types and Values.
67 template <typename T>
DumpToString(const T & entity)68 string DumpToString(const T& entity) {
69   std::string buffer_string;
70   llvm::raw_string_ostream ostream(buffer_string);
71   entity.print(ostream);
72   ostream.flush();
73   return buffer_string;
74 }
75 
76 // Dump the given LLVM module to a string. This requires a function distinct
77 // from DumpToString because the signatures of the print() methods for Values
78 // and Modules are slightly different.
79 string DumpModuleToString(const llvm::Module& module);
80 
81 // Constructs a human-friendly name from the given inputs.  The result is
82 // suitable for use as an llvm::Value's name.
83 //
84 // This is equivalent to
85 //
86 //   - changing the HloInstruction* to its name() (if we called that overload),
87 //   - joining all of the nonempty inputs by '.', and then
88 //   - removing all '%'s.
89 //
90 string IrName(string a);
91 string IrName(absl::string_view a, absl::string_view b);
92 string IrName(const HloInstruction* a, absl::string_view b = "");
93 
94 // Removes special characters from a function name.
95 //
96 // Note that this can cause different inputs to map to the same output, so after
97 // sanitizing a function name, you must run it through a uniquer.
98 string SanitizeFunctionName(string function_name);
99 
100 // Emits a call to the specified intrinsic with the given operands. Overloaded
101 // intrinsics (for example, "minnum") must include a type in overloaded_types
102 // for each overloaded type. Typically, overloaded intrinsics have only a single
103 // overloaded type.
104 llvm::CallInst* EmitCallToIntrinsic(
105     llvm::Intrinsic::ID intrinsic_id, absl::Span<llvm::Value* const> operands,
106     absl::Span<llvm::Type* const> overloaded_types, llvm::IRBuilder<>* b);
107 
108 // Emit float max. Emit maxnum intrinsic is fast math is disabled, or
109 // fcmp+select otherwise
110 llvm::Value* EmitFloatMax(llvm::Value* lhs_value, llvm::Value* rhs_value,
111                           llvm::IRBuilder<>* b);
112 
113 // Emit float min. Emit minnum intrinsic is fast math is disabled, or
114 // fcmp+select otherwise
115 llvm::Value* EmitFloatMin(llvm::Value* lhs_value, llvm::Value* rhs_value,
116                           llvm::IRBuilder<>* b);
117 
118 // Convenience methods for emitting a GEP instruction that indexes into a buffer
119 // (1-dimensional array), equivalent to array[index]. The type is automatically
120 // determined from the element type of the array.  The int64 index overload
121 // wraps the index in a i64 llvm::Value.
122 llvm::Value* EmitBufferIndexingGEP(llvm::Value* array, llvm::Value* index,
123                                    llvm::IRBuilder<>* b);
124 llvm::Value* EmitBufferIndexingGEP(llvm::Value* array, int64 index,
125                                    llvm::IRBuilder<>* b);
126 
127 // Returns the LLVM type which represents the given XLA primitive type.
128 llvm::Type* PrimitiveTypeToIrType(PrimitiveType element_type,
129                                   llvm::Module* module);
130 
131 // Returns the type size in bits. If "type" is a struct, it must be packed.
132 int GetSizeInBits(llvm::Type* type);
133 
134 // Returns the LLVM type which represents the given XLA shape. For example,
135 // if "shape" is [5 x [10 x f32]], the function returns [5 x [10 x float]].
136 llvm::Type* ShapeToIrType(const Shape& shape, llvm::Module* module);
137 
138 // Returns a value that represents a pointer to a global string constant that
139 // encodes the shape as a serialized protobuf.
140 StatusOr<llvm::Value*> EncodeSelfDescribingShapeConstant(const Shape& shape,
141                                                          int32* shape_size,
142                                                          llvm::IRBuilder<>* b);
143 
144 // Inverses the encoding of a Shape protobuf into an LLVM global variable.
145 //
146 // This is intended to be called from the runtime to decode the llvm::Constants
147 // that are created via ConvertShapeToSelfDescribingConstant and subsequently
148 // embedded into the program.
149 StatusOr<Shape> DecodeSelfDescribingShapeConstant(const void* shape_ptr,
150                                                   int32 size_bytes);
151 
152 // Converts a given literal to an IR Constant. Literals have known constant
153 // values at IR emission time.
154 llvm::Constant* ConvertLiteralToIrConstant(const Literal& literal,
155                                            llvm::Module* module);
156 
157 // Allocates a tile of shared memory.
158 llvm::GlobalVariable* AllocateSharedMemoryTile(llvm::Module* module,
159                                                llvm::Type* tile_type,
160                                                absl::string_view name);
161 
162 // Inserts an allocate of the requested type at the entry point of the
163 // function that the builder is currently building. The insert point
164 // of the builder is set to the same place after calling this function
165 // as before.
166 //
167 // This can be useful to avoid e.g. executing an alloca every time
168 // through a loop.
169 llvm::AllocaInst* EmitAllocaAtFunctionEntry(llvm::Type* type,
170                                             absl::string_view name,
171                                             llvm::IRBuilder<>* b,
172                                             int alignment = 0);
173 
174 // As EmitAllocaAtFunctionEntry, but allocates element_count entries
175 // instead of a single element.
176 llvm::AllocaInst* EmitAllocaAtFunctionEntryWithCount(llvm::Type* type,
177                                                      llvm::Value* element_count,
178                                                      absl::string_view name,
179                                                      llvm::IRBuilder<>* b,
180                                                      int alignment = 0);
181 
182 // Creates a basic block with the same context and function as for the
183 // builder. Inserts at the end of the function if insert_before is
184 // null.
185 llvm::BasicBlock* CreateBasicBlock(llvm::BasicBlock* insert_before,
186                                    absl::string_view name,
187                                    llvm::IRBuilder<>* b);
188 
189 // Struct with data on a conditional branch in a diamond shape created
190 // via EmitIfThenElse.
191 struct LlvmIfData {
192   // The block that has the conditional branch.
193   llvm::BasicBlock* if_block;
194 
195   // The block that is executed if the condition is true.
196   llvm::BasicBlock* true_block;
197 
198   // The block that is executed if the condition is false.
199   llvm::BasicBlock* false_block;
200 
201   // The block that follows after both the true_block and the
202   // false_block.
203   llvm::BasicBlock* after_block;
204 };
205 
206 // Inserts a diamond-shaped if-then-else construct at the current
207 // insertion point of the builder. This involves splitting the current
208 // block into two blocks, at the insertion point, and introducing a
209 // true-block and a false-block that connect the two split pieces. The
210 // true-block is executed if the condition parameter evaluates to true
211 // and otherwise the false-block is executed. If `emit_else` is false,
212 // it jumps to the after-block rather than the false-block if the
213 // condition is false, and the returned `false_block` is null.
214 //
215 // Currently the insertion point of the builder must be a well-formed
216 // block with a terminator. If you need to use this for a
217 // non-terminated block, just make the function able to do that too.
218 LlvmIfData EmitIfThenElse(llvm::Value* condition, absl::string_view name,
219                           llvm::IRBuilder<>* b, bool emit_else = true);
220 
221 // Emits a compare operation between "lhs" and "rhs" with the given predicate,
222 // and then converts the result to i8 so that it is addressable.
223 llvm::Value* EmitComparison(llvm::CmpInst::Predicate predicate,
224                             llvm::Value* lhs, llvm::Value* rhs,
225                             llvm::IRBuilder<>* b);
226 
227 // Emits a call that logs the given value with the given tag as a prefix.
228 // The provided tag and value are passed to a runtime logging call that is
229 // embedded in this translation unit when the emitted code is executed.
230 //
231 // This can be very useful for debugging generated programs in short order when
232 // developing new generated routines.
233 //
234 // Precondition: value must be an int64.
235 // Precondition: tag must be a stable pointer for the lifetime of the generated
236 // program (the constant pointer is burned in to the program).
237 void EmitLogging(const char* tag, llvm::Value* value, llvm::IRBuilder<>* b);
238 
239 // Adds alignment metadata to a load instruction using the given alignment.
240 // The alignment refers to the result of the load, not the load itself.
241 void SetAlignmentMetadataForLoad(llvm::LoadInst* load, uint64_t alignment);
242 
243 // Adds dereferenceable metadata to a load instruction using the given
244 // the number of dereferenceable bytes.
245 // Dereferenceable refers to the result of the load, not the load itself.
246 void SetDereferenceableMetadataForLoad(llvm::LoadInst* load,
247                                        uint64_t dereferenceable_bytes);
248 
249 // Tells LLVM `inst >= lower && inst < upper`. Returns `inst` for convenience.
250 llvm::Instruction* AddRangeMetadata(int64 lower, int64 upper,
251                                     llvm::Instruction* inst);
252 
253 void SetToFirstInsertPoint(llvm::BasicBlock* blk, llvm::IRBuilder<>* builder);
254 
255 void SetToLastInsertPoint(llvm::BasicBlock* blk, llvm::IRBuilder<>* builder);
256 
257 // Create a bitwise rotation of `rotand` by `rotor`.
258 llvm::Value* CreateRor(llvm::Value* rotand, llvm::Value* rotor,
259                        llvm::IRBuilder<>* builder);
260 
261 // Returns the number of bytes within the shape.
262 int64 ByteSizeOf(const Shape& shape, const llvm::DataLayout& data_layout);
263 
264 // Gets an llvm::FastMathFlags that reflects the settings in the given
265 // module config.
266 llvm::FastMathFlags GetCpuFastMathFlags(const HloModuleConfig& module_config);
267 
268 // Computes a conservative union of the metadata in "a" and "b".  For
269 // aliasing-related metadata, this means the result can be applied to
270 // instructions whose aliasing relationship can be described either by "a" *or*
271 // by "b".
272 std::map<int, llvm::MDNode*> MergeMetadata(
273     llvm::LLVMContext* context, const std::map<int, llvm::MDNode*>& a,
274     const std::map<int, llvm::MDNode*>& b);
275 
276 // Dumps out `llvm_module` to the path specified in DebugOptions, if dumping is
277 // enabled for the given HLO module.
278 //
279 // A sanitized version of `hlo_module_name` is incorporated into the file name.
280 // If `optimized` is true then a suffix of "-with-opt.ll" is used, else a suffix
281 // of "-no-opt.ll" is used.
282 void DumpIrIfEnabled(const HloModule& hlo_module,
283                      const llvm::Module& llvm_module, bool optimized);
284 
285 llvm::Function* CreateCpuFunction(llvm::FunctionType* function_type,
286                                   llvm::GlobalValue::LinkageTypes linkage,
287                                   const HloModuleConfig& module_config,
288                                   absl::string_view name, llvm::Module* module);
289 
290 // Extracts the xla_backend_extra_options from `config` and passes those that
291 // don't start with xla_ to LLVM.
292 void InitializeLLVMCommandLineOptions(const HloModuleConfig& config);
293 
294 // Zero-extends two 32-bit values to 64 bits, multiplies them, and returns the
295 // result as a pair of (low 32 bits, high 32 bits).
296 std::pair<llvm::Value*, llvm::Value*> UMulLowHigh32(llvm::IRBuilder<>* b,
297                                                     llvm::Value* src0,
298                                                     llvm::Value* src1);
299 // Splits the 64-bit integer value into its high and low 32 bits.
300 std::pair<llvm::Value*, llvm::Value*> SplitInt64ToInt32s(
301     llvm::IRBuilder<>* b, llvm::Value* value_64bits);
302 
303 // Checks whether a global variable is already created to represent a
304 // state passed between RNG calls implemented with Philox algorithm. If not,
305 // creates such a variable. Returns the global variable.
306 llvm::GlobalVariable* GetOrCreateVariableForPhiloxRngState(
307     llvm::Module* module, llvm::IRBuilder<>* b);
308 
309 // Adds a value to the global state variable each time when a RNG hlo is
310 // executed. The value of this global state variable is added to the seed
311 // of the Philox RNG algorithm so that calling the same RNG Hlo multiple times
312 // should rarely produce the same result.
313 void IncrementVariableForPhiloxRngState(int64 value, llvm::Module* module,
314                                         llvm::IRBuilder<>* b);
315 }  // namespace llvm_ir
316 }  // namespace xla
317 
318 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_LLVM_UTIL_H_
319