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

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

38–39

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

139

Capitalize We

[jmmartinez]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

38–39

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

38–39

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

38–39

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

66–97

You can just use MBB.getNumber

[jmmartinez]

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

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

69

I'd expect this to be the post_order loop

101

Just use bool?

102

This looks like a dead variable?

[jmmartinez]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

69

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.

101

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;
  }
};

102

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

102

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

69

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

111

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

[jmmartinez]

llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp

105

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

69

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

70–71

east const?

86

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

 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
 #include "AMDGPUSubtarget.h"
 #include "GCNSubtarget.h"
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
 #include "SIDefines.h"
-#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineOperand.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "release-vgprs"
 
 namespace {
 
 class AMDGPUReleaseVGPRs : public MachineFunctionPass {
 public:
   static char ID;
 
-  const SIInstrInfo *SII;
-  const SIRegisterInfo *TRI;
-
   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
+  // Track if the last instruction referencing a vgpr in a MBB is a VMEM
+  // store. 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});
-    bool &CacheEntry = It->second;
-    if (!Inserted)
-      return CacheEntry;
+  class LastVGPRUseIsVMEMStore {
+    BitVector BlockVMEMStore;
 
+    static Optional<bool> lastVGPRUseIsStore(const MachineBasicBlock &MBB) {
       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 a VMEM store, a VGPR will be used, return true.
+        if ((SIInstrInfo::isVMEM(MI) || SIInstrInfo::isFLAT(MI)) &&
+            MI.mayStore())
+          return 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;
+          return false;
+      }
+      // Wait until the values are propagated from the predecessors
+      return None;
     }
 
-    // Recursive call into parent blocks. Look into predecessors if there is no
-    // vgpr used in this block.
-    return CacheEntry = llvm::any_of(MBB.predecessors(),
-                                     [this](MachineBasicBlock *Parent) {
-                                       return isLastVGPRUseVMEMStore(*Parent);
-                                     });
+  public:
+    LastVGPRUseIsVMEMStore(const MachineFunction &MF)
+        : BlockVMEMStore(MF.getNumBlockIDs()) {
+
+      df_iterator_default_set<const MachineBasicBlock *> Visited;
+      SmallVector<const MachineBasicBlock *> EndWithVMEMStoreBlocks;
+
+      for (const auto &MBB : MF) {
+        auto LastUseIsStore = lastVGPRUseIsStore(MBB);
+        if (!LastUseIsStore.has_value())
+          continue;
+
+        if (*LastUseIsStore) {
+          EndWithVMEMStoreBlocks.push_back(&MBB);
+        } else {
+          Visited.insert(&MBB);
+        }
       }
 
-  bool runOnMachineBasicBlock(MachineBasicBlock &MBB) {
+      for (const auto *MBB : EndWithVMEMStoreBlocks) {
+        for (const auto *Succ : depth_first_ext(MBB, Visited)) {
+          BlockVMEMStore[Succ->getNumber()] = true;
+        }
+      }
+    }
+
+    // Return true if the last instruction referencing a vgpr in this MBB
+    // is a VMEM store, otherwise return false.
+    bool isLastVGPRUseVMEMStore(const MachineBasicBlock &MBB) const {
+      return BlockVMEMStore[MBB.getNumber()];
+    }
+  };
+
+  static bool
+  runOnMachineBasicBlock(MachineBasicBlock &MBB, const SIInstrInfo *SII,
+                         const LastVGPRUseIsVMEMStore &BlockVMEMStore) {
 
     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)) {
+        if (BlockVMEMStore.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();
+    const SIInstrInfo *SII = ST.getInstrInfo();
+    LastVGPRUseIsVMEMStore BlockVMEMStore(MF);
 
     bool Changed = false;
     for (auto &MBB : MF) {
-      Changed |= runOnMachineBasicBlock(MBB);
+      Changed |= runOnMachineBasicBlock(MBB, SII, BlockVMEMStore);
     }
 
-    BlockVMEMStore.clear();
-
     return Changed;
   }
 };
 
 } // namespace
 
 char AMDGPUReleaseVGPRs::ID = 0;
 
 char &llvm::AMDGPUReleaseVGPRsID = AMDGPUReleaseVGPRs::ID;
 
 INITIALIZE_PASS(AMDGPUReleaseVGPRs, DEBUG_TYPE, "Release VGPRs", false, false)