SPIRVPreLegalizer.cpp

Go to the documentation of this file.

//===-- SPIRVPreLegalizer.cpp - prepare IR for legalization -----*- C++ -*-===//
//
// Part of the LLVM Project, 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
//
//===----------------------------------------------------------------------===//
//
// The pass prepares IR for legalization: it assigns SPIR-V types to registers
// and removes intrinsics which holded these types during IR translation.
// Also it processes constants and registers them in GR to avoid duplication.
//
//===----------------------------------------------------------------------===//
 
#include "SPIRV.h"
#include "SPIRVSubtarget.h"
#include "SPIRVUtils.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/CodeGen/GlobalISel/CSEInfo.h"
#include "llvm/CodeGen/GlobalISel/GISelValueTracking.h"
#include "llvm/CodeGen/GlobalISel/MIPatternMatch.h"
#include "llvm/CodeGen/MachineFunctionAnalysisManager.h"
#include "llvm/CodeGen/MachinePassManager.h"
#include "llvm/IR/Analysis.h"
#include "llvm/IR/Attributes.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/IntrinsicsSPIRV.h"
#include "llvm/Support/MathExtras.h"
 
#define DEBUG_TYPE "spirv-prelegalizer"
 
using namespace llvm;
using namespace llvm::MIPatternMatch;
 
namespace {
class SPIRVPreLegalizerLegacy : public MachineFunctionPass {
public:
  static char ID;
  SPIRVPreLegalizerLegacy() : MachineFunctionPass(ID) {}
  bool runOnMachineFunction(MachineFunction &MF) override;
  void getAnalysisUsage(AnalysisUsage &AU) const override;
};
} // namespace
 
void SPIRVPreLegalizerLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
  AU.addPreserved<GISelValueTrackingAnalysisLegacy>();
  MachineFunctionPass::getAnalysisUsage(AU);
}
 
static inline void invalidateAndEraseMI(SPIRVGlobalRegistry *GR,
                                        MachineInstr *MI) {
  GR->invalidateMachineInstr(MI);
  MI->eraseFromParent();
}
 
static void
addConstantsToTrack(MachineFunction &MF, SPIRVGlobalRegistry *GR,
                    const SPIRVSubtarget &STI,
                    DenseMap<MachineInstr *, Type *> &TargetExtConstTypes) {
  MachineRegisterInfo &MRI = MF.getRegInfo();
  DenseMap<MachineInstr *, Register> RegsAlreadyAddedToDT;
  SmallVector<MachineInstr *, 10> ToErase, ToEraseComposites;
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      if (!isSpvIntrinsic(MI, Intrinsic::spv_track_constant))
        continue;
      ToErase.push_back(&MI);
      Register SrcReg = MI.getOperand(2).getReg();
      auto *Const =
          cast<Constant>(cast<ConstantAsMetadata>(
                             MI.getOperand(3).getMetadata()->getOperand(0))
                             ->getValue());
      if (auto *GV = dyn_cast<GlobalValue>(Const)) {
        Register Reg = GR->find(GV, &MF);
        if (!Reg.isValid()) {
          GR->add(GV, MRI.getVRegDef(SrcReg));
          GR->addGlobalObject(GV, &MF, SrcReg);
        } else
          RegsAlreadyAddedToDT[&MI] = Reg;
      } else {
        Register Reg = GR->find(Const, &MF);
        if (!Reg.isValid()) {
          if (auto *ConstVec = dyn_cast<ConstantDataVector>(Const)) {
            auto *BuildVec = MRI.getVRegDef(SrcReg);
            assert(BuildVec &&
                   BuildVec->getOpcode() == TargetOpcode::G_BUILD_VECTOR);
            GR->add(Const, BuildVec);
            for (unsigned i = 0; i < ConstVec->getNumElements(); ++i) {
              // Ensure that OpConstantComposite reuses a constant when it's
              // already created and available in the same machine function.
              Constant *ElemConst = ConstVec->getElementAsConstant(i);
              Register ElemReg = GR->find(ElemConst, &MF);
              if (!ElemReg.isValid())
                GR->add(ElemConst,
                        MRI.getVRegDef(BuildVec->getOperand(1 + i).getReg()));
              else
                BuildVec->getOperand(1 + i).setReg(ElemReg);
            }
          }
          if (Const->getType()->isTargetExtTy()) {
            // remember association so that we can restore it when assign types
            MachineInstr *SrcMI = MRI.getVRegDef(SrcReg);
            if (SrcMI)
              GR->add(Const, SrcMI);
            if (SrcMI && (SrcMI->getOpcode() == TargetOpcode::G_CONSTANT ||
                          SrcMI->getOpcode() == TargetOpcode::G_IMPLICIT_DEF))
              TargetExtConstTypes[SrcMI] = Const->getType();
            if (Const->isNullValue()) {
              MachineBasicBlock &DepMBB = MF.front();
              MachineIRBuilder MIB(DepMBB, DepMBB.getFirstNonPHI());
              SPIRVTypeInst ExtType = GR->getOrCreateSPIRVType(
                  Const->getType(), MIB, SPIRV::AccessQualifier::ReadWrite,
                  true);
              assert(SrcMI && "Expected source instruction to be valid");
              SrcMI->setDesc(STI.getInstrInfo()->get(SPIRV::OpConstantNull));
              SrcMI->addOperand(MachineOperand::CreateReg(
                  GR->getSPIRVTypeID(ExtType), false));
            }
          }
        } else {
          RegsAlreadyAddedToDT[&MI] = Reg;
          // This MI is unused and will be removed. If the MI uses
          // const_composite, it will be unused and should be removed too.
          assert(MI.getOperand(2).isReg() && "Reg operand is expected");
          MachineInstr *SrcMI = MRI.getVRegDef(MI.getOperand(2).getReg());
          if (SrcMI && isSpvIntrinsic(*SrcMI, Intrinsic::spv_const_composite))
            ToEraseComposites.push_back(SrcMI);
        }
      }
    }
  }
  for (MachineInstr *MI : ToErase) {
    Register Reg = MI->getOperand(2).getReg();
    auto It = RegsAlreadyAddedToDT.find(MI);
    if (It != RegsAlreadyAddedToDT.end())
      Reg = It->second;
    auto *RC = MRI.getRegClassOrNull(MI->getOperand(0).getReg());
    if (!MRI.getRegClassOrNull(Reg) && RC)
      MRI.setRegClass(Reg, RC);
    MRI.replaceRegWith(MI->getOperand(0).getReg(), Reg);
    invalidateAndEraseMI(GR, MI);
  }
  for (MachineInstr *MI : ToEraseComposites)
    invalidateAndEraseMI(GR, MI);
}
 
static void foldConstantsIntoIntrinsics(MachineFunction &MF,
                                        SPIRVGlobalRegistry *GR,
                                        MachineIRBuilder MIB) {
  SmallVector<MachineInstr *, 64> ToErase;
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      if (!isSpvIntrinsic(MI, Intrinsic::spv_assign_name))
        continue;
      const MDNode *MD = MI.getOperand(2).getMetadata();
      StringRef ValueName = cast<MDString>(MD->getOperand(0))->getString();
      if (ValueName.size() > 0) {
        MIB.setInsertPt(*MI.getParent(), MI);
        buildOpName(MI.getOperand(1).getReg(), ValueName, MIB);
      }
      ToErase.push_back(&MI);
    }
    for (MachineInstr *MI : ToErase)
      invalidateAndEraseMI(GR, MI);
    ToErase.clear();
  }
}
 
static MachineInstr *findAssignTypeInstr(Register Reg,
                                         MachineRegisterInfo *MRI) {
  for (MachineRegisterInfo::use_instr_iterator I = MRI->use_instr_begin(Reg),
                                               IE = MRI->use_instr_end();
       I != IE; ++I) {
    MachineInstr *UseMI = &*I;
    if ((isSpvIntrinsic(*UseMI, Intrinsic::spv_assign_ptr_type) ||
         isSpvIntrinsic(*UseMI, Intrinsic::spv_assign_type)) &&
        UseMI->getOperand(1).getReg() == Reg)
      return UseMI;
  }
  return nullptr;
}
 
static void buildOpBitcast(SPIRVGlobalRegistry *GR, MachineIRBuilder &MIB,
                           Register ResVReg, Register OpReg) {
  SPIRVTypeInst ResType = GR->getSPIRVTypeForVReg(ResVReg);
  SPIRVTypeInst OpType = GR->getSPIRVTypeForVReg(OpReg);
  assert(ResType && OpType && "Operand types are expected");
  if (!GR->isBitcastCompatible(ResType, OpType))
    report_fatal_error("incompatible result and operand types in a bitcast");
  MachineRegisterInfo *MRI = MIB.getMRI();
  if (!MRI->getRegClassOrNull(ResVReg))
    MRI->setRegClass(ResVReg, GR->getRegClass(ResType));
  if (ResType == OpType)
    MIB.buildInstr(TargetOpcode::COPY).addDef(ResVReg).addUse(OpReg);
  else
    MIB.buildInstr(SPIRV::OpBitcast)
        .addDef(ResVReg)
        .addUse(GR->getSPIRVTypeID(ResType))
        .addUse(OpReg);
}
 
