[InstCombine] Extend stacksave/restore elimination

[nextsilicon-itay-bookstein]

Summary

Previously, InstCombine detected a pair of llvm.stacksave/stackrestore
instructions that are adjacent modulo debug instructions in order to
eliminate the llvm.stackrestore. This precludes situations where
intervening instructions (e.g. loads) preclude the llvm.stacksave and
llvm.stackrestore from becoming adjacent. This commit extends the logic
and allows for eliminating the llvm.stackrestore when the range of
instructions between them does not include any alloca or side-effect
causing instructions.

Signed-off-by: Itay Bookstein <[e-mail]>

[lebedev.ri]

Not really familiar with stack manipulation infra.

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1796–1797

Maybe add a comment that all non-intrinsic calls are considered to be side-effects?

1807

&& SS.getParent() == II.getParent()

1810

What if they are in different blocks?

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1796–1797

Generally, an explanation how exactly "side effects" come into this would be nice. "Side effect" in LLVM is defined as memory write, unwind or divergence, none of which are really related to allocas (an alloca is not a side effect).

[nextsilicon-itay-bookstein]

Added comment about call-with-side-effects, added same-BB-check

[nextsilicon-itay-bookstein]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1796–1797

I added a comment about considering all non-intrinsic calls to be side-effects.
Regarding the allocas, I classify them separately above, so I'm not 100% sure if you meant that I should add something at this line or above?

1807

Great catch, thanks.

[lebedev.ri]

Can this patch be split into two, the NFC refactoring, and the actual change?

[nextsilicon-itay-bookstein]

split from NFC refactor

[nextsilicon-itay-bookstein]

Certainly, done

[lebedev.ri]

Looks good.

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1813–1814

What about stackrestore from the same stacksave?

[nextsilicon-itay-bookstein]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1813–1814

Then it will have been eliminated when it was the subject of the enclosing switch statement, before this one was. I will add an appropriate comment.

[nextsilicon-itay-bookstein]

Improved comments

[lebedev.ri]

Seems fine.

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1813–1814

I'm mainly highlighting that it might not have already happened because of the processing order.
It would be really hard to come up with an appropriate test manually,
but i suspect it is possible.

[nextsilicon-itay-bookstein]

Improved the test to include multiple stackrestores.

[nextsilicon-itay-bookstein]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1813–1814

I see. So it's a matter of not not deferring the work to the next InstCombine iteration (or missing the opportunity if it's we're at the iteration cap). Would you prefer that I replace the comment with the check?

[nextsilicon-itay-bookstein]

Refined logic so that stackrestores of the same stacksave do
not inhibit removal in the same instcombine iteration.

[lebedev.ri]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1813–1814

So it's a matter of not not deferring the work to the next InstCombine iteration (or missing the opportunity if it's we're at the iteration cap).

Right.

Would you prefer that I replace the comment with the check?

I think that would be better, yes.

[nextsilicon-itay-bookstein]

Add comment

[nextsilicon-itay-bookstein]

add comments

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

 //===- InstCombineCalls.cpp -----------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
