[AMDGPU] SDWA operands should not intersect with potential MIs

[SamWot]

Summary

There should be no intesection between SDWA operands and potential MIs. E.g.:

v_and_b32 v0, 0xff, v1 -> src:v1 sel:BYTE_0
v_and_b32 v2, 0xff, v0 -> src:v0 sel:BYTE_0
v_add_u32 v3, v4, v2

In that example it is possible that we would fold 2nd instruction into 3rd (v_add_u32_sdwa) and then try to fold 1st instruction into 2nd (that was already destroyed)
To solve this problem we keep track of every SDWA operand that should be "pulled out" - operand created from MI matched by another SDWA operand.

[arsenm]

lib/Target/AMDGPU/SIPeepholeSDWA.cpp

699

Capitalize, punctuate

715

Don't use shared_ptr

726

This management of the multiple maps seems overcomplicated to me.

Also, since these are only potential cases that could still fail to fold, shouldn't this be handled while folding the operands rather than eliminating them from the potential set in case they fold into other uses?

Shouldn't this be handled as part of the iteration order? If you made SDWAOperandsVector some kind of SetVector and changed the discovery order to post order you should see the final foldable uses first

test/CodeGen/AMDGPU/sdwa-peephole.ll

420–421

Run instnamer on the test

[SamWot]

Resolved some issues

[arsenm]

lib/Target/AMDGPU/SIPeepholeSDWA.cpp

478

shared_ptr should not be used anywhere

[SamWot]

Got rid of shared_ptr

[SamWot]

lib/Target/AMDGPU/SIPeepholeSDWA.cpp

726

Also, since these are only potential cases that could still fail to fold, shouldn't this be handled while folding the operands rather than eliminating them from the potential set in case they fold into other uses?

I moved check if instruction is convertible earlier to protect from failed folds.
Also, I'm eliminating SDWA operand, not potential instruction. This instruction would be able to fold with other SDWA operands

Shouldn't this be handled as part of the iteration order? If you made SDWAOperandsVector some kind of SetVector and changed the discovery order to post order you should see the final foldable uses first

Wouldn't this be even more complicated - to rely on some kind of order to ensure that one instruction folds before that another?
Also, I don't understand what order should this be. Problem is that there are both source and destination SDWA operands. While source operands folds into uses, destination operands folds backward into defs.

[SamWot]

Small fix

[arsenm]

lib/Target/AMDGPU/SIPeepholeSDWA.cpp

726

It's normal for passes to depend on program order for things like this. Multiple maps that need to be kept consistent I think is harder to follow and reason about

[SamWot]

Simplified process. Got rid of PulledOut vector.

[vpykhtin]

I agree if this is a bugfix.

lib/Target/AMDGPU/SIPeepholeSDWA.cpp

 #include "AMDGPUSubtarget.h"
 #include "SIDefines.h"
 #include "SIInstrInfo.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include <unordered_map>
+#include <unordered_set>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "si-peephole-sdwa"
 
 STATISTIC(NumSDWAPatternsFound, "Number of SDWA patterns found.");
 STATISTIC(NumSDWAInstructionsPeepholed,
           "Number of instruction converted to SDWA.");
private:
 
   std::unordered_map<MachineInstr *, std::unique_ptr<SDWAOperand>> SDWAOperands;
 
   Optional<int64_t> foldToImm(const MachineOperand &Op) const;
 
 public:
   static char ID;
 
-  typedef SmallVector<std::unique_ptr<SDWAOperand>, 4> SDWAOperandsVector;
+  typedef SmallVector<SDWAOperand *, 4> SDWAOperandsVector;
 
   SIPeepholeSDWA() : MachineFunctionPass(ID) {
     initializeSIPeepholeSDWAPass(*PassRegistry::getPassRegistry());
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override;
   void matchSDWAOperands(MachineFunction &MF);
+  bool isConvertibleToSDWA(const MachineInstr &MI) const;
   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);
   }
for (MachineInstr &MI : MBB) {
         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;
       }
for (MachineInstr &MI : MBB) {
         ++NumSDWAPatternsFound;
         break;
       }
       }
     }
   }
 }
 