// We lower G_BITCAST to OpBitcast here to avoid a MachineVerifier error.
// The verifier checks if the source and destination LLTs of a G_BITCAST are
// different, but this check is too strict for SPIR-V's typed pointers, which
// may have the same LLT but different SPIRV type (e.g. pointers to different
// pointee types). By lowering to OpBitcast here, we bypass the verifier's
// check. See discussion in https://github.com/llvm/llvm-project/pull/110270
// for more context.
//
// We also handle the llvm.spv.bitcast intrinsic here. If the source and
// destination SPIR-V types are the same, we lower it to a COPY to enable
// further optimizations like copy propagation.
static void lowerBitcasts(MachineFunction &MF, SPIRVGlobalRegistry *GR,
                          MachineIRBuilder MIB) {
  SmallVector<MachineInstr *, 16> ToErase;
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      if (isSpvIntrinsic(MI, Intrinsic::spv_bitcast)) {
        Register DstReg = MI.getOperand(0).getReg();
        Register SrcReg = MI.getOperand(2).getReg();
        SPIRVTypeInst DstType = GR->getSPIRVTypeForVReg(DstReg);
        assert(
            DstType &&
            "Expected destination SPIR-V type to have been assigned already.");
        SPIRVTypeInst SrcType = GR->getSPIRVTypeForVReg(SrcReg);
        assert(SrcType &&
               "Expected source SPIR-V type to have been assigned already.");
        if (DstType == SrcType) {
          MIB.setInsertPt(*MI.getParent(), MI);
          MIB.buildCopy(DstReg, SrcReg);
          ToErase.push_back(&MI);
          continue;
        }
      }
 
      if (MI.getOpcode() != TargetOpcode::G_BITCAST)
        continue;
 
      MIB.setInsertPt(*MI.getParent(), MI);
      buildOpBitcast(GR, MIB, MI.getOperand(0).getReg(),
                     MI.getOperand(1).getReg());
      ToErase.push_back(&MI);
    }
  }
  for (MachineInstr *MI : ToErase)
    invalidateAndEraseMI(GR, MI);
}
 
static void insertBitcasts(MachineFunction &MF, SPIRVGlobalRegistry *GR,
                           MachineIRBuilder MIB) {
  // Get access to information about available extensions
  const SPIRVSubtarget *ST =
      static_cast<const SPIRVSubtarget *>(&MIB.getMF().getSubtarget());
  SmallVector<MachineInstr *, 10> ToErase;
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      if (!isSpvIntrinsic(MI, Intrinsic::spv_ptrcast))
        continue;
      assert(MI.getOperand(2).isReg());
      MIB.setInsertPt(*MI.getParent(), MI);
      ToErase.push_back(&MI);
      Register Def = MI.getOperand(0).getReg();
      Register Source = MI.getOperand(2).getReg();
      Type *ElemTy = getMDOperandAsType(MI.getOperand(3).getMetadata(), 0);
      auto SC =
          isa<FunctionType>(ElemTy) &&
                  ST->canUseExtension(
                      SPIRV::Extension::SPV_INTEL_function_pointers)
              ? SPIRV::StorageClass::CodeSectionINTEL
              : addressSpaceToStorageClass(MI.getOperand(4).getImm(), *ST);
      SPIRVTypeInst AssignedPtrType =
          GR->getOrCreateSPIRVPointerType(ElemTy, MI, SC);
 
      // If the ptrcast would be redundant, replace all uses with the source
      // register.
      MachineRegisterInfo *MRI = MIB.getMRI();
      // For untyped pointers the SPIR-V pointer type does not encode the
      // pointee, so two pointers with different element types share the same
      // pointer type. The element type still matters because it selects the
      // Base Type operand of OpUntyped*AccessChainKHR. Treat the cast as
      // redundant only when the source already carries the same element type.
      // Otherwise keep a distinct register so the element type is preserved.
      bool Redundant =
          AssignedPtrType->getOpcode() == SPIRV::OpTypeUntypedPointerKHR
              ? GR->getUntypedPtrElementType(Source) ==
                    GR->getOrCreateSPIRVType(ElemTy, MIB,
                                             SPIRV::AccessQualifier::ReadWrite,
                                             /*EmitIR=*/true)
              : GR->getSPIRVTypeForVReg(Source) == AssignedPtrType;
      if (Redundant) {
        // Erase Def's assign type instruction if we are going to replace Def.
        if (MachineInstr *AssignMI = findAssignTypeInstr(Def, MRI))
          ToErase.push_back(AssignMI);
        MRI->replaceRegWith(Def, Source);
      } else {
        if (!GR->getSPIRVTypeForVReg(Def, &MF))
          GR->assignSPIRVTypeToVReg(AssignedPtrType, Def, MF);
        MIB.buildBitcast(Def, Source);
      }
    }
  }
  for (MachineInstr *MI : ToErase)
    invalidateAndEraseMI(GR, MI);
}
 
// Translating GV, IRTranslator sometimes generates following IR:
//   %1 = G_GLOBAL_VALUE
//   %2 = COPY %1
//   %3 = G_ADDRSPACE_CAST %2
//
// or
//
//  %1 = G_ZEXT %2
//  G_MEMCPY ... %2 ...
//
// New registers have no SPIRV type and no register class info.
//
// Set SPIRV type for GV, propagate it from GV to other instructions,
// also set register classes.
static SPIRVTypeInst propagateSPIRVType(MachineInstr *MI,
                                        SPIRVGlobalRegistry *GR,
                                        MachineRegisterInfo &MRI,
                                        MachineIRBuilder &MIB) {
  SPIRVTypeInst SpvType = nullptr;
  assert(MI && "Machine instr is expected");
  if (MI->getOperand(0).isReg()) {
    Register Reg = MI->getOperand(0).getReg();
    SpvType = GR->getSPIRVTypeForVReg(Reg);
    if (!SpvType) {
      switch (MI->getOpcode()) {
      case TargetOpcode::G_FCONSTANT:
      case TargetOpcode::G_CONSTANT: {
        MIB.setInsertPt(*MI->getParent(), MI);
        Type *Ty = MI->getOperand(1).getCImm()->getType();
        SpvType = GR->getOrCreateSPIRVType(
            Ty, MIB, SPIRV::AccessQualifier::ReadWrite, true);
        break;
      }
      case TargetOpcode::G_GLOBAL_VALUE: {
        MIB.setInsertPt(*MI->getParent(), MI);
        const GlobalValue *Global = MI->getOperand(1).getGlobal();
        Type *ElementTy = toTypedPointer(GR->getDeducedGlobalValueType(Global));
        unsigned AddrSpace = Global->getType()->getAddressSpace();
        // Function pointers use CodeSectionINTEL storage class in SPIR-V when
        // the SPV_INTEL_function_pointers extension is enabled.
        const SPIRVSubtarget &ST = MIB.getMF().getSubtarget<SPIRVSubtarget>();
        if (isa<Function>(Global) &&
            ST.canUseExtension(SPIRV::Extension::SPV_INTEL_function_pointers))
          AddrSpace =
              storageClassToAddressSpace(SPIRV::StorageClass::CodeSectionINTEL);
        auto *Ty = TypedPointerType::get(ElementTy, AddrSpace);
        SpvType = GR->getOrCreateSPIRVType(
            Ty, MIB, SPIRV::AccessQualifier::ReadWrite, true);
        break;
      }
      case TargetOpcode::G_ANYEXT:
      case TargetOpcode::G_SEXT:
      case TargetOpcode::G_ZEXT: {
        if (MI->getOperand(1).isReg()) {
          if (MachineInstr *DefInstr =
                  MRI.getVRegDef(MI->getOperand(1).getReg())) {
            if (SPIRVTypeInst Def =
                    propagateSPIRVType(DefInstr, GR, MRI, MIB)) {
              unsigned CurrentBW = GR->getScalarOrVectorBitWidth(Def);
              unsigned ExpectedBW =
                  std::max(MRI.getType(Reg).getScalarSizeInBits(), CurrentBW);
              unsigned NumElements = GR->getScalarOrVectorComponentCount(Def);
              SpvType = GR->getOrCreateSPIRVIntegerType(ExpectedBW, MIB);
              if (NumElements > 1)
                SpvType = GR->getOrCreateSPIRVVectorType(SpvType, NumElements,
                                                         MIB, true);
            }
          }
        }
        break;
      }
      case TargetOpcode::G_PTRTOINT:
        SpvType = GR->getOrCreateSPIRVIntegerType(
            MRI.getType(Reg).getScalarSizeInBits(), MIB);
        break;
      case TargetOpcode::G_TRUNC:
      case TargetOpcode::G_ADDRSPACE_CAST:
      case TargetOpcode::G_PTR_ADD:
      case TargetOpcode::COPY: {
        MachineOperand &Op = MI->getOperand(1);
        MachineInstr *Def = Op.isReg() ? MRI.getVRegDef(Op.getReg()) : nullptr;
        if (Def)
          SpvType = propagateSPIRVType(Def, GR, MRI, MIB);
        break;
      }
      default:
        break;
      }
      if (SpvType) {
        // check if the address space needs correction
        LLT RegType = MRI.getType(Reg);
        if (SpvType.isPointer() && RegType.isPointer() &&
            storageClassToAddressSpace(GR->getPointerStorageClass(SpvType)) !=
                RegType.getAddressSpace()) {
          // Don't correct CodeSectionINTEL back to Function for function
          // pointer G_GLOBAL_VALUE - the LLVM register has address space 0
          // but the SPIR-V type was intentionally set to CodeSectionINTEL.
          bool SkipCorrection =
              MI->getOpcode() == TargetOpcode::G_GLOBAL_VALUE &&
              GR->getPointerStorageClass(SpvType) ==
                  SPIRV::StorageClass::CodeSectionINTEL;
          if (!SkipCorrection) {
            const SPIRVSubtarget &ST =
                MI->getParent()->getParent()->getSubtarget<SPIRVSubtarget>();
            auto TSC =
                addressSpaceToStorageClass(RegType.getAddressSpace(), ST);
            SpvType = GR->changePointerStorageClass(SpvType, TSC, *MI);
          }
        }
        GR->assignSPIRVTypeToVReg(SpvType, Reg, MIB.getMF());
      }
      if (!MRI.getRegClassOrNull(Reg))
        MRI.setRegClass(Reg, SpvType ? GR->getRegClass(SpvType)
                                     : &SPIRV::iIDRegClass);
    }
  }
  return SpvType;
}
 
