512 lines · c
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-exception6//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_H14#define MLIR_LIB_TARGET_SPIRV_SERIALIZATION_SERIALIZER_H15 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 23namespace mlir {24namespace spirv {25 26void 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; each32/// 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 the38/// instruction, the 16 low-order bits are the opcode enumerant. The39/// instructions then belong to different sections, which must be laid out in40/// the particular order as specified in "2.4 Logical Layout of a Module" of41/// the SPIR-V spec.42class Serializer {43public: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 NDEBUG55 /// (For debugging) prints each value and its corresponding result <id>.56 void printValueIDMap(raw_ostream &os);57#endif58 59private: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 entity62 // (header, type, op, etc.). They update internal vectors containing63 // different binary sections. They are not meant to be called except the64 // top-level serialization loop.65 // * prepare*() methods are meant to be helpers that prepare for serializing66 // certain entity. They may or may not update internal vectors containing67 // different binary sections. They are meant to be called among themselves68 // 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. Various75 // 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 structure83 //===--------------------------------------------------------------------===//84 85 uint32_t getSpecConstID(StringRef constName) const {86 return specConstIDMap.lookup(constName);87 }88 89 uint32_t getVariableID(StringRef varName) const {90 return globalVarIDMap.lookup(varName);91 }92 93 uint32_t getFunctionID(StringRef fnName) const {94 return funcIDMap.lookup(fnName);95 }96 97 /// Gets the <id> for the function with the given name. Assigns the next98 /// 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 LogicalResult processExtension();106 107 void processMemoryModel();108 109 LogicalResult processConstantOp(spirv::ConstantOp op);110 111 LogicalResult processConstantCompositeReplicateOp(112 spirv::EXTConstantCompositeReplicateOp op);113 114 LogicalResult processSpecConstantOp(spirv::SpecConstantOp op);115 116 LogicalResult117 processSpecConstantCompositeOp(spirv::SpecConstantCompositeOp op);118 119 LogicalResult processSpecConstantCompositeReplicateOp(120 spirv::EXTSpecConstantCompositeReplicateOp op);121 122 LogicalResult123 processSpecConstantOperationOp(spirv::SpecConstantOperationOp op);124 125 LogicalResult processGraphConstantARMOp(spirv::GraphConstantARMOp op);126 127 /// SPIR-V dialect supports OpUndef using spirv.UndefOp that produces a SSA128 /// value to use with other operations. The SPIR-V spec recommends that129 /// OpUndef be generated at module level. The serialization generates an130 /// OpUndef for each type needed at module level.131 LogicalResult processUndefOp(spirv::UndefOp op);132 133 /// Emit OpName for the given `resultID`.134 LogicalResult processName(uint32_t resultID, StringRef name);135 136 /// Processes a SPIR-V function op.137 LogicalResult processFuncOp(spirv::FuncOp op);138 LogicalResult processFuncParameter(spirv::FuncOp op);139 140 /// Processes a SPIR-V GraphARM op.141 LogicalResult processGraphARMOp(spirv::GraphARMOp op);142 143 /// Processes a SPIR-V GraphEntryPointARM op.144 LogicalResult processGraphEntryPointARMOp(spirv::GraphEntryPointARMOp op);145 146 /// Processes a SPIR-V GraphOutputsARMOp op.147 LogicalResult processGraphOutputsARMOp(spirv::GraphOutputsARMOp op);148 149 LogicalResult processVariableOp(spirv::VariableOp op);150 151 /// Process a SPIR-V GlobalVariableOp152 LogicalResult processGlobalVariableOp(spirv::GlobalVariableOp varOp);153 154 /// Process attributes that translate to decorations on the result <id>155 LogicalResult processDecorationAttr(Location loc, uint32_t resultID,156 Decoration decoration, Attribute attr);157 LogicalResult processDecoration(Location loc, uint32_t resultID,158 NamedAttribute attr);159 160 template <typename DType>161 LogicalResult processTypeDecoration(Location loc, DType type,162 uint32_t resultId) {163 return emitError(loc, "unhandled decoration for type:") << type;164 }165 166 /// Process member decoration167 LogicalResult processMemberDecoration(168 uint32_t structID,169 const spirv::StructType::MemberDecorationInfo &memberDecorationInfo);170 171 //===--------------------------------------------------------------------===//172 // Types173 //===--------------------------------------------------------------------===//174 175 uint32_t getTypeID(Type type) const { return typeIDMap.lookup(type); }176 177 Type getVoidType() { return mlirBuilder.getNoneType(); }178 179 bool isVoidType(Type type) const { return isa<NoneType>(type); }180 181 /// Returns true if the given type is a pointer type to a struct in some182 /// interface storage class.183 bool isInterfaceStructPtrType(Type type) const;184 185 /// Main dispatch method for serializing a type. The result <id> of the186 /// serialized type will be returned as `typeID`.187 LogicalResult processType(Location loc, Type type, uint32_t &typeID);188 LogicalResult processTypeImpl(Location loc, Type type, uint32_t &typeID,189 SetVector<StringRef> &serializationCtx);190 191 /// Method for preparing basic SPIR-V type serialization. Returns the type's192 /// opcode and operands for the instruction via `typeEnum` and `operands`.193 LogicalResult prepareBasicType(Location loc, Type type, uint32_t resultID,194 spirv::Opcode &typeEnum,195 SmallVectorImpl<uint32_t> &operands,196 bool &deferSerialization,197 SetVector<StringRef> &serializationCtx);198 199 LogicalResult prepareFunctionType(Location loc, FunctionType type,200 spirv::Opcode &typeEnum,201 SmallVectorImpl<uint32_t> &operands);202 203 LogicalResult prepareGraphType(Location loc, GraphType type,204 spirv::Opcode &typeEnum,205 SmallVectorImpl<uint32_t> &operands);206 207 //===--------------------------------------------------------------------===//208 // Constant209 //===--------------------------------------------------------------------===//210 211 uint32_t getConstantID(Attribute value) const {212 return constIDMap.lookup(value);213 }214 215 uint32_t getConstantCompositeReplicateID(216 std::pair<Attribute, Type> valueTypePair) const {217 return constCompositeReplicateIDMap.lookup(valueTypePair);218 }219 220 /// Main dispatch method for processing a constant with the given `constType`221 /// and `valueAttr`. `constType` is needed here because we can interpret the222 /// `valueAttr` as a different type than the type of `valueAttr` itself; for223 /// example, ArrayAttr, whose type is NoneType, is used for spirv::ArrayType224 /// constants.225 uint32_t prepareConstant(Location loc, Type constType, Attribute valueAttr);226 227 /// Prepares array attribute serialization. This method emits corresponding228 /// OpConstant* and returns the result <id> associated with it. Returns 0 if229 /// failed.230 uint32_t prepareArrayConstant(Location loc, Type constType, ArrayAttr attr);231 232 /// Prepares bool/int/float DenseElementsAttr serialization. This method233 /// iterates the DenseElementsAttr to construct the constant array, and234 /// returns the result <id> associated with it. Returns 0 if failed. Note235 /// that the size of `index` must match the rank.236 /// TODO: Consider to enhance splat elements cases. For splat cases,237 /// we don't need to loop over all elements, especially when the splat value238 /// is zero. We can use OpConstantNull when the value is zero.239 uint32_t prepareDenseElementsConstant(Location loc, Type constType,240 DenseElementsAttr valueAttr, int dim,241 MutableArrayRef<uint64_t> index);242 243 /// Prepares scalar attribute serialization. This method emits corresponding244 /// OpConstant* and returns the result <id> associated with it. Returns 0 if245 /// the attribute is not for a scalar bool/integer/float value. If `isSpec` is246 /// true, then the constant will be serialized as a specialization constant.247 uint32_t prepareConstantScalar(Location loc, Attribute valueAttr,248 bool isSpec = false);249 250 uint32_t prepareConstantBool(Location loc, BoolAttr boolAttr,251 bool isSpec = false);252 253 uint32_t prepareConstantInt(Location loc, IntegerAttr intAttr,254 bool isSpec = false);255 256 uint32_t getGraphConstantARMId(Attribute value) const {257 return graphConstIDMap.lookup(value);258 }259 260 uint32_t prepareGraphConstantId(Location loc, Type graphConstType,261 IntegerAttr intAttr);262 263 uint32_t prepareConstantFp(Location loc, FloatAttr floatAttr,264 bool isSpec = false);265 266 /// Prepares `spirv.EXTConstantCompositeReplicateOp` serialization. This267 /// method emits OpConstantCompositeReplicateEXT and returns the result <id>268 /// associated with it.269 uint32_t prepareConstantCompositeReplicate(Location loc, Type resultType,270 Attribute valueAttr);271 272 //===--------------------------------------------------------------------===//273 // Control flow274 //===--------------------------------------------------------------------===//275 276 /// Returns the result <id> for the given block.277 uint32_t getBlockID(Block *block) const { return blockIDMap.lookup(block); }278 279 /// Returns the result <id> for the given block. If no <id> has been assigned,280 /// assigns the next available <id>281 uint32_t getOrCreateBlockID(Block *block);282 283#ifndef NDEBUG284 /// (For debugging) prints the block with its result <id>.285 void printBlock(Block *block, raw_ostream &os);286#endif287 288 /// Processes the given `block` and emits SPIR-V instructions for all ops289 /// inside. Does not emit OpLabel for this block if `omitLabel` is true.290 /// `emitMerge` is a callback that will be invoked before handling the291 /// terminator op to inject the Op*Merge instruction if this is a SPIR-V292 /// selection/loop header block.293 LogicalResult processBlock(Block *block, bool omitLabel = false,294 function_ref<LogicalResult()> emitMerge = nullptr);295 296 /// Emits OpPhi instructions for the given block if it has block arguments.297 LogicalResult emitPhiForBlockArguments(Block *block);298 299 LogicalResult processSelectionOp(spirv::SelectionOp selectionOp);300 301 LogicalResult processLoopOp(spirv::LoopOp loopOp);302 303 LogicalResult processBranchConditionalOp(spirv::BranchConditionalOp);304 305 LogicalResult processBranchOp(spirv::BranchOp branchOp);306 307 LogicalResult processSwitchOp(spirv::SwitchOp switchOp);308 309 //===--------------------------------------------------------------------===//310 // Operations311 //===--------------------------------------------------------------------===//312 313 LogicalResult encodeExtensionInstruction(Operation *op,314 StringRef extensionSetName,315 uint32_t opcode,316 ArrayRef<uint32_t> operands);317 318 uint32_t getValueID(Value val) const { return valueIDMap.lookup(val); }319 320 LogicalResult processAddressOfOp(spirv::AddressOfOp addressOfOp);321 322 LogicalResult processReferenceOfOp(spirv::ReferenceOfOp referenceOfOp);323 324 /// Main dispatch method for serializing an operation.325 LogicalResult processOperation(Operation *op);326 327 /// Serializes an operation `op` as core instruction with `opcode` if328 /// `extInstSet` is empty. Otherwise serializes it as an extended instruction329 /// with `opcode` from `extInstSet`.330 /// This method is a generic one for dispatching any SPIR-V ops that has no331 /// variadic operands and attributes in TableGen definitions.332 LogicalResult processOpWithoutGrammarAttr(Operation *op, StringRef extInstSet,333 uint32_t opcode);334 335 /// Dispatches to the serialization function for an operation in SPIR-V336 /// dialect that is a mirror of an instruction in the SPIR-V spec. This is337 /// auto-generated from ODS. Dispatch is handled for all operations in SPIR-V338 /// dialect that have hasOpcode == 1.339 LogicalResult dispatchToAutogenSerialization(Operation *op);340 341 /// Serializes an operation in the SPIR-V dialect that is a mirror of an342 /// instruction in the SPIR-V spec. This is auto generated if hasOpcode == 1343 /// and autogenSerialization == 1 in ODS.344 template <typename OpTy>345 LogicalResult processOp(OpTy op) {346 return op.emitError("unsupported op serialization");347 }348 349 //===--------------------------------------------------------------------===//350 // Utilities351 //===--------------------------------------------------------------------===//352 353 /// Emits an OpDecorate instruction to decorate the given `target` with the354 /// given `decoration`.355 LogicalResult emitDecoration(uint32_t target, spirv::Decoration decoration,356 ArrayRef<uint32_t> params = {});357 358 /// Emits an OpLine instruction with the given `loc` location information into359 /// the given `binary` vector.360 LogicalResult emitDebugLine(SmallVectorImpl<uint32_t> &binary, Location loc);361 362private:363 /// The SPIR-V module to be serialized.364 spirv::ModuleOp module;365 366 /// An MLIR builder for getting MLIR constructs.367 mlir::Builder mlirBuilder;368 369 /// Serialization options.370 SerializationOptions options;371 372 /// A flag which indicates if the last processed instruction was a merge373 /// instruction.374 /// According to SPIR-V spec: "If a branch merge instruction is used, the last375 /// OpLine in the block must be before its merge instruction".376 bool lastProcessedWasMergeInst = false;377 378 /// The <id> of the OpString instruction, which specifies a file name, for379 /// use by other debug instructions.380 uint32_t fileID = 0;381 382 /// The next available result <id>.383 uint32_t nextID = 1;384 385 // The following are for different SPIR-V instruction sections. They follow386 // the logical layout of a SPIR-V module.387 388 SmallVector<uint32_t, 4> capabilities;389 SmallVector<uint32_t, 0> extensions;390 SmallVector<uint32_t, 0> extendedSets;391 SmallVector<uint32_t, 3> memoryModel;392 SmallVector<uint32_t, 0> entryPoints;393 SmallVector<uint32_t, 4> executionModes;394 SmallVector<uint32_t, 0> debug;395 SmallVector<uint32_t, 0> names;396 SmallVector<uint32_t, 0> decorations;397 SmallVector<uint32_t, 0> typesGlobalValues;398 SmallVector<uint32_t, 0> functions;399 SmallVector<uint32_t, 0> graphs;400 401 /// Recursive struct references are serialized as OpTypePointer instructions402 /// to the recursive struct type. However, the OpTypePointer instruction403 /// cannot be emitted before the recursive struct's OpTypeStruct.404 /// RecursiveStructPointerInfo stores the data needed to emit such405 /// OpTypePointer instructions after forward references to such types.406 struct RecursiveStructPointerInfo {407 uint32_t pointerTypeID;408 spirv::StorageClass storageClass;409 };410 411 // Maps spirv::StructType to its recursive reference member info.412 DenseMap<Type, SmallVector<RecursiveStructPointerInfo, 0>>413 recursiveStructInfos;414 415 /// `functionHeader` contains all the instructions that must be in the first416 /// block in the function or graph, and `functionBody` contains the rest.417 /// After processing FuncOp/GraphARMOp, the encoded instructions of a function418 /// or graph are appended to `functions` or `graphs` respectively. Examples of419 /// instructions in `functionHeader` in order:420 ///421 /// For a FuncOp:422 /// OpFunction ...423 /// OpFunctionParameter ...424 /// OpFunctionParameter ...425 /// OpLabel ...426 /// OpVariable ...427 /// OpVariable ...428 ///429 /// For a GraphARMOp430 /// OpGraphARM ...431 /// OpGraphInputARM ...432 SmallVector<uint32_t, 0> functionHeader;433 SmallVector<uint32_t, 0> functionBody;434 435 /// Map from type used in SPIR-V module to their <id>s.436 DenseMap<Type, uint32_t> typeIDMap;437 438 /// Map from constant values to their <id>s.439 DenseMap<Attribute, uint32_t> constIDMap;440 441 /// Map from a replicated composite constant's value and type to their <id>s.442 DenseMap<std::pair<Attribute, Type>, uint32_t> constCompositeReplicateIDMap;443 444 /// Map from specialization constant names to their <id>s.445 llvm::StringMap<uint32_t> specConstIDMap;446 447 /// Map from graph constant ID value to their <id>s.448 DenseMap<Attribute, uint32_t> graphConstIDMap;449 450 /// Map from GlobalVariableOps name to <id>s.451 llvm::StringMap<uint32_t> globalVarIDMap;452 453 /// Map from FuncOps name to <id>s.454 llvm::StringMap<uint32_t> funcIDMap;455 456 /// Map from blocks to their <id>s.457 DenseMap<Block *, uint32_t> blockIDMap;458 459 /// Map from the Type to the <id> that represents undef value of that type.460 DenseMap<Type, uint32_t> undefValIDMap;461 462 /// Map from results of normal operations to their <id>s.463 DenseMap<Value, uint32_t> valueIDMap;464 465 /// Map from extended instruction set name to <id>s.466 llvm::StringMap<uint32_t> extendedInstSetIDMap;467 468 /// Map from values used in OpPhi instructions to their offset in the469 /// `functions` section.470 ///471 /// When processing a block with arguments, we need to emit OpPhi472 /// instructions to record the predecessor block <id>s and the values they473 /// send to the block in question. But it's not guaranteed all values are474 /// visited and thus assigned result <id>s. So we need this list to capture475 /// the offsets into `functions` where a value is used so that we can fix it476 /// up later after processing all the blocks in a function.477 ///478 /// More concretely, say if we are visiting the following blocks:479 ///480 /// ```mlir481 /// ^phi(%arg0: i32):482 /// ...483 /// ^parent1:484 /// ...485 /// spirv.Branch ^phi(%val0: i32)486 /// ^parent2:487 /// ...488 /// spirv.Branch ^phi(%val1: i32)489 /// ```490 ///491 /// When we are serializing the `^phi` block, we need to emit at the beginning492 /// of the block OpPhi instructions which has the following parameters:493 ///494 /// OpPhi id-for-i32 id-for-%arg0 id-for-%val0 id-for-^parent1495 /// id-for-%val1 id-for-^parent2496 ///497 /// But we don't know the <id> for %val0 and %val1 yet. One way is to visit498 /// all the blocks twice and use the first visit to assign an <id> to each499 /// value. But it's paying the overheads just for OpPhi emission. Instead,500 /// we still visit the blocks once for emission. When we emit the OpPhi501 /// instructions, we use 0 as a placeholder for the <id>s for %val0 and %val1.502 /// At the same time, we record their offsets in the emitted binary (which is503 /// placed inside `functions`) here. And then after emitting all blocks, we504 /// replace the dummy <id> 0 with the real result <id> by overwriting505 /// `functions[offset]`.506 DenseMap<Value, SmallVector<size_t, 1>> deferredPhiValues;507};508} // namespace spirv509} // namespace mlir510 511#endif // MLIR_LIB_TARGET_SPIRV_SERIALIZATION_SERIALIZER_H512