Update approach:

- Always run V_SET_INACTIVE in WWM if it is touched by WWM marking
- Lower any V_SET_INACTIVE not touched by marking to COPY
- Further optimize the V_SET_INACTIVE expansion based on being in WWM

Author: perlfu

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

@@ -2278,15 +2278,14 @@ bool SIInstrInfo::expandPostRAPseudo(MachineInstr &MI) const {
                            ? AMDGPU::V_MOV_B64_PSEUDO
                            : AMDGPU::V_MOV_B32_e32;
     Register ExecReg = ST.isWave32() ? AMDGPU::EXEC_LO : AMDGPU::EXEC;
-
     Register DstReg = MI.getOperand(0).getReg();
     MachineOperand &ActiveSrc = MI.getOperand(1);
     MachineOperand &InactiveSrc = MI.getOperand(2);
 
-    bool VMov64 = VMovOpc != AMDGPU::V_MOV_B32_e32;
-
-    // Find implicit exec src if this is running in WWM.
-    Register ExecSrcReg = 0;
+    // Find implicit register defining lanes active outside WWM.
+    // Note: default here is set to ExecReg so that functional MIR is still
+    // generated if implicit def is not found and assertions are disabled.
+    Register ExecSrcReg = ExecReg;
     for (auto &Op : MI.implicit_operands()) {
       if (Op.isDef() || !Op.isReg())
         continue;
@@ -2297,120 +2296,135 @@ bool SIInstrInfo::expandPostRAPseudo(MachineInstr &MI) const {
       ExecSrcReg = OpReg;
       break;
     }
+    assert(ExecSrcReg != ExecReg &&
+           "V_SET_INACTIVE must be in known WWM region");
 
     // Ideally in WWM this operation is lowered to V_CNDMASK; however,
     // constant bus constraints and the presence of literal constants
     // present an issue.
     // Fallback to V_MOV base lowering in all but the common cases.
-    bool InWWM = !!ExecSrcReg;
-    bool UseVCndMask = false;
-    if (InWWM) {
-      const MachineFunction *MF = MI.getParent()->getParent();
-      const MachineRegisterInfo &MRI = MF->getRegInfo();
-      const unsigned Opcode = AMDGPU::V_CNDMASK_B32_e64;
-      const MCInstrDesc &Desc = get(Opcode);
-      int Src0Idx = AMDGPU::getNamedOperandIdx(Opcode, AMDGPU::OpName::src0);
-      int Src1Idx = AMDGPU::getNamedOperandIdx(Opcode, AMDGPU::OpName::src1);
-      int ConstantBusLimit = ST.getConstantBusLimit(AMDGPU::V_CNDMASK_B32_e64);
-      int LiteralLimit = ST.hasVOP3Literal() ? 1 : 0;
-      int ConstantBusUses = 1; // Starts at one for ExecRegSrc
-      int LiteralConstants = 0;
-      ConstantBusUses +=
-          usesConstantBus(MRI, ActiveSrc, Desc.operands()[Src1Idx]) ? 1 : 0;
-      ConstantBusUses +=
-          usesConstantBus(MRI, InactiveSrc, Desc.operands()[Src0Idx]) ? 1 : 0;
-      LiteralConstants +=
-          ActiveSrc.isImm() &&
-                  !isInlineConstant(ActiveSrc, Desc.operands()[Src1Idx])
-              ? 1
-              : 0;
-      LiteralConstants +=
-          InactiveSrc.isImm() &&
-                  !isInlineConstant(InactiveSrc, Desc.operands()[Src0Idx])
-              ? 1
-              : 0;
-      UseVCndMask = ConstantBusUses <= ConstantBusLimit &&
-                    LiteralConstants <= LiteralLimit &&
-                    (!VMov64 || (ActiveSrc.isReg() && InactiveSrc.isReg()));
+    const bool VMov64 = VMovOpc != AMDGPU::V_MOV_B32_e32;
+    const MachineFunction *MF = MI.getParent()->getParent();
+    const MachineRegisterInfo &MRI = MF->getRegInfo();
+    const unsigned Opcode = AMDGPU::V_CNDMASK_B32_e64;
+    const MCInstrDesc &Desc = get(Opcode);
+
+    const APInt ActiveImm(64, ActiveSrc.isImm() ? ActiveSrc.getImm() : 0);
+    const APInt InactiveImm(64, InactiveSrc.isImm() ? InactiveSrc.getImm() : 0);
+    const APInt ActiveImmLo(32, ActiveImm.getLoBits(32).getZExtValue());
+    const APInt ActiveImmHi(32, ActiveImm.getHiBits(32).getZExtValue());
+    const APInt InactiveImmLo(32, InactiveImm.getLoBits(32).getZExtValue());
+    const APInt InactiveImmHi(32, InactiveImm.getHiBits(32).getZExtValue());
+
+    int Src0Idx = AMDGPU::getNamedOperandIdx(Opcode, AMDGPU::OpName::src0);
+    int Src1Idx = AMDGPU::getNamedOperandIdx(Opcode, AMDGPU::OpName::src1);
+
+    int ConstantBusLimit = ST.getConstantBusLimit(AMDGPU::V_CNDMASK_B32_e64);
+    int LiteralLimit = ST.hasVOP3Literal() ? 1 : 0;
+    int ConstantBusUses =
+        1 + // Starts at 1 for ExecSrcReg
+        (usesConstantBus(MRI, ActiveSrc, Desc.operands()[Src1Idx]) ? 1 : 0) +
+        (usesConstantBus(MRI, InactiveSrc, Desc.operands()[Src0Idx]) ? 1 : 0);
+    int LiteralConstants =
+        (ActiveSrc.isImm() && !isInlineConstant(ActiveImm) ? 1 : 0) +
+        (InactiveSrc.isImm() && !isInlineConstant(InactiveImm) ? 1 : 0);
+
+    bool UseVCndMask =
+        ConstantBusUses <= ConstantBusLimit && LiteralConstants <= LiteralLimit;
+    if (VMov64 && UseVCndMask) {
+      // Decomposition must not introduce new literals.
+      UseVCndMask &=
+          ActiveSrc.isReg() ||
+          (isInlineConstant(ActiveImmLo) && isInlineConstant(ActiveImmLo)) ||
+          (!isInlineConstant(ActiveImm));
+      UseVCndMask &= InactiveSrc.isReg() ||
+                     (isInlineConstant(InactiveImmLo) &&
+                      isInlineConstant(InactiveImmLo)) ||
+                     (!isInlineConstant(InactiveImm));
     }
 
     if (UseVCndMask && VMov64) {
-      // WWM B64; decompose to two B32 operations.
-      // Test above ensures that both sources are registers.
-      // Note: this is done to avoid falling back to V_MOV multiple times
-      // and introducing exec manipulation for each VGPR separately.
-      assert(ActiveSrc.isReg() && InactiveSrc.isReg());
-      Register ActiveLo = RI.getSubReg(ActiveSrc.getReg(), AMDGPU::sub0);
-      Register ActiveHi = RI.getSubReg(ActiveSrc.getReg(), AMDGPU::sub1);
-      Register InactiveLo = RI.getSubReg(InactiveSrc.getReg(), AMDGPU::sub0);
-      Register InactiveHi = RI.getSubReg(InactiveSrc.getReg(), AMDGPU::sub1);
-      MachineInstr *Tmp;
-      Tmp = BuildMI(MBB, MI, DL, get(AMDGPU::V_SET_INACTIVE_B32),
-                    RI.getSubReg(DstReg, AMDGPU::sub0))
-                .addReg(InactiveLo)
-                .addReg(ActiveLo)
-                .addReg(ExecSrcReg, RegState::Implicit)
-                .addReg(DstReg, RegState::ImplicitDefine);
-      expandPostRAPseudo(*Tmp);
-      Tmp = BuildMI(MBB, MI, DL, get(AMDGPU::V_SET_INACTIVE_B32),
-                    RI.getSubReg(DstReg, AMDGPU::sub1))
-                .addReg(InactiveHi, InactiveSrc.isKill() ? RegState::Kill : 0)
-                .addReg(ActiveHi, ActiveSrc.isKill() ? RegState::Kill : 0)
-                .addReg(ExecSrcReg, RegState::Implicit)
-                .addReg(DstReg, RegState::ImplicitDefine);
-      expandPostRAPseudo(*Tmp);
+      // Dual V_CNDMASK_B32
+      MachineOperand ActiveLo =
+          ActiveSrc.isReg()
+              ? MachineOperand::CreateReg(
+                    RI.getSubReg(ActiveSrc.getReg(), AMDGPU::sub0), false,
+                    /*isImp=*/false, /*isKill*/ false)
+              : MachineOperand::CreateImm(ActiveImmLo.getSExtValue());
+      MachineOperand ActiveHi =
+          ActiveSrc.isReg()
+              ? MachineOperand::CreateReg(
+                    RI.getSubReg(ActiveSrc.getReg(), AMDGPU::sub1), false,
+                    /*isImp=*/false, /*isKill*/ ActiveSrc.isKill())
+              : MachineOperand::CreateImm(ActiveImmHi.getSExtValue());
+      MachineOperand InactiveLo =
+          InactiveSrc.isReg()
+              ? MachineOperand::CreateReg(
+                    RI.getSubReg(InactiveSrc.getReg(), AMDGPU::sub0), false,
+                    /*isImp=*/false, /*isKill*/ false)
+              : MachineOperand::CreateImm(InactiveImmLo.getSExtValue());
+      MachineOperand InactiveHi =
+          InactiveSrc.isReg()
+              ? MachineOperand::CreateReg(
+                    RI.getSubReg(InactiveSrc.getReg(), AMDGPU::sub1), false,
+                    /*isImp=*/false, /*isKill*/ InactiveSrc.isKill())
+              : MachineOperand::CreateImm(InactiveImmHi.getSExtValue());
+      BuildMI(MBB, MI, DL, get(Opcode), RI.getSubReg(DstReg, AMDGPU::sub0))
+          .addImm(0)
+          .add(InactiveLo)
+          .addImm(0)
+          .add(ActiveLo)
+          .addReg(ExecSrcReg)
+          .addReg(DstReg, RegState::ImplicitDefine);
+      BuildMI(MBB, MI, DL, get(Opcode), RI.getSubReg(DstReg, AMDGPU::sub1))
+          .addImm(0)
+          .add(InactiveHi)
+          .addImm(0)
+          .add(ActiveHi)
+          .addReg(ExecSrcReg)
+          .addReg(DstReg, RegState::ImplicitDefine);
     } else if (UseVCndMask) {
-      // WWM B32; use V_CNDMASK.
-      MachineInstr *VCndMask =
-          BuildMI(MBB, MI, DL, get(AMDGPU::V_CNDMASK_B32_e64), DstReg)
-              .addImm(0)
-              .add(InactiveSrc)
-              .addImm(0)
-              .add(ActiveSrc)
-              .addReg(ExecSrcReg);
-      // Copy implicit defs in case this is part of V_SET_INACTIVE_B64.
-      for (auto &Op : MI.implicit_operands()) {
-        if (!Op.isDef())
-          continue;
-        VCndMask->addOperand(Op);
-      }
+      // Single V_CNDMASK_B32
+      BuildMI(MBB, MI, DL, get(Opcode), DstReg)
+          .addImm(0)
+          .add(InactiveSrc)
+          .addImm(0)
+          .add(ActiveSrc)
+          .addReg(ExecSrcReg);
     } else {
       // Fallback V_MOV case.
-      // Avoid unnecessary work if a src is the destination.
+      // Avoid unnecessary work if a source VGPR is also the destination.
       // This can happen if WWM register allocation was efficient.
-      bool SkipActive = ActiveSrc.isReg() && ActiveSrc.getReg() == DstReg;
-      bool SkipInactive = InactiveSrc.isReg() && InactiveSrc.getReg() == DstReg;
-      if (!SkipActive) {
-        if (InWWM) {
-          // Cancel WWM
-          BuildMI(MBB, MI, DL, get(MovOpc), ExecReg).addReg(ExecSrcReg);
+      // Note: this assumes WWM execution.
+      bool DstIsActive = ActiveSrc.isReg() && ActiveSrc.getReg() == DstReg;
+      bool DstIsInactive =
+          InactiveSrc.isReg() && InactiveSrc.getReg() == DstReg;
+      if (!DstIsInactive) {
+        // Set exec mask to inactive lanes,
+        // but only if active lanes would be overwritten.
+        if (DstIsActive) {
+          MachineInstr *ExecMI =
+              BuildMI(MBB, MI, DL, get(NotOpc), ExecReg).addReg(ExecSrcReg);
+          ExecMI->addRegisterDead(AMDGPU::SCC, TRI); // SCC is overwritten
         }
-        // Copy active lanes
+        // Copy inactive lanes
         MachineInstr *VMov =
-            BuildMI(MBB, MI, DL, get(VMovOpc), MI.getOperand(0).getReg())
-                .add(ActiveSrc);
+            BuildMI(MBB, MI, DL, get(VMovOpc), DstReg).add(InactiveSrc);
         if (VMov64)
           expandPostRAPseudo(*VMov);
       }
-      if (!SkipInactive) {
-        // Set exec mask to inactive lanes
-        MachineInstr *ExecMI = BuildMI(MBB, MI, DL, get(NotOpc), ExecReg)
-                                   .addReg(InWWM ? ExecSrcReg : ExecReg);
-        ExecMI->addRegisterDead(AMDGPU::SCC, TRI); // SCC is overwritten
-        // Copy inactive lanes
+      if (!DstIsActive) {
+        // Set exec mask to active lanes
+        BuildMI(MBB, MI, DL, get(MovOpc), ExecReg).addReg(ExecSrcReg);
+        // Copy active lanes
         MachineInstr *VMov =
-            BuildMI(MBB, MI, DL, get(VMovOpc), DstReg).add(InactiveSrc);
+            BuildMI(MBB, MI, DL, get(VMovOpc), MI.getOperand(0).getReg())
+                .add(ActiveSrc);
         if (VMov64)
           expandPostRAPseudo(*VMov);
-        if (!InWWM) {
-          // Restore original exec mask
-          BuildMI(MBB, MI, DL, get(NotOpc), ExecReg).addReg(ExecReg);
-        }
-      }
-      if (InWWM) {
-        // Restore WWM
-        BuildMI(MBB, MI, DL, get(MovOpc), ExecReg).addImm(-1);
       }
+      // Restore WWM
+      BuildMI(MBB, MI, DL, get(MovOpc), ExecReg).addImm(-1);
     }
     MI.eraseFromParent();
     break;