// To support current approach and limitations wrt. bit width here we widen a
// scalar register with a bit width greater than 1 to valid sizes and cap it to
// 128 width.
static unsigned widenBitWidthToNextPow2(unsigned BitWidth) {
  if (BitWidth == 1)
    return 1; // No need to widen 1-bit values
  return std::min(std::max<unsigned>(PowerOf2Ceil(BitWidth), 8u), 128u);
}
 
static std::optional<unsigned>
getNarrowScalarWidth(Register Reg, const MachineRegisterInfo &MRI) {
  LLT Ty = MRI.getType(Reg);
  if (!Ty.isScalar())
    return std::nullopt;
  unsigned W = Ty.getScalarSizeInBits();
  // <= and not == because widenBitWidthToNextPow2 caps at 128.
  if (widenBitWidthToNextPow2(W) <= W)
    return std::nullopt;
  return W;
}
 
static void widenScalarType(Register Reg, MachineRegisterInfo &MRI) {
  LLT RegType = MRI.getType(Reg);
  if (!RegType.isScalar())
    return;
  unsigned CurrentWidth = RegType.getScalarSizeInBits();
  unsigned NewWidth = widenBitWidthToNextPow2(CurrentWidth);
  if (NewWidth != CurrentWidth)
    MRI.setType(Reg, LLT::scalar(NewWidth));
}
 
static void widenCImmType(MachineOperand &MOP) {
  const ConstantInt *CImmVal = MOP.getCImm();
  unsigned CurrentWidth = CImmVal->getBitWidth();
  unsigned NewWidth = widenBitWidthToNextPow2(CurrentWidth);
  if (NewWidth != CurrentWidth) {
    // Replace the immediate value with the widened version
    MOP.setCImm(ConstantInt::get(CImmVal->getType()->getContext(),
                                 CImmVal->getValue().zextOrTrunc(NewWidth)));
  }
}
 
static void setInsertPtAfterDef(MachineIRBuilder &MIB, MachineInstr *Def) {
  MachineBasicBlock &MBB = *Def->getParent();
  MachineBasicBlock::iterator DefIt =
      Def->getNextNode() ? Def->getNextNode()->getIterator() : MBB.end();
  // Skip all the PHI and debug instructions.
  while (DefIt != MBB.end() &&
         (DefIt->isPHI() || DefIt->isDebugOrPseudoInstr()))
    DefIt = std::next(DefIt);
  MIB.setInsertPt(MBB, DefIt);
}
 
namespace llvm {
void updateRegType(Register Reg, Type *Ty, SPIRVTypeInst SpvType,
                   SPIRVGlobalRegistry *GR, MachineIRBuilder &MIB,
                   MachineRegisterInfo &MRI) {
  assert((Ty || SpvType) && "Either LLVM or SPIRV type is expected.");
  MachineInstr *Def = MRI.getVRegDef(Reg);
  setInsertPtAfterDef(MIB, Def);
  if (!SpvType)
    SpvType = GR->getOrCreateSPIRVType(Ty, MIB,
                                       SPIRV::AccessQualifier::ReadWrite, true);
  if (!MRI.getRegClassOrNull(Reg))
    MRI.setRegClass(Reg, GR->getRegClass(SpvType));
  if (!MRI.getType(Reg).isValid())
    MRI.setType(Reg, GR->getRegType(SpvType));
  GR->assignSPIRVTypeToVReg(SpvType, Reg, MIB.getMF());
}
 
void processInstr(MachineInstr &MI, MachineIRBuilder &MIB,
                  MachineRegisterInfo &MRI, SPIRVGlobalRegistry *GR,
                  SPIRVTypeInst KnownResType) {
  MIB.setInsertPt(*MI.getParent(), MI.getIterator());
  for (auto &Op : MI.operands()) {
    if (!Op.isReg() || Op.isDef())
      continue;
    Register OpReg = Op.getReg();
    SPIRVTypeInst SpvType = GR->getSPIRVTypeForVReg(OpReg);
    if (!SpvType && KnownResType) {
      SpvType = KnownResType;
      GR->assignSPIRVTypeToVReg(KnownResType, OpReg, *MI.getMF());
    }
    assert(SpvType);
    if (!MRI.getRegClassOrNull(OpReg))
      MRI.setRegClass(OpReg, GR->getRegClass(SpvType));
    if (!MRI.getType(OpReg).isValid())
      MRI.setType(OpReg, GR->getRegType(SpvType));
  }
}
} // namespace llvm
 
// Sign-sensitive integer ops: their result depends on the value of the input
// sign bit at position (width-1). On sub-pow2 widths the general widening
// loop is a pure LLT relabel, which leaves the sign bit at the *original*
// position instead of the widened MSB. These ops therefore need an explicit
// G_SEXT_INREG on each value operand to move the sign bit up.
//
// Signed-vs-unsigned G_ICMP is distinguished by its predicate operand.
//
// TODO: follow-up PRs will add the remaining sign-sensitive opcodes
// (e.g. G_SADDSAT/G_SSUBSAT, signed overflow ops).
static bool isSignSensitiveOp(const MachineInstr &MI) {
  switch (MI.getOpcode()) {
  case TargetOpcode::G_ASHR:
  case TargetOpcode::G_SDIV:
  case TargetOpcode::G_SREM:
  case TargetOpcode::G_SITOFP:
  case TargetOpcode::G_SMIN:
  case TargetOpcode::G_SMAX:
    return true;
  case TargetOpcode::G_ICMP:
    return CmpInst::isSigned(
        static_cast<CmpInst::Predicate>(MI.getOperand(1).getPredicate()));
  default:
    return false;
  }
}
 
struct NarrowWideningInfo {
  // Width before widening of each sign-sensitive value-operand vreg (one entry
  // per vreg).
  DenseMap<Register, unsigned> OrigWidth;
  // Sign-sensitive ops whose value operand(s) need replacing, ordered for
  // reproducible vreg numbering.
  SmallVector<MachineInstr *> SignSensitiveWorklist;
  // Keyed by instruction, not vreg: G_TRUNC handling can replace the source.
  SmallVector<std::pair<MachineInstr *, unsigned>> BitCountWorklist;
};
 
// G_CTTZ_ZERO_POISON is absent because its low bits are known non-zero, G_CTLS
// because the backend does not select it.
static bool isWidthSensitiveBitCountOp(unsigned Opcode) {
  switch (Opcode) {
  case TargetOpcode::G_CTLZ:
  case TargetOpcode::G_CTLZ_ZERO_POISON:
  case TargetOpcode::G_CTTZ:
  case TargetOpcode::G_CTPOP:
    return true;
  default:
    return false;
  }
}
 
