AMDGPU: Move SIWholeQuadMode pass to after machine scheduling

[nhaehnle]

Summary

Most of this change is about adjusting the analysis, which can now no longer
rely on SSA form and must use LiveIntervals instead.

WQM-related instructions manipulate EXEC and therefore act as a scheduling
barrier for most instructions. With this change, shaders with a mix of WQM
and exact instructions should be scheduled much better because there is no
EXEC interference. This opens the way for additional improvements.

[arsenm]

lib/Target/AMDGPU/SIWholeQuadMode.cpp

151

This should be const or moved to a static helper function. I think it will warn as unused in release builds

236

Single quotes around single character

241

ditto

314–315

You can combine these into one assert

553

Can you move the bit test into a helper function in the BI struct, same for everywhere else

[arsenm]

lib/Target/AMDGPU/SIWholeQuadMode.cpp

188–191

I think you need to add the corresponding passes with INITIALIZE_PASS_DEPENDENCY

[arsenm]

lib/Target/AMDGPU/SIWholeQuadMode.cpp

655

Why are you using SlotIndexes directly? I think in other places LIS->ReplaceMachineInstrInMaps is used

[nhaehnle]

Rebase on top of recent changes, address review comments.

[nhaehnle]

Remove special handling of SI_KILL (slipped back in after rebase).

[arsenm]

lib/Target/AMDGPU/SIWholeQuadMode.cpp

713

You shouldn't need a const_cast (I thought I had a patch that fixed this before)

[nhaehnle]

Rebase and address review comments.

[nhaehnle]

Another rebase and ping.

[nhaehnle]

Looks like some comments I wrote earlier were accidentally not sent...

lib/Target/AMDGPU/SIWholeQuadMode.cpp

151

Changed to const. I see no unused warning yet, perhaps that only applies to static functions.

241

Done.

314–315

I'd personally rather not, since they're logically about different things: the first one is about not changing an already marked instruction, the second about making sure that both flags are not set simultaneously.

553

The tests are all quite different, one would need many different helper functions. I think it would make the code harder to follow.

655

Done.

713

Thanks, that got lost in the rebase. Fixed now.

[michel.daenzer]

lib/Target/AMDGPU/SIWholeQuadMode.cpp

325–326

Also, by using separate assert statements for each condition, it's immediately obvious which condition wasn't satisfied.

[nhaehnle]

ping

[• tstellarAMD]

This is fine with me: LGTM.

[arsenm]

This will save recomputing LiveIntervals an extra time. I also might need this to make SIOptimizeExecMaskingPreRA correct

test/CodeGen/AMDGPU/wqm.ll

516

Intrinsics need updating

lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

if (getOptLevel() > CodeGenOpt::None) {
     // This should be run after scheduling, but before register allocation. It
     // also need extra copies to the address operand to be eliminated.
 
     // FIXME: Move pre-RA and remove extra reg coalescer run.
     insertPass(&MachineSchedulerID, &SILoadStoreOptimizerID);
     insertPass(&MachineSchedulerID, &RegisterCoalescerID);
   }
 
+  insertPass(&MachineSchedulerID, &SIWholeQuadModeID);
+
   addPass(createSIShrinkInstructionsPass());
