AMDGPU: Fix out of bounds indirect indexing errors

[arsenm]

Summary

This was producing acceses to registers beyond the super
register's limits, resulting in verifier failures.

[• tstellarAMD]

LGTM.

[arsenm]

r273977

lib/Target/AMDGPU/SILowerControlFlow.cpp

void splitBlockLiveIns(const MachineBasicBlock &MBB,
                          const MachineOperand &IdxReg);
 
   void emitLoadM0FromVGPRLoop(MachineBasicBlock &LoopBB, DebugLoc DL,
                               MachineInstr *MovRel,
                               const MachineOperand &IdxReg,
                               int Offset);
 
   bool loadM0(MachineInstr &MI, MachineInstr *MovRel, int Offset = 0);
-  void computeIndirectRegAndOffset(unsigned VecReg, unsigned &Reg, int &Offset);
+  std::pair<unsigned, int> computeIndirectRegAndOffset(unsigned VecReg,
+                                                       int Offset) const;
   bool indirectSrc(MachineInstr &MI);
   bool indirectDst(MachineInstr &MI);
 
 public:
   static char ID;
 
   SILowerControlFlow() :
     MachineFunctionPass(ID), TRI(nullptr), TII(nullptr) { }

 
 /// \param @VecReg The register which holds element zero of the vector
 ///                 being addressed into.
 /// \param[out] @Reg The base register to use in the indirect addressing instruction.
 /// \param[in,out] @Offset As an input, this is the constant offset part of the
 //                         indirect Index. e.g. v0 = v[VecReg + Offset]
 //                         As an output, this is a constant value that needs
 //                         to be added to the value stored in M0.
-void SILowerControlFlow::computeIndirectRegAndOffset(unsigned VecReg,
-                                                     unsigned &Reg,
-                                                     int &Offset) {
+std::pair<unsigned, int>
+SILowerControlFlow::computeIndirectRegAndOffset(unsigned VecReg,
+                                                int Offset) const {
   unsigned SubReg = TRI->getSubReg(VecReg, AMDGPU::sub0);
   if (!SubReg)
     SubReg = VecReg;
 
+  const TargetRegisterClass *SuperRC = TRI->getPhysRegClass(VecReg);
   const TargetRegisterClass *RC = TRI->getPhysRegClass(SubReg);
-  int RegIdx = TRI->getHWRegIndex(SubReg) + Offset;
+  int NumElts = SuperRC->getSize() / RC->getSize();
 
+  int BaseRegIdx = TRI->getHWRegIndex(SubReg);
+
+  // Skip out of bounds offsets, or else we would end up using an undefined
+  // register.
+  if (Offset >= NumElts)
+    return std::make_pair(RC->getRegister(BaseRegIdx), Offset);
+
+  int RegIdx = BaseRegIdx + Offset;
   if (RegIdx < 0) {
     Offset = RegIdx;
     RegIdx = 0;
   } else {
     Offset = 0;
   }
 
-  Reg = RC->getRegister(RegIdx);
+  unsigned Reg = RC->getRegister(RegIdx);
+  return std::make_pair(Reg, Offset);
 }
 
 // Return true if a new block was inserted.
 bool SILowerControlFlow::indirectSrc(MachineInstr &MI) {
   MachineBasicBlock &MBB = *MI.getParent();
   DebugLoc DL = MI.getDebugLoc();
 
   unsigned Dst = MI.getOperand(0).getReg();
   const MachineOperand *SrcVec = TII->getNamedOperand(MI, AMDGPU::OpName::src);
   int Off = TII->getNamedOperand(MI, AMDGPU::OpName::offset)->getImm();
   unsigned Reg;
 
-  computeIndirectRegAndOffset(SrcVec->getReg(), Reg, Off);
+  std::tie(Reg, Off) = computeIndirectRegAndOffset(SrcVec->getReg(), Off);
 
   MachineInstr *MovRel =
     BuildMI(*MBB.getParent(), DL, TII->get(AMDGPU::V_MOVRELS_B32_e32), Dst)
     .addReg(Reg, getUndefRegState(SrcVec->isUndef()))
     .addReg(SrcVec->getReg(), RegState::Implicit);
 
   return loadM0(MI, MovRel, Off);
 }
 
 // Return true if a new block was inserted.
 bool SILowerControlFlow::indirectDst(MachineInstr &MI) {
   MachineBasicBlock &MBB = *MI.getParent();
   DebugLoc DL = MI.getDebugLoc();
 
   unsigned Dst = MI.getOperand(0).getReg();
   int Off = TII->getNamedOperand(MI, AMDGPU::OpName::offset)->getImm();
   MachineOperand *Val = TII->getNamedOperand(MI, AMDGPU::OpName::val);
   unsigned Reg;
 
-  computeIndirectRegAndOffset(Dst, Reg, Off);
+  std::tie(Reg, Off) = computeIndirectRegAndOffset(Dst, Off);
 
   MachineInstr *MovRel =
     BuildMI(*MBB.getParent(), DL, TII->get(AMDGPU::V_MOVRELD_B32_e32))
     .addReg(Reg, RegState::Define)
     .addReg(Val->getReg(), getUndefRegState(Val->isUndef()))
     .addReg(Dst, RegState::Implicit);
 
   return loadM0(MI, MovRel, Off);

lib/Target/AMDGPU/SIRegisterInfo.h

public:
 
   const TargetRegisterClass *getPointerRegClass(
     const MachineFunction &MF, unsigned Kind = 0) const override;
 
   void eliminateFrameIndex(MachineBasicBlock::iterator MI, int SPAdj,
                            unsigned FIOperandNum,
                            RegScavenger *RS) const override;
 
-  unsigned getHWRegIndex(unsigned Reg) const override;
+  unsigned getHWRegIndex(unsigned Reg) const override {
+    return getEncodingValue(Reg) & 0xff;
+  }
 
   /// \brief Return the 'base' register class for this register.
   /// e.g. SGPR0 => SReg_32, VGPR => VGPR_32 SGPR0_SGPR1 -> SReg_32, etc.
   const TargetRegisterClass *getPhysRegClass(unsigned Reg) const;
 
   /// \returns true if this class contains only SGPR registers
   bool isSGPRClass(const TargetRegisterClass *RC) const {
     return !hasVGPRs(RC);

lib/Target/AMDGPU/SIRegisterInfo.cpp

default: {
                 TII->get(AMDGPU::V_MOV_B32_e32), TmpReg)
                 .addImm(Offset);
         FIOp.ChangeToRegister(TmpReg, false, false, true);
       }
     }
   }
 }
 
