llvm/lib/Target/RISCV/RISCVVectorPeephole.cpp

//===- RISCVVectorPeephole.cpp - MI Vector Pseudo Peepholes ---------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This pass performs various vector pseudo peephole optimisations after
// instruction selection.
//
// Currently it converts vmerge.vvm to vmv.v.v
// PseudoVMERGE_VVM %false, %false, %true, %allonesmask, %vl, %sew
// ->
// PseudoVMV_V_V %false, %true, %vl, %sew
//
// And masked pseudos to unmasked pseudos
// PseudoVADD_V_V_MASK %passthru, %a, %b, %allonesmask, %vl, sew, policy
// ->
// PseudoVADD_V_V %passthru %a, %b, %vl, sew, policy
//
// It also converts AVLs to VLMAX where possible
// %vl = VLENB * something
// PseudoVADD_V_V %passthru, %a, %b, %vl, sew, policy
// ->
// PseudoVADD_V_V %passthru, %a, %b, -1, sew, policy
//
//===----------------------------------------------------------------------===//

#include "RISCV.h"
#include "RISCVSubtarget.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/RegisterClassInfo.h"
#include "llvm/CodeGen/TargetInstrInfo.h"
#include "llvm/CodeGen/TargetRegisterInfo.h"

using namespace llvm;

#define DEBUG_TYPE "riscv-vector-peephole"

namespace {

class RISCVVectorPeepholeImpl {
public:
  bool run(MachineFunction &MF);

private:
  const TargetInstrInfo *TII;
  MachineRegisterInfo *MRI;
  const TargetRegisterInfo *TRI;
  const RISCVSubtarget *ST;
  bool convertToVLMAX(MachineInstr &MI) const;
  bool convertToWholeRegister(MachineInstr &MI) const;
  bool convertToUnmasked(MachineInstr &MI) const;
  bool convertAllOnesVMergeToVMv(MachineInstr &MI) const;
  bool convertSameMaskVMergeToVMv(MachineInstr &MI);
  bool foldUndefPassthruVMV_V_V(MachineInstr &MI);
  bool foldVMV_V_V(MachineInstr &MI);
  bool foldVMergeToMask(MachineInstr &MI) const;
  bool foldVMANDToMaskedCompare(MachineInstr &MI) const;

  bool hasSameEEW(const MachineInstr &User, const MachineInstr &Src) const;
  bool isAllOnesMask(const MachineInstr *MaskDef) const;
  std::optional<unsigned> getConstant(const MachineOperand &VL) const;
  bool ensureDominates(ArrayRef<const MachineOperand *> Defs,
                       MachineInstr &Use) const;
  Register
  lookThruCopies(Register Reg, bool OneUseOnly = false,
                 SmallVectorImpl<MachineInstr *> *Copies = nullptr) const;
};

class RISCVVectorPeepholeLegacy : public MachineFunctionPass {
public:
  static char ID;

  RISCVVectorPeepholeLegacy() : MachineFunctionPass(ID) {}

  bool runOnMachineFunction(MachineFunction &MF) override;
  MachineFunctionProperties getRequiredProperties() const override {
    return MachineFunctionProperties().setIsSSA();
  }

  StringRef getPassName() const override {
    return "RISC-V Vector Peephole Optimization";
  }

  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.setPreservesCFG();
    AU.addPreserved<MachineRegisterClassInfoWrapperPass>();
    MachineFunctionPass::getAnalysisUsage(AU);
  }
};

} // namespace

char RISCVVectorPeepholeLegacy::ID = 0;

INITIALIZE_PASS(RISCVVectorPeepholeLegacy, DEBUG_TYPE, "RISC-V Fold Masks",
                false, false)

/// Given \p User that has an input operand with EEW=SEW, which uses the dest
/// operand of \p Src with an unknown EEW, return true if their EEWs match.
bool RISCVVectorPeepholeImpl::hasSameEEW(const MachineInstr &User,
                                         const MachineInstr &Src) const {
  unsigned UserLog2SEW =
      User.getOperand(RISCVII::getSEWOpNum(User.getDesc())).getImm();
  unsigned SrcLog2SEW =
      Src.getOperand(RISCVII::getSEWOpNum(Src.getDesc())).getImm();
  unsigned SrcLog2EEW = RISCV::getDestLog2EEW(
      TII->get(RISCV::getRVVMCOpcode(Src.getOpcode())), SrcLog2SEW);
  return SrcLog2EEW == UserLog2SEW;
}

/// Check if an operand is an immediate or a materialized ADDI $x0, imm.
std::optional<unsigned>
RISCVVectorPeepholeImpl::getConstant(const MachineOperand &VL) const {
  if (VL.isImm())
    return VL.getImm();

  if (!VL.getReg().isVirtual())
    return std::nullopt;
  MachineInstr *Def = MRI->getVRegDef(VL.getReg());
  if (!Def || Def->getOpcode() != RISCV::ADDI || !Def->getOperand(1).isReg() ||
      Def->getOperand(1).getReg() != RISCV::X0)
    return std::nullopt;
  return Def->getOperand(2).getImm();
}

