This is an archive of the discontinued LLVM Phabricator instance.

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

Authored by jmmartinez on Sep 26 2022, 6:24 AM.

Diff Detail

Event Timeline

jmmartinez created this revision.Sep 26 2022, 6:24 AM
Herald added a project: Restricted Project. · View Herald TranscriptSep 26 2022, 6:24 AM
jmmartinez requested review of this revision.Sep 26 2022, 6:24 AM
Herald added a project: Restricted Project. · View Herald TranscriptSep 26 2022, 6:24 AM
arsenm added inline comments.Sep 26 2022, 6:51 AM
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 added inline comments.Sep 26 2022, 8:19 AM
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 added inline comments.Sep 26 2022, 8:22 AM
llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp
38–39

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

  • [NFC][AMDGPU][Backend] Make member varialbe a parameter
jmmartinez marked an inline comment as done.Sep 27 2022, 1:56 AM
jmmartinez added inline comments.
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 added inline comments.Sep 27 2022, 12:20 PM
llvm/lib/Target/AMDGPU/AMDGPUReleaseVGPRs.cpp
66–97

You can just use MBB.getNumber

  • Use MachineBasicBlock::getNumber instead of std::distance
arsenm added inline comments.Sep 28 2022, 9:25 AM
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 added inline comments.Sep 29 2022, 12:33 AM
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.

  • Be explicit about BitVector::reference
jmmartinez marked 4 inline comments as done.Sep 30 2022, 12:54 AM
jmmartinez added inline comments.
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 added inline comments.Sep 30 2022, 8:07 AM
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 added inline comments.Oct 3 2022, 7:56 AM
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 updated this revision to Diff 464969.Oct 4 2022, 5:03 AM

Bring back the old algorithm structure but using bitvectors instead

jmmartinez marked an inline comment as done.Oct 4 2022, 5:07 AM
jmmartinez added inline comments.
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 .

  • Propagate the vgpr-is-vmem-store using depth-first-search
jmmartinez marked 2 inline comments as done.Oct 18 2022, 12:28 AM

All remarks addressed.

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

arsenm accepted this revision.Oct 18 2022, 9:38 AM

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

This revision is now accepted and ready to land.Oct 18 2022, 9:38 AM
jpages accepted this revision.Oct 18 2022, 9:42 AM

LGTM, thanks!

West to East const.