[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

1808–1809

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

1811

What if they are in different blocks?

1822–1823

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

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1822–1823

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

1808–1809

Great catch, thanks.

1822–1823

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?

[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

1814–1815

What about stackrestore from the same stacksave?

[nextsilicon-itay-bookstein]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

1814–1815

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

1814–1815

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

1814–1815

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

1814–1815

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

auto Classify = [](const Instruction *I) {
           // Consider all non-intrinsic calls to be side effects
           return ClassifyResult::CallWithSideEffects;
         }
       }
 
       return ClassifyResult::None;
     };
 
-    // If the save is right next to the restore, remove the restore.  This can
-    // happen when variable allocas are DCE'd.
+    // If the stacksave and the stackrestore are in the same BB, and there is
+    // no intervening call, alloca, or stackrestore, 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);
+      if (SS->getIntrinsicID() == Intrinsic::stacksave &&
+          SS->getParent() == II->getParent()) {
+        BasicBlock::iterator BI(SS);
+        bool CannotRemove = false;
+        for (++BI; &*BI != II; ++BI) {
+          // Note that we don't need to skip over stackrestores from the same
+          // stacksave because they will have been eliminated by this code when
+          // they were visited.
+          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;

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