// Collect ops whose semantics depend on the operand width along with their
// pre-widening widths, before later passes retype those vregs to pow2 LLTs
// and the original width is no longer recoverable.
static NarrowWideningInfo
recordNarrowOperandWidths(MachineFunction &MF, const MachineRegisterInfo &MRI) {
  NarrowWideningInfo Info;
  auto RecordIfNarrow = [&](Register Reg) {
    std::optional<unsigned> W = getNarrowScalarWidth(Reg, MRI);
    if (!W)
      return false;
    Info.OrigWidth.try_emplace(Reg, *W);
    return true;
  };
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      if (isWidthSensitiveBitCountOp(MI.getOpcode())) {
        if (std::optional<unsigned> W =
                getNarrowScalarWidth(MI.getOperand(1).getReg(), MRI))
          Info.BitCountWorklist.emplace_back(&MI, *W);
        continue;
      }
      if (!isSignSensitiveOp(MI))
        continue;
      bool NeedsRewrite = false;
      for (const MachineOperand &MO : MI.all_uses())
        NeedsRewrite = RecordIfNarrow(MO.getReg()) || NeedsRewrite;
      if (NeedsRewrite) {
        Info.SignSensitiveWorklist.push_back(&MI);
        // Record the result too so it gets masked.
        RecordIfNarrow(MI.getOperand(0).getReg());
      }
    }
  }
  return Info;
}
 
// For every recorded sign-sensitive op, insert G_SEXT_INREG on each value
// operand whose original width was narrower than the widened pow2 width and
// retype the operand's vreg LLT in place to the widened width.
//
// Info must have been populated by recordNarrowOperandWidths before
// other passes retyped the vregs; otherwise the narrow widths needed here
// are lost.
//
// TODO: handle vector operands.
static void widenSignSensitiveOps(MachineFunction &MF, SPIRVGlobalRegistry *GR,
                                  MachineIRBuilder &MIB,
                                  MachineRegisterInfo &MRI,
                                  const NarrowWideningInfo &Info) {
  // Emit G_SEXT_INREG from Reg's recorded narrow width; retypes Reg to the
  // widened width and returns the sign-extended vreg.
  auto SignExtendReg = [&](Register Reg, unsigned OldW,
                           MachineInstr &MI) -> Register {
    unsigned NewW = widenBitWidthToNextPow2(OldW);
    LLT NewLLT = LLT::scalar(NewW);
    MIB.setInsertPt(*MI.getParent(), MI.getIterator());
    SPIRVTypeInst SpvTy = GR->getOrCreateSPIRVIntegerType(NewW, MIB);
    Register SExted = MRI.createGenericVirtualRegister(NewLLT);
    GR->assignSPIRVTypeToVReg(SpvTy, SExted, MF);
    MRI.setRegClass(SExted, GR->getRegClass(SpvTy));
    MRI.setType(Reg, NewLLT);
    MIB.buildSExtInReg(SExted, Reg, OldW);
    return SExted;
  };
 
  // The wide op yields a sign-extended value, mask it back to OldW bits because
  // later users assume the upper bits are zero.
  auto MaskResult = [&](MachineInstr &MI, unsigned OldW) {
    Register DstReg = MI.getOperand(0).getReg();
    widenScalarType(DstReg, MRI);
    MIB.setInstrAndDebugLoc(MI);
    SPIRVTypeInst SpvTy =
        GR->getOrCreateSPIRVIntegerType(widenBitWidthToNextPow2(OldW), MIB);
    Register Result = createVirtualRegister(SpvTy, GR, MIB);
    MI.getOperand(0).setReg(Result);
    setInsertPtAfterDef(MIB, &MI);
    MIB.buildZExtInReg(DstReg, Result, OldW);
  };
 
  // TODO: when the same narrow vreg feeds multiple sign-sensitive ops (e.g.
  // sdiv %x, %y and srem %x, %y), emit one shared G_SEXT_INREG instead of one
  // per use.
  const TargetRegisterInfo &TRI = *MRI.getTargetRegisterInfo();
  for (MachineInstr *MI : Info.SignSensitiveWorklist) {
    for (const MachineOperand &MO : MI->all_uses()) {
      Register Reg = MO.getReg();
      auto It = Info.OrigWidth.find(Reg);
      if (It == Info.OrigWidth.end())
        continue;
      // substituteRegister fills every slot holding Reg at once: SignExtendReg
      // retypes Reg in place, so a second sext would read the widened width.
      MI->substituteRegister(Reg, SignExtendReg(Reg, It->second, *MI),
                             /*SubIdx=*/0, TRI);
    }
    Register Dst = MI->getOperand(0).getReg();
    if (auto It = Info.OrigWidth.find(Dst); It != Info.OrigWidth.end())
      MaskResult(*MI, It->second);
  }
}
 
// LegalizerHelper::widenScalar has the same cases but cannot be reached: the
// relabel retypes every narrow scalar to a pow2 LLT, so no illegal narrow type
// ever reaches the legalizer.
//
// TODO: handle vector operands.
static void widenBitCountOps(SPIRVGlobalRegistry *GR, MachineIRBuilder &MIB,
                             MachineRegisterInfo &MRI,
                             const NarrowWideningInfo &Info) {
  for (auto [MI, OldWidth] : Info.BitCountWorklist) {
    Register SrcReg = MI->getOperand(1).getReg();
    unsigned NewWidth = widenBitWidthToNextPow2(OldWidth);
    LLT NewTy = LLT::scalar(NewWidth);
    widenScalarType(SrcReg, MRI);
    MIB.setInstrAndDebugLoc(*MI);
    SPIRVTypeInst SpvTy = GR->getOrCreateSPIRVIntegerType(NewWidth, MIB);
 
    // The G_TRUNC lowering masks its result to the narrow width, so a source
    // coming from it needs no second mask.
    APInt Cst;
    bool HighBitsAlreadyZero =
        mi_match(SrcReg, MRI, m_GAnd(m_Reg(), m_ICst(Cst))) &&
        Cst.isSubsetOf(APInt::getLowBitsSet(Cst.getBitWidth(), OldWidth));
    auto ClearHighBits = [&](unsigned Width) -> Register {
      if (HighBitsAlreadyZero)
        return SrcReg;
      Register Masked = createVirtualRegister(SpvTy, GR, MIB);
      MIB.buildZExtInReg(Masked, SrcReg, Width);
      return Masked;
    };
 
    Register Input;
    switch (MI->getOpcode()) {
    case TargetOpcode::G_CTLZ_ZERO_POISON: {
      // Shifting up to the widened MSB moves the poison out too, so no
      // adjustment.
      Input = createVirtualRegister(SpvTy, GR, MIB);
      auto Diff = MIB.buildConstant(NewTy, NewWidth - OldWidth);
      MIB.buildShl(Input, SrcReg, Diff);
      break;
    }
    case TargetOpcode::G_CTTZ: {
      // Keeps an all-zero narrow value counting exactly OldWidth zeros.
      Input = createVirtualRegister(SpvTy, GR, MIB);
      auto TopBit =
          MIB.buildConstant(NewTy, APInt::getOneBitSet(NewWidth, OldWidth));
      MIB.buildOr(Input, SrcReg, TopBit);
      break;
    }
    case TargetOpcode::G_CTPOP:
      Input = ClearHighBits(OldWidth);
      break;
    case TargetOpcode::G_CTLZ: {
      // Clearing the extra bits adds leading zeros the count has to drop.
      Input = ClearHighBits(OldWidth);
      Register DstReg = MI->getOperand(0).getReg();
      widenScalarType(DstReg, MRI);
      Register Count = createVirtualRegister(SpvTy, GR, MIB);
      MI->getOperand(0).setReg(Count);
      setInsertPtAfterDef(MIB, MI);
      auto Diff = MIB.buildConstant(NewTy, NewWidth - OldWidth);
      MIB.buildSub(DstReg, Count, Diff);
      break;
    }
    default:
      llvm_unreachable("unexpected width-sensitive bit-count opcode");
    }
    MI->getOperand(1).setReg(Input);
  }
}
 
