From a16308c2823b37721f786a911323f55b313f6d77 Mon Sep 17 00:00:00 2001
From: Philip Reames <listmail@philipreames.com>
Date: Fri, 25 Mar 2022 09:09:02 -0700
Subject: [PATCH] [SLP] Explicit track required stacksave/alloca dependency
 (try 3)

This is an extension of commit b7806c to handle one last case noticed in test changes for D118538.  Again, this is thought to be a latent bug in the existing code, though this time I have not managed to reduce tests for the original algoritthm.

The prior attempt had failed to account for this case:
  %a = alloca i8
  stacksave
  stackrestore
  store i8 0, i8* %a

If we allow '%a' to reorder into the stacksave/restore region, then the alloca will be deallocated before the use.  We will have taken a well defined program, and introduced a use-after-free bug.

There's also an inverse case where the alloca originally follows the stackrestore, and we need to prevent the reordering it above the restore.

Compile time wise, we potentially do an extra scan of the block for each alloca seen in a bundle.  This is significantly more expensive than the stacksave rooted version and is why I'd tried to avoid this in the initial patch.  There is room to optimize this (by essentially caching a "has stacksave" bit per block), but I'm leaving that to future work if it actually shows up in practice.  Since allocas in bundles should be rare in practice, I suspect we can defer the complexity for a long while.
---
 llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index e8daa02..7296dba 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -8100,11 +8100,14 @@ void BoUpSLP::BlockScheduling::calculateDependencies(ScheduleData *SD,
       }
 
       // If we have an inalloc alloca instruction, it needs to be scheduled
-      // after any preceeding stacksave.
-      if (match(BundleMember->Inst, m_Intrinsic<Intrinsic::stacksave>())) {
+      // after any preceeding stacksave.  We also need to prevent any alloca
+      // from reordering above a preceeding stackrestore.
+      if (match(BundleMember->Inst, m_Intrinsic<Intrinsic::stacksave>()) ||
+          match(BundleMember->Inst, m_Intrinsic<Intrinsic::stackrestore>())) {
         for (Instruction *I = BundleMember->Inst->getNextNode();
              I != ScheduleEnd; I = I->getNextNode()) {
-          if (match(I, m_Intrinsic<Intrinsic::stacksave>()))
+          if (match(I, m_Intrinsic<Intrinsic::stacksave>()) ||
+              match(I, m_Intrinsic<Intrinsic::stackrestore>()))
             // Any allocas past here must be control dependent on I, and I
             // must be memory dependend on BundleMember->Inst.
             break;
@@ -8117,6 +8120,21 @@ void BoUpSLP::BlockScheduling::calculateDependencies(ScheduleData *SD,
         }
       }
 
+      // In addition to the cases handle just above, we need to prevent
+      // allocas from moving below a stacksave.  The stackrestore case
+      // is currently thought to be conservatism.
+      if (isa<AllocaInst>(BundleMember->Inst)) {
+        for (Instruction *I = BundleMember->Inst->getNextNode();
+             I != ScheduleEnd; I = I->getNextNode()) {
+          if (!match(I, m_Intrinsic<Intrinsic::stacksave>()) &&
+              !match(I, m_Intrinsic<Intrinsic::stackrestore>()))
+            continue;
+
+          // Add the dependency
+          makeControlDependent(I);
+          break;
+        }
+      }
 
       // Handle the memory dependencies (if any).
       ScheduleData *DepDest = BundleMember->NextLoadStore;
-- 
2.7.4