[NVPTX] Move NVPTXPeephole after NVPTXPrologEpilogPass

[wengxt]

Summary

Offset of frame index is calculated by NVPTXPrologEpilogPass. Before
that the correct offset of stack objects cannot be obtained, which
leads to wrong offset if there are more than 2 frame objects. This patch
move NVPTXPeephole after NVPTXPrologEpilogPass. Because the frame index
is already replaced by %VRFrame in NVPTXPrologEpilogPass, we check
VRFrame register instead, and try to remove the VRFrame if there
is no usage after NVPTXPeephole pass.

Patched by Xuetian Weng.

[jingyue]

LGTM with the comments addressed.

The commit message could be clearer about what the bug was.

NVPTXPrologEpilogPass relies on frame indices to compute frame offsets. NVPTXPeephole, if run before NVPTXPrologEpilogPass, would lose frame indices too early. Therefore, we move ...

lib/Target/NVPTX/NVPTXPeephole.cpp

25

The comment here needs to be updated.

111–112

isCVTAToLocalCombination should check Prev.getOperand(2) is an immediate; otherwise, things may break.

143–148

Add a sentence explaining what this do.

lib/Target/NVPTX/NVPTXTargetMachine.cpp

213

Although you mentioned why they are in this order in your commit message, you also need to comment here because people won't usually dig into commit messages.

[wengxt]

update according to comments

[wengxt]

lib/Target/NVPTX/NVPTXPeephole.cpp

25

done

111–112

since there's no need to use frame offset here (already replaced by NVPTXPrologEpilogPass), this is changed to directly reuse the operand instead of create a new imm.

143–148

done

[jingyue]

LGTM

lib/Target/NVPTX/NVPTXPeephole.cpp

 // pointer, without this pass the instructions producing alloca'ed address will
 // be based on %SP. NVPTXLowerAlloca tends to help replace store and load on
 // this address with their .local versions, but this may introduce a lot of
 // cvta.to.local instructions. Performance can be improved if we avoid casting
 // address back and forth and directly calculate local address based on %SPL.
 // This peephole pass optimizes these cases, for example
 //
 // It will transform the following pattern
-//    %vreg0<def> = LEA_ADDRi64 <fi#0>, 4
+//    %vreg0<def> = LEA_ADDRi64 %VRFrame, 4
 //    %vreg1<def> = cvta_to_local_yes_64 %vreg0
 //
 // into
 //    %vreg1<def> = LEA_ADDRi64 %VRFrameLocal, 4
 //
 // %VRFrameLocal is the virtual register name of %SPL
 //
 //===----------------------------------------------------------------------===//
 
 #include "NVPTX.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