static void
generateAssignInstrs(MachineFunction &MF, SPIRVGlobalRegistry *GR,
                     MachineIRBuilder MIB,
                     DenseMap<MachineInstr *, Type *> &TargetExtConstTypes) {
  // Get access to information about available extensions
  const SPIRVSubtarget *ST =
      static_cast<const SPIRVSubtarget *>(&MIB.getMF().getSubtarget());
 
  MachineRegisterInfo &MRI = MF.getRegInfo();
  SmallVector<MachineInstr *, 10> ToErase;
  DenseMap<MachineInstr *, Register> RegsAlreadyAddedToDT;
 
  bool IsExtendedInts =
      ST->canUseExtension(
          SPIRV::Extension::SPV_ALTERA_arbitrary_precision_integers) ||
      ST->canUseExtension(SPIRV::Extension::SPV_KHR_bit_instructions) ||
      ST->canUseExtension(SPIRV::Extension::SPV_INTEL_int4);
 
  if (!IsExtendedInts) {
    // Without arbitrary precision integer extensions, SPIR-V only supports
    // integer widths of 8, 16, 32, 64. Non-standard widths (e.g., i24, i40)
    // must be widened to the next power of two.
    //
    // Record the original widths of width-sensitive operands before either
    // the G_TRUNC handling or the general widening loop retypes vregs, then
    // rewrite those ops after G_TRUNC processing using the recorded widths.
    NarrowWideningInfo WideningInfo = recordNarrowOperandWidths(MF, MRI);
 
    // G_TRUNC requires special handling because its semantics depend on the
    // original destination width. For example:
    //   %dst:s24 = G_TRUNC %src:s64
    // After widening s24 to s32, we cannot simply do:
    //   %dst:s32 = G_TRUNC %src:s64
    // because this would keep 32 bits instead of 24. Instead, we insert a
    // G_AND to mask the value to the original width:
    //   %mask:s64 = G_CONSTANT 0xFFFFFF      ; 24-bit mask
    //   %masked:s64 = G_AND %src:s64, %mask
    //   %dst:s32 = G_TRUNC %masked:s64
    // If src and dst widen to the same size, G_TRUNC is replaced entirely:
    //   %mask:s64 = G_CONSTANT 0xFFFFFFFFFF  ; 40-bit mask
    //   %dst:s64 = G_AND %src:s64, %mask
    SmallVector<MachineInstr *, 8> TruncToRemove;
    for (MachineBasicBlock &MBB : MF) {
      for (MachineInstr &MI : MBB) {
        unsigned MIOp = MI.getOpcode();
        if (MIOp != TargetOpcode::G_TRUNC)
          continue;
        assert(MI.getNumOperands() == 2);
        assert(MI.getOperand(0).isReg());
        assert(MI.getOperand(1).isReg());
 
        Register DstReg = MI.getOperand(0).getReg();
        Register SrcReg = MI.getOperand(1).getReg();
 
        LLT DstTy = MRI.getType(DstReg);
        LLT SrcTy = MRI.getType(SrcReg);
        assert((DstTy.isScalar() || DstTy.isVector()) &&
               (SrcTy.isScalar() || SrcTy.isVector()) &&
               "Expected scalar or vector G_TRUNC types");
        assert(DstTy.isVector() == SrcTy.isVector() &&
               "Expected matching scalar/vector G_TRUNC types");
        assert((!DstTy.isVector() ||
                DstTy.getElementCount() == SrcTy.getElementCount()) &&
               "Expected equal vector element counts");
 
        unsigned OriginalDstWidth = DstTy.getScalarSizeInBits();
        unsigned OriginalSrcWidth = SrcTy.getScalarSizeInBits();
 
        unsigned NewDstWidth = widenBitWidthToNextPow2(OriginalDstWidth);
        unsigned NewSrcWidth = widenBitWidthToNextPow2(OriginalSrcWidth);
        LLT NewDstTy = DstTy.changeElementSize(NewDstWidth);
        LLT NewSrcTy = SrcTy.changeElementSize(NewSrcWidth);
 
        // No Dst width change means no truncation semantics change, but the
        // source still needs a legal type.
        if (OriginalDstWidth == NewDstWidth) {
          MRI.setType(SrcReg, NewSrcTy);
          continue;
        }
 
        MRI.setType(SrcReg, NewSrcTy);
        MRI.setType(DstReg, NewDstTy);
 
        MIB.setInsertPt(MBB, MI.getIterator());
        APInt Mask = APInt::getLowBitsSet(NewSrcWidth, OriginalDstWidth);
        MachineInstrBuilder MaskReg =
            DstTy.isVector()
                ? MIB.buildBuildVectorConstant(
                      NewSrcTy,
                      SmallVector<APInt, 4>(DstTy.getNumElements(), Mask))
                : MIB.buildConstant(NewSrcTy, Mask);
        Register MaskedReg = MRI.createGenericVirtualRegister(NewSrcTy);
        MIB.buildAnd(MaskedReg, SrcReg, MaskReg);
 
        if (NewSrcWidth == NewDstWidth) {
          // Rekey OrigWidth from DstReg to MaskedReg so widenSignSensitiveOps
          // still sees the narrow original width after replaceRegWith.
          if (auto It = WideningInfo.OrigWidth.find(DstReg);
              It != WideningInfo.OrigWidth.end()) {
            unsigned W = It->second;
            WideningInfo.OrigWidth.erase(It);
            WideningInfo.OrigWidth.try_emplace(MaskedReg, W);
          }
          MRI.replaceRegWith(DstReg, MaskedReg);
          TruncToRemove.push_back(&MI);
        } else {
          MI.getOperand(1).setReg(MaskedReg);
        }
      }
    }
    for (MachineInstr *MI : TruncToRemove)
      MI->eraseFromParent();
 
    widenSignSensitiveOps(MF, GR, MIB, MRI, WideningInfo);
    widenBitCountOps(GR, MIB, MRI, WideningInfo);
  }
 
  for (MachineBasicBlock *MBB : post_order(&MF)) {
    if (MBB->empty())
      continue;
 
    bool ReachedBegin = false;
    for (auto MII = std::prev(MBB->end()), Begin = MBB->begin();
         !ReachedBegin;) {
      MachineInstr &MI = *MII;
      unsigned MIOp = MI.getOpcode();
 
      if (!IsExtendedInts) {
        // validate bit width of scalar registers and constant immediates
        for (auto &MOP : MI.operands()) {
          if (MOP.isReg())
            widenScalarType(MOP.getReg(), MRI);
          else if (MOP.isCImm())
            widenCImmType(MOP);
        }
      }
 
      if (isSpvIntrinsic(MI, Intrinsic::spv_assign_ptr_type)) {
        Register Reg = MI.getOperand(1).getReg();
        MIB.setInsertPt(*MI.getParent(), MI.getIterator());
        Type *ElementTy = getMDOperandAsType(MI.getOperand(2).getMetadata(), 0);
        auto SC = addressSpaceToStorageClass(MI.getOperand(3).getImm(), *ST);
        if (SC == SPIRV::StorageClass::Function &&
            isa<FunctionType>(ElementTy) &&
            ST->canUseExtension(SPIRV::Extension::SPV_INTEL_function_pointers))
          SC = SPIRV::StorageClass::CodeSectionINTEL;
        SPIRVTypeInst AssignedPtrType =
            GR->getOrCreateSPIRVPointerType(ElementTy, MI, SC);
 
        // For untyped pointers, store the element type for later use.
        if (ST->canUseExtension(SPIRV::Extension::SPV_KHR_untyped_pointers) &&
            !ST->isShader()) {
          SPIRVTypeInst ElemSpvType = GR->getOrCreateSPIRVType(
              ElementTy, MIB, SPIRV::AccessQualifier::ReadWrite,
              /*EmitIR=*/true);
          GR->setUntypedPtrElementType(Reg, ElemSpvType);
        }
 
        // The intrinsic also carries vector-of-pointer values produced by
        // scalarized vector GEPs; wrap the pointer in OpTypeVector to match
        // the vreg's LLT.
        LLT RegTy = MRI.getType(Reg);
        if (RegTy.isValid() && RegTy.isVector())
          AssignedPtrType = GR->getOrCreateSPIRVVectorType(
              AssignedPtrType, RegTy.getNumElements(), MIB,
              /*EmitIR=*/true);
        MachineInstr *Def = MRI.getVRegDef(Reg);
        assert(Def && "Expecting an instruction that defines the register");
        // G_GLOBAL_VALUE already has type info.
        if (Def->getOpcode() != TargetOpcode::G_GLOBAL_VALUE)
          updateRegType(Reg, nullptr, AssignedPtrType, GR, MIB,
                        MF.getRegInfo());
        ToErase.push_back(&MI);
      } else if (isSpvIntrinsic(MI, Intrinsic::spv_assign_type)) {
        Register Reg = MI.getOperand(1).getReg();
        Type *Ty = getMDOperandAsType(MI.getOperand(2).getMetadata(), 0);
        MachineInstr *Def = MRI.getVRegDef(Reg);
        assert(Def && "Expecting an instruction that defines the register");
        // G_GLOBAL_VALUE already has type info.
        if (Def->getOpcode() != TargetOpcode::G_GLOBAL_VALUE)
          updateRegType(Reg, Ty, nullptr, GR, MIB, MF.getRegInfo());
        if (Def->getOpcode() == TargetOpcode::COPY && isVector1(Ty))
          updateRegType(passCopy(Def, &MF.getRegInfo())->getOperand(0).getReg(),
                        Ty, nullptr, GR, MIB, MF.getRegInfo());
        ToErase.push_back(&MI);
      } else if (MIOp == TargetOpcode::FAKE_USE && MI.getNumOperands() > 0) {
        MachineInstr *MdMI = MI.getPrevNode();
        if (MdMI && isSpvIntrinsic(*MdMI, Intrinsic::spv_value_md)) {
          // It's an internal service info from before IRTranslator passes.
          MachineInstr *Def = getVRegDef(MRI, MI.getOperand(0).getReg());
          for (unsigned I = 1, E = MI.getNumOperands(); I != E && Def; ++I)
            if (getVRegDef(MRI, MI.getOperand(I).getReg()) != Def)
              Def = nullptr;
          if (Def) {
            const MDNode *MD = MdMI->getOperand(1).getMetadata();
            StringRef ValueName =
                cast<MDString>(MD->getOperand(1))->getString();
            const MDNode *TypeMD = cast<MDNode>(MD->getOperand(0));
            Type *ValueTy = getMDOperandAsType(TypeMD, 0);
            GR->addValueAttrs(Def, std::make_pair(ValueTy, ValueName.str()));
          }
          ToErase.push_back(MdMI);
        }
        ToErase.push_back(&MI);
      } else if (MIOp == TargetOpcode::G_CONSTANT ||
                 MIOp == TargetOpcode::G_FCONSTANT ||
                 MIOp == TargetOpcode::G_BUILD_VECTOR) {
        // %rc = G_CONSTANT ty Val
        // Ensure %rc has a valid SPIR-V type assigned in the Global Registry.
        Register Reg = MI.getOperand(0).getReg();
        bool NeedAssignType = !GR->getSPIRVTypeForVReg(Reg);
        Type *Ty = nullptr;
        if (MIOp == TargetOpcode::G_CONSTANT) {
          auto TargetExtIt = TargetExtConstTypes.find(&MI);
          Ty = TargetExtIt == TargetExtConstTypes.end()
                   ? MI.getOperand(1).getCImm()->getType()
                   : TargetExtIt->second;
          const ConstantInt *OpCI = MI.getOperand(1).getCImm();
          // TODO: we may wish to analyze here if OpCI is zero and LLT RegType =
          // MRI.getType(Reg); RegType.isPointer() is true, so that we observe
          // at this point not i64/i32 constant but null pointer in the
          // corresponding address space of RegType.getAddressSpace(). This may
          // help to successfully validate the case when a OpConstantComposite's
          // constituent has type that does not match Result Type of
          // OpConstantComposite (see, for example,
          // pointers/PtrCast-null-in-OpSpecConstantOp.ll).
          Register PrimaryReg = GR->find(OpCI, &MF);
          if (!PrimaryReg.isValid()) {
            GR->add(OpCI, &MI);
          } else if (PrimaryReg != Reg &&
                     MRI.getType(Reg) == MRI.getType(PrimaryReg)) {
            auto *RCReg = MRI.getRegClassOrNull(Reg);
            auto *RCPrimary = MRI.getRegClassOrNull(PrimaryReg);
            if (!RCReg || RCPrimary == RCReg) {
              RegsAlreadyAddedToDT[&MI] = PrimaryReg;
              ToErase.push_back(&MI);
              NeedAssignType = false;
            }
          }
        } else if (MIOp == TargetOpcode::G_FCONSTANT) {
          Ty = MI.getOperand(1).getFPImm()->getType();
        } else {
          assert(MIOp == TargetOpcode::G_BUILD_VECTOR);
          Type *ElemTy = nullptr;
          MachineInstr *ElemMI = MRI.getVRegDef(MI.getOperand(1).getReg());
          assert(ElemMI);
 
          if (ElemMI->getOpcode() == TargetOpcode::G_CONSTANT) {
            ElemTy = ElemMI->getOperand(1).getCImm()->getType();
          } else if (ElemMI->getOpcode() == TargetOpcode::G_FCONSTANT) {
            ElemTy = ElemMI->getOperand(1).getFPImm()->getType();
          } else {
            if (SPIRVTypeInst ElemSpvType =
                    GR->getSPIRVTypeForVReg(MI.getOperand(1).getReg(), &MF))
              ElemTy = const_cast<Type *>(GR->getTypeForSPIRVType(ElemSpvType));
          }
          if (ElemTy)
            Ty = VectorType::get(
                ElemTy, MI.getNumExplicitOperands() - MI.getNumExplicitDefs(),
                false);
          else
            NeedAssignType = false;
        }
        if (NeedAssignType)
          updateRegType(Reg, Ty, nullptr, GR, MIB, MRI);
      } else if (MIOp == TargetOpcode::G_GLOBAL_VALUE) {
        propagateSPIRVType(&MI, GR, MRI, MIB);
      }
 
      if (MII == Begin)
        ReachedBegin = true;
      else
        --MII;
    }
  }
  for (MachineInstr *MI : ToErase) {
    auto It = RegsAlreadyAddedToDT.find(MI);
    if (It != RegsAlreadyAddedToDT.end())
      MRI.replaceRegWith(MI->getOperand(0).getReg(), It->second);
    invalidateAndEraseMI(GR, MI);
  }
 
  // Address the case when IRTranslator introduces instructions with new
  // registers without associated SPIRV type.
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      switch (MI.getOpcode()) {
      case TargetOpcode::G_TRUNC:
      case TargetOpcode::G_ANYEXT:
      case TargetOpcode::G_SEXT:
      case TargetOpcode::G_ZEXT:
      case TargetOpcode::G_PTRTOINT:
      case TargetOpcode::COPY:
      case TargetOpcode::G_ADDRSPACE_CAST:
        propagateSPIRVType(&MI, GR, MRI, MIB);
        break;
      }
    }
  }
}
 
