1 | //===- Serializer.h - MLIR SPIR-V Serializer ------------------------------===// |
2 | // |
3 | // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
4 | // See https://llvm.org/LICENSE.txt for license information. |
5 | // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | // |
9 | // This file declares the MLIR SPIR-V module to SPIR-V binary serializer. |
10 | // |
11 | //===----------------------------------------------------------------------===// |
12 | |
13 | #ifndef MLIR_LIB_TARGET_SPIRV_SERIALIZATION_SERIALIZER_H |
14 | #define MLIR_LIB_TARGET_SPIRV_SERIALIZATION_SERIALIZER_H |
15 | |
16 | #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h" |
17 | #include "mlir/IR/Builders.h" |
18 | #include "mlir/Target/SPIRV/Serialization.h" |
19 | #include "llvm/ADT/SetVector.h" |
20 | #include "llvm/ADT/SmallVector.h" |
21 | #include "llvm/Support/raw_ostream.h" |
22 | |
23 | namespace mlir { |
24 | namespace spirv { |
25 | |
26 | void encodeInstructionInto(SmallVectorImpl<uint32_t> &binary, spirv::Opcode op, |
27 | ArrayRef<uint32_t> operands); |
28 | |
29 | /// A SPIR-V module serializer. |
30 | /// |
31 | /// A SPIR-V binary module is a single linear stream of instructions; each |
32 | /// instruction is composed of 32-bit words with the layout: |
33 | /// |
34 | /// | <word-count>|<opcode> | <operand> | <operand> | ... | |
35 | /// | <------ word -------> | <-- word --> | <-- word --> | ... | |
36 | /// |
37 | /// For the first word, the 16 high-order bits are the word count of the |
38 | /// instruction, the 16 low-order bits are the opcode enumerant. The |
39 | /// instructions then belong to different sections, which must be laid out in |
40 | /// the particular order as specified in "2.4 Logical Layout of a Module" of |
41 | /// the SPIR-V spec. |
42 | class Serializer { |
43 | public: |
44 | /// Creates a serializer for the given SPIR-V `module`. |
45 | explicit Serializer(spirv::ModuleOp module, |
46 | const SerializationOptions &options); |
47 | |
48 | /// Serializes the remembered SPIR-V module. |
49 | LogicalResult serialize(); |
50 | |
51 | /// Collects the final SPIR-V `binary`. |
52 | void collect(SmallVectorImpl<uint32_t> &binary); |
53 | |
54 | #ifndef NDEBUG |
55 | /// (For debugging) prints each value and its corresponding result <id>. |
56 | void printValueIDMap(raw_ostream &os); |
57 | #endif |
58 | |
59 | private: |
60 | // Note that there are two main categories of methods in this class: |
61 | // * process*() methods are meant to fully serialize a SPIR-V module entity |
62 | // (header, type, op, etc.). They update internal vectors containing |
63 | // different binary sections. They are not meant to be called except the |
64 | // top-level serialization loop. |
65 | // * prepare*() methods are meant to be helpers that prepare for serializing |
66 | // certain entity. They may or may not update internal vectors containing |
67 | // different binary sections. They are meant to be called among themselves |
68 | // or by other process*() methods for subtasks. |
69 | |
70 | //===--------------------------------------------------------------------===// |
71 | // <id> |
72 | //===--------------------------------------------------------------------===// |
73 | |
74 | // Note that it is illegal to use id <0> in SPIR-V binary module. Various |
75 | // methods in this class, if using SPIR-V word (uint32_t) as interface, |
76 | // check or return id <0> to indicate error in processing. |
77 | |
78 | /// Consumes the next unused <id>. This method will never return 0. |
79 | uint32_t getNextID() { return nextID++; } |
80 | |
81 | //===--------------------------------------------------------------------===// |
82 | // Module structure |
83 | //===--------------------------------------------------------------------===// |
84 | |
85 | uint32_t getSpecConstID(StringRef constName) const { |
86 | return specConstIDMap.lookup(Key: constName); |
87 | } |
88 | |
89 | uint32_t getVariableID(StringRef varName) const { |
90 | return globalVarIDMap.lookup(Key: varName); |
91 | } |
92 | |
93 | uint32_t getFunctionID(StringRef fnName) const { |
94 | return funcIDMap.lookup(Key: fnName); |
95 | } |
96 | |
97 | /// Gets the <id> for the function with the given name. Assigns the next |
98 | /// available <id> if the function haven't been deserialized. |
99 | uint32_t getOrCreateFunctionID(StringRef fnName); |
100 | |
101 | void processCapability(); |
102 | |
103 | void processDebugInfo(); |
104 | |
105 | void processExtension(); |
106 | |
107 | void processMemoryModel(); |
108 | |
109 | LogicalResult processConstantOp(spirv::ConstantOp op); |
110 | |
111 | LogicalResult processSpecConstantOp(spirv::SpecConstantOp op); |
112 | |
113 | LogicalResult |
114 | processSpecConstantCompositeOp(spirv::SpecConstantCompositeOp op); |
115 | |
116 | LogicalResult |
117 | processSpecConstantOperationOp(spirv::SpecConstantOperationOp op); |
118 | |
119 | /// SPIR-V dialect supports OpUndef using spirv.UndefOp that produces a SSA |
120 | /// value to use with other operations. The SPIR-V spec recommends that |
121 | /// OpUndef be generated at module level. The serialization generates an |
122 | /// OpUndef for each type needed at module level. |
123 | LogicalResult processUndefOp(spirv::UndefOp op); |
124 | |
125 | /// Emit OpName for the given `resultID`. |
126 | LogicalResult processName(uint32_t resultID, StringRef name); |
127 | |
128 | /// Processes a SPIR-V function op. |
129 | LogicalResult processFuncOp(spirv::FuncOp op); |
130 | LogicalResult processFuncParameter(spirv::FuncOp op); |
131 | |
132 | LogicalResult processVariableOp(spirv::VariableOp op); |
133 | |
134 | /// Process a SPIR-V GlobalVariableOp |
135 | LogicalResult processGlobalVariableOp(spirv::GlobalVariableOp varOp); |
136 | |
137 | /// Process attributes that translate to decorations on the result <id> |
138 | LogicalResult processDecorationAttr(Location loc, uint32_t resultID, |
139 | Decoration decoration, Attribute attr); |
140 | LogicalResult processDecoration(Location loc, uint32_t resultID, |
141 | NamedAttribute attr); |
142 | |
143 | template <typename DType> |
144 | LogicalResult processTypeDecoration(Location loc, DType type, |
145 | uint32_t resultId) { |
146 | return emitError(loc, message: "unhandled decoration for type:" ) << type; |
147 | } |
148 | |
149 | /// Process member decoration |
150 | LogicalResult processMemberDecoration( |
151 | uint32_t structID, |
152 | const spirv::StructType::MemberDecorationInfo &memberDecorationInfo); |
153 | |
154 | //===--------------------------------------------------------------------===// |
155 | // Types |
156 | //===--------------------------------------------------------------------===// |
157 | |
158 | uint32_t getTypeID(Type type) const { return typeIDMap.lookup(Val: type); } |
159 | |
160 | Type getVoidType() { return mlirBuilder.getNoneType(); } |
161 | |
162 | bool isVoidType(Type type) const { return isa<NoneType>(Val: type); } |
163 | |
164 | /// Returns true if the given type is a pointer type to a struct in some |
165 | /// interface storage class. |
166 | bool isInterfaceStructPtrType(Type type) const; |
167 | |
168 | /// Main dispatch method for serializing a type. The result <id> of the |
169 | /// serialized type will be returned as `typeID`. |
170 | LogicalResult processType(Location loc, Type type, uint32_t &typeID); |
171 | LogicalResult processTypeImpl(Location loc, Type type, uint32_t &typeID, |
172 | SetVector<StringRef> &serializationCtx); |
173 | |
174 | /// Method for preparing basic SPIR-V type serialization. Returns the type's |
175 | /// opcode and operands for the instruction via `typeEnum` and `operands`. |
176 | LogicalResult prepareBasicType(Location loc, Type type, uint32_t resultID, |
177 | spirv::Opcode &typeEnum, |
178 | SmallVectorImpl<uint32_t> &operands, |
179 | bool &deferSerialization, |
180 | SetVector<StringRef> &serializationCtx); |
181 | |
182 | LogicalResult prepareFunctionType(Location loc, FunctionType type, |
183 | spirv::Opcode &typeEnum, |
184 | SmallVectorImpl<uint32_t> &operands); |
185 | |
186 | //===--------------------------------------------------------------------===// |
187 | // Constant |
188 | //===--------------------------------------------------------------------===// |
189 | |
190 | uint32_t getConstantID(Attribute value) const { |
191 | return constIDMap.lookup(Val: value); |
192 | } |
193 | |
194 | /// Main dispatch method for processing a constant with the given `constType` |
195 | /// and `valueAttr`. `constType` is needed here because we can interpret the |
196 | /// `valueAttr` as a different type than the type of `valueAttr` itself; for |
197 | /// example, ArrayAttr, whose type is NoneType, is used for spirv::ArrayType |
198 | /// constants. |
199 | uint32_t prepareConstant(Location loc, Type constType, Attribute valueAttr); |
200 | |
201 | /// Prepares array attribute serialization. This method emits corresponding |
202 | /// OpConstant* and returns the result <id> associated with it. Returns 0 if |
203 | /// failed. |
204 | uint32_t prepareArrayConstant(Location loc, Type constType, ArrayAttr attr); |
205 | |
206 | /// Prepares bool/int/float DenseElementsAttr serialization. This method |
207 | /// iterates the DenseElementsAttr to construct the constant array, and |
208 | /// returns the result <id> associated with it. Returns 0 if failed. Note |
209 | /// that the size of `index` must match the rank. |
210 | /// TODO: Consider to enhance splat elements cases. For splat cases, |
211 | /// we don't need to loop over all elements, especially when the splat value |
212 | /// is zero. We can use OpConstantNull when the value is zero. |
213 | uint32_t prepareDenseElementsConstant(Location loc, Type constType, |
214 | DenseElementsAttr valueAttr, int dim, |
215 | MutableArrayRef<uint64_t> index); |
216 | |
217 | /// Prepares scalar attribute serialization. This method emits corresponding |
218 | /// OpConstant* and returns the result <id> associated with it. Returns 0 if |
219 | /// the attribute is not for a scalar bool/integer/float value. If `isSpec` is |
220 | /// true, then the constant will be serialized as a specialization constant. |
221 | uint32_t prepareConstantScalar(Location loc, Attribute valueAttr, |
222 | bool isSpec = false); |
223 | |
224 | uint32_t prepareConstantBool(Location loc, BoolAttr boolAttr, |
225 | bool isSpec = false); |
226 | |
227 | uint32_t prepareConstantInt(Location loc, IntegerAttr intAttr, |
228 | bool isSpec = false); |
229 | |
230 | uint32_t prepareConstantFp(Location loc, FloatAttr floatAttr, |
231 | bool isSpec = false); |
232 | |
233 | //===--------------------------------------------------------------------===// |
234 | // Control flow |
235 | //===--------------------------------------------------------------------===// |
236 | |
237 | /// Returns the result <id> for the given block. |
238 | uint32_t getBlockID(Block *block) const { return blockIDMap.lookup(Val: block); } |
239 | |
240 | /// Returns the result <id> for the given block. If no <id> has been assigned, |
241 | /// assigns the next available <id> |
242 | uint32_t getOrCreateBlockID(Block *block); |
243 | |
244 | #ifndef NDEBUG |
245 | /// (For debugging) prints the block with its result <id>. |
246 | void printBlock(Block *block, raw_ostream &os); |
247 | #endif |
248 | |
249 | /// Processes the given `block` and emits SPIR-V instructions for all ops |
250 | /// inside. Does not emit OpLabel for this block if `omitLabel` is true. |
251 | /// `emitMerge` is a callback that will be invoked before handling the |
252 | /// terminator op to inject the Op*Merge instruction if this is a SPIR-V |
253 | /// selection/loop header block. |
254 | LogicalResult processBlock(Block *block, bool omitLabel = false, |
255 | function_ref<LogicalResult()> emitMerge = nullptr); |
256 | |
257 | /// Emits OpPhi instructions for the given block if it has block arguments. |
258 | LogicalResult emitPhiForBlockArguments(Block *block); |
259 | |
260 | LogicalResult processSelectionOp(spirv::SelectionOp selectionOp); |
261 | |
262 | LogicalResult processLoopOp(spirv::LoopOp loopOp); |
263 | |
264 | LogicalResult processBranchConditionalOp(spirv::BranchConditionalOp); |
265 | |
266 | LogicalResult processBranchOp(spirv::BranchOp branchOp); |
267 | |
268 | //===--------------------------------------------------------------------===// |
269 | // Operations |
270 | //===--------------------------------------------------------------------===// |
271 | |
272 | LogicalResult encodeExtensionInstruction(Operation *op, |
273 | StringRef extensionSetName, |
274 | uint32_t opcode, |
275 | ArrayRef<uint32_t> operands); |
276 | |
277 | uint32_t getValueID(Value val) const { return valueIDMap.lookup(Val: val); } |
278 | |
279 | LogicalResult processAddressOfOp(spirv::AddressOfOp addressOfOp); |
280 | |
281 | LogicalResult processReferenceOfOp(spirv::ReferenceOfOp referenceOfOp); |
282 | |
283 | /// Main dispatch method for serializing an operation. |
284 | LogicalResult processOperation(Operation *op); |
285 | |
286 | /// Serializes an operation `op` as core instruction with `opcode` if |
287 | /// `extInstSet` is empty. Otherwise serializes it as an extended instruction |
288 | /// with `opcode` from `extInstSet`. |
289 | /// This method is a generic one for dispatching any SPIR-V ops that has no |
290 | /// variadic operands and attributes in TableGen definitions. |
291 | LogicalResult processOpWithoutGrammarAttr(Operation *op, StringRef extInstSet, |
292 | uint32_t opcode); |
293 | |
294 | /// Dispatches to the serialization function for an operation in SPIR-V |
295 | /// dialect that is a mirror of an instruction in the SPIR-V spec. This is |
296 | /// auto-generated from ODS. Dispatch is handled for all operations in SPIR-V |
297 | /// dialect that have hasOpcode == 1. |
298 | LogicalResult dispatchToAutogenSerialization(Operation *op); |
299 | |
300 | /// Serializes an operation in the SPIR-V dialect that is a mirror of an |
301 | /// instruction in the SPIR-V spec. This is auto generated if hasOpcode == 1 |
302 | /// and autogenSerialization == 1 in ODS. |
303 | template <typename OpTy> |
304 | LogicalResult processOp(OpTy op) { |
305 | return op.emitError("unsupported op serialization" ); |
306 | } |
307 | |
308 | //===--------------------------------------------------------------------===// |
309 | // Utilities |
310 | //===--------------------------------------------------------------------===// |
311 | |
312 | /// Emits an OpDecorate instruction to decorate the given `target` with the |
313 | /// given `decoration`. |
314 | LogicalResult emitDecoration(uint32_t target, spirv::Decoration decoration, |
315 | ArrayRef<uint32_t> params = {}); |
316 | |
317 | /// Emits an OpLine instruction with the given `loc` location information into |
318 | /// the given `binary` vector. |
319 | LogicalResult emitDebugLine(SmallVectorImpl<uint32_t> &binary, Location loc); |
320 | |
321 | private: |
322 | /// The SPIR-V module to be serialized. |
323 | spirv::ModuleOp module; |
324 | |
325 | /// An MLIR builder for getting MLIR constructs. |
326 | mlir::Builder mlirBuilder; |
327 | |
328 | /// Serialization options. |
329 | SerializationOptions options; |
330 | |
331 | /// A flag which indicates if the last processed instruction was a merge |
332 | /// instruction. |
333 | /// According to SPIR-V spec: "If a branch merge instruction is used, the last |
334 | /// OpLine in the block must be before its merge instruction". |
335 | bool lastProcessedWasMergeInst = false; |
336 | |
337 | /// The <id> of the OpString instruction, which specifies a file name, for |
338 | /// use by other debug instructions. |
339 | uint32_t fileID = 0; |
340 | |
341 | /// The next available result <id>. |
342 | uint32_t nextID = 1; |
343 | |
344 | // The following are for different SPIR-V instruction sections. They follow |
345 | // the logical layout of a SPIR-V module. |
346 | |
347 | SmallVector<uint32_t, 4> capabilities; |
348 | SmallVector<uint32_t, 0> extensions; |
349 | SmallVector<uint32_t, 0> extendedSets; |
350 | SmallVector<uint32_t, 3> memoryModel; |
351 | SmallVector<uint32_t, 0> entryPoints; |
352 | SmallVector<uint32_t, 4> executionModes; |
353 | SmallVector<uint32_t, 0> debug; |
354 | SmallVector<uint32_t, 0> names; |
355 | SmallVector<uint32_t, 0> decorations; |
356 | SmallVector<uint32_t, 0> typesGlobalValues; |
357 | SmallVector<uint32_t, 0> functions; |
358 | |
359 | /// Recursive struct references are serialized as OpTypePointer instructions |
360 | /// to the recursive struct type. However, the OpTypePointer instruction |
361 | /// cannot be emitted before the recursive struct's OpTypeStruct. |
362 | /// RecursiveStructPointerInfo stores the data needed to emit such |
363 | /// OpTypePointer instructions after forward references to such types. |
364 | struct RecursiveStructPointerInfo { |
365 | uint32_t pointerTypeID; |
366 | spirv::StorageClass storageClass; |
367 | }; |
368 | |
369 | // Maps spirv::StructType to its recursive reference member info. |
370 | DenseMap<Type, SmallVector<RecursiveStructPointerInfo, 0>> |
371 | recursiveStructInfos; |
372 | |
373 | /// `functionHeader` contains all the instructions that must be in the first |
374 | /// block in the function, and `functionBody` contains the rest. After |
375 | /// processing FuncOp, the encoded instructions of a function are appended to |
376 | /// `functions`. An example of instructions in `functionHeader` in order: |
377 | /// OpFunction ... |
378 | /// OpFunctionParameter ... |
379 | /// OpFunctionParameter ... |
380 | /// OpLabel ... |
381 | /// OpVariable ... |
382 | /// OpVariable ... |
383 | SmallVector<uint32_t, 0> ; |
384 | SmallVector<uint32_t, 0> functionBody; |
385 | |
386 | /// Map from type used in SPIR-V module to their <id>s. |
387 | DenseMap<Type, uint32_t> typeIDMap; |
388 | |
389 | /// Map from constant values to their <id>s. |
390 | DenseMap<Attribute, uint32_t> constIDMap; |
391 | |
392 | /// Map from specialization constant names to their <id>s. |
393 | llvm::StringMap<uint32_t> specConstIDMap; |
394 | |
395 | /// Map from GlobalVariableOps name to <id>s. |
396 | llvm::StringMap<uint32_t> globalVarIDMap; |
397 | |
398 | /// Map from FuncOps name to <id>s. |
399 | llvm::StringMap<uint32_t> funcIDMap; |
400 | |
401 | /// Map from blocks to their <id>s. |
402 | DenseMap<Block *, uint32_t> blockIDMap; |
403 | |
404 | /// Map from the Type to the <id> that represents undef value of that type. |
405 | DenseMap<Type, uint32_t> undefValIDMap; |
406 | |
407 | /// Map from results of normal operations to their <id>s. |
408 | DenseMap<Value, uint32_t> valueIDMap; |
409 | |
410 | /// Map from extended instruction set name to <id>s. |
411 | llvm::StringMap<uint32_t> extendedInstSetIDMap; |
412 | |
413 | /// Map from values used in OpPhi instructions to their offset in the |
414 | /// `functions` section. |
415 | /// |
416 | /// When processing a block with arguments, we need to emit OpPhi |
417 | /// instructions to record the predecessor block <id>s and the values they |
418 | /// send to the block in question. But it's not guaranteed all values are |
419 | /// visited and thus assigned result <id>s. So we need this list to capture |
420 | /// the offsets into `functions` where a value is used so that we can fix it |
421 | /// up later after processing all the blocks in a function. |
422 | /// |
423 | /// More concretely, say if we are visiting the following blocks: |
424 | /// |
425 | /// ```mlir |
426 | /// ^phi(%arg0: i32): |
427 | /// ... |
428 | /// ^parent1: |
429 | /// ... |
430 | /// spirv.Branch ^phi(%val0: i32) |
431 | /// ^parent2: |
432 | /// ... |
433 | /// spirv.Branch ^phi(%val1: i32) |
434 | /// ``` |
435 | /// |
436 | /// When we are serializing the `^phi` block, we need to emit at the beginning |
437 | /// of the block OpPhi instructions which has the following parameters: |
438 | /// |
439 | /// OpPhi id-for-i32 id-for-%arg0 id-for-%val0 id-for-^parent1 |
440 | /// id-for-%val1 id-for-^parent2 |
441 | /// |
442 | /// But we don't know the <id> for %val0 and %val1 yet. One way is to visit |
443 | /// all the blocks twice and use the first visit to assign an <id> to each |
444 | /// value. But it's paying the overheads just for OpPhi emission. Instead, |
445 | /// we still visit the blocks once for emission. When we emit the OpPhi |
446 | /// instructions, we use 0 as a placeholder for the <id>s for %val0 and %val1. |
447 | /// At the same time, we record their offsets in the emitted binary (which is |
448 | /// placed inside `functions`) here. And then after emitting all blocks, we |
449 | /// replace the dummy <id> 0 with the real result <id> by overwriting |
450 | /// `functions[offset]`. |
451 | DenseMap<Value, SmallVector<size_t, 1>> deferredPhiValues; |
452 | }; |
453 | } // namespace spirv |
454 | } // namespace mlir |
455 | |
456 | #endif // MLIR_LIB_TARGET_SPIRV_SERIALIZATION_SERIALIZER_H |
457 | |