// Copyright 2020 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

//===- CovertToSPIRV.cpp - Performs the final SPIR-V conversion -----------===//
//
// This file implements a pass to perform the final conversion to SPIR-V.
// This pass converts remaining interface ops into SPIR-V global variables,
// GPU processor ID ops into SPIR-V global variables, loop/standard ops into
// corresponding SPIR-V ops.
//
//===----------------------------------------------------------------------===//

#include <cassert>
#include <cstdint>
#include <tuple>

#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "iree/compiler/Codegen/SPIRV/Passes.h"
#include "iree/compiler/Codegen/SPIRV/Utils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Conversion/ArithToSPIRV/ArithToSPIRV.h"
#include "mlir/Conversion/ComplexToSPIRV/ComplexToSPIRV.h"
#include "mlir/Conversion/ControlFlowToSPIRV/ControlFlowToSPIRVPass.h"
#include "mlir/Conversion/FuncToSPIRV/FuncToSPIRV.h"
#include "mlir/Conversion/GPUToSPIRV/GPUToSPIRV.h"
#include "mlir/Conversion/MathToSPIRV/MathToSPIRV.h"
#include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h"
#include "mlir/Conversion/SCFToSPIRV/SCFToSPIRV.h"
#include "mlir/Conversion/TensorToSPIRV/TensorToSPIRV.h"
#include "mlir/Conversion/UBToSPIRV/UBToSPIRV.h"
#include "mlir/Conversion/VectorToSPIRV/VectorToSPIRV.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Math/Transforms/Passes.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/Visitors.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_CONVERTTOSPIRVPASS
#include "iree/compiler/Codegen/SPIRV/Passes.h.inc"