static void processInstrsWithTypeFolding(MachineFunction &MF,
                                         SPIRVGlobalRegistry *GR,
                                         MachineIRBuilder MIB) {
  MachineRegisterInfo &MRI = MF.getRegInfo();
  for (MachineBasicBlock &MBB : MF)
    for (MachineInstr &MI : MBB)
      if (isTypeFoldingSupported(MI.getOpcode()))
        processInstr(MI, MIB, MRI, GR, nullptr);
}
 
static Register
collectInlineAsmInstrOperands(MachineInstr *MI,
                              SmallVector<unsigned, 4> *Ops = nullptr) {
  Register DefReg;
  unsigned StartOp = InlineAsm::MIOp_FirstOperand,
           AsmDescOp = InlineAsm::MIOp_FirstOperand;
  for (unsigned Idx = StartOp, MISz = MI->getNumOperands(); Idx != MISz;
       ++Idx) {
    const MachineOperand &MO = MI->getOperand(Idx);
    if (MO.isMetadata())
      continue;
    if (Idx == AsmDescOp && MO.isImm()) {
      // compute the index of the next operand descriptor
      const InlineAsm::Flag F(MO.getImm());
      AsmDescOp += 1 + F.getNumOperandRegisters();
      continue;
    }
    if (MO.isReg() && MO.isDef()) {
      if (!Ops)
        return MO.getReg();
      DefReg = MO.getReg();
    } else if (Ops) {
      Ops->push_back(Idx);
    }
  }
  return DefReg;
}
 