/// Convert AVLs that are known to be VLMAX to the VLMAX sentinel.
bool RISCVVectorPeepholeImpl::convertToVLMAX(MachineInstr &MI) const {
  if (!RISCVII::hasVLOp(MI.getDesc().TSFlags) ||
      !RISCVII::hasSEWOp(MI.getDesc().TSFlags))
    return false;

  auto LMUL = RISCVVType::decodeVLMUL(RISCVII::getLMul(MI.getDesc().TSFlags));
  // Fixed-point value, denominator=8
  unsigned LMULFixed = LMUL.second ? (8 / LMUL.first) : 8 * LMUL.first;
  unsigned Log2SEW = MI.getOperand(RISCVII::getSEWOpNum(MI.getDesc())).getImm();
  // A Log2SEW of 0 is an operation on mask registers only
  unsigned SEW = Log2SEW ? 1 << Log2SEW : 8;
  assert(RISCVVType::isValidSEW(SEW) && "Unexpected SEW");
  assert(8 * LMULFixed / SEW > 0);

  // If the exact VLEN is known then we know VLMAX, check if the AVL == VLMAX.
  MachineOperand &VL = MI.getOperand(RISCVII::getVLOpNum(MI.getDesc()));
  if (auto VLen = ST->getRealVLen(), AVL = getConstant(VL);
      VLen && AVL && (*VLen * LMULFixed) / SEW == *AVL * 8) {
    VL.ChangeToImmediate(RISCV::VLMaxSentinel);
    return true;
  }

  // If an AVL is a VLENB that's possibly scaled to be equal to VLMAX, convert
  // it to the VLMAX sentinel value.
  if (!VL.isReg())
    return false;
  MachineInstr *Def = MRI->getVRegDef(VL.getReg());
  if (!Def)
    return false;

  // Fixed-point value, denominator=8
  uint64_t ScaleFixed = 8;
  // Check if the VLENB was potentially scaled with slli/srli
  if (Def->getOpcode() == RISCV::SLLI) {
    assert(Def->getOperand(2).getImm() < 64);
    ScaleFixed <<= Def->getOperand(2).getImm();
    Def = MRI->getVRegDef(Def->getOperand(1).getReg());
  } else if (Def->getOpcode() == RISCV::SRLI) {
    assert(Def->getOperand(2).getImm() < 64);
    ScaleFixed >>= Def->getOperand(2).getImm();
    Def = MRI->getVRegDef(Def->getOperand(1).getReg());
  }

  if (!Def || Def->getOpcode() != RISCV::PseudoReadVLENB)
    return false;

  // AVL = (VLENB * Scale)
  //
  // VLMAX = (VLENB * 8 * LMUL) / SEW
  //
  // AVL == VLMAX
  // -> VLENB * Scale == (VLENB * 8 * LMUL) / SEW
  // -> Scale == (8 * LMUL) / SEW
  if (ScaleFixed != 8 * LMULFixed / SEW)
    return false;

  VL.ChangeToImmediate(RISCV::VLMaxSentinel);

  return true;
}

bool RISCVVectorPeepholeImpl::isAllOnesMask(const MachineInstr *MaskDef) const {
  while (MaskDef->isCopy() && MaskDef->getOperand(1).getReg().isVirtual())
    MaskDef = MRI->getVRegDef(MaskDef->getOperand(1).getReg());

  // TODO: Check that the VMSET is the expected bitwidth? The pseudo has
  // undefined behaviour if it's the wrong bitwidth, so we could choose to
  // assume that it's all-ones? Same applies to its VL.
  switch (MaskDef->getOpcode()) {
  case RISCV::PseudoVMSET_M_B1:
  case RISCV::PseudoVMSET_M_B2:
  case RISCV::PseudoVMSET_M_B4:
  case RISCV::PseudoVMSET_M_B8:
  case RISCV::PseudoVMSET_M_B16:
  case RISCV::PseudoVMSET_M_B32:
  case RISCV::PseudoVMSET_M_B64:
    return true;
  default:
    return false;
  }
}

/// Convert unit strided unmasked loads and stores to whole-register equivalents
/// to avoid the dependency on $vl and $vtype.
///
/// %x = PseudoVLE8_V_M1 %passthru, %ptr, %vlmax, policy
/// PseudoVSE8_V_M1 %v, %ptr, %vlmax
///
/// ->
///
/// %x = VL1RE8_V %ptr
/// VS1R_V %v, %ptr
bool RISCVVectorPeepholeImpl::convertToWholeRegister(MachineInstr &MI) const {
#define CASE_WHOLE_REGISTER_LMUL_SEW(lmul, sew)                                \
  case RISCV::PseudoVLE##sew##_V_M##lmul:                                      \
    NewOpc = RISCV::VL##lmul##RE##sew##_V;                                     \
    break;                                                                     \
  case RISCV::PseudoVSE##sew##_V_M##lmul:                                      \
    NewOpc = RISCV::VS##lmul##R_V;                                             \
    break;
#define CASE_WHOLE_REGISTER_LMUL(lmul)                                         \
  CASE_WHOLE_REGISTER_LMUL_SEW(lmul, 8)                                        \
  CASE_WHOLE_REGISTER_LMUL_SEW(lmul, 16)                                       \
  CASE_WHOLE_REGISTER_LMUL_SEW(lmul, 32)                                       \
  CASE_WHOLE_REGISTER_LMUL_SEW(lmul, 64)

  unsigned NewOpc;
  switch (MI.getOpcode()) {
    CASE_WHOLE_REGISTER_LMUL(1)
    CASE_WHOLE_REGISTER_LMUL(2)
    CASE_WHOLE_REGISTER_LMUL(4)
    CASE_WHOLE_REGISTER_LMUL(8)
  default:
    return false;
  }

  MachineOperand &VLOp = MI.getOperand(RISCVII::getVLOpNum(MI.getDesc()));
  if (!VLOp.isImm() || VLOp.getImm() != RISCV::VLMaxSentinel)
    return false;

  // Whole register instructions aren't pseudos so they don't have
  // policy/SEW/AVL ops, and they don't have passthrus.
  if (RISCVII::hasVecPolicyOp(MI.getDesc().TSFlags))
    MI.removeOperand(RISCVII::getVecPolicyOpNum(MI.getDesc()));
  MI.removeOperand(RISCVII::getSEWOpNum(MI.getDesc()));
  MI.removeOperand(RISCVII::getVLOpNum(MI.getDesc()));
  if (RISCVII::isFirstDefTiedToFirstUse(MI.getDesc()))
    MI.removeOperand(1);

  MI.setDesc(TII->get(NewOpc));

  return true;
}