if (auto Op0 = dyn_cast<IntrinsicInst>(II->getArgOperand(0))) {
       // Check if every byte has common bits in Bytes and Mask.
       uint64_t C = Bytes1 & Mask1;
       if ((C & 0xFF) && (C & 0xFF00) && (C & 0xFF0000) && (C & 0xFF000000))
         return replaceInstUsesWith(*II, Op0->getArgOperand(0));
     }
     break;
   }
   case Intrinsic::stackrestore: {
-    // If the save is right next to the restore, remove the restore.  This can
-    // happen when variable allocas are DCE'd.
+    enum class ClassifyResult {
+      None,
+      Alloca,
+      StackRestore,
+      CallWithSideEffects,
+    };
+    auto Classify = [](const Instruction *I) {
+      if (isa<AllocaInst>(I))
+        return ClassifyResult::Alloca;
+
+      if (auto *CI = dyn_cast<CallInst>(I)) {
+        if (auto *II = dyn_cast<IntrinsicInst>(CI)) {
+          if (II->getIntrinsicID() == Intrinsic::stackrestore)
+            return ClassifyResult::StackRestore;
+
+          if (II->mayHaveSideEffects())
+            return ClassifyResult::CallWithSideEffects;
+        } else {
+          return ClassifyResult::CallWithSideEffects;
+        }
+      }
+
+      return ClassifyResult::None;
+    };
+
+    // If there is no intervening call/alloca between the save and the restore,
+    // remove the restore. This can happen when variable allocas are DCE'd.
     if (IntrinsicInst *SS = dyn_cast<IntrinsicInst>(II->getArgOperand(0))) {
       if (SS->getIntrinsicID() == Intrinsic::stacksave) {
-        // Skip over debug info.
-        if (SS->getNextNonDebugInstruction() == II) {
-          return eraseInstFromFunction(CI);
+        BasicBlock::iterator BI(SS);
+        bool CannotRemove = false;
+        for (++BI; &*BI != II; ++BI) {
+          if (Classify(&*BI) == ClassifyResult::None)
+            continue;
+
+          CannotRemove = true;
+          break;
         }
+
+        if (!CannotRemove)
+          return eraseInstFromFunction(CI);
       }
     }
 
     // Scan down this block to see if there is another stack restore in the
     // same block without an intervening call/alloca.
     BasicBlock::iterator BI(II);
     Instruction *TI = II->getParent()->getTerminator();
     bool CannotRemove = false;
     for (++BI; &*BI != TI; ++BI) {
-      if (isa<AllocaInst>(BI)) {
-        CannotRemove = true;
+      switch (Classify(&*BI)) {
+      case ClassifyResult::None:
         break;
-      }
-      if (CallInst *BCI = dyn_cast<CallInst>(BI)) {
-        if (auto *II2 = dyn_cast<IntrinsicInst>(BCI)) {
-          // If there is a stackrestore below this one, remove this one.
-          if (II2->getIntrinsicID() == Intrinsic::stackrestore)
-            return eraseInstFromFunction(CI);
+
+      case ClassifyResult::StackRestore:
 
-          // Bail if we cross over an intrinsic with side effects, such as
-          // llvm.stacksave, or llvm.read_register.
-          if (II2->mayHaveSideEffects()) {
-            CannotRemove = true;
-            break;
-          }
-        } else {
-          // If we found a non-intrinsic call, we can't remove the stack
-          // restore.
-          CannotRemove = true;
-          break;
-        }
-      }
-    }
+      case ClassifyResult::Alloca:
+      case ClassifyResult::CallWithSideEffects:
+        // If we found an alloca, a non-intrinsic call, or an intrinsic call
+        // with side effects (such as llvm.stacksave and llvm.read_register),
+        // we can't remove the stack restore.
+      if (CannotRemove)
 
     // If the stack restore is in a return, resume, or unwind block and if there
     // are no allocas or calls between the restore and the return, nuke the
     // restore.
     if (!CannotRemove && (isa<ReturnInst>(TI) || isa<ResumeInst>(TI)))
       return eraseInstFromFunction(CI);
     break;
   }

llvm/test/Transforms/InstCombine/stacksaverestore.ll

 ; CHECK: %i = phi i32 [ 0, %entry ], [ %i1, %loop ]
 ; CHECK: %save1 = call i8* @llvm.stacksave()
 ; CHECK: %argmem = alloca inalloca i32
 ; CHECK: store i32 0, i32* %argmem
 ; CHECK: call void @inalloca_callee(i32* {{.*}} inalloca(i32) %argmem)
 ; CHECK: call void @llvm.stackrestore(i8* %save1)
 ; CHECK: br i1 %done, label %loop, label %return
 ; CHECK: ret void
+
+define i32 @test4(i32 %m, i32* nocapture readonly %a) {
+entry:
+  br label %for.body
+
+for.body:
+  %x.012 = phi i32 [ 0, %entry ], [ %add, %for.body ]
+  %i.011 = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %0 = call i8* @llvm.stacksave()
+  %1 = load i32, i32* %a, align 4
+  %mul = mul nsw i32 %1, %m
+  %add = add nsw i32 %mul, %x.012
+  call void @llvm.stackrestore(i8* %0)
+  %inc = add nuw nsw i32 %i.011, 1
+  %exitcond.not = icmp eq i32 %inc, 100
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body
+
+for.cond.cleanup:
+  ret i32 %add
+}
+
+; CHECK-LABEL: define i32 @test4(
+; CHECK-NOT: call void @llvm.stackrestore
+; CHECK: ret i32 %add