-  addPass(createSIWholeQuadModePass());
 }
 
 void GCNPassConfig::addFastRegAlloc(FunctionPass *RegAllocPass) {
   TargetPassConfig::addFastRegAlloc(RegAllocPass);
 }
 
 void GCNPassConfig::addOptimizedRegAlloc(FunctionPass *RegAllocPass) {
   TargetPassConfig::addOptimizedRegAlloc(RegAllocPass);

lib/Target/AMDGPU/SIWholeQuadMode.cpp

 //===-- SIWholeQuadMode.cpp - enter and suspend whole quad mode -----------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 /// \file
 /// \brief This pass adds instructions to enable whole quad mode for pixel
 /// shaders.
 ///
 /// Whole quad mode is required for derivative computations, but it interferes
-/// with shader side effects (stores and atomics). This pass is run on the
-/// scheduled machine IR but before register coalescing, so that machine SSA is
-/// available for analysis. It ensures that WQM is enabled when necessary, but
-/// disabled around stores and atomics.
+/// with shader side effects (stores and atomics). This pass is run after
+/// machine instruction scheduling but before register allocation. It ensures
+/// that WQM is enabled when necessary, but disabled around stores and atomics.
 ///
 /// When necessary, this pass creates a function prolog
 ///
 ///   S_MOV_B64 LiveMask, EXEC
 ///   S_WQM_B64 EXEC, EXEC
 ///
 /// to enter WQM at the top of the function and surrounds blocks of Exact
 /// instructions by

 #include "AMDGPU.h"
 #include "AMDGPUSubtarget.h"
 #include "SIInstrInfo.h"
 #include "SIMachineFunctionInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/Support/Format.h"
+
+#include <map>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "si-wqm"
 
 namespace {
 
 enum {
   StateWQM = 0x1,
   StateExact = 0x2,
 };
 
 struct PrintState {
 public:
-  explicit PrintState(int State) : State(State) {}
+  explicit PrintState(int State, bool Align = false)
+      : State(State), Align(Align) {}
 
   int State;
+  bool Align;
 };
 
 static raw_ostream &operator<<(raw_ostream &OS, const PrintState &PS) {
   if (PS.State & StateWQM)
     OS << "WQM";
   if (PS.State & StateExact) {
     if (PS.State & StateWQM)
       OS << '|';
     OS << "Exact";
   }
 
+  if (PS.Align) {
+    if (!PS.State)
+      OS << "     ";
+    else if (!(PS.State & StateExact))
+      OS << "  ";
+  }
+
   return OS;
 }
 
 struct InstrInfo {
   char Needs = 0;
   char OutNeeds = 0;
 };
 
 struct BlockInfo {
-  char Needs = 0;
-  char InNeeds = 0;
-  char OutNeeds = 0;
+  struct NeedsFlags {
+    char Self = 0;
+    char In = 0;
+    char Out = 0;
+  };
+
+  NeedsFlags Needs;
+  NeedsFlags Propagated;
 };
 
 struct WorkItem {
+  struct Value {
+    SlotIndex Slot;
+    unsigned Reg = 0;
+    LaneBitmask LaneMask = 0;
+  };
+
   MachineBasicBlock *MBB = nullptr;
   MachineInstr *MI = nullptr;
+  Value V;
 
   WorkItem() {}
   WorkItem(MachineBasicBlock *MBB) : MBB(MBB) {}
   WorkItem(MachineInstr *MI) : MI(MI) {}
+  WorkItem(SlotIndex Slot, unsigned Reg, unsigned LaneMask) {
+    V.Slot = Slot;
+    V.Reg = Reg;
+    V.LaneMask = LaneMask;
+  }
 };
 
 class SIWholeQuadMode : public MachineFunctionPass {
 private:
   const SIInstrInfo *TII;
   const SIRegisterInfo *TRI;
   MachineRegisterInfo *MRI;
   LiveIntervals *LIS;
 
   DenseMap<const MachineInstr *, InstrInfo> Instructions;
   DenseMap<MachineBasicBlock *, BlockInfo> Blocks;
+  std::map<std::pair<SlotIndex, unsigned>, LaneBitmask> WQMValues;
   SmallVector<MachineInstr *, 1> LiveMaskQueries;
 
-  void printInfo();
+  void printInfo(MachineFunction &MF) const;
 
+  void markValueWQM(SlotIndex Slot, unsigned Reg, LaneBitmask LaneMask,
+                    std::vector<WorkItem> &Worklist);
   void markInstruction(MachineInstr &MI, char Flag,
                        std::vector<WorkItem> &Worklist);
   void markUsesWQM(const MachineInstr &MI, std::vector<WorkItem> &Worklist);
   char scanInstructions(MachineFunction &MF, std::vector<WorkItem> &Worklist);
+  void propagateValueSub(const WorkItem::Value &V, LaneBitmask LaneMask,
+                         const LiveRange &LR, std::vector<WorkItem> &Worklist);
+  void propagateValue(const WorkItem::Value &V,
+                      std::vector<WorkItem> &Worklist);
   void propagateInstruction(MachineInstr &MI, std::vector<WorkItem> &Worklist);
   void propagateBlock(MachineBasicBlock &MBB, std::vector<WorkItem> &Worklist);
   char analyzeFunction(MachineFunction &MF);
 
   void toExact(MachineBasicBlock &MBB, MachineBasicBlock::iterator Before,
                unsigned SaveWQM, unsigned LiveMaskReg);
   void toWQM(MachineBasicBlock &MBB, MachineBasicBlock::iterator Before,
              unsigned SavedWQM);
public:
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
   const char *getPassName() const override {
     return "SI Whole Quad Mode";
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<LiveIntervals>();
+    AU.addPreserved<LiveIntervals>();
+    AU.addPreserved<SlotIndexes>();
     AU.setPreservesCFG();
     MachineFunctionPass::getAnalysisUsage(AU);
   }
 };
 
 } // End anonymous namespace
 
 char SIWholeQuadMode::ID = 0;
 
 INITIALIZE_PASS_BEGIN(SIWholeQuadMode, DEBUG_TYPE, "SI Whole Quad Mode", false,
                       false)
 INITIALIZE_PASS_DEPENDENCY(LiveIntervals)
 INITIALIZE_PASS_END(SIWholeQuadMode, DEBUG_TYPE, "SI Whole Quad Mode", false,
                     false)
 
 char &llvm::SIWholeQuadModeID = SIWholeQuadMode::ID;
 
 FunctionPass *llvm::createSIWholeQuadModePass() {
   return new SIWholeQuadMode;
 }
 
-void SIWholeQuadMode::printInfo() {
-  for (const auto &BII : Blocks) {
-    dbgs() << "\nBB#" << BII.first->getNumber() << ":\n"
-           << "  InNeeds = " << PrintState(BII.second.InNeeds)
-           << ", Needs = " << PrintState(BII.second.Needs)
-           << ", OutNeeds = " << PrintState(BII.second.OutNeeds) << "\n\n";
+void SIWholeQuadMode::printInfo(MachineFunction &MF) const {
+  for (MachineBasicBlock &MBB : MF) {
+    BlockInfo BI;
+    auto BII = Blocks.find(&MBB);
+    if (BII != Blocks.end())
+      BI = BII->second;
+
+    dbgs() << "\nBB#" << MBB.getNumber() << ":\n"
+           << "  In = " << PrintState(BI.Needs.In)
+           << ", Self = " << PrintState(BI.Needs.Self)
+           << ", Out = " << PrintState(BI.Needs.Out) << '\n';
 
-    for (const MachineInstr &MI : *BII.first) {
+    for (MachineInstr &MI : MBB) {
+      InstrInfo II;
       auto III = Instructions.find(&MI);
-      if (III == Instructions.end())
-        continue;
+      if (III != Instructions.end())
+        II = III->second;
 
-      dbgs() << "  " << MI << "    Needs = " << PrintState(III->second.Needs)
-             << ", OutNeeds = " << PrintState(III->second.OutNeeds) << '\n';
+      dbgs() << ' ' << PrintState(II.Needs, true) << ' ' << MI;
     }
   }
 }
 
+void SIWholeQuadMode::markValueWQM(SlotIndex Slot, unsigned Reg,
+                                   LaneBitmask LaneMask,
+                                   std::vector<WorkItem> &Worklist) {
+  LaneBitmask &Mask = WQMValues[std::make_pair(Slot, Reg)];
+  LaneMask &= ~Mask;
+  if (LaneMask) {
+    Mask |= LaneMask;
+    Worklist.emplace_back(Slot, Reg, LaneMask);
+  }
+}
+
 void SIWholeQuadMode::markInstruction(MachineInstr &MI, char Flag,
                                       std::vector<WorkItem> &Worklist) {
   InstrInfo &II = Instructions[&MI];
 
   assert(Flag == StateWQM || Flag == StateExact);
 
   // Ignore if the instruction is already marked. The typical case is that we
   // mark an instruction WQM multiple times, but for atomics it can happen that
   // Flag is StateWQM, but Needs is already set to StateExact. In this case,
   // letting the atomic run in StateExact is correct as per the relevant specs.
   if (II.Needs)
     return;
 
   II.Needs = Flag;
-  Worklist.push_back(&MI);
+
+  MachineBasicBlock *MBB = MI.getParent();
+  BlockInfo &BI = Blocks[MBB];
+
+  if (!(BI.Needs.Self & Flag)) {
+    BI.Needs.Self |= Flag;
+    BI.Needs.In |= Flag;
+    Worklist.push_back(MBB);
+  }
+
+  if (MachineInstr *PrevMI = MI.getPrevNode()) {
+    InstrInfo &PrevII = Instructions[PrevMI];
+    if (Flag & ~PrevII.OutNeeds) {
+      PrevII.OutNeeds |= Flag;
+      Worklist.push_back(PrevMI);
+    }
+  }
+
+  if (Flag == StateWQM)
+    markUsesWQM(MI, Worklist);
 }
 
 /// Mark all instructions defining the uses in \p MI as WQM.
 void SIWholeQuadMode::markUsesWQM(const MachineInstr &MI,
                                   std::vector<WorkItem> &Worklist) {
+  SlotIndex Idx = LIS->getInstructionIndex(MI);
+
   for (const MachineOperand &Use : MI.uses()) {
     if (!Use.isReg() || !Use.isUse())
       continue;
 
     unsigned Reg = Use.getReg();
 
     // Handle physical registers that we need to track; this is mostly relevant
     // for VCC, which can appear as the (implicit) input of a uniform branch,
     // e.g. when a loop counter is stored in a VGPR.
     if (!TargetRegisterInfo::isVirtualRegister(Reg)) {
       if (Reg == AMDGPU::EXEC)
         continue;
 
       for (MCRegUnitIterator RegUnit(Reg, TRI); RegUnit.isValid(); ++RegUnit) {
         LiveRange &LR = LIS->getRegUnit(*RegUnit);
         const VNInfo *Value = LR.Query(LIS->getInstructionIndex(MI)).valueIn();
         if (!Value)
           continue;
 
-        // Since we're in machine SSA, we do not need to track physical
-        // registers across basic blocks.
+        // We do not need to track physical registers across basic blocks
+        // before register allocation
         if (Value->isPHIDef())
           continue;
 
         markInstruction(*LIS->getInstructionFromIndex(Value->def), StateWQM,
                         Worklist);
       }
 
       continue;
     }
 
-    for (MachineInstr &DefMI : MRI->def_instructions(Use.getReg()))
-      markInstruction(DefMI, StateWQM, Worklist);
+    unsigned SubReg = Use.getSubReg();
+    LaneBitmask LaneMask = SubReg ? TRI->getSubRegIndexLaneMask(SubReg)
+                                  : MRI->getMaxLaneMaskForVReg(Reg);
+
+    markValueWQM(Idx, Reg, LaneMask, Worklist);
   }
 }
 
 // Scan instructions to determine which ones require an Exact execmask and
 // which ones seed WQM requirements.
 char SIWholeQuadMode::scanInstructions(MachineFunction &MF,
                                        std::vector<WorkItem> &Worklist) {
   char GlobalFlags = 0;
   bool WQMOutputs = MF.getFunction()->hasFnAttribute("amdgpu-ps-wqm-outputs");
 
   for (auto BI = MF.begin(), BE = MF.end(); BI != BE; ++BI) {
     MachineBasicBlock &MBB = *BI;
 
     for (auto II = MBB.begin(), IE = MBB.end(); II != IE; ++II) {
       MachineInstr &MI = *II;
       unsigned Opcode = MI.getOpcode();
-      char Flags = 0;
 
       if (TII->isDS(Opcode)) {
-        Flags = StateWQM;
+        markInstruction(MI, StateWQM, Worklist);
+        GlobalFlags |= StateWQM;
       } else if (TII->isWQM(Opcode)) {
         // Sampling instructions don't need to produce results for all pixels
         // in a quad, they just require all inputs of a quad to have been
         // computed for derivatives.
         markUsesWQM(MI, Worklist);
         GlobalFlags |= StateWQM;
-        continue;
       } else if (TII->isDisableWQM(MI)) {
-        Flags = StateExact;
+        markInstruction(MI, StateExact, Worklist);
+        GlobalFlags |= StateExact;
       } else {
         if (Opcode == AMDGPU::SI_PS_LIVE) {
           LiveMaskQueries.push_back(&MI);
         } else if (WQMOutputs) {
-          // The function is in machine SSA form, which means that physical
-          // VGPRs correspond to shader inputs and outputs. Inputs are
+          // Physical VGPRs correspond to shader inputs and outputs. Inputs are
           // only used, outputs are only defined.
           for (const MachineOperand &MO : MI.defs()) {
             if (!MO.isReg())
               continue;
 
             unsigned Reg = MO.getReg();
 
             if (!TRI->isVirtualRegister(Reg) &&
                 TRI->hasVGPRs(TRI->getPhysRegClass(Reg))) {
-              Flags = StateWQM;
+              markInstruction(MI, StateWQM, Worklist);
+              GlobalFlags |= StateWQM;
               break;
             }
           }
         }
-
-        if (!Flags)
-          continue;
       }
-
-      markInstruction(MI, Flags, Worklist);
-      GlobalFlags |= Flags;
     }
   }
 
   return GlobalFlags;
 }
 
-void SIWholeQuadMode::propagateInstruction(MachineInstr &MI,
-                                           std::vector<WorkItem>& Worklist) {
-  MachineBasicBlock *MBB = MI.getParent();
-  InstrInfo II = Instructions[&MI]; // take a copy to prevent dangling references
-  BlockInfo &BI = Blocks[MBB];
+/// Helper function of \ref propagateValue that handles individual subranges.
+void SIWholeQuadMode::propagateValueSub(const WorkItem::Value &V,
+                                        LaneBitmask LaneMask,
+                                        const LiveRange &LR,
+                                        std::vector<WorkItem> &Worklist) {
+  const VNInfo *Value = LR.Query(V.Slot).valueIn();
+  if (!Value)
+    return;
+
+  if (Value->isPHIDef()) {
+    MachineBasicBlock *MBB = LIS->getMBBFromIndex(Value->def);
+
+    for (MachineBasicBlock *Pred : MBB->predecessors()) {
+      SlotIndex PredIndex = LIS->getMBBEndIdx(Pred).getPrevIndex();
 
-  // Control flow-type instructions and stores to temporary memory that are
-  // followed by WQM computations must themselves be in WQM.
-  if ((II.OutNeeds & StateWQM) && !II.Needs &&
-      (MI.isTerminator() || (TII->usesVM_CNT(MI) && MI.mayStore()))) {
-    Instructions[&MI].Needs = StateWQM;
-    II.Needs = StateWQM;
+      markValueWQM(PredIndex, V.Reg, LaneMask, Worklist);
     }
+  } else {
+    MachineInstr *DefMI = LIS->getInstructionFromIndex(Value->def);
 
-  // Propagate to block level
-  BI.Needs |= II.Needs;
-  if ((BI.InNeeds | II.Needs) != BI.InNeeds) {
-    BI.InNeeds |= II.Needs;
-    Worklist.push_back(MBB);
-  }
+    markInstruction(*DefMI, StateWQM, Worklist);
+  }
+}
+
+void SIWholeQuadMode::propagateValue(const WorkItem::Value &V,
+                                     std::vector<WorkItem> &Worklist) {
+  const LiveInterval &LI = LIS->getInterval(V.Reg);
+  LaneBitmask LaneMask = V.LaneMask;
+
+  if (LI.hasSubRanges()) {
+    for (const LiveInterval::SubRange &S : LI.subranges()) {
+      LaneBitmask Common = LaneMask & S.LaneMask;
+
+      if (Common) {
+        LaneMask &= ~Common;
+        propagateValueSub(V, Common, S, Worklist);
+      }
+    }
+  } else {
+    propagateValueSub(V, LaneMask, LI, Worklist);
+  }
 
+// Backwards propagation of OutNeeds and related effects.
+void SIWholeQuadMode::propagateInstruction(MachineInstr &MI,
+                                           std::vector<WorkItem> &Worklist) {
+  const InstrInfo II = Instructions[&MI];
+
+  // Stores to temporary memory (i.e., not marked as Exact) may be relevant for
+  // WQM computations.
+  if (!II.Needs && (II.OutNeeds & StateWQM) &&
+      (TII->usesVM_CNT(MI) && MI.mayStore()))
+    markInstruction(MI, StateWQM, Worklist);
+
   // Propagate backwards within block
   if (MachineInstr *PrevMI = MI.getPrevNode()) {
-    char InNeeds = II.Needs | II.OutNeeds;
-    if (!PrevMI->isPHI()) {
-      InstrInfo &PrevII = Instructions[PrevMI];
-      if ((PrevII.OutNeeds | InNeeds) != PrevII.OutNeeds) {
-        PrevII.OutNeeds |= InNeeds;
-        Worklist.push_back(PrevMI);
-      }
-    }
-  }
+    if (II.OutNeeds & ~PrevII.OutNeeds) {
+      PrevII.OutNeeds |= II.OutNeeds;
 
-  // Propagate WQM flag to instruction inputs
-  assert(II.Needs != (StateWQM | StateExact));
-
-  if (II.Needs == StateWQM)
-    markUsesWQM(MI, Worklist);
-}
-
 void SIWholeQuadMode::propagateBlock(MachineBasicBlock &MBB,
-                                     std::vector<WorkItem>& Worklist) {
+                                     std::vector<WorkItem> &Worklist) {
   BlockInfo BI = Blocks[&MBB]; // Make a copy to prevent dangling references.
 
   // Propagate through instructions
   if (!MBB.empty()) {
     MachineInstr *LastMI = &*MBB.rbegin();
     InstrInfo &LastII = Instructions[LastMI];
-    if ((LastII.OutNeeds | BI.OutNeeds) != LastII.OutNeeds) {
-      LastII.OutNeeds |= BI.OutNeeds;
+    if (BI.Needs.Out & ~LastII.OutNeeds) {
+      LastII.OutNeeds |= BI.Needs.Out;
       Worklist.push_back(LastMI);
     }
   }
 
   // Predecessor blocks must provide for our WQM/Exact needs.
+  if (BI.Needs.In & ~BI.Propagated.In) {
     for (MachineBasicBlock *Pred : MBB.predecessors()) {
       BlockInfo &PredBI = Blocks[Pred];
-    if ((PredBI.OutNeeds | BI.InNeeds) == PredBI.OutNeeds)
-      continue;
-
-    PredBI.OutNeeds |= BI.InNeeds;
-    PredBI.InNeeds |= BI.InNeeds;
+      if (BI.Needs.In & ~PredBI.Needs.Out) {
+        PredBI.Needs.Out |= BI.Needs.In;
+        PredBI.Needs.In |= BI.Needs.In;
         Worklist.push_back(Pred);
       }
+    }
+  }
 
-  // All successors must be prepared to accept the same set of WQM/Exact data.
+  if (BI.Needs.Out & ~BI.Propagated.Out) {
+    // All successors must be prepared to accept the same set of WQM/Exact
+    // data.
     for (MachineBasicBlock *Succ : MBB.successors()) {
       BlockInfo &SuccBI = Blocks[Succ];
-    if ((SuccBI.InNeeds | BI.OutNeeds) == SuccBI.InNeeds)
-      continue;
-
-    SuccBI.InNeeds |= BI.OutNeeds;
+      if (BI.Needs.Out & ~SuccBI.Needs.In) {
+        SuccBI.Needs.In |= BI.Needs.Out;
         Worklist.push_back(Succ);
       }
     }
 
+    // Mark terminators as WQM if required
+    if (BI.Needs.Out & ~BI.Propagated.Out & StateWQM) {
+      for (MachineInstr &Terminator : MBB.terminators())
+        markInstruction(Terminator, StateWQM, Worklist);
+    }
+  }
+
+  Blocks[&MBB].Propagated = BI.Needs;
+}
+
 char SIWholeQuadMode::analyzeFunction(MachineFunction &MF) {
   std::vector<WorkItem> Worklist;
   char GlobalFlags = scanInstructions(MF, Worklist);
 
   while (!Worklist.empty()) {
     WorkItem WI = Worklist.back();
     Worklist.pop_back();
 
     if (WI.MI)
       propagateInstruction(*WI.MI, Worklist);
-    else
+    else if (WI.MBB)
       propagateBlock(*WI.MBB, Worklist);
+    else
+      propagateValue(WI.V, Worklist);
   }
 
   return GlobalFlags;
 }
 
 void SIWholeQuadMode::toExact(MachineBasicBlock &MBB,
                               MachineBasicBlock::iterator Before,
                               unsigned SaveWQM, unsigned LiveMaskReg) {
+  MachineInstr *MI;
+
   if (SaveWQM) {
-    BuildMI(MBB, Before, DebugLoc(), TII->get(AMDGPU::S_AND_SAVEEXEC_B64),
+    MI = BuildMI(MBB, Before, DebugLoc(), TII->get(AMDGPU::S_AND_SAVEEXEC_B64),
                  SaveWQM)
              .addReg(LiveMaskReg);
   } else {
-    BuildMI(MBB, Before, DebugLoc(), TII->get(AMDGPU::S_AND_B64),
+    MI = BuildMI(MBB, Before, DebugLoc(), TII->get(AMDGPU::S_AND_B64),
                  AMDGPU::EXEC)
              .addReg(AMDGPU::EXEC)
              .addReg(LiveMaskReg);
   }
+
+  LIS->InsertMachineInstrInMaps(*MI);
 }
 
 void SIWholeQuadMode::toWQM(MachineBasicBlock &MBB,
                             MachineBasicBlock::iterator Before,
                             unsigned SavedWQM) {
+  MachineInstr *MI;
+
   if (SavedWQM) {
-    BuildMI(MBB, Before, DebugLoc(), TII->get(AMDGPU::COPY), AMDGPU::EXEC)
+    MI = BuildMI(MBB, Before, DebugLoc(), TII->get(AMDGPU::COPY), AMDGPU::EXEC)
              .addReg(SavedWQM);
   } else {
-    BuildMI(MBB, Before, DebugLoc(), TII->get(AMDGPU::S_WQM_B64),
+    MI = BuildMI(MBB, Before, DebugLoc(), TII->get(AMDGPU::S_WQM_B64),
                  AMDGPU::EXEC)
              .addReg(AMDGPU::EXEC);
   }
+
+  LIS->InsertMachineInstrInMaps(*MI);
 }
 
 void SIWholeQuadMode::processBlock(MachineBasicBlock &MBB, unsigned LiveMaskReg,
                                    bool isEntry) {
   auto BII = Blocks.find(&MBB);
   if (BII == Blocks.end())
     return;
 
   const BlockInfo &BI = BII->second;
 
-  if (!(BI.InNeeds & StateWQM))
+  if (!(BI.Needs.In & StateWQM))
     return;
 
   // This is a non-entry block that is WQM throughout, so no need to do
   // anything.
-  if (!isEntry && !(BI.Needs & StateExact) && BI.OutNeeds != StateExact)
+  if (!isEntry && !(BI.Needs.Self & StateExact) && BI.Needs.Out != StateExact)
     return;
 
   DEBUG(dbgs() << "\nProcessing block BB#" << MBB.getNumber() << ":\n");
 
   unsigned SavedWQMReg = 0;
   bool WQMFromExec = isEntry;
   char State = isEntry ? StateExact : StateWQM;
   MachineInstr *FirstNonWQM = nullptr;
if (Needs && State != Needs) {
 
         if (!WQMFromExec && (OutNeeds & StateWQM))
           SavedWQMReg = MRI->createVirtualRegister(&AMDGPU::SReg_64RegClass);
 
         toExact(MBB, FirstNonWQM, SavedWQMReg, LiveMaskReg);
       } else {
         assert(WQMFromExec == (SavedWQMReg == 0));
         toWQM(MBB, &MI, SavedWQMReg);
+
+        if (SavedWQMReg) {
+          LIS->createAndComputeVirtRegInterval(SavedWQMReg);
           SavedWQMReg = 0;
         }
+      }
 
       State = Needs;
     }
 
-    if (MI.getOpcode() == AMDGPU::SI_ELSE && BI.OutNeeds == StateExact)
+    if (MI.getOpcode() == AMDGPU::SI_ELSE && BI.Needs.Out == StateExact)
       MI.getOperand(3).setImm(1);
   }
 
-  if ((BI.OutNeeds & StateWQM) && State != StateWQM) {
+  if ((BI.Needs.Out & StateWQM) && State != StateWQM) {
     assert(WQMFromExec == (SavedWQMReg == 0));
     toWQM(MBB, MBB.end(), SavedWQMReg);
-  } else if (BI.OutNeeds == StateExact && State != StateExact) {
+
+    if (SavedWQMReg)
+      LIS->createAndComputeVirtRegInterval(SavedWQMReg);
+  } else if (BI.Needs.Out == StateExact && State != StateExact) {
     toExact(MBB, FirstNonWQM ? MachineBasicBlock::iterator(FirstNonWQM)
                              : MBB.getFirstTerminator(),
             0, LiveMaskReg);
   }
 }
 
 void SIWholeQuadMode::lowerLiveMaskQueries(unsigned LiveMaskReg) {
   for (MachineInstr *MI : LiveMaskQueries) {
     const DebugLoc &DL = MI->getDebugLoc();
     unsigned Dest = MI->getOperand(0).getReg();
+    MachineInstr *NewMI =
         BuildMI(*MI->getParent(), MI, DL, TII->get(AMDGPU::COPY), Dest)
             .addReg(LiveMaskReg);
+    LIS->ReplaceMachineInstrInMaps(*MI, *NewMI);
+
     MI->eraseFromParent();
   }
 }
 
 bool SIWholeQuadMode::runOnMachineFunction(MachineFunction &MF) {
   if (MF.getFunction()->getCallingConv() != CallingConv::AMDGPU_PS)
     return false;
 
   Instructions.clear();
+  WQMValues.clear();
   Blocks.clear();
   LiveMaskQueries.clear();
 
   const SISubtarget &ST = MF.getSubtarget<SISubtarget>();
 
   TII = ST.getInstrInfo();
   TRI = &TII->getRegisterInfo();
   MRI = &MF.getRegInfo();
   LIS = &getAnalysis<LiveIntervals>();
 
   char GlobalFlags = analyzeFunction(MF);
   if (!(GlobalFlags & StateWQM)) {
     lowerLiveMaskQueries(AMDGPU::EXEC);
     return !LiveMaskQueries.empty();
   }
 
+  DEBUG(printInfo(MF));
+
   // Store a copy of the original live mask when required
   unsigned LiveMaskReg = 0;
   {
     MachineBasicBlock &Entry = MF.front();
     MachineBasicBlock::iterator EntryMI = Entry.getFirstNonPHI();
 
     if (GlobalFlags & StateExact || !LiveMaskQueries.empty()) {
       LiveMaskReg = MRI->createVirtualRegister(&AMDGPU::SReg_64RegClass);
-      BuildMI(Entry, EntryMI, DebugLoc(), TII->get(AMDGPU::COPY), LiveMaskReg)
+      MachineInstr *MI = BuildMI(Entry, EntryMI, DebugLoc(),
+                                 TII->get(AMDGPU::COPY), LiveMaskReg)
                              .addReg(AMDGPU::EXEC);
+      LIS->InsertMachineInstrInMaps(*MI);
     }
 
     if (GlobalFlags == StateWQM) {
       // For a shader that needs only WQM, we can just set it once.
-      BuildMI(Entry, EntryMI, DebugLoc(), TII->get(AMDGPU::S_WQM_B64),
-              AMDGPU::EXEC)
+      MachineInstr *MI = BuildMI(Entry, EntryMI, DebugLoc(),
+                                 TII->get(AMDGPU::S_WQM_B64), AMDGPU::EXEC)
                              .addReg(AMDGPU::EXEC);
-
-      lowerLiveMaskQueries(LiveMaskReg);
-      // EntryMI may become invalid here
-      return true;
+      LIS->InsertMachineInstrInMaps(*MI);
     }
   }
 
-  DEBUG(printInfo());
-
   lowerLiveMaskQueries(LiveMaskReg);
 
+  if (GlobalFlags != StateWQM) {
     // Handle the general case
-  for (auto BII : Blocks)
+    for (auto &BII : Blocks)
       processBlock(*BII.first, LiveMaskReg, BII.first == &*MF.begin());
+  }
+
+  if (LiveMaskReg)
+    LIS->createAndComputeVirtRegInterval(LiveMaskReg);
 
   return true;
 }

test/CodeGen/AMDGPU/llvm.amdgcn.ps.live.ll

 ; RUN: llc -march=amdgcn -mcpu=SI -verify-machineinstrs < %s | FileCheck -check-prefix=CHECK %s
 
 ; CHECK-LABEL: {{^}}test1:
-; CHECK: v_cndmask_b32_e64 v0, 0, 1, exec
+; CHECK: s_mov_b64 [[LIVE:s\[[0-9]+:[0-9]+\]]], exec
+; CHECK: v_cndmask_b32_e64 v0, 0, 1, [[LIVE]]
 ;
 ; Note: We could generate better code here if we recognized earlier that
 ; there is no WQM use and therefore llvm.amdgcn.ps.live is constant. However,
 ; the expectation is that the intrinsic will be used in non-trivial shaders,
 ; so such an optimization doesn't seem worth the effort.
 define amdgpu_ps float @test1() {
   %live = call i1 @llvm.amdgcn.ps.live()
   %live.32 = zext i1 %live to i32
   %r = bitcast i32 %live.32 to float
   ret float %r
 }
 
 ; CHECK-LABEL: {{^}}test2:
 ; CHECK: s_mov_b64 [[LIVE:s\[[0-9]+:[0-9]+\]]], exec
+; CHECK-DAG: s_mov_b64 [[COPY:s\[[0-9]+:[0-9]+\]]], [[LIVE]]
 ; CHECK-DAG: s_wqm_b64 exec, exec
-; CHECK-DAG: v_cndmask_b32_e64 [[VAR:v[0-9]+]], 0, 1, [[LIVE]]
+; CHECK-DAG: v_cndmask_b32_e64 [[VAR:v[0-9]+]], 0, 1, [[COPY]]
 ; CHECK: image_sample v0, [[VAR]],
 define amdgpu_ps float @test2() {
   %live = call i1 @llvm.amdgcn.ps.live()
   %live.32 = zext i1 %live to i32
 
   %t = call <4 x float> @llvm.SI.image.sample.i32(i32 %live.32, <8 x i32> undef, <4 x i32> undef, i32 15, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0)
 
   %r = extractelement <4 x float> %t, i32 0
   ret float %r
 }
 
 ; CHECK-LABEL: {{^}}test3:
 ; CHECK: s_mov_b64 [[LIVE:s\[[0-9]+:[0-9]+\]]], exec
+; CHECK-DAG: s_mov_b64 [[COPY:s\[[0-9]+:[0-9]+\]]], [[LIVE]]
 ; CHECK-DAG: s_wqm_b64 exec, exec
-; CHECK-DAG: s_xor_b64 [[HELPER:s\[[0-9]+:[0-9]+\]]], [[LIVE]], -1
+; CHECK-DAG: s_xor_b64 [[HELPER:s\[[0-9]+:[0-9]+\]]], [[COPY]], -1
 ; CHECK_DAG: s_and_saveexec_b64 [[SAVED:s\[[0-9]+:[0-9]+\]]], [[HELPER]]
 ; CHECK: ; %dead
 define amdgpu_ps float @test3(i32 %in) {
 entry:
   %live = call i1 @llvm.amdgcn.ps.live()
   br i1 %live, label %end, label %dead
 
 dead:

test/CodeGen/AMDGPU/wqm.ll

 
 if:
   unreachable
 
 else:
   ret <4 x float> %dtex
 }
 
+; CHECK-LABEL: {{^}}test_subregs:
+; CHECK: s_mov_b64 [[ORIG:s\[[0-9]+:[0-9]+\]]], exec
+; CHECK: s_wqm_b64 exec, exec
+; CHECK: v_interp_p1_f32
+; CHECK: v_interp_p2_f32
+; CHECK: s_and_b64 exec, exec, [[ORIG]]
+; CHECK: _store
+; CHECK: s_wqm_b64 exec, exec
+; CHECK: image_sample
+;
+; Early coalescing merges %c into a 64 bit VGPR pair, so correctness requires
+; tracking of subregisters.
+;
+define amdgpu_ps <4 x float> @test_subregs(i32 inreg %prims, <2 x i32> %ij, i32 %idx) #1 {
+main_body:
+  %c = call float @llvm.SI.fs.interp(i32 0, i32 0, i32 %prims, <2 x i32> %ij)
+
+  call void @llvm.amdgcn.buffer.store.f32(float 1.0, <4 x i32> undef, i32 %idx, i32 0, i1 0, i1 0)
+
+  %c.i = bitcast float %c to i32
+  %c2.0 = insertelement <2 x i32> undef, i32 %c.i, i32 0
+  %c2.i = insertelement <2 x i32> %c2.0, i32 1, i32 1
+  %tex = call <4 x float> @llvm.SI.image.sample.v2i32(<2 x i32> %c2.i, <8 x i32> undef, <4 x i32> undef, i32 15, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0)
+  ret <4 x float> %tex
+}
+
+; Test tracking of vector condition codes.
+;
+; CHECK-LABEL: {{^}}test_vcc_tracking:
+; CHECK: s_mov_b64 [[ORIG:s\[[0-9]+:[0-9]+\]]], exec
+; CHECK: s_wqm_b64 exec, exec
+; CHECK: v_cmp_gt_i32_e32 vcc,
+; CHECK: ; %else
+; CHECK: image_sample
+; CHECK: ; %if
+; CHECK: image_sample
+; CHECK: ; %end
+; CHECK: s_and_b64 exec, exec, [[ORIG]]
+; CHECK: _store
+define amdgpu_ps <4 x float> @test_vcc_tracking(i32 %sel, i32 %idx) #1 {
+main_body:
+  %cc = icmp sgt i32 %sel, 0
+  br i1 %cc, label %if, label %else
+
+if:
+  %r.if = call <4 x float> @llvm.SI.image.sample.i32(i32 0, <8 x i32> undef, <4 x i32> undef, i32 15, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0)
+  br label %end
+
+else:
+  %r.else = call <4 x float> @llvm.SI.image.sample.v2i32(<2 x i32> <i32 0, i32 1>, <8 x i32> undef, <4 x i32> undef, i32 15, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0)
+  br label %end
+
+end:
+  %r = phi <4 x float> [ %r.if, %if ], [ %r.else, %else ]
+
+  call void @llvm.amdgcn.buffer.store.f32(float 1.0, <4 x i32> undef, i32 %idx, i32 0, i1 0, i1 0)
+
+  ret <4 x float> %r
+}
+
+
 declare void @llvm.amdgcn.image.store.v4i32(<4 x float>, <4 x i32>, <8 x i32>, i32, i1, i1, i1, i1) #1
 declare void @llvm.amdgcn.buffer.store.f32(float, <4 x i32>, i32, i32, i1, i1) #1
 declare void @llvm.amdgcn.buffer.store.v4f32(<4 x float>, <4 x i32>, i32, i32, i1, i1) #1
 
 declare <4 x float> @llvm.amdgcn.image.load.v4i32(<4 x i32>, <8 x i32>, i32, i1, i1, i1, i1) #2
 declare float @llvm.amdgcn.buffer.load.f32(<4 x i32>, i32, i32, i1, i1) #2
 
 declare <4 x float> @llvm.SI.image.sample.i32(i32, <8 x i32>, <4 x i32>, i32, i32, i32, i32, i32, i32, i32, i32) #3
+declare <4 x float> @llvm.SI.image.sample.v2i32(<2 x i32>, <8 x i32>, <4 x i32>, i32, i32, i32, i32, i32, i32, i32, i32) #3
 declare <4 x float> @llvm.SI.image.sample.v4i32(<4 x i32>, <8 x i32>, <4 x i32>, i32, i32, i32, i32, i32, i32, i32, i32) #3
 
+declare float @llvm.SI.fs.interp(i32, i32, i32, <2 x i32>) #3
 declare void @llvm.AMDGPU.kill(float)
 declare void @llvm.SI.export(i32, i32, i32, i32, i32, float, float, float, float)
 
 attributes #1 = { nounwind }
 attributes #2 = { nounwind readonly }
 attributes #3 = { nounwind readnone }
 attributes #4 = { "amdgpu-ps-wqm-outputs" }