[llvm] [RISCV] Fold vmand of mask comparisons into a masked comparison (PR #216264)

[Min-Yih Hsu via llvm-commits]

================
@@ -787,6 +788,169 @@ bool RISCVVectorPeephole::foldVMergeToMask(MachineInstr &MI) const {
   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 RISCVVectorPeephole::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;
+
+    // Only fold comparisons: entries in the masked pseudo table whose unmasked
+    // form has neither a passthru nor a policy operand. This excludes ops like
+    // vmsbf.m/viota.m which share that shape but whose active elements depend
+    // on the mask, so masking them would change their result.
+    const RISCV::RISCVMaskedPseudoInfo *Info =
+        RISCV::lookupMaskedIntrinsicByUnmasked(Cmp.getOpcode());
+    if (!Info)
+      continue;
+    const MCInstrDesc &UnmaskedDesc = Cmp.getDesc();
+    if (RISCVII::isFirstDefTiedToFirstUse(UnmaskedDesc) ||
+        RISCVII::hasVecPolicyOp(UnmaskedDesc.TSFlags))
+      continue;
+    unsigned CmpMCOpc = RISCV::getRVVMCOpcode(Cmp.getOpcode());
+    if (RISCVII::elementsDependOnMask(TII->get(CmpMCOpc).TSFlags))
+      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;
+
+    // Use the smaller of the two VLs for the result. The comparison's inactive
+    // (and tail) elements will be filled from the mask, so if vmand's VL is the
+    // smaller one we still need the comparison's original VL for the elements
+    // between them to come from the mask rather than the raw comparison.
+    const MachineOperand &CmpVL =
+        Cmp.getOperand(RISCVII::getVLOpNum(Cmp.getDesc()));
+    const MachineOperand &MIVL =
+        MI.getOperand(RISCVII::getVLOpNum(MI.getDesc()));
+    MachineOperand MinVL = MachineOperand::CreateImm(0);
+    if (RISCV::isVLKnownLE(CmpVL, MIVL))
+      MinVL = CmpVL;
+    else if (RISCV::isVLKnownLE(MIVL, CmpVL))
+      MinVL = MIVL;
+    else
+      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, &MinVL}, 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);
+    Register PassthruReg =
+        MRI->createVirtualRegister(TII->getRegClass(MaskedDesc, 1));
+    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, Cmp.getDebugLoc(), MaskedDesc)
+            .addDef(DestReg)
----------------
mshockwave wrote:

```suggestion
        BuildMI(*MI.getParent(), Cmp, MIMetadata(Cmp), MaskedDesc, DestReg)
```

https://github.com/llvm/llvm-project/pull/216264