[AMDGPU][Backend] Fix user-after-free in AMDGPUReleaseVGPRs::isLastVGPRUseVMEMStore

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

43

If you're going to resize to the number of blocks, do you really need to use a DenseMap at all? You could use a bit vector indexed by the block number

125

Capitalize We

[jmmartinez]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

43

It doesn't map direclty into a bit-vector since there are 3 states considered (has-a-vmem-store, does-not-have-vmem-store and not-yet-computed).

What do you think about a SmallVector<VMEMStore> with an enum class UsesVMEMStoreTy { UsesVMEMStore, DoesNotUseVMEMStore, NotYetComputed } ?

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

43

Can you avoid the NotYetComputed case if you do this in RPO

[jmmartinez]

• [NFC][AMDGPU][Backend] Make member varialbe a parameter

[jmmartinez]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

43

Way better! Thanks! I changed it to use a BitVector, but also I've encapsulated everything in a small class. The associated object is passed as a parameter to runOnMachineBasicBlock (instead of using a member variable).

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

85

You can just use MBB.getNumber

[jmmartinez]

• Use MachineBasicBlock::getNumber instead of std::distance

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

88

I'd expect this to be the post_order loop

89

Just use bool?

90

This looks like a dead variable?

[jmmartinez]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

88

This loop sets the initialization value for each MachineBasicBlock. This loop does not take into account the predecessors/successors and it can be done in any order. Maybe the name lastBlockVGPRUseIsVMEMStore is misleading? Could be renamed to something such as lastVGPRUseInstInBlockIsVMEMStore or initVGPRUseInstInBlockIsVMEMStore.

The loop after this one is the one that propagates the values from the predecessors to every block and that must be done in reverse-post-order.

89

BitVector's operator[] returns a BitVector::reference and not a bool&. BitVector::reference encapsulates a reference to a single bit.

Maybe it's worth to explicitely state the type BitVector::reference LastUseIsVMEMStore = BlockVMEMStore[MBB->getNumber()];. I think that auto looks confusing.

For context, here's the code of BitVector::reference:

// Encapsulation of a single bit.
class reference {
  
  BitWord *WordRef;
  unsigned BitPos;
  
public:
  reference(BitVector &b, unsigned Idx) {
    WordRef = &b.Bits[Idx / BITWORD_SIZE];
    BitPos = Idx % BITWORD_SIZE;
  }
  
  reference() = delete;
  reference(const reference&) = default;
  
  reference &operator=(reference t) {
    *this = bool(t);
    return *this;
  }
  
  reference& operator=(bool t) {
    if (t)
      *WordRef |= BitWord(1) << BitPos;
    else
      *WordRef &= ~(BitWord(1) << BitPos);
    return *this;
  }
  
  operator bool() const {
    return ((*WordRef) & (BitWord(1) << BitPos)) != 0;
  }
};

90

Addressed in the previous comment, since it's not a reference it's not dead. But it surely looks confusing.

[jmmartinez]

• Be explicit about BitVector::reference

[jmmartinez]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

90

Bad phrasing: reference->value : Since it's not a value , but a reference, it's not dead. But it surely looks confusing.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

88

Maybe I'm confusing my iteration orders again, but I think you can get away with the one post_order (not reverse) loop if you're looking at successors

99

Can't you just inspect the bit vector if these are precomputed without the recursive call?

[jmmartinez]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

93

There is a bug in here: if BlockVMEMStore[MBB->getNumber()] was initialized to false because the last instruction referencing a vgpr instruction is not a store, and for one of the predecessors the last instruction referencing a vgpr is a store, we may override the false value with a true which is not correct.

[jmmartinez]

Bring back the old algorithm structure but using bitvectors instead

[jmmartinez]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

88

I brought back the old algorithm with the recusrive call.

I tried several alternatives but I always found a case where the reverse post order missed propagating the values accordingly :S .

[jmmartinez]

• Propagate the vgpr-is-vmem-store using depth-first-search

[jmmartinez]

All remarks addressed.

[jmmartinez]

@arsenm Could I ask you to take a look to the latest changes ?

[arsenm]

LGTM with nit

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

89–90

east const?

105

Passing a pre-filled set to visited is a bit weird, but I think this is OK

[jpages]

LGTM, thanks!