static unsigned getVMV_V_VOpcodeForVMERGE_VVM(const MachineInstr &MI) {
#define CASE_VMERGE_TO_VMV(lmul)                                               \
  case RISCV::PseudoVMERGE_VVM_##lmul:                                         \
    return RISCV::PseudoVMV_V_V_##lmul;
  switch (MI.getOpcode()) {
  default:
    return 0;
    CASE_VMERGE_TO_VMV(MF8)
    CASE_VMERGE_TO_VMV(MF4)
    CASE_VMERGE_TO_VMV(MF2)
    CASE_VMERGE_TO_VMV(M1)
    CASE_VMERGE_TO_VMV(M2)
    CASE_VMERGE_TO_VMV(M4)
    CASE_VMERGE_TO_VMV(M8)
  }
}

/// Convert a PseudoVMERGE_VVM with an all ones mask to a PseudoVMV_V_V.
///
/// %x = PseudoVMERGE_VVM %passthru, %false, %true, %allones, sew, vl
/// ->
/// %x = PseudoVMV_V_V %passthru, %true, vl, sew, tu_mu
bool RISCVVectorPeepholeImpl::convertAllOnesVMergeToVMv(
    MachineInstr &MI) const {
  unsigned NewOpc = getVMV_V_VOpcodeForVMERGE_VVM(MI);
  if (!NewOpc)
    return false;
  if (!isAllOnesMask(MRI->getVRegDef(MI.getOperand(4).getReg())))
    return false;

  MI.setDesc(TII->get(NewOpc));
  MI.removeOperand(2); // False operand
  MI.removeOperand(3); // Mask operand
  MI.addOperand(
      MachineOperand::CreateImm(RISCVVType::TAIL_UNDISTURBED_MASK_UNDISTURBED));

  // vmv.v.v doesn't have a mask operand, so we may be able to inflate the
  // register class for the destination and passthru operands e.g. VRNoV0 -> VR
  MRI->recomputeRegClass(MI.getOperand(0).getReg());
  if (MI.getOperand(1).getReg().isValid())
    MRI->recomputeRegClass(MI.getOperand(1).getReg());
  return true;
}

// If \p Reg is defined by one or more COPYs of virtual registers, traverses
// the chain and returns the root non-COPY source.
Register RISCVVectorPeepholeImpl::lookThruCopies(
    Register Reg, bool OneUseOnly,
    SmallVectorImpl<MachineInstr *> *Copies) const {
  while (MachineInstr *Def = MRI->getUniqueVRegDef(Reg)) {
    if (!Def->isFullCopy())
      break;
    Register Src = Def->getOperand(1).getReg();
    if (!Src.isVirtual())
      break;
    if (OneUseOnly && !MRI->hasOneNonDBGUse(Reg))
      break;
    if (Copies)
      Copies->push_back(Def);
    Reg = Src;
  }
  return Reg;
}

/// If a PseudoVMERGE_VVM's true operand is a masked pseudo and both have the
/// same mask, and the masked pseudo's passthru is the same as the false
/// operand, we can convert the PseudoVMERGE_VVM to a PseudoVMV_V_V.
///
/// %true = PseudoVADD_VV_M1_MASK %false, %x, %y, %mask, vl1, sew, policy
/// %x = PseudoVMERGE_VVM %passthru, %false, %true, %mask, vl2, sew
/// ->
/// %true = PseudoVADD_VV_M1_MASK %false, %x, %y, %mask, vl1, sew, policy
/// %x = PseudoVMV_V_V %passthru, %true, vl2, sew, tu_mu
bool RISCVVectorPeepholeImpl::convertSameMaskVMergeToVMv(MachineInstr &MI) {
  unsigned NewOpc = getVMV_V_VOpcodeForVMERGE_VVM(MI);
  if (!NewOpc)
    return false;
  MachineInstr *True = MRI->getVRegDef(MI.getOperand(3).getReg());

  if (!True || True->getParent() != MI.getParent())
    return false;

  auto *TrueMaskedInfo = RISCV::getMaskedPseudoInfo(True->getOpcode());
  if (!TrueMaskedInfo || !hasSameEEW(MI, *True))
    return false;

  Register TrueMaskReg = lookThruCopies(
      True->getOperand(TrueMaskedInfo->MaskOpIdx + True->getNumExplicitDefs())
          .getReg());
  Register MIMaskReg = lookThruCopies(MI.getOperand(4).getReg());
  if (!TrueMaskReg.isVirtual() || TrueMaskReg != MIMaskReg)
    return false;

  // Masked off lanes past TrueVL will come from False, and converting to vmv
  // will lose these lanes unless MIVL <= TrueVL.
  // We can relax this when False == Passthru and True's tail policy is TU,
  // because True's tail lanes will preserve its passthru (= False = Passthru).
  const MachineOperand &MIVL = MI.getOperand(RISCVII::getVLOpNum(MI.getDesc()));
  const MachineOperand &TrueVL =
      True->getOperand(RISCVII::getVLOpNum(True->getDesc()));
  Register FalseReg = MI.getOperand(2).getReg();
  if (!RISCV::isVLKnownLE(*MRI, MIVL, TrueVL)) {
    Register PassthruReg = MI.getOperand(1).getReg();
    if (FalseReg.isValid() && FalseReg != PassthruReg)
      return false;
    if (!RISCVII::hasVecPolicyOp(True->getDesc().TSFlags))
      return false;
    uint64_t TruePolicy =
        True->getOperand(RISCVII::getVecPolicyOpNum(True->getDesc())).getImm();
    if (TruePolicy & RISCVVType::TAIL_AGNOSTIC)
      return false;
  }

  // True's passthru needs to be equivalent to False
  Register TruePassthruReg = True->getOperand(1).getReg();
  if (TruePassthruReg != FalseReg) {
    // If True's passthru is undef see if we can change it to False
    if (TruePassthruReg.isValid() ||
        !MRI->hasOneUse(MI.getOperand(3).getReg()) ||
        !ensureDominates(&MI.getOperand(2), *True))
      return false;
    True->getOperand(1).setReg(MI.getOperand(2).getReg());
    // If True is masked then its passthru needs to be in VRNoV0.
    MRI->constrainRegClass(True->getOperand(1).getReg(),
                           TII->getRegClass(True->getDesc(), 1));
  }

  // If True is mask agnostic, we need to make it mask undisturbed.
  if (RISCVII::hasVecPolicyOp(True->getDesc().TSFlags)) {
    MachineOperand &PolicyOp =
        True->getOperand(RISCVII::getVecPolicyOpNum(True->getDesc()));
    PolicyOp.setImm(PolicyOp.getImm() & ~RISCVVType::MASK_AGNOSTIC);
  }

  MI.setDesc(TII->get(NewOpc));
  MI.removeOperand(2); // False operand
  MI.removeOperand(3); // Mask operand
  MI.addOperand(
      MachineOperand::CreateImm(RISCVVType::TAIL_UNDISTURBED_MASK_UNDISTURBED));

  // vmv.v.v doesn't have a mask operand, so we may be able to inflate the
  // register class for the destination and passthru operands e.g. VRNoV0 -> VR
  MRI->recomputeRegClass(MI.getOperand(0).getReg());
  if (MI.getOperand(1).getReg().isValid())
    MRI->recomputeRegClass(MI.getOperand(1).getReg());
  return true;
}

