[AMDGPU] Prevent hang in SIFoldOperands

[choikwa]

In SIFoldOperands::foldOperand, the recursion in REG_SEQUENCE handling could result in infinite loop if UseMI and RSUseMI share a common use operand, flipflopping between two instructions until stack overflows. The fix is to prevent a cycle by using static seenMI set.

@jrbyrnes @bcahoon

[llvmbot]

@llvm/pr-subscribers-backend-amdgpu

Author: choikwa (choikwa)

Changes

In SIFoldOperands::foldOperand, the recursion in REG_SEQUENCE handling could result in infinite loop if UseMI and RSUseMI share a common use operand, flipflopping between two instructions until stack overflows. The fix is to prevent a cycle by using static seenMI set.

@jrbyrnes @bcahoon

---

Full diff: https://github.com/llvm/llvm-project/pull/82099.diff

1 Files Affected:

• (modified) llvm/lib/Target/AMDGPU/SIFoldOperands.cpp (+8-3)

diff --git a/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp b/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp
index 8bf05682cbe7ea..808412809c9a77 100644
--- a/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp
+++ b/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp
@@ -15,6 +15,7 @@
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineOperand.h"
+#include <unordered_set>
 
 #define DEBUG_TYPE "si-fold-operands"
 using namespace llvm;
@@ -74,7 +75,7 @@ class SIFoldOperands : public MachineFunctionPass {
   const SIRegisterInfo *TRI;
   const GCNSubtarget *ST;
   const SIMachineFunctionInfo *MFI;
-
+  
   bool frameIndexMayFold(const MachineInstr &UseMI, int OpNo,
                          const MachineOperand &OpToFold) const;
 
@@ -772,7 +773,7 @@ void SIFoldOperands::foldOperand(
   if (UseMI->isRegSequence()) {
     Register RegSeqDstReg = UseMI->getOperand(0).getReg();
     unsigned RegSeqDstSubReg = UseMI->getOperand(UseOpIdx + 1).getImm();
-
+    static std::unordered_set<MachineInstr*> seenMI;
     for (auto &RSUse : make_early_inc_range(MRI->use_nodbg_operands(RegSeqDstReg))) {
       MachineInstr *RSUseMI = RSUse.getParent();
 
@@ -782,7 +783,11 @@ void SIFoldOperands::foldOperand(
 
       if (RSUse.getSubReg() != RegSeqDstSubReg)
         continue;
-
+      
+      if (seenMI.count(RSUseMI) != 0)
+        continue;
+      seenMI.insert(RSUseMI);
+      
       foldOperand(OpToFold, RSUseMI, RSUseMI->getOperandNo(&RSUse), FoldList,
                   CopiesToReplace);
     }

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[Pierre-vh]

Thanks, needs a testcase as well. Please add a .mir testcase, there's already a few for si-fold-operands I think

[jayfoad]

In SIFoldOperands::foldOperand, the recursion in REG_SEQUENCE handling could result in infinite loop if UseMI and RSUseMI share a common use operand, flipflopping between two instructions until stack overflows. The fix is to prevent a cycle by using static seenMI set.

What does "share a common use operand" mean? Are you saying instruction A uses the result of B, and B uses the result of A? Is there a PHI involved?

[choikwa]

An example would be

UseMI: %49:vgpr_32 = V_MUL_HI_U32_e64 %5:vgpr_32, %35.sub0:sreg_64, implicit $exec

and

RSUseMI: %77:vgpr_32 = V_MUL_HI_U32_e64 %75:vgpr_32, %35.sub0:sreg_64, implicit $exec

[jayfoad]

An example would be

UseMI: %49:vgpr_32 = V_MUL_HI_U32_e64 %5:vgpr_32, %35.sub0:sreg_64, implicit $exec

and

RSUseMI: %77:vgpr_32 = V_MUL_HI_U32_e64 %75:vgpr_32, %35.sub0:sreg_64, implicit $exec

I don't understand how this would cause infinite recursion.

We only go into the RSUseMI code if UseMI is a REG_SEQUENCE instruction.

[choikwa]

You are right, I think I was incorrectly explaining the behaviour:
The first entry point was from this:
%35:sreg_64 = REG_SEQUENCE killed %34:sreg_32, %subreg.sub0, %33:sreg_32, %subreg.sub1

And the subsequent UseMI is just RSUseMI and does not go into the isRegSequence() path. However, the cycle still exists with RSUseMI flipflopping between %77 and %49.

[jayfoad]

Please try to get an actual test case.

[choikwa]

Addressed feedback w/ latest commit.

[jayfoad]

Thanks for the test case! Now can you please try to explain what goes wrong in more detail? I took a very quick look and it seems like this loop never terminates:

    for (auto &RSUse : make_early_inc_range(MRI->use_nodbg_operands(RegSeqDstReg)))

But why? If we understand the problem better, there may be a simpler fix.

[choikwa]

So I did more digging and the search yielded some interesting findings. It turns out that the early_inc iterator (I hand-converted) was incrementing and returning previous use operand iterator after going through foldOperand call. And the reason I think has to do with tryAddToFoldList at the end of foldOperand calling commuteInstruction. There may be some weird interaction going on with commutating the operands for V_MUL instructions and iterating use operands of REG_SEQUENCE that affects the iterator increment to cause infinite loop. It doesn't look like this interaction is intentional -- disabling commutation prevents the infinite loop. Perhaps option to disable commutation in tryAddToFoldList is the answer?

[Pierre-vh]

I think it goes like this:

• defusechain_iterator works on a given MachineOperand *
• We commute operands within an instruction so that pointer stays valid, but now points to a completely different operand
• defusechain_iterator keeps going on that operand (basically iterates a different range)
• We keep commuting further and going back and forth between those two ranges

That makes sense I think, and the better fix would be to stop commuting inside that loop, or commute in bulk after the loop.

[jayfoad]

Yes that explanation makes sense, but this problem has already been solved in a different way in the main loop that calls foldOperand on each use of a register, in foldInstOperand:

  SmallVector<MachineOperand *, 4> UsesToProcess;
  for (auto &Use : MRI->use_nodbg_operands(Dst.getReg()))
    UsesToProcess.push_back(&Use);
  for (auto *U : UsesToProcess) {
    MachineInstr *UseMI = U->getParent();
    foldOperand(OpToFold, UseMI, UseMI->getOperandNo(U), FoldList,
                CopiesToReplace);
  }

So perhaps we should just copy that solution here? (I.e. copy the uses into a temporary vector, to avoid any problem with the list being mutated while we are iterating over it.)

[choikwa]

Updated by caching the uses

[jayfoad]

LGTM, thanks.

[jayfoad]

Nit: don't need braces around a single physical line.

[choikwa]

some NFC's

[piotrAMD]

Rebase issue? This was added in #82435.

[choikwa]

Thanks for catching! not sure if I ever touched that but it seems concerning if rebase did that.

[arsenm]

would be good to compact the register numbers with -run-pass=none