[jmmartinez]

West to East const.

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

public:
   AMDGPUReleaseVGPRs() : MachineFunctionPass(ID) {}
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
     MachineFunctionPass::getAnalysisUsage(AU);
   }
 
   // Used to cache the result of isLastInstructionVMEMStore for each block
   using BlockVMEMStoreType = DenseMap<MachineBasicBlock *, bool>;
   BlockVMEMStoreType BlockVMEMStore;
 
   // Return true if the last instruction referencing a vgpr in this MBB
   // is a VMEM store, otherwise return false.
   // Visit previous basic blocks to find this last instruction if needed.
   // Because this pass is late in the pipeline, it is expected that the
   // last vgpr use will likely be one of vmem store, ds, exp.
   // Loads and others vgpr operations would have been
   // deleted by this point, except for complex control flow involving loops.
   // This is why we are just testing the type of instructions rather
   // than the operands.
   bool isLastVGPRUseVMEMStore(MachineBasicBlock &MBB) {
     // Use the cache to break infinite loop and save some time. Initialize to
     // false in case we have a cycle.
     BlockVMEMStoreType::iterator It;
     bool Inserted;
-    std::tie(It, Inserted) = BlockVMEMStore.insert({&MBB, false});
+    std::tie(It, Inserted) = BlockVMEMStore.try_emplace(&MBB, false);
     bool &CacheEntry = It->second;
     if (!Inserted)
       return CacheEntry;
 
     for (auto &MI : reverse(MBB.instrs())) {
       // If it's a VMEM store, a vgpr will be used, return true.
       if ((SIInstrInfo::isVMEM(MI) || SIInstrInfo::isFLAT(MI)) && MI.mayStore())
         return CacheEntry = true;
 
       // If it's referencing a VGPR but is not a VMEM store, return false.
       if (SIInstrInfo::isDS(MI) || SIInstrInfo::isEXP(MI) ||
           SIInstrInfo::isVMEM(MI) || SIInstrInfo::isFLAT(MI) ||
           SIInstrInfo::isVALU(MI))
         return CacheEntry = false;
     }
 
     // Recursive call into parent blocks. Look into predecessors if there is no
     // vgpr used in this block.
+    // The iterator is not invalidated by the recursive calls since we grew the
+    // dictionary to hold all the entries in advance
     return CacheEntry = llvm::any_of(MBB.predecessors(),
                                      [this](MachineBasicBlock *Parent) {
                                        return isLastVGPRUseVMEMStore(*Parent);
                                      });
   }
 
   bool runOnMachineBasicBlock(MachineBasicBlock &MBB) {
 
     bool Changed = false;
 
     for (auto &MI : MBB.terminators()) {
       // Look for S_ENDPGM instructions
       if (MI.getOpcode() == AMDGPU::S_ENDPGM ||
           MI.getOpcode() == AMDGPU::S_ENDPGM_SAVED) {
         // If the last instruction using a VGPR in the block is a VMEM store,
         // release VGPRs. The VGPRs release will be placed just before ending
         // the program
         if (isLastVGPRUseVMEMStore(MBB)) {
           BuildMI(MBB, MI, DebugLoc(), SII->get(AMDGPU::S_SENDMSG))
               .addImm(AMDGPU::SendMsg::ID_DEALLOC_VGPRS_GFX11Plus);
           Changed = true;
         }
       }
     }
 
     return Changed;
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override {
     Function &F = MF.getFunction();
     if (skipFunction(F) || !AMDGPU::isEntryFunctionCC(F.getCallingConv()))
       return false;
 
     // This pass only runs on GFX11+
     const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
     if (ST.getGeneration() < AMDGPUSubtarget::GFX11)
       return false;
 
     LLVM_DEBUG(dbgs() << "AMDGPUReleaseVGPRs running on " << MF.getName()
                       << "\n");
 
     SII = ST.getInstrInfo();
     TRI = ST.getRegisterInfo();
 
+    // we grow BlockVMEMStore to prevent the invalidation of its iterators
+    BlockVMEMStore.grow(MF.size());
+
     bool Changed = false;
     for (auto &MBB : MF) {
       Changed |= runOnMachineBasicBlock(MBB);
     }
 
     BlockVMEMStore.clear();
 
     return Changed;