bool RISCVVectorPeepholeImpl::convertToUnmasked(MachineInstr &MI) const {
  const RISCV::RISCVMaskedPseudoInfo *I =
      RISCV::getMaskedPseudoInfo(MI.getOpcode());
  if (!I)
    return false;

  if (!isAllOnesMask(MRI->getVRegDef(
          MI.getOperand(I->MaskOpIdx + MI.getNumExplicitDefs()).getReg())))
    return false;

  // There are two classes of pseudos in the table - compares and
  // everything else.  See the comment on RISCVMaskedPseudo for details.
  const unsigned Opc = I->UnmaskedPseudo;
  const MCInstrDesc &MCID = TII->get(Opc);
  [[maybe_unused]] const bool HasPolicyOp =
      RISCVII::hasVecPolicyOp(MCID.TSFlags);
  const bool HasPassthru = RISCVII::isFirstDefTiedToFirstUse(MCID);
  const MCInstrDesc &MaskedMCID = TII->get(MI.getOpcode());
  assert((RISCVII::hasVecPolicyOp(MaskedMCID.TSFlags) ||
          !RISCVII::hasVecPolicyOp(MCID.TSFlags)) &&
         "Unmasked pseudo has policy but masked pseudo doesn't?");
  assert(HasPolicyOp == HasPassthru && "Unexpected pseudo structure");
  assert(!(HasPassthru && !RISCVII::isFirstDefTiedToFirstUse(MaskedMCID)) &&
         "Unmasked with passthru but masked with no passthru?");
  (void)HasPolicyOp;

  MI.setDesc(MCID);

  // Drop the policy operand if unmasked doesn't need it.
  if (RISCVII::hasVecPolicyOp(MaskedMCID.TSFlags) &&
      !RISCVII::hasVecPolicyOp(MCID.TSFlags))
    MI.removeOperand(RISCVII::getVecPolicyOpNum(MaskedMCID));

  // TODO: Increment all MaskOpIdxs in tablegen by num of explicit defs?
  unsigned MaskOpIdx = I->MaskOpIdx + MI.getNumExplicitDefs();
  MI.removeOperand(MaskOpIdx);

  // The unmasked pseudo will no longer be constrained to the vrnov0 reg class,
  // so try and relax it to vr.
  MRI->recomputeRegClass(MI.getOperand(0).getReg());

  // If the original masked pseudo had a passthru, relax it or remove it.
  if (RISCVII::isFirstDefTiedToFirstUse(MaskedMCID)) {
    unsigned PassthruOpIdx = MI.getNumExplicitDefs();
    if (HasPassthru) {
      if (MI.getOperand(PassthruOpIdx).getReg())
        MRI->recomputeRegClass(MI.getOperand(PassthruOpIdx).getReg());
    } else
      MI.removeOperand(PassthruOpIdx);
  }

  return true;
}

/// Given A and B are in the same MBB, returns true if A comes before B.
static bool strictlyDominates(MachineBasicBlock::const_iterator A,
                              MachineBasicBlock::const_iterator B) {
  assert(A->getParent() == B->getParent());
  if (A == B)
    return false;
  const MachineBasicBlock *MBB = A->getParent();
  auto MBBEnd = MBB->end();
  if (B == MBBEnd)
    return true;

  MachineBasicBlock::const_iterator I = MBB->begin();
  for (; &*I != A && &*I != B; ++I)
    ;

  return &*I == A;
}

/// If a register in \p Defs doesn't dominate \p Use, try to move Use so it
/// does. Returns false if any def doesn't dominate and we can't move Use. Each
/// def must be in the same block as Use.
bool RISCVVectorPeepholeImpl::ensureDominates(
    ArrayRef<const MachineOperand *> Defs, MachineInstr &Use) const {
  MachineInstr *Dest = &Use;

  for (const MachineOperand *MO : Defs) {
    assert(MO->getParent()->getParent() == Use.getParent());
    if (!MO->isReg() || !MO->getReg().isValid())
      continue;

    MachineInstr *Def = MRI->getVRegDef(MO->getReg());
    if (Def->getParent() == Dest->getParent() &&
        !strictlyDominates(Def, *Dest)) {
      if (!RISCVInstrInfo::isSafeToMove(*Dest, *Def->getNextNode()))
        return false;
      Dest = Def->getNextNode();
    }
  }

  if (Dest != &Use)
    Use.moveBefore(Dest);

  return true;
}

