[llvm] [AMDGPU] Prevent cyclic behaviour in SIFoldOperands (PR #82099)

[via llvm-commits]

https://github.com/choikwa created https://github.com/llvm/llvm-project/pull/82099

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 

>From 6cc801800fef4bbc0f0f5711145d4097a5381ad9 Mon Sep 17 00:00:00 2001
From: Kevin Choi <kevin.choi at amd.com>
Date: Sat, 17 Feb 2024 01:34:47 -0600
Subject: [PATCH] [AMDGPU] Prevent cyclic behaviour in SIFoldOperands

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.
---
 llvm/lib/Target/AMDGPU/SIFoldOperands.cpp | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

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