[NVPTX] Add VRFrame and VRFrameLocal to integer register classes

[asavonic]

Summary

These registers are used as operands for instructions that expect an
integer register, so they should be added to Int32Regs or Int64Regs
register classes. Otherwise the machine verifier emits an error for
the following LIT tests when LLVM_ENABLE_MACHINE_VERIFIER=1
environment variable is set:

*** Bad machine code: Illegal physical register for instruction ***
- function:    kernel_func
- basic block: %bb.0 entry (0x55c8903d5438)
- instruction: %3:int64regs = LEA_ADDRi64 $vrframelocal, 0
- operand 1:   $vrframelocal
$vrframelocal is not a Int64Regs register.

CodeGen/NVPTX/call-with-alloca-buffer.ll
CodeGen/NVPTX/disable-opt.ll
CodeGen/NVPTX/lower-alloca.ll
CodeGen/NVPTX/lower-args.ll
CodeGen/NVPTX/param-align.ll
CodeGen/NVPTX/reg-types.ll
DebugInfo/NVPTX/dbg-declare-alloca.ll
DebugInfo/NVPTX/dbg-value-const-byref.ll

[asavonic]

While the patch does silence the error, I don't think that assigning
the same register to different register classes (for 32- and 64-bit
triple) is the right thing to do.

Instead, I guess we can define new set of registers (VRFrame64 and
VRFrameLocal64), and use them depending on a triple. Let me know if
this is an acceptable solution, or we should do something else.

[tra]

I agree. Those registers should match the pointer size. Adding a 64-bit variant is probably the way to go.

[tra]

LGTM in general.

It would be great to add two test runs (for 32 and 64-bit modes) in test/CodeGen/NVPTX/local-stack-frame.ll with -verify-machineinstrs to make sure it does work as intended.

[asavonic]

• Added -verify-machineinstrs to local-stack-frame.ll

[asavonic]

• Updated debug-info.ll LIT test.

Not sure why it is changed though.

[tra]

In D110164#3060863, @asavonic wrote:

• Updated debug-info.ll LIT test.

Not sure why it is changed though.

The order of .reg directives is not guaranteed. The tests should've used CHECK-DAG for them. We may as well fix it now, too.

[asavonic]

• Updated debug-info.ll to use CHECK-DAG for .reg checks.

