// 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";
  std::string targetFeatures;
  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<std::string>(
        "iree-vulkan-target-features", targetFeatures,
        llvm::cl::desc(
            "Vulkan target features (SPIR-V version, capabilities, "
            "extensions, etc.) to use. If provided, replaces the default "
            "features for the target."));
    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 (!options_.targetFeatures.empty()) {
        target = IREE::GPU::TargetAttr::get(context, target.getArch(),
                                            options_.targetFeatures,
                                            target.getWgp(), target.getChip());
      }
      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 &registry) 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);