/// If a PseudoVMV_V_V's passthru is undef then we can replace it with its input
bool RISCVVectorPeepholeImpl::foldUndefPassthruVMV_V_V(MachineInstr &MI) {
  if (RISCV::getRVVMCOpcode(MI.getOpcode()) != RISCV::VMV_V_V)
    return false;
  if (MI.getOperand(1).getReg().isValid())
    return false;

  // If the input was a pseudo with a policy operand, we can give it a tail
  // agnostic policy if MI's undef tail subsumes the input's.
  MachineInstr *Src = MRI->getVRegDef(MI.getOperand(2).getReg());
  if (Src && !Src->hasUnmodeledSideEffects() &&
      MRI->hasOneUse(MI.getOperand(2).getReg()) &&
      RISCVII::hasVLOp(Src->getDesc().TSFlags) &&
      RISCVII::hasVecPolicyOp(Src->getDesc().TSFlags) && hasSameEEW(MI, *Src)) {
    const MachineOperand &MIVL = MI.getOperand(3);
    const MachineOperand &SrcVL =
        Src->getOperand(RISCVII::getVLOpNum(Src->getDesc()));

    MachineOperand &SrcPolicy =
        Src->getOperand(RISCVII::getVecPolicyOpNum(Src->getDesc()));

    if (RISCV::isVLKnownLE(*MRI, MIVL, SrcVL))
      SrcPolicy.setImm(SrcPolicy.getImm() | RISCVVType::TAIL_AGNOSTIC);
  }

  MRI->constrainRegClass(MI.getOperand(2).getReg(),
                         MRI->getRegClass(MI.getOperand(0).getReg()));
  MRI->replaceRegWith(MI.getOperand(0).getReg(), MI.getOperand(2).getReg());
  MRI->clearKillFlags(MI.getOperand(2).getReg());
  MI.eraseFromParent();
  return true;
}

/// If a PseudoVMV_V_V is the only user of its input, fold its passthru and VL
/// into it.
///
/// %x = PseudoVADD_V_V_M1 %passthru, %a, %b, %vl1, sew, policy
/// %y = PseudoVMV_V_V_M1 %passthru, %x, %vl2, sew, policy
///    (where %vl1 <= %vl2)
///
/// ->
///
/// %y = PseudoVADD_V_V_M1 %passthru, %a, %b, vl1, sew, policy
bool RISCVVectorPeepholeImpl::foldVMV_V_V(MachineInstr &MI) {
  if (RISCV::getRVVMCOpcode(MI.getOpcode()) != RISCV::VMV_V_V)
    return false;

  MachineOperand &Passthru = MI.getOperand(1);

  if (!MRI->hasOneUse(MI.getOperand(2).getReg()))
    return false;

  MachineInstr *Src = MRI->getVRegDef(MI.getOperand(2).getReg());
  if (!Src || Src->hasUnmodeledSideEffects() ||
      Src->getParent() != MI.getParent() ||
      !RISCVII::isFirstDefTiedToFirstUse(Src->getDesc()) ||
      !RISCVII::hasVLOp(Src->getDesc().TSFlags))
    return false;

  // Src's dest needs to have the same EEW as MI's input.
  if (!hasSameEEW(MI, *Src))
    return false;

  std::optional<std::pair<unsigned, unsigned>> NeedsCommute;

  // Src needs to have the same passthru as VMV_V_V
  MachineOperand &SrcPassthru = Src->getOperand(Src->getNumExplicitDefs());
  if (SrcPassthru.getReg().isValid() &&
      SrcPassthru.getReg() != Passthru.getReg()) {
    // If Src's passthru != Passthru, check if it uses Passthru in another
    // operand and try to commute it.
    int OtherIdx = Src->findRegisterUseOperandIdx(Passthru.getReg(), TRI);
    if (OtherIdx == -1)
      return false;
    unsigned OpIdx1 = OtherIdx;
    unsigned OpIdx2 = Src->getNumExplicitDefs();
    if (!TII->findCommutedOpIndices(*Src, OpIdx1, OpIdx2))
      return false;
    NeedsCommute = {OpIdx1, OpIdx2};
  }

  // Src VL will have already been reduced if legal by RISCVVLOptimizer,
  // so we don't need to handle a smaller source VL here.  However, the
  // user's VL may be larger
  MachineOperand &SrcVL = Src->getOperand(RISCVII::getVLOpNum(Src->getDesc()));
  if (!RISCV::isVLKnownLE(*MRI, SrcVL, MI.getOperand(3)))
    return false;

  // If the new passthru doesn't dominate Src, try to move Src so it does.
  if (!ensureDominates(&Passthru, *Src))
    return false;

  if (NeedsCommute) {
    auto [OpIdx1, OpIdx2] = *NeedsCommute;
    [[maybe_unused]] bool Commuted =
        TII->commuteInstruction(*Src, /*NewMI=*/false, OpIdx1, OpIdx2);
    assert(Commuted && "Failed to commute Src?");
  }

  if (SrcPassthru.getReg() != Passthru.getReg()) {
    SrcPassthru.setReg(Passthru.getReg());
    // If Src is masked then its passthru needs to be in VRNoV0.
    if (Passthru.getReg().isValid())
      MRI->constrainRegClass(
          Passthru.getReg(),
          TII->getRegClass(Src->getDesc(), SrcPassthru.getOperandNo()));
  }

  if (RISCVII::hasVecPolicyOp(Src->getDesc().TSFlags)) {
    // If MI was tail agnostic and the VL didn't increase, preserve it.
    int64_t Policy = RISCVVType::TAIL_UNDISTURBED_MASK_UNDISTURBED;
    if ((MI.getOperand(5).getImm() & RISCVVType::TAIL_AGNOSTIC) &&
        RISCV::isVLKnownLE(*MRI, MI.getOperand(3), SrcVL))
      Policy |= RISCVVType::TAIL_AGNOSTIC;
    Src->getOperand(RISCVII::getVecPolicyOpNum(Src->getDesc())).setImm(Policy);
  }

  MRI->constrainRegClass(Src->getOperand(0).getReg(),
                         MRI->getRegClass(MI.getOperand(0).getReg()));
  MRI->replaceRegWith(MI.getOperand(0).getReg(), Src->getOperand(0).getReg());
  MI.eraseFromParent();

  return true;
}