llvm/lib/Target/NVPTX/NVPTXFrameLowering.cpp

 
 void NVPTXFrameLowering::emitPrologue(MachineFunction &MF,
                                       MachineBasicBlock &MBB) const {
   if (MF.getFrameInfo().hasStackObjects()) {
     assert(&MF.front() == &MBB && "Shrink-wrapping not yet supported");
     MachineInstr *MI = &MBB.front();
     MachineRegisterInfo &MR = MF.getRegInfo();
 
+    const NVPTXRegisterInfo *NRI =
+        MF.getSubtarget<NVPTXSubtarget>().getRegisterInfo();
+
     // This instruction really occurs before first instruction
     // in the BB, so giving it no debug location.
     DebugLoc dl = DebugLoc();
 
     // Emits
     //   mov %SPL, %depot;
     //   cvta.local %SP, %SPL;
     // for local address accesses in MF.
     bool Is64Bit =
         static_cast<const NVPTXTargetMachine &>(MF.getTarget()).is64Bit();
     unsigned CvtaLocalOpcode =
         (Is64Bit ? NVPTX::cvta_local_yes_64 : NVPTX::cvta_local_yes);
     unsigned MovDepotOpcode =
         (Is64Bit ? NVPTX::MOV_DEPOT_ADDR_64 : NVPTX::MOV_DEPOT_ADDR);
-    if (!MR.use_empty(NVPTX::VRFrame)) {
+    if (!MR.use_empty(NRI->getFrameRegister(MF))) {
       // If %SP is not used, do not bother emitting "cvta.local %SP, %SPL".
       MI = BuildMI(MBB, MI, dl,
                    MF.getSubtarget().getInstrInfo()->get(CvtaLocalOpcode),
-                   NVPTX::VRFrame)
-               .addReg(NVPTX::VRFrameLocal);
+                   NRI->getFrameRegister(MF))
+               .addReg(NRI->getFrameLocalRegister(MF));
     }
     BuildMI(MBB, MI, dl, MF.getSubtarget().getInstrInfo()->get(MovDepotOpcode),
-            NVPTX::VRFrameLocal)
+            NRI->getFrameLocalRegister(MF))
         .addImm(MF.getFunctionNumber());
   }
 }
 
 StackOffset
 NVPTXFrameLowering::getFrameIndexReference(const MachineFunction &MF, int FI,
                                            Register &FrameReg) const {
   const MachineFrameInfo &MFI = MF.getFrameInfo();

llvm/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
-//    %0 = LEA_ADDRi64 %VRFrame, 4
+//    %0 = LEA_ADDRi64 %VRFrame64, 4
 //    %1 = cvta_to_local_yes_64 %0
 //
 // into
-//    %1 = LEA_ADDRi64 %VRFrameLocal, 4
+//    %1 = LEA_ADDRi64 %VRFrameLocal64, 4
 //
-// %VRFrameLocal is the virtual register name of %SPL
+// %VRFrameLocal64 is the virtual register name of %SPL
 //
 //===----------------------------------------------------------------------===//
 
 #include "NVPTX.h"
+#include "NVPTXRegisterInfo.h"
+#include "NVPTXSubtarget.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 
 using namespace llvm;
 
static bool isCVTAToLocalCombinationCandidate(MachineInstr &Root) {
 
   // Check the register operand is uniquely defined by LEA_ADDRi instruction
   if (!GenericAddrDef || GenericAddrDef->getParent() != &MBB ||
       (GenericAddrDef->getOpcode() != NVPTX::LEA_ADDRi64 &&
        GenericAddrDef->getOpcode() != NVPTX::LEA_ADDRi)) {
     return false;
   }
 
+  const NVPTXRegisterInfo *NRI =
+      MF.getSubtarget<NVPTXSubtarget>().getRegisterInfo();
+
   // Check the LEA_ADDRi operand is Frame index
   auto &BaseAddrOp = GenericAddrDef->getOperand(1);
-  if (BaseAddrOp.isReg() && BaseAddrOp.getReg() == NVPTX::VRFrame) {
+  if (BaseAddrOp.isReg() && BaseAddrOp.getReg() == NRI->getFrameRegister(MF)) {
     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());
 
+  const NVPTXRegisterInfo *NRI =
+      MF.getSubtarget<NVPTXSubtarget>().getRegisterInfo();
+
   MachineInstrBuilder MIB =
       BuildMI(MF, Root.getDebugLoc(), TII->get(Prev.getOpcode()),
               Root.getOperand(0).getReg())
-          .addReg(NVPTX::VRFrameLocal)
+          .addReg(NRI->getFrameLocalRegister(MF))
           .add(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();
   }
while (BlockIter != MBB.end()) {
       auto &MI = *BlockIter++;
       if (isCVTAToLocalCombinationCandidate(MI)) {
         CombineCVTAToLocal(MI);
         Changed = true;
       }
     }  // Instruction
   }    // Basic Block
 
+  const NVPTXRegisterInfo *NRI =
+      MF.getSubtarget<NVPTXSubtarget>().getRegisterInfo();
+
   // Remove unnecessary %VRFrame = cvta.local %VRFrameLocal
   const auto &MRI = MF.getRegInfo();
-  if (MRI.use_empty(NVPTX::VRFrame)) {
-    if (auto MI = MRI.getUniqueVRegDef(NVPTX::VRFrame)) {
+  if (MRI.use_empty(NRI->getFrameRegister(MF))) {
+    if (auto MI = MRI.getUniqueVRegDef(NRI->getFrameRegister(MF))) {
       MI->eraseFromParentAndMarkDBGValuesForRemoval();
     }
   }
 
   return Changed;
 }
 
 MachineFunctionPass *llvm::createNVPTXPeephole() { return new NVPTXPeephole(); }

llvm/lib/Target/NVPTX/NVPTXRegisterInfo.h

public:
 
   BitVector getReservedRegs(const MachineFunction &MF) const override;
 
   void eliminateFrameIndex(MachineBasicBlock::iterator MI, int SPAdj,
                            unsigned FIOperandNum,
                            RegScavenger *RS = nullptr) const override;
 
   Register getFrameRegister(const MachineFunction &MF) const override;
+  Register getFrameLocalRegister(const MachineFunction &MF) const;
 
   ManagedStringPool *getStrPool() const {
     return const_cast<ManagedStringPool *>(&ManagedStrPool);
   }
 
   const char *getName(unsigned RegNo) const {
     std::stringstream O;
     O << "reg" << RegNo;

llvm/lib/Target/NVPTX/NVPTXRegisterInfo.cpp

 //===- NVPTXRegisterInfo.cpp - NVPTX Register Information -----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the NVPTX implementation of the TargetRegisterInfo class.
 //
 //===----------------------------------------------------------------------===//
 
 #include "NVPTXRegisterInfo.h"
 #include "NVPTX.h"
 #include "NVPTXSubtarget.h"
+#include "NVPTXTargetMachine.h"
 #include "llvm/ADT/BitVector.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/MC/MachineLocation.h"
 
 using namespace llvm;
void NVPTXRegisterInfo::eliminateFrameIndex(MachineBasicBlock::iterator II,
   MachineInstr &MI = *II;
   int FrameIndex = MI.getOperand(FIOperandNum).getIndex();
 
   MachineFunction &MF = *MI.getParent()->getParent();
   int Offset = MF.getFrameInfo().getObjectOffset(FrameIndex) +
                MI.getOperand(FIOperandNum + 1).getImm();
 
   // Using I0 as the frame pointer
-  MI.getOperand(FIOperandNum).ChangeToRegister(NVPTX::VRFrame, false);
+  MI.getOperand(FIOperandNum).ChangeToRegister(getFrameRegister(MF), false);
   MI.getOperand(FIOperandNum + 1).ChangeToImmediate(Offset);
 }
 
 Register NVPTXRegisterInfo::getFrameRegister(const MachineFunction &MF) const {
-  return NVPTX::VRFrame;
+  const NVPTXTargetMachine &TM =
+      static_cast<const NVPTXTargetMachine &>(MF.getTarget());
+  return TM.is64Bit() ? NVPTX::VRFrame64 : NVPTX::VRFrame32;
+}
+
+Register
+NVPTXRegisterInfo::getFrameLocalRegister(const MachineFunction &MF) const {
+  const NVPTXTargetMachine &TM =
+      static_cast<const NVPTXTargetMachine &>(MF.getTarget());
+  return TM.is64Bit() ? NVPTX::VRFrameLocal64 : NVPTX::VRFrameLocal32;
 }

llvm/lib/Target/NVPTX/NVPTXRegisterInfo.td

 class NVPTXRegClass<list<ValueType> regTypes, int alignment, dag regList>
      : RegisterClass <"NVPTX", regTypes, alignment, regList>;
 
 //===----------------------------------------------------------------------===//
 //  Registers
 //===----------------------------------------------------------------------===//
 
 // Special Registers used as stack pointer
-def VRFrame         : NVPTXReg<"%SP">;
-def VRFrameLocal    : NVPTXReg<"%SPL">;
+def VRFrame32         : NVPTXReg<"%SP">;
+def VRFrame64         : NVPTXReg<"%SP">;
+def VRFrameLocal32    : NVPTXReg<"%SPL">;
+def VRFrameLocal64    : NVPTXReg<"%SPL">;
 
 // Special Registers used as the stack
 def VRDepot  : NVPTXReg<"%Depot">;
 
 // We use virtual registers, but define a few physical registers here to keep
 // SDAG and the MachineInstr layers happy.
 foreach i = 0...4 in {
   def P#i  : NVPTXReg<"%p"#i>;  // Predicate
foreach i = 0...31 in {
   def ENVREG#i : NVPTXReg<"%envreg"#i>;
 }
 
 //===----------------------------------------------------------------------===//
 //  Register classes
 //===----------------------------------------------------------------------===//
 def Int1Regs : NVPTXRegClass<[i1], 8, (add (sequence "P%u", 0, 4))>;
 def Int16Regs : NVPTXRegClass<[i16], 16, (add (sequence "RS%u", 0, 4))>;
-def Int32Regs : NVPTXRegClass<[i32], 32, (add (sequence "R%u", 0, 4))>;
-def Int64Regs : NVPTXRegClass<[i64], 64, (add (sequence "RL%u", 0, 4))>;
+def Int32Regs : NVPTXRegClass<[i32], 32, (add (sequence "R%u", 0, 4), VRFrame32, VRFrameLocal32)>;
+def Int64Regs : NVPTXRegClass<[i64], 64, (add (sequence "RL%u", 0, 4), VRFrame64, VRFrameLocal64)>;
 def Float16Regs : NVPTXRegClass<[f16], 16, (add (sequence "H%u", 0, 4))>;
 def Float16x2Regs : NVPTXRegClass<[v2f16], 32, (add (sequence "HH%u", 0, 4))>;
 def Float32Regs : NVPTXRegClass<[f32], 32, (add (sequence "F%u", 0, 4))>;
 def Float64Regs : NVPTXRegClass<[f64], 64, (add (sequence "FL%u", 0, 4))>;
 def Int32ArgRegs : NVPTXRegClass<[i32], 32, (add (sequence "ia%u", 0, 4))>;
 def Int64ArgRegs : NVPTXRegClass<[i64], 64, (add (sequence "la%u", 0, 4))>;
 def Float32ArgRegs : NVPTXRegClass<[f32], 32, (add (sequence "fa%u", 0, 4))>;
 def Float64ArgRegs : NVPTXRegClass<[f64], 64, (add (sequence "da%u", 0, 4))>;
 
 // Read NVPTXRegisterInfo.cpp to see how VRFrame and VRDepot are used.
-def SpecialRegs : NVPTXRegClass<[i32], 32, (add VRFrame, VRFrameLocal, VRDepot,
+def SpecialRegs : NVPTXRegClass<[i32], 32, (add VRFrame32, VRFrameLocal32, VRDepot,
                                             (sequence "ENVREG%u", 0, 31))>;

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

-; RUN: llc < %s -march=nvptx -mcpu=sm_20 | FileCheck %s --check-prefix=PTX32
-; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 | FileCheck %s --check-prefix=PTX64
+; RUN: llc < %s -march=nvptx -mcpu=sm_20 -verify-machineinstrs | FileCheck %s --check-prefix=PTX32
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_20 -verify-machineinstrs | FileCheck %s --check-prefix=PTX64
 
 ; Ensure we access the local stack properly
 
 ; PTX32:        mov.u32          %SPL, __local_depot{{[0-9]+}};
 ; PTX32:        cvta.local.u32   %SP, %SPL;
 ; PTX32:        ld.param.u32     %r{{[0-9]+}}, [foo_param_0];
 ; PTX32:        st.volatile.u32  [%SP+0], %r{{[0-9]+}};
 ; PTX64:        mov.u64          %SPL, __local_depot{{[0-9]+}};