[AMDGPU] SDWA: make pass global

[SamWot]

Summary

Remove checks for basic blocks.

[rampitec]

A virtual register can have multiple definitions. I think you need to make sure you are using a right one. Can you please forge a test like:

if (...) vregx = some_sdwa_cand;
else vregx = other_sdwa_cand;
use vregx;

and see what is folded at use?

You can conservatively use getUniqueVRegDef to make sure there is only one def, or make sure which def is really used. If you use getUniqueVRegDef it may decrease even current optimization scope since in many cases you have subreg candidates, which have multiple defs.

[arsenm]

In D31935#723907, @rampitec wrote:

A virtual register can have multiple definitions. I think you need to make sure you are using a right one. Can you please forge a test like:

if (...) vregx = some_sdwa_cand;
else vregx = other_sdwa_cand;
use vregx;

and see what is folded at use?

You can conservatively use getUniqueVRegDef to make sure there is only one def, or make sure which def is really used. If you use getUniqueVRegDef it may decrease even current optimization scope since in many cases you have subreg candidates, which have multiple defs.

This is an SSA pass, there is only one reg def. Using getUniqueVRegDef is not necessary

[rampitec]

In D31935#723964, @arsenm wrote:

This is an SSA pass, there is only one reg def. Using getUniqueVRegDef is not necessary

Missed that, thanks.

[rampitec]

LGTM

llvm/trunk/lib/Target/AMDGPU/SIPeepholeSDWA.cpp

public:
 
   typedef SmallVector<std::unique_ptr<SDWAOperand>, 4> SDWAOperandsVector;
 
   SIPeepholeSDWA() : MachineFunctionPass(ID) {
     initializeSIPeepholeSDWAPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override;
-  void matchSDWAOperands(MachineBasicBlock &MBB);
+  void matchSDWAOperands(MachineFunction &MF);
   bool convertToSDWA(MachineInstr &MI, const SDWAOperandsVector &SDWAOperands);
 
   StringRef getPassName() const override { return "SI Peephole SDWA"; }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesCFG();
     MachineFunctionPass::getAnalysisUsage(AU);
   }
static raw_ostream& operator<<(raw_ostream &OS, const SDWADstOperand &Dst) {
   OS << "SDWA dst: " << *Dst.getTargetOperand()
      << " dst_sel:" << Dst.getDstSel()
      << " dst_unused:" << Dst.getDstUnused() << '\n';
   return OS;
 }
 
 #endif
 