-unsigned SIRegisterInfo::getHWRegIndex(unsigned Reg) const {
-  return getEncodingValue(Reg) & 0xff;
-}
-
 // FIXME: This is very slow. It might be worth creating a map from physreg to
 // register class.
 const TargetRegisterClass *SIRegisterInfo::getPhysRegClass(unsigned Reg) const {
   assert(!TargetRegisterInfo::isVirtualRegister(Reg));
 
   static const TargetRegisterClass *const BaseClasses[] = {
     &AMDGPU::VGPR_32RegClass,
     &AMDGPU::SReg_32RegClass,

test/CodeGen/AMDGPU/indirect-addressing-si.ll

bb:
   %tmp6 = extractelement <6 x i32> %tmp5, i32 1
   %tmp7 = bitcast <6 x float> %tmp4 to <6 x i32>
   %tmp8 = extractelement <6 x i32> %tmp7, i32 5
   store volatile i32 %tmp6, i32 addrspace(3)* undef, align 4
   store volatile i32 %tmp8, i32 addrspace(3)* undef, align 4
   ret void
 }
 
+; offset puts outside of superegister bounaries, so clamp to 1st element.
+; CHECK-LABEL: {{^}}extract_largest_inbounds_offset:
+; CHECK: buffer_load_dwordx4 v{{\[}}[[LO_ELT:[0-9]+]]:[[HI_ELT:[0-9]+]]{{\]}}
+; CHECK: s_load_dword [[IDX:s[0-9]+]]
+; CHECK: s_mov_b32 m0, [[IDX]]
+; CHECK-NEXT: v_movrels_b32_e32 [[EXTRACT:v[0-9]+]], v[[HI_ELT]]
+; CHECK: buffer_store_dword [[EXTRACT]]
+define void @extract_largest_inbounds_offset(i32 addrspace(1)* %out, <4 x i32> addrspace(1)* %in, i32 %idx) {
+entry:
+  %ld = load volatile <4 x i32>, <4  x i32> addrspace(1)* %in
+  %offset = add i32 %idx, 3
+  %value = extractelement <4 x i32> %ld, i32 %offset
+  store i32 %value, i32 addrspace(1)* %out
+  ret void
+}
+
+; CHECK-LABL: {{^}}extract_out_of_bounds_offset:
+; CHECK: buffer_load_dwordx4 v{{\[}}[[LO_ELT:[0-9]+]]:[[HI_ELT:[0-9]+]]{{\]}}
+; CHECK: s_load_dword [[IDX:s[0-9]+]]
+; CHECK: s_add_i32 m0, [[IDX]], 4
+; CHECK-NEXT: v_movrels_b32_e32 [[EXTRACT:v[0-9]+]], v[[LO_ELT]]
+; CHECK: buffer_store_dword [[EXTRACT]]
+define void @extract_out_of_bounds_offset(i32 addrspace(1)* %out, <4 x i32> addrspace(1)* %in, i32 %idx) {
+entry:
+  %ld = load volatile <4 x i32>, <4  x i32> addrspace(1)* %in
+  %offset = add i32 %idx, 4
+  %value = extractelement <4 x i32> %ld, i32 %offset
+  store i32 %value, i32 addrspace(1)* %out
+  ret void
+}
+
 declare i32 @llvm.amdgcn.workitem.id.x() #1
 
 attributes #0 = { nounwind }
 attributes #1 = { nounwind readnone }