/// Try to fold away VMERGE_VVM instructions into their operands:
///
/// %true = PseudoVADD_VV ...
/// %x = PseudoVMERGE_VVM_M1 %false, %false, %true, %mask
/// ->
/// %x = PseudoVADD_VV_M1_MASK %false, ..., %mask
///
/// We can only fold if vmerge's passthru operand, vmerge's false operand and
/// %true's passthru operand (if it has one) are the same. This is because we
/// have to consolidate them into one passthru operand in the result.
///
/// If %true is masked, then we can use its mask instead of vmerge's if vmerge's
/// mask is all ones.
///
/// The resulting VL is the minimum of the two VLs.
///
/// The resulting policy is the effective policy the vmerge would have had,
/// i.e. whether or not it's passthru operand was implicit-def.
bool RISCVVectorPeepholeImpl::foldVMergeToMask(MachineInstr &MI) const {
  if (RISCV::getRVVMCOpcode(MI.getOpcode()) != RISCV::VMERGE_VVM)
    return false;

  // Collect chain of COPYs on True's result for later cleanup.
  SmallVector<MachineInstr *, 4> TrueCopies;
  Register PassthruReg = lookThruCopies(MI.getOperand(1).getReg());
  const MachineOperand &FalseOp = MI.getOperand(2);
  Register FalseReg = lookThruCopies(FalseOp.getReg());
  Register TrueReg = lookThruCopies(MI.getOperand(3).getReg(),
                                    /*OneUseOnly=*/true, &TrueCopies);
  if (!TrueReg.isVirtual() || !MRI->hasOneUse(TrueReg))
    return false;
  MachineInstr *TrueDef = MRI->getVRegDef(TrueReg);
  if (!TrueDef)
    return false;
  MachineInstr &True = *TrueDef;
  if (True.getParent() != MI.getParent())
    return false;
  const MachineOperand &MaskOp = MI.getOperand(4);
  MachineInstr *Mask = MRI->getUniqueVRegDef(MaskOp.getReg());
  assert(Mask);

  const RISCV::RISCVMaskedPseudoInfo *Info =
      RISCV::lookupMaskedIntrinsicByUnmasked(True.getOpcode());
  if (!Info)
    return false;

  // If the EEW of True is different from vmerge's SEW, then we can't fold.
  if (!hasSameEEW(MI, True))
    return false;

  // We require that either passthru and false are the same, or that passthru
  // is undefined.
  if (PassthruReg && !(PassthruReg.isVirtual() && PassthruReg == FalseReg))
    return false;

  std::optional<std::pair<unsigned, unsigned>> NeedsCommute;

  // If True has a passthru operand then it needs to be the same as vmerge's
  // False, since False will be used for the result's passthru operand.
  Register TruePassthru;
  if (RISCVII::isFirstDefTiedToFirstUse(True.getDesc()))
    TruePassthru =
        lookThruCopies(True.getOperand(True.getNumExplicitDefs()).getReg());
  if (TruePassthru && !(TruePassthru.isVirtual() && TruePassthru == FalseReg)) {
    // If True's passthru != False, check if it uses False in another operand
    // and try to commute it.
    int OtherIdx = True.findRegisterUseOperandIdx(FalseReg, TRI);
    if (OtherIdx == -1)
      return false;
    unsigned OpIdx1 = OtherIdx;
    unsigned OpIdx2 = True.getNumExplicitDefs();
    if (!TII->findCommutedOpIndices(True, OpIdx1, OpIdx2))
      return false;
    NeedsCommute = {OpIdx1, OpIdx2};
  }

  // Make sure it doesn't raise any observable fp exceptions, since changing the
  // active elements will affect how fflags is set.
  if (True.hasUnmodeledSideEffects() || True.mayRaiseFPException())
    return false;

  const MachineOperand &VMergeVL =
      MI.getOperand(RISCVII::getVLOpNum(MI.getDesc()));
  const MachineOperand &TrueVL =
      True.getOperand(RISCVII::getVLOpNum(True.getDesc()));

  MachineOperand MinVL = MachineOperand::CreateImm(0);
  if (RISCV::isVLKnownLE(*MRI, TrueVL, VMergeVL))
    MinVL = TrueVL;
  else if (RISCV::isVLKnownLE(*MRI, VMergeVL, TrueVL))
    MinVL = VMergeVL;
  else if (!TruePassthru && !True.mayLoadOrStore())
    // If True's passthru is undef, we can use vmerge's vl.
    MinVL = VMergeVL;
  else
    return false;

  unsigned RVVTSFlags =
      TII->get(RISCV::getRVVMCOpcode(True.getOpcode())).TSFlags;
  if (RISCVII::elementsDependOnVL(RVVTSFlags) && !TrueVL.isIdenticalTo(MinVL))
    return false;
  if (RISCVII::elementsDependOnMask(RVVTSFlags) && !isAllOnesMask(Mask))
    return false;

  // Use a tumu policy, relaxing it to tail agnostic provided that the passthru
  // operand is undefined.
  //
  // However, if the VL became smaller than what the vmerge had originally, then
  // elements past VL that were previously in the vmerge's body will have moved
  // to the tail. In that case we always need to use tail undisturbed to
  // preserve them.
  uint64_t Policy = RISCVVType::TAIL_UNDISTURBED_MASK_UNDISTURBED;
  if (!PassthruReg && RISCV::isVLKnownLE(*MRI, VMergeVL, MinVL))
    Policy |= RISCVVType::TAIL_AGNOSTIC;

  assert(RISCVII::hasVecPolicyOp(True.getDesc().TSFlags) &&
         "Foldable unmasked pseudo should have a policy op already");

  // Make sure Mask, False and MinVL dominate True and its copies, otherwise
  // move down True so it does.
  if (!ensureDominates({&MaskOp, &FalseOp, &MinVL}, True))
    return false;

  if (NeedsCommute) {
    auto [OpIdx1, OpIdx2] = *NeedsCommute;
    [[maybe_unused]] bool Commuted =
        TII->commuteInstruction(True, /*NewMI=*/false, OpIdx1, OpIdx2);
    assert(Commuted && "Failed to commute True?");
    Info = RISCV::lookupMaskedIntrinsicByUnmasked(True.getOpcode());
  }

  True.setDesc(TII->get(Info->MaskedPseudo));

  // Insert the mask operand.
  // TODO: Increment MaskOpIdx by number of explicit defs?
  True.insert(True.operands_begin() + Info->MaskOpIdx +
                  True.getNumExplicitDefs(),
              MachineOperand::CreateReg(MaskOp.getReg(), false));

  // Update the passthru, AVL and policy.
  True.getOperand(True.getNumExplicitDefs()).setReg(FalseReg);
  True.removeOperand(RISCVII::getVLOpNum(True.getDesc()));
  True.insert(True.operands_begin() + RISCVII::getVLOpNum(True.getDesc()),
              MinVL);
  True.getOperand(RISCVII::getVecPolicyOpNum(True.getDesc())).setImm(Policy);

  MRI->replaceRegWith(True.getOperand(0).getReg(), MI.getOperand(0).getReg());
  // Now that True is masked, constrain its operands from vr -> vrnov0.
  for (MachineOperand &MO : True.explicit_operands()) {
    if (!MO.isReg() || !MO.getReg().isVirtual())
      continue;
    MRI->constrainRegClass(
        MO.getReg(), True.getRegClassConstraint(MO.getOperandNo(), TII, TRI));
  }
  // We should clear the IsKill flag since we have a new use now.
  MRI->clearKillFlags(FalseReg);
  MI.eraseFromParent();

  // Cleanup all the COPYs on True's value. We have to manually do this because
  // sometimes sinking True causes these COPY to be invalid (use before define).
  for (MachineInstr *TrueCopy : TrueCopies)
    TrueCopy->eraseFromParent();

  return true;
}