-static bool isSameBB(const MachineInstr *FirstMI, const MachineInstr *SecondMI) {
-  assert(FirstMI && SecondMI);
-  return FirstMI->getParent() == SecondMI->getParent();
-}
-
 static void copyRegOperand(MachineOperand &To, const MachineOperand &From) {
   assert(To.isReg() && From.isReg());
   To.setReg(From.getReg());
   To.setSubReg(From.getSubReg());
   To.setIsUndef(From.isUndef());
   if (To.isUse()) {
     To.setIsKill(From.isKill());
   } else {
MachineInstr *SDWASrcOperand::potentialToConvert(const SIInstrInfo *TII) {
   assert(Replaced->isReg());
 
   MachineInstr *PotentialMI = nullptr;
   for (MachineOperand &PotentialMO : MRI->use_operands(Replaced->getReg())) {
     // If this is use of another subreg of dst reg then do nothing
     if (!isSubregOf(*Replaced, PotentialMO, MRI->getTargetRegisterInfo()))
       continue;
 
-    // If there exist use of dst in another basic block or use of superreg of
-    // dst then we should not combine this opernad
-    if (!isSameBB(PotentialMO.getParent(), getParentInst()) ||
-        !isSameReg(PotentialMO, *Replaced))
+    // If there exist use of superreg of dst then we should not combine this
+    // opernad
+    if (!isSameReg(PotentialMO, *Replaced))
       return nullptr;
 
     // Check that PotentialMI is only instruction that uses dst reg
     if (PotentialMI == nullptr) {
       PotentialMI = PotentialMO.getParent();
     } else if (PotentialMI != PotentialMO.getParent()) {
       return nullptr;
     }
MachineInstr *SDWADstOperand::potentialToConvert(const SIInstrInfo *TII) {
   MachineInstr *ParentMI = getParentInst();
   MachineOperand *Replaced = getReplacedOperand();
   assert(Replaced->isReg());
 
   for (MachineOperand &PotentialMO : MRI->def_operands(Replaced->getReg())) {
     if (!isSubregOf(*Replaced, PotentialMO, MRI->getTargetRegisterInfo()))
       continue;
 
-    if (!isSameBB(getParentInst(), PotentialMO.getParent()) ||
-        !isSameReg(*Replaced, PotentialMO))
+    if (!isSameReg(*Replaced, PotentialMO))
       return nullptr;
 
     // Check that ParentMI is the only instruction that uses replaced register
     for (MachineOperand &UseMO : MRI->use_operands(PotentialMO.getReg())) {
       if (isSubregOf(UseMO, PotentialMO, MRI->getTargetRegisterInfo()) &&
           UseMO.getParent() != ParentMI) {
         return nullptr;
       }
Optional<int64_t> SIPeepholeSDWA::foldToImm(const MachineOperand &Op) const {
   // If this is not immediate then it can be copy of immediate value, e.g.:
   // %vreg1<def> = S_MOV_B32 255;
   if (Op.isReg()) {
     for (const MachineOperand &Def : MRI->def_operands(Op.getReg())) {
       if (!isSameReg(Op, Def))
         continue;
 
       const MachineInstr *DefInst = Def.getParent();
-      if (!TII->isFoldableCopy(*DefInst) || !isSameBB(Op.getParent(), DefInst))
+      if (!TII->isFoldableCopy(*DefInst))
         return None;
 
       const MachineOperand &Copied = DefInst->getOperand(1);
       if (!Copied.isImm())
         return None;
 
       return Copied.getImm();
     }
   }
 
   return None;
 }
 
-void SIPeepholeSDWA::matchSDWAOperands(MachineBasicBlock &MBB) {
+void SIPeepholeSDWA::matchSDWAOperands(MachineFunction &MF) {
+  for (MachineBasicBlock &MBB : MF) {
     for (MachineInstr &MI : MBB) {
       unsigned Opcode = MI.getOpcode();
       switch (Opcode) {
       case AMDGPU::V_LSHRREV_B32_e32:
       case AMDGPU::V_ASHRREV_I32_e32:
       case AMDGPU::V_LSHLREV_B32_e32: {
         // from: v_lshrrev_b32_e32 v1, 16/24, v0
         // to SDWA src:v0 src_sel:WORD_1/BYTE_3
 
         // from: v_ashrrev_i32_e32 v1, 16/24, v0
         // to SDWA src:v0 src_sel:WORD_1/BYTE_3 sext:1
 
         // from: v_lshlrev_b32_e32 v1, 16/24, v0
         // to SDWA dst:v1 dst_sel:WORD_1/BYTE_3 dst_unused:UNUSED_PAD
         MachineOperand *Src0 = TII->getNamedOperand(MI, AMDGPU::OpName::src0);
         auto Imm = foldToImm(*Src0);
         if (!Imm)
           break;
 
         if (*Imm != 16 && *Imm != 24)
           break;
 
         MachineOperand *Src1 = TII->getNamedOperand(MI, AMDGPU::OpName::src1);
         MachineOperand *Dst = TII->getNamedOperand(MI, AMDGPU::OpName::vdst);
         if (TRI->isPhysicalRegister(Src1->getReg()) ||
             TRI->isPhysicalRegister(Dst->getReg()))
           break;
 
         if (Opcode == AMDGPU::V_LSHLREV_B32_e32) {
           auto SDWADst = make_unique<SDWADstOperand>(
               Dst, Src1, *Imm == 16 ? WORD_1 : BYTE_3, UNUSED_PAD);
           DEBUG(dbgs() << "Match: " << MI << "To: " << *SDWADst << '\n');
           SDWAOperands[&MI] = std::move(SDWADst);
           ++NumSDWAPatternsFound;
         } else {
           auto SDWASrc = make_unique<SDWASrcOperand>(
               Src1, Dst, *Imm == 16 ? WORD_1 : BYTE_3, false, false,
               Opcode == AMDGPU::V_LSHRREV_B32_e32 ? false : true);
           DEBUG(dbgs() << "Match: " << MI << "To: " << *SDWASrc << '\n');
           SDWAOperands[&MI] = std::move(SDWASrc);
           ++NumSDWAPatternsFound;
         }
         break;
       }
 
       case AMDGPU::V_LSHRREV_B16_e32:
       case AMDGPU::V_ASHRREV_I16_e32:
       case AMDGPU::V_LSHLREV_B16_e32: {
         // from: v_lshrrev_b16_e32 v1, 8, v0
         // to SDWA src:v0 src_sel:BYTE_1
 
         // from: v_ashrrev_i16_e32 v1, 8, v0
         // to SDWA src:v0 src_sel:BYTE_1 sext:1
 
         // from: v_lshlrev_b16_e32 v1, 8, v0
         // to SDWA dst:v1 dst_sel:BYTE_1 dst_unused:UNUSED_PAD
         MachineOperand *Src0 = TII->getNamedOperand(MI, AMDGPU::OpName::src0);
         auto Imm = foldToImm(*Src0);
         if (!Imm || *Imm != 8)
           break;
 
         MachineOperand *Src1 = TII->getNamedOperand(MI, AMDGPU::OpName::src1);
         MachineOperand *Dst = TII->getNamedOperand(MI, AMDGPU::OpName::vdst);
 
         if (TRI->isPhysicalRegister(Src1->getReg()) ||
             TRI->isPhysicalRegister(Dst->getReg()))
           break;
 
         if (Opcode == AMDGPU::V_LSHLREV_B16_e32) {
           auto SDWADst =
               make_unique<SDWADstOperand>(Dst, Src1, BYTE_1, UNUSED_PAD);
           DEBUG(dbgs() << "Match: " << MI << "To: " << *SDWADst << '\n');
           SDWAOperands[&MI] = std::move(SDWADst);
           ++NumSDWAPatternsFound;
         } else {
           auto SDWASrc = make_unique<SDWASrcOperand>(
               Src1, Dst, BYTE_1, false, false,
               Opcode == AMDGPU::V_LSHRREV_B16_e32 ? false : true);
           DEBUG(dbgs() << "Match: " << MI << "To: " << *SDWASrc << '\n');
           SDWAOperands[&MI] = std::move(SDWASrc);
           ++NumSDWAPatternsFound;
         }
         break;
       }
 
       case AMDGPU::V_BFE_I32:
       case AMDGPU::V_BFE_U32: {
         // e.g.:
         // from: v_bfe_u32 v1, v0, 8, 8
         // to SDWA src:v0 src_sel:BYTE_1
 
         // offset | width | src_sel
         // ------------------------
         // 0      | 8     | BYTE_0
         // 0      | 16    | WORD_0
         // 0      | 32    | DWORD ?
         // 8      | 8     | BYTE_1
         // 16     | 8     | BYTE_2
         // 16     | 16    | WORD_1
         // 24     | 8     | BYTE_3
 
         MachineOperand *Src1 = TII->getNamedOperand(MI, AMDGPU::OpName::src1);
         auto Offset = foldToImm(*Src1);
         if (!Offset)
           break;
 
         MachineOperand *Src2 = TII->getNamedOperand(MI, AMDGPU::OpName::src2);
         auto Width = foldToImm(*Src2);
         if (!Width)
           break;
 
         SdwaSel SrcSel = DWORD;
 
         if (*Offset == 0 && *Width == 8)
           SrcSel = BYTE_0;
         else if (*Offset == 0 && *Width == 16)
           SrcSel = WORD_0;
         else if (*Offset == 0 && *Width == 32)
           SrcSel = DWORD;
         else if (*Offset == 8 && *Width == 8)
           SrcSel = BYTE_1;
         else if (*Offset == 16 && *Width == 8)
           SrcSel = BYTE_2;
         else if (*Offset == 16 && *Width == 16)
           SrcSel = WORD_1;
         else if (*Offset == 24 && *Width == 8)
           SrcSel = BYTE_3;
         else
           break;
 
         MachineOperand *Src0 = TII->getNamedOperand(MI, AMDGPU::OpName::src0);
         MachineOperand *Dst = TII->getNamedOperand(MI, AMDGPU::OpName::vdst);
       
         if (TRI->isPhysicalRegister(Src0->getReg()) ||
             TRI->isPhysicalRegister(Dst->getReg()))
           break;
 
         auto SDWASrc = make_unique<SDWASrcOperand>(
             Src0, Dst, SrcSel, false, false,
             Opcode == AMDGPU::V_BFE_U32 ? false : true);
         DEBUG(dbgs() << "Match: " << MI << "To: " << *SDWASrc << '\n');
         SDWAOperands[&MI] = std::move(SDWASrc);
         ++NumSDWAPatternsFound;
         break;
       }
       case AMDGPU::V_AND_B32_e32: {
         // e.g.:
         // from: v_and_b32_e32 v1, 0x0000ffff/0x000000ff, v0
         // to SDWA src:v0 src_sel:WORD_0/BYTE_0
 
         MachineOperand *Src0 = TII->getNamedOperand(MI, AMDGPU::OpName::src0);
         auto Imm = foldToImm(*Src0);
         if (!Imm)
           break;
 
         if (*Imm != 0x0000ffff && *Imm != 0x000000ff)
           break;
 
         MachineOperand *Src1 = TII->getNamedOperand(MI, AMDGPU::OpName::src1);
         MachineOperand *Dst = TII->getNamedOperand(MI, AMDGPU::OpName::vdst);
       
         if (TRI->isPhysicalRegister(Src1->getReg()) ||
             TRI->isPhysicalRegister(Dst->getReg()))
           break;
 
         auto SDWASrc = make_unique<SDWASrcOperand>(
             Src1, Dst, *Imm == 0x0000ffff ? WORD_0 : BYTE_0);
         DEBUG(dbgs() << "Match: " << MI << "To: " << *SDWASrc << '\n');
         SDWAOperands[&MI] = std::move(SDWASrc);
         ++NumSDWAPatternsFound;
         break;
       }
       }
     }
   }
+}
 
 bool SIPeepholeSDWA::convertToSDWA(MachineInstr &MI,
                                    const SDWAOperandsVector &SDWAOperands) {
   // Check if this instruction can be converted to SDWA:
   // 1. Does this opcode support SDWA
   if (AMDGPU::getSDWAOp(MI.getOpcode()) == -1)
     return false;
 
bool SIPeepholeSDWA::runOnMachineFunction(MachineFunction &MF) {
   }
 
   MRI = &MF.getRegInfo();
   TRI = ST.getRegisterInfo();
   TII = ST.getInstrInfo();
 
   std::unordered_map<MachineInstr *, SDWAOperandsVector> PotentialMatches;
 
-  // FIXME: For now we only combine instructions in one basic block
-  for (MachineBasicBlock &MBB : MF) {
-    SDWAOperands.clear();
-    matchSDWAOperands(MBB);
+  matchSDWAOperands(MF);
 
-    PotentialMatches.clear();
-    for (auto &OperandPair : SDWAOperands) {
-      auto &Operand = OperandPair.second;
-      MachineInstr *PotentialMI = Operand->potentialToConvert(TII);
-      if (PotentialMI) {
-        PotentialMatches[PotentialMI].push_back(std::move(Operand));
-      }
-    }
 
-    for (auto &PotentialPair : PotentialMatches) {
-      MachineInstr &PotentialMI = *PotentialPair.first;
-      convertToSDWA(PotentialMI, PotentialPair.second);
-    }
-  }
+
+  SDWAOperands.clear();
   return false;
 }

llvm/trunk/test/CodeGen/AMDGPU/sdwa-peephole.ll

entry:
   %a = load <2 x i16>, <2 x i16> addrspace(1)* %ina, align 4
   %b = load <2 x i16>, <2 x i16> addrspace(1)* %inb, align 4
   %mul = mul <2 x i16> %a, %b
   %mul2 = mul <2 x i16> %mul, %b
   store <2 x i16> %mul2, <2 x i16> addrspace(1)* %out, align 4
   ret void
 }
 
-; GCN-LABEL: {{^}}mul_add_shr_i32:
-; NOSDWA-NOT: v_mul_u32_u24_sdwa
+; GCN-LABEL: {{^}}add_bb_v2i16:
 ; NOSDWA-NOT: v_add_i32_sdwa
-; SDWA-NOT: v_mul_u32_u24_sdwa
-; SDWA-NOT: v_add_i32_sdwa
 
-define void @mul_add_shr_i32(i32 addrspace(1)* %out, i32 addrspace(1)* %ina, i32 addrspace(1)* %inb, i1 addrspace(1)* %incond) {
+; SDWA: v_add_i32_sdwa v{{[0-9]+}}, vcc, v{{[0-9]+}}, v{{[0-9]+}} dst_sel:WORD_1 dst_unused:UNUSED_PAD src0_sel:WORD_1 src1_sel:WORD_1
+
+define amdgpu_kernel void @add_bb_v2i16(<2 x i16> addrspace(1)* %out, <2 x i16> addrspace(1)* %ina, <2 x i16> addrspace(1)* %inb) {
 entry:
-  %a = load i32, i32 addrspace(1)* %ina, align 4
-  %b = load i32, i32 addrspace(1)* %inb, align 4
-  %cond = load i1, i1 addrspace(1)* %incond, align 4
-  %shra = lshr i32 %a, 16
-  %shrb = lshr i32 %b, 16
-  br i1 %cond, label %mul_label, label %add_label
-mul_label:
-  %mul = mul i32 %shra, %shrb
-  br label %store_label
+  %a = load <2 x i16>, <2 x i16> addrspace(1)* %ina, align 4
+  %b = load <2 x i16>, <2 x i16> addrspace(1)* %inb, align 4
+  br label %add_label
 add_label:
-  %add = add i32 %shra, %shrb
+  %add = add <2 x i16> %a, %b
   br label %store_label
 store_label:
-  %store = phi i32 [%mul, %mul_label], [%add, %add_label]
-  store i32 %store, i32 addrspace(1)* %out, align 4
+  store <2 x i16> %add, <2 x i16> addrspace(1)* %out, align 4
   ret void
-}
+}