[AMDGPU] Fold immediates in the optimizeCompareInstr

[rampitec]

Summary

Peephole works before the first SIFoldOperands so most of
the immediates are in registers.

[rampitec]

Besides return value of getFoldableImm() was inconsistent.

[rampitec]

Fix tidy warnings.

llvm/lib/Target/AMDGPU/SIInstrInfo.h

public:
 
   unsigned getAddressSpaceForPseudoSourceKind(
              unsigned Kind) const override;
 
   bool
   areMemAccessesTriviallyDisjoint(const MachineInstr &MIa,
                                   const MachineInstr &MIb) const override;
 
-  bool isFoldableCopy(const MachineInstr &MI) const;
+  static bool isFoldableCopy(const MachineInstr &MI);
 
   bool FoldImmediate(MachineInstr &UseMI, MachineInstr &DefMI, Register Reg,
                      MachineRegisterInfo *MRI) const final;
 
   unsigned getMachineCSELookAheadLimit() const override { return 500; }
 
   MachineInstr *convertToThreeAddress(MachineFunction::iterator &MBB,
                                       MachineInstr &MI,

llvm/lib/Target/AMDGPU/SIInstrInfo.cpp

for (int Idx = 0; Idx != NElts; ++Idx) {
     preserveCondRegFlags(Select->getOperand(3), Cond[1]);
     fixImplicitOperands(*Select);
 
     MIB.addReg(DstElt)
        .addImm(SubIdx);
   }
 }
 