/// Fold a mask-register AND of a mask comparison into a mask-undisturbed
/// masked comparison, saving an instruction:
///
///   %cmp1 = PseudoVMSLT_VV_M1 %a, %b, %vl, %sew
///   %cmp2 = PseudoVMSLT_VV_M1 %c, %d, %vl, %sew
///   %and  = PseudoVMAND_MM %cmp1, %cmp2, %vl, 0
/// ->
///   %cmp1 = PseudoVMSLT_VV_M1 %a, %b, %vl, %sew
///   %and  = PseudoVMSLT_VV_M1_MASK %cmp1, %c, %d, %cmp1, %vl, %sew, mu
///
/// This works because for a mask-undisturbed masked compare whose passthru is
/// the same register as its mask %m, the result is %m[i] ? (c cmp d)[i] :
/// %m[i], which is exactly %m[i] & (c cmp d)[i], i.e. vmand(%m, vmscmp(c, d)).
///
/// Since vmand is commutative it's enough for either operand to be a foldable
/// comparison; the other operand becomes both the mask and the passthru.
bool RISCVVectorPeepholeImpl::foldVMANDToMaskedCompare(MachineInstr &MI) const {
  if (RISCV::getRVVMCOpcode(MI.getOpcode()) != RISCV::VMAND_MM)
    return false;

  // The masked comparison we create needs its mask (and passthru) in v0, which
  // the original vmand did not require. If the vmand's result has more than one
  // use then it is an interior mask value rather than a final result feeding
  // v0, and introducing the v0 requirement tends to add vmv1r.v moves. Only
  // fold single-use results, where the value coalesces onto v0 for free.
  if (!MRI->hasOneUse(MI.getOperand(0).getReg()))
    return false;

  // Try each operand as the comparison to be masked; the other becomes the
  // mask/passthru.
  for (unsigned CmpIdx : {1, 2}) {
    unsigned MaskIdx = CmpIdx == 1 ? 2 : 1;

    // The comparison must be single use so that folding it into MI doesn't
    // leave an extra unmasked comparison behind.
    SmallVector<MachineInstr *, 4> CmpCopies;
    Register CmpReg = lookThruCopies(MI.getOperand(CmpIdx).getReg(),
                                     /*OneUseOnly=*/true, &CmpCopies);
    if (!CmpReg.isVirtual() || !MRI->hasOneUse(CmpReg))
      continue;
    MachineInstr &Cmp = *MRI->getUniqueVRegDef(CmpReg);
    if (Cmp.getParent() != MI.getParent())
      continue;
    if (!RISCVInstrInfo::isRVVCompare(Cmp))
      continue;

    // Find the masked pseudo corresponding to the unmasked comparison.
    const RISCV::RISCVMaskedPseudoInfo *Info =
        RISCV::lookupMaskedIntrinsicByUnmasked(Cmp.getOpcode());
    if (!Info)
      continue;

    // The EEW of the comparison's dest must match vmand's SEW.
    if (!hasSameEEW(MI, Cmp))
      continue;

    // Masking restricts the comparison to the mask's active elements, so any FP
    // exceptions raised on inactive elements would be lost.
    if (Cmp.hasUnmodeledSideEffects() || Cmp.mayRaiseFPException())
      continue;

    // All active elements of vmand must also be active in the comparison. If
    // the comparison's VL were smaller, elements in between the two VLs would
    // become tail elements of the masked comparison and could not be preserved
    // from the mask because mask results are always tail agnostic.
    const MachineOperand &CmpVL =
        Cmp.getOperand(RISCVII::getVLOpNum(Cmp.getDesc()));
    const MachineOperand &MIVL =
        MI.getOperand(RISCVII::getVLOpNum(MI.getDesc()));
    if (!RISCV::isVLKnownLE(*MRI, MIVL, CmpVL))
      continue;

    const MachineOperand &MaskOp = MI.getOperand(MaskIdx);
    Register MaskReg = MaskOp.getReg();

    unsigned MaskedOpc = Info->MaskedPseudo;
    const MCInstrDesc &MaskedDesc = TII->get(MaskedOpc);
    unsigned SEW = Cmp.getOperand(RISCVII::getSEWOpNum(Cmp.getDesc())).getImm();

    // Only fold if the masked comparison's dest can live in v0. Its mask
    // operand must be v0, and we reuse the mask as the passthru, so if the dest
    // can also be v0 the whole thing coalesces onto v0 and we save the vmand
    // for free. For LMUL >= 2 the dest is earlyclobbered into vrnov0, which
    // would force extra vmv1r.v moves for the mask and result and make this a
    // regression, so bail out in that case. This check must happen before we
    // mutate any instructions below.
    if (!TII->getRegClass(MaskedDesc, 0)->contains(RISCV::V0))
      continue;

    // Make sure the mask and VL dominate the comparison, sinking it if needed.
    if (!ensureDominates({&MaskOp, &MIVL}, Cmp))
      continue;

    // The masked comparison's mask operand lives in the VMV0 (v0) class, and
    // its passthru operand shares the dest's class. Copy the vmand mask into
    // both; the coalescer collapses these back onto v0, matching the
    // two-instruction ideal.
    Register MaskV0Reg = MRI->createVirtualRegister(&RISCV::VMV0RegClass);
    BuildMI(*MI.getParent(), Cmp, Cmp.getDebugLoc(),
            TII->get(TargetOpcode::COPY), MaskV0Reg)
        .addReg(MaskReg);
    // The passthru shares the dest's class, which the V0 check above restricts
    // to LMUL <= 1, so it is always a single vector register.
    Register PassthruReg = MRI->createVirtualRegister(&RISCV::VRRegClass);
    BuildMI(*MI.getParent(), Cmp, Cmp.getDebugLoc(),
            TII->get(TargetOpcode::COPY), PassthruReg)
        .addReg(MaskReg);

    // Build the masked comparison. Its dest reuses vmand's dest; the passthru
    // (tied to the dest) and mask are both the other vmand operand. Preserve
    // the source comparison's MI flags (e.g. nofpexcept), which still hold
    // since the masked comparison operates on a subset of the original active
    // elements.
    Register DestReg = MI.getOperand(0).getReg();
    MachineInstr *Masked =
        BuildMI(*MI.getParent(), Cmp, MIMetadata(Cmp), MaskedDesc, DestReg)
            .addReg(PassthruReg)
            .add(Cmp.getOperand(1))
            .add(Cmp.getOperand(2))
            .addReg(MaskV0Reg)
            .add(MIVL)
            .addImm(SEW)
            // The result is a mask register, whose tail is always agnostic, so
            // we only need mask-undisturbed (MASK_AGNOSTIC clear) to preserve
            // the inactive elements from the mask/passthru.
            .addImm(RISCVVType::TAIL_AGNOSTIC)
            .setMIFlags(Cmp.getFlags());

    // Now that the comparison is masked, constrain its operands to the masked
    // pseudo's register classes (e.g. vr -> vrnov0 for LMUL >= 2).
    for (MachineOperand &MO : Masked->explicit_operands()) {
      if (!MO.isReg() || !MO.getReg().isVirtual())
        continue;
      if (const TargetRegisterClass *RC =
              Masked->getRegClassConstraint(MO.getOperandNo(), TII, TRI))
        MRI->constrainRegClass(MO.getReg(), RC);
    }
    MRI->clearKillFlags(MaskReg);
    MI.eraseFromParent();
    Cmp.eraseFromParent();
    for (MachineInstr *CmpCopy : CmpCopies)
      CmpCopy->eraseFromParent();

    return true;
  }

  return false;
}