-#include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/Target/TargetRegisterInfo.h"
 #include "llvm/Target/TargetInstrInfo.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "nvptx-peephole"
 
 namespace llvm {
static bool isCVTAToLocalCombinationCandidate(MachineInstr &Root) {
   if (!GenericAddrDef || GenericAddrDef->getParent() != &MBB ||
       (GenericAddrDef->getOpcode() != NVPTX::LEA_ADDRi64 &&
        GenericAddrDef->getOpcode() != NVPTX::LEA_ADDRi)) {
     return false;
   }
 
   // Check the LEA_ADDRi operand is Frame index
   auto &BaseAddrOp = GenericAddrDef->getOperand(1);
-  if (BaseAddrOp.getType() == MachineOperand::MO_FrameIndex) {
+  if (BaseAddrOp.isReg() && BaseAddrOp.getReg() == NVPTX::VRFrame) {
     return true;
   }
 
   return false;
 }
 
 static void CombineCVTAToLocal(MachineInstr &Root) {
   auto &MBB = *Root.getParent();
   auto &MF = *MBB.getParent();
   const auto &MRI = MF.getRegInfo();
   const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
   auto &Prev = *MRI.getUniqueVRegDef(Root.getOperand(1).getReg());
 
-  // Get the correct offset
-  int FrameIndex = Prev.getOperand(1).getIndex();
-  int Offset = MF.getFrameInfo()->getObjectOffset(FrameIndex) +
-               Prev.getOperand(2).getImm();
-
   MachineInstrBuilder MIB =
       BuildMI(MF, Root.getDebugLoc(), TII->get(Prev.getOpcode()),
               Root.getOperand(0).getReg())
           .addReg(NVPTX::VRFrameLocal)
-          .addOperand(MachineOperand::CreateImm(Offset));
+          .addOperand(Prev.getOperand(2));
 
   MBB.insert((MachineBasicBlock::iterator)&Root, MIB);
 
   // Check if MRI has only one non dbg use, which is Root
   if (MRI.hasOneNonDBGUse(Prev.getOperand(0).getReg())) {
     Prev.eraseFromParentAndMarkDBGValuesForRemoval();
   }
   Root.eraseFromParentAndMarkDBGValuesForRemoval();
for (auto &MBB : MF) {
     while (BlockIter != MBB.end()) {
       auto &MI = *BlockIter++;
       if (isCVTAToLocalCombinationCandidate(MI)) {
         CombineCVTAToLocal(MI);
         Changed = true;
       }
     }  // Instruction
   }    // Basic Block
+
+  // Remove unnecessary %VRFrame = cvta.local %VRFrameLocal
+  const auto &MRI = MF.getRegInfo();
+  if (MRI.use_empty(NVPTX::VRFrame)) {
+    if (auto MI = MRI.getUniqueVRegDef(NVPTX::VRFrame)) {
+      MI->eraseFromParentAndMarkDBGValuesForRemoval();
+    }
+  }
+
   return Changed;
 }
 
 MachineFunctionPass *llvm::createNVPTXPeephole() { return new NVPTXPeephole(); }

lib/Target/NVPTX/NVPTXTargetMachine.cpp

bool NVPTXPassConfig::addInstSelector() {
 
   addPass(createLowerAggrCopies());
   addPass(createAllocaHoisting());
   addPass(createNVPTXISelDag(getNVPTXTargetMachine(), getOptLevel()));
 
   if (!ST.hasImageHandles())
     addPass(createNVPTXReplaceImageHandlesPass());
 
-  addPass(createNVPTXPeephole());
-
   return false;
 }
 
 void NVPTXPassConfig::addPostRegAlloc() {
   addPass(createNVPTXPrologEpilogPass(), false);
+  // NVPTXPrologEpilogPass calculates frame object offset and replace frame
+  // index with VRFrame register. NVPTXPeephole need to be run after that and
+  // will replace VRFrame with VRFrameLocal when possible.
+  addPass(createNVPTXPeephole());
 }
 
 FunctionPass *NVPTXPassConfig::createTargetRegisterAllocator(bool) {
   return nullptr; // No reg alloc
 }
 
 void NVPTXPassConfig::addFastRegAlloc(FunctionPass *RegAllocPass) {
   assert(!RegAllocPass && "NVPTX uses no regalloc!");

test/CodeGen/NVPTX/local-stack-frame.ll

define void @foo3(i32 %a) {
   %1 = bitcast [3 x i32]* %local to i32*
   %2 = getelementptr inbounds i32, i32* %1, i32 %a
   store i32 %a, i32* %2
   ret void
 }
 
 ; PTX32:        cvta.local.u32   %SP, %SPL;
 ; PTX32:        add.u32          {{%r[0-9]+}}, %SP, 0;
+; PTX32:        add.u32          {{%r[0-9]+}}, %SPL, 0;
+; PTX32:        add.u32          {{%r[0-9]+}}, %SP, 4;
+; PTX32:        add.u32          {{%r[0-9]+}}, %SPL, 4;
 ; PTX32:        st.local.u32     [{{%r[0-9]+}}], {{%r[0-9]+}}
 ; PTX32:        st.local.u32     [{{%r[0-9]+}}], {{%r[0-9]+}}
 ; PTX64:        cvta.local.u64   %SP, %SPL;
 ; PTX64:        add.u64          {{%rd[0-9]+}}, %SP, 0;
+; PTX64:        add.u64          {{%rd[0-9]+}}, %SPL, 0;
+; PTX64:        add.u64          {{%rd[0-9]+}}, %SP, 4;
+; PTX64:        add.u64          {{%rd[0-9]+}}, %SPL, 4;
 ; PTX64:        st.local.u32     [{{%rd[0-9]+}}], {{%r[0-9]+}}
 ; PTX64:        st.local.u32     [{{%rd[0-9]+}}], {{%r[0-9]+}}
 define void @foo4() {
   %A = alloca i32
   %B = alloca i32
   store i32 0, i32* %A
   store i32 0, i32* %B
   call void @bar(i32* %A)
   call void @bar(i32* %B)
   ret void
 }