namespace {

//===----------------------------------------------------------------------===//
// Resource utilities
//===----------------------------------------------------------------------===//

/// Resource info describing a hal.interface.binding.subspan op.
struct SubspanResourceInfo {
  uint32_t set = 0;
  uint32_t binding = 0;
  Type type = nullptr;
  bool aliased = false;
  spirv::GlobalVariableOp var = nullptr;
};

/// Map from hal.interface.binding.subspan ops to their corresponding
/// spirv.GlobalVariable ops.
using InterfaceResourceMap = llvm::DenseMap<Operation *, SubspanResourceInfo>;

constexpr unsigned kBdaDispatchRootDwordCount = 8;
constexpr uint32_t kIndirectBindingsSetIndex = 3;

/// Creates a resource variable using the `globalVariableType` at the beginning
/// of `moduleOp`'s block via `symbolTable` and binds it to `set` and `binding`.
static spirv::GlobalVariableOp
createResourceVariable(Location loc, const SubspanResourceInfo &resource,
                       spirv::PointerType globalVariableType, ModuleOp moduleOp,
                       SymbolTable *symbolTable, bool isIndirect) {
  OpBuilder builder(moduleOp.getContext());
  spirv::GlobalVariableOp variable;
  if (!isIndirect) {
    std::string name =
        llvm::formatv("__resource_var_{}_{}_", resource.set, resource.binding);
    variable = spirv::GlobalVariableOp::create(
        builder, loc, globalVariableType, name, resource.set, resource.binding);
    if (resource.aliased) {
      variable->setAttr("aliased", builder.getUnitAttr());
    }
  } else {
    std::string name =
        llvm::formatv("__resource_var_indirect_{}_", resource.set);
    variable = spirv::GlobalVariableOp::create(builder, loc, globalVariableType,
                                               name, kIndirectBindingsSetIndex,
                                               resource.set);
  }
  assert(variable);

  symbolTable->insert(variable, moduleOp.getBody()->begin());
  return variable;
}

/// Returns the (set, binding) pair for the given interface op.
static std::pair<uint32_t, uint32_t>
getInterfaceSetAndBinding(IREE::HAL::InterfaceBindingSubspanOp op) {
  return {0, op.getBinding().getSExtValue()};
}

/// Scans all hal.interface.binding.subspan ops in `module`, creates their
/// corresponding spirv.GlobalVariables when needed, and returns the map.
/// The created variables need to have their types fixed later.
/// Assumes direct bindings and creates a global variable for each (set,
/// binding, type) tuple. These globals may alias.
static InterfaceResourceMap createResourceVariables(mlir::ModuleOp module) {
  SymbolTable symbolTable(module);
  InterfaceResourceMap interfaceToResourceInfo;

  // We insert each new global variable at the beginning of the module,
  // therefore, to preserve the original order, we process all functions and all
  // subspan ops in the reverse order.
  auto functions = llvm::to_vector(module.getOps<mlir::FunctionOpInterface>());
  for (auto func : llvm::reverse(functions)) {
    // Collect all interface ops and their (set, binding) pairs in this
    // function. Use SmallVector here for a deterministic order.
    SmallVector<IREE::HAL::InterfaceBindingSubspanOp> subspanOps;
    SmallVector<std::pair<uint32_t, uint32_t>> setBindings;

    // Use a map to see if we have different types for one (set, binding) pair,
    // which will require creating multiple SPIR-V global variables.
    llvm::DenseMap<std::pair<uint32_t, uint32_t>, llvm::DenseSet<Type>>
        setBindingTypes;

    func.walk([&](IREE::HAL::InterfaceBindingSubspanOp subspanOp) {
      subspanOps.emplace_back(subspanOp);
      setBindings.emplace_back(getInterfaceSetAndBinding(subspanOp));
      setBindingTypes[setBindings.back()].insert(subspanOp.getType());
    });

    // Keep track of created SPIR-V global variables. This allows us to
    // deduplicate when possible to reduce generated SPIR-V blob size.
    llvm::DenseMap<std::tuple<uint32_t, uint32_t, Type>,
                   spirv::GlobalVariableOp>
        resourceVars;

    // We are using a none type for creating the global variable. It's fine.
    // The correctness boundary is the pass. We will fix it up during
    // conversion so it won't leak.
    auto placeholderType = spirv::PointerType::get(
        NoneType::get(module.getContext()), spirv::StorageClass::StorageBuffer);

    for (auto [subspanOp, setBinding] :
         llvm::reverse(llvm::zip_equal(subspanOps, setBindings))) {
      const auto [set, binding] = setBinding;
      Type type = subspanOp.getType();

      // In the *direct* bindings mode, if we have multiple SPIR-V global
      // variables bound to the same (set, binding) pair and they are used in
      // the same function, those variables need to have alias decoration.
      bool aliases = setBindingTypes[{set, binding}].size() > 1;
      std::tuple<uint32_t, uint32_t, Type> key = {set, binding, type};

      // Use placeholder value for the type to create the variable. We will
      // overwrite it with the actual type later.
      spirv::GlobalVariableOp existingVar = resourceVars.lookup(key);
      SubspanResourceInfo resource = {set, binding, type, aliases, existingVar};
      if (!resource.var) {
        resource.var = createResourceVariable(
            subspanOp.getLoc(), resource, placeholderType, module, &symbolTable,
            /*isIndirect=*/false);
        resourceVars[key] = resource.var;
      }
      assert(resource.var);
      interfaceToResourceInfo[subspanOp] = resource;
    }
  }

  return interfaceToResourceInfo;
}

static spirv::PointerType
getBlockDecoratedStructPointerType(spirv::PointerType pointerType) {
  auto structType = dyn_cast<spirv::StructType>(pointerType.getPointeeType());
  if (!structType || structType.hasDecoration(spirv::Decoration::Block)) {
    return pointerType;
  }

  SmallVector<spirv::StructType::OffsetInfo> offsets;
  if (structType.hasOffset()) {
    offsets.reserve(structType.getNumElements());
    for (unsigned i = 0; i < structType.getNumElements(); ++i) {
      offsets.push_back(structType.getMemberOffset(i));
    }
  }
  SmallVector<spirv::StructType::MemberDecorationInfo> memberDecorations;
  structType.getMemberDecorations(memberDecorations);
  SmallVector<spirv::StructType::StructDecorationInfo> structDecorations;
  structType.getStructDecorations(structDecorations);
  structDecorations.emplace_back(spirv::Decoration::Block,
                                 UnitAttr::get(pointerType.getContext()));
  SmallVector<Type> elementTypes(structType.getElementTypes());
  auto blockStructType = spirv::StructType::get(
      elementTypes, offsets, memberDecorations, structDecorations);
  return spirv::PointerType::get(blockStructType,
                                 pointerType.getStorageClass());
}

// Push-constant blocks need an explicit Block decoration before serialization.
// Otherwise a push constant and a descriptor buffer with the same struct body
// can be uniqued to the same SPIR-V type id and receive duplicate decorations.
static spirv::PointerType
getBlockDecoratedPushConstantStorageType(unsigned elementCount,
                                         OpBuilder &builder, Type integerType) {
  auto arrayType =
      spirv::ArrayType::get(integerType, elementCount, /*stride=*/4);
  auto structType = spirv::StructType::get({arrayType}, /*offsetInfo=*/0);
  auto pointerType =
      spirv::PointerType::get(structType, spirv::StorageClass::PushConstant);
  return getBlockDecoratedStructPointerType(pointerType);
}

static spirv::GlobalVariableOp getPushConstantVariable(Block &body,
                                                       unsigned elementCount) {
  for (auto varOp : body.getOps<spirv::GlobalVariableOp>()) {
    auto pointerType = dyn_cast<spirv::PointerType>(varOp.getType());
    if (!pointerType ||
        pointerType.getStorageClass() != spirv::StorageClass::PushConstant) {
      continue;
    }
    auto structType = dyn_cast<spirv::StructType>(pointerType.getPointeeType());
    if (!structType || structType.getNumElements() != 1) {
      continue;
    }
    auto arrayType = dyn_cast<spirv::ArrayType>(structType.getElementType(0));
    if (arrayType && arrayType.getNumElements() == elementCount) {
      return varOp;
    }
  }
  return nullptr;
}

static spirv::GlobalVariableOp
getOrInsertPushConstantVariable(Location loc, Block &block,
                                unsigned elementCount, OpBuilder &builder,
                                Type integerType) {
  auto pointerType = getBlockDecoratedPushConstantStorageType(
      elementCount, builder, integerType);
  if (auto varOp = getPushConstantVariable(block, elementCount)) {
    varOp.setTypeAttr(TypeAttr::get(pointerType));
    return varOp;
  }

  OpBuilder moduleBuilder =
      OpBuilder::atBlockBegin(&block, builder.getListener());
  return spirv::GlobalVariableOp::create(moduleBuilder, loc, pointerType,
                                         "__push_constant_var__",
                                         /*initializer=*/nullptr);
}

static Value getPushConstantValue(Operation *op, unsigned elementCount,
                                  unsigned offset, Type integerType,
                                  OpBuilder &builder) {
  Operation *parent = SymbolTable::getNearestSymbolTable(op->getParentOp());
  if (!parent) {
    op->emitError("expected operation to be within a module-like op");
    return nullptr;
  }
  spirv::GlobalVariableOp varOp = getOrInsertPushConstantVariable(
      op->getLoc(), parent->getRegion(0).front(), elementCount, builder,
      integerType);
  Value zeroOp = spirv::ConstantOp::getZero(integerType, op->getLoc(), builder);
  Value offsetOp = spirv::ConstantOp::create(builder, op->getLoc(), integerType,
                                             builder.getI32IntegerAttr(offset));
  auto addrOp = spirv::AddressOfOp::create(builder, op->getLoc(), varOp);
  auto accessChainOp = spirv::AccessChainOp::create(
      builder, op->getLoc(), addrOp, llvm::ArrayRef({zeroOp, offsetOp}));
  return spirv::LoadOp::create(builder, op->getLoc(), accessChainOp);
}

static spirv::PointerType
getIndirectBindingTableType(MLIRContext *ctx, uint32_t maxBinding,
                            spirv::StorageClass tableStorageClass) {
  auto placeholderResourceType = IntegerType::get(ctx, 32);
  auto memberPtr = spirv::PointerType::get(
      placeholderResourceType, spirv::StorageClass::PhysicalStorageBuffer);
  SmallVector<Type> members(maxBinding + 1, memberPtr);
  // We are creating a top-level Block decorated interface type here, it needs
  // explicit layout per SPIR-V requirements.
  SmallVector<uint32_t> offsets(maxBinding + 1);
  for (size_t i = 0; i < offsets.size(); ++i) {
    offsets[i] = static_cast<uint32_t>(8 * i); // Each pointer takes 8 bytes.
  }
  auto structType = spirv::StructType::get(members, offsets);
  return spirv::PointerType::get(structType, tableStorageClass);
}

/// Scans all hal.interface.binding.subspan ops in `module` and returns the map.
/// Descriptor-table indirect bindings create one spirv.GlobalVariable for each
/// set, while BDA root bindings leave the table address in push constants.
static InterfaceResourceMap
createIndirectResourceVariables(mlir::ModuleOp module,
                                bool useBdaRootBindings) {
  SymbolTable symbolTable(module);
  InterfaceResourceMap interfaceToResourceInfo;

  // We insert each new global variable at the beginning of the module,
  // therefore, to preserve the original order, we process all functions and all
  // subspan ops in the reverse order.
  auto functions = llvm::to_vector(module.getOps<func::FuncOp>());
  for (func::FuncOp func : llvm::reverse(functions)) {
    // Collect all interface ops and their (set, binding) pairs in this
    // function. Use SmallVector here for a deterministic order.
    SmallVector<IREE::HAL::InterfaceBindingSubspanOp> subspanOps;

    // Keep track of the maximum binding index for each set. We need this to
    // know how many members pointers to add to each global variable struct.
    DenseMap<uint32_t, uint32_t> setToMaxBindingIdx;

    func.walk([&](IREE::HAL::InterfaceBindingSubspanOp subspanOp) {
      subspanOps.push_back(subspanOp);
      auto [set, binding] = getInterfaceSetAndBinding(subspanOp);
      auto [it, inserted] = setToMaxBindingIdx.try_emplace(set, binding);
      if (!inserted) {
        it->second = std::max(it->second, binding);
      }
    });

    DenseMap<uint32_t, spirv::GlobalVariableOp> setToVar;
    for (IREE::HAL::InterfaceBindingSubspanOp subspanOp :
         llvm::reverse(subspanOps)) {
      auto [set, binding] = getInterfaceSetAndBinding(subspanOp);
      spirv::GlobalVariableOp existingVar = setToVar.lookup(set);
      SubspanResourceInfo resource = {set, binding, subspanOp.getType(), false,
                                      existingVar};
      if (useBdaRootBindings) {
        interfaceToResourceInfo[subspanOp] = resource;
        continue;
      }
      if (!resource.var) {
        uint32_t maxBindingIdx = setToMaxBindingIdx[set];
        spirv::PointerType globalVarType =
            getIndirectBindingTableType(module->getContext(), maxBindingIdx,
                                        spirv::StorageClass::StorageBuffer);

        resource.var =
            createResourceVariable(subspanOp.getLoc(), resource, globalVarType,
                                   module, &symbolTable, /*isIndirect=*/true);
        setToVar[set] = resource.var;
      }
      assert(resource.var);
      interfaceToResourceInfo[subspanOp] = resource;
    }
  }

  return interfaceToResourceInfo;
}

} // namespace