bool RISCVVectorPeepholeImpl::run(MachineFunction &MF) {
  // Skip if the vector extension is not enabled.
  ST = &MF.getSubtarget<RISCVSubtarget>();
  if (!ST->hasVInstructions())
    return false;

  TII = ST->getInstrInfo();
  MRI = &MF.getRegInfo();
  TRI = MRI->getTargetRegisterInfo();

  bool Changed = false;

  for (MachineBasicBlock &MBB : MF) {
    for (MachineInstr &MI : make_early_inc_range(MBB))
      Changed |= foldVMergeToMask(MI);

    for (MachineInstr &MI : make_early_inc_range(MBB))
      Changed |= foldVMANDToMaskedCompare(MI);

    for (MachineInstr &MI : make_early_inc_range(MBB)) {
      Changed |= convertToVLMAX(MI);
      Changed |= convertToUnmasked(MI);
      Changed |= convertToWholeRegister(MI);
      Changed |= convertAllOnesVMergeToVMv(MI);
      Changed |= convertSameMaskVMergeToVMv(MI);
      if (foldUndefPassthruVMV_V_V(MI)) {
        Changed |= true;
        continue; // MI is erased
      }
      Changed |= foldVMV_V_V(MI);
    }
  }

  return Changed;
}

bool RISCVVectorPeepholeLegacy::runOnMachineFunction(MachineFunction &MF) {
  if (skipFunction(MF.getFunction()))
    return false;
  return RISCVVectorPeepholeImpl().run(MF);
}

PreservedAnalyses
RISCVVectorPeepholePass::run(MachineFunction &MF,
                             MachineFunctionAnalysisManager &MFAM) {
  MFPropsModifier _(*this, MF);
  bool Changed = RISCVVectorPeepholeImpl().run(MF);
  if (!Changed)
    return PreservedAnalyses::all();

  PreservedAnalyses PA = getMachineFunctionPassPreservedAnalyses();
  PA.preserveSet<CFGAnalyses>();
  PA.preserve<MachineRegisterClassAnalysis>();
  return PA;
}

FunctionPass *llvm::createRISCVVectorPeepholeLegacyPass() {
  return new RISCVVectorPeepholeLegacy();
}