static void
insertInlineAsmProcess(MachineFunction &MF, SPIRVGlobalRegistry *GR,
                       const SPIRVSubtarget &ST, MachineIRBuilder MIRBuilder,
                       const SmallVector<MachineInstr *> &ToProcess) {
  MachineRegisterInfo &MRI = MF.getRegInfo();
  Register AsmTargetReg;
  for (unsigned i = 0, Sz = ToProcess.size(); i + 1 < Sz; i += 2) {
    MachineInstr *I1 = ToProcess[i], *I2 = ToProcess[i + 1];
    assert(isSpvIntrinsic(*I1, Intrinsic::spv_inline_asm) && I2->isInlineAsm());
    MIRBuilder.setInsertPt(*I2->getParent(), *I2);
 
    if (!AsmTargetReg.isValid()) {
      // define vendor specific assembly target or dialect
      AsmTargetReg = MRI.createGenericVirtualRegister(LLT::scalar(32));
      MRI.setRegClass(AsmTargetReg, &SPIRV::iIDRegClass);
      auto AsmTargetMIB =
          MIRBuilder.buildInstr(SPIRV::OpAsmTargetINTEL).addDef(AsmTargetReg);
      addStringImm(ST.getTargetTripleAsStr(), AsmTargetMIB);
      GR->add(AsmTargetMIB.getInstr(), AsmTargetMIB);
    }
 
    // create types
    const MDNode *IAMD = I1->getOperand(1).getMetadata();
    FunctionType *FTy = cast<FunctionType>(getMDOperandAsType(IAMD, 0));
    SmallVector<SPIRVTypeInst, 4> ArgTypes;
    for (const auto &ArgTy : FTy->params())
      ArgTypes.push_back(GR->getOrCreateSPIRVType(
          ArgTy, MIRBuilder, SPIRV::AccessQualifier::ReadWrite, true));
    SPIRVTypeInst RetType =
        GR->getOrCreateSPIRVType(FTy->getReturnType(), MIRBuilder,
                                 SPIRV::AccessQualifier::ReadWrite, true);
    SPIRVTypeInst FuncType = GR->getOrCreateOpTypeFunctionWithArgs(
        FTy, RetType, ArgTypes, MIRBuilder);
 
    // define vendor specific assembly instructions string
    Register AsmReg = MRI.createGenericVirtualRegister(LLT::scalar(32));
    MRI.setRegClass(AsmReg, &SPIRV::iIDRegClass);
    auto AsmMIB = MIRBuilder.buildInstr(SPIRV::OpAsmINTEL)
                      .addDef(AsmReg)
                      .addUse(GR->getSPIRVTypeID(RetType))
                      .addUse(GR->getSPIRVTypeID(FuncType))
                      .addUse(AsmTargetReg);
    // inline asm string:
    addStringImm(I2->getOperand(InlineAsm::MIOp_AsmString).getSymbolName(),
                 AsmMIB);
    // inline asm constraint string:
    addStringImm(cast<MDString>(I1->getOperand(2).getMetadata()->getOperand(0))
                     ->getString(),
                 AsmMIB);
    GR->add(AsmMIB.getInstr(), AsmMIB);
 
    // calls the inline assembly instruction
    unsigned ExtraInfo = I2->getOperand(InlineAsm::MIOp_ExtraInfo).getImm();
    if (ExtraInfo & InlineAsm::Extra_HasSideEffects)
      MIRBuilder.buildInstr(SPIRV::OpDecorate)
          .addUse(AsmReg)
          .addImm(static_cast<uint32_t>(SPIRV::Decoration::SideEffectsINTEL));
 
    Register DefReg = collectInlineAsmInstrOperands(I2);
    if (!DefReg.isValid()) {
      DefReg = MRI.createGenericVirtualRegister(LLT::scalar(32));
      MRI.setRegClass(DefReg, &SPIRV::iIDRegClass);
      SPIRVTypeInst VoidType = GR->getOrCreateSPIRVType(
          Type::getVoidTy(MF.getFunction().getContext()), MIRBuilder,
          SPIRV::AccessQualifier::ReadWrite, true);
      GR->assignSPIRVTypeToVReg(VoidType, DefReg, MF);
    }
 
    auto AsmCall = MIRBuilder.buildInstr(SPIRV::OpAsmCallINTEL)
                       .addDef(DefReg)
                       .addUse(GR->getSPIRVTypeID(RetType))
                       .addUse(AsmReg);
    for (unsigned IntrIdx = 3; IntrIdx < I1->getNumOperands(); ++IntrIdx)
      AsmCall.addUse(I1->getOperand(IntrIdx).getReg());
 
    // IRTranslator gets a bit confused when lowering inline ASM with outputs
    // and inserts a spurious COPY & TRUNC as registers are assumed to be i64;
    // we have to clean that up here to prevent erroneous trunc casts either on
    // a struct (for multiple outputs) or same width integers to get lowered
    // into SPIR-V
    if (MRI.hasOneUse(DefReg)) {
      MachineInstr &CopyMI = *MRI.use_instr_begin(DefReg);
      if (CopyMI.getOpcode() == TargetOpcode::COPY) {
        Register CopyDst = CopyMI.getOperand(0).getReg();
        if (MRI.hasOneUse(CopyDst)) {
          MachineInstr &TruncMI = *MRI.use_instr_begin(CopyDst);
          if (TruncMI.getOpcode() == TargetOpcode::G_TRUNC) {
            MRI.setType(DefReg, GR->getRegType(RetType));
            Register TruncReg = TruncMI.defs().begin()->getReg();
            MRI.replaceRegWith(TruncReg, DefReg);
            invalidateAndEraseMI(GR, &TruncMI);
            invalidateAndEraseMI(GR, &CopyMI);
          }
        }
      }
    }
  }
  for (MachineInstr *MI : ToProcess)
    invalidateAndEraseMI(GR, MI);
}
 
static void insertInlineAsm(MachineFunction &MF, SPIRVGlobalRegistry *GR,
                            const SPIRVSubtarget &ST,
                            MachineIRBuilder MIRBuilder) {
  SmallVector<MachineInstr *> ToProcess;
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      if (isSpvIntrinsic(MI, Intrinsic::spv_inline_asm) ||
          MI.getOpcode() == TargetOpcode::INLINEASM)
        ToProcess.push_back(&MI);
    }
  }
  if (ToProcess.size() == 0)
    return;
 
  if (!ST.canUseExtension(SPIRV::Extension::SPV_INTEL_inline_assembly))
    report_fatal_error("Inline assembly instructions require the "
                       "following SPIR-V extension: SPV_INTEL_inline_assembly",
                       false);
 
  insertInlineAsmProcess(MF, GR, ST, MIRBuilder, ToProcess);
}
 
static void insertSpirvDecorations(MachineFunction &MF, SPIRVGlobalRegistry *GR,
                                   MachineIRBuilder MIB) {
  const SPIRVSubtarget &ST = cast<SPIRVSubtarget>(MIB.getMF().getSubtarget());
  SmallVector<MachineInstr *, 10> ToErase;
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      if (!isSpvIntrinsic(MI, Intrinsic::spv_assign_decoration) &&
          !isSpvIntrinsic(MI, Intrinsic::spv_assign_aliasing_decoration) &&
          !isSpvIntrinsic(MI, Intrinsic::spv_assign_fpmaxerror_decoration))
        continue;
      MIB.setInsertPt(*MI.getParent(), MI.getNextNode());
      if (isSpvIntrinsic(MI, Intrinsic::spv_assign_decoration)) {
        buildOpSpirvDecorations(MI.getOperand(1).getReg(), MIB,
                                MI.getOperand(2).getMetadata(), ST);
      } else if (isSpvIntrinsic(MI,
                                Intrinsic::spv_assign_fpmaxerror_decoration)) {
        ConstantFP *OpV = mdconst::dyn_extract<ConstantFP>(
            MI.getOperand(2).getMetadata()->getOperand(0));
        uint32_t OpValue = OpV->getValueAPF().bitcastToAPInt().getZExtValue();
 
        buildOpDecorate(MI.getOperand(1).getReg(), MIB,
                        SPIRV::Decoration::FPMaxErrorDecorationINTEL,
                        {OpValue});
      } else {
        GR->buildMemAliasingOpDecorate(MI.getOperand(1).getReg(), MIB,
                                       MI.getOperand(2).getImm(),
                                       MI.getOperand(3).getMetadata());
      }
 
      ToErase.push_back(&MI);
    }
  }
  for (MachineInstr *MI : ToErase)
    invalidateAndEraseMI(GR, MI);
}
 
// Returns the value of the switch case operand in Reg. The case value stays a
// G_CONSTANT until the module emits a SPIR-V constant for the same value, at
// which point the case register is replaced with the one defining that
// constant, which keeps its value in literal operands rather than in a CImm.
static const ConstantInt *getSwitchCaseValue(Register Reg,
                                             const MachineRegisterInfo &MRI,
                                             LLVMContext &Ctx) {
  APInt Val;
  if (mi_match(Reg, MRI, m_ICst(Val)))
    return ConstantInt::get(Ctx, Val);
 
  const MachineInstr *Def = nullptr;
  if (!mi_match(Reg, MRI, m_MInstr(Def)))
    llvm_unreachable("Switch case operand has no definition");
 
  LLT Ty = MRI.getType(Reg);
  assert(Ty.isValid() && "Expected a typed switch case value");
  Val = APInt(Ty.getScalarSizeInBits(), 0);
 
  switch (Def->getOpcode()) {
  case SPIRV::OpConstantNull:
  case SPIRV::OpConstantI:
    // The operands after the type are 32-bit literal words, least significant
    // first, as written by addNumImm(). OpConstantNull carries none, so it
    // decodes to zero without a case of its own.
    for (unsigned I = 2, E = Def->getNumExplicitOperands(); I != E; ++I) {
      uint32_t Word = static_cast<uint32_t>(Def->getOperand(I).getImm());
      Val |= APInt(Val.getBitWidth(), Word).shl((I - 2) * 32);
    }
    break;
  default:
    llvm_unreachable("Unexpected definition of a switch case value");
  }
  return ConstantInt::get(Ctx, Val);
}
 
// LLVM allows the switches to use registers as cases, while SPIR-V required
// those to be immediate values. This function replaces such operands with the
// equivalent immediate constant.
static void processSwitchesConstants(MachineFunction &MF,
                                     SPIRVGlobalRegistry *GR,
                                     MachineIRBuilder MIB) {
  MachineRegisterInfo &MRI = MF.getRegInfo();
  LLVMContext &Ctx = MF.getFunction().getContext();
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      if (!isSpvIntrinsic(MI, Intrinsic::spv_switch))
        continue;
 
      SmallVector<MachineOperand, 8> NewOperands;
      NewOperands.push_back(MI.getOperand(0)); // Opcode
      NewOperands.push_back(MI.getOperand(1)); // Condition
      NewOperands.push_back(MI.getOperand(2)); // Default
      for (unsigned i = 3; i < MI.getNumOperands(); i += 2) {
        Register Reg = MI.getOperand(i).getReg();
        NewOperands.push_back(
            MachineOperand::CreateCImm(getSwitchCaseValue(Reg, MRI, Ctx)));
 
        NewOperands.push_back(MI.getOperand(i + 1));
      }
 
      assert(MI.getNumOperands() == NewOperands.size());
      while (MI.getNumOperands() > 0)
        MI.removeOperand(0);
      for (auto &MO : NewOperands)
        MI.addOperand(MO);
    }
  }
}
 