//===----------------------------------------------------------------------===//
// Conversion patterns
//===----------------------------------------------------------------------===//

namespace {
/// A pattern to convert hal.interface.constant.load into a sequence of SPIR-V
/// ops to load from a global variable representing the push constant storage.
struct HALInterfaceLoadConstantConverter final
    : OpConversionPattern<IREE::HAL::InterfaceConstantLoadOp> {
  bool supportsAssume = false;
  bool useBdaRootBindings = false;

  HALInterfaceLoadConstantConverter(TypeConverter &typeConverter,
                                    MLIRContext *context, bool supportsAssume,
                                    bool useBdaRootBindings)
      : OpConversionPattern(typeConverter, context),
        supportsAssume(supportsAssume), useBdaRootBindings(useBdaRootBindings) {
  }

  LogicalResult
  matchAndRewrite(IREE::HAL::InterfaceConstantLoadOp loadOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Get the push constant count from the layout attribute on the op itself.
    auto layoutAttr = loadOp.getLayout();
    uint64_t elementCount = layoutAttr.getConstants();
    unsigned index = loadOp.getOrdinal().getZExtValue();
    if (useBdaRootBindings) {
      elementCount += kBdaDispatchRootDwordCount;
      index += kBdaDispatchRootDwordCount;
    }

    // The following function generates SPIR-V ops with i32 types. So it does
    // type "conversion" (index -> i32) implicitly. This is expected to be
    // paired with a cast (i32 -> index) afterwards.
    IntegerType i32Type = rewriter.getIntegerType(32);
    Value value =
        getPushConstantValue(loadOp, elementCount, index, i32Type, rewriter);

    if (loadOp.getResult().hasOneUse() && supportsAssume) {
      OpOperand *operand = loadOp.getResult().getUses().begin().getOperand();
      auto assumeOp = dyn_cast<IREE::Util::AssumeIntOp>(operand->getOwner());
      if (assumeOp) {
        Location loc = assumeOp.getLoc();
        unsigned opIdx = operand->getOperandNumber();

        auto [min, max] = assumeOp.getUnionedUnsignedRange(opIdx);
        if (min.has_value() && max.has_value()) {
          Value minConst = spirv::ConstantOp::create(
              rewriter, loc, i32Type, rewriter.getI32IntegerAttr(*min));
          Value maxConst = spirv::ConstantOp::create(
              rewriter, loc, i32Type, rewriter.getI32IntegerAttr(*max));
          Value minBound = spirv::UGreaterThanEqualOp::create(rewriter, loc,
                                                              value, minConst);
          spirv::KHRAssumeTrueOp::create(rewriter, loc, minBound);
          Value maxBound =
              spirv::ULessThanEqualOp::create(rewriter, loc, value, maxConst);
          spirv::KHRAssumeTrueOp::create(rewriter, loc, maxBound);
        }

        std::optional<uint64_t> divisibility =
            assumeOp.getUnionedUnsignedDivisor(opIdx);
        if (divisibility.has_value() && *divisibility > 1) {
          Value divisor = spirv::ConstantOp::create(
              rewriter, loc, i32Type,
              rewriter.getI32IntegerAttr(*divisibility));
          Value zero = spirv::ConstantOp::create(rewriter, loc, i32Type,
                                                 rewriter.getI32IntegerAttr(0));
          Value lowPart = spirv::UModOp::create(rewriter, loc, value, divisor);
          Value dividesExactly =
              spirv::IEqualOp::create(rewriter, loc, lowPart, zero);
          spirv::KHRAssumeTrueOp::create(rewriter, loc, dividesExactly);
        }
      }
    }

    rewriter.replaceOp(loadOp, value);
    return success();
  }
};

/// A pattern to convert util.assume.int into a noop, since we're using it
/// to annotate push constants.
struct UtilAssumeIntConverter final
    : OpConversionPattern<IREE::Util::AssumeIntOp> {
  using Base::Base;

  LogicalResult
  matchAndRewrite(IREE::Util::AssumeIntOp assumeOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    rewriter.replaceOp(assumeOp, adaptor.getOperands());
    return success();
  }
};

/// A pattern to convert hal.interface.workgroup.id/count/size into
/// corresponding SPIR-V Builtin ops.
template <typename InterfaceOpTy, spirv::BuiltIn builtin>
struct HALInterfaceWorkgroupOpsConverter final
    : OpConversionPattern<InterfaceOpTy> {
  using OpConversionPattern<InterfaceOpTy>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(InterfaceOpTy op, typename InterfaceOpTy::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    int32_t index = static_cast<int32_t>(op.getDimension().getSExtValue());
    auto i32Type = rewriter.getIntegerType(32);
    Value spirvBuiltin =
        spirv::getBuiltinVariableValue(op, builtin, i32Type, rewriter);
    Value spirvId = spirv::CompositeExtractOp::create(
        rewriter, spirvBuiltin.getLoc(), i32Type, spirvBuiltin,
        rewriter.getI32ArrayAttr({index}));

    // Casting if Indexing type not 32-bit.
    auto &typeConverter =
        *this->template getTypeConverter<SPIRVTypeConverter>();
    Type indexType = typeConverter.getIndexType();
    if (indexType != i32Type) {
      spirvId = spirv::UConvertOp::create(rewriter, spirvId.getLoc(), indexType,
                                          spirvId);
    }
    rewriter.replaceOp(op, spirvId);
    return success();
  }
};

/// A pattern to convert hal.interface.binding.subspan into a sequence of SPIR-V
/// ops to get the address to a global variable representing the resource
/// buffer.
struct HALInterfaceBindingSubspanConverter final
    : OpConversionPattern<IREE::HAL::InterfaceBindingSubspanOp> {
  HALInterfaceBindingSubspanConverter(
      TypeConverter &typeConverter, MLIRContext *context,
      const InterfaceResourceMap &interfaceToResourceVars,
      bool useIndirectBindings, bool useBdaRootBindings,
      PatternBenefit benefit = 1)
      : OpConversionPattern(typeConverter, context, benefit),
        interfaceToResourceVars(interfaceToResourceVars),
        useIndirectBindings(useIndirectBindings),
        useBdaRootBindings(useBdaRootBindings) {}

  LogicalResult
  matchAndRewrite(IREE::HAL::InterfaceBindingSubspanOp subspanOp,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (subspanOp.use_empty()) {
      rewriter.eraseOp(subspanOp);
      return success();
    }

    Value offset = subspanOp.getByteOffset();
    APInt offsetInt;
    if (offset && matchPattern(offset, m_ConstantInt(&offsetInt)) &&
        !offsetInt.isZero()) {
      return subspanOp.emitOpError() << "should have no or zero byte offset";
    }

    Type resultType = subspanOp.getType();
    auto convertedType =
        getTypeConverter()->convertType<spirv::PointerType>(resultType);
    if (!convertedType) {
      return subspanOp.emitError()
             << "failed to convert SPIR-V type: " << resultType;
    }
    SubspanResourceInfo info = interfaceToResourceVars.lookup(subspanOp);
    spirv::GlobalVariableOp varOp = info.var;

    if (!useIndirectBindings) {
      assert(varOp);
      // Fix up the variable's type.
      varOp.setTypeAttr(TypeAttr::get(convertedType));

      rewriter.replaceOpWithNewOp<spirv::AddressOfOp>(subspanOp, varOp);
      return success();
    }

    // Handle indirect bindings.
    if (convertedType.getStorageClass() !=
        spirv::StorageClass::PhysicalStorageBuffer) {
      return subspanOp->emitError(
          "indirect bindings require PhysicalStorageBuffer storage class");
    }

    Location loc = subspanOp.getLoc();
    auto i32Ty = rewriter.getI32Type();
    Value idx = spirv::ConstantOp::create(
        rewriter, loc, i32Ty, rewriter.getI32IntegerAttr(info.binding));
    Value bindingTable = nullptr;
    if (useBdaRootBindings) {
      Value tableAddress = loadBdaRootBindingTableAddress(subspanOp, rewriter);
      auto tablePtrType = getIndirectBindingTableType(
          rewriter.getContext(), info.binding,
          spirv::StorageClass::PhysicalStorageBuffer);
      bindingTable = spirv::ConvertUToPtrOp::create(rewriter, loc, tablePtrType,
                                                    tableAddress);
    } else {
      assert(varOp);
      bindingTable = spirv::AddressOfOp::create(rewriter, loc, varOp);
    }
    auto ptr = spirv::AccessChainOp::create(rewriter, loc, bindingTable, idx);
    auto pointerAlignedAccess = spirv::MemoryAccessAttr::get(
        rewriter.getContext(), spirv::MemoryAccess::Aligned);
    auto pointerAlignment =
        rewriter.getIntegerAttr(rewriter.getI32Type(), sizeof(uint64_t));
    auto addr = spirv::LoadOp::create(rewriter, loc, ptr, pointerAlignedAccess,
                                      pointerAlignment);
    assert(cast<spirv::PointerType>(addr.getType()).getStorageClass() ==
               spirv::StorageClass::PhysicalStorageBuffer &&
           "Expected a physical storage buffer pointer");

    // Bitcast the pointer to the correct pointer type. This is allowed for
    // physical storage buffer addresses.
    Value ptrInt = spirv::ConvertPtrToUOp::create(rewriter, loc,
                                                  rewriter.getI64Type(), addr);
    rewriter.replaceOpWithNewOp<spirv::ConvertUToPtrOp>(subspanOp,
                                                        convertedType, ptrInt);
    return success();
  }

private:
  Value
  loadBdaRootBindingTableAddress(IREE::HAL::InterfaceBindingSubspanOp subspanOp,
                                 ConversionPatternRewriter &rewriter) const {
    Location loc = subspanOp.getLoc();
    auto i32Type = rewriter.getI32Type();
    auto i64Type = rewriter.getI64Type();
    const unsigned pushConstantDwordCount =
        kBdaDispatchRootDwordCount +
        static_cast<unsigned>(subspanOp.getLayout().getConstants());
    Value low32 = getPushConstantValue(subspanOp, pushConstantDwordCount,
                                       /*offset=*/0, i32Type, rewriter);
    Value high32 = getPushConstantValue(subspanOp, pushConstantDwordCount,
                                        /*offset=*/1, i32Type, rewriter);
    Value low64 = spirv::UConvertOp::create(rewriter, loc, i64Type, low32);
    Value high64 = spirv::UConvertOp::create(rewriter, loc, i64Type, high32);
    Value shiftAmount = spirv::ConstantOp::create(
        rewriter, loc, i64Type, rewriter.getI64IntegerAttr(32));
    Value shiftedHigh64 =
        spirv::ShiftLeftLogicalOp::create(rewriter, loc, high64, shiftAmount);
    return spirv::BitwiseOrOp::create(rewriter, loc, low64, shiftedHigh64);
  }

  const InterfaceResourceMap &interfaceToResourceVars;
  const bool useIndirectBindings;
  const bool useBdaRootBindings;
};

/// Pattern to lower operations that become a no-ops at this level.
template <typename OpTy>
struct FoldAsNoOp final : OpConversionPattern<OpTy> {
  using OpConversionPattern<OpTy>::OpConversionPattern;
  LogicalResult
  matchAndRewrite(OpTy op, typename OpTy::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    rewriter.replaceOp(op, adaptor.getOperands());
    return success();
  }
};

/// Removes memref.cast that converts static and dynamic shapes.
struct RemoveStaticDynamicCast final : OpRewritePattern<memref::CastOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(memref::CastOp castOp,
                                PatternRewriter &rewriter) const override {
    auto srcType = cast<MemRefType>(castOp.getSource().getType());
    auto dstType = cast<MemRefType>(castOp.getType());
    if (srcType.getRank() == 1 && dstType.getRank() == 1 &&
        srcType.hasStaticShape() != dstType.hasStaticShape()) {
      rewriter.replaceOp(castOp, castOp.getSource());
      return success();
    }
    return failure();
  }
};

