| // Copyright 2019 The IREE Authors |
| // |
| // Licensed under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| |
| #include "iree/compiler/Codegen/Common/Passes.h" |
| #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h" |
| #include "iree/compiler/Codegen/Dialect/GPU/TargetUtils/KnownTargets.h" |
| #include "iree/compiler/Codegen/SPIRV/Passes.h" |
| #include "iree/compiler/Codegen/Utils/GPUUtils.h" |
| #include "iree/compiler/Dialect/Encoding/IR/EncodingDialect.h" |
| #include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h" |
| #include "iree/compiler/Dialect/HAL/Utils/ExecutableDebugInfoUtils.h" |
| #include "iree/compiler/PluginAPI/Client.h" |
| #include "iree/compiler/Utils/FlatbufferUtils.h" |
| #include "iree/compiler/Utils/ModuleUtils.h" |
| #include "iree/schemas/vulkan_executable_def_builder.h" |
| #include "llvm/ADT/STLExtras.h" |
| #include "llvm/Support/CommandLine.h" |
| #include "llvm/Support/raw_ostream.h" |
| #include "mlir/Dialect/GPU/IR/GPUDialect.h" |
| #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h" |
| #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h" |
| #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h" |
| #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h" |
| #include "mlir/Dialect/SPIRV/Linking/ModuleCombiner.h" |
| #include "mlir/IR/Builders.h" |
| #include "mlir/IR/BuiltinAttributes.h" |
| #include "mlir/IR/BuiltinOps.h" |
| #include "mlir/Interfaces/FunctionInterfaces.h" |
| #include "mlir/Pass/Pass.h" |
| #include "mlir/Target/SPIRV/Serialization.h" |
| |
| namespace mlir::iree_compiler::IREE::HAL { |
| namespace { |
| constexpr unsigned kBdaDispatchRootDwordCount = 8; |
| constexpr uint32_t kBdaDispatchRootLength = |
| kBdaDispatchRootDwordCount * sizeof(uint32_t); |
| constexpr uint32_t kBdaDispatchRootOffset = 0; |
| constexpr uint32_t kBdaDispatchConstantOffset = kBdaDispatchRootLength; |
| |
| constexpr char kVulkanSpirvFlatbufferFormat[] = "vulkan-spirv-fb"; |
| constexpr char kVulkanSpirvBdaFormat[] = "vulkan-spirv-bda-v1"; |
| |
| enum class VulkanDispatchAbi { |
| Descriptors, |
| Bda, |
| All, |
| }; |
| |
| class PropagateExecutableTargetPass final |
| : public PassWrapper<PropagateExecutableTargetPass, |
| OperationPass<ModuleOp>> { |
| public: |
| MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PropagateExecutableTargetPass) |
| |
| explicit PropagateExecutableTargetPass( |
| IREE::HAL::ExecutableTargetAttr targetAttr) |
| : targetAttr(targetAttr) {} |
| |
| StringRef getArgument() const final { |
| return "iree-vulkan-spirv-propagate-executable-target"; |
| } |
| |
| StringRef getDescription() const final { |
| return "Propagates the selected Vulkan executable target into the inner " |
| "module."; |
| } |
| |
| void runOnOperation() final { |
| ModuleOp moduleOp = getOperation(); |
| if (failed(setOrVerifyTarget(moduleOp))) { |
| return signalPassFailure(); |
| } |
| for (auto funcOp : moduleOp.getOps<FunctionOpInterface>()) { |
| if (failed(setOrVerifyTarget(funcOp))) { |
| return signalPassFailure(); |
| } |
| } |
| } |
| |
| private: |
| LogicalResult setOrVerifyTarget(Operation *op) { |
| Attribute existingAttr = op->getAttr(IREE::HAL::ExecutableTargetAttr::name); |
| if (!existingAttr) { |
| op->setAttr(IREE::HAL::ExecutableTargetAttr::name, targetAttr); |
| return success(); |
| } |
| if (existingAttr == targetAttr) { |
| return success(); |
| } |
| return op->emitError() << "conflicting Vulkan executable target attribute"; |
| } |
| |
| // Target selected by the HAL target backend for this translation pipeline. |
| IREE::HAL::ExecutableTargetAttr targetAttr; |
| }; |
| |
| struct VulkanSPIRVTargetOptions { |
| // Use vp_android_baseline_2022 profile as the default target--it's a good |
| // lowest common denominator to guarantee the generated SPIR-V is widely |
| // accepted for now. Eventually we want to use a list for multi-targeting. |
| std::string target = "vp_android_baseline_2022"; |
| VulkanDispatchAbi dispatchAbi = VulkanDispatchAbi::Descriptors; |
| |
| void bindOptions(OptionsBinder &binder) { |
| static llvm::cl::OptionCategory category("VulkanSPIRV HAL Target"); |
| binder.opt<std::string>( |
| "iree-vulkan-target", target, |
| llvm::cl::desc( |
| "Vulkan target controlling the SPIR-V environment. Given the wide " |
| "support of Vulkan, this option supports a few schemes: 1) LLVM " |
| "CodeGen backend style: e.g., 'gfx*' for AMD GPUs and 'sm_*' for " |
| "NVIDIA GPUs; 2) architecture code name style: e.g., " |
| "'rdna3'/'valhall4'/'ampere'/'adreno' for AMD/ARM/NVIDIA/Qualcomm " |
| "GPUs; 3) product name style: 'rx7900xtx'/'rtx4090' for AMD/NVIDIA " |
| "GPUs. See " |
| "https://iree.dev/guides/deployment-configurations/gpu-vulkan/ for " |
| "more details.")); |
| binder.opt<VulkanDispatchAbi>( |
| "iree-vulkan-dispatch-abi", dispatchAbi, |
| llvm::cl::desc("Selects the Vulkan dispatch ABI emitted for generated " |
| "executables."), |
| llvm::cl::values( |
| clEnumValN(VulkanDispatchAbi::Descriptors, "descriptors", |
| "Emit descriptor-set executables."), |
| clEnumValN(VulkanDispatchAbi::Bda, "bda", |
| "Emit BDA root binding table executables."), |
| clEnumValN(VulkanDispatchAbi::All, "all", |
| "Emit all supported executable ABI variants ordered by " |
| "runtime preference."))); |
| } |
| }; |
| |
| using DescriptorSetLayout = std::pair<unsigned, ArrayRef<PipelineBindingAttr>>; |
| |
| static IREE::GPU::TargetAttr |
| getTargetAttrWithBdaRootAbiFeatures(IREE::GPU::TargetAttr target) { |
| StringRef features = target.getFeatures(); |
| std::string bdaFeatures = features.str(); |
| auto appendFeatureIfMissing = [&](StringRef feature) { |
| if (features.contains(feature)) { |
| return; |
| } |
| if (!bdaFeatures.empty()) { |
| bdaFeatures += ","; |
| } |
| bdaFeatures += feature; |
| }; |
| appendFeatureIfMissing("cap:Int64"); |
| appendFeatureIfMissing("cap:PhysicalStorageBufferAddresses"); |
| appendFeatureIfMissing("ext:SPV_KHR_physical_storage_buffer"); |
| |
| if (features == StringRef(bdaFeatures)) { |
| return target; |
| } |
| return IREE::GPU::TargetAttr::get(target.getContext(), target.getArch(), |
| bdaFeatures, target.getWgp(), |
| target.getChip()); |
| } |
| |
| static iree_hal_vulkan_BdaDispatchLayoutDef_ref_t |
| createBdaDispatchLayoutDef(IREE::HAL::ExecutableExportOp exportOp, |
| FlatbufferBuilder &fbb) { |
| iree_hal_vulkan_BdaDispatchLayoutDef_start(fbb); |
| iree_hal_vulkan_BdaDispatchLayoutDef_abi_version_add(fbb, 1); |
| iree_hal_vulkan_BdaDispatchLayoutDef_root_push_constant_offset_add( |
| fbb, kBdaDispatchRootOffset); |
| iree_hal_vulkan_BdaDispatchLayoutDef_root_push_constant_length_add( |
| fbb, kBdaDispatchRootLength); |
| iree_hal_vulkan_BdaDispatchLayoutDef_constant_push_constant_offset_add( |
| fbb, kBdaDispatchConstantOffset); |
| iree_hal_vulkan_BdaDispatchLayoutDef_constant_count_add( |
| fbb, static_cast<uint32_t>(exportOp.getLayout().getConstants())); |
| iree_hal_vulkan_BdaDispatchLayoutDef_binding_table_entry_type_add( |
| fbb, iree_hal_vulkan_BdaBindingTableEntryType_ADDRESS64); |
| iree_hal_vulkan_BdaDispatchLayoutDef_binding_count_add( |
| fbb, static_cast<uint32_t>(exportOp.getLayout().getBindings().size())); |
| return iree_hal_vulkan_BdaDispatchLayoutDef_end(fbb); |
| } |
| |
| static std::tuple<iree_hal_vulkan_DescriptorSetLayoutDef_vec_ref_t, |
| iree_hal_vulkan_PipelineLayoutDef_vec_ref_t, |
| DenseMap<IREE::HAL::PipelineLayoutAttr, uint32_t>> |
| createPipelineLayoutDefs(ArrayRef<IREE::HAL::ExecutableExportOp> exportOps, |
| bool useBdaRootAbi, FlatbufferBuilder &fbb) { |
| DenseMap<DescriptorSetLayout, size_t> descriptorSetLayoutMap; |
| DenseMap<IREE::HAL::PipelineLayoutAttr, uint32_t> pipelineLayoutMap; |
| SmallVector<iree_hal_vulkan_DescriptorSetLayoutDef_ref_t> |
| descriptorSetLayoutRefs; |
| SmallVector<iree_hal_vulkan_PipelineLayoutDef_ref_t> pipelineLayoutRefs; |
| for (auto exportOp : exportOps) { |
| auto pipelineLayoutAttr = exportOp.getLayout(); |
| if (pipelineLayoutMap.contains(pipelineLayoutAttr)) { |
| continue; // already present |
| } |
| |
| SmallVector<uint32_t> descriptorSetLayoutOrdinals; |
| if (!useBdaRootAbi) { |
| // Currently only one descriptor set on the compiler side. We could |
| // partition it by binding type (direct vs indirect, etc). |
| auto descriptorSetLayout = |
| DescriptorSetLayout(0, pipelineLayoutAttr.getBindings()); |
| auto it = descriptorSetLayoutMap.find(descriptorSetLayout); |
| if (it != descriptorSetLayoutMap.end()) { |
| descriptorSetLayoutOrdinals.push_back(it->second); |
| } else { |
| SmallVector<iree_hal_vulkan_DescriptorSetLayoutBindingDef_ref_t> |
| bindingRefs; |
| for (auto [i, bindingAttr] : |
| llvm::enumerate(pipelineLayoutAttr.getBindings())) { |
| uint32_t ordinal = static_cast<uint32_t>(i); |
| iree_hal_vulkan_VkDescriptorType_enum_t descriptorType = 0; |
| switch (bindingAttr.getType()) { |
| case IREE::HAL::DescriptorType::UniformBuffer: |
| descriptorType = iree_hal_vulkan_VkDescriptorType_UNIFORM_BUFFER; |
| break; |
| case IREE::HAL::DescriptorType::StorageBuffer: |
| descriptorType = iree_hal_vulkan_VkDescriptorType_STORAGE_BUFFER; |
| break; |
| } |
| uint32_t descriptorCount = 1; |
| uint32_t stageFlags = 0x00000020u; // VK_SHADER_STAGE_COMPUTE_BIT |
| bindingRefs.push_back( |
| iree_hal_vulkan_DescriptorSetLayoutBindingDef_create( |
| fbb, ordinal, descriptorType, descriptorCount, stageFlags)); |
| } |
| auto bindingsRef = fbb.createOffsetVecDestructive(bindingRefs); |
| |
| descriptorSetLayoutOrdinals.push_back(descriptorSetLayoutRefs.size()); |
| descriptorSetLayoutMap[descriptorSetLayout] = |
| descriptorSetLayoutRefs.size(); |
| descriptorSetLayoutRefs.push_back( |
| iree_hal_vulkan_DescriptorSetLayoutDef_create(fbb, bindingsRef)); |
| } |
| } |
| auto descriptorSetLayoutOrdinalsRef = |
| fbb.createInt32Vec(descriptorSetLayoutOrdinals); |
| |
| iree_hal_vulkan_PushConstantRange_vec_ref_t pushConstantRangesRef = 0; |
| int64_t pushConstantCount = pipelineLayoutAttr.getConstants(); |
| if (useBdaRootAbi) { |
| pushConstantCount += kBdaDispatchRootDwordCount; |
| } |
| if (pushConstantCount) { |
| SmallVector<iree_hal_vulkan_PushConstantRange> pushConstantRanges; |
| iree_hal_vulkan_PushConstantRange range0; |
| range0.stage_flags = 0x00000020u; // VK_SHADER_STAGE_COMPUTE_BIT |
| range0.offset = 0; |
| range0.size = pushConstantCount * sizeof(uint32_t); |
| pushConstantRanges.push_back(range0); |
| pushConstantRangesRef = iree_hal_vulkan_PushConstantRange_vec_create( |
| fbb, pushConstantRanges.data(), pushConstantRanges.size()); |
| } |
| |
| pipelineLayoutMap[pipelineLayoutAttr] = |
| static_cast<uint32_t>(pipelineLayoutRefs.size()); |
| iree_hal_vulkan_PipelineLayoutDef_start(fbb); |
| iree_hal_vulkan_PipelineLayoutDef_descriptor_set_layout_ordinals_add( |
| fbb, descriptorSetLayoutOrdinalsRef); |
| if (pushConstantRangesRef) { |
| iree_hal_vulkan_PipelineLayoutDef_push_constant_ranges_add( |
| fbb, pushConstantRangesRef); |
| } |
| pipelineLayoutRefs.push_back(iree_hal_vulkan_PipelineLayoutDef_end(fbb)); |
| } |
| |
| auto descriptorSetLayoutsRef = |
| fbb.createOffsetVecDestructive(descriptorSetLayoutRefs); |
| auto pipelineLayoutsRef = fbb.createOffsetVecDestructive(pipelineLayoutRefs); |
| return std::make_tuple(descriptorSetLayoutsRef, pipelineLayoutsRef, |
| pipelineLayoutMap); |
| } |
| |
| // TODO: VulkanOptions for choosing the Vulkan version and extensions/features. |
| class VulkanTargetDevice final : public TargetDevice { |
| public: |
| VulkanTargetDevice(const VulkanSPIRVTargetOptions & /*options*/) {} |
| |
| IREE::HAL::DeviceTargetAttr |
| getDefaultDeviceTarget(MLIRContext *context, |
| const TargetRegistry &targetRegistry) const final { |
| Builder b(context); |
| auto deviceConfigAttr = b.getDictionaryAttr({}); |
| auto executableConfigAttr = b.getDictionaryAttr({}); |
| |
| SmallVector<IREE::HAL::ExecutableTargetAttr> executableTargetAttrs; |
| targetRegistry.getTargetBackend("vulkan-spirv") |
| ->getDefaultExecutableTargets(context, "vulkan", executableConfigAttr, |
| executableTargetAttrs); |
| |
| return IREE::HAL::DeviceTargetAttr::get(context, b.getStringAttr("vulkan"), |
| deviceConfigAttr, |
| executableTargetAttrs); |
| } |
| }; |
| |
| class VulkanSPIRVTargetBackend final : public TargetBackend { |
| public: |
| VulkanSPIRVTargetBackend(const VulkanSPIRVTargetOptions &options) |
| : options_(options) {} |
| |
| std::string getLegacyDefaultDeviceID() const final { return "vulkan"; } |
| |
| void getDefaultExecutableTargets( |
| MLIRContext *context, StringRef deviceID, DictionaryAttr deviceConfigAttr, |
| SmallVectorImpl<IREE::HAL::ExecutableTargetAttr> &executableTargetAttrs) |
| const final { |
| switch (options_.dispatchAbi) { |
| case VulkanDispatchAbi::Descriptors: |
| executableTargetAttrs.push_back(getExecutableTarget(context, false)); |
| break; |
| case VulkanDispatchAbi::Bda: |
| executableTargetAttrs.push_back(getExecutableTarget(context, true)); |
| break; |
| case VulkanDispatchAbi::All: |
| executableTargetAttrs.push_back(getExecutableTarget(context, true)); |
| executableTargetAttrs.push_back(getExecutableTarget(context, false)); |
| break; |
| } |
| } |
| |
| IREE::HAL::ExecutableTargetAttr |
| getExecutableTarget(MLIRContext *context, bool useBdaRootAbi) const { |
| Builder b(context); |
| SmallVector<NamedAttribute, 1> configItems; |
| if (auto target = GPU::getVulkanTargetDetails(options_.target, context)) { |
| if (useBdaRootAbi) { |
| target = getTargetAttrWithBdaRootAbiFeatures(target); |
| } |
| addConfigGPUTarget(context, target, configItems); |
| } else { |
| emitError(b.getUnknownLoc(), "Unknown Vulkan target '") |
| << options_.target << "'"; |
| return nullptr; |
| } |
| |
| return IREE::HAL::ExecutableTargetAttr::get( |
| context, b.getStringAttr("vulkan-spirv"), |
| useBdaRootAbi ? b.getStringAttr(kVulkanSpirvBdaFormat) |
| : b.getStringAttr(kVulkanSpirvFlatbufferFormat), |
| b.getDictionaryAttr(configItems)); |
| } |
| |
| void getDependentDialects(DialectRegistry ®istry) const final { |
| registry.insert<IREE::Codegen::IREECodegenDialect, |
| IREE::Encoding::IREEEncodingDialect, spirv::SPIRVDialect, |
| gpu::GPUDialect, IREE::GPU::IREEGPUDialect>(); |
| } |
| |
| void |
| buildConfigurationPassPipeline(IREE::HAL::ExecutableTargetAttr targetAttr, |
| OpPassManager &passManager) final { |
| buildCodegenConfigurationPreProcessingPassPipeline(passManager); |
| buildSPIRVCodegenConfigurationPassPipeline(passManager.nest<ModuleOp>()); |
| } |
| |
| void buildTranslationPassPipeline(IREE::HAL::ExecutableTargetAttr targetAttr, |
| OpPassManager &passManager) final { |
| OpPassManager &modulePassManager = passManager.nest<ModuleOp>(); |
| modulePassManager.addPass( |
| std::make_unique<PropagateExecutableTargetPass>(targetAttr)); |
| buildSPIRVCodegenPassPipeline(modulePassManager); |
| buildCodegenTranslationPostProcessingPassPipeline(passManager); |
| } |
| |
| void buildLinkingPassPipeline(OpPassManager &passManager) final { |
| buildSPIRVLinkingPassPipeline(passManager); |
| } |
| |
| LogicalResult serializeExecutable(const SerializationOptions &options, |
| IREE::HAL::ExecutableVariantOp variantOp, |
| OpBuilder &executableBuilder) final { |
| // Today we special-case external variants but in the future we could allow |
| // for a linking approach allowing both code generation and external .spv |
| // files to be combined together. |
| if (variantOp.isExternal()) { |
| return serializeExternalExecutable(options, variantOp, executableBuilder); |
| } |
| |
| ModuleOp innerModuleOp = variantOp.getInnerModule(); |
| auto spirvModuleOps = innerModuleOp.getOps<spirv::ModuleOp>(); |
| if (spirvModuleOps.empty()) { |
| return variantOp.emitError() << "should contain some spirv.module ops"; |
| } |
| |
| // Create a list of executable exports (by ordinal) to the SPIR-V module and |
| // entry point defining them. |
| auto unsortedExportOps = |
| llvm::to_vector(variantOp.getOps<IREE::HAL::ExecutableExportOp>()); |
| DenseMap<StringRef, std::tuple<IREE::HAL::ExecutableExportOp, uint64_t>> |
| exportOrdinalMap; |
| for (auto exportOp : variantOp.getOps<IREE::HAL::ExecutableExportOp>()) { |
| uint64_t ordinal = 0; |
| if (std::optional<APInt> optionalOrdinal = exportOp.getOrdinal()) { |
| ordinal = optionalOrdinal->getZExtValue(); |
| } else { |
| // For executables with only one entry point linking doesn't kick in at |
| // all. So the ordinal can be missing for this case. |
| if (!llvm::hasSingleElement(unsortedExportOps)) { |
| return exportOp.emitError() << "should have ordinal attribute"; |
| } |
| } |
| exportOrdinalMap[exportOp.getSymName()] = |
| std::make_tuple(exportOp, ordinal); |
| } |
| SmallVector<IREE::HAL::ExecutableExportOp> sortedExportOps; |
| sortedExportOps.resize(unsortedExportOps.size()); |
| SmallVector<std::tuple<IREE::HAL::ExecutableExportOp, spirv::ModuleOp, |
| spirv::EntryPointOp>> |
| exportOps; |
| exportOps.resize(unsortedExportOps.size()); |
| for (spirv::ModuleOp spirvModuleOp : spirvModuleOps) { |
| for (spirv::EntryPointOp spirvEntryPointOp : |
| spirvModuleOp.getOps<spirv::EntryPointOp>()) { |
| auto it = exportOrdinalMap.find(spirvEntryPointOp.getFn()); |
| if (it == exportOrdinalMap.end()) { |
| continue; |
| } |
| auto [exportOp, ordinal] = it->second; |
| sortedExportOps[ordinal] = exportOp; |
| exportOps[ordinal] = |
| std::make_tuple(exportOp, spirvModuleOp, spirvEntryPointOp); |
| } |
| } |
| |
| FlatbufferBuilder builder; |
| iree_hal_vulkan_ExecutableDef_start_as_root(builder); |
| |
| // Attach embedded source file contents. |
| auto sourceFilesRef = createSourceFilesVec( |
| options.debugLevel, variantOp.getSourcesAttr(), builder); |
| |
| // Generate optional per-export debug information. |
| // May be empty if no debug information was requested. |
| auto exportDebugInfos = |
| createExportDefs(options.debugLevel, sortedExportOps, builder); |
| |
| // Create a serialized SPIR-V module for each entry point. When a |
| // spirv.module contains multiple entry points each gets its own copy of |
| // the binary — deduplicating shared modules is left as a future |
| // optimization (N:M mapping via specialization constants). |
| DenseMap<spirv::EntryPointOp, uint32_t> entryPointToModuleMap; |
| SmallVector<iree_hal_vulkan_ShaderModuleDef_ref_t> shaderModuleRefs; |
| for (auto [exportOp, spirvModuleOp, spirvEntryPointOp] : exportOps) { |
| if (!options.dumpIntermediatesPath.empty()) { |
| std::string assembly; |
| llvm::raw_string_ostream os(assembly); |
| spirvModuleOp.print(os, OpPrintingFlags().useLocalScope()); |
| dumpDataToPath(options.dumpIntermediatesPath, options.dumpBaseName, |
| spirvEntryPointOp.getFn(), ".spirv.mlir", assembly); |
| } |
| |
| // Serialize the spirv::ModuleOp into the binary blob. |
| SmallVector<uint32_t, 0> spirvBinary; |
| if (failed(spirv::serialize(spirvModuleOp, spirvBinary)) || |
| spirvBinary.empty()) { |
| return spirvModuleOp.emitError() << "failed to serialize"; |
| } |
| if (!options.dumpBinariesPath.empty()) { |
| dumpDataToPath<uint32_t>(options.dumpBinariesPath, options.dumpBaseName, |
| spirvEntryPointOp.getFn(), ".spv", |
| spirvBinary); |
| } |
| auto spirvCodeRef = flatbuffers_uint32_vec_create( |
| builder, spirvBinary.data(), spirvBinary.size()); |
| entryPointToModuleMap[spirvEntryPointOp] = |
| static_cast<uint32_t>(shaderModuleRefs.size()); |
| shaderModuleRefs.push_back( |
| iree_hal_vulkan_ShaderModuleDef_create(builder, spirvCodeRef)); |
| } |
| auto shaderModulesRef = |
| builder.createOffsetVecDestructive(shaderModuleRefs); |
| |
| const bool useBdaRootAbi = |
| variantOp.getTarget().getFormat() == kVulkanSpirvBdaFormat; |
| |
| // Create unique descriptor and pipeline layouts for each entry point. |
| auto [descriptorSetLayoutsRef, pipelineLayoutsRef, pipelineLayoutMap] = |
| createPipelineLayoutDefs(sortedExportOps, useBdaRootAbi, builder); |
| |
| // Create pipelines representing entry points. |
| // Note that the element at index #i is for entry point with ordinal #i. |
| SmallVector<iree_hal_vulkan_PipelineDef_ref_t> pipelineRefs; |
| for (auto [exportOp, spirvModuleOp, spirvEntryPointOp] : exportOps) { |
| int64_t ordinal = exportOp.getOrdinalAttr().getInt(); |
| |
| uint32_t shaderModuleOrdinal = |
| entryPointToModuleMap.at(spirvEntryPointOp); |
| uint32_t pipelineLayoutOrdinal = |
| pipelineLayoutMap.at(exportOp.getLayout()); |
| |
| // Subgroup size requests are optional. |
| auto spirvFuncOp = |
| spirvModuleOp.lookupSymbol<spirv::FuncOp>(spirvEntryPointOp.getFn()); |
| auto abiAttr = spirvFuncOp->getAttrOfType<spirv::EntryPointABIAttr>( |
| spirv::getEntryPointABIAttrName()); |
| uint32_t subgroupSize = |
| abiAttr ? abiAttr.getSubgroupSize().value_or(0) : 0; |
| |
| iree_hal_vulkan_BdaDispatchLayoutDef_ref_t bdaDispatchLayoutRef = |
| useBdaRootAbi ? createBdaDispatchLayoutDef(exportOp, builder) : 0; |
| |
| auto entryPointRef = builder.createString(spirvEntryPointOp.getFn()); |
| iree_hal_vulkan_PipelineDef_start(builder); |
| iree_hal_vulkan_PipelineDef_shader_module_ordinal_add( |
| builder, shaderModuleOrdinal); |
| iree_hal_vulkan_PipelineDef_entry_point_add(builder, entryPointRef); |
| iree_hal_vulkan_PipelineDef_pipeline_layout_ordinal_add( |
| builder, pipelineLayoutOrdinal); |
| iree_hal_vulkan_PipelineDef_subgroup_size_add(builder, subgroupSize); |
| iree_hal_vulkan_PipelineDef_debug_info_add(builder, |
| exportDebugInfos[ordinal]); |
| if (useBdaRootAbi) { |
| iree_hal_vulkan_PipelineDef_dispatch_abi_add( |
| builder, iree_hal_vulkan_DispatchAbi_BDA_V1); |
| iree_hal_vulkan_PipelineDef_bda_dispatch_layout_add( |
| builder, bdaDispatchLayoutRef); |
| } |
| pipelineRefs.push_back(iree_hal_vulkan_PipelineDef_end(builder)); |
| } |
| auto pipelinesRef = builder.createOffsetVecDestructive(pipelineRefs); |
| |
| // Add top-level executable fields following their order of definition. |
| iree_hal_vulkan_ExecutableDef_pipelines_add(builder, pipelinesRef); |
| iree_hal_vulkan_ExecutableDef_descriptor_set_layouts_add( |
| builder, descriptorSetLayoutsRef); |
| iree_hal_vulkan_ExecutableDef_pipeline_layouts_add(builder, |
| pipelineLayoutsRef); |
| iree_hal_vulkan_ExecutableDef_shader_modules_add(builder, shaderModulesRef); |
| iree_hal_vulkan_ExecutableDef_source_files_add(builder, sourceFilesRef); |
| |
| iree_hal_vulkan_ExecutableDef_end_as_root(builder); |
| |
| // Add the binary data to the target executable. |
| auto binaryOp = IREE::HAL::ExecutableBinaryOp::create( |
| executableBuilder, variantOp.getLoc(), variantOp.getSymName(), |
| variantOp.getTarget().getFormat(), |
| builder.getHeaderPrefixedBufferAttr( |
| executableBuilder.getContext(), |
| /*magic=*/iree_hal_vulkan_ExecutableDef_file_identifier, |
| /*version=*/0)); |
| binaryOp.setMimeTypeAttr( |
| executableBuilder.getStringAttr("application/x-flatbuffers")); |
| |
| return success(); |
| } |
| |
| LogicalResult |
| serializeExternalExecutable(const SerializationOptions &options, |
| IREE::HAL::ExecutableVariantOp variantOp, |
| OpBuilder &executableBuilder) { |
| if (!variantOp.getObjects().has_value()) { |
| return variantOp.emitOpError() |
| << "no objects defined for external variant"; |
| } else if (variantOp.getObjects()->getValue().size() != 1) { |
| // For now we assume there will be exactly one object file. |
| // TODO(#7824): support multiple .spv files in a single flatbuffer archive |
| // so that we can combine executables. |
| return variantOp.emitOpError() << "only one object reference is " |
| "supported for external variants"; |
| } |
| |
| // Load .spv object file. |
| auto objectAttr = cast<IREE::HAL::ExecutableObjectAttr>( |
| variantOp.getObjects()->getValue().front()); |
| std::string spirvBinary; |
| if (auto data = objectAttr.loadData()) { |
| spirvBinary = data.value(); |
| } else { |
| return variantOp.emitOpError() |
| << "object file could not be loaded: " << objectAttr; |
| } |
| if (spirvBinary.size() % 4 != 0) { |
| return variantOp.emitOpError() |
| << "object file is not 4-byte aligned as expected for SPIR-V"; |
| } |
| |
| FlatbufferBuilder builder; |
| iree_hal_vulkan_ExecutableDef_start_as_root(builder); |
| |
| // Wrap and embed shader module binary. |
| auto spirvCodeRef = flatbuffers_uint32_vec_create( |
| builder, reinterpret_cast<const uint32_t *>(spirvBinary.data()), |
| spirvBinary.size() / sizeof(uint32_t)); |
| SmallVector<iree_hal_vulkan_ShaderModuleDef_ref_t> shaderModuleRefs; |
| shaderModuleRefs.push_back( |
| iree_hal_vulkan_ShaderModuleDef_create(builder, spirvCodeRef)); |
| auto shaderModulesRef = |
| builder.createOffsetVecDestructive(shaderModuleRefs); |
| |
| const bool useBdaRootAbi = |
| variantOp.getTarget().getFormat() == kVulkanSpirvBdaFormat; |
| |
| // Generate descriptor set and pipeline layouts from export ops. |
| auto exportOps = llvm::to_vector(variantOp.getExportOps()); |
| auto [descriptorSetLayoutsRef, pipelineLayoutsRef, pipelineLayoutMap] = |
| createPipelineLayoutDefs(exportOps, useBdaRootAbi, builder); |
| |
| // Create a pipeline for each export. |
| SmallVector<iree_hal_vulkan_PipelineDef_ref_t> pipelineRefs; |
| for (auto exportOp : exportOps) { |
| uint32_t shaderModuleOrdinal = 0; // only one today |
| uint32_t pipelineLayoutOrdinal = |
| pipelineLayoutMap.at(exportOp.getLayout()); |
| |
| // Subgroup size requests are optional. |
| // TODO: support annotation on an attribute to allow users to specify. |
| uint32_t subgroupSize = 0; |
| |
| iree_hal_vulkan_BdaDispatchLayoutDef_ref_t bdaDispatchLayoutRef = |
| useBdaRootAbi ? createBdaDispatchLayoutDef(exportOp, builder) : 0; |
| |
| auto entryPointRef = builder.createString(exportOp.getName()); |
| iree_hal_vulkan_PipelineDef_start(builder); |
| iree_hal_vulkan_PipelineDef_shader_module_ordinal_add( |
| builder, shaderModuleOrdinal); |
| iree_hal_vulkan_PipelineDef_entry_point_add(builder, entryPointRef); |
| iree_hal_vulkan_PipelineDef_pipeline_layout_ordinal_add( |
| builder, pipelineLayoutOrdinal); |
| iree_hal_vulkan_PipelineDef_subgroup_size_add(builder, subgroupSize); |
| if (useBdaRootAbi) { |
| iree_hal_vulkan_PipelineDef_dispatch_abi_add( |
| builder, iree_hal_vulkan_DispatchAbi_BDA_V1); |
| iree_hal_vulkan_PipelineDef_bda_dispatch_layout_add( |
| builder, bdaDispatchLayoutRef); |
| } |
| pipelineRefs.push_back(iree_hal_vulkan_PipelineDef_end(builder)); |
| } |
| auto pipelinesRef = builder.createOffsetVecDestructive(pipelineRefs); |
| |
| // Add top-level executable fields following their order of definition. |
| iree_hal_vulkan_ExecutableDef_pipelines_add(builder, pipelinesRef); |
| iree_hal_vulkan_ExecutableDef_descriptor_set_layouts_add( |
| builder, descriptorSetLayoutsRef); |
| iree_hal_vulkan_ExecutableDef_pipeline_layouts_add(builder, |
| pipelineLayoutsRef); |
| iree_hal_vulkan_ExecutableDef_shader_modules_add(builder, shaderModulesRef); |
| |
| iree_hal_vulkan_ExecutableDef_end_as_root(builder); |
| |
| // Add the binary data to the target executable. |
| auto binaryOp = IREE::HAL::ExecutableBinaryOp::create( |
| executableBuilder, variantOp.getLoc(), variantOp.getSymName(), |
| variantOp.getTarget().getFormat(), |
| builder.getHeaderPrefixedBufferAttr( |
| executableBuilder.getContext(), |
| /*magic=*/iree_hal_vulkan_ExecutableDef_file_identifier, |
| /*version=*/0)); |
| binaryOp.setMimeTypeAttr( |
| executableBuilder.getStringAttr("application/x-flatbuffers")); |
| |
| return success(); |
| } |
| |
| private: |
| const VulkanSPIRVTargetOptions &options_; |
| }; |
| |
| struct VulkanSPIRVSession final |
| : PluginSession<VulkanSPIRVSession, VulkanSPIRVTargetOptions, |
| PluginActivationPolicy::DefaultActivated> { |
| void populateHALTargetDevices(IREE::HAL::TargetDeviceList &targets) final { |
| // #hal.device.target<"vulkan", ... |
| targets.add("vulkan", [&]() { |
| return std::make_shared<VulkanTargetDevice>(options); |
| }); |
| } |
| void populateHALTargetBackends(IREE::HAL::TargetBackendList &targets) final { |
| // #hal.executable.target<"vulkan-spirv", ... |
| targets.add("vulkan-spirv", [&]() { |
| return std::make_shared<VulkanSPIRVTargetBackend>(options); |
| }); |
| } |
| }; |
| |
| } // namespace |
| } // namespace mlir::iree_compiler::IREE::HAL |
| |
| extern "C" bool iree_register_compiler_plugin_hal_target_vulkan_spirv( |
| mlir::iree_compiler::PluginRegistrar *registrar) { |
| registrar->registerPlugin<mlir::iree_compiler::IREE::HAL::VulkanSPIRVSession>( |
| "hal_target_vulkan_spirv"); |
| return true; |
| } |
| |
| IREE_DEFINE_COMPILER_OPTION_FLAGS( |
| mlir::iree_compiler::IREE::HAL::VulkanSPIRVTargetOptions); |