-bool SIInstrInfo::isFoldableCopy(const MachineInstr &MI) const {
+bool SIInstrInfo::isFoldableCopy(const MachineInstr &MI) {
   switch (MI.getOpcode()) {
   case AMDGPU::V_MOV_B32_e32:
   case AMDGPU::V_MOV_B32_e64:
   case AMDGPU::V_MOV_B64_PSEUDO: {
     // If there are additional implicit register operands, this may be used for
     // register indexing so the source register operand isn't simply copied.
     unsigned NumOps = MI.getDesc().getNumOperands() +
       MI.getDesc().getNumImplicitUses();
if (isFLAT(MIb))
       return checkInstOffsetsDoNotOverlap(MIa, MIb);
 
     return false;
   }
 
   return false;
 }
 
-static int64_t getFoldableImm(const MachineOperand* MO) {
+static bool getFoldableImm(Register Reg, const MachineRegisterInfo &MRI,
+                           int64_t &Imm) {
+  if (Reg.isPhysical())
+    return false;
+  auto *Def = MRI.getUniqueVRegDef(Reg);
+  if (Def && SIInstrInfo::isFoldableCopy(*Def) && Def->getOperand(1).isImm()) {
+    Imm = Def->getOperand(1).getImm();
+    return true;
+  }
+  return false;
+}
+
+static bool getFoldableImm(const MachineOperand *MO, int64_t &Imm) {
   if (!MO->isReg())
     return false;
   const MachineFunction *MF = MO->getParent()->getParent()->getParent();
   const MachineRegisterInfo &MRI = MF->getRegInfo();
-  auto Def = MRI.getUniqueVRegDef(MO->getReg());
-  if (Def && Def->getOpcode() == AMDGPU::V_MOV_B32_e32 &&
-      Def->getOperand(1).isImm())
-    return Def->getOperand(1).getImm();
-  return AMDGPU::NoRegister;
+  return getFoldableImm(MO->getReg(), MRI, Imm);
 }
 
 static void updateLiveVariables(LiveVariables *LV, MachineInstr &MI,
                                 MachineInstr &NewMI) {
   if (LV) {
     unsigned NumOps = MI.getNumOperands();
     for (unsigned I = 1; I < NumOps; ++I) {
       MachineOperand &Op = MI.getOperand(I);
MachineInstr *SIInstrInfo::convertToThreeAddress(MachineFunction::iterator &MBB,
   const MachineOperand *Clamp = getNamedOperand(MI, AMDGPU::OpName::clamp);
   const MachineOperand *Omod = getNamedOperand(MI, AMDGPU::OpName::omod);
   MachineInstrBuilder MIB;
 
   if (!Src0Mods && !Src1Mods && !Clamp && !Omod && !IsF64 &&
       // If we have an SGPR input, we will violate the constant bus restriction.
       (ST.getConstantBusLimit(Opc) > 1 || !Src0->isReg() ||
        !RI.isSGPRReg(MBB->getParent()->getRegInfo(), Src0->getReg()))) {
-    if (auto Imm = getFoldableImm(Src2)) {
+    int64_t Imm;
+    if (getFoldableImm(Src2, Imm)) {
       unsigned NewOpc =
           IsFMA ? (IsF16 ? AMDGPU::V_FMAAK_F16 : AMDGPU::V_FMAAK_F32)
                 : (IsF16 ? AMDGPU::V_MADAK_F16 : AMDGPU::V_MADAK_F32);
       if (pseudoToMCOpcode(NewOpc) != -1) {
         MIB = BuildMI(*MBB, MI, MI.getDebugLoc(), get(NewOpc))
                   .add(*Dst)
                   .add(*Src0)
                   .add(*Src1)
                   .addImm(Imm);
         updateLiveVariables(LV, MI, *MIB);
         return MIB;
       }
     }
     unsigned NewOpc = IsFMA
                           ? (IsF16 ? AMDGPU::V_FMAMK_F16 : AMDGPU::V_FMAMK_F32)
                           : (IsF16 ? AMDGPU::V_MADMK_F16 : AMDGPU::V_MADMK_F32);
-    if (auto Imm = getFoldableImm(Src1)) {
+    if (getFoldableImm(Src1, Imm)) {
       if (pseudoToMCOpcode(NewOpc) != -1) {
         MIB = BuildMI(*MBB, MI, MI.getDebugLoc(), get(NewOpc))
                   .add(*Dst)
                   .add(*Src0)
                   .addImm(Imm)
                   .add(*Src2);
         updateLiveVariables(LV, MI, *MIB);
         return MIB;
       }
     }
-    if (auto Imm = getFoldableImm(Src0)) {
+    if (getFoldableImm(Src0, Imm)) {
       if (pseudoToMCOpcode(NewOpc) != -1 &&
           isOperandLegal(
               MI, AMDGPU::getNamedOperandIdx(NewOpc, AMDGPU::OpName::src0),
               Src1)) {
         MIB = BuildMI(*MBB, MI, MI.getDebugLoc(), get(NewOpc))
                   .add(*Dst)
                   .add(*Src1)
                   .addImm(Imm)
bool SIInstrInfo::analyzeCompare(const MachineInstr &MI, Register &SrcReg,
 
   return false;
 }
 
 bool SIInstrInfo::optimizeCompareInstr(MachineInstr &CmpInstr, Register SrcReg,
                                        Register SrcReg2, int64_t CmpMask,
                                        int64_t CmpValue,
                                        const MachineRegisterInfo *MRI) const {
-  if (SrcReg2 || SrcReg.isPhysical())
+  if (!SrcReg || SrcReg.isPhysical())
+    return false;
+
+  if (SrcReg2 && !getFoldableImm(SrcReg2, *MRI, CmpValue))
     return false;
 
   const auto optimizeCmpAnd = [&CmpInstr, SrcReg, CmpValue, MRI,
                                this](int64_t ExpectedValue,
                                      unsigned SrcSize,
                                      bool IsReversable) -> bool {
     // s_cmp_eq_u32 (s_and_b32 $src, 1), 1 => s_and_b32 $src, 1
     // s_cmp_eq_i32 (s_and_b32 $src, 1), 1 => s_and_b32 $src, 1
const auto optimizeCmpAnd = [&CmpInstr, SrcReg, CmpValue, MRI,
     MachineInstr *Def = MRI->getUniqueVRegDef(SrcReg);
     if (!Def || Def->getParent() != CmpInstr.getParent())
       return false;
 
     if (Def->getOpcode() != AMDGPU::S_AND_B32 &&
         Def->getOpcode() != AMDGPU::S_AND_B64)
       return false;
 
+    const auto isMask = [](const MachineOperand *MO) -> bool {
+      int64_t Mask;
+      if (MO->isImm())
+        Mask = MO->getImm();
+      else if (!getFoldableImm(MO, Mask))
+        return false;
+      return Mask == 1;
+    };
+
     MachineOperand *SrcOp = &Def->getOperand(1);
-    if (SrcOp->isImm() && SrcOp->getImm() == 1)
+    if (isMask(SrcOp))
       SrcOp = &Def->getOperand(2);
-    else if (!Def->getOperand(2).isImm() || Def->getOperand(2).getImm() != 1)
+    else if (isMask(&Def->getOperand(2)))
+      SrcOp = &Def->getOperand(1);
+    else
       return false;
 
     Register DefReg = Def->getOperand(0).getReg();
     if (IsReversedCC && !MRI->hasOneNonDBGUse(DefReg))
       return false;
 
     for (auto I = std::next(Def->getIterator()), E = CmpInstr.getIterator();
          I != E; ++I) {

llvm/test/CodeGen/AMDGPU/basic-branch.ll

 end:
   ret void
 }
 
 ; GCN-LABEL: {{^}}test_brcc_i1:
 ; GCN: s_load_dword [[VAL:s[0-9]+]]
 ; GCNNOOPT: s_mov_b32 [[ONE:s[0-9]+]], 1{{$}}
 ; GCNNOOPT: s_and_b32 s{{[0-9]+}}, [[VAL]], [[ONE]]
-; GCNOPT: s_and_b32 s{{[0-9]+}}, [[VAL]], 1
-; GCN: s_cmp_eq_u32
+; GCNOPT:   s_bitcmp0_b32 [[VAL]], 0
+; GCNNOOPT: s_cmp_eq_u32
 ; GCN: s_cbranch_scc1 [[END:BB[0-9]+_[0-9]+]]
 
 ; GCN: buffer_store_dword
 
 ; GCN: {{^}}[[END]]:
 ; GCN: s_endpgm
 define amdgpu_kernel void @test_brcc_i1(i32 addrspace(1)* noalias %out, i32 addrspace(1)* noalias %in, i1 %val) #0 {
   %cmp0 = icmp ne i1 %val, 0

llvm/test/CodeGen/AMDGPU/optimize-compare.mir

body:             |
 
   bb.1:
     successors: %bb.2(0x80000000)
 
   bb.2:
     S_ENDPGM 0
 
 ...
+
+---
+name:            and_1_folded_src0_cmp_eq_u32_1_folded_src2
+body:             |
+  ; GCN-LABEL: name: and_1_folded_src0_cmp_eq_u32_1_folded_src2
+  ; GCN: bb.0:
+  ; GCN:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; GCN:   [[COPY:%[0-9]+]]:sreg_32 = COPY $sgpr0
+  ; GCN:   [[S_MOV_B32_:%[0-9]+]]:sreg_32 = S_MOV_B32 1
+  ; GCN:   S_BITCMP1_B32 killed [[COPY]], 0, implicit-def $scc
+  ; GCN:   S_CBRANCH_SCC0 %bb.2, implicit $scc
+  ; GCN:   S_BRANCH %bb.1
+  ; GCN: bb.1:
+  ; GCN:   successors: %bb.2(0x80000000)
+  ; GCN: bb.2:
+  ; GCN:   S_ENDPGM 0
+  bb.0:
+    successors: %bb.1(0x40000000), %bb.2(0x40000000)
+    liveins: $sgpr0, $vgpr0_vgpr1
+
+    %0:sreg_32 = COPY $sgpr0
+    %1:sreg_32 = S_MOV_B32 1
+    %2:sreg_32 = S_AND_B32 %1, killed %0, implicit-def dead $scc
+    S_CMP_EQ_U32 killed %2:sreg_32, %1, implicit-def $scc
+    S_CBRANCH_SCC0 %bb.2, implicit $scc
+    S_BRANCH %bb.1
+
+  bb.1:
+    successors: %bb.2(0x80000000)
+
+  bb.2:
+    S_ENDPGM 0
+
+...
+
+---
+name:            and_1_folded_src1_cmp_eq_u32_1_folded_src2
+body:             |
+  ; GCN-LABEL: name: and_1_folded_src1_cmp_eq_u32_1_folded_src2
+  ; GCN: bb.0:
+  ; GCN:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; GCN:   [[COPY:%[0-9]+]]:sreg_32 = COPY $sgpr0
+  ; GCN:   [[S_MOV_B32_:%[0-9]+]]:sreg_32 = S_MOV_B32 1
+  ; GCN:   S_BITCMP1_B32 killed [[COPY]], 0, implicit-def $scc
+  ; GCN:   S_CBRANCH_SCC0 %bb.2, implicit $scc
+  ; GCN:   S_BRANCH %bb.1
+  ; GCN: bb.1:
+  ; GCN:   successors: %bb.2(0x80000000)
+  ; GCN: bb.2:
+  ; GCN:   S_ENDPGM 0
+  bb.0:
+    successors: %bb.1(0x40000000), %bb.2(0x40000000)
+    liveins: $sgpr0, $vgpr0_vgpr1
+
+    %0:sreg_32 = COPY $sgpr0
+    %1:sreg_32 = S_MOV_B32 1
+    %2:sreg_32 = S_AND_B32 killed %0, %1, implicit-def dead $scc
+    S_CMP_EQ_U32 killed %2:sreg_32, %1, implicit-def $scc
+    S_CBRANCH_SCC0 %bb.2, implicit $scc
+    S_BRANCH %bb.1
+
+  bb.1:
+    successors: %bb.2(0x80000000)
+
+  bb.2:
+    S_ENDPGM 0
+
+...
+
+---
+name:            and_1_folded_src1_cmp_eq_u64_1_folded_src2
+body:             |
+  ; GCN-LABEL: name: and_1_folded_src1_cmp_eq_u64_1_folded_src2
+  ; GCN: bb.0:
+  ; GCN:   successors: %bb.1(0x40000000), %bb.2(0x40000000)
+  ; GCN:   [[COPY:%[0-9]+]]:sreg_64 = COPY $sgpr0_sgpr1
+  ; GCN:   [[S_MOV_B64_:%[0-9]+]]:sreg_64 = S_MOV_B64 1
+  ; GCN:   S_BITCMP1_B64 killed [[COPY]], 0, implicit-def $scc
+  ; GCN:   S_CBRANCH_SCC0 %bb.2, implicit $scc
+  ; GCN:   S_BRANCH %bb.1
+  ; GCN: bb.1:
+  ; GCN:   successors: %bb.2(0x80000000)
+  ; GCN: bb.2:
+  ; GCN:   S_ENDPGM 0
+  bb.0:
+    successors: %bb.1(0x40000000), %bb.2(0x40000000)
+    liveins: $sgpr0_sgpr1, $vgpr0_vgpr1
+
+    %0:sreg_64 = COPY $sgpr0_sgpr1
+    %1:sreg_64 = S_MOV_B64 1
+    %2:sreg_64 = S_AND_B64 killed %0, %1, implicit-def dead $scc
+    S_CMP_EQ_U64 killed %2:sreg_64, %1, implicit-def $scc
+    S_CBRANCH_SCC0 %bb.2, implicit $scc
+    S_BRANCH %bb.1
+
+  bb.1:
+    successors: %bb.2(0x80000000)
+
+  bb.2:
+    S_ENDPGM 0
+
+...

llvm/test/CodeGen/AMDGPU/setcc.ll

define amdgpu_kernel void @v3i8_eq(<3 x i8> addrspace(1)* %out, <3 x i8> addrspace(1)* %ptra, <3 x i8> addrspace(1)* %ptrb) #0 {
   %cmp = icmp eq <3 x i8> %a, %b
   %ext = sext <3 x i1> %cmp to <3 x i8>
   store <3 x i8> %ext, <3 x i8> addrspace(1)* %gep.out
   ret void
 }
 
 ; Make sure we don't try to emit i1 setcc ops
 ; FUNC-LABEL: setcc-i1
-; GCN: s_and_b32 [[AND:s[0-9]+]], s{{[0-9]+}}, 1
-; GCN: s_cmp_eq_u32 [[AND]], 0
+; GCN: s_bitcmp0_b32 s{{[0-9]+}}, 0
 define amdgpu_kernel void @setcc-i1(i32 %in) #0 {
   %and = and i32 %in, 1
   %cmp = icmp eq i32 %and, 0
   br i1 %cmp, label %endif, label %if
 if:
   unreachable
 endif:
   ret void

llvm/test/CodeGen/AMDGPU/wave32.ll

one:
   ret void
 
 two:
   store half %b.val, half addrspace(1)* %r
   ret void
 }
 
 ; GCN-LABEL: {{^}}test_brcc_i1:
-; GCN:      s_cmp_eq_u32 s{{[0-9]+}}, 0
+; GCN:      s_bitcmp0_b32 s{{[0-9]+}}, 0
 ; GCN-NEXT: s_cbranch_scc1
 define amdgpu_kernel void @test_brcc_i1(i32 addrspace(1)* noalias %out, i32 addrspace(1)* noalias %in, i1 %val) #0 {
   %cmp0 = icmp ne i1 %val, 0
   br i1 %cmp0, label %store, label %end
 
 store:
   store i32 222, i32 addrspace(1)* %out
   ret void