/// Removes unrealized_conversion_cast ops introduced during progressive
/// lowering when possible.
struct RemoveIdentityConversionCast final
    : OpConversionPattern<UnrealizedConversionCastOp> {
  using Base::Base;
  LogicalResult
  matchAndRewrite(UnrealizedConversionCastOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (op->getNumOperands() == 1 && op->getNumResults() == 1 &&
        adaptor.getOperands().front().getType() ==
            op->getResultTypes().front()) {
      rewriter.replaceOp(op, adaptor.getOperands());
      return success();
    }

    return failure();
  }
};

//===----------------------------------------------------------------------===//
// Conversion pass
//===----------------------------------------------------------------------===//

/// A pass to perform the SPIR-V conversion.
///
/// Converts remaining interface ops into SPIR-V global variables, GPU processor
/// ID ops into SPIR-V global variables, loop/standard ops into corresponding
/// SPIR-V ops.
class ConvertToSPIRVPass final
    : public impl::ConvertToSPIRVPassBase<ConvertToSPIRVPass> {
public:
  using Base::Base;
  explicit ConvertToSPIRVPass(unsigned indexBits) : indexBits(indexBits) {}

  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<spirv::SPIRVDialect>();
  }

  LogicalResult initializeOptions(
      StringRef options,
      function_ref<LogicalResult(const Twine &)> errorHandler) override {
    if (failed(Pass::initializeOptions(options, errorHandler))) {
      return failure();
    }
    indexBits = indexBitsOption;
    return success();
  }

  void runOnOperation() override;