// Some instructions are used during CodeGen but should never be emitted.
// Cleaning up those.
static void cleanupHelperInstructions(MachineFunction &MF,
                                      SPIRVGlobalRegistry *GR) {
  SmallVector<MachineInstr *, 8> ToEraseMI;
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      if (isSpvIntrinsic(MI, Intrinsic::spv_track_constant) ||
          MI.getOpcode() == TargetOpcode::G_BRINDIRECT)
        ToEraseMI.push_back(&MI);
    }
  }
 
  for (MachineInstr *MI : ToEraseMI)
    invalidateAndEraseMI(GR, MI);
}
 
// Find all usages of G_BLOCK_ADDR in our intrinsics and replace those
// operands/registers by the actual MBB it references.
static void processBlockAddr(MachineFunction &MF, SPIRVGlobalRegistry *GR,
                             MachineIRBuilder MIB) {
  // Gather the reverse-mapping BB -> MBB.
  DenseMap<const BasicBlock *, MachineBasicBlock *> BB2MBB;
  for (MachineBasicBlock &MBB : MF)
    BB2MBB[MBB.getBasicBlock()] = &MBB;
 
  // Gather instructions requiring patching. For now, only those can use
  // G_BLOCK_ADDR.
  SmallVector<MachineInstr *, 8> InstructionsToPatch;
  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : MBB) {
      if (isSpvIntrinsic(MI, Intrinsic::spv_switch) ||
          isSpvIntrinsic(MI, Intrinsic::spv_loop_merge) ||
          isSpvIntrinsic(MI, Intrinsic::spv_selection_merge))
        InstructionsToPatch.push_back(&MI);
    }
  }
 
  // For each instruction to fix, we replace all the G_BLOCK_ADDR operands by
  // the actual MBB it references. Once those references have been updated, we
  // can cleanup remaining G_BLOCK_ADDR references.
  SmallPtrSet<MachineBasicBlock *, 8> ClearAddressTaken;
  SmallPtrSet<MachineInstr *, 8> ToEraseMI;
  MachineRegisterInfo &MRI = MF.getRegInfo();
  for (MachineInstr *MI : InstructionsToPatch) {
    SmallVector<MachineOperand, 8> NewOps;
    for (unsigned i = 0; i < MI->getNumOperands(); ++i) {
      // The operand is not a register, keep as-is.
      if (!MI->getOperand(i).isReg()) {
        NewOps.push_back(MI->getOperand(i));
        continue;
      }
 
      Register Reg = MI->getOperand(i).getReg();
      MachineInstr *BuildMBB = MRI.getVRegDef(Reg);
      // The register is not the result of G_BLOCK_ADDR, keep as-is.
      if (!BuildMBB || BuildMBB->getOpcode() != TargetOpcode::G_BLOCK_ADDR) {
        NewOps.push_back(MI->getOperand(i));
        continue;
      }
 
      assert(BuildMBB && BuildMBB->getOpcode() == TargetOpcode::G_BLOCK_ADDR &&
             BuildMBB->getOperand(1).isBlockAddress() &&
             BuildMBB->getOperand(1).getBlockAddress());
      BasicBlock *BB =
          BuildMBB->getOperand(1).getBlockAddress()->getBasicBlock();
      auto It = BB2MBB.find(BB);
      if (It == BB2MBB.end())
        report_fatal_error("cannot find a machine basic block by a basic block "
                           "in a switch statement");
      MachineBasicBlock *ReferencedBlock = It->second;
      NewOps.push_back(MachineOperand::CreateMBB(ReferencedBlock));
 
      ClearAddressTaken.insert(ReferencedBlock);
      ToEraseMI.insert(BuildMBB);
    }
 
    // Replace the operands.
    assert(MI->getNumOperands() == NewOps.size());
    while (MI->getNumOperands() > 0)
      MI->removeOperand(0);
    for (auto &MO : NewOps)
      MI->addOperand(MO);
 
    if (MachineInstr *Next = MI->getNextNode()) {
      if (isSpvIntrinsic(*Next, Intrinsic::spv_track_constant)) {
        ToEraseMI.insert(Next);
        Next = MI->getNextNode();
      }
      if (Next && Next->getOpcode() == TargetOpcode::G_BRINDIRECT)
        ToEraseMI.insert(Next);
    }
  }
 
  // BlockAddress operands were used to keep information between passes,
  // let's undo the "address taken" status to reflect that Succ doesn't
  // actually correspond to an IR-level basic block.
  for (MachineBasicBlock *Succ : ClearAddressTaken)
    Succ->setAddressTakenIRBlock(nullptr);
 
  // If we just delete G_BLOCK_ADDR instructions with BlockAddress operands,
  // this leaves their BasicBlock counterparts in a "address taken" status. This
  // would make AsmPrinter to generate a series of unneeded labels of a "Address
  // of block that was removed by CodeGen" kind. Let's first ensure that we
  // don't have a dangling BlockAddress constants by zapping the BlockAddress
  // nodes, and only after that proceed with erasing G_BLOCK_ADDR instructions.
  Constant *Replacement =
      ConstantInt::get(Type::getInt32Ty(MF.getFunction().getContext()), 1);
  for (MachineInstr *BlockAddrI : ToEraseMI) {
    if (BlockAddrI->getOpcode() == TargetOpcode::G_BLOCK_ADDR) {
      BlockAddress *BA = const_cast<BlockAddress *>(
          BlockAddrI->getOperand(1).getBlockAddress());
      BA->replaceAllUsesWith(
          ConstantExpr::getIntToPtr(Replacement, BA->getType()));
      BA->destroyConstant();
    }
    invalidateAndEraseMI(GR, BlockAddrI);
  }
}
 
static bool isImplicitFallthrough(MachineBasicBlock &MBB) {
  if (MBB.empty())
    return MBB.getNextNode() != nullptr;
 
  // Branching SPIR-V intrinsics are not detected by this generic method.
  // Thus, we can only trust negative result.
  if (!MBB.canFallThrough())
    return false;
 
  // Otherwise, we must manually check if we have a SPIR-V intrinsic which
  // prevent an implicit fallthrough.
  for (MachineBasicBlock::reverse_iterator It = MBB.rbegin(), E = MBB.rend();
       It != E; ++It) {
    if (isSpvIntrinsic(*It, Intrinsic::spv_switch))
      return false;
  }
  return true;
}
 
static void removeImplicitFallthroughs(MachineFunction &MF,
                                       MachineIRBuilder MIB) {
  // It is valid for MachineBasicBlocks to not finish with a branch instruction.
  // In such cases, they will simply fallthrough their immediate successor.
  for (MachineBasicBlock &MBB : MF) {
    if (!isImplicitFallthrough(MBB))
      continue;
 
    assert(MBB.succ_size() == 1);
    MIB.setInsertPt(MBB, MBB.end());
    MIB.buildBr(**MBB.successors().begin());
  }
}
 
static bool runPreLegalizer(MachineFunction &MF) {
  // Initialize the type registry.
  const SPIRVSubtarget &ST = MF.getSubtarget<SPIRVSubtarget>();
  SPIRVGlobalRegistry *GR = ST.getSPIRVGlobalRegistry();
  GR->setCurrentFunc(MF);
  MachineIRBuilder MIB(MF);
  // a registry of target extension constants
  DenseMap<MachineInstr *, Type *> TargetExtConstTypes;
  // to keep record of tracked constants
  addConstantsToTrack(MF, GR, ST, TargetExtConstTypes);
  foldConstantsIntoIntrinsics(MF, GR, MIB);
  insertBitcasts(MF, GR, MIB);
  generateAssignInstrs(MF, GR, MIB, TargetExtConstTypes);
 
  processSwitchesConstants(MF, GR, MIB);
  processBlockAddr(MF, GR, MIB);
  cleanupHelperInstructions(MF, GR);
 
  processInstrsWithTypeFolding(MF, GR, MIB);
  removeImplicitFallthroughs(MF, MIB);
  insertSpirvDecorations(MF, GR, MIB);
  insertInlineAsm(MF, GR, ST, MIB);
  lowerBitcasts(MF, GR, MIB);
 
  return true;
}
 
INITIALIZE_PASS(SPIRVPreLegalizerLegacy, DEBUG_TYPE, "SPIRV pre legalizer",
                false, false)
 
char SPIRVPreLegalizerLegacy::ID = 0;
 
FunctionPass *llvm::createSPIRVPreLegalizerLegacyPass() {
  return new SPIRVPreLegalizerLegacy();
}
 
bool SPIRVPreLegalizerLegacy::runOnMachineFunction(MachineFunction &MF) {
  return runPreLegalizer(MF);
}
 
PreservedAnalyses
SPIRVPreLegalizerPass::run(MachineFunction &MF,
                           MachineFunctionAnalysisManager &MFAM) {
  bool Changed = runPreLegalizer(MF);
  if (!Changed)
    return PreservedAnalyses::all();
 
  return getMachineFunctionPassPreservedAnalyses()
      .preserve<GISelValueTrackingAnalysis>();
}