-bool SIPeepholeSDWA::convertToSDWA(MachineInstr &MI,
-                                   const SDWAOperandsVector &SDWAOperands) {
+bool SIPeepholeSDWA::isConvertibleToSDWA(const MachineInstr &MI) const {
   // Check if this instruction can be converted to SDWA:
   // 1. Does this opcode support SDWA
   if (AMDGPU::getSDWAOp(MI.getOpcode()) == -1)
     return false;
 
   // 2. Are all operands - VGPRs
   for (const MachineOperand &Operand : MI.explicit_operands()) {
     if (!Operand.isReg() || !TRI->isVGPR(*MRI, Operand.getReg()))
       return false;
   }
+}
 
+bool SIPeepholeSDWA::convertToSDWA(MachineInstr &MI,
+                                   const SDWAOperandsVector &SDWAOperands) {
   // Convert to sdwa
   int SDWAOpcode = AMDGPU::getSDWAOp(MI.getOpcode());
   assert(SDWAOpcode != -1);
 
   const MCInstrDesc &SDWADesc = TII->get(SDWAOpcode);
 
   // Create SDWA version of instruction MI and initialize its operands
   MachineInstrBuilder SDWAInst =
if (!ST.hasSDWA() ||
       !AMDGPU::isVI(ST)) { // TODO: Add support for SDWA on gfx9
     return false;
   }
 
   MRI = &MF.getRegInfo();
   TRI = ST.getRegisterInfo();
   TII = ST.getInstrInfo();
 
-  std::unordered_map<MachineInstr *, SDWAOperandsVector> PotentialMatches;
   
+  // Find all SDWA operands in MF.
   matchSDWAOperands(MF);
 
-  for (auto &OperandPair : SDWAOperands) {
-    auto &Operand = OperandPair.second;
+  std::unordered_map<MachineInstr *, SDWAOperandsVector> PotentialMatches;
+
+  // There should be no intesection between SDWA operands and potential MIs
+  // e.g.:
+  // v_and_b32 v0, 0xff, v1 -> src:v1 sel:BYTE_0
+  // v_and_b32 v2, 0xff, v0 -> src:v0 sel:BYTE_0
+  // v_add_u32 v3, v4, v2
+  //
+  // In that example it is possible that we would fold 2nd instruction into 3rd
+  // (v_add_u32_sdwa) and then try to fold 1st instruction into 2nd (that was
+  // already destroyed).
+  // We keep track of every SDWA operand that should be "pulled out" - operand
+  // created from MI matched by another SDWA operand.
+  std::unordered_set<SDWAOperand *> PulledOut;
+
+  for (const auto &OperandPair : SDWAOperands) {
+    const auto &Operand = OperandPair.second;
     MachineInstr *PotentialMI = Operand->potentialToConvert(TII);
-    if (PotentialMI) {
-      PotentialMatches[PotentialMI].push_back(std::move(Operand));
+    if (PotentialMI && isConvertibleToSDWA(*PotentialMI)) {
+      if (SDWAOperands.count(PotentialMI) > 0) {
+        PulledOut.insert(SDWAOperands[PotentialMI].get());
+      }
+      PotentialMatches[PotentialMI].push_back(Operand.get());
+    }
+  }
+
+  // Remove all potential matches that should be pulled out
+  for (auto PotentialIt = PotentialMatches.begin();
+       PotentialIt != PotentialMatches.end(); ) {
+    auto &Ops = PotentialIt->second;
+    for (auto OpsIt = Ops.begin(); OpsIt != Ops.end(); ) {
+      if (PulledOut.count(*OpsIt) > 0) {
+        OpsIt = Ops.erase(OpsIt);
+      } else {
+        ++OpsIt;
+      }
+    }
+
+    if (Ops.empty()) {
+      PotentialIt = PotentialMatches.erase(PotentialIt);
+    } else {
+      ++PotentialIt;
     }
   }
 
   for (auto &PotentialPair : PotentialMatches) {
     MachineInstr &PotentialMI = *PotentialPair.first;
     convertToSDWA(PotentialMI, PotentialPair.second);
   }
 
   SDWAOperands.clear();
   return false;
 }

test/CodeGen/AMDGPU/sdwa-peephole.ll

entry:
   br label %add_label
 add_label:
   %add = add <2 x i16> %a, %b
   br label %store_label
 store_label:
   store <2 x i16> %add, <2 x i16> addrspace(1)* %out, align 4
   ret void
 }
+
+
+; Check that "pulling out" SDWA operands works correctly.
+; GCN-LABEL: {{^}}pulled_out_test:
+; NOSDWA-DAG: v_and_b32_e32 v{{[0-9]+}}, v{{[0-9]+}}, v{{[0-9]+}}
+; NOSDWA-DAG: v_lshlrev_b16_e32 v{{[0-9]+}}, 8, v{{[0-9]+}}
+; NOSDWA-DAG: v_and_b32_e32 v{{[0-9]+}}, v{{[0-9]+}}, v{{[0-9]+}}
+; NOSDWA-DAG: v_lshlrev_b16_e32 v{{[0-9]+}}, 8, v{{[0-9]+}}
+; NOSDWA: v_or_b32_e32 v{{[0-9]+}}, v{{[0-9]+}}, v{{[0-9]+}}
+; NOSDWA-NOT: v_and_b32_sdwa
+; NOSDWA-NOT: v_or_b32_sdwa
+
+; SDWA-DAG: v_and_b32_sdwa v{{[0-9]+}}, v{{[0-9]+}}, v{{[0-9]+}} dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:DWORD src1_sel:WORD_1
+; SDWA-DAG: v_lshlrev_b16_e32 v{{[0-9]+}}, 8, v{{[0-9]+}}
+; SDWA-DAG: v_and_b32_sdwa v{{[0-9]+}}, v{{[0-9]+}}, v{{[0-9]+}} dst_sel:DWORD dst_unused:UNUSED_PAD src0_sel:DWORD src1_sel:WORD_1
+; SDWA-DAG: v_lshlrev_b16_e32 v{{[0-9]+}}, 8, v{{[0-9]+}}
+; SDWA: v_or_b32_sdwa v{{[0-9]+}}, v{{[0-9]+}}, v{{[0-9]+}} dst_sel:WORD_1 dst_unused:UNUSED_PAD src0_sel:DWORD src1_sel:DWORD
+
+define amdgpu_kernel void @pulled_out_test(<8 x i8> addrspace(1)* %sourceA, <8 x i8> addrspace(1)* %destValues) {
+entry:
+  %idxprom = ashr exact i64 15, 32
+  %arrayidx = getelementptr inbounds <8 x i8>, <8 x i8> addrspace(1)* %sourceA, i64 %idxprom
+  %tmp = load <8 x i8>, <8 x i8> addrspace(1)* %arrayidx, align 8
+
+  %tmp1 = extractelement <8 x i8> %tmp, i32 0
+  %tmp2 = extractelement <8 x i8> %tmp, i32 1
+  %tmp3 = extractelement <8 x i8> %tmp, i32 2
+  %tmp4 = extractelement <8 x i8> %tmp, i32 3
+  %tmp5 = extractelement <8 x i8> %tmp, i32 4
+  %tmp6 = extractelement <8 x i8> %tmp, i32 5
+  %tmp7 = extractelement <8 x i8> %tmp, i32 6
+  %tmp8 = extractelement <8 x i8> %tmp, i32 7
+
+  %tmp9 = insertelement <2 x i8> undef, i8 %tmp1, i32 0
+  %tmp10 = insertelement <2 x i8> %tmp9, i8 %tmp2, i32 1
+  %tmp11 = insertelement <2 x i8> undef, i8 %tmp3, i32 0
+  %tmp12 = insertelement <2 x i8> %tmp11, i8 %tmp4, i32 1
+  %tmp13 = insertelement <2 x i8> undef, i8 %tmp5, i32 0
+  %tmp14 = insertelement <2 x i8> %tmp13, i8 %tmp6, i32 1
+  %tmp15 = insertelement <2 x i8> undef, i8 %tmp7, i32 0
+  %tmp16 = insertelement <2 x i8> %tmp15, i8 %tmp8, i32 1
+
+  %tmp17 = shufflevector <2 x i8> %tmp10, <2 x i8> %tmp12, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+  %tmp18 = shufflevector <2 x i8> %tmp14, <2 x i8> %tmp16, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+  %tmp19 = shufflevector <4 x i8> %tmp17, <4 x i8> %tmp18, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+  
+  %arrayidx5 = getelementptr inbounds <8 x i8>, <8 x i8> addrspace(1)* %destValues, i64 %idxprom
+  store <8 x i8> %tmp19, <8 x i8> addrspace(1)* %arrayidx5, align 8
+  ret void
+}