private:
  // Use 64 bits for index widths.
  unsigned indexBits;
};
} // namespace

void ConvertToSPIRVPass::runOnOperation() {
  MLIRContext *context = &getContext();
  ModuleOp moduleOp = getOperation();

  if (moduleOp.getBody()->empty()) {
    return;
  }

  bool useIndirectBindings = usesIndirectBindingsAttr(moduleOp);
  bool useBdaRootBindings = usesBdaRootDispatchAbiAttr(moduleOp);

  // Build a map from function name to dispatch_config for workgroup_size
  // and subgroup_size lookup.
  DenseMap<StringRef, IREE::Codegen::DispatchConfigOp> configMap;
  for (auto configOp : moduleOp.getOps<IREE::Codegen::DispatchConfigOp>()) {
    configMap[configOp.getFunctionRef()] = configOp;
  }

  for (auto funcOp : moduleOp.getOps<mlir::FunctionOpInterface>()) {
    if (!configMap.contains(funcOp.getName())) {
      continue;
    }
    IREE::Codegen::DispatchConfigOp configOp = configMap[funcOp.getName()];
    if (funcOp->hasAttr(spirv::getEntryPointABIAttrName())) {
      continue;
    }
    SmallVector<int> workgroupSize;
    if (std::optional<ArrayRef<int64_t>> wgSize = configOp.getWorkgroupSize()) {
      workgroupSize = llvm::map_to_vector(wgSize.value(),
                                          [](int64_t v) -> int { return v; });
    }
    if (workgroupSize.empty()) {
      funcOp->emitOpError(
          "expected workgroup_size to be set for SPIR-V lowering");
      return signalPassFailure();
    }
    std::optional<int> subgroupSize;
    if (std::optional<uint64_t> sg = configOp.getSubgroupSize()) {
      subgroupSize = sg.value();
    }
    funcOp->setAttr(
        spirv::getEntryPointABIAttrName(),
        spirv::getEntryPointABIAttr(context, workgroupSize, subgroupSize));
  }

  for (auto funcOp : moduleOp.getOps<mlir::FunctionOpInterface>()) {
    RewritePatternSet shapePatterns(context);
    shapePatterns.add<RemoveStaticDynamicCast>(context);
    if (failed(applyPatternsGreedily(funcOp, std::move(shapePatterns)))) {
      funcOp.emitOpError() << "failed running shape patterns";
      return signalPassFailure();
    }
  }

  /// Rewrite extui/si(bitcast) as a mix of vector.shuffle + bitwise arithmetic.
  /// This handles cases like `vector.bitcast i8 to vector<2xi4>` that come from
  /// narrow load emulation by never materializing the sub-byte values. SPIR-V
  /// does not have support for arithmetic on sub-byte types so we currently
  /// rely on this rewrite for the cases seen today.
  /// TODO: Support general emulation of compute on sub-byte types. This is
  /// not mutually exclusive with this pattern, but does mean it is no longer
  /// load bearing.  Also these patterns are already run during
  /// `EmulateNarrotType` pass but dont trigger there due to missing support for
  /// emulation of `vector.transfer_read` in the emulation path. Remove the
  /// patterns from here after that is done.
  for (auto funcOp : moduleOp.getOps<mlir::FunctionOpInterface>()) {
    RewritePatternSet narrowingPatterns(context);
    vector::populateVectorNarrowTypeRewritePatterns(narrowingPatterns);
    if (failed(applyPatternsGreedily(funcOp, std::move(narrowingPatterns)))) {
      funcOp.emitOpError() << "failed running narrowing patterns";
      return signalPassFailure();
    }
  }

  // Expand any remaining `bf16` `extf` and `trunc` patterns.
  {
    RewritePatternSet patterns(context);
    arith::populateExpandBFloat16Patterns(patterns);
    arith::populateExpandF4E2M1Patterns(patterns);
    arith::populateExpandF8E8M0Patterns(patterns);
    arith::BitcastOp::getCanonicalizationPatterns(patterns, context);
    if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) {
      moduleOp.emitOpError() << "failed running bf16 extf/trunc patterns";
      return signalPassFailure();
    }
  }

  // Indirect bindings pass device addresses through 64-bit scalars regardless
  // of the module's normal index width preference.
  const unsigned effectiveIndexBits = useIndirectBindings ? 64 : indexBits;
  if (effectiveIndexBits != 32 && effectiveIndexBits != 64) {
    moduleOp.emitOpError(
        "only 32-bit or 64-bit indices are supported for SPIR-V");
    return signalPassFailure();
  }
  bool use64bitIndex = effectiveIndexBits == 64;

  auto targetAttr = moduleOp->getAttrOfType<spirv::TargetEnvAttr>(
      spirv::getTargetEnvAttrName());
  if (!targetAttr) {
    moduleOp.emitOpError("should contain a spirv.target_env attribute");
    return signalPassFailure();
  }
  spirv::TargetEnv targetEnv(targetAttr);

  if (use64bitIndex && !targetEnv.allows(spirv::Capability::Int64)) {
    moduleOp.emitOpError(
        "64-bit indices are not supported for the specified target "
        "environment");
    return signalPassFailure();
  }
  if (useIndirectBindings &&
      (!targetEnv.allows({spirv::Capability::PhysicalStorageBufferAddresses,
                          spirv::Capability::Int64}) ||
       !targetEnv.allows(spirv::Extension::SPV_KHR_physical_storage_buffer))) {
    moduleOp.emitOpError("indirect bindings are not supported for the "
                         "specified target environment");
    return signalPassFailure();
  }

  bool supportsAssume = targetEnv.allows(spirv::Capability::ExpectAssumeKHR);

  SPIRVConversionOptions options = {};
  options.use64bitIndex = use64bitIndex;

  SPIRVTypeConverter typeConverter(targetAttr, options);

  // Additionally pull in conversion rules for GPU subgroup MMA ops.
  populateMMAToSPIRVCoopMatrixTypeConversion(typeConverter);
  RewritePatternSet patterns(&getContext());
  ScfToSPIRVContext scfToSPIRVContext;

  // Pull in GPU patterns to convert processor ID ops and loop ops.
  populateGPUToSPIRVPatterns(typeConverter, patterns);
  populateGpuWMMAToSPIRVCoopMatrixKHRConversionPatterns(typeConverter,
                                                        patterns);

  // Pull in SCF patterns to convert control flow ops.
  populateSCFToSPIRVPatterns(typeConverter, scfToSPIRVContext, patterns);

  // Pull in MemRef patterns to convert load/store ops.
  populateMemRefToSPIRVPatterns(typeConverter, patterns);

  // Pull in standard/math patterns to convert arithmetic ops and others.
  arith::populateCeilFloorDivExpandOpsPatterns(patterns);
  arith::populateArithToSPIRVPatterns(typeConverter, patterns);
  populateFuncToSPIRVPatterns(typeConverter, patterns);
  populateMathToSPIRVPatterns(typeConverter, patterns);
  populateComplexToSPIRVPatterns(typeConverter, patterns);

  // Pull in standard patterns to convert tensor operations to SPIR-V. These are
  // primarily used to handle tensor-type constants and contain a
  // threshold. Only those constants that are below the threshold are converted
  // to SPIR-V. In IREE we want to control this threshold at Flow level. So set
  // this value arbitrarily high to make sure that everything within a dispatch
  // region is converted.
  mlir::populateTensorToSPIRVPatterns(
      typeConverter, std::numeric_limits<int64_t>::max() / 8, patterns);

  // Pull in vector patterns to convert vector ops.
  mlir::populateVectorToSPIRVPatterns(typeConverter, patterns);

  // Pull in builtin func to spirv.func conversion.
  populateBuiltinFuncToSPIRVPatterns(typeConverter, patterns);

  ub::populateUBToSPIRVConversionPatterns(typeConverter, patterns);

  // Add IREE HAL interface op conversions.
  patterns.add<
      HALInterfaceWorkgroupOpsConverter<IREE::HAL::InterfaceWorkgroupIDOp,
                                        spirv::BuiltIn::WorkgroupId>,
      HALInterfaceWorkgroupOpsConverter<IREE::HAL::InterfaceWorkgroupSizeOp,
                                        spirv::BuiltIn::WorkgroupSize>,
      HALInterfaceWorkgroupOpsConverter<IREE::HAL::InterfaceWorkgroupCountOp,
                                        spirv::BuiltIn::NumWorkgroups>,
      UtilAssumeIntConverter>(typeConverter, context);

  patterns.add<HALInterfaceLoadConstantConverter>(
      typeConverter, context, supportsAssume, useBdaRootBindings);

  // Performs a preliminary step to analyze all hal.interface.binding.subspan
  // ops and creates spirv.GlobalVariables.
  InterfaceResourceMap interfaceToResourceVars =
      useIndirectBindings
          ? createIndirectResourceVariables(moduleOp, useBdaRootBindings)
          : createResourceVariables(moduleOp);
  // For using use them in conversion.
  patterns.add<HALInterfaceBindingSubspanConverter>(
      typeConverter, context, interfaceToResourceVars, useIndirectBindings,
      useBdaRootBindings);

  /// Fold certain operations as no-ops:
  /// - linalg.reshape becomes a no-op since all memrefs are linearized in
  ///   SPIR-V.
  /// - tensor_to_memref can become a no-op since tensors are lowered to
  ///   !spirv.array.
  /// - unrealized_conversion_cast with the same source and target type.
  patterns.add<
      FoldAsNoOp<memref::CollapseShapeOp>, FoldAsNoOp<memref::ExpandShapeOp>,
      FoldAsNoOp<bufferization::ToBufferOp>, RemoveIdentityConversionCast>(
      typeConverter, context);

  std::unique_ptr<ConversionTarget> target =
      SPIRVConversionTarget::get(targetAttr);
  // Disallow all other ops.
  target->markUnknownOpDynamicallyLegal([](Operation *) { return false; });

  SmallVector<mlir::FunctionOpInterface, 1> functions;
  for (auto fn : moduleOp.getOps<mlir::FunctionOpInterface>()) {
    if (!fn.isPublic()) {
      continue;
    }
    functions.push_back(fn);
  }

  FrozenRewritePatternSet frozenPatterns(std::move(patterns));
  for (auto fn : functions) {
    if (failed(applyFullConversion(fn, *target, frozenPatterns))) {
      return signalPassFailure();
    }
  }

  auto addressingModel = spirv::AddressingModel::Logical;
  if (useIndirectBindings) {
    addressingModel = spirv::AddressingModel::PhysicalStorageBuffer64;
  }

  // Collect all SPIR-V ops into a spirv.module.
  OpBuilder builder = OpBuilder::atBlockBegin(moduleOp.getBody());
  auto spvModule = spirv::ModuleOp::create(
      builder, moduleOp.getLoc(), addressingModel, spirv::MemoryModel::GLSL450);
  Block *body = spvModule.getBody();
  Dialect *spvDialect = spvModule->getDialect();
  for (Operation &op : llvm::make_early_inc_range(*moduleOp.getBody())) {
    // Skip the newly created spirv.module itself.
    if (&op == spvModule) {
      continue;
    }
    if (op.getDialect() == spvDialect) {
      op.moveBefore(body, body->end());
    }
  }
}

//===----------------------------------------------------------------------===//
// Pass entry point and registration
//===----------------------------------------------------------------------===//

std::unique_ptr<OperationPass<ModuleOp>>
createConvertToSPIRVPass(unsigned indexBits) {
  return std::make_unique<ConvertToSPIRVPass>(indexBits);
}

} // namespace mlir::iree_compiler
